Repository: vectorch-ai/ScaleLLM (language: python)

Issue:
load dtype from config
removed hardcode dtype and use the one loaded from config.


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 81e13aaa..6a007fd7 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -36,7 +36,7 @@ jobs:
     env:
         BUILD_TYPE: ${{ matrix.build_type }}
         # Tells vcpkg where binary packages are stored.
-        VCPKG_DEFAULT_BINARY_CACHE: ${{ github.workspace }}/.vcpkg_bincache
+        VCPKG_DEFAULT_BINARY_CACHE: ${{ github.workspace }}/.vcpkg/bincache
 
     steps:
     - name: Install toolkits
@@ -58,15 +58,15 @@ jobs:
 
     - name: CMake Build in the 'build' subdirectory
       run: |
-        cmake -G Ninja -S . -B build
+        cmake -G Ninja -S . -B build -DCMAKE_CUDA_ARCHITECTURES=80 -DTORCH_CUDA_ARCH_LIST=8.0
         cmake --build build --config ${{env.BUILD_TYPE}} --target all
 
     # - name: benchmark
     #   if: ${{ env.BUILD_TYPE == 'Release' }}
     #   run: ./build/src/benchmark/micro_benchmark
 
-    # - name: Test
-    #   working-directory: ${{github.workspace}}/build
-    #   # Execute tests defined by the CMake configuration.
-    #   run: ctest --output-on-failure -C ${{env.BUILD_TYPE}}
+    - name: Run unit tests
+      working-directory: ${{github.workspace}}/build
+      # Execute tests defined by the CMake configuration.
+      run: ctest --output-on-failure -C ${{env.BUILD_TYPE}}
 
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 4a9635dd..be570864 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -15,9 +15,13 @@ endif()
 
 
 # set architecture for CUDA
