Repository: li-plus/chatglm.cpp (language: python)

Issue:
chatGLM2 & llama.cpp
请问这个REPO可以直接支持chatGLM2吗？
另外最近2周llama.cpp里很多关于cublas的优化提速代码是否已经merge进来了？
MacBook Pro 编译失败
在苹果芯片下，clang: error: the clang compiler does not support '-march=native'，请问怎么解决呢


Discussion / hints:

Fix patch:
diff --git a/.gitmodules b/.gitmodules
index a69386f..8843620 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,6 +1,3 @@
-[submodule "third_party/ggml"]
-	path = third_party/ggml
-	url = https://github.com/li-plus/ggml.git
 [submodule "third_party/sentencepiece"]
 	path = third_party/sentencepiece
 	url = https://github.com/google/sentencepiece.git
@@ -8,3 +5,6 @@
 	path = third_party/pybind11
 	url = https://github.com/pybind/pybind11.git
 	branch = stable
+[submodule "third_party/ggml"]
+	path = third_party/ggml
+	url = https://github.com/ggerganov/ggml.git
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 278c45c..b30a07b 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -7,7 +7,6 @@ set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin CACHE STRING "")
 
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -Wall")
-set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -march=native")
 
 if (NOT CMAKE_BUILD_TYPE)
     set(CMAKE_BUILD_TYPE Release)
@@ -78,6 +77,7 @@ set(PY_SOURCES
     ${PROJECT_SOURCE_DIR}/setup.py
     ${PROJECT_SOURCE_DIR}/cli_chat.py
     ${PROJECT_SOURCE_DIR}/web_demo.py
+    ${PROJECT_SOURCE_DIR}/test_convert.py
     ${PROJECT_SOURCE_DIR}/test_chatglm_cpp.py)
 add_custom_target(lint
     COMMAND clang-format -i ${CPP_SOURCES}
diff --git a/README.md b/README.md
index 6f5c735..ffa8280 100644
--- a/README.md
+++ b/README.md
@@ -3,7 +3,7 @@
 [![CMake](https://github.com/li-plus/chatglm.cpp/actions/workflows/cmake.yml/badge.svg)](https://github.com/li-plus/chatglm.cpp/actions/workflows/cmake.yml)
 [![License: MIT](https://img.shields.io/badge/license-MIT-blue)](LICENSE)
 
-A C++ implementation of [ChatGLM-6B](https://github.com/THUDM/ChatGLM-6B). Run int4 model inference on your macBook in a minute.
+C++ implementation of [ChatGLM-6B](https://github.com/THUDM/ChatGLM-6B) and [ChatGLM2-6B](https://github.com/THUDM/ChatGLM2-6B) for real-time chatting on your MacBook.
 
 ![demo](docs/demo.gif)
 
@@ -30,9 +30,12 @@ git submodule update --init --recursive
 
 **Quantize Model**
 
-Use `convert.py` to transform ChatGLM-6B to quantized GGML format. For example, to convert the fp16 ChatGLM-6B model to q4_0 (quantized int4) GGML model, run:
+Use `convert.py` to transform ChatGLM-6B or ChatGLM2-6B into quantized GGML format. For example, to convert the fp16 base model to q4_0 (quantized int4) GGML model, run:
 ```sh
+# For ChatGLM-6B
 python3 convert.py -i THUDM/chatglm-6b -t q4_0 -o chatglm-ggml.bin
+# For ChatGLM2-6B
+python3 convert.py -i THUDM/chatglm2-6b -t q4_0 -o chatglm2-ggml.bin
 ```
 
 For LoRA model, specify `-l <lora_model_name_or_path>` flag to merge your LoRA weights into the base model.
@@ -45,13 +48,15 @@ cmake -B build
 cmake --build build -j
 ```
 
-Now you may chat with the quantized ChatGLM model by running:
+Now you may chat with the quantized ChatGLM-6B model by running:
 ```sh
-./build/bin/main -m chatglm-ggml.bin -p 你好
+./build/bin/main -m chatglm-ggml.bin -p 你好                            # ChatGLM-6B
 # 你好👋！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。
+./build/bin/main -m chatglm2-ggml.bin -p 你好 --top_p 0.8 --temp 0.8    # ChatGLM2-6B
+# 你好👋！我是人工智能助手 ChatGLM2-6B，很高兴见到你，欢迎问我任何问题。
 ```
 
-To run ChatGLM in interactive mode, add the `-i` flag:
+To run the model in interactive mode, add the `-i` flag. For example:
 ```sh
 ./build/bin/main -m chatglm-ggml.bin -i
 ```
@@ -95,7 +100,7 @@ pip install git+https://github.com/li-plus/chatglm.cpp.git@main
 pip install .
 ```
 
-Run the Python example to chat with the quantized ChatGLM model:
+Run the Python example to chat with the quantized model:
 ```sh
 python3 cli_chat.py -m chatglm-ggml.bin -i
 ```
@@ -105,18 +110,30 @@ You may also launch a web demo to chat in your browser:
 python3 web_demo.py -m chatglm-ggml.bin
 ```
 
+For ChatGLM2, change the model path to `chatglm2-ggml.bin` and everything works fine.
+
 ![web_demo](docs/web_demo.jpg)
 
 ## Performance
 
 Measured on a Linux server with Intel(R) Xeon(R) Platinum 8260 CPU @ 2.40GHz using 16 threads.
 
+ChatGLM-6B:
+
 |           | Q4_0  | Q8_0  | F16  | F32  |
 |-----------|-------|-------|------|------|
-| ms/token  | 92    | 130   | 217  | 399  |
+| ms/token  | 74    | 114   | 189  | 357  |
 | file size | 3.3GB | 6.2GB | 12GB | 23GB |
 | mem usage | 4.0GB | 6.9GB | 13GB | 24GB |
 
+ChatGLM2-6B:
+
+|           | Q4_0  | Q8_0  | F16  | F32  |
+|-----------|-------|-------|------|------|
+| ms/token  | 64    | 106   | 189  | 372  |
+| file size | 3.3GB | 6.2GB | 12GB | 24GB |
+| mem usage | 3.4GB | 6.2GB | 12GB | 23GB |
+
 ## Development
 
 * To perform unit tests, add the CMake flag `-DCHATGLM_ENABLE_TESTING=ON`, recompile, and run `./build/bin/chatglm_test`. For benchmark only, run `./build/bin/chatglm_test --gtest_filter=ChatGLM.benchmark`.
@@ -125,5 +142,5 @@ Measured on a Linux server with Intel(R) Xeon(R) Platinum 8260 CPU @ 2.40GHz usi
 
 ## Acknowledgements
 
-* This project is greatly inspired by [ggerganov](https://github.com/ggerganov)'s [llama.cpp](https://github.com/ggerganov/llama.cpp) and is based on his NN library [ggml](https://github.com/ggerganov/ggml).
-* Thank [THUDM](https://github.com/THUDM) for the amazing [ChatGLM-6B](https://github.com/THUDM/ChatGLM-6B) and for releasing the model sources and checkpoints.
+* This project is greatly inspired by [@ggerganov](https://github.com/ggerganov)'s [llama.cpp](https://github.com/ggerganov/llama.cpp) and is based on his NN library [ggml](https://github.com/ggerganov/ggml).
+* Thank [@THUDM](https://github.com/THUDM) for the amazing [ChatGLM-6B](https://github.com/THUDM/ChatGLM-6B) and [ChatGLM2-6B](https://github.com/THUDM/ChatGLM2-6B) and for releasing the model sources and checkpoints.
diff --git a/chatglm.cpp b/chatglm.cpp
index 2f92d6d..95769bc 100644
--- a/chatglm.cpp
+++ b/chatglm.cpp
@@ -22,13 +22,6 @@
 
 namespace chatglm {
 
-template <typename T>
-static inline T read_as(std::istream &is) {
-    T data;
-    is.read((char *)&data, sizeof(data));
-    return data;
-}
-
 static std::string shape_to_string(ggml_tensor *tensor) {
     std::ostringstream oss;
     oss << '[';
@@ -54,10 +47,16 @@ std::string to_string(ggml_tensor *tensor, bool with_data) {
     oss << "ggml_tensor(";
 
     if (with_data) {
-        oss << "[";
+        if (tensor->n_dims > 3)
+            oss << "[";
         for (int i3 = 0; i3 < tensor->ne[3]; i3++) {
+            if (tensor->n_dims > 2)
+                oss << (i3 > 0 ? ",\n\n[" : "[");
             for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
+                if (tensor->n_dims > 1)
+                    oss << (i2 > 0 ? ",\n\n[" : "[");
                 for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
+                    oss << (i1 > 0 ? ",\n[" : "[");
                     for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
                         auto ptr = (char *)tensor->data + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] +
                                    i0 * tensor->nb[0];
@@ -69,15 +68,21 @@ std::string to_string(ggml_tensor *tensor, bool with_data) {
                         } else {
                             CHATGLM_THROW << "unimplemented";
                         }
-                        oss << (i3 + i2 + i1 + i0 > 0 ? ", " : "") << std::fixed << std::setprecision(4) << val;
+                        oss << (i0 > 0 ? ", " : "") << std::setw(7) << std::fixed << std::setprecision(4) << val;
                     }
+                    oss << "]";
                 }
+                if (tensor->n_dims > 1)
+                    oss << "]";
             }
+            if (tensor->n_dims > 2)
+                oss << "]";
         }
-        oss << "], ";
+        if (tensor->n_dims > 3)
+            oss << "]";
     }
 
-    oss << "shape=" << shape_to_string(tensor) << ", stride=" << strides_to_string(tensor) << ")";
+    oss << ", shape=" << shape_to_string(tensor) << ", stride=" << strides_to_string(tensor) << ")";
     return oss.str();
 }
 
@@ -124,29 +129,161 @@ void TextStreamer::end() {
     print_len_ = 0;
 }
 
-// ===== tokenizer =====
+MappedFile::MappedFile(const std::string &path) {
+    int fd = open(path.c_str(), O_RDONLY);
+    CHATGLM_CHECK(fd > 0) << "cannot open file " << path << ": " << strerror(errno);
+
+    struct stat sb;
+    CHATGLM_CHECK(fstat(fd, &sb) == 0) << strerror(errno);
+    size = sb.st_size;
+
+    data = (char *)mmap(nullptr, size, PROT_READ, MAP_SHARED, fd, 0);
+    CHATGLM_CHECK(data != MAP_FAILED) << strerror(errno);
+
+    CHATGLM_CHECK(close(fd) == 0) << strerror(errno);
+}
+
+MappedFile::~MappedFile() { CHATGLM_CHECK(munmap(data, size) == 0) << strerror(errno); }
+
+void ModelLoader::seek(int64_t offset, int whence) {
+    if (whence == SEEK_SET) {
+        ptr = data + offset;
+    } else if (whence == SEEK_CUR) {
+        ptr += offset;
+    } else if (whence == SEEK_END) {
+        ptr = data + size + offset;
+    } else {
+        CHATGLM_THROW << "invalid seek mode " << whence;
+    }
+}
+
+std::string ModelLoader::read_string(size_t length) {
+    std::string s(ptr, ptr + length);
+    ptr += length;
+    return s;
+}
+
+void ModelLoader::read_tensor(const std::string &name, ggml_tensor *tensor) {
+    // read and check tensor name
+    {
+        int name_size = read_basic<int>();
+        CHATGLM_CHECK(name_size == (int)name.size())
+            << "tensor " << name << " name size mismatch: expect " << name.size() << " but got " << name_size;
+        std::string weight_name = read_string(name_size);
+        CHATGLM_CHECK(weight_name == name) << "tensor name mismatch: expect " << name << " but got " << weight_name;
+    }
+
+    // read and check tensor shape
+    {
+        int ndim = read_basic<int>();
+        CHATGLM_CHECK(ndim == tensor->n_dims)
+            << "tensor " << name << " ndim mismatch: expect " << tensor->n_dims << " but got " << ndim;
+        for (int i = ndim - 1; i >= 0; i--) {
+            int dim_size = read_basic<int>();
+            CHATGLM_CHECK(dim_size == tensor->ne[i]) << "tensor " << name << " shape mismatch at dim " << i
+                                                     << ": expect " << tensor->ne[i] << " but got " << dim_size;
+        }
+    }
+
+    // read and check tensor dtype
+    {
+        ggml_type dtype = (ggml_type)read_basic<int>();
+        CHATGLM_CHECK(dtype == tensor->type)
+            << "tensor " << name << " dtype mismatch: expect " << tensor->type << " but got " << dtype;
+    }
+
+    // map tensor data
+    {
+        constexpr int64_t MEM_ALIGNED = 16;
+        const int64_t data_offset = (tell() + (MEM_ALIGNED - 1)) & ~(MEM_ALIGNED - 1);
+        tensor->data = const_cast<char *const>(data) + data_offset;
+        seek(data_offset + ggml_nbytes(tensor), SEEK_SET);
+    }
+}
+
+// ===== modules =====
+
+ggml_tensor *Embedding::forward(ForwardContext *ctx, ggml_tensor *input) const {
+    ggml_tensor *output = ggml_get_rows(ctx->gctx.get(), weight, input);
+    return output;
+}
+
+ggml_tensor *Linear::forward(ForwardContext *ctx, ggml_tensor *input) const {
+    // input: [seqlen, in_features]
+    ggml_tensor *output = ggml_mul_mat(ctx->gctx.get(), weight, input); // [seqlen, out_features]
+    if (bias) {
+        ggml_tensor *bcast_bias = ggml_view_2d(ctx->gctx.get(), bias, output->ne[0], output->ne[1], 0, 0);
+        output = ggml_add_inplace(ctx->gctx.get(), output, bcast_bias);
+    }
+    return output;
+}
+
+ggml_tensor *LayerNorm::forward(ForwardContext *ctx, ggml_tensor *input) const {
+    // input: [seqlen, normalized_shape]
+    ggml_tensor *output = ggml_norm_inplace(ctx->gctx.get(), input);
+    output = ggml_mul_inplace(ctx->gctx.get(), output, weight);
+    ggml_tensor *bcast_bias = ggml_view_2d(ctx->gctx.get(), bias, output->ne[0], output->ne[1], 0, 0);
+    output = ggml_add_inplace(ctx->gctx.get(), output, bcast_bias);
+    return output;
+}
 
-ChatGLMTokenizer::ChatGLMTokenizer(std::string_view serialized_model_proto, int bos_token_id, int eos_token_id,
-                                   int mask_token_id, int gmask_token_id, int pad_token_id)
-    : bos_token_id(bos_token_id), eos_token_id(eos_token_id), mask_token_id(mask_token_id),
-      gmask_token_id(gmask_token_id), pad_token_id(pad_token_id) {
-    sp = std::make_unique<sentencepiece::SentencePieceProcessor>();
-    const auto status = sp->LoadFromSerializedProto(serialized_model_proto);
+ggml_tensor *RMSNorm::forward(ForwardContext *ctx, ggml_tensor *input) const {
+    ggml_tensor *output = ggml_rms_norm_inplace(ctx->gctx.get(), input);
+    output = ggml_mul_inplace(ctx->gctx.get(), output, weight);
+    return output;
+}
+
+// ===== ChatGLM-6B =====
+
+ChatGLMTokenizer::ChatGLMTokenizer(std::string_view serialized_model_proto) {
+    const auto status = sp.LoadFromSerializedProto(serialized_model_proto);
     CHATGLM_CHECK(status.ok()) << status.ToString();
+
+    bos_token_id = sp.PieceToId("<sop>");
+    eos_token_id = sp.PieceToId("<eop>");
+    mask_token_id = sp.PieceToId("[MASK]");
+    gmask_token_id = sp.PieceToId("[gMASK]");
+    pad_token_id = sp.PieceToId("<pad>");
 }
 
 std::vector<int> ChatGLMTokenizer::encode(const std::string &text) const {
     std::string input = preprocess(text);
     std::vector<int> ids;
-    sp->Encode(input, &ids);
-    ids.push_back(gmask_token_id);
-    ids.push_back(bos_token_id);
+    sp.Encode(input, &ids);
+    ids.insert(ids.end(), {gmask_token_id, bos_token_id});
     return ids;
 }
 
+std::vector<int> ChatGLMTokenizer::encode_history(const std::vector<std::string> &history, int max_length) const {
+    std::string prompt = build_prompt(history);
+    std::vector<int> input_ids = encode(prompt);
+    if ((int)input_ids.size() > max_length) {
+        // sliding window: always take the last max_length tokens
+        input_ids.erase(input_ids.begin(), input_ids.end() - max_length);
+    }
+    return input_ids;
+}
+
+std::string ChatGLMTokenizer::build_prompt(const std::vector<std::string> &history) {
+    CHATGLM_CHECK(history.size() % 2 == 1) << "invalid history size " << history.size();
+
+    std::ostringstream oss_prompt;
+    if (history.size() == 1) {
+        oss_prompt << history.front();
+    } else {
+        for (size_t i = 0; i < history.size(); i += 2) {
+            oss_prompt << "[Round " << i / 2 << "]\n问：" << history[i] << "\n答：";
+            if (i < history.size() - 1) {
+                oss_prompt << history[i + 1] << "\n";
+            }
+        }
+    }
+    return oss_prompt.str();
+}
+
 std::string ChatGLMTokenizer::decode(const std::vector<int> &ids) const {
     std::string text;
-    sp->Decode(ids, &text);
+    sp.Decode(ids, &text);
     text = postprocess(text);
     return text;
 }
@@ -235,210 +372,58 @@ std::string ChatGLMTokenizer::postprocess(const std::string &text) {
     return output;
 }
 
-// ===== model =====
-
-ggml_tensor *Embedding::forward(ForwardContext *ctx, ggml_tensor *input) const {
-    ggml_tensor *output = ggml_get_rows(ctx->gctx, weight, input);
-    return output;
-}
-
-ggml_tensor *Linear::forward(ForwardContext *ctx, ggml_tensor *input) const {
-    // input: [seqlen, in_features]
-    ggml_tensor *output = ggml_mul_mat(ctx->gctx, weight, input); // [seqlen, out_features]
-    ggml_tensor *bcast_bias = ggml_view_2d(ctx->gctx, bias, output->ne[0], output->ne[1], 0, 0);
-    output = ggml_add_inplace(ctx->gctx, output, bcast_bias);
-    return output;
-}
-
-ggml_tensor *LayerNorm::forward(ForwardContext *ctx, ggml_tensor *input) const {
-    // input: [seqlen, normalized_shape]
-    ggml_tensor *output = ggml_norm_inplace(ctx->gctx, input);
-    ggml_tensor *bcast_weight = ggml_view_2d(ctx->gctx, weight, output->ne[0], output->ne[1], 0, 0);
-    output = ggml_mul_inplace(ctx->gctx, output, bcast_weight);
-    ggml_tensor *bcast_bias = ggml_view_2d(ctx->gctx, bias, output->ne[0], output->ne[1], 0, 0);
-    output = ggml_add_inplace(ctx->gctx, output, bcast_bias);
-    return output;
-}
-
-ggml_tensor *GLU::forward(ForwardContext *ctx, ggml_tensor *hidden_states) const {
-    ggml_tensor *output = dense_h_to_4h.forward(ctx, hidden_states);
-    output = ggml_gelu_inplace(ctx->gctx, output);
-    output = dense_4h_to_h.forward(ctx, output);
-    return output;
-}
-
-ggml_tensor *SelfAttention::forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past, int n_ctx) const {
-    int hidden_size = hidden_states->ne[0];
-    int qlen = hidden_states->ne[1];
-    int head_size = hidden_size / num_attention_heads;
-    int rope_dim = head_size / 2;
-
-    ggml_tensor *qkv = query_key_value.forward(ctx, hidden_states); // [qlen, 3 * hidden]
-
-    ggml_tensor *query_layer = ggml_view_3d(ctx->gctx, qkv, head_size, num_attention_heads, qlen,
-                                            3 * head_size * ggml_element_size(qkv), qkv->nb[1], 0);
-    query_layer = ggml_rope_inplace(ctx->gctx, query_layer, n_past, rope_dim, 4, n_ctx); // [qlen, n_head, head_size]
-    query_layer = ggml_permute(ctx->gctx, query_layer, 0, 2, 1, 3);                      // [n_head, qlen, head_size]
-
-    ggml_tensor *key_layer =
-        ggml_view_3d(ctx->gctx, qkv, head_size, num_attention_heads, qlen, 3 * head_size * ggml_element_size(qkv),
-                     qkv->nb[1], head_size * ggml_element_size(qkv));
-    key_layer = ggml_rope_inplace(ctx->gctx, key_layer, n_past, rope_dim, 4, n_ctx); // [qlen, n_head, head_size]
-    key_layer = ggml_permute(ctx->gctx, key_layer, 0, 2, 1, 3);                      // [n_head, qlen, head_size]
-
-    ggml_tensor *value_layer = ggml_view_3d(ctx->gctx, qkv, head_size, num_attention_heads, qlen,
-                                            3 * head_size * ggml_element_size(qkv), qkv->nb[1],
-                                            2 * head_size * ggml_element_size(qkv)); // [qlen, n_head, head_size]
-    value_layer = ggml_permute(ctx->gctx, value_layer, 1, 2, 0, 3);                  // [n_head, head_size, qlen]
-
-    // store key & value to cache
-    ggml_tensor *k_cache_view =
-        ggml_view_3d(ctx->gctx, k_cache, head_size, qlen, num_attention_heads, k_cache->nb[1], k_cache->nb[2],
-                     n_past * head_size * ggml_element_size(k_cache)); // [n_head, qlen, head_size]
-    ggml_build_forward_expand(&ctx->gf, ggml_cpy(ctx->gctx, key_layer, k_cache_view));
-    ggml_tensor *v_cache_view =
-        ggml_view_3d(ctx->gctx, v_cache, qlen, head_size, num_attention_heads, v_cache->nb[1], v_cache->nb[2],
-                     n_past * ggml_element_size(v_cache)); // [n_head, head_size, qlen]
-    ggml_build_forward_expand(&ctx->gf, ggml_cpy(ctx->gctx, value_layer, v_cache_view));
-
-    key_layer = ggml_view_3d(ctx->gctx, k_cache, head_size, n_past + qlen, num_attention_heads, k_cache->nb[1],
-                             k_cache->nb[2], 0); // [n_head, klen, head_size]
-    value_layer = ggml_view_3d(ctx->gctx, v_cache, n_past + qlen, head_size, num_attention_heads, v_cache->nb[1],
-                               v_cache->nb[2], 0); // [n_head, head_size, klen]
-
-    ggml_tensor *attn_scores = ggml_mul_mat(ctx->gctx, key_layer, query_layer); // [n_head, qlen, klen]
-    if (n_past == 0) {
-        // build attention mask for context input
-        ggml_tensor *inf = ggml_new_tensor_3d(ctx->gctx, attn_scores->type, 1, qlen - 1, num_attention_heads);
-        ggml_set_f32(inf, -INFINITY);
-        ggml_tensor *masked_attn_scores = ggml_view_3d(
-            ctx->gctx, attn_scores, 1, qlen - 1, num_attention_heads, qlen * ggml_element_size(attn_scores),
-            qlen * qlen * ggml_element_size(attn_scores), (qlen - 1) * ggml_element_size(attn_scores));
-        ggml_build_forward_expand(&ctx->gf, ggml_cpy(ctx->gctx, inf, masked_attn_scores));
-    }
-    attn_scores = ggml_scale_inplace(ctx->gctx, attn_scores, ggml_new_f32(ctx->gctx, 1.f / std::sqrt(head_size)));
-    ggml_tensor *attn_probs = ggml_soft_max_inplace(ctx->gctx, attn_scores); // [n_head, qlen, klen]
-
-    ggml_tensor *context_layer = ggml_mul_mat(ctx->gctx, value_layer, attn_probs); // [n_head, qlen, head_size]
-    context_layer = ggml_reshape_2d(ctx->gctx, ggml_cont(ctx->gctx, ggml_permute(ctx->gctx, context_layer, 0, 2, 1, 3)),
-                                    hidden_size, qlen);
-
-    ggml_tensor *attn_output = dense.forward(ctx, context_layer);
-    return attn_output;
-}
-
-ggml_tensor *GLMBlock::forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past, int n_ctx) const {
-    ggml_tensor *alpha = ggml_new_f32(ctx->gctx, std::sqrt(2.f * num_layers));
-
-    ggml_tensor *attn_input = input_layernorm.forward(ctx, hidden_states);
-    ggml_tensor *attn_output = attention.forward(ctx, attn_input, n_past, n_ctx);
-    ggml_build_forward_expand(&ctx->gf, attn_output);
-    hidden_states = ggml_add_inplace(ctx->gctx, ggml_scale_inplace(ctx->gctx, attn_input, alpha), attn_output);
-
-    ggml_tensor *mlp_input = post_attention_layernorm.forward(ctx, hidden_states);
-    ggml_tensor *mlp_output = mlp.forward(ctx, mlp_input);
-    ggml_build_forward_expand(&ctx->gf, mlp_output);
-    ggml_tensor *output = ggml_add_inplace(ctx->gctx, ggml_scale_inplace(ctx->gctx, mlp_input, alpha), mlp_output);
-    return output;
-}
-
-ChatGLMModel::ChatGLMModel(InitContext *ctx, const ChatGLMConfig &config)
-    : word_embeddings(ctx, config.vocab_size, config.hidden_size), final_layernorm(ctx, config.hidden_size) {
-    layers.reserve(config.num_layers);
-    for (int layer_id = 0; layer_id < config.num_layers; layer_id++) {
-        layers.emplace_back(ctx, config.hidden_size, config.num_attention_heads, config.num_layers,
-                            config.max_sequence_length);
-    }
-}
-
-ggml_tensor *ChatGLMModel::forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past, int n_ctx) const {
-    ggml_tensor *hidden_states = word_embeddings.forward(ctx, input_ids);
-    for (const GLMBlock &layer : layers) {
-        ggml_set_scratch(ctx->gctx, ctx->scratch);
-        hidden_states = layer.forward(ctx, hidden_states, n_past, n_ctx);
-    }
-    ggml_set_scratch(ctx->gctx, {.offs = 0, .size = 0, .data = nullptr});
-    hidden_states = final_layernorm.forward(ctx, hidden_states);
-    return hidden_states;
-}
-
-ggml_tensor *ChatGLMForConditionalGeneration::forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past,
-                                                      int n_ctx) const {
-    ggml_tensor *transformer_outputs = transformer.forward(ctx, input_ids, n_past, n_ctx);
-    // NOTE: only compute next_token_logits for the last token
-    if (input_ids->ne[0] > 1) {
-        transformer_outputs =
-            ggml_view_1d(ctx->gctx, transformer_outputs, config.hidden_size,
-                         (input_ids->ne[0] - 1) * config.hidden_size * ggml_element_size(transformer_outputs));
-    }
-    ggml_tensor *lm_head_weight = transformer.word_embeddings.weight; // tied weight
-    ggml_tensor *lm_logits = ggml_mul_mat(ctx->gctx, lm_head_weight, transformer_outputs);
-    return lm_logits;
-}
-
-struct TokenIdScore {
-    int id;
-    float score;
-
-    friend bool operator<(const TokenIdScore &a, const TokenIdScore &b) { return a.score < b.score; }
-    friend bool operator>(const TokenIdScore &a, const TokenIdScore &b) { return a.score > b.score; }
-};
-
-static void sampling_softmax_inplace(TokenIdScore *first, TokenIdScore *last) {
-    float max_score = std::max_element(first, last)->score;
-    float sum = 0.f;
-    for (TokenIdScore *p = first; p != last; p++) {
-        float s = std::exp(p->score - max_score);
-        p->score = s;
-        sum += s;
-    }
-    float inv_sum = 1.f / sum;
-    for (TokenIdScore *p = first; p != last; p++) {
-        p->score *= inv_sum;
-    }
-}
-
 // Adapted from https://github.com/ggerganov/llama.cpp/blob/master/examples/common.cpp
 static inline int get_num_physical_cores() {
     unsigned int n_threads = std::thread::hardware_concurrency();
     return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4;
 }
 
-int ChatGLMForConditionalGeneration::generate_next_token(const std::vector<int> &input_ids,
-                                                         const GenerationConfig &gen_config, int n_past,
-                                                         int n_ctx) const {
+std::string to_string(ModelType model_type) {
+    switch (model_type) {
+    case MODEL_TYPE_CHATGLM:
+        return "ChatGLM";
+    case MODEL_TYPE_CHATGLM2:
+        return "ChatGLM2";
+    default:
+        CHATGLM_THROW << "unknown model type " << model_type;
+    }
+}
+
+int BaseModelForConditionalGeneration::generate_next_token(const std::vector<int> &input_ids,
+                                                           const GenerationConfig &gen_config, int n_past,
+                                                           int n_ctx) const {
     ForwardContext ctx;
-    ctx.gctx = ggml_init({.mem_size = MEM_SIZE, .mem_buffer = mem_buffer.get(), .no_alloc = false});
-    CHATGLM_CHECK(ctx.gctx) << "failed to init ggml context";
+    ctx.gctx = GGMLContext({.mem_size = mem_size_, .mem_buffer = mem_buffer_.get(), .no_alloc = false});
     int n_threads = gen_config.num_threads > 0 ? gen_config.num_threads : get_num_physical_cores();
     ctx.gf = {};
     ctx.gf.n_threads = input_ids.size() >= 32 && ggml_cpu_has_blas() && !ggml_cpu_has_gpublas() ? 1 : n_threads;
-    ctx.scratch = {.offs = 0, .size = SCRATCH_SIZE, .data = scratch_buffer.get()};
+    ctx.scratch = {.offs = 0, .size = scratch_size_, .data = scratch_buffer_.get()};
 
-    ggml_tensor *input_ids_tensor = ggml_new_tensor_1d(ctx.gctx, GGML_TYPE_I32, input_ids.size());
+    ggml_tensor *input_ids_tensor = ggml_new_tensor_1d(ctx.gctx.get(), GGML_TYPE_I32, input_ids.size());
     memcpy(input_ids_tensor->data, input_ids.data(), ggml_nbytes(input_ids_tensor));
 
     ggml_tensor *lm_logits = forward(&ctx, input_ids_tensor, n_past, n_ctx);
 
     ggml_build_forward_expand(&ctx.gf, lm_logits);
-    ggml_graph_compute(ctx.gctx, &ctx.gf);
+    ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
 
 #ifdef GGML_PERF
     ggml_graph_print(&ctx.gf);
 #endif
 
+    int vocab_size = lm_logits->ne[0];
     float *next_token_logits = (float *)lm_logits->data;
 
     int next_token_id;
     if (gen_config.do_sample) {
         // temperature sampling
         float inv_temp = 1.f / gen_config.temperature;
-        for (int i = 0; i < config.vocab_size; i++) {
+        for (int i = 0; i < vocab_size; i++) {
             next_token_logits[i] *= inv_temp;
         }
 
-        std::vector<TokenIdScore> token_scores(config.vocab_size);
-        for (int i = 0; i < config.vocab_size; i++) {
+        std::vector<TokenIdScore> token_scores(vocab_size);
+        for (int i = 0; i < vocab_size; i++) {
             token_scores[i] = {.id = i, .score = next_token_logits[i]};
         }
 
@@ -477,19 +462,33 @@ int ChatGLMForConditionalGeneration::generate_next_token(const std::vector<int>
         next_token_id = token_scores[dist(gen)].id;
     } else {
         // greedy search
-        next_token_id = std::max_element(next_token_logits, next_token_logits + config.vocab_size) - next_token_logits;
+        next_token_id = std::max_element(next_token_logits, next_token_logits + vocab_size) - next_token_logits;
     }
-    ggml_free(ctx.gctx);
 
     return next_token_id;
 }
 
-std::vector<int> ChatGLMForConditionalGeneration::generate(const std::vector<int> &input_ids,
-                                                           const GenerationConfig &gen_config,
-                                                           BaseStreamer *streamer) const {
-    CHATGLM_CHECK(gen_config.max_length <= config.max_sequence_length)
-        << "max_length (" << gen_config.max_length << ") is larger than model max_sequence_length ("
-        << config.max_sequence_length << ")";
+void BaseModelForConditionalGeneration::sampling_softmax_inplace(TokenIdScore *first, TokenIdScore *last) {
+    float max_score = std::max_element(first, last)->score;
+    float sum = 0.f;
+    for (TokenIdScore *p = first; p != last; p++) {
+        float s = std::exp(p->score - max_score);
+        p->score = s;
+        sum += s;
+    }
+    float inv_sum = 1.f / sum;
+    for (TokenIdScore *p = first; p != last; p++) {
+        p->score *= inv_sum;
+    }
+}
+
+std::vector<int> BaseModelForConditionalGeneration::generate(const std::vector<int> &input_ids,
+                                                             const GenerationConfig &gen_config,
+                                                             BaseStreamer *streamer) const {
+    CHATGLM_CHECK(gen_config.max_length <= config_.max_length)
+        << "requested max_length (" << gen_config.max_length << ") is larger than model's max_length ("
+        << config_.max_length << ")";
+
     std::vector<int> curr_input_ids(input_ids);
 
     std::vector<int> output_ids;
@@ -513,7 +512,7 @@ std::vector<int> ChatGLMForConditionalGeneration::generate(const std::vector<int
             streamer->put({next_token_id});
         }
 
-        if (next_token_id == config.eos_token_id) {
+        if (next_token_id == config_.eos_token_id) {
             break;
         }
     }
@@ -525,158 +524,497 @@ std::vector<int> ChatGLMForConditionalGeneration::generate(const std::vector<int
     return output_ids;
 }
 
-MappedFile::MappedFile(const std::string &path) {
-    int fd = open(path.c_str(), O_RDONLY);
-    CHATGLM_CHECK(fd > 0) << "cannot open file " << path << ": " << strerror(errno);
+ggml_tensor *GLMMLP::forward(ForwardContext *ctx, ggml_tensor *hidden_states) const {
+    ggml_tensor *output = dense_h_to_4h.forward(ctx, hidden_states);
+    output = ggml_gelu_inplace(ctx->gctx.get(), output);
+    output = dense_4h_to_h.forward(ctx, output);
+    return output;
+}
 
-    struct stat sb;
-    CHATGLM_CHECK(fstat(fd, &sb) == 0) << strerror(errno);
-    size = sb.st_size;
+ggml_tensor *GLMSelfAttention::forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past, int n_ctx) const {
+    int hidden_size = hidden_states->ne[0];
+    int qlen = hidden_states->ne[1];
+    int head_size = hidden_size / num_attention_heads;
+    int rope_dim = head_size / 2;
 
-    data = mmap(nullptr, size, PROT_READ, MAP_SHARED, fd, 0);
-    CHATGLM_CHECK(data != MAP_FAILED) << strerror(errno);
+    ggml_tensor *qkv = query_key_value.forward(ctx, hidden_states); // [qlen, 3 * hidden]
 
-    CHATGLM_CHECK(close(fd) == 0) << strerror(errno);
-}
+    ggml_tensor *query_layer = ggml_view_3d(ctx->gctx.get(), qkv, head_size, num_attention_heads, qlen,
+                                            3 * head_size * ggml_element_size(qkv), qkv->nb[1], 0);
+    query_layer =
+        ggml_rope_inplace(ctx->gctx.get(), query_layer, n_past, rope_dim, 4, n_ctx); // [qlen, heads, head_size]
+    query_layer = ggml_permute(ctx->gctx.get(), query_layer, 0, 2, 1, 3);            // [heads, qlen, head_size]
 
-MappedFile::~MappedFile() { CHATGLM_CHECK(munmap(data, size) == 0) << strerror(errno); }
+    ggml_tensor *key_layer =
+        ggml_view_3d(ctx->gctx.get(), qkv, head_size, num_attention_heads, qlen, 3 * head_size * ggml_element_size(qkv),
+                     qkv->nb[1], head_size * ggml_element_size(qkv));
+    key_layer = ggml_rope_inplace(ctx->gctx.get(), key_layer, n_past, rope_dim, 4, n_ctx); // [qlen, heads, head_size]
+    key_layer = ggml_permute(ctx->gctx.get(), key_layer, 0, 2, 1, 3);                      // [heads, qlen, head_size]
 
-ChatGLMPipeline::ChatGLMPipeline(const std::string &path) {
-    mapped_file = std::make_unique<MappedFile>(path);
+    ggml_tensor *value_layer = ggml_view_3d(ctx->gctx.get(), qkv, head_size, num_attention_heads, qlen,
+                                            3 * head_size * ggml_element_size(qkv), qkv->nb[1],
+                                            2 * head_size * ggml_element_size(qkv)); // [qlen, heads, head_size]
+    value_layer = ggml_permute(ctx->gctx.get(), value_layer, 1, 2, 0, 3);            // [heads, head_size, qlen]
 
-    std::ifstream fin(path, std::ios::binary);
-    CHATGLM_CHECK(fin) << "failed to open model file " << path;
+    // store key & value to cache
+    ggml_tensor *k_cache_view =
+        ggml_view_3d(ctx->gctx.get(), k_cache, head_size, qlen, num_attention_heads, k_cache->nb[1], k_cache->nb[2],
+                     n_past * head_size * ggml_element_size(k_cache)); // [heads, qlen, head_size]
+    ggml_build_forward_expand(&ctx->gf, ggml_cpy(ctx->gctx.get(), key_layer, k_cache_view));
+    ggml_tensor *v_cache_view =
+        ggml_view_3d(ctx->gctx.get(), v_cache, qlen, head_size, num_attention_heads, v_cache->nb[1], v_cache->nb[2],
+                     n_past * ggml_element_size(v_cache)); // [heads, head_size, qlen]
+    ggml_build_forward_expand(&ctx->gf, ggml_cpy(ctx->gctx.get(), value_layer, v_cache_view));
 
-    // load magic
-    std::string magic(4, '\0');
-    fin.read(&magic[0], magic.size());
-    CHATGLM_CHECK(magic == "ggml") << "model file is broken (bad magic)";
+    key_layer = ggml_view_3d(ctx->gctx.get(), k_cache, head_size, n_past + qlen, num_attention_heads, k_cache->nb[1],
+                             k_cache->nb[2], 0); // [heads, klen, head_size]
+    value_layer = ggml_view_3d(ctx->gctx.get(), v_cache, n_past + qlen, head_size, num_attention_heads, v_cache->nb[1],
+                               v_cache->nb[2], 0); // [heads, head_size, klen]
+
+    ggml_tensor *attn_scores = ggml_mul_mat(ctx->gctx.get(), key_layer, query_layer); // [heads, qlen, klen]
+    if (n_past == 0) {
+        // build attention mask for context input
+        ggml_tensor *inf = ggml_new_tensor_3d(ctx->gctx.get(), attn_scores->type, 1, qlen - 1, num_attention_heads);
+        ggml_set_f32(inf, -INFINITY);
+        ggml_tensor *masked_attn_scores = ggml_view_3d(
+            ctx->gctx.get(), attn_scores, 1, qlen - 1, num_attention_heads, qlen * ggml_element_size(attn_scores),
+            qlen * qlen * ggml_element_size(attn_scores), (qlen - 1) * ggml_element_size(attn_scores));
+        ggml_build_forward_expand(&ctx->gf, ggml_cpy(ctx->gctx.get(), inf, masked_attn_scores));
+    }
+    attn_scores =
+        ggml_scale_inplace(ctx->gctx.get(), attn_scores, ggml_new_f32(ctx->gctx.get(), 1.f / std::sqrt(head_size)));
+    ggml_tensor *attn_probs = ggml_soft_max_inplace(ctx->gctx.get(), attn_scores); // [heads, qlen, klen]
+
+    ggml_tensor *context_layer = ggml_mul_mat(ctx->gctx.get(), value_layer, attn_probs); // [heads, qlen, head_size]
+    context_layer = ggml_reshape_2d(
+        ctx->gctx.get(), ggml_cont(ctx->gctx.get(), ggml_permute(ctx->gctx.get(), context_layer, 0, 2, 1, 3)),
+        hidden_size, qlen);
+
+    ggml_tensor *attn_output = dense.forward(ctx, context_layer);
+    return attn_output;
+}
+
+ggml_tensor *GLMBlock::forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past, int n_ctx) const {
+    ggml_tensor *alpha = ggml_new_f32(ctx->gctx.get(), std::sqrt(2.f * num_hidden_layers));
+
+    ggml_tensor *attn_input = input_layernorm.forward(ctx, hidden_states);
+    ggml_tensor *attn_output = attention.forward(ctx, attn_input, n_past, n_ctx);
+    ggml_build_forward_expand(&ctx->gf, attn_output);
+    hidden_states =
+        ggml_add_inplace(ctx->gctx.get(), ggml_scale_inplace(ctx->gctx.get(), attn_input, alpha), attn_output);
+
+    ggml_tensor *mlp_input = post_attention_layernorm.forward(ctx, hidden_states);
+    ggml_tensor *mlp_output = mlp.forward(ctx, mlp_input);
+    ggml_build_forward_expand(&ctx->gf, mlp_output);
+    ggml_tensor *output =
+        ggml_add_inplace(ctx->gctx.get(), ggml_scale_inplace(ctx->gctx.get(), mlp_input, alpha), mlp_output);
+    return output;
+}
 
-    // load config
-    ChatGLMConfig config = read_as<ChatGLMConfig>(fin);
+ChatGLMModel::ChatGLMModel(InitContext *ctx, const ChatGLMConfig &config)
+    : word_embeddings(ctx, config.vocab_size, config.hidden_size), final_layernorm(ctx, config.hidden_size) {
+    layers.reserve(config.num_hidden_layers);
+    for (int layer_id = 0; layer_id < config.num_hidden_layers; layer_id++) {
+        layers.emplace_back(ctx, config.hidden_size, config.num_attention_heads, config.num_hidden_layers,
+                            config.max_length);
+    }
+}
 
-    // load tokenizer
-    int proto_size = read_as<int>(fin);
-    std::string_view serialized_model_proto((char *)mapped_file->data + fin.tellg(), proto_size);
-    fin.seekg(proto_size, fin.cur);
-    tokenizer = std::make_unique<ChatGLMTokenizer>(serialized_model_proto, config.bos_token_id, config.eos_token_id,
-                                                   config.mask_token_id, config.gmask_token_id, config.pad_token_id);
+ggml_tensor *ChatGLMModel::forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past, int n_ctx) const {
+    ggml_tensor *hidden_states = word_embeddings.forward(ctx, input_ids);
+    for (const GLMBlock &layer : layers) {
+        ggml_set_scratch(ctx->gctx.get(), ctx->scratch);
+        hidden_states = layer.forward(ctx, hidden_states, n_past, n_ctx);
+    }
+    ggml_set_scratch(ctx->gctx.get(), {.offs = 0, .size = 0, .data = nullptr});
+    hidden_states = final_layernorm.forward(ctx, hidden_states);
+    return hidden_states;
+}
 
-    // load model
+ChatGLMForConditionalGeneration::ChatGLMForConditionalGeneration(const ChatGLMConfig &config)
+    : BaseModelForConditionalGeneration(MODEL_TYPE_CHATGLM, config, MEM_SIZE, SCRATCH_SIZE), config(config) {
     constexpr size_t tensor_ovhd = GGML_TENSOR_SIZE + GGML_OBJECT_SIZE;
-    const size_t num_tensors = 3 + config.num_layers * 14;
+    const size_t num_tensors = 3 + config.num_hidden_layers * 14;
     const size_t ctx_size = num_tensors * tensor_ovhd;
-    ctx.gctx = ggml_init({.mem_size = ctx_size, .mem_buffer = nullptr, .no_alloc = true});
-    CHATGLM_CHECK(ctx.gctx) << "failed to init weight context";
-    ctx.dtype = config.dtype;
-
-    const size_t kvcache_size =
-        config.num_layers * 2 * config.max_sequence_length * config.hidden_size * ggml_type_size(GGML_TYPE_F16);
-    kvcache_buffer.reset(new char[kvcache_size]);
-    char *kvcache_ptr = kvcache_buffer.get();
+    w_ctx_.gctx = GGMLContext({.mem_size = ctx_size, .mem_buffer = nullptr, .no_alloc = true});
+    w_ctx_.dtype = config.dtype;
 
-    auto load_tensor = [this, &fin](const std::string &name, ggml_tensor *tensor) {
-        int ndim = read_as<int>(fin);
-        CHATGLM_CHECK(ndim == tensor->n_dims)
-            << "tensor " << name << " ndim mismatch: expect " << tensor->n_dims << ", got " << ndim;
-        int name_size = read_as<int>(fin);
-        ggml_type dtype = (ggml_type)read_as<int>(fin);
-        CHATGLM_CHECK(dtype == tensor->type)
-            << "tensor " << name << " dtype mismatch: expect " << tensor->type << ", got " << dtype;
-        for (int i = 0; i < ndim; i++) {
-            int dim = read_as<int>(fin);
-            CHATGLM_CHECK(dim == tensor->ne[ndim - 1 - i]) << "tensor " << name << " dim " << i << " mismatch: expect "
-                                                           << tensor->ne[ndim - 1 - i] << ", got " << dim;
-        }
-        std::string tensor_name(name_size, 0);
-        fin.read(&tensor_name[0], name_size);
-        CHATGLM_CHECK(tensor_name == name)
-            << "tensor " << name << " name mismatch: expect " << name << ", got " << tensor_name;
+    transformer = ChatGLMModel(&w_ctx_, config);
 
-        constexpr int64_t MEM_ALIGNED = 16;
-        const int64_t data_offset = (fin.tellg() + (MEM_ALIGNED - 1)) & ~(MEM_ALIGNED - 1);
-        tensor->data = (char *)mapped_file->data + data_offset;
-        fin.seekg(data_offset + ggml_nbytes(tensor));
-    };
+    const size_t kv_cache_size =
+        config.num_hidden_layers * 2 * config.max_length * config.hidden_size * ggml_type_size(GGML_TYPE_F16);
+    kv_cache_buffer_.reset(new char[kv_cache_size]);
+    char *kv_cache_ptr = kv_cache_buffer_.get();
+    for (int i = 0; i < config.num_hidden_layers; i++) {
+        std::string layer_prefix = "transformer.layers." + std::to_string(i) + '.';
+        transformer.layers[i].attention.k_cache->data = kv_cache_ptr;
+        kv_cache_ptr += ggml_nbytes(transformer.layers[i].attention.k_cache);
+        transformer.layers[i].attention.v_cache->data = kv_cache_ptr;
+        kv_cache_ptr += ggml_nbytes(transformer.layers[i].attention.v_cache);
+    }
+    CHATGLM_CHECK(kv_cache_ptr == kv_cache_buffer_.get() + kv_cache_size) << "corrupted kv cache";
+}
 
-    model = std::make_unique<ChatGLMForConditionalGeneration>(&ctx, config);
-    load_tensor("transformer.word_embeddings.weight", model->transformer.word_embeddings.weight);
-    for (int i = 0; i < config.num_layers; i++) {
+void ChatGLMForConditionalGeneration::load(ModelLoader &loader) {
+    loader.read_tensor("transformer.word_embeddings.weight", transformer.word_embeddings.weight);
+    for (int i = 0; i < config.num_hidden_layers; i++) {
         std::string layer_prefix = "transformer.layers." + std::to_string(i) + '.';
-        load_tensor(layer_prefix + "input_layernorm.weight", model->transformer.layers[i].input_layernorm.weight);
-        load_tensor(layer_prefix + "input_layernorm.bias", model->transformer.layers[i].input_layernorm.bias);
-        load_tensor(layer_prefix + "attention.query_key_value.weight",
-                    model->transformer.layers[i].attention.query_key_value.weight);
-        load_tensor(layer_prefix + "attention.query_key_value.bias",
-                    model->transformer.layers[i].attention.query_key_value.bias);
-        load_tensor(layer_prefix + "attention.dense.weight", model->transformer.layers[i].attention.dense.weight);
-        load_tensor(layer_prefix + "attention.dense.bias", model->transformer.layers[i].attention.dense.bias);
-        model->transformer.layers[i].attention.k_cache->data = kvcache_ptr;
-        kvcache_ptr += ggml_nbytes(model->transformer.layers[i].attention.k_cache);
-        model->transformer.layers[i].attention.v_cache->data = kvcache_ptr;
-        kvcache_ptr += ggml_nbytes(model->transformer.layers[i].attention.v_cache);
-        load_tensor(layer_prefix + "post_attention_layernorm.weight",
-                    model->transformer.layers[i].post_attention_layernorm.weight);
-        load_tensor(layer_prefix + "post_attention_layernorm.bias",
-                    model->transformer.layers[i].post_attention_layernorm.bias);
-        load_tensor(layer_prefix + "mlp.dense_h_to_4h.weight", model->transformer.layers[i].mlp.dense_h_to_4h.weight);
-        load_tensor(layer_prefix + "mlp.dense_h_to_4h.bias", model->transformer.layers[i].mlp.dense_h_to_4h.bias);
-        load_tensor(layer_prefix + "mlp.dense_4h_to_h.weight", model->transformer.layers[i].mlp.dense_4h_to_h.weight);
-        load_tensor(layer_prefix + "mlp.dense_4h_to_h.bias", model->transformer.layers[i].mlp.dense_4h_to_h.bias);
-    }
-    load_tensor("transformer.final_layernorm.weight", model->transformer.final_layernorm.weight);
-    load_tensor("transformer.final_layernorm.bias", model->transformer.final_layernorm.bias);
-    CHATGLM_CHECK(kvcache_ptr == kvcache_buffer.get() + kvcache_size) << "corrupted kv cache";
-    CHATGLM_CHECK(ggml_used_mem(ctx.gctx) == ggml_get_mem_size(ctx.gctx)) << "corrupted model weights";
+        loader.read_tensor(layer_prefix + "input_layernorm.weight", transformer.layers[i].input_layernorm.weight);
+        loader.read_tensor(layer_prefix + "input_layernorm.bias", transformer.layers[i].input_layernorm.bias);
+        loader.read_tensor(layer_prefix + "attention.query_key_value.weight",
+                           transformer.layers[i].attention.query_key_value.weight);
+        loader.read_tensor(layer_prefix + "attention.query_key_value.bias",
+                           transformer.layers[i].attention.query_key_value.bias);
+        loader.read_tensor(layer_prefix + "attention.dense.weight", transformer.layers[i].attention.dense.weight);
+        loader.read_tensor(layer_prefix + "attention.dense.bias", transformer.layers[i].attention.dense.bias);
+        loader.read_tensor(layer_prefix + "post_attention_layernorm.weight",
+                           transformer.layers[i].post_attention_layernorm.weight);
+        loader.read_tensor(layer_prefix + "post_attention_layernorm.bias",
+                           transformer.layers[i].post_attention_layernorm.bias);
+        loader.read_tensor(layer_prefix + "mlp.dense_h_to_4h.weight", transformer.layers[i].mlp.dense_h_to_4h.weight);
+        loader.read_tensor(layer_prefix + "mlp.dense_h_to_4h.bias", transformer.layers[i].mlp.dense_h_to_4h.bias);
+        loader.read_tensor(layer_prefix + "mlp.dense_4h_to_h.weight", transformer.layers[i].mlp.dense_4h_to_h.weight);
+        loader.read_tensor(layer_prefix + "mlp.dense_4h_to_h.bias", transformer.layers[i].mlp.dense_4h_to_h.bias);
+    }
+    loader.read_tensor("transformer.final_layernorm.weight", transformer.final_layernorm.weight);
+    loader.read_tensor("transformer.final_layernorm.bias", transformer.final_layernorm.bias);
+    CHATGLM_CHECK(ggml_used_mem(w_ctx_.gctx.get()) == ggml_get_mem_size(w_ctx_.gctx.get()))
+        << "corrupted model weights";
 
 #ifdef GGML_USE_CUBLAS
-    for (int i = 0; i < config.num_layers; i++) {
-        ggml_cuda_transform_tensor(model->transformer.layers[i].attention.query_key_value.weight);
-        ggml_cuda_transform_tensor(model->transformer.layers[i].attention.dense.weight);
-        ggml_cuda_transform_tensor(model->transformer.layers[i].mlp.dense_h_to_4h.weight);
-        ggml_cuda_transform_tensor(model->transformer.layers[i].mlp.dense_4h_to_h.weight);
+    auto transform_tensor_to_cuda = [](ggml_tensor *tensor) {
+        tensor->backend = GGML_BACKEND_GPU;
+        ggml_cuda_transform_tensor(tensor->data, tensor);
+    };
+
+    for (int i = 0; i < config.num_hidden_layers; i++) {
+        transform_tensor_to_cuda(transformer.layers[i].attention.query_key_value.weight);
+        transform_tensor_to_cuda(transformer.layers[i].attention.dense.weight);
+        transform_tensor_to_cuda(transformer.layers[i].mlp.dense_h_to_4h.weight);
+        transform_tensor_to_cuda(transformer.layers[i].mlp.dense_4h_to_h.weight);
     }
 #endif
 }
 
-ChatGLMPipeline::~ChatGLMPipeline() {
-    if (ctx.gctx) {
-        ggml_free(ctx.gctx);
-        ctx.gctx = nullptr;
+ggml_tensor *ChatGLMForConditionalGeneration::forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past,
+                                                      int n_ctx) const {
+    ggml_tensor *transformer_outputs = transformer.forward(ctx, input_ids, n_past, n_ctx);
+    // NOTE: only compute next_token_logits for the last token
+    if (input_ids->ne[0] > 1) {
+        transformer_outputs =
+            ggml_view_1d(ctx->gctx.get(), transformer_outputs, config.hidden_size,
+                         (input_ids->ne[0] - 1) * config.hidden_size * ggml_element_size(transformer_outputs));
     }
+    ggml_tensor *lm_head_weight = transformer.word_embeddings.weight; // tied weight
+    ggml_tensor *lm_logits = ggml_mul_mat(ctx->gctx.get(), lm_head_weight, transformer_outputs);
+    return lm_logits;
 }
 
-std::string ChatGLMPipeline::chat(const std::vector<std::string> &history, const GenerationConfig &gen_config,
-                                  BaseStreamer *streamer) const {
-    std::string prompt = build_prompt(history);
-    std::vector<int> input_ids = tokenizer->encode(prompt);
-    if ((int)input_ids.size() > gen_config.max_context_length) {
-        // sliding window: always take the last max_context_length tokens
-        input_ids.erase(input_ids.begin(), input_ids.end() - gen_config.max_context_length);
-    }
+// ===== ChatGLM2-6B =====
 
-    std::vector<int> output_ids = model->generate(input_ids, gen_config, streamer);
-    std::vector<int> new_output_ids(output_ids.begin() + input_ids.size(), output_ids.end());
+ChatGLM2Tokenizer::ChatGLM2Tokenizer(std::string_view serialized_model_proto) {
+    const auto status = sp.LoadFromSerializedProto(serialized_model_proto);
+    CHATGLM_CHECK(status.ok()) << status.ToString();
 
-    std::string output = tokenizer->decode(new_output_ids);
-    return output;
+    int special_id = sp.GetPieceSize();
+    mask_token_id = special_id++;
+    gmask_token_id = special_id++;
+    smask_token_id = special_id++;
+    sop_token_id = special_id++;
+    eop_token_id = special_id++;
+}
+
+std::vector<int> ChatGLM2Tokenizer::encode(const std::string &text) const {
+    std::vector<int> ids;
+    sp.Encode(text, &ids);
+    ids.insert(ids.begin(), {gmask_token_id, sop_token_id}); // special prefix
+    return ids;
 }
 
-std::string ChatGLMPipeline::build_prompt(const std::vector<std::string> &history) {
+std::string ChatGLM2Tokenizer::decode(const std::vector<int> &ids) const {
+    // filter out special tokens
+    std::vector<int> normal_ids(ids);
+    normal_ids.erase(std::remove_if(normal_ids.begin(), normal_ids.end(), [this](int id) { return is_special_id(id); }),
+                     normal_ids.end());
+
+    std::string text;
+    sp.Decode(normal_ids, &text);
+    return text;
+}
+
+std::vector<int> ChatGLM2Tokenizer::encode_history(const std::vector<std::string> &history, int max_length) const {
+    std::string prompt = build_prompt(history);
+    std::vector<int> input_ids = encode(prompt);
+    if ((int)input_ids.size() > max_length) {
+        // sliding window: drop the least recent history while keeping the special prefix tokens
+        int num_drop = (int)input_ids.size() - max_length;
+        input_ids.erase(input_ids.begin() + 2, input_ids.begin() + 2 + num_drop);
+    }
+    return input_ids;
+}
+
+std::string ChatGLM2Tokenizer::build_prompt(const std::vector<std::string> &history) {
     CHATGLM_CHECK(history.size() % 2 == 1) << "invalid history size " << history.size();
 
     std::ostringstream oss_prompt;
-    if (history.size() == 1) {
-        oss_prompt << history.front();
-    } else {
-        for (size_t i = 0; i < history.size(); i += 2) {
-            oss_prompt << "[Round " << i / 2 << "]\n问：" << history[i] << "\n答：";
-            if (i < history.size() - 1) {
-                oss_prompt << history[i + 1] << "\n";
-            }
+    for (size_t i = 0; i < history.size(); i += 2) {
+        oss_prompt << "[Round " << i / 2 + 1 << "]\n\n问：" << history[i] << "\n\n答：";
+        if (i < history.size() - 1) {
+            oss_prompt << history[i + 1] << "\n\n";
         }
     }
     return oss_prompt.str();
 }
 
+bool ChatGLM2Tokenizer::is_special_id(int id) const {
+    return id == mask_token_id || id == gmask_token_id || id == smask_token_id || id == sop_token_id ||
+           id == eop_token_id;
+}
+
+ggml_tensor *GLM2SelfAttention::forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past) const {
+    const int hidden_size = hidden_states->ne[0];
+    const int qlen = hidden_states->ne[1];
+    const int head_size = hidden_size / num_attention_heads;
+    const int rope_dim = head_size / 2;
+    const int mqa_scale = num_attention_heads / num_kv_heads;
+
+    ggml_tensor *qkv = query_key_value.forward(ctx, hidden_states); // [qlen, hidden + 2 * kv_hidden]
+
+    ggml_tensor *query_layer =
+        ggml_view_3d(ctx->gctx.get(), qkv, head_size, num_attention_heads, qlen, head_size * ggml_element_size(qkv),
+                     qkv->nb[1], 0); // [qlen, heads, head_size]
+    query_layer = ggml_rope_inplace(ctx->gctx.get(), query_layer, n_past, rope_dim, 0, 0);
+    query_layer = ggml_view_4d(ctx->gctx.get(), query_layer, head_size, mqa_scale, num_kv_heads, qlen,
+                               query_layer->nb[1], query_layer->nb[1] * mqa_scale, query_layer->nb[2],
+                               0);                                        // [qlen, kv_heads, mqa_scale, head_size]
+    query_layer = ggml_permute(ctx->gctx.get(), query_layer, 0, 2, 3, 1); // [kv_heads, mqa_scale, qlen, head_size]
+
+    ggml_tensor *key_layer =
+        ggml_view_3d(ctx->gctx.get(), qkv, head_size, num_kv_heads, qlen, head_size * ggml_element_size(qkv),
+                     qkv->nb[1], hidden_size * ggml_element_size(qkv)); // [qlen, kv_heads, head_size]
+    key_layer = ggml_rope_inplace(ctx->gctx.get(), key_layer, n_past, rope_dim, 0, 0);
+    key_layer = ggml_permute(ctx->gctx.get(), key_layer, 0, 2, 1, 3); // [kv_heads, qlen, head_size]
+
+    ggml_tensor *value_layer = ggml_view_3d(
+        ctx->gctx.get(), qkv, head_size, num_kv_heads, qlen, head_size * ggml_element_size(qkv), qkv->nb[1],
+        (hidden_size + head_size * num_kv_heads) * ggml_element_size(qkv)); // [qlen, kv_heads, head_size]
+    value_layer = ggml_permute(ctx->gctx.get(), value_layer, 1, 2, 0, 3);   // [kv_heads, head_size, qlen]
+
+    // store key & value to cache
+    ggml_tensor *k_cache_view =
+        ggml_view_3d(ctx->gctx.get(), k_cache, head_size, qlen, num_kv_heads, k_cache->nb[1], k_cache->nb[2],
+                     n_past * head_size * ggml_element_size(k_cache)); // [kv_heads, qlen, head_size]
+    ggml_build_forward_expand(&ctx->gf, ggml_cpy(ctx->gctx.get(), key_layer, k_cache_view));
+    ggml_tensor *v_cache_view =
+        ggml_view_3d(ctx->gctx.get(), v_cache, qlen, head_size, num_kv_heads, v_cache->nb[1], v_cache->nb[2],
+                     n_past * ggml_element_size(v_cache)); // [kv_heads, head_size, qlen]
+    ggml_build_forward_expand(&ctx->gf, ggml_cpy(ctx->gctx.get(), value_layer, v_cache_view));
+
+    // concat key & value with past kv
+    key_layer = ggml_view_4d(ctx->gctx.get(), k_cache, head_size, n_past + qlen, mqa_scale, num_kv_heads,
+                             k_cache->nb[1], 0, k_cache->nb[2],
+                             0); // [kv_heads, mqa_scale, klen, head_size]
+    value_layer = ggml_view_4d(ctx->gctx.get(), v_cache, n_past + qlen, head_size, mqa_scale, num_kv_heads,
+                               v_cache->nb[1], 0, v_cache->nb[2],
+                               0); // [kv_heads, mqa_scale, head_size, klen]
+
+    // flash attention
+    ggml_tensor *context_layer = ggml_flash_attn(ctx->gctx.get(), query_layer, key_layer, value_layer,
+                                                 true); // [mqa_scale, kv_heads, qlen, head_size]
+    context_layer = ggml_reshape_2d(
+        ctx->gctx.get(), ggml_cont(ctx->gctx.get(), ggml_permute(ctx->gctx.get(), context_layer, 0, 3, 1, 2)),
+        hidden_size, qlen); // [qlen, hidden]
+
+    ggml_tensor *attn_output = dense.forward(ctx, context_layer);
+    return attn_output;
+}
+
+ggml_tensor *GLM2MLP::forward(ForwardContext *ctx, ggml_tensor *hidden_states) const {
+    ggml_tensor *output = dense_h_to_4h.forward(ctx, hidden_states);
+
+    // swiglu activation
+    ggml_tensor *x0 = ggml_view_2d(ctx->gctx.get(), output, output->ne[0] / 2, output->ne[1], output->nb[1], 0);
+    ggml_tensor *x1 = ggml_view_2d(ctx->gctx.get(), output, output->ne[0] / 2, output->ne[1], output->nb[1],
+                                   output->ne[0] / 2 * ggml_element_size(output));
+    output = ggml_mul_inplace(ctx->gctx.get(), ggml_silu_inplace(ctx->gctx.get(), ggml_cont(ctx->gctx.get(), x0)), x1);
+
+    output = dense_4h_to_h.forward(ctx, output);
+    return output;
+}
+
+ggml_tensor *GLM2Block::forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past) const {
+    ggml_tensor *residual = ggml_dup(ctx->gctx.get(), hidden_states);
+    ggml_build_forward_expand(&ctx->gf, residual);
+
+    hidden_states = input_layernorm.forward(ctx, hidden_states);
+    hidden_states = attention.forward(ctx, hidden_states, n_past);
+    hidden_states = ggml_add_inplace(ctx->gctx.get(), hidden_states, residual);
+
+    residual = ggml_dup(ctx->gctx.get(), hidden_states);
+    ggml_build_forward_expand(&ctx->gf, residual);
+
+    hidden_states = post_attention_layernorm.forward(ctx, hidden_states);
+    hidden_states = mlp.forward(ctx, hidden_states);
+    hidden_states = ggml_add_inplace(ctx->gctx.get(), hidden_states, residual);
+
+    return hidden_states;
+}
+
+ChatGLM2Model::ChatGLM2Model(InitContext *ctx, const ChatGLM2Config &config)
+    : word_embeddings(ctx, config.vocab_size, config.hidden_size), final_layernorm(ctx, config.hidden_size) {
+    // TODO: reduce max length? 32k might be too large for cpu inference
+    layers.reserve(config.num_hidden_layers);
+    for (int layer_id = 0; layer_id < config.num_hidden_layers; layer_id++) {
+        layers.emplace_back(ctx, config.hidden_size, config.num_attention_heads, config.num_kv_heads,
+                            config.intermediate_size, config.max_length);
+    }
+}
+
+ggml_tensor *ChatGLM2Model::forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past) const {
+    ggml_tensor *hidden_states = word_embeddings.forward(ctx, input_ids);
+    for (const auto &layer : layers) {
+        ggml_set_scratch(ctx->gctx.get(), ctx->scratch);
+        hidden_states = layer.forward(ctx, hidden_states, n_past);
+    }
+    ggml_set_scratch(ctx->gctx.get(), {.offs = 0, .size = 0, .data = nullptr});
+    hidden_states = final_layernorm.forward(ctx, hidden_states);
+    return hidden_states;
+}
+
+ChatGLM2ForConditionalGeneration::ChatGLM2ForConditionalGeneration(const ChatGLM2Config &config)
+    : BaseModelForConditionalGeneration(MODEL_TYPE_CHATGLM2, config, MEM_SIZE, SCRATCH_SIZE), config(config) {
+    constexpr size_t tensor_ovhd = GGML_TENSOR_SIZE + GGML_OBJECT_SIZE;
+    const size_t num_tensors = 3 + config.num_hidden_layers * 9;
+    const size_t ctx_size = num_tensors * tensor_ovhd;
+    w_ctx_.gctx = GGMLContext({.mem_size = ctx_size, .mem_buffer = nullptr, .no_alloc = true});
+    w_ctx_.dtype = config.dtype;
+
+    transformer = ChatGLM2Model(&w_ctx_, config);
+    lm_head = Linear(&w_ctx_, config.hidden_size, config.vocab_size, false);
+
+    const size_t kv_cache_size = config.num_hidden_layers * 2ull * config.max_length * config.hidden_size /
+                                 config.num_attention_heads * config.num_kv_heads * ggml_type_size(GGML_TYPE_F32);
+    kv_cache_buffer_.reset(new char[kv_cache_size]);
+
+    char *kv_cache_ptr = kv_cache_buffer_.get();
+    for (int i = 0; i < config.num_hidden_layers; i++) {
+        std::string layer_prefix = "transformer.layers." + std::to_string(i) + '.';
+        transformer.layers[i].attention.k_cache->data = kv_cache_ptr;
+        kv_cache_ptr += ggml_nbytes(transformer.layers[i].attention.k_cache);
+        transformer.layers[i].attention.v_cache->data = kv_cache_ptr;
+        kv_cache_ptr += ggml_nbytes(transformer.layers[i].attention.v_cache);
+    }
+    CHATGLM_CHECK(kv_cache_ptr == kv_cache_buffer_.get() + kv_cache_size) << "corrupted kv cache";
+}
+
+void ChatGLM2ForConditionalGeneration::load(ModelLoader &loader) {
+    loader.read_tensor("transformer.embedding.word_embeddings.weight", transformer.word_embeddings.weight);
+    for (int i = 0; i < config.num_hidden_layers; i++) {
+        std::string layer_prefix = "transformer.encoder.layers." + std::to_string(i) + '.';
+        loader.read_tensor(layer_prefix + "input_layernorm.weight", transformer.layers[i].input_layernorm.weight);
+        loader.read_tensor(layer_prefix + "self_attention.query_key_value.weight",
+                           transformer.layers[i].attention.query_key_value.weight);
+        loader.read_tensor(layer_prefix + "self_attention.query_key_value.bias",
+                           transformer.layers[i].attention.query_key_value.bias);
+        loader.read_tensor(layer_prefix + "self_attention.dense.weight", transformer.layers[i].attention.dense.weight);
+        loader.read_tensor(layer_prefix + "post_attention_layernorm.weight",
+                           transformer.layers[i].post_attention_layernorm.weight);
+        loader.read_tensor(layer_prefix + "mlp.dense_h_to_4h.weight", transformer.layers[i].mlp.dense_h_to_4h.weight);
+        loader.read_tensor(layer_prefix + "mlp.dense_4h_to_h.weight", transformer.layers[i].mlp.dense_4h_to_h.weight);
+    }
+    loader.read_tensor("transformer.encoder.final_layernorm.weight", transformer.final_layernorm.weight);
+    loader.read_tensor("transformer.output_layer.weight", lm_head.weight);
+    CHATGLM_CHECK(ggml_used_mem(w_ctx_.gctx.get()) == ggml_get_mem_size(w_ctx_.gctx.get()))
+        << "corrupted model weights";
+
+#ifdef GGML_USE_CUBLAS
+    auto transform_tensor_to_cuda = [](ggml_tensor *tensor) {
+        tensor->backend = GGML_BACKEND_GPU;
+        ggml_cuda_transform_tensor(tensor->data, tensor);
+    };
+
+    for (int i = 0; i < config.num_hidden_layers; i++) {
+        transform_tensor_to_cuda(transformer.layers[i].attention.query_key_value.weight);
+        transform_tensor_to_cuda(transformer.layers[i].attention.dense.weight);
+        transform_tensor_to_cuda(transformer.layers[i].mlp.dense_h_to_4h.weight);
+        transform_tensor_to_cuda(transformer.layers[i].mlp.dense_4h_to_h.weight);
+    }
+    transform_tensor_to_cuda(lm_head.weight);
+#endif
+}
+
+ggml_tensor *ChatGLM2ForConditionalGeneration::forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past,
+                                                       int n_ctx) const {
+    ggml_tensor *transformer_outputs = transformer.forward(ctx, input_ids, n_past);
+    // NOTE: only compute next_token_logits for the last token
+    if (input_ids->ne[0] > 1) {
+        transformer_outputs =
+            ggml_view_1d(ctx->gctx.get(), transformer_outputs, config.hidden_size,
+                         (input_ids->ne[0] - 1) * config.hidden_size * ggml_element_size(transformer_outputs));
+    }
+    ggml_tensor *lm_logits = lm_head.forward(ctx, transformer_outputs);
+    return lm_logits;
+}
+
+// ===== pipeline =====
+
+Pipeline::Pipeline(const std::string &path) {
+    mapped_file = std::make_unique<MappedFile>(path);
+    ModelLoader loader(std::string_view((char *)mapped_file->data, mapped_file->size));
+
+    // load magic
+    std::string magic = loader.read_string(4);
+    CHATGLM_CHECK(magic == "ggml") << "model file is broken (bad magic)";
+
+    // load model type
+    ModelType model_type = (ModelType)loader.read_basic<int>();
+    if (model_type == MODEL_TYPE_CHATGLM) {
+        // load version
+        int version = loader.read_basic<int>();
+        CHATGLM_CHECK(version == 1) << "only support version 1 for now but got " << version;
+
+        // load config
+        ChatGLMConfig config = loader.read_basic<ChatGLMConfig>();
+
+        // load tokenizer
+        int proto_size = loader.read_basic<int>();
+        std::string_view serialized_model_proto((char *)mapped_file->data + loader.tell(), proto_size);
+        loader.seek(proto_size, SEEK_CUR);
+        tokenizer = std::make_unique<ChatGLMTokenizer>(serialized_model_proto);
+
+        // load model
+        model = std::make_unique<ChatGLMForConditionalGeneration>(config);
+        model->load(loader);
+    } else if (model_type == MODEL_TYPE_CHATGLM2) {
+        // load version
+        int version = loader.read_basic<int>();
+        CHATGLM_CHECK(version == 1) << "only support version 1 for now but got " << version;
+
+        // load config
+        ChatGLM2Config config = loader.read_basic<ChatGLM2Config>();
+
+        // load tokenizer
+        int proto_size = loader.read_basic<int>();
+        std::string_view serialized_model_proto((char *)mapped_file->data + loader.tell(), proto_size);
+        loader.seek(proto_size, SEEK_CUR);
+        tokenizer = std::make_unique<ChatGLM2Tokenizer>(serialized_model_proto);
+
+        // load model
+        model = std::make_unique<ChatGLM2ForConditionalGeneration>(config);
+        model->load(loader);
+    } else {
+        CHATGLM_THROW << "invalid model type " << model_type;
+    }
+}
+
+std::string Pipeline::chat(const std::vector<std::string> &history, const GenerationConfig &gen_config,
+                           BaseStreamer *streamer) const {
+    std::vector<int> input_ids = tokenizer->encode_history(history, gen_config.max_context_length);
+    std::vector<int> output_ids = model->generate(input_ids, gen_config, streamer);
+
+    std::vector<int> new_output_ids(output_ids.begin() + input_ids.size(), output_ids.end());
+    std::string output = tokenizer->decode(new_output_ids);
+    return output;
+}
+
 } // namespace chatglm
diff --git a/chatglm.h b/chatglm.h
index 9354a57..a18acd0 100644
--- a/chatglm.h
+++ b/chatglm.h
@@ -3,19 +3,23 @@
 #include <ggml.h>
 #include <sentencepiece_processor.h>
 #include <sstream>
+#include <unordered_map>
 #include <vector>
 
 namespace chatglm {
 
-// ===== utils =====
+// ===== common =====
 
-struct LogMessageFatal {
-    std::ostringstream oss;
+class LogMessageFatal {
+  public:
+    LogMessageFatal(const char *file, int line) { oss_ << file << ':' << line << ' '; }
+    [[noreturn]] ~LogMessageFatal() noexcept(false) { throw std::runtime_error(oss_.str()); }
+    std::ostringstream &stream() { return oss_; }
 
-    LogMessageFatal(const char *file, int line) { oss << file << ':' << line << ' '; }
-    [[noreturn]] ~LogMessageFatal() noexcept(false) { throw std::runtime_error(oss.str()); }
-    std::ostringstream &stream() { return oss; }
+  private:
+    std::ostringstream oss_;
 };
+
 #define CHATGLM_THROW ::chatglm::LogMessageFatal(__FILE__, __LINE__).stream()
 #define CHATGLM_CHECK(cond)                                                                                            \
     if (!(cond))                                                                                                       \
@@ -23,232 +27,502 @@ struct LogMessageFatal {
 
 std::string to_string(ggml_tensor *tensor, bool with_data = true);
 
-// ===== config =====
-
-struct ChatGLMConfig {
+struct BaseConfig {
+    // common attributes
+    ggml_type dtype;
     int vocab_size;
     int hidden_size;
     int num_attention_heads;
-    int num_layers;
-    int max_sequence_length;
+    int num_hidden_layers;
+    int intermediate_size;
+    // for sequence generation
+    int max_length;
+    // for tokenizer
     int bos_token_id;
     int eos_token_id;
-    int gmask_token_id;
-    int mask_token_id;
     int pad_token_id;
+    int sep_token_id;
+};
+
+class BaseTokenizer {
+  public:
+    virtual ~BaseTokenizer() = default;
+    virtual std::vector<int> encode(const std::string &text) const = 0;
+    virtual std::string decode(const std::vector<int> &ids) const = 0;
+    virtual std::vector<int> encode_history(const std::vector<std::string> &history, int max_length) const = 0;
+};
+
+class GGMLContext {
+  public:
+    GGMLContext() : gctx_(nullptr) {}
+
+    GGMLContext(ggml_init_params params) : gctx_(ggml_init(params)) {
+        CHATGLM_CHECK(gctx_) << "failed to init ggml context";
+    }
+    // copy constructor
+    GGMLContext(const GGMLContext &) = delete;
+    // move constructor
+    GGMLContext(GGMLContext &&other) : gctx_(other.gctx_) { other.gctx_ = nullptr; }
+
+    ~GGMLContext() { reset(); }
+
+    // copy assignment
+    GGMLContext &operator=(const GGMLContext &) = delete;
+    // move assignment
+    GGMLContext &operator=(GGMLContext &&other) {
+        reset();
+        gctx_ = other.gctx_;
+        other.gctx_ = nullptr;
+        return *this;
+    }
+
+    ggml_context *get() const { return gctx_; }
+
+    void reset() {
+        if (gctx_) {
+            ggml_free(gctx_);
+            gctx_ = nullptr;
+        }
+    }
+
+  private:
+    ggml_context *gctx_;
+};
+
+struct InitContext {
+    GGMLContext gctx;
     ggml_type dtype;
 };
 
-// ===== tokenizer =====
+struct ForwardContext {
+    GGMLContext gctx;
+    ggml_cgraph gf;
+    ggml_scratch scratch;
+};
 
-struct ChatGLMTokenizer {
-    std::unique_ptr<sentencepiece::SentencePieceProcessor> sp;
-    int bos_token_id;
-    int eos_token_id;
-    int mask_token_id;
-    int gmask_token_id;
-    int pad_token_id;
+class Embedding {
+  public:
+    Embedding() : weight(nullptr) {}
+    Embedding(InitContext *ctx, int num_embeddings, int embedding_dim)
+        : weight(ggml_new_tensor_2d(ctx->gctx.get(), ctx->dtype, embedding_dim, num_embeddings)) {}
+
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input) const;
 
-    ChatGLMTokenizer(std::string_view serialized_model_proto, int bos_token_id, int eos_token_id, int mask_token_id,
-                     int gmask_token_id, int pad_token_id);
+  public:
+    ggml_tensor *weight;
+};
 
-    std::vector<int> encode(const std::string &text) const;
+class Linear {
+  public:
+    Linear() : weight(nullptr), bias(nullptr) {}
+    Linear(InitContext *ctx, int in_features, int out_features, bool use_bias = true)
+        : weight(ggml_new_tensor_2d(ctx->gctx.get(), ctx->dtype, in_features, out_features)),
+          bias(use_bias ? ggml_new_tensor_1d(ctx->gctx.get(), GGML_TYPE_F32, out_features) : nullptr) {}
 
-    std::string decode(const std::vector<int> &ids) const;
+    int in_features() const { return weight->ne[0]; }
+    int out_features() const { return weight->ne[1]; }
 
-  private:
-    static std::string preprocess(const std::string &text);
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input) const;
 
-    static std::string postprocess(const std::string &text);
+  public:
+    ggml_tensor *weight; // [out_features, in_features]
+    ggml_tensor *bias;   // [out_features]
+};
+
+class LayerNorm {
+  public:
+    LayerNorm() : weight(nullptr), bias(nullptr) {}
+    LayerNorm(InitContext *ctx, int normalized_shape)
+        : weight(ggml_new_tensor_1d(ctx->gctx.get(), GGML_TYPE_F32, normalized_shape)),
+          bias(ggml_new_tensor_1d(ctx->gctx.get(), GGML_TYPE_F32, normalized_shape)) {}
+
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input) const;
+
+  public:
+    ggml_tensor *weight; // [normalized_shape]
+    ggml_tensor *bias;   // [normalized_shape]
 };
 
-// ===== streamer =====
+class RMSNorm {
+  public:
+    RMSNorm() : weight(nullptr) {}
+    RMSNorm(InitContext *ctx, int normalized_shape)
+        : weight(ggml_new_tensor_1d(ctx->gctx.get(), GGML_TYPE_F32, normalized_shape)) {}
 
-struct BaseStreamer {
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input) const;
+
+  public:
+    ggml_tensor *weight;
+};
+
+class BaseStreamer {
+  public:
+    virtual ~BaseStreamer() = default;
     virtual void put(const std::vector<int> &output_ids) = 0;
     virtual void end() = 0;
 };
 
 // reference: https://github.com/huggingface/transformers/blob/main/src/transformers/generation/streamers.py
-struct TextStreamer : public BaseStreamer {
-    TextStreamer(ChatGLMTokenizer *tokenizer) : tokenizer_(tokenizer), is_prompt_(true), print_len_(0) {}
+class TextStreamer : public BaseStreamer {
+  public:
+    TextStreamer(BaseTokenizer *tokenizer) : tokenizer_(tokenizer), is_prompt_(true), print_len_(0) {}
     void put(const std::vector<int> &output_ids) override;
     void end() override;
 
   private:
-    ChatGLMTokenizer *tokenizer_;
+    BaseTokenizer *tokenizer_;
     bool is_prompt_;
     std::vector<int> token_cache_;
     int print_len_;
 };
 
-// ===== model =====
+class MappedFile {
+  public:
+    MappedFile(const std::string &path);
+    ~MappedFile();
 
-struct InitContext {
-    ggml_context *gctx;
-    ggml_type dtype;
+  public:
+    char *data;
+    size_t size;
 };
 
-struct ForwardContext {
-    ggml_context *gctx;
-    ggml_cgraph gf;
-    ggml_scratch scratch;
-};
+class ModelLoader {
+  public:
+    ModelLoader(std::string_view buffer) : data(buffer.data()), size(buffer.size()), ptr(buffer.data()) {}
 
-struct Embedding {
-    ggml_tensor *weight;
+    int64_t tell() const { return ptr - data; }
 
-    Embedding() : weight(nullptr) {}
-    Embedding(InitContext *ctx, int num_embeddings, int embedding_dim)
-        : weight(ggml_new_tensor_2d(ctx->gctx, ctx->dtype, embedding_dim, num_embeddings)) {}
+    void seek(int64_t offset, int whence);
 
-    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input) const;
+    template <typename T>
+    T read_basic() {
+        T obj = *(T *)ptr;
+        ptr += sizeof(T);
+        return obj;
+    }
+
+    std::string read_string(size_t length);
+
+    void read_tensor(const std::string &name, ggml_tensor *tensor);
+
+  public:
+    const char *const data;
+    size_t size;
+    const char *ptr;
 };
 
-struct Linear {
-    ggml_tensor *weight; // [out_features, in_features]
-    ggml_tensor *bias;   // [out_features]
+// ===== generation =====
 
-    Linear() : weight(nullptr), bias(nullptr) {}
-    Linear(InitContext *ctx, int in_features, int out_features)
-        : weight(ggml_new_tensor_2d(ctx->gctx, ctx->dtype, in_features, out_features)),
-          bias(ggml_new_tensor_1d(ctx->gctx, GGML_TYPE_F32, out_features)) {}
+struct GenerationConfig {
+    int max_length;
+    int max_context_length;
+    bool do_sample;
+    int top_k;
+    float top_p;
+    float temperature;
+    int num_threads;
 
-    int in_features() const { return weight->ne[0]; }
-    int out_features() const { return weight->ne[1]; }
+    GenerationConfig(int max_length = 2048, int max_context_length = 512, bool do_sample = true, int top_k = 0,
+                     float top_p = 0.7, float temperature = 0.95, int num_threads = 0)
+        : max_length(max_length), max_context_length(max_context_length), do_sample(do_sample), top_k(top_k),
+          top_p(top_p), temperature(temperature), num_threads(num_threads) {}
+};
 
-    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input) const;
+enum ModelType {
+    MODEL_TYPE_CHATGLM = 1,
+    MODEL_TYPE_CHATGLM2 = 2,
 };
 
-struct LayerNorm {
-    ggml_tensor *weight; // [normalized_shape]
-    ggml_tensor *bias;   // [normalized_shape]
+std::string to_string(ModelType model_type);
 
-    LayerNorm() : weight(nullptr), bias(nullptr) {}
-    LayerNorm(InitContext *ctx, int normalized_shape)
-        : weight(ggml_new_tensor_1d(ctx->gctx, GGML_TYPE_F32, normalized_shape)),
-          bias(ggml_new_tensor_1d(ctx->gctx, GGML_TYPE_F32, normalized_shape)) {}
+class BaseModelForConditionalGeneration {
+  public:
+    BaseModelForConditionalGeneration(ModelType model_type, BaseConfig config, size_t mem_size, size_t scratch_size)
+        : model_type_(model_type), config_(config), mem_size_(mem_size), mem_buffer_(new char[mem_size]),
+          scratch_size_(scratch_size), scratch_buffer_(new char[scratch_size]) {}
+    virtual ~BaseModelForConditionalGeneration() = default;
 
-    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input) const;
+    virtual void load(ModelLoader &loader) = 0;
+    virtual ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past, int n_ctx) const = 0;
+
+    ModelType type() const { return model_type_; }
+    std::string type_name() const { return to_string(model_type_); }
+
+    std::vector<int> generate(const std::vector<int> &input_ids, const GenerationConfig &gen_config,
+                              BaseStreamer *streamer = nullptr) const;
+
+    int generate_next_token(const std::vector<int> &input_ids, const GenerationConfig &gen_config, int n_past,
+                            int n_ctx) const;
+
+    struct TokenIdScore {
+        int id;
+        float score;
+
+        bool operator<(const TokenIdScore &other) const { return score < other.score; }
+        bool operator>(const TokenIdScore &other) const { return score > other.score; }
+    };
+
+  private:
+    static void sampling_softmax_inplace(TokenIdScore *first, TokenIdScore *last);
+
+  private:
+    ModelType model_type_;
+    BaseConfig config_;
+    size_t mem_size_;
+    std::unique_ptr<char[]> mem_buffer_; // BLAS buffer
+    size_t scratch_size_;
+    std::unique_ptr<char[]> scratch_buffer_; // intermediate tensor buffer
 };
 
-struct GLU {
-    Linear dense_h_to_4h;
-    Linear dense_4h_to_h;
+// ===== ChatGLM-6B =====
+
+struct ChatGLMConfig : public BaseConfig {};
 
-    GLU() = default;
-    GLU(InitContext *ctx, int hidden_size)
+class ChatGLMTokenizer : public BaseTokenizer {
+  public:
+    ChatGLMTokenizer(std::string_view serialized_model_proto);
+
+    std::vector<int> encode(const std::string &text) const override;
+
+    std::string decode(const std::vector<int> &ids) const override;
+
+    std::vector<int> encode_history(const std::vector<std::string> &history, int max_length) const override;
+
+    static std::string build_prompt(const std::vector<std::string> &history);
+
+  private:
+    static std::string preprocess(const std::string &text);
+
+    static std::string postprocess(const std::string &text);
+
+  public:
+    sentencepiece::SentencePieceProcessor sp;
+    int bos_token_id;
+    int eos_token_id;
+    int mask_token_id;
+    int gmask_token_id;
+    int pad_token_id;
+};
+
+class GLMMLP {
+  public:
+    GLMMLP() = default;
+    GLMMLP(InitContext *ctx, int hidden_size)
         : dense_h_to_4h(ctx, hidden_size, 4 * hidden_size), dense_4h_to_h(ctx, 4 * hidden_size, hidden_size) {}
 
     ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *hidden_states) const;
+
+  public:
+    Linear dense_h_to_4h;
+    Linear dense_4h_to_h;
 };
 
-struct SelfAttention {
+class GLMSelfAttention {
+  public:
+    // TODO: kv cache type
+    GLMSelfAttention() : num_attention_heads(0) {}
+    GLMSelfAttention(InitContext *ctx, int hidden_size, int num_attention_heads, int max_length)
+        : query_key_value(ctx, hidden_size, 3 * hidden_size), dense(ctx, hidden_size, hidden_size),
+          num_attention_heads(num_attention_heads),
+          k_cache(ggml_new_tensor_3d(ctx->gctx.get(), GGML_TYPE_F16, hidden_size / num_attention_heads, max_length,
+                                     num_attention_heads)),
+          v_cache(ggml_new_tensor_3d(ctx->gctx.get(), GGML_TYPE_F16, max_length, hidden_size / num_attention_heads,
+                                     num_attention_heads)) {}
+
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past, int n_ctx) const;
+
+  public:
     Linear query_key_value;
     Linear dense;
     int num_attention_heads;
     ggml_tensor *k_cache; // [n_head, maxlen, head_size]
     ggml_tensor *v_cache; // [n_head, head_size, maxlen]
+};
 
-    // TODO: kvcache type
-    SelfAttention() : num_attention_heads(0) {}
-    SelfAttention(InitContext *ctx, int hidden_size, int num_attention_heads, int max_length)
-        : query_key_value(ctx, hidden_size, 3 * hidden_size), dense(ctx, hidden_size, hidden_size),
-          num_attention_heads(num_attention_heads),
-          k_cache(ggml_new_tensor_3d(ctx->gctx, GGML_TYPE_F16, hidden_size / num_attention_heads, max_length,
-                                     num_attention_heads)),
-          v_cache(ggml_new_tensor_3d(ctx->gctx, GGML_TYPE_F16, max_length, hidden_size / num_attention_heads,
-                                     num_attention_heads)) {}
+class GLMBlock {
+  public:
+    GLMBlock() : num_hidden_layers(0) {}
+    GLMBlock(InitContext *ctx, int hidden_size, int num_attention_heads, int num_hidden_layers, int max_length)
+        : input_layernorm(ctx, hidden_size), attention(ctx, hidden_size, num_attention_heads, max_length),
+          post_attention_layernorm(ctx, hidden_size), mlp(ctx, hidden_size), num_hidden_layers(num_hidden_layers) {}
 
     ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past, int n_ctx) const;
-};
 
-struct GLMBlock {
+  public:
     LayerNorm input_layernorm;
-    SelfAttention attention;
+    GLMSelfAttention attention;
     LayerNorm post_attention_layernorm;
-    GLU mlp;
-    int num_layers;
+    GLMMLP mlp;
+    int num_hidden_layers;
+};
 
-    GLMBlock() : num_layers(0) {}
-    GLMBlock(InitContext *ctx, int hidden_size, int num_attention_heads, int num_layers, int max_length)
-        : input_layernorm(ctx, hidden_size), attention(ctx, hidden_size, num_attention_heads, max_length),
-          post_attention_layernorm(ctx, hidden_size), mlp(ctx, hidden_size), num_layers(num_layers) {}
+class ChatGLMModel {
+  public:
+    ChatGLMModel() = default;
+    ChatGLMModel(InitContext *ctx, const ChatGLMConfig &config);
 
-    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past, int n_ctx) const;
-};
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past, int n_ctx) const;
 
-struct ChatGLMModel {
+  public:
     Embedding word_embeddings;
     std::vector<GLMBlock> layers;
     LayerNorm final_layernorm;
+};
 
-    ChatGLMModel() = default;
-    ChatGLMModel(InitContext *ctx, const ChatGLMConfig &config);
+class ChatGLMForConditionalGeneration : public BaseModelForConditionalGeneration {
+  public:
+    ChatGLMForConditionalGeneration() = default;
+    ChatGLMForConditionalGeneration(const ChatGLMConfig &config);
 
-    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past, int n_ctx) const;
+    void load(ModelLoader &loader) override;
+
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past, int n_ctx) const override;
+
+  public:
+    static constexpr size_t MEM_SIZE = 272ull * 1024 * 1024;
+    static constexpr size_t SCRATCH_SIZE = 128ull * 1024 * 1024;
+    ChatGLMConfig config;
+    ChatGLMModel transformer;
+
+  private:
+    // hold ggml_context & kv_cache
+    InitContext w_ctx_; // weight context
+    // InitContext kv_ctx_; // TODO: kv cache context
+    std::unique_ptr<char[]> kv_cache_buffer_;
 };
 
-struct GenerationConfig {
-    int max_length;
-    int max_context_length;
-    bool do_sample;
-    int top_k;
-    float top_p;
-    float temperature;
-    int num_threads;
+// ===== ChatGLM2-6B =====
 
-    GenerationConfig(int max_length = 2048, int max_context_length = 512, bool do_sample = true, int top_k = 0,
-                     float top_p = 0.7, float temperature = 0.95, int num_threads = 0)
-        : max_length(max_length), max_context_length(max_context_length), do_sample(do_sample), top_k(top_k),
-          top_p(top_p), temperature(temperature), num_threads(num_threads) {}
+struct ChatGLM2Config : public BaseConfig {
+    int num_kv_heads;
 };
 
-struct ChatGLMForConditionalGeneration {
-    ChatGLMConfig config;
-    ChatGLMModel transformer;
+class ChatGLM2Tokenizer : public BaseTokenizer {
+  public:
+    ChatGLM2Tokenizer(std::string_view serialized_model_proto);
 
-    static constexpr size_t MEM_SIZE = 272ull * 1024 * 1024;
-    std::unique_ptr<char[]> mem_buffer; // BLAS buffer
-    static constexpr size_t SCRATCH_SIZE = 128ull * 1024 * 1024;
-    std::unique_ptr<char[]> scratch_buffer; // intermediate tensor buffer
+    std::vector<int> encode(const std::string &text) const override;
 
-    ChatGLMForConditionalGeneration() = default;
-    ChatGLMForConditionalGeneration(InitContext *ctx, const ChatGLMConfig &config)
-        : config(config), transformer(ctx, config), mem_buffer(new char[MEM_SIZE]),
-          scratch_buffer(new char[SCRATCH_SIZE]) {}
+    std::string decode(const std::vector<int> &ids) const override;
 
-    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past, int n_ctx) const;
+    std::vector<int> encode_history(const std::vector<std::string> &history, int max_length) const override;
 
-    std::vector<int> generate(const std::vector<int> &input_ids, const GenerationConfig &gen_config,
-                              BaseStreamer *streamer = nullptr) const;
+    static std::string build_prompt(const std::vector<std::string> &history);
 
-    int generate_next_token(const std::vector<int> &input_ids, const GenerationConfig &gen_config, int n_past,
-                            int n_ctx) const;
+    bool is_special_id(int id) const;
+
+  public:
+    sentencepiece::SentencePieceProcessor sp;
+    int mask_token_id;
+    int gmask_token_id;
+    int smask_token_id;
+    int sop_token_id;
+    int eop_token_id;
 };
 
-struct MappedFile {
-    void *data;
-    size_t size;
+class GLM2SelfAttention {
+  public:
+    GLM2SelfAttention() : num_attention_heads(0), num_kv_heads(0) {}
+    GLM2SelfAttention(InitContext *ctx, int hidden_size, int num_attention_heads, int num_kv_heads, int max_length)
+        : num_attention_heads(num_attention_heads), num_kv_heads(num_kv_heads),
+          query_key_value(ctx, hidden_size, hidden_size + 2 * (hidden_size / num_attention_heads) * num_kv_heads),
+          dense(ctx, hidden_size, hidden_size, false),
+          k_cache(ggml_new_tensor_3d(ctx->gctx.get(), GGML_TYPE_F32, hidden_size / num_attention_heads, max_length,
+                                     num_kv_heads)),
+          v_cache(ggml_new_tensor_3d(ctx->gctx.get(), GGML_TYPE_F32, max_length, hidden_size / num_attention_heads,
+                                     num_kv_heads)) {}
+
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past) const;
+
+  public:
+    int num_attention_heads;
+    int num_kv_heads;
+    Linear query_key_value;
+    Linear dense;
+    ggml_tensor *k_cache; // [mqa_n_head, maxlen, head_size]
+    ggml_tensor *v_cache; // [mqa_n_head, head_size, maxlen]
+};
 
-    MappedFile(const std::string &path);
-    ~MappedFile();
+class GLM2MLP {
+  public:
+    GLM2MLP(InitContext *ctx, int hidden_size, int intermediate_size)
+        : dense_h_to_4h(ctx, hidden_size, intermediate_size * 2, false),
+          dense_4h_to_h(ctx, intermediate_size, hidden_size, false) {}
+
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *hidden_states) const;
+
+  public:
+    Linear dense_h_to_4h;
+    Linear dense_4h_to_h;
 };
 
-struct ChatGLMPipeline {
-    std::unique_ptr<ChatGLMTokenizer> tokenizer;
-    std::unique_ptr<ChatGLMForConditionalGeneration> model;
-    InitContext ctx;
-    std::unique_ptr<char[]> kvcache_buffer;
-    std::unique_ptr<MappedFile> mapped_file;
+class GLM2Block {
+  public:
+    GLM2Block() = default;
+    GLM2Block(InitContext *ctx, int hidden_size, int num_attention_heads, int num_kv_heads, int intermediate_size,
+              int max_length)
+        : input_layernorm(ctx, hidden_size), attention(ctx, hidden_size, num_attention_heads, num_kv_heads, max_length),
+          post_attention_layernorm(ctx, hidden_size), mlp(ctx, hidden_size, intermediate_size) {}
+
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past) const;
+
+  public:
+    RMSNorm input_layernorm;
+    GLM2SelfAttention attention;
+    RMSNorm post_attention_layernorm;
+    GLM2MLP mlp;
+};
 
-    ChatGLMPipeline(const std::string &path);
-    ~ChatGLMPipeline();
+class ChatGLM2Model {
+  public:
+    ChatGLM2Model() = default;
+    ChatGLM2Model(InitContext *ctx, const ChatGLM2Config &config);
+
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past) const;
+
+  public:
+    Embedding word_embeddings;
+    std::vector<GLM2Block> layers;
+    RMSNorm final_layernorm;
+};
+
+class ChatGLM2ForConditionalGeneration : public BaseModelForConditionalGeneration {
+  public:
+    ChatGLM2ForConditionalGeneration() = default;
+    ChatGLM2ForConditionalGeneration(const ChatGLM2Config &config);
+
+    void load(ModelLoader &loader) override;
+
+    ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past, int n_ctx) const override;
+
+  public:
+    static constexpr size_t MEM_SIZE = 512ull * 1024 * 1024;
+    static constexpr size_t SCRATCH_SIZE = 144ull * 1024 * 1024;
+
+    ChatGLM2Config config;
+    ChatGLM2Model transformer;
+    Linear lm_head;
+
+  private:
+    // hold ggml_context & kv_cache
+    InitContext w_ctx_; // weight context
+    // InitContext kv_ctx_; // TODO: kv cache context
+    std::unique_ptr<char[]> kv_cache_buffer_;
+};
+
+// ===== pipeline =====
+
+class Pipeline {
+  public:
+    Pipeline(const std::string &path);
 
     std::string chat(const std::vector<std::string> &history, const GenerationConfig &gen_config,
                      BaseStreamer *streamer = nullptr) const;
 
-    static std::string build_prompt(const std::vector<std::string> &history);
+  public:
+    std::unique_ptr<BaseTokenizer> tokenizer;
+    std::unique_ptr<BaseModelForConditionalGeneration> model;
+    std::unique_ptr<MappedFile> mapped_file;
 };
 
 } // namespace chatglm
diff --git a/chatglm_cpp/__init__.py b/chatglm_cpp/__init__.py
index 45729be..dda216d 100644
--- a/chatglm_cpp/__init__.py
+++ b/chatglm_cpp/__init__.py
@@ -4,10 +4,10 @@
 
 import chatglm_cpp._C as _C
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 
-class ChatGLMPipeline(_C.ChatGLMPipeline):
+class Pipeline(_C.Pipeline):
     def __init__(self, model_path: os.PathLike) -> None:
         model_path = Path(model_path)
         super().__init__(str(model_path))
@@ -22,6 +22,7 @@ def stream_chat(
         top_k: int = 0,
         top_p: float = 0.7,
         temperature: float = 0.95,
+        num_threads: int = 0,
     ) -> Iterator[str]:
         gen_config = _C.GenerationConfig(
             max_length=max_length,
@@ -30,11 +31,10 @@ def stream_chat(
             top_k=top_k,
             top_p=top_p,
             temperature=temperature,
+            num_threads=num_threads,
         )
 
-        prompt = self.build_prompt(history)
-        input_ids = self.tokenizer.encode(prompt)
-        input_ids = input_ids[-max_context_length:]  # sliding window
+        input_ids = self.tokenizer.encode_history(history, max_context_length)
 
         output_ids = input_ids
         n_past = 0
@@ -77,6 +77,7 @@ def chat(
         top_k: int = 0,
         top_p: float = 0.7,
         temperature: float = 0.95,
+        num_threads: int = 0,
     ) -> str:
         gen_config = _C.GenerationConfig(
             max_length=max_length,
@@ -85,11 +86,10 @@ def chat(
             top_k=top_k,
             top_p=top_p,
             temperature=temperature,
+            num_threads=num_threads,
         )
 
-        prompt = self.build_prompt(history)
-        input_ids = self.tokenizer.encode(prompt)
-        input_ids = input_ids[-max_context_length:]  # sliding window
+        input_ids = self.tokenizer.encode_history(history, max_context_length)
 
         output_ids = input_ids
         n_past = 0
diff --git a/chatglm_pybind.cpp b/chatglm_pybind.cpp
index caba8d9..f4269c4 100644
--- a/chatglm_pybind.cpp
+++ b/chatglm_pybind.cpp
@@ -10,9 +10,18 @@ using namespace pybind11::literals;
 PYBIND11_MODULE(_C, m) {
     m.doc() = "ChatGLM.cpp python binding";
 
-    py::class_<ChatGLMTokenizer>(m, "ChatGLMTokenizer")
-        .def("encode", &ChatGLMTokenizer::encode)
-        .def("decode", &ChatGLMTokenizer::decode);
+    py::class_<BaseConfig>(m, "BaseConfig")
+        .def_readonly("dtype", &BaseConfig::dtype)
+        .def_readonly("vocab_size", &BaseConfig::vocab_size)
+        .def_readonly("hidden_size", &BaseConfig::hidden_size)
+        .def_readonly("num_attention_heads", &BaseConfig::num_attention_heads)
+        .def_readonly("num_hidden_layers", &BaseConfig::num_hidden_layers)
+        .def_readonly("intermediate_size", &BaseConfig::intermediate_size)
+        .def_readonly("max_length", &BaseConfig::max_length)
+        .def_readonly("bos_token_id", &BaseConfig::bos_token_id)
+        .def_readonly("eos_token_id", &BaseConfig::eos_token_id)
+        .def_readonly("pad_token_id", &BaseConfig::pad_token_id)
+        .def_readonly("sep_token_id", &BaseConfig::sep_token_id);
 
     py::class_<GenerationConfig>(m, "GenerationConfig")
         .def(py::init<int, int, bool, int, float, float, int>(), "max_length"_a = 2048, "max_context_length"_a = 512,
@@ -25,28 +34,35 @@ PYBIND11_MODULE(_C, m) {
         .def_readwrite("temperature", &GenerationConfig::temperature)
         .def_readwrite("num_threads", &GenerationConfig::num_threads);
 
-    py::class_<ChatGLMConfig>(m, "ChatGLMConfig")
-        .def_readonly("vocab_size", &ChatGLMConfig::vocab_size)
-        .def_readonly("hidden_size", &ChatGLMConfig::hidden_size)
-        .def_readonly("num_attention_heads", &ChatGLMConfig::num_attention_heads)
-        .def_readonly("num_layers", &ChatGLMConfig::num_layers)
-        .def_readonly("max_sequence_length", &ChatGLMConfig::max_sequence_length)
-        .def_readonly("bos_token_id", &ChatGLMConfig::bos_token_id)
-        .def_readonly("eos_token_id", &ChatGLMConfig::eos_token_id)
-        .def_readonly("gmask_token_id", &ChatGLMConfig::gmask_token_id)
-        .def_readonly("mask_token_id", &ChatGLMConfig::mask_token_id)
-        .def_readonly("pad_token_id", &ChatGLMConfig::pad_token_id)
-        .def_readonly("dtype", &ChatGLMConfig::dtype);
+    py::class_<ChatGLMConfig, BaseConfig>(m, "ChatGLMConfig");
+
+    py::class_<ChatGLMTokenizer>(m, "ChatGLMTokenizer")
+        .def("encode", &ChatGLMTokenizer::encode)
+        .def("decode", &ChatGLMTokenizer::decode)
+        .def("encode_history", &ChatGLMTokenizer::encode_history);
 
     py::class_<ChatGLMForConditionalGeneration>(m, "ChatGLMForConditionalGeneration")
         .def_readonly("config", &ChatGLMForConditionalGeneration::config)
+        .def_property_readonly("type_name", &ChatGLMForConditionalGeneration::type_name)
         .def("generate_next_token", &ChatGLMForConditionalGeneration::generate_next_token);
 
-    py::class_<ChatGLMPipeline>(m, "ChatGLMPipeline")
+    py::class_<ChatGLM2Config, BaseConfig>(m, "ChatGLM2Config")
+        .def_readonly("num_kv_heads", &ChatGLM2Config::num_kv_heads);
+
+    py::class_<ChatGLM2Tokenizer>(m, "ChatGLM2Tokenizer")
+        .def("encode", &ChatGLM2Tokenizer::encode)
+        .def("decode", &ChatGLM2Tokenizer::decode)
+        .def("encode_history", &ChatGLM2Tokenizer::encode_history);
+
+    py::class_<ChatGLM2ForConditionalGeneration>(m, "ChatGLM2ForConditionalGeneration")
+        .def_readonly("config", &ChatGLM2ForConditionalGeneration::config)
+        .def_property_readonly("type_name", &ChatGLM2ForConditionalGeneration::type_name)
+        .def("generate_next_token", &ChatGLM2ForConditionalGeneration::generate_next_token);
+
+    py::class_<Pipeline>(m, "Pipeline")
         .def(py::init<const std::string &>())
-        .def_property_readonly("model", [](const ChatGLMPipeline &self) { return self.model.get(); })
-        .def_property_readonly("tokenizer", [](const ChatGLMPipeline &self) { return self.tokenizer.get(); })
-        .def_static("build_prompt", &ChatGLMPipeline::build_prompt);
+        .def_property_readonly("model", [](const Pipeline &self) { return self.model.get(); })
+        .def_property_readonly("tokenizer", [](const Pipeline &self) { return self.tokenizer.get(); });
 }
 
 } // namespace chatglm
diff --git a/cli_chat.py b/cli_chat.py
index 53fd61d..1750fbf 100644
--- a/cli_chat.py
+++ b/cli_chat.py
@@ -30,7 +30,7 @@ def main():
     parser.add_argument("-t", "--threads", default=0, type=int)
     args = parser.parse_args()
 
-    pipeline = chatglm_cpp.ChatGLMPipeline(args.model)
+    pipeline = chatglm_cpp.Pipeline(args.model)
 
     if not args.interactive:
         for piece in pipeline.stream_chat(
@@ -50,13 +50,13 @@ def main():
     history = []
     while True:
         try:
-            prompt = input("Prompt  > ")
+            prompt = input(f"{'Prompt':{len(pipeline.model.type_name)}} > ")
         except EOFError:
             break
         if not prompt:
             continue
         history.append(prompt)
-        print("ChatGLM > ", sep="", end="")
+        print(f"{pipeline.model.type_name} > ", sep="", end="")
         output = ""
         for piece in pipeline.stream_chat(
             history,
diff --git a/convert.py b/convert.py
index 37b84e1..5816a72 100644
--- a/convert.py
+++ b/convert.py
@@ -1,5 +1,5 @@
 """
-Convert Hugging Face ChatGLM models to GGML format
+Convert Hugging Face ChatGLM/ChatGLM2 models to GGML format
 """
 import argparse
 import platform
@@ -23,41 +23,16 @@
     sys.modules["cpm_kernels"] = object()
 
 
-class GgmlType(Enum):
+class GGMLType(Enum):
     F32 = 0
     F16 = 1
     Q4_0 = 2
     Q8_0 = 8
 
 
-TORCH_TO_GGML_DTYPE = {
-    torch.float32: GgmlType.F32,
-    torch.float16: GgmlType.F16,
-    torch.int8: GgmlType.Q8_0,
-}
-
-
-def dump_config(f, config, ggml_type):
-    config_values = [
-        config.vocab_size,
-        config.hidden_size,
-        config.num_attention_heads,
-        config.num_layers,
-        config.max_sequence_length,
-        config.bos_token_id,
-        config.eos_token_id,
-        config.gmask_token_id,
-        config.mask_token_id,
-        config.pad_token_id,
-        ggml_type.value,
-    ]
-    f.write(struct.pack("i" * len(config_values), *config_values))
-
-
-def dump_tokenizer(f, tokenizer):
-    serialized_model_proto = tokenizer.sp_tokenizer.text_tokenizer.sp.serialized_model_proto()
-    f.write(struct.pack("i", len(serialized_model_proto)))
-    f.write(serialized_model_proto)
+class ModelType(Enum):
+    CHATGLM = 1
+    CHATGLM2 = 2
 
 
 def quantize_q8_0(tensor: torch.Tensor) -> torch.CharTensor:
@@ -86,29 +61,24 @@ def quantize_q4_0(tensor: torch.Tensor) -> torch.CharTensor:
     return tensor
 
 
-def dump_tensor(f, name: str, tensor: torch.Tensor, ggml_type: GgmlType):
+def dump_tensor(f, name: str, tensor: torch.Tensor, ggml_type: GGMLType):
     assert tensor.dtype == torch.float32
 
-    # header
-    f.write(
-        struct.pack(
-            "i" * (3 + tensor.ndim),
-            tensor.ndim,
-            len(name.encode()),
-            ggml_type.value,
-            *tensor.shape,
-        )
-    )
+    # tensor name
+    f.write(struct.pack("i", len(name.encode())))
     f.write(name.encode())
 
-    # data
-    if ggml_type == GgmlType.F32:
+    # tensor shape & dtype
+    f.write(struct.pack("i" * (2 + tensor.ndim), tensor.ndim, *tensor.shape, ggml_type.value))
+
+    # tensor data
+    if ggml_type == GGMLType.F32:
         tensor = tensor.float()
-    elif ggml_type == GgmlType.F16:
+    elif ggml_type == GGMLType.F16:
         tensor = tensor.half()
-    elif ggml_type == GgmlType.Q8_0:
+    elif ggml_type == GGMLType.Q8_0:
         tensor = quantize_q8_0(tensor)
-    elif ggml_type == GgmlType.Q4_0:
+    elif ggml_type == GGMLType.Q4_0:
         tensor = quantize_q4_0(tensor)
     else:
         raise NotImplementedError(f"Cannot dump tensor of dtype {tensor.dtype}")
@@ -119,30 +89,8 @@ def dump_tensor(f, name: str, tensor: torch.Tensor, ggml_type: GgmlType):
     tensor.numpy().tofile(f)
 
 
-def dump_model(f, model, ggml_type):
-    weight_names = ["transformer.word_embeddings.weight"]
-    for i in range(model.config.num_layers):
-        weight_names += [
-            f"transformer.layers.{i}.input_layernorm.weight",
-            f"transformer.layers.{i}.input_layernorm.bias",
-            f"transformer.layers.{i}.attention.query_key_value.weight",
-            f"transformer.layers.{i}.attention.query_key_value.bias",
-            f"transformer.layers.{i}.attention.dense.weight",
-            f"transformer.layers.{i}.attention.dense.bias",
-            f"transformer.layers.{i}.post_attention_layernorm.weight",
-            f"transformer.layers.{i}.post_attention_layernorm.bias",
-            f"transformer.layers.{i}.mlp.dense_h_to_4h.weight",
-            f"transformer.layers.{i}.mlp.dense_h_to_4h.bias",
-            f"transformer.layers.{i}.mlp.dense_4h_to_h.weight",
-            f"transformer.layers.{i}.mlp.dense_4h_to_h.bias",
-        ]
-    weight_names += [
-        "transformer.final_layernorm.weight",
-        "transformer.final_layernorm.bias",
-    ]
-
+def dump_state_dict(f, weight_names, state_dict, quantization_bit, ggml_type):
     tensor_info = []
-    state_dict = model.state_dict()
     for name in tqdm(weight_names, desc="Dumping model state"):
         tensor = state_dict[name]
         if tensor.ndim == 2:
@@ -150,10 +98,10 @@ def dump_model(f, model, ggml_type):
 
             # step 1: de-quantize it back to float32
             if tensor.dtype == torch.int8:
-                assert model.config.quantization_bit in [4, 8]
+                assert quantization_bit in [4, 8]
                 scale = state_dict[f"{name}_scale"].float()  # channel-wise scale
 
-                if model.config.quantization_bit == 4:
+                if quantization_bit == 4:
                     # convert int4 weight to int8
                     low_bits = ((tensor << 4) & 0xF0) >> 4
                     high_bits = (tensor & 0xF0) >> 4
@@ -168,7 +116,7 @@ def dump_model(f, model, ggml_type):
             # 1d weight: convert it to float32
             assert tensor.ndim == 1
             tensor = tensor.float()
-            tensor_ggml_type = GgmlType.F32
+            tensor_ggml_type = GGMLType.F32
 
         dump_tensor(f, name, tensor, tensor_ggml_type)
         tensor_info.append((name, tensor.shape, tensor_ggml_type.name))
@@ -176,6 +124,140 @@ def dump_model(f, model, ggml_type):
     print(tabulate(tensor_info, headers=["name", "shape", "dtype"], tablefmt="psql"))
 
 
+class BaseConverter:
+    @classmethod
+    def convert(cls, model, tokenizer, ggml_type, save_path):
+        # convert all weights to fp16
+        with open(save_path, "wb") as f:
+            f.write(b"ggml")  # magic
+            f.write(struct.pack("ii", cls.MODEL_TYPE.value, 1))  # model type & version
+            cls.dump_config(f, model.config, ggml_type)
+            cls.dump_tokenizer(f, tokenizer)
+            cls.dump_model(f, model, ggml_type)
+
+        print(f"{cls.MODEL_TYPE.name} GGML model saved to {save_path}")
+
+
+class ChatGLMConverter(BaseConverter):
+    MODEL_TYPE = ModelType.CHATGLM
+
+    @staticmethod
+    def dump_config(f, config, ggml_type):
+        assert config.position_encoding_2d, "unimplemented: position_encoding_2d should be True"
+        assert (
+            config.inner_hidden_size == 4 * config.hidden_size
+        ), "unimplemented: inner_hidden_size should be 4 times hidden_size"
+        config_values = [
+            ggml_type.value,
+            config.vocab_size,
+            config.hidden_size,
+            config.num_attention_heads,
+            config.num_layers,
+            config.inner_hidden_size,
+            config.max_sequence_length,
+            config.bos_token_id,
+            config.eos_token_id,
+            config.pad_token_id,
+            config.sep_token_id if config.sep_token_id is not None else -1,
+        ]
+        f.write(struct.pack("i" * len(config_values), *config_values))
+
+    @staticmethod
+    def dump_tokenizer(f, tokenizer):
+        serialized_model_proto = tokenizer.sp_tokenizer.text_tokenizer.sp.serialized_model_proto()
+        f.write(struct.pack("i", len(serialized_model_proto)))
+        f.write(serialized_model_proto)
+
+    @staticmethod
+    def dump_model(f, model, ggml_type):
+        assert torch.allclose(
+            model.state_dict()["transformer.word_embeddings.weight"], model.state_dict()["lm_head.weight"]
+        ), "unimplemented: lm_head weight must be tied to input embedding"
+
+        weight_names = ["transformer.word_embeddings.weight"]
+        for i in range(model.config.num_layers):
+            weight_names += [
+                f"transformer.layers.{i}.input_layernorm.weight",
+                f"transformer.layers.{i}.input_layernorm.bias",
+                f"transformer.layers.{i}.attention.query_key_value.weight",
+                f"transformer.layers.{i}.attention.query_key_value.bias",
+                f"transformer.layers.{i}.attention.dense.weight",
+                f"transformer.layers.{i}.attention.dense.bias",
+                f"transformer.layers.{i}.post_attention_layernorm.weight",
+                f"transformer.layers.{i}.post_attention_layernorm.bias",
+                f"transformer.layers.{i}.mlp.dense_h_to_4h.weight",
+                f"transformer.layers.{i}.mlp.dense_h_to_4h.bias",
+                f"transformer.layers.{i}.mlp.dense_4h_to_h.weight",
+                f"transformer.layers.{i}.mlp.dense_4h_to_h.bias",
+            ]
+        weight_names += [
+            "transformer.final_layernorm.weight",
+            "transformer.final_layernorm.bias",
+        ]
+        dump_state_dict(f, weight_names, model.state_dict(), model.config.quantization_bit, ggml_type)
+
+
+class ChatGLM2Converter(BaseConverter):
+    MODEL_TYPE = ModelType.CHATGLM2
+
+    @staticmethod
+    def dump_config(f, config, ggml_type):
+        assert config.add_bias_linear is False, "unimplemented: add_bias_linear must be false"
+        assert config.add_qkv_bias is True, "unimplemented: add_qkv_bias must be true"
+        assert (
+            config.apply_residual_connection_post_layernorm is False
+        ), "unimplemented: apply_residual_connection_post_layernorm must be false"
+        assert (
+            config.kv_channels * config.num_attention_heads == config.hidden_size
+        ), "unimplemented: invalid kv_channels"
+        assert config.multi_query_attention is True, "unimplemented: multi_query_attention must be true"
+        assert config.original_rope is True, "unimplemented: original_rope must be true"
+        assert config.post_layer_norm is True, "unimplemented: post_layer_norm must be true"
+        assert config.rmsnorm is True, "unimplemented: rmsnorm must be true"
+
+        config_values = [
+            ggml_type.value,
+            config.padded_vocab_size,
+            config.hidden_size,
+            config.num_attention_heads,
+            config.num_layers,
+            config.ffn_hidden_size,
+            config.seq_length,
+            config.bos_token_id if config.bos_token_id is not None else -1,
+            config.eos_token_id,
+            config.pad_token_id,
+            config.sep_token_id if config.sep_token_id is not None else -1,
+            config.multi_query_group_num,
+        ]
+
+        f.write(struct.pack("i" * len(config_values), *config_values))
+
+    @staticmethod
+    def dump_tokenizer(f, tokenizer):
+        serialized_model_proto = tokenizer.tokenizer.sp_model.serialized_model_proto()
+        f.write(struct.pack("i", len(serialized_model_proto)))
+        f.write(serialized_model_proto)
+
+    @staticmethod
+    def dump_model(f, model, ggml_type):
+        weight_names = ["transformer.embedding.word_embeddings.weight"]
+        for i in range(model.config.num_layers):
+            weight_names += [
+                f"transformer.encoder.layers.{i}.input_layernorm.weight",
+                f"transformer.encoder.layers.{i}.self_attention.query_key_value.weight",
+                f"transformer.encoder.layers.{i}.self_attention.query_key_value.bias",
+                f"transformer.encoder.layers.{i}.self_attention.dense.weight",
+                f"transformer.encoder.layers.{i}.post_attention_layernorm.weight",
+                f"transformer.encoder.layers.{i}.mlp.dense_h_to_4h.weight",
+                f"transformer.encoder.layers.{i}.mlp.dense_4h_to_h.weight",
+            ]
+        weight_names += [
+            "transformer.encoder.final_layernorm.weight",
+            "transformer.output_layer.weight",
+        ]
+        dump_state_dict(f, weight_names, model.state_dict(), model.config.quantization_bit, ggml_type)
+
+
 def main():
     parser = argparse.ArgumentParser("chatglm-convert")
     parser.add_argument("-i", "--model_name_or_path", type=str, default="THUDM/chatglm-6b")
@@ -184,7 +266,7 @@ def main():
     parser.add_argument("-t", "--type", type=str, default="q4_0", choices=["f32", "f16", "q8_0", "q4_0"])
     args = parser.parse_args()
 
-    ggml_type = GgmlType[args.type.upper()]
+    ggml_type = GGMLType[args.type.upper()]
     tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, trust_remote_code=True)
     model = AutoModel.from_pretrained(args.model_name_or_path, trust_remote_code=True)
     if args.lora_model_name_or_path is not None:
@@ -193,22 +275,10 @@ def main():
         model = PeftModel.from_pretrained(model, args.lora_model_name_or_path)
         model = model.merge_and_unload()
 
-    # check config
-    assert model.config.position_encoding_2d, "unimplemented: position_encoding_2d should be True"
-    assert (
-        model.config.inner_hidden_size == 4 * model.config.hidden_size
-    ), "unimplemented: inner_hidden_size should be 4 times hidden_size"
-    assert torch.allclose(
-        model.state_dict()["transformer.word_embeddings.weight"], model.state_dict()["lm_head.weight"]
-    ), "unimplemented: lm_head weight must be tied to input embedding"
-
-    # convert all weights to fp16
-    with open(args.save_path, "wb") as f:
-        f.write(b"ggml")  # magic
-        dump_config(f, model.config, ggml_type)
-        dump_tokenizer(f, tokenizer)
-        dump_model(f, model, ggml_type)
-    print(f"GGML model saved to {args.save_path}")
+    if hasattr(model.config, "multi_query_attention"):
+        ChatGLM2Converter.convert(model, tokenizer, ggml_type, args.save_path)
+    else:
+        ChatGLMConverter.convert(model, tokenizer, ggml_type, args.save_path)
 
 
 if __name__ == "__main__":
diff --git a/main.cpp b/main.cpp
index a423929..268ae95 100644
--- a/main.cpp
+++ b/main.cpp
@@ -1,5 +1,6 @@
 #include "chatglm.h"
 #include <getopt.h>
+#include <iomanip>
 #include <iostream>
 
 struct Args {
@@ -102,7 +103,8 @@ static Args parse_args(int argc, char **argv) {
 }
 
 void chat(const Args &args) {
-    chatglm::ChatGLMPipeline pipeline(args.model_path);
+    chatglm::Pipeline pipeline(args.model_path);
+    std::string model_name = pipeline.model->type_name();
 
     chatglm::TextStreamer streamer(pipeline.tokenizer.get());
     chatglm::GenerationConfig gen_config(args.max_length, args.max_context_length, args.temp > 0, args.top_k,
@@ -118,7 +120,8 @@ void chat(const Args &args) {
 
         std::vector<std::string> history;
         while (1) {
-            std::cout << "Prompt  > " << std::flush;
+            std::cout << std::setw(model_name.size()) << std::left << "Prompt"
+                      << " > " << std::flush;
             std::string prompt;
             if (!std::getline(std::cin, prompt)) {
                 break;
@@ -127,7 +130,7 @@ void chat(const Args &args) {
                 continue;
             }
             history.emplace_back(std::move(prompt));
-            std::cout << "ChatGLM > ";
+            std::cout << model_name << " > ";
             std::string output = pipeline.chat(history, gen_config, &streamer);
             history.emplace_back(std::move(output));
         }
diff --git a/setup.py b/setup.py
index 56a5719..25e23e5 100644
--- a/setup.py
+++ b/setup.py
@@ -41,7 +41,7 @@ def build_extension(self, ext: CMakeExtension) -> None:
 version = re.search(r'__version__ = "(.*?)"', (HERE / "chatglm_cpp/__init__.py").read_text()).group(1)
 
 setup(
-    name="chatglm_cpp",
+    name="chatglm-cpp",
     version=version,
     author="Jiahao Li",
     author_email="liplus17@163.com",
diff --git a/third_party/ggml b/third_party/ggml
index 46603d1..ee1b372 160000
--- a/third_party/ggml
+++ b/third_party/ggml
@@ -1,1 +1,1 @@
-Subproject commit 46603d14b84cb1c40a05016ef90c029635d43583
+Subproject commit ee1b3727e60403012dd2b57d35b60558f4db66d8
diff --git a/web_demo.py b/web_demo.py
index 180f41f..5f7993d 100644
--- a/web_demo.py
+++ b/web_demo.py
@@ -11,7 +11,7 @@
 parser.add_argument("-t", "--threads", default=0, type=int)
 args = parser.parse_args()
 
-pipeline = chatglm_cpp.ChatGLMPipeline(args.model)
+pipeline = chatglm_cpp.Pipeline(args.model)
 
 """Override Chatbot.postprocess"""
 
@@ -68,7 +68,12 @@ def predict(input, chatbot, max_length, top_p, temperature, history):
     response = ""
     history.append(input)
     for response_piece in pipeline.stream_chat(
-        history, max_length=max_length, do_sample=temperature > 0, top_p=top_p, temperature=temperature
+        history,
+        max_length=max_length,
+        do_sample=temperature > 0,
+        top_p=top_p,
+        temperature=temperature,
+        num_threads=args.threads,
     ):
         response += response_piece
         chatbot[-1] = (parse_text(input), parse_text(response))


Test patch:
diff --git a/chatglm_test.cpp b/chatglm_test.cpp
index a51fd76..71b4136 100644
--- a/chatglm_test.cpp
+++ b/chatglm_test.cpp
@@ -32,111 +32,46 @@ class ChatGLMTest : public ::testing::Test {
 
     void SetUp() override {
         ictx.dtype = GGML_TYPE_F32;
-        ictx.gctx = ggml_init({1024 * 1024, nullptr, false});
+        ictx.gctx = GGMLContext({1024 * 1024, nullptr, false});
 
         scratch_buf.resize(1024 * 1024);
 
-        ctx.gctx = ggml_init({1024 * 1024, nullptr, false});
+        ctx.gctx = GGMLContext({1024 * 1024, nullptr, false});
         ctx.scratch = {.offs = 0, .size = scratch_buf.size(), .data = scratch_buf.data()};
         ctx.gf = {};
         ctx.gf.n_threads = 1;
     }
-
-    void TearDown() override {
-        ggml_free(ictx.gctx);
-        ggml_free(ctx.gctx);
-    }
 };
 
 TEST_F(ChatGLMTest, Embedding) {
-    /* test case generated by:
-import torch
-
-torch.manual_seed(0)
-m = torch.nn.Embedding(4, 3)
-x = torch.tensor([1, 3, 0, 2, 3])
-y = m(x)
-print("x", x.flatten())
-print("weight", m.weight.flatten())
-print("y", y.flatten())
-    */
+    float w_data[]{1.5410, -0.2934, -2.1788, 0.5684,  -1.0845, -1.3986,
+                   0.4033, 0.8380,  -0.7193, -0.4033, -0.5966, 0.1820};
     int x_data[]{1, 3, 0, 2, 3};
-    float weight_data[]{1.5410, -0.2934, -2.1788, 0.5684,  -1.0845, -1.3986,
-                        0.4033, 0.8380,  -0.7193, -0.4033, -0.5966, 0.1820};
     float y_data[]{0.5684,  -1.0845, -1.3986, -0.4033, -0.5966, 0.1820,  1.5410, -0.2934,
                    -2.1788, 0.4033,  0.8380,  -0.7193, -0.4033, -0.5966, 0.1820};
 
-    ggml_tensor *x = ggml_new_tensor_1d(ctx.gctx, GGML_TYPE_I32, 5);
+    ggml_tensor *x = ggml_new_tensor_1d(ctx.gctx.get(), GGML_TYPE_I32, 5);
     memcpy(x->data, x_data, sizeof(x_data));
     Embedding model(&ictx, 4, 3);
-    memcpy(model.weight->data, weight_data, sizeof(weight_data));
-    ggml_tensor *y = model.forward(&ctx, x);
-
-    ggml_build_forward_expand(&ctx.gf, y);
-    ggml_graph_compute(ctx.gctx, &ctx.gf);
-
-    expect_all_close((float *)y->data, y_data, ggml_nelements(y));
-}
-
-TEST_F(ChatGLMTest, LayerNorm) {
-    /* test case generated by:
-import torch
-import torch.nn.functional as F
-
-torch.manual_seed(0)
-x = torch.randn(2, 9)
-print('x', x.flatten())
-w = torch.randn(9)
-print('w', w.flatten())
-b = torch.randn(9)
-print('b', b)
-y = F.layer_norm(x, [9], w, b)
-print('y', y.flatten())
-    */
-    float x_data[]{-1.1258, -1.1524, 0.5667, 0.7935, 0.5988, -1.5551, -0.3414, 1.8530, 0.4681,
-                   -0.1577, -0.1734, 0.1835, 1.3894, 1.5863, 0.9463,  -0.8437, 0.9318, 1.2590};
-    float w_data[]{0.4087, 1.4214, 0.1494, -0.6709, -0.2142, -0.4320, -0.7079, -0.1064, -1.2427};
-    float b_data[]{-0.4762, -0.6859, -1.5051, 0.2550, 1.8042, -0.4502, -0.4830, 0.6577, 0.9387};
-    float y_data[]{-0.9153, -2.2485, -1.4268, -0.2403, 1.6855, 0.1889,  -0.2470, 0.4726, 0.4030,
-                   -0.8501, -2.0141, -1.5776, -0.4376, 1.5300, -0.6553, 0.7756,  0.6091, -0.1405};
-
-    ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx, GGML_TYPE_F32, 9, 2);
-    memcpy(x->data, x_data, sizeof(x_data));
-    LayerNorm model(&ictx, 9);
     memcpy(model.weight->data, w_data, sizeof(w_data));
-    memcpy(model.bias->data, b_data, sizeof(b_data));
     ggml_tensor *y = model.forward(&ctx, x);
 
     ggml_build_forward_expand(&ctx.gf, y);
-    ggml_graph_compute(ctx.gctx, &ctx.gf);
+    ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
 
     expect_all_close((float *)y->data, y_data, ggml_nelements(y));
 }
 
 TEST_F(ChatGLMTest, Linear) {
-    /* test case generated by:
-import torch
-import torch.nn.functional as F
-
-torch.manual_seed(0)
-x = torch.randn(2, 3)
-print('x', x.flatten())
-w = torch.randn(8, 3)
-print('w', w.flatten())
-b = torch.randn(8)
-print('b', b)
-y = F.linear(x, w, b)
-print('y', y.flatten())
-    */
-    float x_data[]{1.5410, -0.2934, -2.1788, 0.5684, -1.0845, -1.3986};
-    float w_data[]{-0.0885, 0.5239, -0.6659, 0.8504,  -1.3527, -1.6959, 0.5667,  0.7935,
-                   0.4397,  0.1124, 0.6408,  0.4412,  -0.2159, -0.7425, 0.5627,  0.2596,
-                   0.5229,  2.3022, -1.4689, -1.5867, 1.2032,  0.0845,  -1.2001, -0.0048};
-    float b_data[]{-1.2904, -0.7911, -0.0209, -0.7185, 0.5186, -1.3125, 0.1920, 0.5428};
-    float y_data[]{-0.1298, 4.6113, -0.3385, -1.6945, -0.8222, -6.0819, -4.2276, 1.0356,
-                   -0.9776, 3.5312, -1.1743, -1.9665, 0.4142,  -4.9518, -0.6050, 1.8991};
-
-    ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx, GGML_TYPE_F32, 3, 2);
+    float w_data[]{-1.1258, -1.1524, -0.2506, -0.4339, 0.8487, 0.6920,  -0.3160, -2.1152,
+                   0.4681,  -0.1577, 1.4437,  0.2660,  0.1665, 0.8744,  -0.1435, -0.1116,
+                   0.9318,  1.2590,  2.0050,  0.0537,  0.6181, -0.4128, -0.8411, -2.3160};
+    float b_data[]{0.3704, 1.4565, 0.9398, 0.7748, 0.1919, 1.2638, -1.2904, -0.7911};
+    float x_data[]{-0.0209, -0.7185, 0.5186, -1.3125, 0.1920, 0.5428};
+    float y_data[]{1.0919, 1.2147, 2.7089, -0.1211, -0.5142, 1.2496, -1.0504, -1.3794,
+                   1.4908, 2.5645, 1.2025, 1.4034,  0.0634,  2.2725, -3.5762, -1.6679};
+
+    ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 3, 2);
     memcpy(x->data, x_data, sizeof(x_data));
 
     // fp32
@@ -149,7 +84,7 @@ print('y', y.flatten())
         ggml_tensor *y = model.forward(&ctx, x);
 
         ggml_build_forward_expand(&ctx.gf, y);
-        ggml_graph_compute(ctx.gctx, &ctx.gf);
+        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
 
         expect_all_close((float *)y->data, y_data, ggml_nelements(y));
     }
@@ -163,122 +98,55 @@ print('y', y.flatten())
         ggml_tensor *y = model.forward(&ctx, x);
 
         ggml_build_forward_expand(&ctx.gf, y);
-        ggml_graph_compute(ctx.gctx, &ctx.gf);
+        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
 
         EXPECT_EQ(y->type, GGML_TYPE_F32);
         expect_all_close((float *)y->data, y_data, ggml_nelements(y), 5e-3);
     }
 }
 
-TEST_F(ChatGLMTest, GLU) {
-    /* test case generated by:
-import torch
-from chatglm.modeling_chatglm import GLU
-
-torch.manual_seed(0)
-m = GLU(3, empty_init=False).eval()
-x = torch.randn(2, 3)
-y = m(x)
-print("x", x.flatten())
-print("dense_h_to_4h.weight", m.dense_h_to_4h.weight.flatten())
-print("dense_h_to_4h.bias", m.dense_h_to_4h.bias.flatten())
-print("dense_4h_to_h.weight", m.dense_4h_to_h.weight.flatten())
-print("dense_4h_to_h.bias", m.dense_4h_to_h.bias.flatten())
-print("y", y.flatten())
-    */
-    float x_data[]{0.6880, 1.2326, 1.1580, 0.1673, -1.5773, 1.0689};
-    float dense_h_to_4h_weight_data[]{-0.0043, 0.3097,  -0.4752, -0.4249, -0.2224, 0.1548,  -0.0114, 0.4578,  -0.0512,
-                                      0.1528,  -0.1745, -0.1135, -0.5516, -0.3824, -0.2380, 0.0214,  0.2282,  0.3464,
-                                      -0.3914, -0.2514, 0.2097,  0.4794,  -0.1188, 0.4320,  -0.0931, 0.0611,  0.5228,
-                                      -0.5356, -0.3635, -0.1462, -0.2251, 0.4988,  -0.3742, -0.2658, -0.4034, -0.5407};
-    float dense_h_to_4h_bias_data[]{-0.3370, 0.4963,  0.2576,  0.2798,  0.0304, -0.2960,
-                                    0.0977,  -0.5391, -0.4172, -0.2976, 0.3643, 0.3385};
-    float dense_4h_to_h_weight_data[]{-0.1280, -0.0104, 0.1846,  0.2870, 0.1146,  0.0390,  0.1936, -0.1700, 0.0538,
-                                      -0.2238, -0.2001, -0.1491, 0.1306, 0.1161,  -0.1710, 0.0872, 0.1585,  -0.0364,
-                                      0.0110,  0.0669,  0.1791,  0.2772, -0.2225, -0.1058, 0.1135, 0.2392,  0.2512,
-                                      0.2547,  0.0575,  -0.2510, 0.0266, -0.1806, -0.2690, 0.2565, 0.2195,  -0.2880};
-    float dense_4h_to_h_bias_data[]{0.0540, -0.0486, -0.0475};
-    float y_data[]{0.1555, -0.2438, 0.0476, 0.2104, 0.1448, 0.0937};
-
-    ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx, GGML_TYPE_F32, 3, 2);
+TEST_F(ChatGLMTest, LayerNorm) {
+    float w_data[]{1.5410, -0.2934, -2.1788, 0.5684, -1.0845, -1.3986, 0.4033, 0.8380, -0.7193};
+    float b_data[]{-0.4033, -0.5966, 0.1820, -0.8567, 1.1006, -1.0712, 0.1227, -0.5663, 0.3731};
+    float x_data[]{0.4397, 0.1124, 0.5433,  -0.3952, 0.2055, -0.4503, -0.5731, -0.5554, 0.5943,
+                   1.5419, 1.8197, -0.5515, -1.3253, 0.1886, -0.0691, -0.4949, -1.4959, -0.1938};
+    float y_data[]{1.1039, -0.6742, -2.4414, -1.3358, 0.5938, 0.2759,  -0.3738, -1.5655, -0.5730,
+                   1.9137, -1.1141, 1.1753,  -1.5275, 0.8437, -1.0652, -0.0398, -1.6891, 0.4602};
+
+    ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 9, 2);
     memcpy(x->data, x_data, sizeof(x_data));
+    LayerNorm model(&ictx, 9);
+    memcpy(model.weight->data, w_data, sizeof(w_data));
+    memcpy(model.bias->data, b_data, sizeof(b_data));
+    ggml_tensor *y = model.forward(&ctx, x);
 
-    GLU model(&ictx, 3);
-    memcpy(model.dense_h_to_4h.weight->data, dense_h_to_4h_weight_data, sizeof(dense_h_to_4h_weight_data));
-    memcpy(model.dense_h_to_4h.bias->data, dense_h_to_4h_bias_data, sizeof(dense_h_to_4h_bias_data));
-    memcpy(model.dense_4h_to_h.weight->data, dense_4h_to_h_weight_data, sizeof(dense_4h_to_h_weight_data));
-    memcpy(model.dense_4h_to_h.bias->data, dense_4h_to_h_bias_data, sizeof(dense_4h_to_h_bias_data));
+    ggml_build_forward_expand(&ctx.gf, y);
+    ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+
+    expect_all_close((float *)y->data, y_data, ggml_nelements(y));
+}
+
+TEST_F(ChatGLMTest, RMSNorm) {
+    float w_data[]{0.4963, 0.7682, 0.0885, 0.1320, 0.3074, 0.6341, 0.4901};
+    float x_data[]{0.1642, 0.3058, 0.2100,  0.9056,  0.6035, 0.8110,  -0.0451,
+                   0.8797, 1.0482, -0.0445, -0.7229, 2.8663, -0.5655, 0.1604};
+    float y_data[]{0.1521, 0.4385, 0.0347,  0.2232,  0.3464, 0.9599,  -0.0412,
+                   0.3489, 0.6436, -0.0031, -0.0763, 0.7043, -0.2866, 0.0628};
+
+    RMSNorm model(&ictx, 7);
+    memcpy(model.weight->data, w_data, sizeof(w_data));
+
+    ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 7, 2);
+    memcpy(x->data, x_data, sizeof(x_data));
     ggml_tensor *y = model.forward(&ctx, x);
 
     ggml_build_forward_expand(&ctx.gf, y);
-    ggml_graph_compute(ctx.gctx, &ctx.gf);
+    ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
 
     expect_all_close((float *)y->data, y_data, ggml_nelements(y));
 }
 
-TEST_F(ChatGLMTest, SelfAttention) {
-    /* test case generated by:
-import torch
-from chatglm.modeling_chatglm import SelfAttention
-
-torch.manual_seed(0)
-m = SelfAttention(16, 2, layer_id=3, empty_init=False).float().eval()
-x = torch.randn(4, 1, 16)  # [seqlen, bs, hidden]
-position_ids = torch.tensor([[[0, 1, 2, 2], [0, 0, 0, 1]]])
-attention_mask = torch.tensor(
-    [
-        [
-            [
-                [0, 0, 0, 1],
-                [0, 0, 0, 1],
-                [0, 0, 0, 1],
-                [0, 0, 0, 0],
-            ]
-        ]
-    ]
-).bool()
-y, layer_past = m(
-    x,
-    position_ids=position_ids,
-    attention_mask=attention_mask,
-    layer_id=m.layer_id,
-    use_cache=True,
-)
-print("x", x.flatten())
-print("query_key_value.weight", m.query_key_value.weight.flatten())
-print("query_key_value.bias", m.query_key_value.bias.flatten())
-print("dense.weight", m.dense.weight.flatten())
-print("dense.bias", m.dense.bias.flatten())
-print("y", y.flatten())
-
-x2 = torch.randn(1, 1, 16)
-position_ids = torch.tensor([[[2], [2]]])
-attention_mask = torch.zeros(1, 1, dtype=torch.bool)
-y2, layer_past = m(
-    x2,
-    position_ids=position_ids,
-    attention_mask=attention_mask,
-    layer_id=m.layer_id,
-    layer_past=layer_past,
-    use_cache=True,
-)
-print("x2", x2.flatten())
-print("y2", y2.flatten())
-
-x3 = torch.randn(1, 1, 16)
-position_ids = torch.tensor([[[2], [3]]])
-attention_mask = torch.zeros(1, 1, dtype=torch.bool)
-y3, layer_past = m(
-    x3,
-    position_ids=position_ids,
-    attention_mask=attention_mask,
-    layer_id=m.layer_id,
-    layer_past=layer_past,
-    use_cache=True,
-)
-print("x3", x3.flatten())
-print("y3", y3.flatten())
-    */
+TEST_F(ChatGLMTest, GLMSelfAttention) {
     float query_key_value_weight_data[]{
         -1.8717e-03, 1.3411e-01,  -2.0576e-01, -1.8398e-01, -9.6289e-02, 6.7039e-02,  -4.9533e-03, 1.9822e-01,
         -2.2186e-02, 6.6153e-02,  -7.5553e-02, -4.9141e-02, -2.3884e-01, -1.6557e-01, -1.0306e-01, 9.2609e-03,
@@ -409,7 +277,7 @@ print("y3", y3.flatten())
     float dense_bias_data[]{0.1240, -0.1587, 0.2446, -0.2486, -0.2395, -0.0591, 0.2042, 0.0250,
                             0.0960, -0.1833, 0.0912, -0.0279, 0.1002,  0.1766,  0.1087, -0.0213};
 
-    SelfAttention model(&ictx, 16, 2, 16);
+    GLMSelfAttention model(&ictx, 16, 2, 16);
     memcpy(model.query_key_value.weight->data, query_key_value_weight_data, sizeof(query_key_value_weight_data));
     memcpy(model.query_key_value.bias->data, query_key_value_bias_data, sizeof(query_key_value_bias_data));
     memcpy(model.dense.weight->data, dense_weight_data, sizeof(dense_weight_data));
@@ -434,86 +302,56 @@ print("y3", y3.flatten())
             1.5804e-01, -1.7552e-01, 2.2854e-01,  -1.2723e-01, -2.5867e-01, -3.2931e-02, 2.1397e-01, 2.2805e-01,
             1.5255e-01, 1.1970e-02,  4.6676e-02,  -9.0893e-02, 1.3114e-01,  5.0670e-02,  2.1247e-01, 4.2616e-02};
 
-        ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx, GGML_TYPE_F32, 16, 4);
+        ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 16, 4);
         memcpy(x->data, x_data, sizeof(x_data));
         ggml_tensor *y = model.forward(&ctx, x, 0, 4);
 
         ggml_build_forward_expand(&ctx.gf, y);
-        ggml_graph_compute(ctx.gctx, &ctx.gf);
+        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
 
         expect_all_close((float *)y->data, y_data, ggml_nelements(y));
     }
 
     // cross attention
     {
+        ctx.gf = {};
+        ctx.gf.n_threads = 1;
+
         float x_data[]{0.4047, -0.6549, 0.0521, 0.3401,  -0.2124, 1.5629, -0.9072, -1.5662,
                        0.0485, 0.9935,  2.1180, -0.1827, -0.7768, 1.7764, -0.5033, 0.0566};
         float y_data[]{0.1011, -0.1058, 0.1434,  -0.1980, -0.1304, 0.0632, 0.2109, 0.0619,
                        0.0893, 0.0146,  -0.1291, -0.1430, 0.1655,  0.0916, 0.1488, -0.0282};
-        ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx, GGML_TYPE_F32, 16, 1);
+        ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 16, 1);
         memcpy(x->data, x_data, sizeof(x_data));
 
         ggml_tensor *y = model.forward(&ctx, x, 4, 4);
 
         ggml_build_forward_expand(&ctx.gf, y);
-        ggml_graph_compute(ctx.gctx, &ctx.gf);
+        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
 
         expect_all_close((float *)y->data, y_data, ggml_nelements(y));
     }
     {
+        ctx.gf = {};
+        ctx.gf.n_threads = 1;
+
         float x_data[]{-0.5481, -0.6917, -0.6559, 0.6949, -0.4671, -0.1680, 0.7585,  1.1881,
                        -0.6305, -0.0654, 0.6188,  2.0020, 0.2952,  0.5314,  -0.5227, 0.0995};
         float y_data[]{0.1479, -0.1224, 0.2174,  -0.1259, -0.1662, 0.0400, 0.2235, 0.1599,
                        0.0633, 0.0484,  -0.0541, -0.1475, 0.1395,  0.0448, 0.2216, 0.0838};
-        ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx, GGML_TYPE_F32, 16, 1);
+        ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 16, 1);
         memcpy(x->data, x_data, sizeof(x_data));
 
         ggml_tensor *y = model.forward(&ctx, x, 5, 4);
 
         ggml_build_forward_expand(&ctx.gf, y);
-        ggml_graph_compute(ctx.gctx, &ctx.gf);
+        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
 
         expect_all_close((float *)y->data, y_data, ggml_nelements(y));
     }
 }
 
 TEST_F(ChatGLMTest, GLMBlock) {
-    /* test case generated by:
-import torch
-from chatglm.modeling_chatglm import GLMBlock
-
-torch.manual_seed(0)
-m = GLMBlock(hidden_size=8, num_attention_heads=2, layernorm_epsilon=1e-5, layer_id=3, empty_init=False).float().eval()
-x = torch.randn(4, 1, 8)  # [seqlen, bs, hidden]
-position_ids = torch.tensor([[[0, 1, 2, 2], [0, 0, 0, 1]]])
-attention_mask = torch.tensor(
-    [
-        [
-            [
-                [0, 0, 0, 1],
-                [0, 0, 0, 1],
-                [0, 0, 0, 1],
-                [0, 0, 0, 0],
-            ]
-        ]
-    ]
-).bool()
-y, = m(x, position_ids=position_ids, attention_mask=attention_mask, layer_id=m.layer_id)
-print("x", x.flatten())
-print("input_layernorm.weight", m.input_layernorm.weight.flatten())
-print("input_layernorm.bias", m.input_layernorm.bias.flatten())
-print("query_key_value.weight", m.attention.query_key_value.weight.flatten())
-print("query_key_value.bias", m.attention.query_key_value.bias.flatten())
-print("dense.weight", m.attention.dense.weight.flatten())
-print("dense.bias", m.attention.dense.bias.flatten())
-print("post_attention_layernorm.weight", m.post_attention_layernorm.weight.flatten())
-print("post_attention_layernorm.bias", m.post_attention_layernorm.bias.flatten())
-print("dense_h_to_4h.weight", m.mlp.dense_h_to_4h.weight.flatten())
-print("dense_h_to_4h.bias", m.mlp.dense_h_to_4h.bias.flatten())
-print("dense_4h_to_h.weight", m.mlp.dense_4h_to_h.weight.flatten())
-print("dense_4h_to_h.bias", m.mlp.dense_4h_to_h.bias.flatten())
-print("y", y.flatten())
-    */
     float x_data[]{-1.0159, -1.5296, -1.2180, 0.8683,  -0.3850, 0.0964, -0.3182, 1.8511,  -0.1479, -1.0800, -0.1396,
                    -1.2574, 1.5435,  0.9686,  -1.5552, 0.7989,  1.1896, -0.1795, -1.8232, -0.0514, -1.6361, -2.5106,
                    -1.2383, -0.9296, -0.0585, 0.7729,  0.4689,  0.2599, 0.7576,  -2.0418, -0.5524, 0.6270};
@@ -611,7 +449,7 @@ print("y", y.flatten())
                    -7.4960, 11.8125, 7.5786,  -9.6781, 7.1442,  14.2660, 4.8352,  -7.2181,  6.3933,  -4.4885, -10.8612,
                    -2.3747, 0.0401,  -0.9771, 6.1194,  2.9300,  2.2652,  6.4303,  -17.4357, -5.0949, 5.3586};
 
-    ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx, GGML_TYPE_F32, 8, 4);
+    ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 8, 4);
     memcpy(x->data, x_data, sizeof(x_data));
 
     GLMBlock model(&ictx, 8, 2, 28, 16);
@@ -632,38 +470,201 @@ print("y", y.flatten())
     ggml_tensor *y = model.forward(&ctx, x, 0, 4);
 
     ggml_build_forward_expand(&ctx.gf, y);
-    ggml_graph_compute(ctx.gctx, &ctx.gf);
+    ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
 
     expect_all_close((float *)y->data, y_data, ggml_nelements(y), 5e-4);
 }
 
-TEST_F(ChatGLMTest, ChatGLMForConditionalGeneration) {
-    ChatGLMConfig config{
-        .vocab_size = 64,
-        .hidden_size = 16,
-        .num_attention_heads = 2,
-        .num_layers = 7,
-        .max_sequence_length = 128,
-        .bos_token_id = 130004,
-        .eos_token_id = 130005,
-        .gmask_token_id = 130001,
-        .mask_token_id = 130000,
-        .pad_token_id = 3,
-        .dtype = GGML_TYPE_F32,
-    };
-    int x_data[]{1, 3, 7, 4};
+TEST_F(ChatGLMTest, GLM2Block) {
+    float input_layernorm_weight[]{0.4692, 0.1864, 0.3191, 0.8249, 0.2995, 0.8105, 0.3017, 0.3836,
+                                   0.5106, 0.0412, 0.4950, 0.4496, 0.4551, 0.4000, 0.8942, 0.8690};
+    float attn_qkv_weight[]{
+        -0.0019, 0.1341,  -0.2058, -0.1840, -0.0963, 0.0670,  -0.0050, 0.1982,  -0.0222, 0.0662,  -0.0756, -0.0491,
+        -0.2388, -0.1656, -0.1031, 0.0093,  0.0988,  0.1500,  -0.1695, -0.1089, 0.0908,  0.2076,  -0.0515, 0.1871,
+        -0.0403, 0.0265,  0.2264,  -0.2319, -0.1574, -0.0633, -0.0974, 0.2160,  -0.1620, -0.1151, -0.1747, -0.2341,
+        -0.1459, 0.2149,  0.1116,  0.1212,  0.0131,  -0.1282, 0.0423,  -0.2334, -0.1806, -0.1289, 0.1577,  0.1466,
+        -0.1109, -0.0090, 0.1599,  0.2485,  0.0992,  0.0338,  0.1676,  -0.1472, 0.0466,  -0.1938, -0.1733, -0.1291,
+        0.1131,  0.1005,  -0.1481, 0.0755,  0.1372,  -0.0316, 0.0095,  0.0579,  0.1551,  0.2400,  -0.1927, -0.0916,
+        0.0983,  0.2071,  0.2176,  0.2206,  0.0498,  -0.2174, 0.0230,  -0.1564, -0.2330, 0.2221,  0.1901,  -0.2494,
+        0.0468,  -0.0421, -0.0411, -0.1144, 0.0961,  -0.1481, 0.0916,  0.1264,  0.1790,  0.0935,  -0.2474, -0.1622,
+        0.1248,  0.0523,  -0.1950, -0.1440, 0.2352,  0.1685,  -0.1090, -0.0629, -0.2381, -0.0045, -0.1883, -0.1928,
+        -0.0138, 0.0375,  -0.1024, 0.1483,  -0.1521, 0.2268,  0.1713,  -0.2108, -0.0622, 0.0113,  0.0365,  0.0593,
+        0.0981,  0.0150,  -0.1220, 0.1183,  -0.2398, -0.1482, -0.0626, -0.1218, -0.0875, -0.2049, -0.0532, 0.0534,
+        -0.1629, -0.0128, 0.1790,  -0.0257, 0.0069,  -0.0216, 0.0506,  0.1590,  0.2368,  0.1588,  0.2374,  -0.0181,
+        -0.2246, -0.1185, 0.1702,  -0.0016, -0.1243, -0.1916, -0.2340, -0.2110, -0.0507, 0.1371,  0.1352,  -0.2411,
+        0.1559,  -0.1956, -0.0529, -0.1014, -0.0482, -0.0491, -0.2243, -0.2159, -0.0391, 0.0032,  -0.1136, 0.0942,
+        -0.2250, -0.0169, 0.2199,  -0.1020, 0.2258,  0.0905,  -0.2256, 0.1582,  -0.0288, -0.1116, 0.1999,  -0.2020,
+        0.0268,  -0.0523, 0.1785,  0.0698,  0.1201,  0.0883,  -0.0601, -0.0526, -0.2060, 0.1355,  0.1985,  0.1711,
+        -0.1763, 0.0111,  -0.1762, -0.1376, -0.1457, 0.0854,  -0.1490, -0.0055, 0.0105,  0.1612,  -0.1890, -0.1716,
+        -0.1452, 0.1750,  -0.0899, 0.2109,  0.0904,  0.0317,  -0.0019, -0.0494, 0.0314,  -0.0571, -0.0018, 0.0319,
+        -0.1956, -0.1310, 0.2019,  -0.2029, -0.0180, 0.2473,  0.0903,  0.0071,  -0.2167, 0.1238,  -0.1781, -0.0710,
+        -0.0839, -0.0370, 0.0027,  0.2062,  0.0312,  0.2239,  0.1529,  -0.1581, 0.1121,  -0.1767, -0.1060, 0.0735,
+        0.0825,  0.1876,  -0.0805, 0.0004,  0.1287,  -0.2418, 0.1807,  -0.2067, 0.0034,  -0.0425, -0.1317, 0.0330,
+        0.2067,  -0.0731, -0.1484, -0.0925, -0.2478, 0.1128,  -0.1201, -0.1668, -0.1440, 0.1437,  0.1324,  0.1919,
+        0.0907,  -0.0835, -0.0699, 0.0739,  0.2055,  0.0680,  -0.1183, -0.1175, -0.2364, 0.0540,  -0.1403, -0.2229,
+        0.2192,  -0.1624, -0.0284, 0.0716,  0.0080,  -0.1682, -0.2021, 0.1993,  0.0407,  0.2074,  -0.0838, 0.0736,
+        -0.0572, -0.0111, -0.1523, 0.0846,  0.0790,  -0.0052, -0.0562, -0.1541, 0.1729,  -0.1861, 0.1024,  -0.0841,
+        -0.1206, 0.0449,  -0.1299, 0.0576,  0.0491,  -0.1856, 0.0416,  0.1065,  0.0990,  -0.0315, -0.2049, -0.0385,
+        0.0868,  -0.0912, 0.0949,  0.1665,  -0.1305, 0.0025,  0.1034,  0.0196,  0.0209,  0.0312,  -0.1965, 0.0196,
+        0.1731,  0.2253,  0.1469,  0.0335,  0.1168,  -0.1216, -0.2072, -0.2150, 0.2494,  0.1587,  -0.1728, 0.0978,
+        0.1888,  0.2499,  0.2186,  0.1937,  -0.0573, -0.0877, 0.2053,  0.1401,  -0.1504, 0.2248,  0.1208,  0.1363,
+        -0.1567, 0.0717,  -0.0876, 0.1953,  -0.0450, 0.0973,  0.0444,  0.1064,  -0.0850, 0.1219,  -0.1746, 0.0565,
+        -0.1691, -0.2466, -0.2008, 0.1974,  0.1353,  0.2346,  0.2003,  -0.2233, -0.1706, -0.0404, -0.1624, 0.1736,
+        -0.1890, -0.1220, -0.2415, -0.1419, 0.2056,  0.2047,  0.1790,  0.1930,  0.2223,  -0.0640, 0.1100,  0.2227,
+        0.0827,  0.2499,  0.1297,  0.1554,  -0.0875, 0.1200,  0.0287,  -0.0597, -0.1409, -0.1403, -0.1924, 0.1678,
+        0.1777,  -0.0285, -0.1447, 0.1932,  0.1599,  0.0186,  -0.1180, 0.2298,  0.1022,  -0.1898, 0.2393,  0.1898,
+        -0.0911, 0.1405,  -0.1420, -0.0392, 0.2123,  0.0103,  -0.1768, -0.0836, -0.0679, -0.0482, 0.0239,  0.2312,
+        0.0134,  -0.1544, 0.0128,  0.1199,  0.1240,  -0.2285, -0.0447, -0.1858, -0.1067, 0.0901,  -0.1775, 0.0929,
+        0.2122,  0.0164,  -0.1666, -0.0896, 0.0546,  -0.1906, 0.1242,  -0.2270, -0.2403, -0.2429, -0.0507, 0.1681,
+        -0.2366, 0.2078,  -0.1000, 0.0732,  0.0114,  -0.2254, 0.2073,  0.1346,  0.2485,  0.1263,  -0.1650, 0.2086,
+        0.0134,  0.1186,  -0.2005, -0.0719, -0.2455, -0.0974, 0.0539,  -0.1963, 0.0797,  0.1342,  0.0348,  -0.1673,
+        -0.1938, -0.0771, 0.1097,  0.2466,  0.1438,  -0.0282, 0.0877,  -0.2453, -0.2135, 0.1167,  -0.1416, 0.1203,
+        -0.1765, -0.1238, -0.2059, 0.1305,  -0.0255, 0.1924,  0.1547,  0.1383,  0.0080,  -0.0773, -0.0544, 0.0332,
+        0.1239,  -0.1751, 0.2098,  -0.0272, -0.2095, -0.1353, 0.2212,  0.2286,  -0.2316, 0.1763,  0.1253,  0.1480,
+        0.2116,  -0.1347, 0.0789,  0.1023,  -0.0739, 0.0837,  -0.0719, 0.1546};
+    float attn_qkv_bias[]{-0.0694, -0.0932, 0.0629,  0.0887,  -0.1221, 0.0221,  0.1449,  -0.0249,
+                          0.0761,  -0.0603, 0.0876,  -0.1811, -0.1470, -0.1269, 0.2298,  -0.0673,
+                          -0.0007, -0.1211, 0.2496,  0.2442,  -0.1885, -0.2027, -0.1895, -0.0012,
+                          -0.0637, -0.1636, -0.0897, 0.0472,  -0.1306, 0.0554,  -0.0573, -0.1211};
+    float attn_dense_weight[]{
+        3.4346e-02,  2.0556e-01,  -1.6902e-01, 1.1609e-02,  -9.2192e-02, 2.4533e-01,  -2.3719e-01, -2.3967e-01,
+        2.4635e-01,  -1.5817e-01, 4.7931e-02,  -2.1577e-02, -5.2664e-02, -5.5837e-02, 1.5886e-01,  1.1949e-02,
+        -2.4340e-01, -1.4759e-01, -8.5225e-02, 1.2580e-01,  -1.6178e-01, 2.3573e-01,  -5.5682e-02, -4.4890e-02,
+        1.9589e-01,  1.2565e-01,  2.1203e-01,  1.4461e-01,  -7.5838e-02, -1.6587e-01, -1.8601e-02, 2.0691e-01,
+        -8.3905e-02, -2.3184e-01, 1.0248e-01,  2.4337e-01,  -7.1171e-02, -2.0701e-01, -2.2676e-01, 6.2650e-02,
+        -1.8930e-02, -1.2625e-01, 5.0533e-02,  9.4940e-02,  1.9883e-01,  1.9410e-01,  -3.7421e-02, -2.2044e-01,
+        -2.2591e-01, 2.3342e-01,  1.1051e-01,  1.0898e-01,  -2.1631e-01, 2.3150e-01,  2.3684e-01,  2.2572e-01,
+        -2.1090e-01, -9.4331e-02, -1.7194e-01, 2.3674e-01,  -1.0742e-01, -1.1414e-01, 1.3098e-01,  -1.1565e-01,
+        -1.2313e-01, -2.1868e-02, -2.4030e-02, -1.9474e-01, 2.0842e-01,  -1.1028e-01, 8.8679e-02,  2.1745e-01,
+        1.2608e-01,  3.5395e-02,  2.1271e-01,  3.3607e-02,  -1.1566e-01, 2.3650e-01,  5.9170e-02,  -2.4392e-01,
+        -7.1173e-02, -1.7029e-01, 2.1919e-01,  -4.1286e-02, -2.2788e-01, -1.5716e-02, 1.5700e-01,  6.4950e-02,
+        7.9051e-02,  2.3218e-02,  9.3205e-02,  -6.0913e-02, -9.9450e-02, -2.3368e-01, -1.8833e-01, 1.0835e-01,
+        -1.4803e-01, 3.5867e-02,  7.9798e-02,  1.7701e-02,  -1.6209e-01, 2.3908e-01,  -1.4537e-01, 2.0563e-01,
+        -1.9888e-01, -6.0137e-02, 1.3599e-01,  -1.0215e-01, 2.1001e-01,  -1.7204e-01, -2.0994e-01, -1.1273e-01,
+        4.0413e-02,  2.3019e-01,  -1.1935e-01, 8.9407e-02,  -6.2677e-02, -5.4216e-02, 1.8383e-01,  -1.9375e-01,
+        2.6553e-02,  2.3510e-01,  -3.4353e-02, 1.9411e-01,  -7.6989e-02, 2.0124e-01,  -2.4183e-01, -3.6031e-02,
+        -4.3912e-02, 8.1021e-02,  9.8082e-02,  1.9195e-01,  -3.7246e-02, -9.9138e-03, 1.7120e-01,  -6.7643e-02,
+        2.1916e-01,  -1.6645e-01, -2.7052e-02, -1.3428e-02, 1.1155e-01,  1.7092e-01,  -3.9626e-02, -2.0712e-01,
+        1.2387e-01,  7.4783e-02,  1.0042e-01,  -1.5422e-01, 1.6090e-01,  2.3679e-01,  2.1693e-02,  -2.3351e-01,
+        1.7550e-01,  -1.8536e-01, 5.7468e-02,  3.6329e-02,  -1.1701e-01, 8.7049e-02,  -2.2361e-01, 5.6940e-02,
+        -1.5849e-01, -2.7033e-02, 3.2157e-02,  2.1299e-01,  -1.1928e-01, 1.6016e-01,  -3.1758e-02, -1.1873e-01,
+        -2.1770e-01, -2.2938e-01, 2.4415e-01,  -6.2350e-02, 1.2487e-02,  6.7778e-02,  1.6993e-01,  2.1337e-01,
+        2.0275e-01,  -1.8522e-01, -4.0054e-02, -1.4793e-01, -1.4284e-01, 5.9302e-02,  2.3466e-01,  -2.0028e-01,
+        1.5130e-01,  -1.2962e-01, -4.8694e-02, 1.9844e-01,  -5.6543e-02, 2.2764e-02,  -1.7476e-01, 2.1281e-01,
+        -3.2317e-02, -1.8285e-01, 7.3203e-02,  -1.7775e-01, -1.9838e-01, 1.5230e-02,  1.9821e-01,  -7.0746e-02,
+        1.1767e-01,  2.1483e-01,  1.6582e-01,  -1.3114e-01, -2.7405e-02, -7.8671e-02, -2.0103e-01, 1.0329e-04,
+        1.8811e-01,  2.1063e-01,  2.3325e-02,  5.6781e-02,  -1.0823e-01, 1.8871e-01,  -1.0403e-01, -1.7366e-01,
+        3.8498e-02,  1.4985e-01,  -2.2540e-01, 2.2599e-01,  8.9933e-02,  -1.7515e-01, -5.3852e-02, 2.1689e-01,
+        -1.9181e-01, -7.3069e-02, 8.1995e-02,  -2.1902e-01, 1.3705e-01,  1.3014e-01,  1.5505e-01,  -1.5939e-01,
+        2.4900e-01,  -1.4819e-01, 2.4959e-01,  -2.3992e-01, -2.2274e-01, 1.5355e-01,  2.6128e-02,  1.4422e-02,
+        -1.3844e-01, -1.0487e-01, -7.3097e-02, -2.4354e-01, 1.2992e-02,  4.4215e-02,  -2.1166e-04, 8.0733e-02,
+        2.3721e-01,  6.6471e-02,  -9.1518e-02, -1.0289e-01, -1.5995e-01, -1.7330e-01, -4.0262e-02, -4.4216e-02};
+    float post_attention_layernorm_weight[]{0.1611, 0.7323, 0.1078, 0.0743, 0.6528, 0.5073, 0.2667, 0.0177,
+                                            0.3064, 0.6670, 0.0372, 0.0143, 0.5634, 0.1398, 0.0620, 0.3074};
+    float mlp_dense_h_to_4h_weight[]{
+        0.1112,  -0.1069, 0.1993,  -0.1754, 0.0007,  0.2247,  0.2486,  -0.1448, 0.0445,  0.0295,  -0.1172, -0.0864,
+        0.0677,  -0.1738, 0.0412,  0.1082,  -0.0985, 0.2077,  -0.0165, 0.1134,  0.2476,  -0.0764, 0.1359,  -0.0715,
+        -0.0365, -0.0424, -0.0016, -0.0944, 0.0586,  0.0094,  0.1585,  -0.0506, 0.0251,  -0.0930, -0.2094, 0.1012,
+        0.0320,  -0.1001, -0.0845, 0.0654,  -0.0452, 0.1634,  0.0142,  0.0944,  0.1089,  -0.0613, 0.1082,  0.1845,
+        0.0115,  0.0489,  0.0091,  0.1731,  -0.1055, -0.1329, 0.1089,  -0.2176, 0.0040,  -0.1150, 0.1650,  -0.2328,
+        0.1516,  0.2473,  0.0506,  -0.0167, 0.2461,  -0.1068, -0.0200, -0.1095, -0.0345, 0.0543,  0.0334,  -0.0472,
+        -0.2500, 0.0128,  -0.0074, 0.0376,  0.1916,  0.2430,  -0.1483, -0.0156, -0.1002, -0.2304, -0.1819, 0.2346,
+        0.1458,  0.1343,  0.1389,  -0.1986, 0.0263,  0.2327,  -0.1395, 0.2224,  0.0657,  0.1759,  -0.1071, 0.1153,
+        -0.2216, -0.0159, 0.0834,  0.0750,  0.2092,  0.2457,  0.2274,  0.1678,  0.1758,  -0.0323, -0.1940, -0.0928,
+        -0.0281, -0.1351, 0.1278,  0.0878,  0.0403,  0.0604,  0.2196,  0.0911,  -0.2192, -0.1815, 0.1102,  0.0341,
+        0.1219,  -0.2497, -0.2307, 0.1533,  0.1599,  -0.2262, 0.0949,  -0.1952, 0.1893,  0.0784,  0.2472,  -0.2465,
+        -0.2407, 0.1690,  0.2220,  0.1532,  0.1662,  -0.2106, 0.1810,  -0.2360, 0.0559,  -0.1638, -0.0993, 0.0443,
+        -0.0983, -0.0396, -0.2370, -0.2357, 0.0968,  -0.1058, -0.1521, -0.0986, -0.2044, -0.1999, 0.1496,  0.1271,
+        0.0333,  -0.1366, 0.0832,  0.0112,  -0.1787, 0.0538,  0.2276,  -0.2459, -0.2486, 0.0320,  0.1883,  -0.1521,
+        0.0550,  0.1757,  0.0771,  0.0541,  0.2494,  0.1629,  -0.1406, 0.0216,  0.1277,  0.2205,  -0.2487, 0.1354,
+        -0.0726, -0.2448, 0.0555,  -0.1361, 0.0354,  -0.1623, -0.1881, -0.0212, -0.0840, 0.1462,  0.0216,  0.1951};
+    float mlp_dense_4h_to_h_weight[]{
+        0.0765,  -0.1313, 0.2765,  0.0223,  0.0766,  -0.1313, 0.0015,  -0.0785, -0.0347, -0.3184, -0.0025, -0.3389,
+        -0.0878, 0.2366,  -0.2839, -0.2752, -0.0997, -0.1119, 0.0691,  0.0678,  0.2070,  -0.2812, -0.0288, 0.2283,
+        0.3531,  -0.0298, -0.2362, 0.1174,  0.2325,  -0.2551, -0.2369, 0.1804,  -0.0286, -0.3856, -0.2354, 0.1656,
+        -0.1385, 0.2818,  0.3227,  0.0594,  -0.0323, -0.1296, -0.0206, 0.0743,  -0.3118, -0.0976, -0.3396, 0.2506,
+        -0.2600, 0.3730,  -0.1052, -0.2339, 0.1961,  0.0609,  0.2827,  0.1704,  -0.3933, 0.2582,  -0.0769, -0.1804,
+        0.2593,  0.2977,  -0.3588, -0.0369, 0.3353,  0.1581,  0.3439,  -0.1399, -0.2252, 0.3511,  0.1702,  0.3919,
+        -0.1705, -0.2621, -0.0479, -0.3844, 0.1600,  0.3011,  0.0980,  -0.0403, 0.2024,  -0.2591, 0.3994,  -0.4059,
+        -0.3911, -0.0965, 0.3335,  0.0409,  0.1568,  -0.2992, 0.1489,  -0.0456, 0.1636,  0.2883,  0.1775,  -0.0347};
+
+    constexpr int seq_len = 3;
+    constexpr int hidden_size = 16;
+    constexpr int num_attention_heads = 4;
+    constexpr int num_kv_heads = 2;
+    constexpr int ffn_hidden_size = 6;
+    constexpr int max_length = 8;
+
+    GLM2Block model(&ictx, hidden_size, num_attention_heads, num_kv_heads, ffn_hidden_size, max_length);
+    memcpy(model.input_layernorm.weight->data, input_layernorm_weight, sizeof(input_layernorm_weight));
+    memcpy(model.attention.query_key_value.weight->data, attn_qkv_weight, sizeof(attn_qkv_weight));
+    memcpy(model.attention.query_key_value.bias->data, attn_qkv_bias, sizeof(attn_qkv_bias));
+    memcpy(model.attention.dense.weight->data, attn_dense_weight, sizeof(attn_dense_weight));
+    memcpy(model.post_attention_layernorm.weight->data, post_attention_layernorm_weight,
+           sizeof(post_attention_layernorm_weight));
+    memcpy(model.mlp.dense_h_to_4h.weight->data, mlp_dense_h_to_4h_weight, sizeof(mlp_dense_h_to_4h_weight));
+    memcpy(model.mlp.dense_4h_to_h.weight->data, mlp_dense_4h_to_h_weight, sizeof(mlp_dense_4h_to_h_weight));
 
-    ggml_tensor *x = ggml_new_tensor_1d(ctx.gctx, GGML_TYPE_I32, 4);
-    memcpy(x->data, x_data, sizeof(x_data));
-    ChatGLMForConditionalGeneration model(&ictx, config);
-    ggml_tensor *y = model.forward(&ctx, x, 0, 4);
+    // self attention
+    {
+        float x_data[]{-0.3439, 0.2291,  0.4366,  0.2702,  -0.2002, 1.0184, -0.4169, -0.0308, -0.9846, 0.2024,
+                       0.4156,  -1.2037, -0.1333, -0.2934, 0.4649,  0.5762, -0.0578, 1.1715,  -1.5185, 0.0475,
+                       -0.0103, 1.4353,  0.2518,  0.7408,  -1.0982, 0.1548, -1.3989, 0.0408,  -1.1174, 0.7029,
+                       -0.6675, 0.1254,  0.4047,  -0.6549, 0.0521,  0.3401, -0.2124, 1.5629,  -0.9072, -1.5662,
+                       0.0485,  0.9935,  2.1180,  -0.1827, -0.7768, 1.7764, -0.5033, 0.0566};
+        float y_data[]{-0.3964, 0.2281,  0.6048,  -0.3200, -0.0622, 0.9805,  -0.7734, -0.1910, -1.1877, 0.1297,
+                       0.4010,  -1.0275, -0.0449, -0.5022, 0.5744,  0.6155,  -0.1103, 1.1213,  -1.3659, -0.3552,
+                       0.1080,  1.3184,  0.0149,  0.6601,  -1.1811, 0.0690,  -1.4537, 0.0565,  -1.1463, 0.5818,
+                       -0.6814, 0.0777,  0.2416,  -0.6671, 0.2555,  -0.0109, -0.1029, 1.5266,  -1.0859, -1.6493,
+                       -0.0626, 0.8874,  2.0539,  -0.1481, -0.8140, 1.6529,  -0.5594, 0.0603};
+
+        ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, hidden_size, seq_len);
+        memcpy(x->data, x_data, sizeof(x_data));
 
-    ggml_build_forward_expand(&ctx.gf, y);
-    ggml_graph_compute(ctx.gctx, &ctx.gf);
+        ggml_tensor *y = model.forward(&ctx, x, 0);
+        ggml_build_forward_expand(&ctx.gf, y);
+        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+
+        expect_all_close((float *)y->data, y_data, ggml_nelements(y));
+    }
+
+    // cross attention
+    {
+        ctx.gf = {};
+        ctx.gf.n_threads = 1;
+
+        float x2_data[]{-0.5481, -0.6917, -0.6559, 0.6949, -0.4671, -0.1680, 0.7585,  1.1881,
+                        -0.6305, -0.0654, 0.6188,  2.0020, 0.2952,  0.5314,  -0.5227, 0.0995};
+        float y2_data[]{-0.6584, -0.6923, -0.5611, 0.5060, -0.4226, -0.2116, 0.5962,  1.1427,
+                        -0.7516, -0.1508, 0.6161,  2.0421, 0.2349,  0.4791,  -0.6062, 0.0516};
+
+        ggml_tensor *x2 = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, hidden_size, 1);
+        memcpy(x2->data, x2_data, sizeof(x2_data));
+
+        ggml_tensor *y2 = model.forward(&ctx, x2, seq_len);
+        ggml_build_forward_expand(&ctx.gf, y2);
+        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+
+        expect_all_close((float *)y2->data, y2_data, ggml_nelements(y2));
+    }
+    {
+        ctx.gf = {};
+        ctx.gf.n_threads = 1;
+
+        float x3_data[]{-0.1169, 0.3728,  -1.0211, 1.0795, 1.1469,  -0.1733, 0.0637,  -1.2699,
+                        -0.6212, -0.2381, 0.0892,  1.8008, -2.0627, 0.3222,  -1.1390, 1.2418};
+        float y3_data[]{-0.2959, 0.3463,  -0.9365, 0.9941, 1.1817,  -0.1974, -0.0675, -1.3237,
+                        -0.7128, -0.3565, 0.0678,  1.8228, -2.1314, 0.3017,  -1.2696, 1.1742};
+
+        ggml_tensor *x3 = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, hidden_size, 1);
+        memcpy(x3->data, x3_data, sizeof(x3_data));
 
-    // just check shape & dtype
-    EXPECT_TRUE(has_shape(y, {64}));
-    EXPECT_EQ(y->type, GGML_TYPE_F32);
+        ggml_tensor *y3 = model.forward(&ctx, x3, seq_len + 1);
+        ggml_build_forward_expand(&ctx.gf, y3);
+        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+
+        expect_all_close((float *)y3->data, y3_data, ggml_nelements(y3));
+    }
 }
 
 TEST_F(ChatGLMTest, quantize) {
@@ -702,12 +703,12 @@ TEST_F(ChatGLMTest, quantize) {
         -1.4777e+00, -1.7557e+00, 7.6166e-02,  -1.0786e+00, 1.4403e+00,  -1.1059e-01, 5.7686e-01,  -1.6917e-01,
         -6.4025e-02, 1.0384e+00,  9.0682e-01,  -4.7551e-01, -8.7074e-01, 1.4474e-01,  1.9029e+00,  3.9040e-01};
 
-    ggml_tensor *src = ggml_new_tensor_2d(ctx.gctx, GGML_TYPE_F32, 128, 2);
+    ggml_tensor *src = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 128, 2);
     memcpy(src->data, src_data, sizeof(src_data));
 
     // q8_0
     {
-        ggml_tensor *q8_dst = ggml_new_tensor_2d(ctx.gctx, GGML_TYPE_Q8_0, 128, 2);
+        ggml_tensor *q8_dst = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_Q8_0, 128, 2);
         int64_t hist[16]{};
         ggml_quantize_q8_0((float *)src->data, q8_dst->data, ggml_nelements(src), src->ne[0], hist);
 
@@ -719,7 +720,7 @@ TEST_F(ChatGLMTest, quantize) {
     }
     // q4_0
     {
-        ggml_tensor *q4_dst = ggml_new_tensor_2d(ctx.gctx, GGML_TYPE_Q4_0, 128, 2);
+        ggml_tensor *q4_dst = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_Q4_0, 128, 2);
         int64_t hist[16]{};
         ggml_quantize_q4_0((float *)src->data, q4_dst->data, ggml_nelements(src), src->ne[0], hist);
 
@@ -743,62 +744,125 @@ static bool equal(const std::vector<int> &a, const std::vector<int> &b) {
     return true;
 }
 
-TEST(ChatGLM, chat) {
+TEST(Pipeline, ChatGLM) {
     fs::path model_path = fs::path(__FILE__).parent_path() / "chatglm-ggml.bin";
     if (!fs::exists(model_path)) {
         GTEST_SKIP() << "Skipping ChatGLM e2e test (ggml model not found)";
     }
-    ChatGLMPipeline pipeline(model_path);
+    Pipeline pipeline(model_path);
+    EXPECT_TRUE(dynamic_cast<ChatGLMForConditionalGeneration *>(pipeline.model.get()));
 
     // ===== tokenization =====
 
-    // encoding
-    {
-        std::string prompt("你好");
-        std::vector<int> input_ids = pipeline.tokenizer->encode(prompt);
-        std::vector<int> target_ids{5, 74874, 130001, 130004};
-        EXPECT_TRUE(equal(input_ids, target_ids));
+    struct TokenizerTestCase {
+        std::string prompt;
+        std::vector<int> input_ids;
+    };
+    std::vector<TokenizerTestCase> cases{
+        {"你好", {5, 74874, 130001, 130004}},
+        {"[Round 0]\n问：你好\n答：你好👋！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。\n[Round "
+         "1]\n问：晚上睡不着应该怎么办\n答：",
+         {53,     6945,   5,      8,     42,    4,     64286,  12,    74874, 4,   67342,  12,    74874, 130328,
+          130247, 130233, 130227, 35,    65806, 68241, 75890,  14132, 5388,  340, 11,     21,    222,   6,
+          76693,  66877,  63852,  6,     66430, 68747, 102501, 63823, 4,     52,  6945,   5,     9,     42,
+          4,      64286,  12,     65450, 83400, 64213, 66846,  4,     67342, 12,  130001, 130004}},
+        {"def main():\n    print('hello world')\t# greeting",
+         {1616, 594, 125936, 4, 130011, 2274, 89, 7283, 398, 125686, 130008, 61, 25672, 130001, 130004}}};
+
+    for (const auto &c : cases) {
+        // encode
+        std::vector<int> input_ids = pipeline.tokenizer->encode(c.prompt);
+        EXPECT_TRUE(equal(input_ids, c.input_ids));
+        // decode
+        std::string output = pipeline.tokenizer->decode(c.input_ids);
+        EXPECT_EQ(output, c.prompt);
     }
+
+    // ===== prompter =====
     {
-        std::string prompt("[Round 0]\n问：你好\n答：你好👋！我是人工智能助手 "
-                           "ChatGLM-6B，很高兴见到你，欢迎问我任何问题。\n[Round 1]\n问：晚上睡不着应该怎么办\n答：");
-        std::vector<int> input_ids = pipeline.tokenizer->encode(prompt);
-        std::vector<int> target_ids{53,    6945,   5,      8,      42,     4,      64286, 12,    74874,  4,     67342,
-                                    12,    74874,  130328, 130247, 130233, 130227, 35,    65806, 68241,  75890, 14132,
-                                    5388,  340,    11,     21,     222,    6,      76693, 66877, 63852,  6,     66430,
-                                    68747, 102501, 63823,  4,      52,     6945,   5,     9,     42,     4,     64286,
-                                    12,    65450,  83400,  64213,  66846,  4,      67342, 12,    130001, 130004};
-        EXPECT_TRUE(equal(input_ids, target_ids));
+        EXPECT_EQ(ChatGLMTokenizer::build_prompt({"你好"}), "你好");
+        EXPECT_EQ(ChatGLMTokenizer::build_prompt(
+                      {"你好", "你好👋！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。",
+                       "晚上睡不着应该怎么办"}),
+                  "[Round 0]\n问：你好\n答：你好👋！我是人工智能助手 "
+                  "ChatGLM-6B，很高兴见到你，欢迎问我任何问题。\n[Round 1]\n问：晚上睡不着应该怎么办\n答：");
     }
+
+    // ===== generation =====
+
+    // memory test
     {
-        // white space
-        std::string prompt("def main():\n    print('hello world')\t# greeting");
-        std::vector<int> input_ids = pipeline.tokenizer->encode(prompt);
-        std::vector<int> target_ids{1616, 594,    125936, 4,  130011, 2274,   89,    7283,
-                                    398,  125686, 130008, 61, 25672,  130001, 130004};
-        EXPECT_TRUE(equal(input_ids, target_ids));
+        GenerationConfig gen_config;
+        gen_config.max_length = gen_config.max_context_length + 1;
+        gen_config.do_sample = false;
+
+        std::ostringstream oss;
+        for (int i = 0; i < gen_config.max_context_length; i++) {
+            oss << "你好";
+        }
+        std::vector<std::string> history{oss.str()};
+        pipeline.chat(history, gen_config);
     }
-    // decoding
+
+    // chat
     {
-        std::vector<int> output_ids{1616, 594,    125936, 4,  130011, 2274,   89,    7283,
-                                    398,  125686, 130008, 61, 25672,  130001, 130004};
-        std::string output = pipeline.tokenizer->decode(output_ids);
-        std::string target("def main():\n    print('hello world')\t# greeting");
-        EXPECT_EQ(target, output);
+        GenerationConfig gen_config;
+        gen_config.do_sample = false;
+        std::vector<std::string> history{"你好"};
+        std::string output = pipeline.chat(history, gen_config);
+        EXPECT_EQ(output, "你好👋！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。");
     }
+}
 
-    // ===== generation =====
+TEST(Pipeline, ChatGLM2) {
+    fs::path model_path = fs::path(__FILE__).parent_path() / "chatglm2-ggml.bin";
+    if (!fs::exists(model_path)) {
+        GTEST_SKIP() << "Skipping ChatGLM2 e2e test (ggml model not found)";
+    }
+    Pipeline pipeline(model_path);
+    EXPECT_TRUE(dynamic_cast<ChatGLM2ForConditionalGeneration *>(pipeline.model.get()));
+
+    // ===== tokenization =====
 
-    // build prompt
+    struct TokenizerTestCase {
+        std::string prompt;
+        std::vector<int> input_ids;
+    };
+    std::vector<TokenizerTestCase> cases{
+        {"你好", {64790, 64792, 36474, 54591}},
+        {"[Round 1]\n\n问：你好\n\n答：",
+         {64790, 64792, 790, 30951, 517, 30910, 30939, 30996, 13, 13, 54761, 31211, 39701, 13, 13, 55437, 31211}},
+        {"[Round 1]\n\n问：你好\n\n答：你好👋！我是人工智能助手 "
+         "ChatGLM2-6B，很高兴见到你，欢迎问我任何问题。\n\n[Round 2]\n\n问：晚上睡不着应该怎么办\n\n答：",
+         {64790, 64792, 790,   30951, 517,   30910, 30939, 30996, 13,    13,    54761, 31211, 39701,
+          13,    13,    55437, 31211, 39701, 243,   162,   148,   142,   31404, 33030, 34797, 42481,
+          22011, 10461, 30944, 30943, 30941, 30978, 30949, 31123, 48895, 35214, 54622, 31123, 32616,
+          39905, 31901, 31639, 31155, 13,    13,    30995, 30951, 517,   30910, 30943, 30996, 13,
+          13,    54761, 31211, 32820, 54266, 31876, 35153, 13,    13,    55437, 31211}},
+        {"def main():\n    print('hello world')\t# greeting",
+         {64790, 64792, 884, 1301, 9427, 13, 296, 4466, 2029, 15616, 30914, 993, 3387, 12, 31010, 30174}}};
+
+    for (const auto &c : cases) {
+        // encode
+        std::vector<int> input_ids = pipeline.tokenizer->encode(c.prompt);
+        EXPECT_TRUE(equal(input_ids, c.input_ids));
+        // decode
+        std::string output = pipeline.tokenizer->decode(c.input_ids);
+        EXPECT_EQ(output, c.prompt);
+    }
+
+    // ===== prompter =====
     {
-        EXPECT_EQ(ChatGLMPipeline::build_prompt({"你好"}), "你好");
-        EXPECT_EQ(ChatGLMPipeline::build_prompt({"你好",
-                                                 "你好👋！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。",
-                                                 "晚上睡不着应该怎么办"}),
-                  "[Round 0]\n问：你好\n答：你好👋！我是人工智能助手 "
-                  "ChatGLM-6B，很高兴见到你，欢迎问我任何问题。\n[Round 1]\n问：晚上睡不着应该怎么办\n答：");
+        EXPECT_EQ(ChatGLM2Tokenizer::build_prompt({"你好"}), "[Round 1]\n\n问：你好\n\n答：");
+        EXPECT_EQ(ChatGLM2Tokenizer::build_prompt(
+                      {"你好", "你好👋！我是人工智能助手 ChatGLM2-6B，很高兴见到你，欢迎问我任何问题。",
+                       "晚上睡不着应该怎么办"}),
+                  "[Round 1]\n\n问：你好\n\n答：你好👋！我是人工智能助手 "
+                  "ChatGLM2-6B，很高兴见到你，欢迎问我任何问题。\n\n[Round 2]\n\n问：晚上睡不着应该怎么办\n\n答：");
     }
 
+    // ===== generation =====
+
     // memory test
     {
         GenerationConfig gen_config;
@@ -819,7 +883,7 @@ TEST(ChatGLM, chat) {
         gen_config.do_sample = false;
         std::vector<std::string> history{"你好"};
         std::string output = pipeline.chat(history, gen_config);
-        EXPECT_EQ(output, "你好👋！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。");
+        EXPECT_EQ(output, "你好👋！我是人工智能助手 ChatGLM2-6B，很高兴见到你，欢迎问我任何问题。");
     }
 }
 
@@ -851,15 +915,14 @@ struct PerfStreamer : public BaseStreamer {
     int64_t ms_per_new_token() const { return ms_new_tokens() / num_new_tokens; }
 };
 
-TEST(ChatGLM, benchmark) {
-    fs::path model_path = fs::path(__FILE__).parent_path() / "chatglm-ggml.bin";
+static void run_benchmark(const fs::path &model_path) {
     if (!fs::exists(model_path)) {
-        GTEST_SKIP() << "Skipping ChatGLM benchmark (ggml model not found)";
+        GTEST_SKIP() << "Skipping benchmark test (model " << model_path << " not found)";
     }
 
     ggml_time_init();
     int64_t start_ms = ggml_time_ms();
-    ChatGLMPipeline pipeline(model_path);
+    Pipeline pipeline(model_path);
     int64_t load_model_ms = ggml_time_ms() - start_ms;
 
     start_ms = ggml_time_ms();
@@ -878,7 +941,8 @@ TEST(ChatGLM, benchmark) {
     pipeline.chat(history, gen_config, &streamer);
     int64_t gen_s = (ggml_time_ms() - start_ms) / 1000.f;
 
-    std::cout << "======================== benchmark ========================\n"
+    std::cout << "======== benchmark results for " << model_path.filename() << " ========\n"
+              << "Using #threads: " << gen_config.num_threads << "\n"
               << "Model loaded within: " << load_model_ms << " ms\n"
               << "Generation finished within: " << gen_s << " s\n"
               << "Per ctx token latency: " << streamer.ms_per_ctx_token() << " ms/token (" << streamer.ms_ctx_tokens()
@@ -888,4 +952,14 @@ TEST(ChatGLM, benchmark) {
               << "===========================================================\n";
 }
 
+TEST(Benchmark, ChatGLM) {
+    fs::path model_path = fs::path(__FILE__).parent_path() / "chatglm-ggml.bin";
+    run_benchmark(model_path);
+}
+
+TEST(Benchmark, ChatGLM2) {
+    fs::path model_path = fs::path(__FILE__).parent_path() / "chatglm2-ggml.bin";
+    run_benchmark(model_path);
+}
+
 } // namespace chatglm
diff --git a/test_chatglm_cpp.py b/test_chatglm_cpp.py
index 6c923f9..9a56166 100644
--- a/test_chatglm_cpp.py
+++ b/test_chatglm_cpp.py
@@ -3,16 +3,33 @@
 import chatglm_cpp
 import pytest
 
-MODEL_PATH = Path(__file__).resolve().parent / "chatglm-ggml.bin"
+CHATGLM_MODEL_PATH = Path(__file__).resolve().parent / "chatglm-ggml.bin"
+CHATGLM2_MODEL_PATH = Path(__file__).resolve().parent / "chatglm2-ggml.bin"
 
 
-@pytest.mark.skipif(not MODEL_PATH.is_file(), reason="model file not found")
+@pytest.mark.skipif(not CHATGLM_MODEL_PATH.exists(), reason="model file not found")
 def test_chatglm_pipeline():
-    pipeline = chatglm_cpp.ChatGLMPipeline(MODEL_PATH)
     history = ["你好"]
+    target = "你好👋！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。"
+
+    pipeline = chatglm_cpp.Pipeline(CHATGLM_MODEL_PATH)
+    output = pipeline.chat(history, do_sample=False)
+    assert output == target
+
+    stream_output = pipeline.stream_chat(history, do_sample=False)
+    stream_output = "".join(stream_output)
+    assert stream_output == target
+
+
+@pytest.mark.skipif(not CHATGLM2_MODEL_PATH.exists(), reason="model file not found")
+def test_chatglm2_pipeline():
+    history = ["你好"]
+    target = "你好👋！我是人工智能助手 ChatGLM2-6B，很高兴见到你，欢迎问我任何问题。"
+
+    pipeline = chatglm_cpp.Pipeline(CHATGLM2_MODEL_PATH)
     output = pipeline.chat(history, do_sample=False)
-    assert output == "你好👋！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。"
+    assert output == target
 
     stream_output = pipeline.stream_chat(history, do_sample=False)
     stream_output = "".join(stream_output)
-    assert stream_output == "你好👋！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。"
+    assert stream_output == target
diff --git a/test_convert.py b/test_convert.py
index a50f5fc..2626b77 100644
--- a/test_convert.py
+++ b/test_convert.py
@@ -1,71 +1,319 @@
-from convert import quantize_q8_0, quantize_q4_0
+import sys
+from pathlib import Path
+
 import torch
+import torch.nn.functional as F
+from convert import quantize_q4_0, quantize_q8_0
 
 # generated by:
 #   torch.manual_seed(0)
 #   weight = torch.randn(2, 128)
-weight = torch.tensor(
-        [[-1.1258e+00, -1.1524e+00, -2.5058e-01, -4.3388e-01,  8.4871e-01,
-          6.9201e-01, -3.1601e-01, -2.1152e+00,  3.2227e-01, -1.2633e+00,
-          3.4998e-01,  3.0813e-01,  1.1984e-01,  1.2377e+00,  1.1168e+00,
-         -2.4728e-01, -1.3527e+00, -1.6959e+00,  5.6665e-01,  7.9351e-01,
-          5.9884e-01, -1.5551e+00, -3.4136e-01,  1.8530e+00,  7.5019e-01,
-         -5.8550e-01, -1.7340e-01,  1.8348e-01,  1.3894e+00,  1.5863e+00,
-          9.4630e-01, -8.4368e-01, -6.1358e-01,  3.1593e-02, -4.9268e-01,
-          2.4841e-01,  4.3970e-01,  1.1241e-01,  6.4079e-01,  4.4116e-01,
-         -1.0231e-01,  7.9244e-01, -2.8967e-01,  5.2507e-02,  5.2286e-01,
-          2.3022e+00, -1.4689e+00, -1.5867e+00, -6.7309e-01,  8.7283e-01,
-          1.0554e+00,  1.7784e-01, -2.3034e-01, -3.9175e-01,  5.4329e-01,
-         -3.9516e-01, -4.4622e-01,  7.4402e-01,  1.5210e+00,  3.4105e+00,
-         -1.5312e+00, -1.2341e+00,  1.8197e+00, -5.5153e-01, -5.6925e-01,
-          9.1997e-01,  1.1108e+00,  1.2899e+00, -1.4782e+00,  2.5672e+00,
-         -4.7312e-01,  3.3555e-01, -1.6293e+00, -5.4974e-01, -4.7983e-01,
-         -4.9968e-01, -1.0670e+00,  1.1149e+00, -1.4067e-01,  8.0575e-01,
-         -9.3348e-02,  6.8705e-01, -8.3832e-01,  8.9182e-04,  8.4189e-01,
-         -4.0003e-01,  1.0395e+00,  3.5815e-01, -2.4600e-01,  2.3025e+00,
-         -1.8817e+00, -4.9727e-02, -1.0450e+00, -9.5650e-01,  3.3532e-02,
-          7.1009e-01,  1.6459e+00, -1.3602e+00,  3.4457e-01,  5.1987e-01,
-         -2.6133e+00, -1.6965e+00, -2.2824e-01,  2.7995e-01,  2.4693e-01,
-          7.6887e-02,  3.3801e-01,  4.5440e-01,  4.5694e-01, -8.6537e-01,
-          7.8131e-01, -9.2679e-01, -2.1883e-01, -2.4351e+00, -7.2915e-02,
-         -3.3986e-02,  9.6252e-01,  3.4917e-01, -9.2146e-01, -5.6195e-02,
-         -6.2270e-01, -4.6372e-01,  1.9218e+00, -4.0255e-01,  1.2390e-01,
-          1.1648e+00,  9.2337e-01,  1.3873e+00],
-        [-8.8338e-01, -4.1891e-01, -8.0483e-01,  5.6561e-01,  6.1036e-01,
-          4.6688e-01,  1.9507e+00, -1.0631e+00, -7.7326e-02,  1.1640e-01,
-         -5.9399e-01, -1.2439e+00, -1.0209e-01, -1.0335e+00, -3.1264e-01,
-          2.4579e-01, -2.5964e-01,  1.1834e-01,  2.4396e-01,  1.1646e+00,
-          2.8858e-01,  3.8660e-01, -2.0106e-01, -1.1793e-01,  1.9220e-01,
-         -7.7216e-01, -1.9003e+00,  1.3068e-01, -7.0429e-01,  3.1472e-01,
-          1.5739e-01,  3.8536e-01,  9.6715e-01, -9.9108e-01,  3.0161e-01,
-         -1.0732e-01,  9.9846e-01, -4.9871e-01,  7.6111e-01,  6.1830e-01,
-          3.1405e-01,  2.1333e-01, -1.2005e-01,  3.6046e-01, -3.1403e-01,
-         -1.0787e+00,  2.4081e-01, -1.3962e+00, -6.6144e-02, -3.5836e-01,
-         -1.5616e+00, -3.5464e-01,  1.0811e+00,  1.3148e-01,  1.5735e+00,
-          7.8143e-01, -1.0787e+00, -7.2091e-01,  1.4708e+00,  2.7564e-01,
-          6.6678e-01, -9.9439e-01, -1.1894e+00, -1.1959e+00, -5.5963e-01,
-          5.3347e-01,  4.0689e-01,  3.9459e-01,  1.7151e-01,  8.7604e-01,
-         -2.8709e-01,  1.0216e+00, -7.4395e-02, -1.0922e+00,  3.9203e-01,
-          5.9453e-01,  6.6227e-01, -1.2063e+00,  6.0744e-01, -5.4716e-01,
-          1.1711e+00,  9.7496e-02,  9.6337e-01,  8.4032e-01, -1.2537e+00,
-          9.8684e-01, -4.9466e-01, -1.2830e+00,  9.5522e-01,  1.2836e+00,
-         -6.6586e-01,  5.6513e-01,  2.8770e-01, -3.3375e-02, -1.0619e+00,
-         -1.1443e-01, -3.4334e-01,  1.5713e+00,  1.9161e-01,  3.7994e-01,
-         -1.4476e-01,  6.3762e-01, -2.8129e-01, -1.3299e+00, -1.4201e-01,
-         -5.3415e-01, -5.2338e-01,  8.6150e-01, -8.8696e-01,  8.3877e-01,
-          1.1529e+00, -1.7611e+00, -1.4777e+00, -1.7557e+00,  7.6166e-02,
-         -1.0786e+00,  1.4403e+00, -1.1059e-01,  5.7686e-01, -1.6917e-01,
-         -6.4025e-02,  1.0384e+00,  9.0682e-01, -4.7551e-01, -8.7074e-01,
-          1.4474e-01,  1.9029e+00,  3.9040e-01]]
-)
+
+# fmt: off
+weight = torch.tensor([[
+    -1.1258e+00, -1.1524e+00, -2.5058e-01, -4.3388e-01,  8.4871e-01,  6.9201e-01, -3.1601e-01, -2.1152e+00,
+     3.2227e-01, -1.2633e+00,  3.4998e-01,  3.0813e-01,  1.1984e-01,  1.2377e+00,  1.1168e+00, -2.4728e-01,
+    -1.3527e+00, -1.6959e+00,  5.6665e-01,  7.9351e-01,  5.9884e-01, -1.5551e+00, -3.4136e-01,  1.8530e+00,
+     7.5019e-01, -5.8550e-01, -1.7340e-01,  1.8348e-01,  1.3894e+00,  1.5863e+00,  9.4630e-01, -8.4368e-01,
+    -6.1358e-01,  3.1593e-02, -4.9268e-01,  2.4841e-01,  4.3970e-01,  1.1241e-01,  6.4079e-01,  4.4116e-01,
+    -1.0231e-01,  7.9244e-01, -2.8967e-01,  5.2507e-02,  5.2286e-01,  2.3022e+00, -1.4689e+00, -1.5867e+00,
+    -6.7309e-01,  8.7283e-01,  1.0554e+00,  1.7784e-01, -2.3034e-01, -3.9175e-01,  5.4329e-01, -3.9516e-01,
+    -4.4622e-01,  7.4402e-01,  1.5210e+00,  3.4105e+00, -1.5312e+00, -1.2341e+00,  1.8197e+00, -5.5153e-01,
+    -5.6925e-01,  9.1997e-01,  1.1108e+00,  1.2899e+00, -1.4782e+00,  2.5672e+00, -4.7312e-01,  3.3555e-01,
+    -1.6293e+00, -5.4974e-01, -4.7983e-01, -4.9968e-01, -1.0670e+00,  1.1149e+00, -1.4067e-01,  8.0575e-01,
+    -9.3348e-02,  6.8705e-01, -8.3832e-01,  8.9182e-04,  8.4189e-01, -4.0003e-01,  1.0395e+00,  3.5815e-01,
+    -2.4600e-01,  2.3025e+00, -1.8817e+00, -4.9727e-02, -1.0450e+00, -9.5650e-01,  3.3532e-02,  7.1009e-01,
+     1.6459e+00, -1.3602e+00,  3.4457e-01,  5.1987e-01, -2.6133e+00, -1.6965e+00, -2.2824e-01,  2.7995e-01,
+     2.4693e-01,  7.6887e-02,  3.3801e-01,  4.5440e-01,  4.5694e-01, -8.6537e-01,  7.8131e-01, -9.2679e-01,
+    -2.1883e-01, -2.4351e+00, -7.2915e-02, -3.3986e-02,  9.6252e-01,  3.4917e-01, -9.2146e-01, -5.6195e-02,
+    -6.2270e-01, -4.6372e-01,  1.9218e+00, -4.0255e-01,  1.2390e-01,  1.1648e+00,  9.2337e-01,  1.3873e+00],
+   [-8.8338e-01, -4.1891e-01, -8.0483e-01,  5.6561e-01,  6.1036e-01,  4.6688e-01,  1.9507e+00, -1.0631e+00,
+    -7.7326e-02,  1.1640e-01, -5.9399e-01, -1.2439e+00, -1.0209e-01, -1.0335e+00, -3.1264e-01,  2.4579e-01,
+    -2.5964e-01,  1.1834e-01,  2.4396e-01,  1.1646e+00,  2.8858e-01,  3.8660e-01, -2.0106e-01, -1.1793e-01,
+     1.9220e-01, -7.7216e-01, -1.9003e+00,  1.3068e-01, -7.0429e-01,  3.1472e-01,  1.5739e-01,  3.8536e-01,
+     9.6715e-01, -9.9108e-01,  3.0161e-01, -1.0732e-01,  9.9846e-01, -4.9871e-01,  7.6111e-01,  6.1830e-01,
+     3.1405e-01,  2.1333e-01, -1.2005e-01,  3.6046e-01, -3.1403e-01, -1.0787e+00,  2.4081e-01, -1.3962e+00,
+    -6.6144e-02, -3.5836e-01, -1.5616e+00, -3.5464e-01,  1.0811e+00,  1.3148e-01,  1.5735e+00,  7.8143e-01,
+    -1.0787e+00, -7.2091e-01,  1.4708e+00,  2.7564e-01,  6.6678e-01, -9.9439e-01, -1.1894e+00, -1.1959e+00,
+    -5.5963e-01,  5.3347e-01,  4.0689e-01,  3.9459e-01,  1.7151e-01,  8.7604e-01, -2.8709e-01,  1.0216e+00,
+    -7.4395e-02, -1.0922e+00,  3.9203e-01,  5.9453e-01,  6.6227e-01, -1.2063e+00,  6.0744e-01, -5.4716e-01,
+     1.1711e+00,  9.7496e-02,  9.6337e-01,  8.4032e-01, -1.2537e+00,  9.8684e-01, -4.9466e-01, -1.2830e+00,
+     9.5522e-01,  1.2836e+00, -6.6586e-01,  5.6513e-01,  2.8770e-01, -3.3375e-02, -1.0619e+00, -1.1443e-01,
+    -3.4334e-01,  1.5713e+00,  1.9161e-01,  3.7994e-01, -1.4476e-01,  6.3762e-01, -2.8129e-01, -1.3299e+00,
+    -1.4201e-01, -5.3415e-01, -5.2338e-01,  8.6150e-01, -8.8696e-01,  8.3877e-01,  1.1529e+00, -1.7611e+00,
+    -1.4777e+00, -1.7557e+00,  7.6166e-02, -1.0786e+00,  1.4403e+00, -1.1059e-01,  5.7686e-01, -1.6917e-01,
+    -6.4025e-02,  1.0384e+00,  9.0682e-01, -4.7551e-01, -8.7074e-01,  1.4474e-01,  1.9029e+00,  3.9040e-01]])
+# fmt: on
+
 
 def test_quantize_q8_0():
     q8_tensor = quantize_q8_0(weight).int()
-    ggml_q8_tensor = torch.tensor([68, 36, -68, -69, -15, -26, 51, 42, -19, -127, 19, -76, 21, 19, 7, 74, 67, -15, -81, -102, 34, 48, 36, -93, -20, 111, 45, -35, -10, 11, 83, 95, 57, -51, -32, 38, -23, 1, -18, 9, 16, 4, 24, 16, -4, 30, -11, 2, 19, 86, -55, -59, -25, 33, 39, 7, -9, -15, 20, -15, -17, 28, 57, 127, -57, -46, 68, -21, 45, 37, -28, 46, 55, 64, -73, 127, -23, 17, -81, -27, -24, -25, -53, 55, -7, 40, -5, 34, -41, 0, 42, -20, 51, 18, -12, 114, -93, -2, -52, -47, 2, 35, 69, 37, 80, -66, 17, 25, -127, -82, -11, 14, 12, 4, 16, 22, 22, -42, 38, -45, -11, -118, -4, -2, 47, 17, -45, -3, -30, -23, 93, -20, 6, 57, 45, 67, -35, 35, -58, -27, -52, 37, 40, 30, 127, -69, -5, 8, -39, -81, -7, -67, -20, 16, -17, 8, 16, 76, 19, 25, -13, -8, 13, -50, -124, 9, -46, 20, 10, 25, 88, 34, 78, -80, 24, -9, 81, -40, 61, 50, 25, 17, -10, 29, -25, -87, 19, -113, -5, -29, -126, -29, 87, 11, 127, 63, -87, -58, 119, 22, 54, -80, -96, -97, 45, 33, -55, 53, 40, 39, 17, 87, -28, 101, -7, -108, 39, 59, 66, -119, 60, -54, 116, 10, 95, 83, -124, 98, -49, -127, 95, 127, -66, 56, 28, -3, -105, -11, -84, 35, -23, 105, 13, 25, -10, 43, -19, -89, -9, -36, -35, 57, -59, 56, 77, -118, -99, -117, 5, -72, 96, -7, 38, -11, -4, 69, 61, -32, -58, 10, 127, 26]).view(q8_tensor.shape)
+    # fmt: off
+    ggml_q8_tensor = torch.tensor([
+          68,   36,  -68,  -69,  -15,  -26,   51,   42,  -19, -127,   19,  -76,
+          21,   19,    7,   74,   67,  -15,  -81, -102,   34,   48,   36,  -93,
+         -20,  111,   45,  -35,  -10,   11,   83,   95,   57,  -51,  -32,   38,
+         -23,    1,  -18,    9,   16,    4,   24,   16,   -4,   30,  -11,    2,
+          19,   86,  -55,  -59,  -25,   33,   39,    7,   -9,  -15,   20,  -15,
+         -17,   28,   57,  127,  -57,  -46,   68,  -21,   45,   37,  -28,   46,
+          55,   64,  -73,  127,  -23,   17,  -81,  -27,  -24,  -25,  -53,   55,
+          -7,   40,   -5,   34,  -41,    0,   42,  -20,   51,   18,  -12,  114,
+         -93,   -2,  -52,  -47,    2,   35,   69,   37,   80,  -66,   17,   25,
+        -127,  -82,  -11,   14,   12,    4,   16,   22,   22,  -42,   38,  -45,
+         -11, -118,   -4,   -2,   47,   17,  -45,   -3,  -30,  -23,   93,  -20,
+           6,   57,   45,   67,  -35,   35,  -58,  -27,  -52,   37,   40,   30,
+         127,  -69,   -5,    8,  -39,  -81,   -7,  -67,  -20,   16,  -17,    8,
+          16,   76,   19,   25,  -13,   -8,   13,  -50, -124,    9,  -46,   20,
+          10,   25,   88,   34,   78,  -80,   24,   -9,   81,  -40,   61,   50,
+          25,   17,  -10,   29,  -25,  -87,   19, -113,   -5,  -29, -126,  -29,
+          87,   11,  127,   63,  -87,  -58,  119,   22,   54,  -80,  -96,  -97,
+          45,   33,  -55,   53,   40,   39,   17,   87,  -28,  101,   -7, -108,
+          39,   59,   66, -119,   60,  -54,  116,   10,   95,   83, -124,   98,
+         -49, -127,   95,  127,  -66,   56,   28,   -3, -105,  -11,  -84,   35,
+         -23,  105,   13,   25,  -10,   43,  -19,  -89,   -9,  -36,  -35,   57,
+         -59,   56,   77, -118,  -99, -117,    5,  -72,   96,   -7,   38,  -11,
+          -4,   69,   61,  -32,  -58,   10,  127,   26]).view(q8_tensor.shape)
+    # fmt: on
     assert (q8_tensor == ggml_q8_tensor).all()
 
 
 def test_quantize_q4_0():
     q4_tensor = quantize_q4_0(weight).int()
-    ggml_q4_tensor = torch.tensor([59, 52, 52, 36, -89, -74, -85, 43, 119, -16, -71, 99, 121, -103, -40, -19, -52, 87, -46, -74, -87, 104, 105, -121, -105, -104, 118, -105, -104, 102, 73, 8, -57, -77, 75, -100, 34, -75, -118, 101, -75, -124, 93, -112, 89, 119, -99, 26, -23, -118, -69, -75, -120, 101, 58, 53, 125, 20, -119, -118, -80, -109, 87, -119, 105, 120, -23, 121, -119, -59, -70, -59, -50, -77, -100, -118, 123, 54, 117, 102, -112, -116, 120, -72, -6, 125, -72, 124, 121, 103, 75, -78, -125, -83, -10, -87, 51, 123, 4, 69, -42, -57, 25, 118, 90, -35, -25, -17, 34, -79, 27, 117, 37, 54, -9, 35, -70, -14, 40, 15, -58, 68, 100, -113, -12, -101, -99, -77, -23, -15, -121, -42, 41, -123, 105, -98, -119, 74, 74, -92, -52, 116, 3, 111]).view(q4_tensor.shape)
+    # fmt: off
+    ggml_q4_tensor = torch.tensor([
+          59,   52,   52,   36,  -89,  -74,  -85,   43,  119,  -16,  -71,   99,
+         121, -103,  -40,  -19,  -52,   87,  -46,  -74,  -87,  104,  105, -121,
+        -105, -104,  118, -105, -104,  102,   73,    8,  -57,  -77,   75, -100,
+          34,  -75, -118,  101,  -75, -124,   93, -112,   89,  119,  -99,   26,
+         -23, -118,  -69,  -75, -120,  101,   58,   53,  125,   20, -119, -118,
+         -80, -109,   87, -119,  105,  120,  -23,  121, -119,  -59,  -70,  -59,
+         -50,  -77, -100, -118,  123,   54,  117,  102, -112, -116,  120,  -72,
+          -6,  125,  -72,  124,  121,  103,   75,  -78, -125,  -83,  -10,  -87,
+          51,  123,    4,   69,  -42,  -57,   25,  118,   90,  -35,  -25,  -17,
+          34,  -79,   27,  117,   37,   54,   -9,   35,  -70,  -14,   40,   15,
+         -58,   68,  100, -113,  -12, -101,  -99,  -77,  -23,  -15, -121,  -42,
+          41, -123,  105,  -98, -119,   74,   74,  -92,  -52,  116,    3,  111]).view(q4_tensor.shape)
+    # fmt: on
     assert (q4_tensor == ggml_q4_tensor).all()
+
+
+CHATGLM_MODEL_PATH = Path(
+    "~/.cache/huggingface/hub/models--THUDM--chatglm-6b/snapshots/a70fe6b0a3cf1675b3aec07e3b7bb7a8ce73c6ae"
+).expanduser()
+
+CHATGLM2_MODEL_PATH = Path(
+    "~/.cache/huggingface/hub/models--THUDM--chatglm2-6b/snapshots/fc442f7e7cf3ac073433cef0f301b4744c25edb6"
+).expanduser()
+
+
+def make_data_embedding():
+    m = torch.nn.Embedding(4, 3)
+    x = torch.tensor([1, 3, 0, 2, 3])
+    y = m(x)
+    print("w", m.weight.data.flatten())
+    print("x", x.flatten())
+    print("y", y.flatten())
+
+
+def make_data_linear():
+    w = torch.randn(8, 3)
+    b = torch.randn(8)
+    x = torch.randn(2, 3)
+    y = F.linear(x, w, b)
+
+    print("w", w.flatten())
+    print("b", b.flatten())
+    print("x", x.flatten())
+    print("y", y.flatten())
+
+
+def make_data_layernorm():
+    w = torch.randn(9)
+    b = torch.randn(9)
+    x = torch.randn(2, 9)
+    y = F.layer_norm(x, [9], w, b)
+    print("w", w.flatten())
+    print("b", b.flatten())
+    print("x", x.flatten())
+    print("y", y.flatten())
+
+
+def make_data_rms_norm():
+    from modeling_chatglm import RMSNorm
+
+    m = RMSNorm(7, eps=1e-6).eval()
+    m.weight.data.uniform_()
+
+    x = torch.randn(2, 7)
+    with torch.no_grad():
+        y = m(x)
+    print("weight", m.weight.data.flatten())
+    print("x", x.flatten())
+    print("y", y.flatten())
+
+
+def make_data_glm_self_attention():
+    from modeling_chatglm import SelfAttention
+
+    m = SelfAttention(16, 2, layer_id=3, empty_init=False).float().eval()
+    x = torch.randn(4, 1, 16)  # [seqlen, bs, hidden]
+    position_ids = torch.tensor([[[0, 1, 2, 2], [0, 0, 0, 1]]])
+    attention_mask = torch.tensor(
+        [
+            [
+                [
+                    [0, 0, 0, 1],
+                    [0, 0, 0, 1],
+                    [0, 0, 0, 1],
+                    [0, 0, 0, 0],
+                ]
+            ]
+        ]
+    ).bool()
+    y, layer_past = m(
+        x,
+        position_ids=position_ids,
+        attention_mask=attention_mask,
+        layer_id=m.layer_id,
+        use_cache=True,
+    )
+    print("x", x.flatten())
+    print("query_key_value.weight", m.query_key_value.weight.flatten())
+    print("query_key_value.bias", m.query_key_value.bias.flatten())
+    print("dense.weight", m.dense.weight.flatten())
+    print("dense.bias", m.dense.bias.flatten())
+    print("y", y.flatten())
+
+    x2 = torch.randn(1, 1, 16)
+    position_ids = torch.tensor([[[2], [2]]])
+    attention_mask = torch.zeros(1, 1, dtype=torch.bool)
+    y2, layer_past = m(
+        x2,
+        position_ids=position_ids,
+        attention_mask=attention_mask,
+        layer_id=m.layer_id,
+        layer_past=layer_past,
+        use_cache=True,
+    )
+    print("x2", x2.flatten())
+    print("y2", y2.flatten())
+
+    x3 = torch.randn(1, 1, 16)
+    position_ids = torch.tensor([[[2], [3]]])
+    attention_mask = torch.zeros(1, 1, dtype=torch.bool)
+    y3, layer_past = m(
+        x3,
+        position_ids=position_ids,
+        attention_mask=attention_mask,
+        layer_id=m.layer_id,
+        layer_past=layer_past,
+        use_cache=True,
+    )
+    print("x3", x3.flatten())
+    print("y3", y3.flatten())
+
+
+def make_data_glm_block():
+    from modeling_chatglm import GLMBlock
+
+    m = (
+        GLMBlock(hidden_size=8, num_attention_heads=2, layernorm_epsilon=1e-5, layer_id=3, empty_init=False)
+        .float()
+        .eval()
+    )
+    x = torch.randn(4, 1, 8)  # [seqlen, bs, hidden]
+    position_ids = torch.tensor([[[0, 1, 2, 2], [0, 0, 0, 1]]])
+    attention_mask = torch.tensor(
+        [
+            [
+                [
+                    [0, 0, 0, 1],
+                    [0, 0, 0, 1],
+                    [0, 0, 0, 1],
+                    [0, 0, 0, 0],
+                ]
+            ]
+        ]
+    ).bool()
+    (y,) = m(x, position_ids=position_ids, attention_mask=attention_mask, layer_id=m.layer_id)
+    print("x", x.flatten())
+    print("input_layernorm.weight", m.input_layernorm.weight.data.flatten())
+    print("input_layernorm.bias", m.input_layernorm.bias.data.flatten())
+    print("query_key_value.weight", m.attention.query_key_value.weight.data.flatten())
+    print("query_key_value.bias", m.attention.query_key_value.bias.data.flatten())
+    print("dense.weight", m.attention.dense.weight.data.flatten())
+    print("dense.bias", m.attention.dense.bias.data.flatten())
+    print("post_attention_layernorm.weight", m.post_attention_layernorm.weight.data.flatten())
+    print("post_attention_layernorm.bias", m.post_attention_layernorm.bias.data.flatten())
+    print("dense_h_to_4h.weight", m.mlp.dense_h_to_4h.weight.data.flatten())
+    print("dense_h_to_4h.bias", m.mlp.dense_h_to_4h.bias.data.flatten())
+    print("dense_4h_to_h.weight", m.mlp.dense_4h_to_h.weight.data.flatten())
+    print("dense_4h_to_h.bias", m.mlp.dense_4h_to_h.bias.data.flatten())
+    print("y", y.flatten())
+
+
+def make_data_glm2_block():
+    from modeling_chatglm import GLMBlock, RotaryEmbedding
+    from transformers import AutoConfig
+
+    config = AutoConfig.from_pretrained(CHATGLM2_MODEL_PATH, trust_remote_code=True)
+    config.layernorm_epsilon = 1e-6
+    config.hidden_size = 16
+    config.num_attention_heads = 4
+    config.multi_query_group_num = 2
+    config.ffn_hidden_size = 6
+    config.kv_channels = config.hidden_size // config.num_attention_heads
+    config.torch_dtype = torch.float32
+    m = GLMBlock(config, layer_number=3).eval()
+    m.input_layernorm.weight.data.uniform_()
+    m.post_attention_layernorm.weight.data.uniform_()
+
+    seq_length = 3
+    rotary_dim = config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
+    rotary_pos_emb_module = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope)
+    rotary_pos_emb = rotary_pos_emb_module(8)[None, :seq_length].transpose(0, 1).contiguous()
+
+    # self attention
+    x = torch.randn(seq_length, 1, config.hidden_size)
+    with torch.no_grad():
+        y, kv_cache = m(x, attention_mask=None, rotary_pos_emb=rotary_pos_emb)
+
+    print(m)
+
+    print("input_layernorm.weight", m.input_layernorm.weight.data.flatten())
+    print("attn.qkv.weight", m.self_attention.query_key_value.weight.data.flatten())
+    print("attn.qkv.bias", m.self_attention.query_key_value.bias.data.flatten())
+    print("attn.dense.weight", m.self_attention.dense.weight.data.flatten())
+    print("post_attention_layernorm.weight", m.post_attention_layernorm.weight.data.flatten())
+    print("mlp.dense_h_to_4h.weight", m.mlp.dense_h_to_4h.weight.data.flatten())
+    print("mlp.dense_4h_to_h.weight", m.mlp.dense_4h_to_h.weight.data.flatten())
+
+    print("x", x.flatten())
+    print("y", y.flatten())
+
+    # cross attention
+    position_ids = torch.tensor([[seq_length]])
+    rotary_pos_emb = rotary_pos_emb_module(8)[position_ids].transpose(0, 1).contiguous()
+    x = torch.randn(1, 1, config.hidden_size)
+    with torch.no_grad():
+        y, kv_cache = m(x, attention_mask=None, rotary_pos_emb=rotary_pos_emb, kv_cache=kv_cache)
+    print("x2", x.flatten())
+    print("y2", y.flatten())
+
+    position_ids = torch.tensor([[seq_length + 1]])
+    rotary_pos_emb = rotary_pos_emb_module(8)[position_ids].transpose(0, 1).contiguous()
+    x = torch.randn(1, 1, config.hidden_size)
+    with torch.no_grad():
+        y, kv_cache = m(x, attention_mask=None, rotary_pos_emb=rotary_pos_emb, kv_cache=kv_cache)
+    print("x3", x.flatten())
+    print("y3", y.flatten())
+
+
+def main():
+    sys.path.append(str(CHATGLM_MODEL_PATH))
+    torch.manual_seed(0)
+    make_data_glm2_block()
+
+
+if __name__ == "__main__":
+    main()