-set(CMAKE_CUDA_ARCHITECTURES 80 86 89 90)
-set(TORCH_CUDA_ARCH_LIST "8.0 8.6 8.9 9.0")
-message(STATUS "Using CUDA Architecture: ${CMAKE_CUDA_ARCHITECTURES}")
+if (NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+  set(CMAKE_CUDA_ARCHITECTURES 80 86 89 90)
+  set(TORCH_CUDA_ARCH_LIST "8.0 8.6 8.9 9.0")
+  message(STATUS "Using CUDA Architecture: ${CMAKE_CUDA_ARCHITECTURES}")
+else()
+  message(STATUS "Using CUDA Architecture: ${CMAKE_CUDA_ARCHITECTURES}")
+endif()
 
 # configure vcpkg
 # have to set CMAKE_TOOLCHAIN_FILE before first project call.
diff --git a/Dockerfile b/Dockerfile
index 74f7be27..c0a17f70 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -20,7 +20,7 @@ RUN apt-get update -q -y && \
 
 # install rust
 RUN curl https://sh.rustup.rs -sSf | sh -s -- -y
-ENV source "$HOME/.cargo/env"
+ENV PATH=$HOME/.cargo/bin:$PATH
 
 WORKDIR /ScaleLLM
 
diff --git a/README.md b/README.md
index d2cb8ce6..83269dad 100644
--- a/README.md
+++ b/README.md
@@ -1,14 +1,280 @@
-# ScaleLLM
+# ScaleLLM: An efficient LLM Inference solution
+[![build and test](https://github.com/vectorch-ai/ScaleLLM/actions/workflows/build.yml/badge.svg?branch=main)](https://github.com/vectorch-ai/ScaleLLM/actions/workflows/build.yml) [![GitHub Repo stars](https://img.shields.io/github/stars/vectorch-ai/ScaleLLM?style=social)](https://github.com/vectorch-ai/ScaleLLM/stargazers)
+[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 
----
-ScaleLLM is a high-performance inference system for large language models, designed for production environments.
+[![Discord](https://dcbadge.vercel.app/api/server/PKe5gvBZfn)](https://discord.gg/PKe5gvBZfn)
 
-Key features:
 
-- **High performance**:
-- **Tensor parallelism**:
-- **Customizable**:
-- **Production ready**:
+> **Warning**<br />
+> ScaleLLM is currently in the active development stage and may not yet provide the optimal level of inference efficiency. We are fully dedicated to continuously enhancing its efficiency while also adding more features.
 
 
-If you have any questions, please don't hesitate to ask in ["Discussions" forum](https://github.com/vectorch-ai/ScaleLLM/discussions) or ["Discord" chat room](https://discord.gg/PKe5gvBZfn).
+In the coming weeks, we have exciting plans to focus on [**_speculative decoding_**](https://github.com/orgs/vectorch-ai/projects/1) and [**_stateful conversation_**](https://github.com/orgs/vectorch-ai/projects/2), alongside further kernel optimizations. We appreciate your understanding and look forward to delivering an even better solution.
+
+
+## Latest News:
+* [11/2023] - First [official release](https://github.com/vectorch-ai/ScaleLLM/releases/tag/v0.0.1) with support for popular open-source models.
+
+
+## Table of contents
+
+- [Overview](#overview)
+- [Get Started](#get-started)
+  - [Docker Container](#docker-container)
+  - [Docker Compose](#docker-compose)
+- [Usage Examples](#usage-examples)
+- [Supported Models](#supported-models)
+- [Quatization](#quatization)
+- [Limitations](#limitations)
+- [Contributing](#Contributing)
+- [Acknowledgements](#acknowledgements)
+- [License](#license)
+
+
+## Overview
+
+ScaleLLM is a cutting-edge inference system engineered for large language models (LLMs), meticulously designed to meet the demands of production environments. It extends its support to a wide range of popular open-source models, including Llama2, Bloom, GPT-NeoX, and more. 
+
+## Key Features
+
+- [High Efficiency](): Excels in high-performance LLM inference, leveraging state-of-the-art techniques and technologies like [Flash Attention](https://github.com/Dao-AILab/flash-attention), [Paged Attention](https://github.com/vllm-project/vllm), [Continuous batching](https://www.anyscale.com/blog/continuous-batching-llm-inference), and more.
+- [Tensor Parallelism](): Utilizes tensor parallelism for efficient model execution.
+- [OpenAI-compatible API](): An efficient [golang](https://en.wikipedia.org/wiki/Go_(programming_language)) rest api server that compatible with OpenAI.
+- [Huggingface models](): Seamless integration with most popular [HF models](#supported-models), supporting safetensors.
+- [Customizable](): Offers flexibility for customization to meet your specific needs, and provides an easy way to add new models.
+- [Production Ready](): Engineered with production environments in mind, ScaleLLM is equipped with robust system monitoring and management features to ensure a seamless deployment experience.
+
+## Getting Started
+
+The easiest way to get started with our project is by using the official Docker images. If you don't have Docker installed, please follow the installation instructions for your platform.
+
+### Docker Installation
+
+You can download and install Docker from the official website: [Docker Installation](https://docs.docker.com/get-docker/).
+
+> **Note**<br />
+> To use GPUs, you also need to install the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html).
+
+### Docker Container
+
+Once you have Docker installed, you can run our project's Docker container using the following command:
+
+```bash
+docker run -it --gpus=all --net=host \
+  -v $HOME/.cache/huggingface/hub:/models \
+  -e HF_MODEL_ID=TheBloke/Llama-2-7B-chat-AWQ \
+  -e DEVICE=cuda:0 \
+  docker.io/vectorchai/scalellm:latest --logtostderr
+``` 
+
+This command starts the Docker container with GPU support and various configuration options.
+
+- `HF_MODEL_ID` specifies which Hugging Face model you want to run.
+- `HF_MODEL_REVISION` specifies which Hugging Face model revision you want to run. by default, it is set to `"main"`.
+- `HF_MODEL_ALLOW_PATTERN` specifies which types of files are allowed to be downloaded. by default, it is set to `"*.json,*.safetensors,*.model"`.
+- `DEVICE` specifies the device on which this model should run. by default, it is set to `"auto"`.
+- `HUGGING_FACE_HUB_TOKEN` specifies the token from [huggingface](https://huggingface.co/settings/tokens) for gated models.
+
+> **Note**<br />
+> Although ScaleLLM supports both `CPU` and `GPU`, we recommend using GPU for better performance. CPU support is mainly for debugging and testing purposes, so the performance might be sub-optimal. If you want to use CPU, please set `DEVICE=cpu` in the command.
+
+### Ports and Endpoints
+
+After running the Docker container, two ports are exposed:
+
+1. **Port 8888 for gRPC Server**:
+
+   The gRPC server is served on 0.0.0.0:8888 by default. You can use gRPC to interact with the service.
+
+2. **Port 9999 for HTTP Server**:
+
+   The simple HTTP server for instrument will be served on 0.0.0.0:9999 by default. This server provides various endpoints for managing and monitoring the service:
+
+   - Use `curl localhost:9999/health` to check the health status of the service.
+   - Use `curl localhost:9999/metrics` to export Prometheus metrics.
+   - Use `curl localhost:9999/gflags` to list all available gflags for configuration.
+   - add more to come...
+
+### Rest API Server
+
+You can also start a REST API gateway using the following command:
+
+```bash
+docker run -it --net=host \
+  docker.io/vectorchai/scalellm-gateway:latest --logtostderr
+```
+
+The REST API Server is available on `localhost:8080`. You can use REST API requests to interact with the system. Check out the [Usage Examples](#usage-examples) section for more details.
+
+### Local Chatbot UI
+
+A local Chatbot UI is also available on [localhost:3000](localhost:3000). You can start it with the following command:
+
+```bash
+docker run -it --net=host \
+  -e OPENAI_API_HOST=http://127.0.0.1:8080 \
+  -e OPENAI_API_KEY=YOUR_API_KEY \
+  docker.io/vectorchai/chatbot-ui:latest
+```
+
+## Docker Compose
+
+Using Docker Compose is the easiest way to run ScaleLLM with all the services together. If you don't have Docker Compose installed, please follow the [installation doc](https://docs.docker.com/compose/install/) for your platform.
+
+```bash
+curl https://raw.githubusercontent.com/vectorch-ai/ScaleLLM/main/scalellm.yml -sSf > scalellm_compose.yml
+HF_MODEL_ID=TheBloke/Llama-2-7B-chat-AWQ DEVICE=cuda docker compose -f ./scalellm_compose.yml up
+```
+
+you will get following running services:
+* Chatbot UI on port 3000: [localhost:3000](localhost:3000)
+* ScaleLLM gRPC server on port 8888: `localhost:8888`
+* ScaleLLM HTTP server for monitoring on port 9999: `localhost:9999`
+* ScaleLLM REST API server on port 8080: `localhost:8080`
+
+## Usage Examples
+### Chat Completions
+
+You can get chat completions with the following example:
+
+```bash
+curl http://localhost:8080/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "TheBloke/Llama-2-7B-chat-AWQ",
+    "messages": [
+      {
+        "role": "system",
+        "content": "You are a helpful assistant."
+      },
+      {
+        "role": "user",
+        "content": "Hello!"
+      }
+    ]
+  }'
+```
+
+```python
+import os
+import sys
+import openai
+
+openai.api_base = "http://localhost:8080/v1"
+
+# List available models
+print("==== Available models ====")
+models = openai.Model.list()
+
+model = "TheBloke/Llama-2-7B-chat-AWQ"
+
+completion = openai.ChatCompletion.create(
+    model=model,
+    messages=[
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Hello"},
+    ],
+    max_tokens=256,
+    stream=True,
+)
+
+print(f"==== Model: {model} ====")
+for chunk in completion:
+    content = chunk["choices"][0]["delta"].get("content")
+    if content:
+        print(content, end="")
+``` 
+
+### Completions
+
+For regular completions, you can use this example:
+
+```bash
+curl http://localhost:8080/v1/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "TheBloke/Llama-2-7B-chat-AWQ",
+    "prompt": "hello",
+    "max_tokens": 32,
+    "temperature": 0.7,
+    "stream": true
+  }'
+```
+
+```python
+import os
+import sys
+import openai
+
+openai.api_base = "http://localhost:8080/v1"
+
+# List available models
+print("==== Available models ====")
+models = openai.Model.list()
+
+model = "TheBloke/Llama-2-7B-chat-AWQ"
+
+completion = openai.Completion.create(
+    model=model,
+    prompt="hello",
+    max_tokens=256,
+    temperature=0.7,
+    stream=True,
+)
+
+print(f"==== Model: {model} ====")
+for chunk in completion:
+    content = chunk["choices"][0].get("text")
+    if content:
+        print(content, end="")
+```
+
+## Supported Models
+
+|   Models   | Tensor Parallel | Quantization | HF models examples |
+| :--------: | :-------------: | :----------: | :---------------------------:|
+|    Yi      |       Yes       |     Yes      | [01-ai/Yi-6B](https://huggingface.co/01-ai/Yi-6B), [01-ai/Yi-6B-200K](https://huggingface.co/01-ai/Yi-6B-200K), [casperhansen/yi-6b-awq](https://huggingface.co/casperhansen/yi-6b-awq), [TheBloke/Yi-34B-GPTQ](https://huggingface.co/TheBloke/Yi-34B-GPTQ) |
+|   Llama2   |       Yes       |     Yes      | [meta-llama/Llama-2-7b](https://huggingface.co/meta-llama/Llama-2-7b), [TheBloke/Llama-2-13B-chat-GPTQ](https://huggingface.co/TheBloke/Llama-2-13B-chat-GPTQ), [TheBloke/Llama-2-70B-AWQ](https://huggingface.co/TheBloke/Llama-2-70B-AWQ) |
+|   Aquila   |       Yes       |     Yes      | [BAAI/Aquila-7B](https://huggingface.co/BAAI/Aquila-7B), [BAAI/AquilaChat-7B](https://huggingface.co/BAAI/AquilaChat-7B) |
+|   Bloom    |       Yes       |     Yes      | [bigscience/bloom](https://huggingface.co/bigscience/bloom) |
+|   GPT_j    |       Yes       |     Yes      | [EleutherAI/gpt-j-6b](https://huggingface.co/EleutherAI/gpt-j-6b) |
+|  GPT_NeoX  |       Yes       |     --      | [EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b) |
+|    GPT2    |       Yes       |     --       | [gpt2](https://huggingface.co/gpt2)|
+| InternLM   |       Yes       |     Yes      | [internlm/internlm-7b](https://huggingface.co/internlm/internlm-7b) |
+|  Mistral   |       Yes       |     Yes      | [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1) |
+|    MPT     |       Yes       |     Yes      | [mosaicml/mpt-30b](https://huggingface.co/mosaicml/mpt-30b) |
+
+If your model is not included in the supported list, we are more than willing to assist you. Please feel free to create a request for adding a new model on [GitHub Issues](https://github.com/vectorch-ai/ScaleLLM/issues).
+
+## Quantization
+Quantization is a crucial process for reducing the memory footprint of models. ScaleLLM offers support for two quantization techniques: Accurate Post-Training Quantization ([APTQ](https://arxiv.org/abs/2210.17323)) and Activation-aware Weight Quantization ([AWQ](https://arxiv.org/abs/2306.00978)), with seamless integration into the following libraries: autogptq, exllama, exllamav2, and awq. 
+
+>By default, **exllamav2** is employed for GPTQ 4-bit quantization. However, you have the flexibility to choose a specific implementation by configuring the **"--qlinear_gptq_impl"** option, which allows you to select from exllama, exllamav2, or auto option.
+
+## Limitations
+
+There are several known limitations we are looking to address in the coming months, including:
+
+- Only supports Hugging Face models with [fast tokenizers](https://github.com/huggingface/tokenizers).
+
+## Contributing
+
+If you have any questions or want to contribute, please don't hesitate to ask in our ["Discussions" forum](https://github.com/vectorch-ai/ScaleLLM/discussions) or join our ["Discord" chat room](https://discord.gg/PKe5gvBZfn). We welcome your input and contributions to make ScaleLLM even better. Please follow the [Contributing.md](https://github.com/vectorch-ai/ScaleLLM/blob/main/CONTRIBUTING.md) to get started.
+
+## Acknowledgements
+The following open-source projects have been used in this project, either in their original form or modified to meet our needs:
+* [pytorch](https://github.com/pytorch/pytorch)
+* [FasterTransformer](https://github.com/NVIDIA/FasterTransformer)
+* [vllm](https://github.com/vllm-project/vllm)
+* [AutoGPTQ](https://github.com/PanQiWei/AutoGPTQ)
+* [llm-awq](https://github.com/mit-han-lab/llm-awq)
+* [flash-attn](https://github.com/Dao-AILab/flash-attention)
+* [exllama](https://github.com/turboderp/exllamav2)
+* [tokenizers](https://github.com/huggingface/tokenizers)
+* [safetensors](https://github.com/huggingface/safetensors/)
+* [sentencepiece](https://github.com/google/sentencepiece)
+* [grpc-gateway](https://github.com/grpc-ecosystem/grpc-gateway)
+* [chatbot-ui](https://github.com/mckaywrigley/chatbot-ui)
+
+## License
+This project is released under the [Apache 2.0 license](https://github.com/vectorch-ai/ScaleLLM/blob/main/LICENSE).
\ No newline at end of file
diff --git a/cmake/cargo_library.cmake b/cmake/cargo_library.cmake
index 343dffd9..1992c486 100644
--- a/cmake/cargo_library.cmake
+++ b/cmake/cargo_library.cmake
@@ -44,12 +44,10 @@ function(cargo_library)
     endif()
   endif()
 
-  if(NOT CMAKE_BUILD_TYPE)
+  if(CMAKE_BUILD_TYPE STREQUAL "Debug")
     set(LIB_BUILD_TYPE "debug")
-  elseif(${CMAKE_BUILD_TYPE} STREQUAL "Release" OR ${CMAKE_BUILD_TYPE} STREQUAL "RelWithDebInfo")
-    set(LIB_BUILD_TYPE "release")
   else()
-    set(LIB_BUILD_TYPE "debug")
+    set(LIB_BUILD_TYPE "release")
   endif()
 
   if(IOS)
diff --git a/cmake/cargo_shared_library.cmake b/cmake/cargo_shared_library.cmake
index 1a275acc..8e3fa005 100644
--- a/cmake/cargo_shared_library.cmake
+++ b/cmake/cargo_shared_library.cmake
@@ -44,12 +44,10 @@ function(cargo_shared_library)
     endif()
   endif()
 
-  if(NOT CMAKE_BUILD_TYPE)
+  if(CMAKE_BUILD_TYPE STREQUAL "Debug")
     set(LIB_BUILD_TYPE "debug")
-  elseif(${CMAKE_BUILD_TYPE} STREQUAL "Release" OR ${CMAKE_BUILD_TYPE} STREQUAL "RelWithDebInfo")
-    set(LIB_BUILD_TYPE "release")
   else()
-    set(LIB_BUILD_TYPE "debug")
+    set(LIB_BUILD_TYPE "release")
   endif()
 
   if(IOS)
diff --git a/entrypoint.sh b/entrypoint.sh
index b0480278..cda907eb 100755
--- a/entrypoint.sh
+++ b/entrypoint.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-
+DEVICE=${DEVICE:-"auto"}
 # Set default values for HF_MODEL_REVISION and HF_MODEL_ALLOW_PATTERN
 HF_MODEL_REVISION=${HF_MODEL_REVISION:-main}
 # Define allowed file patterns for config, tokenizer, and model weights
@@ -26,5 +26,7 @@ elif [ -n "$HF_MODEL_PATH" ]; then
     ARGS+=" --model_path "$HF_MODEL_PATH""
 fi
 
+ARGS+=" --device "$DEVICE""
+
 # Run the 'scalellm' command with the specified arguments
 LD_LIBRARY_PATH=/app/lib:$LD_LIBRARY_PATH /app/bin/scalellm $ARGS "$@"
diff --git a/gateway/chat_handler.go b/gateway/chat_handler.go
index 7b89180e..386bae45 100644
--- a/gateway/chat_handler.go
+++ b/gateway/chat_handler.go
@@ -33,7 +33,8 @@ func SendChatRequest(ctx context.Context, marshaler gw.Marshaler, client scalell
 	if err != nil {
 		return nil, false, err
 	}
-	return stream, *protoReq.Stream, nil
+	isStream := protoReq.Stream != nil && *protoReq.Stream
+	return stream, isStream, nil
 
 }
 
diff --git a/gateway/completion_handler.go b/gateway/completion_handler.go
index 880324fb..79f9bfa2 100644
--- a/gateway/completion_handler.go
+++ b/gateway/completion_handler.go
@@ -19,7 +19,7 @@ import (
 	gw "github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
 )
 
-func SendChateRequest(ctx context.Context, marshaler gw.Marshaler, client scalellm.CompletionClient, req *http.Request) (scalellm.Completion_CompleteClient, bool, error) {
+func SendCompleteRequest(ctx context.Context, marshaler gw.Marshaler, client scalellm.CompletionClient, req *http.Request) (scalellm.Completion_CompleteClient, bool, error) {
 	var protoReq scalellm.CompletionRequest
 	newReader, berr := utilities.IOReaderFactory(req.Body)
 	if berr != nil {
@@ -33,7 +33,8 @@ func SendChateRequest(ctx context.Context, marshaler gw.Marshaler, client scalel
 	if err != nil {
 		return nil, false, err
 	}
-	return stream, *protoReq.Stream, nil
+	isStream := protoReq.Stream != nil && *protoReq.Stream
+	return stream, isStream, nil
 
 }
 
@@ -42,7 +43,7 @@ func RegisterCompletionHandlerClient(ctx context.Context, handler *HttpHandler,
 	handler.Handle("POST", "/v1/completions", func(w http.ResponseWriter, req *http.Request) {
 		ctx, cancel := context.WithCancel(req.Context())
 		defer cancel()
-		resp, isStream, err := SendChateRequest(ctx, handler.marshaler, client, req)
+		resp, isStream, err := SendCompleteRequest(ctx, handler.marshaler, client, req)
 		if err != nil {
 			DefaultErrorHandler(ctx, handler.marshaler, w, req, err)
 			return
diff --git a/gateway/main.go b/gateway/main.go
index fee53132..e75cbebc 100644
--- a/gateway/main.go
+++ b/gateway/main.go
@@ -15,8 +15,8 @@ import (
 var (
 	// command-line options:
 	// gRPC server endpoint
-	grpcServerEndpoint = flag.String("grpc-server", "localhost:8888", "gRPC server endpoint")
-	httpServerEndpoint = flag.String("http-server", "localhost:8080", "HTTP server endpoint")
+	grpcServerEndpoint = flag.String("grpc-server", "127.0.0.1:8888", "gRPC server endpoint")
+	httpServerEndpoint = flag.String("http-server", "0.0.0.0:8080", "HTTP server endpoint")
 )
 
 func run() error {
diff --git a/scalellm.yml b/scalellm.yml
index 87266c94..95219173 100644
--- a/scalellm.yml
+++ b/scalellm.yml
@@ -5,9 +5,12 @@ services:
     image: vectorchai/scalellm:latest
     hostname: scalellm
     container_name: scalellm
-    network_mode: "host"
+    ports:
+      - 8888:8888
+      - 9999:9999
     environment:
-      - HF_MODEL_ID=${HF_MODEL_ID:-TheBloke/Llama-2-13B-chat-AWQ}
+      - HF_MODEL_ID=${HF_MODEL_ID:-TheBloke/Llama-2-7B-chat-AWQ}
+      - DEVICE=${DEVICE:-auto}
       - HF_MODEL_REVISION=${HF_MODEL_REVISION}
       - HF_MODEL_ALLOW_PATTERN=${HF_MODEL_ALLOW_PATTERN}
       - HF_MODEL_CACHE_DIR=${HF_MODEL_CACHE_DIR}
@@ -28,8 +31,9 @@ services:
     image: vectorchai/scalellm-gateway:latest
     hostname: scalellm-gateway
     container_name: scalellm-gateway
-    network_mode: "host"
-    command: --grpc-server=localhost:8888 --logtostderr
+    ports:
+      - 8080:8080
+    command: --grpc-server=scalellm:8888 --logtostderr
     depends_on:
       - scalellm
 
@@ -37,9 +41,10 @@ services:
     image: vectorchai/chatbot-ui:latest
     hostname: chatbot-ui
     container_name: chatbot-ui
-    network_mode: "host"
+    ports:
+      - 3000:3000
     environment:
-      - OPENAI_API_HOST=http://127.0.0.1:8080
+      - OPENAI_API_HOST=http://scalellm-gateway:8080
       - OPENAI_API_KEY=YOUR_API_KEY
       - DEFAULT_SYSTEM_PROMPT=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.
     depends_on:
diff --git a/scripts/start_scalellm.sh b/scripts/start_scalellm.sh
new file mode 100755
index 00000000..81cd85e1
--- /dev/null
+++ b/scripts/start_scalellm.sh
@@ -0,0 +1,30 @@
+#!/bin/bash
+DEVICE=${DEVICE:-"auto"}
+# Set default values for HF_MODEL_REVISION and HF_MODEL_ALLOW_PATTERN
+HF_MODEL_REVISION=${HF_MODEL_REVISION:-main}
+# Define allowed file patterns for config, tokenizer, and model weights
+HF_MODEL_ALLOW_PATTERN=${HF_MODEL_ALLOW_PATTERN:-"*.json,*.safetensors,*.model"}
+
+ARGS=""
+
+# Check if HF_MODEL_ID is defined; if so, download the model from the Hugging Face hub
+if [ -n "$HF_MODEL_ID" ]; then
+    echo "Downloading model from the Hugging Face hub for model id: "$HF_MODEL_ID" and revision: "$HF_MODEL_REVISION""
+
+    MODEL_PATH=$(python3 -c 'from huggingface_hub import snapshot_download; path = snapshot_download("'"$HF_MODEL_ID"'", revision="'"$HF_MODEL_REVISION"'", allow_patterns="'"$HF_MODEL_ALLOW_PATTERN"'".split(",")); print(path)')
+    # return if error
+    if [ $? -ne 0 ]; then
+        echo "Error downloading model from the Hugging Face hub for model id: "$HF_MODEL_ID" and revision: "$HF_MODEL_REVISION""
+        exit 1
+    fi
+    ARGS+=" --model_path "$MODEL_PATH" --model_id "$HF_MODEL_ID""
+elif [ -n "$HF_MODEL_PATH" ]; then
+    echo "Using model from the specified path "$HF_MODEL_PATH""
+
+    ARGS+=" --model_path "$HF_MODEL_PATH""
+fi
+
+ARGS+=" --device "$DEVICE""
+
+# Run the 'scalellm' command with the specified arguments
+$HOME/code/ScaleLLM/build/src/server/scalellm $ARGS "$@"
diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp
index 95316470..b7c9c142 100644
--- a/src/engine/engine.cpp
+++ b/src/engine/engine.cpp
@@ -1,5 +1,6 @@
 #include "engine.h"
 
+#include <boost/algorithm/string.hpp>
 #include <memory>
 
 #include "common/logging.h"
@@ -21,10 +22,33 @@ DEFINE_double(max_memory_utilization,
               "maximum memory utilization allowed, default 0.9");
 
 namespace llm {
+namespace {
+torch::ScalarType parse_dtype(const std::string& dtype_str,
+                              const torch::Device& device) {
+  if (boost::iequals(dtype_str, "half") ||
+      boost::iequals(dtype_str, "float16")) {
+    return torch::kHalf;
+  }
+  if (boost::iequals(dtype_str, "bfloat16")) {
+    return torch::kBFloat16;
+  }
+  if ((boost::iequals(dtype_str, "float") ||
+       boost::iequals(dtype_str, "float32"))) {
+    return torch::kFloat32;
+  }
+
+  if (dtype_str.empty() || boost::iequals(dtype_str, "auto")) {
+    if (device.is_cuda()) {
+      return torch::kFloat16;
+    }
+    return torch::kFloat32;
+  }
+  GCHECK(false) << "Unsupported dtype: " << dtype_str << " on device "
+                << device;
+}
+}  // namespace
 
-Engine::Engine(torch::ScalarType dtype,
-               const std::vector<torch::Device>& devices)
-    : dtype_(dtype), devices_(devices) {
+Engine::Engine(const std::vector<torch::Device>& devices) : devices_(devices) {
   CHECK_GT(devices.size(), 0) << "At least one device is required";
 
   const int32_t world_size = static_cast<int32_t>(devices.size());
@@ -37,8 +61,7 @@ Engine::Engine(torch::ScalarType dtype,
     const int32_t rank = static_cast<int32_t>(i);
     ProcessGroup* pg = world_size > 1 ? process_groups_[i].get() : nullptr;
     ParallelArgs parallel_args(rank, world_size, pg);
-    workers_.emplace_back(
-        std::make_unique<Worker>(parallel_args, dtype, devices[i]));
+    workers_.emplace_back(std::make_unique<Worker>(parallel_args, devices[i]));
   }
 }
 
@@ -63,13 +86,22 @@ bool Engine::init_model(const std::string& model_weights_path) {
   GCHECK(tokenizer_ != nullptr);
 
   args_ = model_loader->model_args();
+  dtype_ = parse_dtype(args_.dtype(), devices_[0]);
+  GLOG(INFO) << "Initializing model with dtype: " << dtype_;
+
+  if (tokenizer_->vocab_size() != args_.vocab_size()) {
+    GLOG(WARNING) << "Vocab size mismatch: tokenizer: "
+                  << tokenizer_->vocab_size()
+                  << ", model: " << args_.vocab_size();
+  }
+
   const auto& quant_args = model_loader->quant_args();
   GLOG(INFO) << "Initializing model with " << args_ << ", " << quant_args;
 
   if (workers_.size() == 1) {
     Worker* worker = workers_[0].get();
     // only one worker, call init_model in current thread
-    if (!worker->init_model(args_, quant_args)) {
+    if (!worker->init_model(dtype_, args_, quant_args)) {
       return false;
     }
     // load the weights from the checkpoint
@@ -85,7 +117,7 @@ bool Engine::init_model(const std::string& model_weights_path) {
   std::vector<folly::SemiFuture<bool>> futures;
   futures.reserve(workers_.size());
   for (auto& worker : workers_) {
-    futures.push_back(worker->init_model_async(args_, quant_args));
+    futures.push_back(worker->init_model_async(dtype_, args_, quant_args));
   }
   // wait for all futures to complete
   auto results = folly::collectAll(futures).get();
diff --git a/src/engine/engine.h b/src/engine/engine.h
index 4adc98ec..b6966cef 100644
--- a/src/engine/engine.h
+++ b/src/engine/engine.h
@@ -33,7 +33,7 @@ namespace llm {
 class Engine {
  public:
   // create an engine with the given devices
-  Engine(torch::ScalarType dtype, const std::vector<torch::Device>& devices);
+  Engine(const std::vector<torch::Device>& devices);
 
   virtual ~Engine() = default;
 
diff --git a/src/engine/worker.cpp b/src/engine/worker.cpp
index f1649a12..a17b0013 100644
--- a/src/engine/worker.cpp
+++ b/src/engine/worker.cpp
@@ -19,15 +19,15 @@
 
 namespace llm {
 
-Worker::Worker(const ParallelArgs& parallel_args,
-               torch::ScalarType dtype,
-               const torch::Device& device)
-    : parallel_args_(parallel_args), dtype_(dtype), device_(device) {}
+Worker::Worker(const ParallelArgs& parallel_args, const torch::Device& device)
+    : parallel_args_(parallel_args), device_(device) {}
 
-bool Worker::init_model(const ModelArgs& args,
+bool Worker::init_model(torch::ScalarType dtype,
+                        const ModelArgs& args,
                         const QuantizationArgs& quant_args) {
   // initialize model
   args_ = args;
+  dtype_ = dtype;
   model_ = CausalLM::create(args, quant_args, parallel_args_, dtype_, device_);
   GCHECK(model_ != nullptr) << "Failed to create model.";
   return true;
@@ -114,15 +114,19 @@ folly::SemiFuture<OutputParameters> Worker::execute_model_async(
 
 // initialize model, cache manager. async call
 folly::SemiFuture<bool> Worker::init_model_async(
+    torch::ScalarType dtype,
     const ModelArgs& args,
     const QuantizationArgs& quant_args) {
   folly::Promise<bool> promise;
   auto future = promise.getSemiFuture();
-  executor_.schedule(
-      [this, &args, &quant_args, promise = std::move(promise)]() mutable {
-        const bool success = this->init_model(args, quant_args);
-        promise.setValue(success);
-      });
+  executor_.schedule([this,
+                      dtype,
+                      &args,
+                      &quant_args,
+                      promise = std::move(promise)]() mutable {
+    const bool success = this->init_model(dtype, args, quant_args);
+    promise.setValue(success);
+  });
   return future;
 }
 
diff --git a/src/engine/worker.h b/src/engine/worker.h
index 9f594cf0..a91f57b7 100644
--- a/src/engine/worker.h
+++ b/src/engine/worker.h
@@ -32,14 +32,14 @@ struct OutputParameters {
 
 class Worker final {
  public:
-  Worker(const ParallelArgs& parallel_args,
-         torch::ScalarType dtype,
-         const torch::Device& device);
+  Worker(const ParallelArgs& parallel_args, const torch::Device& device);
 
   ~Worker() = default;
 
   // initialize model, cache manager. blocking call
-  bool init_model(const ModelArgs& args, const QuantizationArgs& quant_args);
+  bool init_model(torch::ScalarType dtype,
+                  const ModelArgs& args,
+                  const QuantizationArgs& quant_args);
 
   // Load the model weights from state_dict. blocking call
   // can be called multiple times to reload the model with different parameters
@@ -59,7 +59,8 @@ class Worker final {
                                  const SamplingParameters& sampling_params);
 
   // initialize model, cache manager. async call
-  folly::SemiFuture<bool> init_model_async(const ModelArgs& args,
+  folly::SemiFuture<bool> init_model_async(torch::ScalarType dtype,
+                                           const ModelArgs& args,
                                            const QuantizationArgs& quant_args);
 
   // Load the model weights from state_dict. async call
diff --git a/src/kernels/awq/gemm_cuda_gen.cu b/src/kernels/awq/gemm_cuda_gen.cu
index ec2e1648..ee8790f1 100644
--- a/src/kernels/awq/gemm_cuda_gen.cu
+++ b/src/kernels/awq/gemm_cuda_gen.cu
@@ -16,14 +16,6 @@ Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023}
 
 #include "dequantize.cuh"
 
-// Pack two half values.
-static inline __device__ __host__ unsigned __pack_half2(const half x,
-                                                        const half y) {
-  unsigned v0 = *((unsigned short*)&x);
-  unsigned v1 = *((unsigned short*)&y);
-  return (v1 << 16) | v0;
-}
-
 __global__ void __launch_bounds__(64)
     gemm_forward_4bit_cuda_m16n128k32(int G,
                                       int split_k_iters,
@@ -40,11 +32,7 @@ __global__ void __launch_bounds__(64)
   __shared__ half A_shared[16 * (32 + 8)];
   __shared__ half B_shared[32 * (128 + 8)];
 
-  __shared__ half scaling_factors_shared[128];
-  __shared__ half zeros_shared[128];
-
   int j_factors1 = ((OC + 128 - 1) / 128);
-  int blockIdx_x = 0;
   int blockIdx_y = blockIdx.x % ((M + 16 - 1) / 16 * j_factors1);
   int blockIdx_z = blockIdx.x / ((M + 16 - 1) / 16 * j_factors1);
 
@@ -58,7 +46,6 @@ __global__ void __launch_bounds__(64)
 
   static constexpr int row_stride_warp = 32 * 8 / 32;
   static constexpr int row_stride = 2 * 32 * 8 / 128;
-  bool ld_zero_flag = (threadIdx.y * 32 + threadIdx.x) * 8 < 128;
   // TODO: Haotian: blockIdx_y / j_factors1 in A loading to support bsz > 16
   bool ld_A_flag =
       (blockIdx_y / j_factors1 * 16 + threadIdx.y * row_stride_warp +
@@ -95,7 +82,7 @@ __global__ void __launch_bounds__(64)
                               (((int)blockIdx_y) % j_factors1) * (128) +
                               (((int)threadIdx.x) % (128 / 8)) * 8;
 
-  half* C_ptr = C + blockIdx_z * M * OC  // blockIdz.x -> split_k dim
+  half* C_ptr = C + (long long)blockIdx_z * M * OC  // blockIdz.x -> split_k dim
                 + (((int)blockIdx_y) % j_factors1) * 128 +
                 ((int)threadIdx.y) * 64 + (((int)threadIdx.x) % 4) * 2;
 
@@ -304,12 +291,8 @@ __global__ void __launch_bounds__(64)
   __shared__ half A_shared[16 * (32 + 8)];
   __shared__ half B_shared[32 * (64 + 8)];
 
-  __shared__ half scaling_factors_shared[64];
-  __shared__ half zeros_shared[64];
-
   int j_factors1 = ((OC + 64 - 1) / 64);
 
-  int blockIdx_x = 0;
   int blockIdx_y = blockIdx.x % ((M + 16 - 1) / 16 * j_factors1);
   int blockIdx_z = blockIdx.x / ((M + 16 - 1) / 16 * j_factors1);
 
@@ -323,7 +306,6 @@ __global__ void __launch_bounds__(64)
 
   static constexpr int row_stride_warp = 32 * 8 / 32;
   static constexpr int row_stride = 2 * 32 * 8 / 64;
-  bool ld_zero_flag = (threadIdx.y * 32 + threadIdx.x) * 8 < 64;
   // TODO: Haotian: blockIdx_y / j_factors1 in A loading to support bsz > 16
   bool ld_A_flag =
       (blockIdx_y / j_factors1 * 16 + threadIdx.y * row_stride_warp +
@@ -360,7 +342,7 @@ __global__ void __launch_bounds__(64)
                               (((int)blockIdx_y) % j_factors1) * (64) +
                               (((int)threadIdx.x) % (64 / 8)) * 8;
 
-  half* C_ptr = C + blockIdx_z * M * OC  // blockIdz.x -> split_k dim
+  half* C_ptr = C + (long long)blockIdx_z * M * OC  // blockIdz.x -> split_k dim
                 + (((int)blockIdx_y) % j_factors1) * 64 +
                 ((int)threadIdx.y) * 32 + (((int)threadIdx.x) % 4) * 2;
 
@@ -593,13 +575,14 @@ torch::Tensor gemm_forward_cuda(torch::Tensor _in_feats,
   if (num_out_channels % group_size != 0)
     throw std::invalid_argument("OC is not multiple of Group size");
 
+  const auto stream = at::cuda::getCurrentCUDAStream();
   if (num_out_channels % 128 == 0) {
     int j_factors1 = num_out_channels / 128 / 1;
     dim3 num_blocks((num_out_feats + 16 - 1) / 16 * j_factors1 * split_k_iters);
     // threadIdx.x: 32
     // threadIdx.y: i_factors[2] * j_factors[2]
     dim3 threads_per_block(32, 2);
-    gemm_forward_4bit_cuda_m16n128k32<<<num_blocks, threads_per_block>>>(
+    gemm_forward_4bit_cuda_m16n128k32<<<num_blocks, threads_per_block, 0, stream>>>(
         group_size,
         split_k_iters,
         in_feats,
@@ -618,7 +601,7 @@ torch::Tensor gemm_forward_cuda(torch::Tensor _in_feats,
     // threadIdx.x: 32
     // threadIdx.y: i_factors[2] * j_factors[2]
     dim3 threads_per_block(32, 2);
-    gemm_forward_4bit_cuda_m16n64k32<<<num_blocks, threads_per_block>>>(
+    gemm_forward_4bit_cuda_m16n64k32<<<num_blocks, threads_per_block, 0, stream>>>(
         group_size,
         split_k_iters,
         in_feats,
diff --git a/src/kernels/exllama/cuda_func/column_remap.cu b/src/kernels/exllama/cuda_func/column_remap.cu
index 30e4039d..76fafd51 100644
--- a/src/kernels/exllama/cuda_func/column_remap.cu
+++ b/src/kernels/exllama/cuda_func/column_remap.cu
@@ -2,6 +2,7 @@
 
 #include "column_remap.cuh"
 #include "../util.cuh"
+#include <c10/cuda/CUDAStream.h>
 
 const int SHUF_BLOCKSIZE_X = 256;
 const int SHUF_BLOCKSIZE_Y = 16;
@@ -59,5 +60,5 @@ void column_remap_cuda
         1
     );
 
-    column_remap_kernel<<<blocks, threads>>>(x, x_new, x_width, x_height, x_map);
+    column_remap_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(x, x_new, x_width, x_height, x_map);
 }
diff --git a/src/kernels/exllama/cuda_func/q4_matmul.cu b/src/kernels/exllama/cuda_func/q4_matmul.cu
index 0ee6e16d..cf7002c5 100644
--- a/src/kernels/exllama/cuda_func/q4_matmul.cu
+++ b/src/kernels/exllama/cuda_func/q4_matmul.cu
@@ -10,6 +10,8 @@
 #include "../hip_compat.cuh"
 #endif
 
+#include <c10/cuda/CUDAStream.h>
+
 const int THREADS_X = 32;       // Block size and thread count along columns in w and out
 const int THREADS_Y = 1;        // Block size and thread count along rows in x and out
 
diff --git a/src/kernels/exllama/cuda_func/q4_matrix.cu b/src/kernels/exllama/cuda_func/q4_matrix.cu
index 2ad49bc3..2ae2bb01 100644
--- a/src/kernels/exllama/cuda_func/q4_matrix.cu
+++ b/src/kernels/exllama/cuda_func/q4_matrix.cu
@@ -4,6 +4,7 @@
 #include <vector>
 #include "../util.cuh"
 #include "../matrix.cuh"
+#include <c10/cuda/CUDAStream.h>
 
 using namespace std;
 
@@ -153,7 +154,7 @@ void Q4Matrix::make_sequential(const uint32_t* cpu_g_idx)
         1
     );
 
-    make_sequential_kernel<<<blocks, threads>>>(cuda_qweight, cuda_new_qweight, cuda_x_map, height / 8, width);
+    make_sequential_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(cuda_qweight, cuda_new_qweight, cuda_x_map, height / 8, width);
 
     // Replace qweights
 
@@ -221,5 +222,5 @@ void Q4Matrix::reconstruct(half* out)
         1
     );
 
-    reconstruct_kernel<<<blocks, threads>>>(cuda_qweight, out, cuda_scales, cuda_qzeros, height / 8, width, groupsize);
+    reconstruct_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(cuda_qweight, out, cuda_scales, cuda_qzeros, height / 8, width, groupsize);
 }
diff --git a/src/kernels/exllamav2/cuda/q_gemm.cu b/src/kernels/exllamav2/cuda/q_gemm.cu
index 0ffd2063..f6dc83c7 100644
--- a/src/kernels/exllamav2/cuda/q_gemm.cu
+++ b/src/kernels/exllamav2/cuda/q_gemm.cu
@@ -3,6 +3,8 @@
 #include "matrix_view.cuh"
 #include "../config.h"
 
+#include <c10/cuda/CUDAStream.h>
+
 #include "quant/qdq_2.cuh"
 #include "quant/qdq_3.cuh"
 #include "quant/qdq_4.cuh"
@@ -59,6 +61,7 @@ void gemm_half_q_half_cuda_part
     bool clear
 )
 {
+    const auto stream = at::cuda::getCurrentCUDAStream();
     if (!b->is_gptq)
     {
         dim3 blockDim, gridDim;
@@ -71,7 +74,7 @@ void gemm_half_q_half_cuda_part
 
         fp_gemm_half_q_half_kernel kernel = pick_gemm_half_q_half_kernel(true, m_count);
 
-        kernel<<<gridDim, blockDim>>>
+        kernel<<<gridDim, blockDim, 0, stream>>>
         (
             a,
             b->cuda_q_weight,
@@ -109,7 +112,7 @@ void gemm_half_q_half_cuda_part
 //         DBGI(b->rows_4);
 //         DBGI(b->height);
 
-        kernel<<<gridDim, blockDim>>>
+        kernel<<<gridDim, blockDim, 0, stream>>>
         (
             a,
             b->cuda_q_weight,
diff --git a/src/kernels/exllamav2/cuda/q_matrix.cu b/src/kernels/exllamav2/cuda/q_matrix.cu
index e166d8e9..0335135c 100644
--- a/src/kernels/exllamav2/cuda/q_matrix.cu
+++ b/src/kernels/exllamav2/cuda/q_matrix.cu
@@ -2,6 +2,8 @@
 #include "matrix_view.cuh"
 #include "util.cuh"
 
+#include <c10/cuda/CUDAStream.h>
+
 #include "quant/qdq_2.cuh"
 #include "quant/qdq_3.cuh"
 #include "quant/qdq_4.cuh"
@@ -136,7 +138,7 @@ QMatrix::QMatrix
     gridDim.x = DIVIDE(width, THREADS_X);
     gridDim.y = 1;
 
-    shuffle_kernel<<<gridDim, blockDim>>>(cuda_q_weight, height, width, rows_8, rows_6, rows_5, rows_4, rows_3, rows_2);
+    shuffle_kernel<<<gridDim, blockDim, 0, at::cuda::getCurrentCUDAStream()>>>(cuda_q_weight, height, width, rows_8, rows_6, rows_5, rows_4, rows_3, rows_2);
 }
 
 
@@ -442,7 +444,7 @@ void QMatrix::reconstruct(half* out)
 
     if (!is_gptq)
     {
-        reconstruct_kernel<<<gridDim, blockDim>>>
+        reconstruct_kernel<<<gridDim, blockDim, 0, at::cuda::getCurrentCUDAStream()>>>
         (
             cuda_q_weight,
             cuda_q_perm,
@@ -464,7 +466,7 @@ void QMatrix::reconstruct(half* out)
     }
     else
     {
-        reconstruct_gptq_kernel<<<gridDim, blockDim>>>
+        reconstruct_gptq_kernel<<<gridDim, blockDim, 0, at::cuda::getCurrentCUDAStream()>>>
         (
             cuda_q_weight,
             cuda_q_perm,
@@ -579,7 +581,7 @@ void QMatrix::make_sequential(const uint32_t* cpu_g_idx)
     gridDim.x = DIVIDE(width, THREADS_X);
     gridDim.y = height / 8;
 
-    make_sequential_kernel<<<gridDim, blockDim>>>
+    make_sequential_kernel<<<gridDim, blockDim, 0, at::cuda::getCurrentCUDAStream()>>>
     (
         cuda_q_weight,
         cuda_new_qweight,
diff --git a/src/kernels/gptq/gptq_kernel.cu b/src/kernels/gptq/gptq_kernel.cu
index f1ee4a50..464ce4ef 100644
--- a/src/kernels/gptq/gptq_kernel.cu
+++ b/src/kernels/gptq/gptq_kernel.cu
@@ -1,11 +1,13 @@
 // adapted from
 // https://github.com/PanQiWei/AutoGPTQ/blob/main/autogptq_extension/cuda_64/autogptq_cuda_kernel_64.cu
 
+#include <c10/cuda/CUDAStream.h>
 #include <cuda.h>
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 #include <torch/all.h>
 #include <torch/torch.h>
+
 #include "common/logging.h"
 
 namespace llm {
@@ -93,17 +95,18 @@ void vec_quant_matmul_launch_kernel(torch::Tensor vec,
   AT_DISPATCH_FLOATING_TYPES(
       vec.scalar_type(), "vec_quant_matmul_kernel", ([&] {
         vec_quant_matmul_kernel<scalar_t, BLOCK_HEIGHT, BLOCK_WIDTH, BITS>
-            <<<blocks, threads>>>(vec.data_ptr<scalar_t>(),
-                                  mat.data_ptr<int>(),
-                                  mul.data_ptr<scalar_t>(),
-                                  scales.data_ptr<scalar_t>(),
-                                  zeros.data_ptr<int>(),
-                                  g_idx.data_ptr<int>(),
-                                  batch,
-                                  vec_height,
-                                  height,
-                                  width,
-                                  zero_width);
+            <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
+                vec.data_ptr<scalar_t>(),
+                mat.data_ptr<int>(),
+                mul.data_ptr<scalar_t>(),
+                scales.data_ptr<scalar_t>(),
+                zeros.data_ptr<int>(),
+                g_idx.data_ptr<int>(),
+                batch,
+                vec_height,
+                height,
+                width,
+                zero_width);
       }));
 }
 
@@ -244,17 +247,18 @@ void vec_quant3_matmul_launch_kernel(torch::Tensor vec,
   AT_DISPATCH_FLOATING_TYPES(
       vec.scalar_type(), "vec_quant3_matmul_kernel", ([&] {
         vec_quant3_matmul_kernel<scalar_t, BLOCK_HEIGHT, BLOCK_WIDTH>
-            <<<blocks, threads>>>(vec.data_ptr<scalar_t>(),
-                                  mat.data_ptr<int>(),
-                                  mul.data_ptr<scalar_t>(),
-                                  scales.data_ptr<scalar_t>(),
-                                  zeros.data_ptr<int>(),
-                                  g_idx.data_ptr<int>(),
-                                  batch,
-                                  vec_height,
-                                  height,
-                                  width,
-                                  zero_width);
+            <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
+                vec.data_ptr<scalar_t>(),
+                mat.data_ptr<int>(),
+                mul.data_ptr<scalar_t>(),
+                scales.data_ptr<scalar_t>(),
+                zeros.data_ptr<int>(),
+                g_idx.data_ptr<int>(),
+                batch,
+                vec_height,
+                height,
+                width,
+                zero_width);
       }));
 }
 
diff --git a/src/kernels/gptq/gptq_kernel_64.cu b/src/kernels/gptq/gptq_kernel_64.cu
deleted file mode 100644
index a32e213a..00000000
--- a/src/kernels/gptq/gptq_kernel_64.cu
+++ /dev/null
@@ -1,554 +0,0 @@
-// adapted from
-// https://github.com/PanQiWei/AutoGPTQ/blob/main/autogptq_extension/cuda_64/autogptq_cuda_kernel_64.cu
-
-#include <cuda.h>
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-#include <torch/all.h>
-#include <torch/torch.h>
-
-// atomicAdd for double-precision floating-point numbers on hardware with
-// compute capability < 6.0 from:
-// https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#atomic-functions
-// #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 600
-// __device__ double atomicAdd(
-//     double* address,
-//     double val
-// ) {
-//   unsigned long long int* address_as_ull = (unsigned long long int*)address;
-//   unsigned long long int old = *address_as_ull, assumed;
-//
-//   do {
-//     assumed = old;
-//     old = atomicCAS(
-//       address_as_ull,
-//       assumed,
-//       __double_as_longlong(val + __longlong_as_double(assumed))
-//     );
-//
-//   // Note: uses integer comparison to avoid hang in case of NaN (since NaN !=
-//   NaN) } while (assumed != old);
-//
-//   return __longlong_as_double(old);
-// }
-// #endif
-
-#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 700) || defined(USE_ROCM)
-// adapted from
-// https://github.com/torch/cutorch/blob/master/lib/THC/THCAtomics.cuh
-__device__ __forceinline__ void atomicAdd(c10::Half* address, c10::Half val) {
-  unsigned int* address_as_ui =
-      reinterpret_cast<unsigned int*>(reinterpret_cast<char*>(address) -
-                                      (reinterpret_cast<size_t>(address) & 2));
-  unsigned int old = *address_as_ui;
-  unsigned int assumed;
-
-  do {
-    assumed = old;
-    unsigned short hsum =
-        reinterpret_cast<size_t>(address) & 2 ? (old >> 16) : (old & 0xffff);
-    hsum += val;
-    old = reinterpret_cast<size_t>(address) & 2 ? (old & 0xffff) | (hsum << 16)
-                                                : (old & 0xffff0000) | hsum;
-    old = atomicCAS(address_as_ui, assumed, old);
-
-    // Note: uses integer comparison to avoid hang in case of NaN (since NaN !=
-    // NaN)
-  } while (assumed != old);
-}
-__device__ __forceinline__ void atomicAdd(__half* address, c10::Half val) {
-  unsigned int* address_as_ui =
-      (unsigned int*)((char*)address - ((size_t)address & 2));
-  unsigned int old = *address_as_ui;
-  unsigned int assumed;
-
-  do {
-    assumed = old;
-    __half_raw hsum;
-    hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff);
-    half tmpres = __hadd(hsum, val);
-    hsum = __half_raw(tmpres);
-    old = (size_t)address & 2 ? (old & 0xffff) | (hsum.x << 16)
-                              : (old & 0xffff0000) | hsum.x;
-    old = atomicCAS(address_as_ui, assumed, old);
-  } while (assumed != old);
-}
-#endif
-
-template <typename scalar_t>
-__global__ void VecQuant2MatMulKernel(const scalar_t* __restrict__ vec,
-                                      const int* __restrict__ mat,
-                                      scalar_t* __restrict__ mul,
-                                      const scalar_t* __restrict__ scales,
-                                      const int* __restrict__ zeros,
-                                      const int* __restrict__ g_idx,
-                                      int batch,
-                                      int vec_height,
-                                      int height,
-                                      int width,
-                                      int zero_width);
-
-template <typename scalar_t>
-__global__ void VecQuant3MatMulKernel(const scalar_t* __restrict__ vec,
-                                      const int* __restrict__ mat,
-                                      scalar_t* __restrict__ mul,
-                                      const scalar_t* __restrict__ scales,
-                                      const int* __restrict__ zeros,
-                                      const int* __restrict__ g_idx,
-                                      int batch,
-                                      int vec_height,
-                                      int height,
-                                      int width,
-                                      int zero_width);
-
-template <typename scalar_t>
-__global__ void VecQuant4MatMulKernel(const scalar_t* __restrict__ vec,
-                                      const int* __restrict__ mat,
-                                      scalar_t* __restrict__ mul,
-                                      const scalar_t* __restrict__ scales,
-                                      const int* __restrict__ zeros,
-                                      const int* __restrict__ g_idx,
-                                      int batch,
-                                      int vec_height,
-                                      int height,
-                                      int width,
-                                      int zero_width);
-
-template <typename scalar_t>
-__global__ void VecQuant8MatMulKernel(const scalar_t* __restrict__ vec,
-                                      const int* __restrict__ mat,
-                                      scalar_t* __restrict__ mul,
-                                      const scalar_t* __restrict__ scales,
-                                      const int* __restrict__ zeros,
-                                      const int* __restrict__ g_idx,
-                                      int batch,
-                                      int vec_height,
-                                      int height,
-                                      int width,
-                                      int zero_width);
-
-const int BLOCKWIDTH = 64;
-const int BLOCKHEIGHT2 = 4;
-const int BLOCKHEIGHT3 = 6;
-const int BLOCKHEIGHT4 = 8;
-const int BLOCKHEIGHT8 = 16;
-
-__device__ inline unsigned int as_unsigned(int i) {
-  return *reinterpret_cast<unsigned int*>(&i);
-}
-
-__device__ inline int as_int(int i) { return *reinterpret_cast<int*>(&i); }
-
-void vecquant2matmul_cuda(torch::Tensor vec,
-                          torch::Tensor mat,
-                          torch::Tensor mul,
-                          torch::Tensor scales,
-                          torch::Tensor zeros,
-                          torch::Tensor g_idx) {
-  int batch = vec.size(0);
-  int vec_height = vec.size(1);
-  int height = mat.size(0);
-  int width = mat.size(1);
-  int zero_width = zeros.size(1);
-
-  dim3 blocks((height + BLOCKHEIGHT2 - 1) / BLOCKHEIGHT2,
-              (width + BLOCKWIDTH - 1) / BLOCKWIDTH);
-  dim3 threads(BLOCKWIDTH);
-
-  AT_DISPATCH_FLOATING_TYPES(vec.scalar_type(), "vecquant2matmul_cuda", ([&] {
-                               VecQuant2MatMulKernel<<<blocks, threads>>>(
-                                   vec.data_ptr<scalar_t>(),
-                                   mat.data_ptr<int>(),
-                                   mul.data_ptr<scalar_t>(),
-                                   scales.data_ptr<scalar_t>(),
-                                   zeros.data_ptr<int>(),
-                                   g_idx.data_ptr<int>(),
-                                   batch,
-                                   vec_height,
-                                   height,
-                                   width,
-                                   zero_width);
-                             }));
-}
-
-template <typename scalar_t>
-__global__ void VecQuant2MatMulKernel(const scalar_t* __restrict__ vec,
-                                      const int* __restrict__ mat,
-                                      scalar_t* __restrict__ mul,
-                                      const scalar_t* __restrict__ scales,
-                                      const int* __restrict__ zeros,
-                                      const int* __restrict__ g_idx,
-                                      int batch,
-                                      int vec_height,
-                                      int height,
-                                      int width,
-                                      int zero_width) {
-  int h = BLOCKHEIGHT2 * blockIdx.x;
-  int w = BLOCKWIDTH * blockIdx.y + threadIdx.x;
-
-  __shared__ scalar_t blockvec[BLOCKWIDTH];
-  int i = width * h + w;
-  int g_h = h * 16;
-  int k;
-  unsigned int g;
-  scalar_t w_tmp;
-
-  int z_w = w / 16;
-  int z_mod = (w % 16) * 2;
-
-  float weight[BLOCKWIDTH];
-
-  for (k = 0; k < BLOCKWIDTH; ++k) {
-    int k_w = (k / 16);
-    int k_bit = (k % 16) * 2;
-
-    g = as_int(g_idx[g_h + k]);
-    scalar_t scale = scales[g * width + w];
-    scalar_t zero =
-        scalar_t((as_unsigned(zeros[g * zero_width + z_w]) >> z_mod & 0x3) + 1);
-
-    w_tmp = ((as_unsigned(mat[i + (k_w * width)]) >> k_bit) & 0x3);
-
-    weight[k] = scale * (w_tmp - zero);
-  }
-
-  scalar_t res;
-  for (int b = 0; b < batch; ++b) {
-    res = 0;
-
-    blockvec[threadIdx.x] =
-        vec[b * vec_height + blockIdx.x * BLOCKWIDTH + threadIdx.x];
-    __syncthreads();
-    for (k = 0; k < BLOCKWIDTH; ++k) {
-      res += weight[k] * blockvec[k];
-    }
-    atomicAdd(&mul[b * width + w], res);
-    __syncthreads();
-  }
-}
-
-void vecquant3matmul_cuda(torch::Tensor vec,
-                          torch::Tensor mat,
-                          torch::Tensor mul,
-                          torch::Tensor scales,
-                          torch::Tensor zeros,
-                          torch::Tensor g_idx) {
-  int batch = vec.size(0);
-  int vec_height = vec.size(1);
-  int height = mat.size(0);
-  int width = mat.size(1);
-  int zero_width = zeros.size(1);
-
-  dim3 blocks((height + BLOCKHEIGHT3 - 1) / BLOCKHEIGHT3,
-              (width + BLOCKWIDTH - 1) / BLOCKWIDTH);
-  dim3 threads(BLOCKWIDTH);
-
-  AT_DISPATCH_FLOATING_TYPES(vec.scalar_type(), "vecquant3matmul_cuda", ([&] {
-                               VecQuant3MatMulKernel<<<blocks, threads>>>(
-                                   vec.data_ptr<scalar_t>(),
-                                   mat.data_ptr<int>(),
-                                   mul.data_ptr<scalar_t>(),
-                                   scales.data_ptr<scalar_t>(),
-                                   zeros.data_ptr<int>(),
-                                   g_idx.data_ptr<int>(),
-                                   batch,
-                                   vec_height,
-                                   height,
-                                   width,
-                                   zero_width);
-                             }));
-}
-
-template <typename scalar_t>
-__global__ void VecQuant3MatMulKernel(const scalar_t* __restrict__ vec,
-                                      const int* __restrict__ mat,
-                                      scalar_t* __restrict__ mul,
-                                      const scalar_t* __restrict__ scales,
-                                      const int* __restrict__ zeros,
-                                      const int* __restrict__ g_idx,
-                                      int batch,
-                                      int vec_height,
-                                      int height,
-                                      int width,
-                                      int zero_width) {
-  int h = BLOCKHEIGHT3 * blockIdx.x;
-  int w = BLOCKWIDTH * blockIdx.y + threadIdx.x;
-
-  __shared__ scalar_t blockvec[BLOCKWIDTH];
-  int i = width * h + w;
-  int g_h = (h / 3) * 32;
-  int k;
-  unsigned int g;
-  scalar_t w_tmp;
-
-  int z_w = (w / 32) * 3;
-  int z_mod = w % 32;
-  int z_bit;
-  unsigned int z_tmp;
-  if (z_mod != 10) {
-    if (z_mod != 21) {
-      z_bit = z_mod;
-      if (z_bit > 21) {
-        z_bit -= 22;
-        z_bit *= 3;
-        z_bit += 2;
-        z_w += 2;
-      } else if (z_bit > 10) {
-        z_bit -= 11;
-        z_bit *= 3;
-        z_bit += 1;
-        z_w += 1;
-      } else {
-        z_bit *= 3;
-      }
-    } else {
-      z_w += 1;
-    }
-  }
-
-  float weight[BLOCKWIDTH];
-
-  for (k = 0; k < BLOCKWIDTH; ++k) {
-    int k_w = (k / 32) * 3;
-    int k_mod = k % 32;
-    int k_bit;
-
-    if (k_mod != 10) {
-      if (k_mod != 21) {
-        k_bit = k_mod;
-        if (k_bit > 21) {
-          k_bit -= 22;
-          k_bit *= 3;
-          k_bit += 2;
-          k_w += 2;
-        } else if (k_bit > 10) {
-          k_bit -= 11;
-          k_bit *= 3;
-          k_bit += 1;
-          k_w += 1;
-        } else {
-          k_bit *= 3;
-        }
-      } else {
-        k_w += 1;
-      }
-    }
-
-    g = as_int(g_idx[g_h + k]);
-    scalar_t scale = scales[g * width + w];
-    scalar_t zero;
-    if (z_mod == 10) {
-      z_tmp = (as_unsigned(zeros[g * zero_width + z_w]) >> 30) |
-              ((as_unsigned(zeros[g * zero_width + (z_w + 1)]) << 2) & 0x4);
-      zero = scalar_t((z_tmp) + 1);
-    } else if (z_mod == 21) {
-      z_tmp = (as_unsigned(zeros[g * zero_width + z_w]) >> 31) |
-              ((as_unsigned(zeros[g * zero_width + (z_w + 1)]) << 1) & 0x6);
-      zero = scalar_t((z_tmp) + 1);
-    } else {
-      zero = scalar_t(
-          ((as_unsigned(zeros[g * zero_width + z_w]) >> z_bit) & 0x7) + 1);
-    }
-
-    if (k_mod == 10) {
-      w_tmp = (as_unsigned(mat[i + (k_w * width)]) >> 30) |
-              ((as_unsigned(mat[i + ((k_w + 1) * width)]) << 2) & 0x4);
-    } else if (k_mod == 21) {
-      w_tmp = (as_unsigned(mat[i + (k_w * width)]) >> 31) |
-              ((as_unsigned(mat[i + ((k_w + 1) * width)]) << 1) & 0x6);
-    } else {
-      w_tmp = ((as_unsigned(mat[i + (k_w * width)]) >> k_bit) & 0x7);
-    }
-    weight[k] = scale * (w_tmp - zero);
-  }
-
-  scalar_t res;
-  for (int b = 0; b < batch; ++b) {
-    res = 0;
-
-    blockvec[threadIdx.x] =
-        vec[b * vec_height + blockIdx.x * BLOCKWIDTH + threadIdx.x];
-    __syncthreads();
-    for (k = 0; k < BLOCKWIDTH; ++k) {
-      res += weight[k] * blockvec[k];
-    }
-    atomicAdd(&mul[b * width + w], res);
-    __syncthreads();
-  }
-}
-
-void vecquant4matmul_cuda(torch::Tensor vec,
-                          torch::Tensor mat,
-                          torch::Tensor mul,
-                          torch::Tensor scales,
-                          torch::Tensor zeros,
-                          torch::Tensor g_idx) {
-  int batch = vec.size(0);
-  int vec_height = vec.size(1);
-  int height = mat.size(0);
-  int width = mat.size(1);
-  int zero_width = zeros.size(1);
-
-  dim3 blocks((height + BLOCKHEIGHT4 - 1) / BLOCKHEIGHT4,
-              (width + BLOCKWIDTH - 1) / BLOCKWIDTH);
-  dim3 threads(BLOCKWIDTH);
-
-  AT_DISPATCH_FLOATING_TYPES(vec.scalar_type(), "vecquant4matmul_cuda", ([&] {
-                               VecQuant4MatMulKernel<<<blocks, threads>>>(
-                                   vec.data_ptr<scalar_t>(),
-                                   mat.data_ptr<int>(),
-                                   mul.data_ptr<scalar_t>(),
-                                   scales.data_ptr<scalar_t>(),
-                                   zeros.data_ptr<int>(),
-                                   g_idx.data_ptr<int>(),
-                                   batch,
-                                   vec_height,
-                                   height,
-                                   width,
-                                   zero_width);
-                             }));
-}
-
-template <typename scalar_t>
-__global__ void VecQuant4MatMulKernel(const scalar_t* __restrict__ vec,
-                                      const int* __restrict__ mat,
-                                      scalar_t* __restrict__ mul,
-                                      const scalar_t* __restrict__ scales,
-                                      const int* __restrict__ zeros,
-                                      const int* __restrict__ g_idx,
-                                      int batch,
-                                      int vec_height,
-                                      int height,
-                                      int width,
-                                      int zero_width) {
-  int h = BLOCKHEIGHT4 * blockIdx.x;
-  int w = BLOCKWIDTH * blockIdx.y + threadIdx.x;
-
-  __shared__ scalar_t blockvec[BLOCKWIDTH];
-  int i = width * h + w;
-  int g_h = h * 8;
-  int k;
-  unsigned int g;
-  scalar_t w_tmp;
-
-  int z_w = w / 8;
-  int z_mod = (w % 8) * 4;
-
-  float weight[BLOCKWIDTH];
-
-  for (k = 0; k < BLOCKWIDTH; ++k) {
-    int k_w = (k / 8);
-    int k_bit = (k % 8) * 4;
-
-    g = as_int(g_idx[g_h + k]);
-    scalar_t scale = scales[g * width + w];
-    scalar_t zero = scalar_t(
-        ((as_unsigned(zeros[g * zero_width + z_w]) >> z_mod) & 0xF) + 1);
-
-    w_tmp = ((as_unsigned(mat[i + (k_w * width)]) >> k_bit) & 0xF);
-
-    weight[k] = scale * (w_tmp - zero);
-  }
-
-  scalar_t res;
-  for (int b = 0; b < batch; ++b) {
-    res = 0;
-
-    blockvec[threadIdx.x] =
-        vec[b * vec_height + blockIdx.x * BLOCKWIDTH + threadIdx.x];
-    __syncthreads();
-    for (k = 0; k < BLOCKWIDTH; ++k) {
-      res += weight[k] * blockvec[k];
-    }
-    atomicAdd(&mul[b * width + w], res);
-    __syncthreads();
-  }
-}
-
-void vecquant8matmul_cuda(torch::Tensor vec,
-                          torch::Tensor mat,
-                          torch::Tensor mul,
-                          torch::Tensor scales,
-                          torch::Tensor zeros,
-                          torch::Tensor g_idx) {
-  int batch = vec.size(0);
-  int vec_height = vec.size(1);
-  int height = mat.size(0);
-  int width = mat.size(1);
-  int zero_width = zeros.size(1);
-
-  dim3 blocks((height + BLOCKHEIGHT8 - 1) / BLOCKHEIGHT8,
-              (width + BLOCKWIDTH - 1) / BLOCKWIDTH);
-  dim3 threads(BLOCKWIDTH);
-
-  AT_DISPATCH_FLOATING_TYPES(vec.scalar_type(), "vecquant8matmul_cuda", ([&] {
-                               VecQuant8MatMulKernel<<<blocks, threads>>>(
-                                   vec.data_ptr<scalar_t>(),
-                                   mat.data_ptr<int>(),
-                                   mul.data_ptr<scalar_t>(),
-                                   scales.data_ptr<scalar_t>(),
-                                   zeros.data_ptr<int>(),
-                                   g_idx.data_ptr<int>(),
-                                   batch,
-                                   vec_height,
-                                   height,
-                                   width,
-                                   zero_width);
-                             }));
-}
-
-template <typename scalar_t>
-__global__ void VecQuant8MatMulKernel(const scalar_t* __restrict__ vec,
-                                      const int* __restrict__ mat,
-                                      scalar_t* __restrict__ mul,
-                                      const scalar_t* __restrict__ scales,
-                                      const int* __restrict__ zeros,
-                                      const int* __restrict__ g_idx,
-                                      int batch,
-                                      int vec_height,
-                                      int height,
-                                      int width,
-                                      int zero_width) {
-  int h = BLOCKHEIGHT8 * blockIdx.x;
-  int w = BLOCKWIDTH * blockIdx.y + threadIdx.x;
-
-  __shared__ scalar_t blockvec[BLOCKWIDTH];
-  int i = width * h + w;
-  int g_h = h * 4;
-  int k;
-  unsigned int g;
-  scalar_t w_tmp;
-
-  int z_w = w / 4;
-  int z_mod = (w % 4) * 8;
-
-  float weight[BLOCKWIDTH];
-
-  for (k = 0; k < BLOCKWIDTH; ++k) {
-    int k_w = (k / 4);
-    int k_bit = (k % 4) * 8;
-
-    g = as_int(g_idx[g_h + k]);
-    scalar_t scale = scales[g * width + w];
-    scalar_t zero = scalar_t(
-        ((as_unsigned(zeros[g * zero_width + z_w]) >> z_mod) & 0xFF) + 1);
-
-    w_tmp = ((as_unsigned(mat[i + (k_w * width)]) >> k_bit) & 0xFF);
-
-    weight[k] = scale * (w_tmp - zero);
-  }
-
-  scalar_t res;
-  for (int b = 0; b < batch; ++b) {
-    res = 0;
-
-    blockvec[threadIdx.x] =
-        vec[b * vec_height + blockIdx.x * BLOCKWIDTH + threadIdx.x];
-    __syncthreads();
-    for (k = 0; k < BLOCKWIDTH; ++k) {
-      res += weight[k] * blockvec[k];
-    }
-    atomicAdd(&mul[b * width + w], res);
-    __syncthreads();
-  }
-}
diff --git a/src/layers/attention.cpp b/src/layers/attention.cpp
index 0d1397c8..570fa7e6 100644
--- a/src/layers/attention.cpp
+++ b/src/layers/attention.cpp
@@ -1,6 +1,7 @@
 #include "attention.h"
 
 #include <gflags/gflags.h>
+#include <torch/csrc/autograd/generated/variable_factories.h>
 #include <torch/torch.h>
 
 #include "common/logging.h"
@@ -16,6 +17,11 @@ DEFINE_string(
     "type of attention to use for single_query_masked_self_attention, slow, "
     "cuda, or empty for auto");
 
+DEFINE_bool(
+    force_use_paged_attention_v2,
+    false,
+    "force to use paged attention v2 for single_query_masked_self_attention");
+
 // ref to flash_attn in third_party/flash_attn
 extern std::vector<at::Tensor> mha_varlen_fwd(
     const at::Tensor& q,
@@ -36,8 +42,8 @@ extern std::vector<at::Tensor> mha_varlen_fwd(
     bool return_softmax,
     torch::optional<at::Generator> gen_);
 
-// ref to single_query_cached_kv_attention in third_party/vllm
-extern void single_query_cached_kv_attention(
+// ref to paged_attention_v1 in third_party/vllm
+extern void paged_attention_v1(
     torch::Tensor& out,
     torch::Tensor& query,
     torch::Tensor& key_cache,
@@ -50,6 +56,22 @@ extern void single_query_cached_kv_attention(
     int max_context_len,
     const c10::optional<torch::Tensor>& alibi_slopes);
 
+extern void paged_attention_v2(
+    torch::Tensor& out,
+    torch::Tensor& exp_sums,
+    torch::Tensor& max_logits,
+    torch::Tensor& tmp_out,
+    torch::Tensor& query,
+    torch::Tensor& key_cache,
+    torch::Tensor& value_cache,
+    torch::Tensor& head_mapping,
+    float scale,
+    torch::Tensor& block_tables,
+    torch::Tensor& context_lens,
+    int block_size,
+    int max_context_len,
+    const c10::optional<torch::Tensor>& alibi_slopes);
+
 namespace llm {
 AttentionImpl::AttentionImpl(int64_t n_heads,
                              int64_t n_kv_heads,
@@ -372,6 +394,10 @@ void single_query_masked_self_attention_cuda(
     float scale,                                  // scale for softmax
     torch::Tensor& output) {
   auto [key_cache, value_cache] = kv_cache.get_kv_cache();
+  static constexpr int32_t kPartitionSize = 512;
+  const auto n_seq = query.size(0);
+  const auto n_heads = query.size(1);
+  const auto head_dim = query.size(2);
   const auto block_size = key_cache.size(3);
 
   // make a 'copy' of variable since the api is using non-const reference
@@ -379,17 +405,50 @@ void single_query_masked_self_attention_cuda(
   torch::Tensor _block_tables = block_tables;
   torch::Tensor _context_lens = context_lens;
   torch::Tensor _kv_head_mapping = kv_head_mapping;
-  single_query_cached_kv_attention(output,
-                                   _query,
-                                   key_cache,
-                                   value_cache,
-                                   _kv_head_mapping,
-                                   scale,
-                                   _block_tables,
-                                   _context_lens,
-                                   block_size,
-                                   max_context_len,
-                                   alibi_slopes);
+
+  // Adapted from:
+  // https://github.com/vllm-project/vllm/blob/f8a1e39fae05ca610be8d5a78be9d40f5274e5fc/vllm/model_executor/layers/attention.py#L159
+  // round up partition number
+  const auto num_partitions =
+      (max_context_len + kPartitionSize - 1) / kPartitionSize;
+  const bool use_v1 = num_partitions == 1 || (n_seq * n_heads) > kPartitionSize;
+  // Use the same simple heuristic as vllm to decide whether to use v1.
+  if (!FLAGS_force_use_paged_attention_v2 && use_v1) {
+    paged_attention_v1(output,
+                       _query,
+                       key_cache,
+                       value_cache,
+                       _kv_head_mapping,
+                       scale,
+                       _block_tables,
+                       _context_lens,
+                       block_size,
+                       max_context_len,
+                       alibi_slopes);
+  } else {
+    GCHECK(kPartitionSize % block_size == 0);
+    auto tmp_out = torch::empty({n_seq, n_heads, num_partitions, head_dim},
+                                output.options());
+    auto exp_sums =
+        torch::empty({n_seq, n_heads, num_partitions},
+                     torch::dtype(torch::kFloat32).device(output.device()));
+    auto max_logits = torch::empty_like(exp_sums);
+
+    paged_attention_v2(output,
+                       exp_sums,
+                       max_logits,
+                       tmp_out,
+                       _query,
+                       key_cache,
+                       value_cache,
+                       _kv_head_mapping,
+                       scale,
+                       _block_tables,
+                       _context_lens,
+                       block_size,
+                       max_context_len,
+                       alibi_slopes);
+  }
 }
 
 }  // namespace detail
diff --git a/src/layers/quantization/qlinear_awq_impl.cpp b/src/layers/quantization/qlinear_awq_impl.cpp
index ce1762a8..65fd80b2 100644
--- a/src/layers/quantization/qlinear_awq_impl.cpp
+++ b/src/layers/quantization/qlinear_awq_impl.cpp
@@ -37,7 +37,6 @@ ColumnParallelQLinearAWQImpl::ColumnParallelQLinearAWQImpl(
   const auto bits = quant_args.bits();
   const auto group_size = quant_args.group_size();
   GCHECK(bits == 4) << "Only 4 bits are supported for AWQ";
-  GCHECK(group_size > 0) << "group_size must be positive";
   pack_factor_ = 32 / bits;
 }
 
@@ -72,9 +71,7 @@ RowParallelQLinearAWQImpl::RowParallelQLinearAWQImpl(
                              dtype,
                              device) {
   const auto bits = quant_args.bits();
-  const auto group_size = quant_args.group_size();
   GCHECK(bits == 4) << "Only 4 bits are supported for AWQ";
-  GCHECK(group_size > 0) << "group_size must be positive";
   pack_factor_ = 32 / bits;
 }
 
diff --git a/src/layers/quantization/qlinear_exllama_impl.cpp b/src/layers/quantization/qlinear_exllama_impl.cpp
index 1847d75b..91642440 100644
--- a/src/layers/quantization/qlinear_exllama_impl.cpp
+++ b/src/layers/quantization/qlinear_exllama_impl.cpp
@@ -44,9 +44,9 @@ ColumnParallelQLinearExllamaImpl::ColumnParallelQLinearExllamaImpl(
                                 dtype,
                                 device) {
   const auto bits = quant_args.bits();
-  const auto group_size = quant_args.group_size();
   GCHECK(bits == 4) << "Only 4 bits are supported";
-  GCHECK(group_size > 0) << "group_size must be positive";
+  const auto group_size =
+      quant_args.group_size() > 0 ? quant_args.group_size() : in_features;
 
   // using torch::aarange to create g_idx
   std::vector<int32_t> g_idx_data;
@@ -101,9 +101,9 @@ RowParallelQLinearExllamaImpl::RowParallelQLinearExllamaImpl(
                              dtype,
                              device) {
   const auto bits = quant_args.bits();
-  const auto group_size = quant_args.group_size();
   GCHECK(bits == 4) << "Only 4 bits are supported";
-  GCHECK(group_size > 0) << "group_size must be positive";
+  const auto group_size =
+      quant_args.group_size() > 0 ? quant_args.group_size() : in_features;
 
   std::vector<int32_t> g_idx_data;
   g_idx_data.reserve(in_features);
diff --git a/src/layers/quantization/qlinear_exllamav2_impl.cpp b/src/layers/quantization/qlinear_exllamav2_impl.cpp
index 04c9dc52..bfd1be9a 100644
--- a/src/layers/quantization/qlinear_exllamav2_impl.cpp
+++ b/src/layers/quantization/qlinear_exllamav2_impl.cpp
@@ -74,9 +74,10 @@ ColumnParallelQLinearExllamav2Impl::ColumnParallelQLinearExllamav2Impl(
                                 dtype,
                                 device) {
   const auto bits = quant_args.bits();
-  const auto group_size = quant_args.group_size();
   GCHECK(bits == 4) << "Only 4 bits are supported";
-  GCHECK(group_size > 0) << "group_size must be positive";
+
+  const auto group_size =
+      quant_args.group_size() > 0 ? quant_args.group_size() : in_features;
 
   if (quant_args.desc_act()) {
     std::vector<int32_t> g_idx_data;
@@ -162,10 +163,11 @@ RowParallelQLinearExllamav2Impl::RowParallelQLinearExllamav2Impl(
                              dtype,
                              device) {
   const auto bits = quant_args.bits();
-  const auto group_size = quant_args.group_size();
   GCHECK(bits == 2 || bits == 3 || bits == 4 || bits == 8)
       << "Only 2,3,4,8 bits are supported";
-  GCHECK(group_size > 0) << "group_size must be positive";
+
+  const auto group_size =
+      quant_args.group_size() > 0 ? quant_args.group_size() : in_features;
 
   if (quant_args.desc_act()) {
     std::vector<int32_t> g_idx_data;
diff --git a/src/layers/quantization/qlinear_gptq_impl.cpp b/src/layers/quantization/qlinear_gptq_impl.cpp
index 5560c3ea..2ad093c2 100644
--- a/src/layers/quantization/qlinear_gptq_impl.cpp
+++ b/src/layers/quantization/qlinear_gptq_impl.cpp
@@ -47,10 +47,11 @@ ColumnParallelQLinearGPTQImpl::ColumnParallelQLinearGPTQImpl(
                                 device),
       bits_(quant_args.bits()) {
   const auto bits = quant_args.bits();
-  const auto group_size = quant_args.group_size();
   GCHECK(bits == 2 || bits == 3 || bits == 4 || bits == 8)
       << "Only 2,3,4,8 bits are supported";
-  GCHECK(group_size > 0) << "group_size must be positive";
+
+  const auto group_size =
+      quant_args.group_size() > 0 ? quant_args.group_size() : in_features;
 
   std::vector<int32_t> g_idx_data;
   g_idx_data.reserve(in_features);
@@ -110,10 +111,12 @@ RowParallelQLinearGPTQImpl::RowParallelQLinearGPTQImpl(
                              device),
       bits_(quant_args.bits()) {
   const auto bits = quant_args.bits();
-  const auto group_size = quant_args.group_size();
   GCHECK(bits == 2 || bits == 3 || bits == 4 || bits == 8)
       << "Only 2,3,4,8 bits are supported";
-  GCHECK(group_size > 0) << "group_size must be positive";
+
+  const auto group_size = quant_args.group_size() > 0
+                              ? quant_args.group_size()
+                              : in_features;
 
   std::vector<int32_t> g_idx_data;
   g_idx_data.reserve(in_features);
diff --git a/src/layers/quantization/qlinear_impl.cpp b/src/layers/quantization/qlinear_impl.cpp
index a450d9eb..3929f039 100644
--- a/src/layers/quantization/qlinear_impl.cpp
+++ b/src/layers/quantization/qlinear_impl.cpp
@@ -118,8 +118,8 @@ ColumnParallelQLinearImpl::ColumnParallelQLinearImpl(
       gather_output_(gather_output),
       parallel_args_(parallel_args) {
   const auto bits = quant_args.bits();
-  const auto group_size = quant_args.group_size();
-  GCHECK(group_size > 0) << "group_size must be positive";
+  const auto group_size =
+      quant_args.group_size() > 0 ? quant_args.group_size() : in_features;
   GCHECK(qweight_pack_dim == 0 || qweight_pack_dim == 1)
       << "qweight_pack_dim must be 0 or 1";
   const int64_t world_size = parallel_args.world_size();
@@ -350,8 +350,6 @@ RowParallelQLinearImpl::RowParallelQLinearImpl(
       input_is_parallelized_(input_is_parallelized),
       parallel_args_(parallel_args) {
   const auto bits = quant_args.bits();
-  const auto group_size = quant_args.group_size();
-  GCHECK(group_size > 0) << "group_size must be positive";
   GCHECK(qweight_pack_dim == 0 || qweight_pack_dim == 1)
       << "qweight_pack_dim must be 0 or 1";
   const int64_t world_size = parallel_args.world_size();
@@ -360,6 +358,8 @@ RowParallelQLinearImpl::RowParallelQLinearImpl(
       << world_size;
   const int64_t in_features_per_partition = in_features / world_size;
   const int64_t pack_factor = 32 / bits;
+  const auto group_size =
+      quant_args.group_size() > 0 ? quant_args.group_size() : in_features;
 
   if (qweight_pack_dim == 0) {
     qweight_ = register_parameter(
diff --git a/src/model_loader/model_loader.cpp b/src/model_loader/model_loader.cpp
index 72fc2ce0..a723aba4 100644
--- a/src/model_loader/model_loader.cpp
+++ b/src/model_loader/model_loader.cpp
@@ -78,7 +78,8 @@ std::unique_ptr<Tokenizer> PTModelLoader::tokenizer() const {
     GLOG(ERROR) << "Failed to find tokenizer file: " << tokenizer_path;
     return nullptr;
   }
-  return std::make_unique<SentencePieceTokenizer>(tokenizer_path);
+  return std::make_unique<SentencePieceTokenizer>(tokenizer_path,
+                                                  /*prepend_bos=*/true);
 }
 
 PTModelLoader::PTModelLoader(const std::string& model_weights_path)
@@ -150,12 +151,14 @@ std::unique_ptr<Tokenizer> HFModelLoader::tokenizer() const {
     return HFTokenizer::from_file(tokenizer_path);
   }
 
+  // fallback to tokenizer.model if tokenizer.json does not exist
   const std::string vocab_path = model_weights_path_ + "/tokenizer.model";
   if (std::filesystem::exists(vocab_path)) {
     GLOG(WARNING) << "Failed to find tokenizer.json, use tokenizer.model "
                      "instead. Please consider to convert the tokenizer.model "
                      "to tokenizer.json for better performance.";
-    return std::make_unique<SentencePieceTokenizer>(vocab_path);
+    return std::make_unique<SentencePieceTokenizer>(vocab_path,
+                                                    /*prepend_bos=*/false);
   }
 
   GLOG(ERROR)
@@ -253,6 +256,13 @@ bool HFModelLoader::load_model_args(const std::string& model_weights_path) {
     quant_args_.quant_method() = FLAGS_quant_method;
   }
 
+  // always use float16 for quantization
+  if (!quant_args_.quant_method().empty() && args_.dtype() != "float16") {
+    LOG(WARNING) << "Overwriting dtype from " << args_.dtype() << " to float16 "
+                 << "for quantization";
+    args_.dtype() = "float16";
+  }
+
   return true;
 }
 
diff --git a/src/model_loader/state_dict.cpp b/src/model_loader/state_dict.cpp
index e5899c6a..9d04748a 100644
--- a/src/model_loader/state_dict.cpp
+++ b/src/model_loader/state_dict.cpp
@@ -215,6 +215,11 @@ torch::Tensor StateDict::get_sharded_tensor(const std::string_view& tensor_name,
   }
   // chunk tensor along the dim
   const int64_t dim_size = tensor.size(dim);
+  if (dim_size <= num_ranks_per_shard) {
+    // too small to shard, return the whole tensor instead
+    return tensor;
+  }
+
   GCHECK(dim_size % num_ranks_per_shard == 0)
       << "can't devide tensor evenly on " << dim << " with dim: " << dim_size
       << " ranks_per_shard: " << num_ranks_per_shard;
diff --git a/src/models/CMakeLists.txt b/src/models/CMakeLists.txt
index 78dbd3d6..d7415842 100644
--- a/src/models/CMakeLists.txt
+++ b/src/models/CMakeLists.txt
@@ -22,15 +22,3 @@ cc_library(
     torch
 )
 
-cc_test(
-  NAME
-    models_test
-  SRCS
-    llama_test.cpp
-    huggingface/mpt_test.cpp
-  DEPS
-    :models
-    :state_dict
-    GTest::gtest_main
-)
-
diff --git a/src/models/args.h b/src/models/args.h
index 846c62ed..ac4a1616 100644
--- a/src/models/args.h
+++ b/src/models/args.h
@@ -10,6 +10,8 @@ namespace llm {
 struct ModelArgs {
   DEFINE_ARG(std::string, model_type);
 
+  DEFINE_ARG(std::string, dtype);
+
   // dimension of the encoder layer.
   DEFINE_ARG(int64_t, hidden_size) = 0;
 
@@ -84,6 +86,7 @@ struct ModelArgs {
 
 inline std::ostream& operator<<(std::ostream& os, const ModelArgs& args) {
   os << "ModelArgs: [model_type: " << args.model_type();
+  os << ", dtype: " << args.dtype();
   os << ", hidden_size: " << args.hidden_size();
   os << ", hidden_act: " << args.hidden_act();
   os << ", intermediate_size: " << args.intermediate_size();
@@ -106,6 +109,7 @@ inline std::ostream& operator<<(std::ostream& os, const ModelArgs& args) {
   os << ", attn_alibi: " << args.attn_alibi();
   os << ", alibi_bias_max: " << args.alibi_bias_max();
   os << ", no_bias: " << args.no_bias();
+  os << ", residual_post_layernorm: " << args.residual_post_layernorm();
   os << "]";
   return os;
 }
diff --git a/src/models/dialog.h b/src/models/dialog.h
index 78414d10..2e0a3ef1 100644
--- a/src/models/dialog.h
+++ b/src/models/dialog.h
@@ -36,7 +36,7 @@ class Dialog {
 };
 
 // dialog conversation for llama2 model
-class Llama2Dialog : public Dialog {
+class Llama2Dialog final : public Dialog {
  public:
   // generate prompt from dialogs
   std::optional<std::string> get_prompt() const override;
diff --git a/src/models/huggingface/aquila.h b/src/models/huggingface/aquila.h
index b8068a87..c4a31625 100644
--- a/src/models/huggingface/aquila.h
+++ b/src/models/huggingface/aquila.h
@@ -9,6 +9,7 @@
 #include "layers/normalization.h"
 #include "memory/kv_cache.h"
 #include "models/args.h"
+#include "models/dialog.h"
 #include "models/input_parameters.h"
 #include "models/model_registry.h"
 
@@ -362,13 +363,45 @@ class AquilaForCausalLMImpl : public torch::nn::Module {
 };
 TORCH_MODULE(AquilaForCausalLM);
 
+class AquilaDialog final : public Dialog {
+ public:
+  // generate prompt from dialogs
+  // Prompt template for Aquila:
+  // System: A chat between a curious ...
+  // Human: {prompt}
+  // Assistant:
+  std::optional<std::string> get_prompt() const override {
+    // at least one user message
+    if (messages_.size() % 2 == 0) {
+      return std::nullopt;
+    }
+
+    std::stringstream ss;
+    // start with system message
+    if (!system_message_.empty()) {
+      ss << "System: " << system_message_ << "\n";
+    }
+
+    // then user and assistant message pairs (u/a/u/a/u...)
+    for (size_t i = 0; i < messages_.size(); ++i) {
+      const char* role = (i % 2) == 0 ? "Human: " : "Assistant: ";
+      ss << role << messages_[i] << "\n";
+    }
+    // end with assistant message
+    ss << "Assistant:";
+    return ss.str();
+  }
+};
+
 // register the model to make it available
 REGISTER_CAUSAL_MODEL(aquila, AquilaForCausalLM);
+REGISTER_DIALOG(aquila, AquilaDialog);
 REGISTER_MODEL_ARGS(aquila, [&] {
   // example config:
   // https://huggingface.co/BAAI/Aquila-7B/blob/main/config.json.
   // set default values for args explicitly with values from:
   // https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral/configuration_mistral.py#L104
+  LOAD_ARG_OR(dtype, "torch_dtype", "");
   LOAD_ARG_OR(vocab_size, "vocab_size", 100008);
   LOAD_ARG_OR(hidden_size, "hidden_size", 4096);
   LOAD_ARG_OR(n_layers, "num_hidden_layers", 32);
diff --git a/src/models/huggingface/bloom.h b/src/models/huggingface/bloom.h
index 5a8321c4..2cc77d2a 100644
--- a/src/models/huggingface/bloom.h
+++ b/src/models/huggingface/bloom.h
@@ -430,6 +430,7 @@ REGISTER_CAUSAL_MODEL(bloom, BloomForCausalLM);
 REGISTER_MODEL_ARGS(bloom, [&] {
   // example config:
   // https://huggingface.co/bigscience/bloom/blob/main/config.json
+  LOAD_ARG_OR(dtype, "torch_dtype", "");
   LOAD_ARG_OR(vocab_size, "vocab_size", 250880);
   LOAD_ARG_OR(hidden_size, "n_embed", 14336);
   LOAD_ARG_OR(n_layers, "n_layer", 70);
diff --git a/src/models/huggingface/gpt2.h b/src/models/huggingface/gpt2.h
index b692ca1d..874418a6 100644
--- a/src/models/huggingface/gpt2.h
+++ b/src/models/huggingface/gpt2.h
@@ -365,6 +365,7 @@ TORCH_MODULE(GPT2ForCausalLM);
 // register the model to make it available
 REGISTER_CAUSAL_MODEL(gpt2, GPT2ForCausalLM);
 REGISTER_MODEL_ARGS(gpt2, [&] {
+  LOAD_ARG_OR(dtype, "torch_dtype", "");
   LOAD_ARG_OR(vocab_size, "vocab_size", 50257);
   LOAD_ARG_OR(hidden_size, "n_embd", 768);
   LOAD_ARG_OR(n_layers, "n_layer", 12);
diff --git a/src/models/huggingface/gpt_j.h b/src/models/huggingface/gpt_j.h
index ceff86de..7468cfe7 100644
--- a/src/models/huggingface/gpt_j.h
+++ b/src/models/huggingface/gpt_j.h
@@ -351,6 +351,7 @@ REGISTER_MODEL_ARGS(gptj, [&] {
   // https://huggingface.co/EleutherAI/gpt-j-6b/blob/main/config.json set
   // default values for args explicitly with values from:
   // https://github.com/huggingface/transformers/blob/main/src/transformers/models/gptj/configuration_gptj.py#L98
+  LOAD_ARG_OR(dtype, "torch_dtype", "");
   LOAD_ARG_OR(vocab_size, "vocab_size", 50400);
   LOAD_ARG_OR(hidden_size, "n_embd", 4096);
   LOAD_ARG_OR(n_layers, "n_layer", 28);
diff --git a/src/models/huggingface/gpt_neox.h b/src/models/huggingface/gpt_neox.h
index 6f0b6c66..79f7d97f 100644
--- a/src/models/huggingface/gpt_neox.h
+++ b/src/models/huggingface/gpt_neox.h
@@ -403,6 +403,7 @@ REGISTER_MODEL_ARGS(gpt_neox, [&] {
   // https://huggingface.co/EleutherAI/gpt-neox-20b/blob/main/config.json set
   // set default values for args explicitly with values from:
   // https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/configuration_gpt_neox.py#L106
+  LOAD_ARG_OR(dtype, "torch_dtype", "");
   LOAD_ARG_OR(vocab_size, "vocab_size", 50432);
   LOAD_ARG_OR(hidden_size, "hidden_size", 6144);
   LOAD_ARG_OR(n_layers, "num_hidden_layers", 44);
diff --git a/src/models/huggingface/internlm.h b/src/models/huggingface/internlm.h
index 2efcc274..cdddacd0 100644
--- a/src/models/huggingface/internlm.h
+++ b/src/models/huggingface/internlm.h
@@ -350,9 +350,43 @@ class InternlmForCausalLMImpl : public torch::nn::Module {
 };
 TORCH_MODULE(InternlmForCausalLM);
 
+class InternlmDialog final : public Dialog {
+ public:
+  // generate prompt from dialogs
+  // Prompt template:
+  // 

:{messages[0]}<eoh>\n
:{messages[1]}<eoa>\n
+  std::optional<std::string> get_prompt() const override {
+    // at least one user message
+    if (messages_.size() % 2 == 0) {
+      return std::nullopt;
+    }
+
+    std::stringstream ss;
+    // start with system message
+    if (!system_message_.empty()) {
+      ss << system_message_;
+    }
+    // then user and assistant message pairs (u/a/u/a/u...)
+    for (size_t i = 0; i < messages_.size(); ++i) {
+      if (i % 2 == 0) {
+        // user message
+        ss << "

:" << messages_[i] << "<eoh>\n";
+      } else {
+        // assistant message
+        ss << "
:" << messages_[i] << "<eoa>\n";
+      }
+    }
+    // end with assistant message
+    ss << "
:";
+    return ss.str();
+  }
+};
+
 // register the model to make it available
 REGISTER_CAUSAL_MODEL(internlm, InternlmForCausalLM);
+REGISTER_DIALOG(internlm, InternlmDialog);
 REGISTER_MODEL_ARGS(internlm, [&] {
+  LOAD_ARG_OR(dtype, "torch_dtype", "");
   LOAD_ARG_OR(vocab_size, "vocab_size", 103168);
   LOAD_ARG_OR(hidden_size, "hidden_size", 5120);
   LOAD_ARG_OR(n_layers, "num_hidden_layers", 60);
diff --git a/src/models/huggingface/llama.h b/src/models/huggingface/llama.h
index ebcfc06e..f6e1dc2d 100644
--- a/src/models/huggingface/llama.h
+++ b/src/models/huggingface/llama.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <c10/core/ScalarType.h>
 #include <torch/torch.h>
 
 #include "layers/activation.h"
@@ -372,6 +373,7 @@ REGISTER_DIALOG(llama, Llama2Dialog);
 // default values for args explicitly with values from:
 // https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/configuration_llama.py#L112
 REGISTER_MODEL_ARGS(llama, [&] {
+  LOAD_ARG_OR(dtype, "torch_dtype", "");
   LOAD_ARG_OR(vocab_size, "vocab_size", 32000);
   LOAD_ARG_OR(hidden_size, "hidden_size", 4096);
   LOAD_ARG_OR(n_layers, "num_hidden_layers", 32);
diff --git a/src/models/huggingface/mistral.h b/src/models/huggingface/mistral.h
index ced48e71..abe72dbe 100644
--- a/src/models/huggingface/mistral.h
+++ b/src/models/huggingface/mistral.h
@@ -361,9 +361,45 @@ class MistralForCausalLMImpl : public torch::nn::Module {
 };
 TORCH_MODULE(MistralForCausalLM);
 
+class MistralDialog final : public Dialog {
+ public:
+  // generate prompt from dialogs
+  // Prompt template:
+  // 

 Instruction 
 Model answer

 Follow-up instruction 

+  std::optional<std::string> get_prompt() const override {
+    // at least one user message
+  if (messages_.size() % 2 == 0) {
+    return std::nullopt;
+  }
+
+  std::stringstream ss;
+  // start with system message
+  // N.B. tokenizer would add 
 to the beginning of the prompt
+  if (!system_message_.empty()) {
+    ss << system_message_;
+  }
+
+  // then user and assistant message pairs (u/a/u/a/u...)
+  for (size_t i = 0; i < messages_.size(); ++i) {
+    if (i % 2 == 0) {
+      // user message
+      ss << "
 " << messages_[i] << " ";
+    } else {
+      // assistant message
+      ss << "
 " << messages_[i] << "
";
+    }
+  }
+  // end with assistant message
+  ss << "
";
+  return ss.str();
+  }
+};
+
 // register the model to make it available
 REGISTER_CAUSAL_MODEL(mistral, MistralForCausalLM);
+REGISTER_DIALOG(mistral, MistralDialog);
 REGISTER_MODEL_ARGS(mistral, [&] {
+  LOAD_ARG_OR(dtype, "torch_dtype", "");
   LOAD_ARG_OR(vocab_size, "vocab_size", 32000);
   LOAD_ARG_OR(hidden_size, "hidden_size", 4096);
   LOAD_ARG_OR(n_layers, "num_hidden_layers", 32);
diff --git a/src/models/huggingface/mpt.h b/src/models/huggingface/mpt.h
index b803b287..60f18122 100644
--- a/src/models/huggingface/mpt.h
+++ b/src/models/huggingface/mpt.h
@@ -456,6 +456,7 @@ TORCH_MODULE(MPTForCausalLM);
 
 REGISTER_CAUSAL_MODEL(mpt, MPTForCausalLM);
 REGISTER_MODEL_ARGS(mpt, [&] {
+  LOAD_ARG_OR(dtype, "torch_dtype", "");
   LOAD_ARG_OR(vocab_size, "vocab_size", 50368);
   LOAD_ARG_OR(hidden_size, "d_model", 2048);
   LOAD_ARG_OR(n_layers, "num_hidden_layers", 32);
diff --git a/src/models/huggingface/yi.h b/src/models/huggingface/yi.h
new file mode 100644
index 00000000..301419d2
--- /dev/null
+++ b/src/models/huggingface/yi.h
@@ -0,0 +1,382 @@
+#pragma once
+
+#include <torch/torch.h>
+
+#include "layers/activation.h"
+#include "layers/attention_rope.h"
+#include "layers/embedding.h"
+#include "layers/linear.h"
+#include "layers/normalization.h"
+#include "memory/kv_cache.h"
+#include "models/args.h"
+#include "models/dialog.h"
+#include "models/input_parameters.h"
+#include "models/model_registry.h"
+
+// Yi model compatible with huggingface weights
+namespace llm::hf {
+
+class YiMLPImpl : public torch::nn::Module {
+ public:
+  YiMLPImpl(const ModelArgs& args,
+            const QuantizationArgs& quant_args,
+            const ParallelArgs& parallel_args,
+            torch::ScalarType dtype,
+            const torch::Device& device) {
+    act_with_mul_ = Activation::get_act_with_mul_func("silu", device);
+    GCHECK(act_with_mul_ != nullptr);
+
+    const int64_t hidden_size = args.hidden_size();
+    const int64_t intermediate_size = args.intermediate_size();
+
+    // register the weight parameter
+    gate_up_proj_ =
+        register_module("gate_up_proj",
+                        ColumnParallelLinear(hidden_size,
+                                             intermediate_size * 2,
+                                             /*bias=*/false,
+                                             /*gather_output=*/false,
+                                             quant_args,
+                                             parallel_args,
+                                             dtype,
+                                             device));
+    down_proj_ =
+        register_module("down_proj",
+                        RowParallelLinear(intermediate_size,
+                                          hidden_size,
+                                          /*bias=*/false,
+                                          /*input_is_parallelized=*/true,
+                                          quant_args,
+                                          parallel_args,
+                                          dtype,
+                                          device));
+  }
+
+  torch::Tensor forward(torch::Tensor x) {
+    return down_proj_(act_with_mul_(gate_up_proj_(x)));
+  }
+
+  // load the weight from the checkpoint
+  void load_state_dict(const StateDict& state_dict) {
+    // call each submodule's load_state_dict function
+    gate_up_proj_->load_state_dict(state_dict, {"gate_proj.", "up_proj."});
+    down_proj_->load_state_dict(state_dict.select("down_proj."));
+  }
+
+  void verify_loaded_weights(const std::string& prefix) const {
+    gate_up_proj_->verify_loaded_weights(prefix + "[gate_proj,up_proj].");
+    down_proj_->verify_loaded_weights(prefix + "down_proj.");
+  }
+
+ private:
+  // parameter members, must be registered
+  ColumnParallelLinear gate_up_proj_{nullptr};
+  RowParallelLinear down_proj_{nullptr};
+
+  // calculate act(x) * y
+  ActFunc act_with_mul_{nullptr};
+};
+TORCH_MODULE(YiMLP);
+
+class YiAttentionImpl : public torch::nn::Module {
+ public:
+  YiAttentionImpl(const ModelArgs& args,
+                  const QuantizationArgs& quant_args,
+                  const ParallelArgs& parallel_args,
+                  torch::ScalarType dtype,
+                  const torch::Device& device) {
+    const int32_t world_size = parallel_args.world_size();
+    const int64_t hidden_size = args.hidden_size();
+    const int64_t n_heads = args.n_heads();
+    const int64_t n_kv_heads = args.n_kv_heads().value_or(n_heads);
+    const int64_t head_dim = hidden_size / n_heads;
+    const int64_t n_local_heads = n_heads / world_size;
+    const int64_t n_local_kv_heads = n_kv_heads / world_size;
+
+    // size for q, k, v
+    qkv_sizes_ = {n_local_heads * head_dim,
+                  n_local_kv_heads * head_dim,
+                  n_local_kv_heads * head_dim};
+
+    // register submodules
+    qkv_proj_ = register_module(
+        "qkv_proj",
+        ColumnParallelLinear(hidden_size,
+                             (n_heads + 2 * n_kv_heads) * head_dim,
+                             /*bias=*/false,
+                             /*gather_output=*/false,
+                             quant_args,
+                             parallel_args,
+                             dtype,
+                             device));
+
+    o_proj_ = register_module("o_proj",
+                              RowParallelLinear(hidden_size,
+                                                hidden_size,
+                                                /*bias=*/false,
+                                                /*input_is_parallelized=*/true,
+                                                quant_args,
+                                                parallel_args,
+                                                dtype,
+                                                device));
+
+    // initialize attention
+    const float scale = 1.0f / std::sqrt(static_cast<float>(head_dim));
+    atten_ = register_module("atten",
+                             AttentionWithRoPE(n_local_heads,
+                                               n_local_kv_heads,
+                                               head_dim,
+                                               scale,
+                                               /*rotary_dim=*/head_dim,
+                                               args.rope_scaling(),
+                                               args.rope_theta(),
+                                               args.max_position_embeddings(),
+                                               /*interleaved=*/false,
+                                               dtype,
+                                               device));
+  }
+
+  torch::Tensor forward(torch::Tensor x,
+                        torch::Tensor positions,
+                        KVCache& kv_cache,
+                        const InputParameters& input_params) {
+    // (num_tokens, dim) x (dim, n_local_heads * head_dim)
+    // => (num_tokens, n_local_heads * head_dim)
+    auto qkv = qkv_proj_(x).split(/*split_size=*/qkv_sizes_, /*dim=*/-1);
+    DCHECK_EQ(qkv.size(), 3);
+
+    // calculate attention, output: (num_tokens, n_local_heads * head_dim)
+    auto output =
+        atten_(qkv[0], qkv[1], qkv[2], positions, kv_cache, input_params);
+    return o_proj_(output);
+  }
+
+  // load the weight from the checkpoint
+  void load_state_dict(const StateDict& state_dict) {
+    // call each submodule's load_state_dict function
+    qkv_proj_->load_state_dict(state_dict, {"q_proj.", "k_proj.", "v_proj."});
+    o_proj_->load_state_dict(state_dict.select("o_proj."));
+  }
+
+  void verify_loaded_weights(const std::string& prefix) const {
+    qkv_proj_->verify_loaded_weights(prefix + "[q_proj,k_proj,v_proj].");
+    o_proj_->verify_loaded_weights(prefix + "o_proj.");
+  }
+
+ private:
+  // parameter members, must be registered
+  ColumnParallelLinear qkv_proj_{nullptr};
+
+  RowParallelLinear o_proj_{nullptr};
+
+  // module members without parameters
+  AttentionWithRoPE atten_{nullptr};
+
+  // size for q, k, v
+  std::vector<int64_t> qkv_sizes_;
+};
+TORCH_MODULE(YiAttention);
+
+class YiDecoderLayerImpl : public torch::nn::Module {
+ public:
+  YiDecoderLayerImpl(const ModelArgs& args,
+                     const QuantizationArgs& quant_args,
+                     const ParallelArgs& parallel_args,
+                     torch::ScalarType dtype,
+                     const torch::Device& device) {
+    // register submodules
+    self_attn_ = register_module(
+        "self_attn",
+        YiAttention(args, quant_args, parallel_args, dtype, device));
+    mlp_ = register_module(
+        "mlp", YiMLP(args, quant_args, parallel_args, dtype, device));
+    ln1_ = register_module(
+        "ln1", RMSNorm(args.hidden_size(), args.rms_norm_eps(), dtype, device));
+    ln2_ = register_module(
+        "ln2", RMSNorm(args.hidden_size(), args.rms_norm_eps(), dtype, device));
+  }
+
+  torch::Tensor forward(torch::Tensor x,
+                        torch::Tensor positions,
+                        KVCache& kv_cache,
+                        const InputParameters& input_params) {
+    auto h = x + self_attn_(ln1_(x), positions, kv_cache, input_params);
+    return h + mlp_(ln2_(h));
+  }
+
+  // load the weight from the checkpoint
+  void load_state_dict(const StateDict& state_dict) {
+    // call each submodule's load_state_dict function
+    self_attn_->load_state_dict(state_dict.select("self_attn."));
+    mlp_->load_state_dict(state_dict.select("mlp."));
+    ln1_->load_state_dict(state_dict.select("ln1."));
+    ln2_->load_state_dict(state_dict.select("ln2."));
+  }
+
+  void verify_loaded_weights(const std::string& prefix) const {
+    self_attn_->verify_loaded_weights(prefix + "self_attn.");
+    mlp_->verify_loaded_weights(prefix + "mlp.");
+    ln1_->verify_loaded_weights(prefix + "ln1.");
+    ln2_->verify_loaded_weights(prefix + "ln2.");
+  }
+
+ private:
+  // parameter members, must be registered
+  YiAttention self_attn_{nullptr};
+
+  YiMLP mlp_{nullptr};
+
+  RMSNorm ln1_{nullptr};
+
+  RMSNorm ln2_{nullptr};
+};
+TORCH_MODULE(YiDecoderLayer);
+
+class YiModelImpl : public torch::nn::Module {
+ public:
+  YiModelImpl(const ModelArgs& args,
+              const QuantizationArgs& quant_args,
+              const ParallelArgs& parallel_args,
+              torch::ScalarType dtype,
+              const torch::Device& device) {
+    // register submodules
+    embed_tokens_ = register_module("embed_tokens",
+                                    ParallelEmbedding(args.vocab_size(),
+                                                      args.hidden_size(),
+                                                      parallel_args,
+                                                      dtype,
+                                                      device));
+    blocks_ = register_module("layers", torch::nn::ModuleList());
+    layers_.reserve(args.n_layers());
+    for (int32_t i = 0; i < args.n_layers(); i++) {
+      auto block =
+          YiDecoderLayer(args, quant_args, parallel_args, dtype, device);
+      layers_.push_back(block);
+      blocks_->push_back(block);
+    }
+    norm_ = register_module(
+        "norm",
+        RMSNorm(args.hidden_size(), args.rms_norm_eps(), dtype, device));
+  }
+
+  // tokens: [num_tokens]
+  // positions: [num_tokens] token pos in the sequence
+  torch::Tensor forward(torch::Tensor tokens,
+                        torch::Tensor positions,
+                        std::vector<KVCache>& kv_caches,
+                        const InputParameters& input_params) {
+    auto h = embed_tokens_(tokens);
+    for (size_t i = 0; i < layers_.size(); i++) {
+      auto& layer = layers_[i];
+      h = layer(h, positions, kv_caches[i], input_params);
+    }
+    return norm_(h);
+  }
+
+  // load the weight from the checkpoint
+  void load_state_dict(const StateDict& state_dict) {
+    embed_tokens_->load_state_dict(state_dict.select("embed_tokens."));
+    // call each layer's load_state_dict function
+    for (int i = 0; i < layers_.size(); i++) {
+      layers_[i]->load_state_dict(
+          state_dict.select("layers." + std::to_string(i) + "."));
+    }
+    norm_->load_state_dict(state_dict.select("norm."));
+  }
+
+  void verify_loaded_weights(const std::string& prefix) const {
+    embed_tokens_->verify_loaded_weights(prefix + "embed_tokens.");
+    for (int i = 0; i < layers_.size(); i++) {
+      layers_[i]->verify_loaded_weights(prefix + "layers." + std::to_string(i) +
+                                        ".");
+    }
+    norm_->verify_loaded_weights(prefix + "norm.");
+  }
+
+ private:
+  // parameter members, must be registered
+  ParallelEmbedding embed_tokens_{nullptr};
+
+  torch::nn::ModuleList blocks_{nullptr};
+  // hold same data but different type as blocks_ to avoid type cast
+  std::vector<YiDecoderLayer> layers_;
+
+  RMSNorm norm_{nullptr};
+};
+TORCH_MODULE(YiModel);
+
+class YiForCausalLMImpl : public torch::nn::Module {
+ public:
+  YiForCausalLMImpl(const ModelArgs& args,
+                    const QuantizationArgs& quant_args,
+                    const ParallelArgs& parallel_args,
+                    torch::ScalarType dtype,
+                    const torch::Device& device) {
+    // register submodules
+    model_ = register_module(
+        "model", YiModel(args, quant_args, parallel_args, dtype, device));
+
+    lm_head_ = register_module("lm_head",
+                               ColumnParallelLinear(args.hidden_size(),
+                                                    args.vocab_size(),
+                                                    /*bias=*/false,
+                                                    /*gather_output=*/true,
+                                                    parallel_args,
+                                                    dtype,
+                                                    device));
+  }
+
+  // tokens: [num_tokens]
+  // positions: [num_tokens] token pos in the sequence
+  torch::Tensor forward(torch::Tensor tokens,
+                        torch::Tensor positions,
+                        std::vector<KVCache>& kv_caches,
+                        const InputParameters& input_params) {
+    auto h = model_(tokens, positions, kv_caches, input_params);
+    // select last token for each sequence
+    h = h.index_select(/*dim=*/0, input_params.last_token_indicies);
+    return lm_head_(h);
+  }
+
+  // load the weight from the checkpoint
+  void load_state_dict(const StateDict& state_dict) {
+    model_->load_state_dict(state_dict.select("model."));
+    lm_head_->load_state_dict(state_dict.select("lm_head."));
+  }
+
+  void verify_loaded_weights() const {
+    model_->verify_loaded_weights("model.");
+    lm_head_->verify_loaded_weights("lm_head.");
+  }
+
+ private:
+  // parameter members, must be registered
+  YiModel model_{nullptr};
+
+  ColumnParallelLinear lm_head_{nullptr};
+};
+TORCH_MODULE(YiForCausalLM);
+
+// register the causal model
+REGISTER_CAUSAL_MODEL(Yi, YiForCausalLM);
+// register the model args
+// example config:
+// https://huggingface.co/01-ai/Yi-6B/blob/main/config.json
+REGISTER_MODEL_ARGS(Yi, [&] {
+  LOAD_ARG_OR(dtype, "torch_dtype", "");
+  LOAD_ARG_OR(vocab_size, "vocab_size", 64000);
+  LOAD_ARG_OR(hidden_size, "hidden_size", 4096);
+  LOAD_ARG_OR(n_layers, "num_hidden_layers", 32);
+  LOAD_ARG_OR(n_heads, "num_attention_heads", 32);
+  LOAD_OPTIONAL_ARG(n_kv_heads, "num_key_value_heads");
+  LOAD_ARG_OR(intermediate_size, "intermediate_size", 11008);
+  LOAD_ARG_OR(hidden_act, "hidden_act", "silu");
+  LOAD_ARG_OR(max_position_embeddings, "max_position_embeddings", 2048);
+  LOAD_ARG_OR(rms_norm_eps, "rms_norm_eps", 1e-5);
+  LOAD_ARG_OR(bos_token_id, "bos_token_id", 1);
+  LOAD_ARG_OR(eos_token_id, "eos_token_id", 2);
+  LOAD_ARG_OR(rope_theta, "rope_theta", 5000000.0f);
+  LOAD_ARG_OR(rope_scaling, "rope_scaling", 1.0f);
+});
+
+}  // namespace llm::hf
diff --git a/src/models/llama.h b/src/models/llama.h
index b6a70678..904d2255 100644
--- a/src/models/llama.h
+++ b/src/models/llama.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <c10/core/ScalarType.h>
 #include <torch/torch.h>
 
 #include "layers/activation.h"
@@ -364,6 +365,7 @@ REGISTER_CAUSAL_MODEL(llama2, LlamaForCausalLM);
 REGISTER_DIALOG(llama2, Llama2Dialog);
 
 REGISTER_MODEL_ARGS(llama2, [&] {
+  LOAD_ARG_OR(dtype, "torch_dtype", torch::toString(torch::kBFloat16));
   LOAD_ARG_OR(vocab_size, "vocab_size", 32000);
   LOAD_ARG_OR(hidden_size, "dim", 4096);
   LOAD_ARG_OR(n_layers, "n_layers", 32);
diff --git a/src/models/model_registry.cpp b/src/models/model_registry.cpp
index 142ed9db..eb10cf68 100644
--- a/src/models/model_registry.cpp
+++ b/src/models/model_registry.cpp
@@ -12,6 +12,7 @@
 #include "huggingface/llama.h"
 #include "huggingface/mistral.h"
 #include "huggingface/mpt.h"
+#include "huggingface/yi.h"
 #include "llama.h"
 
 namespace llm {
diff --git a/src/sampling/logits_processor.h b/src/sampling/logits_processor.h
index f65164b6..c367730b 100644
--- a/src/sampling/logits_processor.h
+++ b/src/sampling/logits_processor.h
@@ -175,6 +175,7 @@ class TemperatureLogitsProcessor : public LogitsProcessor {
   torch::Tensor temperatures_;
 };
 
+// TODO: move topk and topp into sampler
 class TopPLogitsProcessor : public LogitsProcessor {
  public:
   TopPLogitsProcessor(
@@ -182,7 +183,7 @@ class TopPLogitsProcessor : public LogitsProcessor {
       torch::ScalarType dtype,
       const torch::Device& device,
       float filter_value = -std::numeric_limits<float>::infinity(),
-      int min_tokens_to_keep = 1)
+      int64_t min_tokens_to_keep = 1)
       : filter_value(filter_value), min_tokens_to_keep(min_tokens_to_keep) {
     top_p_opposite =
         1.0 -
@@ -221,7 +222,7 @@ class TopPLogitsProcessor : public LogitsProcessor {
   // are filtered
   float filter_value;
   // the minimum number of tokens to keep
-  int min_tokens_to_keep;
+  int64_t min_tokens_to_keep;
 };
 
 class TopKLogitsProcessor : public LogitsProcessor {
@@ -229,7 +230,7 @@ class TopKLogitsProcessor : public LogitsProcessor {
   // top_k: input is 1-based, 0 means no filtering or disable filtering
   TopKLogitsProcessor(
       const std::vector<int64_t>& top_k,
-      torch::ScalarType dtype,
+      torch::ScalarType /*dtype*/,
       const torch::Device& device,
       float filter_value = -std::numeric_limits<float>::infinity(),
       int64_t min_tokens_to_keep = 1)
@@ -249,7 +250,8 @@ class TopKLogitsProcessor : public LogitsProcessor {
       disabled.push_back(val == 0 ? 1 : 0);
     }
 
-    top_k_ = torch::tensor(adjusted_top_k, torch::dtype(dtype).device(device))
+    top_k_ = torch::tensor(adjusted_top_k,
+                           torch::dtype(torch::kInt64).device(device))
                  .unsqueeze(1);
 
     if (std::any_of(
diff --git a/src/server/handlers/chat_handler.cpp b/src/server/handlers/chat_handler.cpp
index 9388ac3f..70b9dccd 100644
--- a/src/server/handlers/chat_handler.cpp
+++ b/src/server/handlers/chat_handler.cpp
@@ -92,6 +92,8 @@ std::unique_ptr<Request> grpc_request_to_request(ChatCallData* call_data,
   auto dialog_factory =
       ModelRegistry::get_dialog_factory(model_args.model_type());
   if (dialog_factory == nullptr) {
+    // TODO: return meaningful error message to client
+    // Error: Chat template has not configured, please use /completion API
     GLOG(ERROR) << "Failed to get dialog factory for model type: "
                 << model_args.model_type();
     return nullptr;
diff --git a/src/server/http_server.cpp b/src/server/http_server.cpp
index f5eb1843..eaf7eb22 100644
--- a/src/server/http_server.cpp
+++ b/src/server/http_server.cpp
@@ -62,6 +62,7 @@ bool HttpServer::Start(uint16_t port, int32_t num_threads) {
   break_ev_ = event_new(evbase_, fds_[0], EV_READ, StopCallback, evbase_);
   event_add(break_ev_, nullptr);
   thread_ = std::thread(event_base_loop, evbase_, 0);
+  GLOG(INFO) << "Started http server on 0.0.0.0:" << port;
   return true;
 }
 
diff --git a/src/server/main.cpp b/src/server/main.cpp
index ad7581fd..cc50fbdb 100644
--- a/src/server/main.cpp
+++ b/src/server/main.cpp
@@ -1,4 +1,5 @@
 #include <absl/strings/str_split.h>
+#include <c10/core/Device.h>
 #include <folly/init/Init.h>
 #include <gflags/gflags.h>
 #include <torch/torch.h>
@@ -25,7 +26,10 @@ DEFINE_string(model_id, "", "hf model name.");
 
 DEFINE_string(model_path, "", "hf model path to the model file.");
 
-DEFINE_string(device, "cuda:0", "Device to run the model on.");
+DEFINE_string(device,
+              "auto",
+              "Device to run the model on, e.g. cpu, cuda:0, cuda:0,cuda:1, or "
+              "auto to use all available gpus.");
 
 DEFINE_int32(http_port, 9999, "Port for http server.");
 DEFINE_int32(grpc_port, 8888, "Port for grpc server.");
@@ -37,6 +41,37 @@ void shutdown_handler(int signal) {
   running.store(false, std::memory_order_relaxed);
 }
 
+std::vector<torch::Device> parse_devices(const std::string& device_str) {
+  std::vector<torch::Device> devices;
+  if (device_str == "auto") {
+    // use all available gpus if any
+    const auto num_gpus = torch::cuda::device_count();
+    if (num_gpus == 0) {
+      GLOG(INFO) << "no gpus found, using cpu.";
+      return {torch::kCPU};
+    }
+    devices.reserve(num_gpus);
+    for (int i = 0; i < num_gpus; ++i) {
+      devices.emplace_back(torch::kCUDA, i);
+    }
+    GLOG(INFO) << "using " << num_gpus << " gpus.";
+    return devices;
+  }
+
+  // parse device string
+  const std::vector<std::string> device_strs = absl::StrSplit(device_str, ',');
+  std::set<torch::DeviceType> device_types;
+  devices.reserve(device_strs.size());
+  for (const auto& device_str : device_strs) {
+    devices.emplace_back(device_str);
+    device_types.insert(devices.back().type());
+  }
+  GCHECK(!devices.empty()) << "No devices specified.";
+  GCHECK(device_types.size() == 1)
+      << "All devices must be of the same type. Got: " << FLAGS_device;
+  return devices;
+}
+
 int main(int argc, char** argv) {
   // glog and glfag will be initialized in folly::init
   folly::Init init(&argc, &argv);
@@ -83,34 +118,11 @@ int main(int argc, char** argv) {
     GLOG(ERROR) << "Failed to start http server on port " << FLAGS_http_port;
     return -1;
   }
-  GLOG(INFO) << "Started http server on localhost:" << FLAGS_http_port;
-
-  // split device into chunks
-  const std::vector<std::string> device_strs =
-      absl::StrSplit(FLAGS_device, ',');
-  std::vector<torch::Device> devices;
-  devices.reserve(device_strs.size());
-  std::set<torch::DeviceType> device_types;
-  for (const auto& device_str : device_strs) {
-    devices.emplace_back(device_str);
-    device_types.insert(devices.back().type());
-  }
-  GCHECK(!devices.empty()) << "No devices specified.";
-  GCHECK(device_types.size() == 1)
-      << "All devices must be of the same type. Got: " << FLAGS_device;
-
-  // set the default dtype
-  torch::ScalarType dtype{};
-  if (devices[0].is_cpu()) {
-    // always use float32 on CPU since float16 is not supported
-    dtype = torch::kFloat;
-    GLOG(INFO) << "Using float32 on CPU.";
-  } else {
-    dtype = torch::kHalf;
-  }
 
+  // parse devices
+  const auto devices = parse_devices(FLAGS_device);
   // create engine
-  auto engine = std::make_unique<Engine>(dtype, devices);
+  auto engine = std::make_unique<Engine>(devices);
   GCHECK(engine->init(FLAGS_model_path));
 
   // create scheduler and grpc handlers
@@ -126,11 +138,11 @@ int main(int argc, char** argv) {
                          std::move(chat_handler),
                          std::move(models_handler));
   GrpcServer::Options options;
-  options.address = "localhost";
+  options.address = "0.0.0.0";
   options.port = FLAGS_grpc_port;
 
   if (!grpc_server.start(options)) {
-    GLOG(ERROR) << "failed to start grpc server";
+    GLOG(ERROR) << "failed to start grpc server on port " << FLAGS_grpc_port;
     return -1;
   }
   // install graceful shutdown handler
diff --git a/src/server/simple.cpp b/src/server/simple.cpp
index 09d280c0..9a31b435 100644
--- a/src/server/simple.cpp
+++ b/src/server/simple.cpp
@@ -55,16 +55,6 @@ int main(int argc, char* argv[]) {
   GCHECK(device_types.size() == 1)
       << "All devices must be of the same type. Got: " << FLAGS_device;
 
-  // set the default dtype
-  torch::ScalarType dtype{};
-  if (devices[0].is_cpu()) {
-    // always use float32 on CPU since float16 is not supported
-    dtype = torch::kFloat;
-    GLOG(INFO) << "Using float32 on CPU.";
-  } else {
-    dtype = torch::kHalf;
-  }
-
   // check if model path exists
   std::string model_path = FLAGS_model_name_or_path;
   if (!std::filesystem::exists(model_path)) {
@@ -72,7 +62,7 @@ int main(int argc, char* argv[]) {
     model_path = llm::hf::download_model(FLAGS_model_name_or_path);
   }
 
-  llm::Engine engine(dtype, devices);
+  llm::Engine engine(devices);
   GCHECK(engine.init(model_path));
   auto tokenizer = engine.tokenizer();
   llm::BlockManager* block_manager = engine.block_manager();
diff --git a/src/tokenizer/sentencepiece_tokenizer.cpp b/src/tokenizer/sentencepiece_tokenizer.cpp
index b29e969a..b41c3c3e 100644
--- a/src/tokenizer/sentencepiece_tokenizer.cpp
+++ b/src/tokenizer/sentencepiece_tokenizer.cpp
@@ -6,8 +6,9 @@
 namespace llm {
 
 SentencePieceTokenizer::SentencePieceTokenizer(
-    const std::string& vocab_file_path)
-    : vocab_file_path_(vocab_file_path) {
+    const std::string& vocab_file_path,
+    bool prepend_bos)
+    : vocab_file_path_(vocab_file_path), prepend_bos_(prepend_bos) {
   const auto status = sp_processor_.Load(vocab_file_path);
   if (!status.ok()) {
     GLOG(FATAL) << "Failed to load SentencePiece model from " << vocab_file_path
@@ -16,7 +17,8 @@ SentencePieceTokenizer::SentencePieceTokenizer(
 }
 
 std::unique_ptr<Tokenizer> SentencePieceTokenizer::clone() const {
-  return std::make_unique<SentencePieceTokenizer>(this->vocab_file_path_);
+  return std::make_unique<SentencePieceTokenizer>(this->vocab_file_path_,
+                                                  this->prepend_bos_);
 }
 
 bool SentencePieceTokenizer::encode(const std::string_view& text,
@@ -28,7 +30,9 @@ bool SentencePieceTokenizer::encode(const std::string_view& text,
     return false;
   }
   // prepend bos token
-  ids->insert(ids->begin(), sp_processor_.bos_id());
+  if (prepend_bos_) {
+    ids->insert(ids->begin(), sp_processor_.bos_id());
+  }
   return true;
 }
 
diff --git a/src/tokenizer/sentencepiece_tokenizer.h b/src/tokenizer/sentencepiece_tokenizer.h
index 3f349b3c..2ff366f8 100644
--- a/src/tokenizer/sentencepiece_tokenizer.h
+++ b/src/tokenizer/sentencepiece_tokenizer.h
@@ -8,7 +8,7 @@ namespace llm {
 // a tokenizer that uses google/SentencePiece
 class SentencePieceTokenizer : public Tokenizer {
  public:
-  explicit SentencePieceTokenizer(const std::string& vocab_file_path);
+  explicit SentencePieceTokenizer(const std::string& vocab_file_path, bool prepend_bos);
 
   bool encode(const std::string_view& text,
               std::vector<int32_t>* ids) const override;
@@ -23,6 +23,8 @@ class SentencePieceTokenizer : public Tokenizer {
   std::string vocab_file_path_;
 
   sentencepiece::SentencePieceProcessor sp_processor_;
+
+  bool prepend_bos_ = false;
 };
 
 }  // namespace llm
diff --git a/third_party/vllm b/third_party/vllm
index 90eb3f43..651c614a 160000
--- a/third_party/vllm
+++ b/third_party/vllm
@@ -1,1 +1,1 @@
-Subproject commit 90eb3f43ca5228647c834243d98881d53c7745d0
+Subproject commit 651c614aa43e497a2e2aab473493ba295201ab20
diff --git a/tools/run_in_docker.sh b/tools/run_in_docker.sh
index 3d039e41..783a82c2 100755
--- a/tools/run_in_docker.sh
+++ b/tools/run_in_docker.sh
@@ -60,7 +60,7 @@ RUN_OPTS+=(--rm -it --network=host)
 # output data that might be used by other scripts/binaries
 RUN_OPTS+=("-v $(pwd):$(pwd)")
 RUN_OPTS+=("-v /tmp:/tmp")
-CMD="bash -c 'cd $(pwd); VCPKG_DEFAULT_BINARY_CACHE=$(pwd)/.vcpkg_bincache $@'"
+CMD="bash -c 'cd $(pwd); VCPKG_DEFAULT_BINARY_CACHE=$(pwd)/.vcpkg/bincache $@'"
 
 [[ "${CMD}" = "" ]] && usage
 [[ ! -x $(command -v docker) ]] && echo "ERROR: 'docker' command missing from PATH." && usage


Test patch:
diff --git a/src/models/huggingface/mpt_test.cpp b/src/models/huggingface/mpt_test.cpp
deleted file mode 100644
index cdd383ab..00000000
--- a/src/models/huggingface/mpt_test.cpp
+++ /dev/null
@@ -1,14 +0,0 @@
-#include "mpt.h"
-
-#include <gtest/gtest.h>
-
-namespace llm {
-TEST(MPTTest, ModelInit) {
-  // TODO: test if we can load the model with different device types and dtypes
-  // hf::MPTModel mpt(const ModelArgs& args,
-  //                  const ParallelArgs& parallel_args,
-  //                  torch::ScalarType dtype,
-  //                  const torch::Device& device);
-}
-
-}  // namespace llm
diff --git a/src/models/llama_test.cpp b/src/models/llama_test.cpp
deleted file mode 100644
index 1e29a61c..00000000
--- a/src/models/llama_test.cpp
+++ /dev/null
@@ -1,14 +0,0 @@
-#include "llama.h"
-
-#include <gtest/gtest.h>
-
-namespace llm {
-TEST(LlamaTest, ModelInit) {
-  // TODO: test if we can load the model with different device types and dtypes
-  // Llama2Impl(const ModelArgs& args,
-  //            const ParallelArgs& parallel_args,
-  //            torch::ScalarType dtype,
-  //            const torch::Device& device)
-}
-
-}  // namespace llm
diff --git a/src/sampling/logits_processor_test.cpp b/src/sampling/logits_processor_test.cpp
index 2103af87..25be1e41 100644
--- a/src/sampling/logits_processor_test.cpp
+++ b/src/sampling/logits_processor_test.cpp
@@ -30,8 +30,86 @@ TEST(LogitsProcessorTest, Temperature) {
   }
 }
 
-TEST(LogitsProcessorTest, TemperatureKernel) {
-  // TODO: add test for TemperatureLogitsProcessorKernel
+TEST(LogitsProcessorTest, DISABLED_TopK) {
+  // Test TopKLogitsProcessor
+  // Set the random seed
+  torch::manual_seed(100);
+  torch::ScalarType dtype(torch::kHalf);
+  torch::Device device(torch::kCUDA);
+  int64_t batch_size = 5;
+  int64_t vocab_size = 100;
+  const std::vector<int64_t> top_k = {50, 60, 70, 80, 100};
+  const float filter_value = -std::numeric_limits<float>::infinity();
+  TopKLogitsProcessor processor(top_k, dtype, device, filter_value);
+
+  auto logits = torch::randn({batch_size, vocab_size},
+                             torch::dtype(dtype).device(device));
+  torch::Tensor token_ids;
+  auto logits_output = processor(token_ids, logits.clone());
+
+  for (int64_t i = 0; i < batch_size; ++i) {
+    const int64_t k = top_k[i];
+
+    // calculate top k values one by one
+    auto [top_k_values, top_k_indices] =
+        logits[i].topk(k, /*dim=*/-1, /*largest=*/true, /*sorted=*/true);
+    auto [output_sorted, output_indices] =
+        logits_output[i].sort(/*dim=*/-1, /*descending=*/true);
+    // top k values should be the same
+    ASSERT_TRUE(
+        torch::equal(output_sorted.slice(/*dim=*/0, 0, k), top_k_values));
+    // all remaining values should be filter_value
+    auto masked_output = output_sorted.slice(/*dim=*/0, k);
+    ASSERT_TRUE(torch::equal(masked_output,
+                             torch::full_like(masked_output, filter_value)));
+  }
+}
+
+TEST(LogitsProcessorTest, DISABLED_TopP) {
+  // Test TopPLogitsProcessor
+  // Set the random seed
+  torch::manual_seed(100);
+  torch::ScalarType dtype(torch::kHalf);
+  torch::Device device(torch::kCUDA);
+  int64_t batch_size = 4;
+  int64_t vocab_size = 100;
+  int64_t min_tokens_to_keep = 1;
+  std::vector<float> top_p = {0.5, 0.7, 0.9, 1.0};
+  const float filter_value = -std::numeric_limits<float>::infinity();
+  TopPLogitsProcessor processor(
+      top_p, dtype, device, filter_value, min_tokens_to_keep);
+
+  auto logits = torch::randn({batch_size, vocab_size},
+                             torch::dtype(dtype).device(device));
+  torch::Tensor token_ids;
+  auto logits_output = processor(token_ids, logits.clone());
+
+  // verify result one by one
+  for (int64_t i = 0; i < batch_size; ++i) {
+    const float p = top_p[i];
+
+    // calculate number of values to keep (k)
+    auto probs = torch::softmax(logits[i], /*dim=*/-1);
+    // calculate top p values one by one
+    auto [sorted_probs, sorted_indices] =
+        probs.sort(/*dim=*/-1, /*descending=*/true);
+    auto probs_sum = sorted_probs.cumsum(/*dim=*/-1);
+    torch::Tensor mask = (probs_sum - sorted_probs) <= p;
+    const int64_t k = std::max(mask.sum().item<int64_t>(), min_tokens_to_keep);
+
+    // gather top logits value
+    auto [sorted_logits, sorted_logits_indices] =
+        logits[i].sort(/*dim=*/-1, /*descending=*/true);
+
+    auto [output_sorted, output_indices] =
+        logits_output[i].sort(/*dim=*/-1, /*descending=*/true);
+    ASSERT_TRUE(torch::equal(output_sorted.slice(/*dim=*/0, 0, k),
+                             sorted_logits.slice(/*dim=*/0, 0, k)));
+    // all remaining values should be filter_value
+    auto masked_output = output_sorted.slice(/*dim=*/0, k);
+    ASSERT_TRUE(torch::equal(masked_output,
+                             torch::full_like(masked_output, filter_value)));
+  }
 }
 
 }  // namespace llm
diff --git a/src/tokenizer/sentencepiece_tokenizer_test.cpp b/src/tokenizer/sentencepiece_tokenizer_test.cpp
index e88d5c70..1d19019e 100644
--- a/src/tokenizer/sentencepiece_tokenizer_test.cpp
+++ b/src/tokenizer/sentencepiece_tokenizer_test.cpp
@@ -5,7 +5,7 @@
 namespace llm {
 
 TEST(SentencePieceTokenizerTest, EncodeDecodeTest) {
-  SentencePieceTokenizer tokenizer("data/tokenizer.model");
+  SentencePieceTokenizer tokenizer("data/tokenizer.model", /*prepend_bos=*/true);
   std::vector<int> ids;
   ASSERT_TRUE(tokenizer.encode("Hello, world!", &ids));
   const std::vector<int> desired_ids = {1, 15043, 29892, 3186, 29991};
@@ -16,7 +16,7 @@ TEST(SentencePieceTokenizerTest, EncodeDecodeTest) {
 }
 
 TEST(SentencePieceTokenizerTest, CJKTest) {
-  SentencePieceTokenizer tokenizer("data/tokenizer.model");
+  SentencePieceTokenizer tokenizer("data/tokenizer.model", /*prepend_bos=*/true);
   const std::string test_text = "你好，世界！";
   std::vector<int> ids;
   ASSERT_TRUE(tokenizer.encode(test_text, &ids));