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

Issue:
GPU推理速度相比于llama.cpp慢了一倍
我的机器配置是AMD Ryzen 5950x, NVidia RTX A6000, CUDA 11.7
我目前有2套测试配置，都是截止到7月6日的最新代码，都使用同样的参数 t=6, l/n=128, prompt="how to build a house in 10 steps“
C1:  chatglm2-6B,  使用 chatglm.cpp
C2:  vicuna_7b_v1.3, 使用 llama.cpp
**在CPU下：**
 **FP16** 

- C1: 2.5 t/s
- C2: 2.3 t/s

**Q4_0** 

- C1: 8.5 t/s
- C2: 7.5 t/s
所以在CPU下C1是稳定比C2快10%的，这应该是符合6BVS7B的

**在GPU下:**
**FP16**
- C1: 34.6 t/s
- C2: 44.3 t/s

**Q4_0** 

- C1: 49.8 t/s
- C2: 108.1 t/s

    所以在GPU下反而C1比C2慢很多，low bit差别更加明显。用GPU的时候我是设置--threads 1，跟llama.cpp一致。
    

chatglm.cpp里没有提供-ngl所以我不确定是不是一些层没有完全offload到GPU导致的，比较难DEBUG

UPDATE:
附带一张我最近和fastllm的性能测试， 可以看到对于fastllm, chatglm2 6B在GPU上是比vicua 7B快40%的，证明这个架构没什么问题。但在chatglm.cpp当前反而是比vicuna 7B慢了50%
![image](https://github.com/li-plus/chatglm.cpp/assets/15835199/322c98de-1e47-404a-a867-80ee83373c68)





Discussion / hints:

Fix patch:
diff --git a/.github/workflows/cmake.yml b/.github/workflows/cmake.yml
index 46f511b..c34e03c 100644
--- a/.github/workflows/cmake.yml
+++ b/.github/workflows/cmake.yml
@@ -38,5 +38,5 @@ jobs:
       working-directory: ${{github.workspace}}/build
       # Execute tests defined by the CMake configuration.
       # See https://cmake.org/cmake/help/latest/manual/ctest.1.html for more detail
-      run: ctest -C ${{env.BUILD_TYPE}}
+      run: ctest -C ${{env.BUILD_TYPE}} -E Benchmark
 
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 1e7a239..28c255c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -14,6 +14,7 @@ if (NOT CMAKE_BUILD_TYPE)
 endif ()
 
 # third-party libraries
+add_compile_definitions(GGML_CUDA_MMV_Y=2)  # for large vocab
 include_directories(third_party/ggml/include/ggml third_party/ggml/src)
 add_subdirectory(third_party/ggml)
 
@@ -38,7 +39,7 @@ set(CPP_SOURCES
 set_source_files_properties(${CPP_SOURCES} PROPERTIES COMPILE_FLAGS "-pedantic-errors")
 
 add_library(chatglm STATIC chatglm.cpp)
-target_link_libraries(chatglm PRIVATE ggml sentencepiece-static)
+target_link_libraries(chatglm PUBLIC ggml sentencepiece-static)
 
 add_executable(main main.cpp)
 target_link_libraries(main PRIVATE chatglm)
@@ -58,11 +59,12 @@ if (CHATGLM_ENABLE_TESTING)
     # For Windows: Prevent overriding the parent project's compiler/linker settings
     set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)
     FetchContent_MakeAvailable(googletest)
+    include(GoogleTest)
 
     # Now simply link against gtest or gtest_main as needed. Eg
     add_executable(chatglm_test chatglm_test.cpp)
     target_link_libraries(chatglm_test PRIVATE chatglm gtest_main)
-    add_test(NAME chatglm_test COMMAND chatglm_test)
+    gtest_discover_tests(chatglm_test)
 endif ()
 
 option(CHATGLM_ENABLE_PYBIND, "chatglm: enable python binding" OFF)
diff --git a/README.md b/README.md
index a7f49ae..0b41de2 100644
--- a/README.md
+++ b/README.md
@@ -141,23 +141,27 @@ docker run -it --rm -v [model path]:/opt/ chulinx/chatglm /chatglm -m /opt/chatg
 
 ## Performance
 
-Measured on a Linux server with Intel(R) Xeon(R) Platinum 8260 CPU @ 2.40GHz using 16 threads.
+Environment:
+* CPU performance is measured on a Linux server with Intel(R) Xeon(R) Platinum 8260 CPU @ 2.40GHz using 16 threads.
+* CUDA performance is measured on a V100-SXM2-32GB GPU using 1 thread.
 
 ChatGLM-6B:
 
-|           | Q4_0  | Q4_1  | Q5_0  | Q5_1  | Q8_0  | F16   | F32   |
-|-----------|-------|-------|-------|-------|-------|-------|-------|
-| ms/token  | 74    | 77    | 86    | 89    | 114   | 189   | 357   |
-| file size | 3.3GB | 3.7GB | 4.0GB | 4.4GB | 6.2GB | 12GB  | 23GB  |
-| mem usage | 4.0GB | 4.4GB | 4.7GB | 5.1GB | 6.9GB | 13GB  | 24GB  |
+|                 | Q4_0  | Q4_1  | Q5_0  | Q5_1  | Q8_0  | F16   | F32   |
+|-----------------|-------|-------|-------|-------|-------|-------|-------|
+| ms/token (CPU)  | 74    | 77    | 86    | 89    | 114   | 189   | 357   |
+| ms/token (V100) | 10.0  | 9.8   | 10.7  | 10.6  | 14.6  | 19.8  | 34.2  |
+| file size       | 3.3GB | 3.7GB | 4.0GB | 4.4GB | 6.2GB | 12GB  | 23GB  |
+| mem usage       | 4.0GB | 4.4GB | 4.7GB | 5.1GB | 6.9GB | 13GB  | 24GB  |
 
 ChatGLM2-6B:
 
-|           | Q4_0  | Q4_1  | Q5_0  | Q5_1  | Q8_0  | F16   | F32   |
-|-----------|-------|-------|-------|-------|-------|-------|-------|
-| ms/token  | 64    | 71    | 79    | 83    | 106   | 189   | 372   |
-| file size | 3.3GB | 3.7GB | 4.0GB | 4.4GB | 6.2GB | 12GB  | 24GB  |
-| mem usage | 3.4GB | 3.8GB | 4.1GB | 4.5GB | 6.2GB | 12GB  | 23GB  |
+|                 | Q4_0  | Q4_1  | Q5_0  | Q5_1  | Q8_0  | F16   | F32   |
+|-----------------|-------|-------|-------|-------|-------|-------|-------|
+| ms/token (CPU)  | 64    | 71    | 79    | 83    | 106   | 189   | 372   |
+| ms/token (V100) | 9.7   | 9.4   | 10.3  | 10.2  | 14.0  | 19.1  | 33.0  |
+| file size       | 3.3GB | 3.7GB | 4.0GB | 4.4GB | 6.2GB | 12GB  | 24GB  |
+| mem usage       | 3.4GB | 3.8GB | 4.1GB | 4.5GB | 6.2GB | 12GB  | 23GB  |
 
 ## Development
 
diff --git a/chatglm.cpp b/chatglm.cpp
index d995b03..94ed9bb 100644
--- a/chatglm.cpp
+++ b/chatglm.cpp
@@ -9,6 +9,7 @@
 #include <iomanip>
 #include <iostream>
 #include <locale>
+#include <numeric>
 #include <random>
 #include <regex>
 #include <string>
@@ -101,12 +102,59 @@ std::string to_string(ggml_tensor *tensor, bool with_data) {
         }
         if (tensor->n_dims > 3)
             oss << "]";
+        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();
 }
 
+void tensor_assign_buffers(ggml_tensor *tensor, bool scratch, bool force_inplace) {
+#ifdef GGML_USE_CUBLAS
+    if (scratch) {
+        CHATGLM_CHECK(!force_inplace);
+        ggml_cuda_assign_buffers(tensor);
+    } else {
+        if (force_inplace) {
+            ggml_cuda_assign_buffers_force_inplace(tensor);
+        } else {
+            // BE CAREFUL TO USE THIS!
+            ggml_cuda_assign_buffers_no_scratch(tensor);
+        }
+    }
+#endif
+}
+
+void tensor_to_device(ggml_tensor *tensor) {
+#ifdef GGML_USE_CUBLAS
+    if (tensor->backend == GGML_BACKEND_GPU || tensor->backend == GGML_BACKEND_GPU_SPLIT) {
+        return;
+    }
+    tensor->backend = GGML_BACKEND_GPU;
+    ggml_cuda_transform_tensor(tensor->data, tensor);
+#endif
+}
+
+void tensor_to_cpu(ggml_tensor *tensor) {
+#ifdef GGML_USE_CUBLAS
+    if (tensor->backend == GGML_BACKEND_CPU) {
+        return;
+    }
+    ggml_cuda_free_data(tensor);
+    tensor->backend = GGML_BACKEND_CPU;
+#endif
+}
+
+// for debugging purpose
+static inline ggml_tensor *add_zero(ggml_context *ctx, ggml_tensor *tensor) {
+    ggml_tensor *zeros = ggml_new_tensor(ctx, tensor->type, tensor->n_dims, tensor->ne);
+    ggml_set_f32(zeros, 0);
+    tensor_to_device(zeros);
+    ggml_tensor *out = ggml_add(ctx, tensor, zeros);
+    tensor_assign_buffers(out);
+    return out;
+}
+
 // ===== streamer =====
 
 void StreamerGroup::put(const std::vector<int> &output_ids) {
@@ -300,9 +348,10 @@ ggml_tensor *Embedding::forward(ForwardContext *ctx, ggml_tensor *input) const {
 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]
+    tensor_assign_buffers(output);
     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);
+        output = ggml_add_inplace(ctx->gctx.get(), output, bias);
+        tensor_assign_buffers(output);
     }
     return output;
 }
@@ -310,15 +359,20 @@ ggml_tensor *Linear::forward(ForwardContext *ctx, ggml_tensor *input) const {
 ggml_tensor *LayerNorm::forward(ForwardContext *ctx, ggml_tensor *input) const {
     // input: [seqlen, normalized_shape]
     ggml_tensor *output = ggml_norm_inplace(ctx->gctx.get(), input);
+    tensor_assign_buffers(output);
     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);
+    tensor_assign_buffers(output);
+    output = ggml_add_inplace(ctx->gctx.get(), output, bias);
+    tensor_assign_buffers(output);
     return output;
 }
 
 ggml_tensor *RMSNorm::forward(ForwardContext *ctx, ggml_tensor *input) const {
-    ggml_tensor *output = ggml_rms_norm_inplace(ctx->gctx.get(), input);
+    auto ggml_rms_norm_fn = inplace ? ggml_rms_norm_inplace : ggml_rms_norm;
+    ggml_tensor *output = ggml_rms_norm_fn(ctx->gctx.get(), input);
+    tensor_assign_buffers(output);
     output = ggml_mul_inplace(ctx->gctx.get(), output, weight);
+    tensor_assign_buffers(output);
     return output;
 }
 
@@ -462,7 +516,7 @@ std::string ChatGLMTokenizer::postprocess(const std::string &text) {
 }
 
 // Adapted from https://github.com/ggerganov/llama.cpp/blob/master/examples/common.cpp
-static inline int get_num_physical_cores() {
+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;
 }
@@ -483,18 +537,26 @@ int BaseModelForConditionalGeneration::generate_next_token(const std::vector<int
                                                            int n_ctx) const {
     ForwardContext ctx;
     ctx.gctx = GGMLContext(mem_size_, mem_buffer_.get(), 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 = {0, scratch_size_, scratch_buffer_.get()};
 
+    int n_threads = gen_config.num_threads; // user defined
+    if (n_threads <= 0) {
+        n_threads = ggml_cpu_has_gpublas() ? 1 : get_num_physical_cores(); // default thread num
+    }
+    if (input_ids.size() >= 32 && ggml_cpu_has_blas() && !ggml_cpu_has_gpublas()) {
+        n_threads = 1; // BLAS enabled
+    }
+
     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);
+    lm_logits->backend = GGML_BACKEND_CPU;
 
     ggml_build_forward_expand(&ctx.gf, lm_logits);
-    ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+    // TODO: upgrade to ggml_graph_compute with cplan
+    ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, n_threads);
 
 #ifdef GGML_PERF
     ggml_graph_print(&ctx.gf);
@@ -506,9 +568,8 @@ int BaseModelForConditionalGeneration::generate_next_token(const std::vector<int
     int next_token_id;
     if (gen_config.do_sample) {
         // temperature sampling
-        float inv_temp = 1.f / gen_config.temperature;
-        for (int i = 0; i < vocab_size; i++) {
-            next_token_logits[i] *= inv_temp;
+        if (gen_config.temperature > 0) {
+            sampling_temperature(next_token_logits, next_token_logits + vocab_size, gen_config.temperature);
         }
 
         std::vector<TokenIdScore> token_scores(vocab_size);
@@ -518,24 +579,15 @@ int BaseModelForConditionalGeneration::generate_next_token(const std::vector<int
 
         // top_k sampling
         if (0 < gen_config.top_k && gen_config.top_k < (int)token_scores.size()) {
-            std::nth_element(token_scores.begin(), token_scores.begin() + gen_config.top_k, token_scores.end(),
-                             std::greater<TokenIdScore>());
+            sampling_top_k(token_scores.data(), token_scores.data() + gen_config.top_k,
+                           token_scores.data() + token_scores.size());
             token_scores.resize(gen_config.top_k);
         }
 
         // top_p sampling
         if (0.f < gen_config.top_p && gen_config.top_p < 1.f) {
-            std::sort(token_scores.begin(), token_scores.end(), std::greater<TokenIdScore>()); // hot code!
-            sampling_softmax_inplace(token_scores.data(), token_scores.data() + token_scores.size());
-
-            float cumsum = 0.f;
-            for (size_t i = 0; i < token_scores.size(); i++) {
-                cumsum += token_scores[i].score;
-                if (cumsum >= gen_config.top_p) {
-                    token_scores.resize(i + 1);
-                    break;
-                }
-            }
+            auto pos = sampling_top_p(token_scores.data(), token_scores.data() + token_scores.size(), gen_config.top_p);
+            token_scores.resize(pos - token_scores.data());
         }
 
         // sample next token
@@ -557,6 +609,41 @@ int BaseModelForConditionalGeneration::generate_next_token(const std::vector<int
     return next_token_id;
 }
 
+void BaseModelForConditionalGeneration::sampling_temperature(float *first, float *last, float temp) {
+    float inv_temp = 1.f / temp;
+    for (float *it = first; it != last; it++) {
+        *it *= inv_temp;
+    }
+}
+
+void BaseModelForConditionalGeneration::sampling_top_k(TokenIdScore *first, TokenIdScore *kth, TokenIdScore *last) {
+    std::nth_element(first, kth, last, std::greater<TokenIdScore>());
+}
+
+TokenIdScore *BaseModelForConditionalGeneration::sampling_top_p(TokenIdScore *first, TokenIdScore *last, float top_p) {
+    // fast top_p in expected O(n) time complexity
+    sampling_softmax_inplace(first, last);
+
+    while (first + 1 < last) {
+        float pivot_score = (last - 1)->score; // use mid score?
+        TokenIdScore *mid =
+            std::partition(first, last - 1, [pivot_score](const TokenIdScore &x) { return x.score > pivot_score; });
+        std::swap(*mid, *(last - 1));
+
+        float prefix_sum =
+            std::accumulate(first, mid, 0.f, [](float sum, const TokenIdScore &x) { return sum + x.score; });
+        if (prefix_sum >= top_p) {
+            last = mid;
+        } else if (prefix_sum + mid->score < top_p) {
+            first = mid + 1;
+            top_p -= prefix_sum + mid->score;
+        } else {
+            return mid + 1;
+        }
+    }
+    return last;
+}
+
 void BaseModelForConditionalGeneration::sampling_softmax_inplace(TokenIdScore *first, TokenIdScore *last) {
     float max_score = std::max_element(first, last)->score;
     float sum = 0.f;
@@ -616,10 +703,19 @@ std::vector<int> BaseModelForConditionalGeneration::generate(const std::vector<i
 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);
+    tensor_assign_buffers(output);
     output = dense_4h_to_h.forward(ctx, output);
     return output;
 }
 
+GLMSelfAttention::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 *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];
@@ -630,54 +726,85 @@ ggml_tensor *GLMSelfAttention::forward(ForwardContext *ctx, ggml_tensor *hidden_
 
     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);
+#ifdef GGML_USE_CUBLAS
+    // dst for inplace ops should be contiguous for cuda
+    query_layer = ggml_cont(ctx->gctx.get(), query_layer);
+    tensor_assign_buffers(query_layer);
+#endif
     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]
+    tensor_assign_buffers(query_layer);
+    query_layer = ggml_permute(ctx->gctx.get(), query_layer, 0, 2, 1, 3); // [heads, qlen, head_size]
+    tensor_assign_buffers(query_layer, false, true);
 
     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));
+#ifdef GGML_USE_CUBLAS
+    key_layer = ggml_cont(ctx->gctx.get(), key_layer);
+    tensor_assign_buffers(key_layer);
+#endif
     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]
+    tensor_assign_buffers(key_layer);
+    key_layer = ggml_permute(ctx->gctx.get(), key_layer, 0, 2, 1, 3); // [heads, qlen, head_size]
+    tensor_assign_buffers(key_layer, false, true);
 
     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]
+    tensor_assign_buffers(value_layer, false, true);
 
     // 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]
+    tensor_assign_buffers(k_cache_view);
     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]
+    tensor_assign_buffers(v_cache_view);
     ggml_build_forward_expand(&ctx->gf, ggml_cpy(ctx->gctx.get(), value_layer, v_cache_view));
 
     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]
+    tensor_assign_buffers(key_layer);
     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]
+    tensor_assign_buffers(value_layer);
+
+#ifdef GGML_USE_CUBLAS
+    // make query contiguous to speed up cuda gemm
+    query_layer = ggml_cont(ctx->gctx.get(), query_layer);
+    tensor_assign_buffers(query_layer);
+#endif
 
     ggml_tensor *attn_scores = ggml_mul_mat(ctx->gctx.get(), key_layer, query_layer); // [heads, qlen, klen]
+    tensor_assign_buffers(attn_scores);
     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);
+        tensor_to_device(inf); // TODO: optimize
         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));
+        tensor_assign_buffers(masked_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)));
+    tensor_assign_buffers(attn_scores);
     ggml_tensor *attn_probs = ggml_soft_max_inplace(ctx->gctx.get(), attn_scores); // [heads, qlen, klen]
+    tensor_assign_buffers(attn_probs);
 
     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);
+    tensor_assign_buffers(context_layer);
+    context_layer = ggml_cont(ctx->gctx.get(), ggml_permute(ctx->gctx.get(), context_layer, 0, 2, 1, 3));
+    tensor_assign_buffers(context_layer);
+    context_layer = ggml_reshape_2d(ctx->gctx.get(), context_layer, hidden_size, qlen);
+    tensor_assign_buffers(context_layer);
 
     ggml_tensor *attn_output = dense.forward(ctx, context_layer);
     return attn_output;
@@ -689,14 +816,19 @@ ggml_tensor *GLMBlock::forward(ForwardContext *ctx, ggml_tensor *hidden_states,
     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);
+    attn_input = ggml_scale_inplace(ctx->gctx.get(), attn_input, alpha);
+    tensor_assign_buffers(attn_input);
+    hidden_states = ggml_add_inplace(ctx->gctx.get(), attn_input, attn_output);
+    tensor_assign_buffers(hidden_states);
 
     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);
+    mlp_input = ggml_scale_inplace(ctx->gctx.get(), mlp_input, alpha);
+    tensor_assign_buffers(mlp_input);
+    ggml_tensor *output = ggml_add_inplace(ctx->gctx.get(), mlp_input, mlp_output);
+    tensor_assign_buffers(output);
+
     return output;
 }
 
@@ -724,66 +856,71 @@ ggml_tensor *ChatGLMModel::forward(ForwardContext *ctx, ggml_tensor *input_ids,
 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_hidden_layers * 14;
+    const size_t num_tensors = 4 + config.num_hidden_layers * 14;
     const size_t ctx_size = num_tensors * tensor_ovhd;
     w_ctx_.gctx = GGMLContext(ctx_size, nullptr, true);
     w_ctx_.dtype = config.dtype;
 
     transformer = ChatGLMModel(&w_ctx_, config);
+    lm_head = Linear(&w_ctx_, config.hidden_size, config.vocab_size, false);
+    CHATGLM_CHECK(ggml_used_mem(w_ctx_.gctx.get()) == ggml_get_mem_size(w_ctx_.gctx.get()))
+        << "corrupted model weights";
 
     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);
+    for (auto &layer : transformer.layers) {
+        layer.attention.k_cache->data = kv_cache_ptr;
+        kv_cache_ptr += ggml_nbytes(layer.attention.k_cache);
+        tensor_to_device(layer.attention.k_cache);
+        layer.attention.v_cache->data = kv_cache_ptr;
+        kv_cache_ptr += ggml_nbytes(layer.attention.v_cache);
+        tensor_to_device(layer.attention.v_cache);
     }
     CHATGLM_CHECK(kv_cache_ptr == kv_cache_buffer_.get() + kv_cache_size) << "corrupted kv cache";
 }
 
 void ChatGLMForConditionalGeneration::load(ModelLoader &loader) {
-    loader.read_tensor("transformer.word_embeddings.weight", transformer.word_embeddings.weight);
+    std::vector<std::pair<std::string, ggml_tensor *>> state_dict;
+    state_dict.reserve(3 + config.num_hidden_layers * 12);
+
+    state_dict.emplace_back("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) + '.';
-        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);
+        state_dict.emplace_back(layer_prefix + "input_layernorm.weight", transformer.layers[i].input_layernorm.weight);
+        state_dict.emplace_back(layer_prefix + "input_layernorm.bias", transformer.layers[i].input_layernorm.bias);
+        state_dict.emplace_back(layer_prefix + "attention.query_key_value.weight",
+                                transformer.layers[i].attention.query_key_value.weight);
+        state_dict.emplace_back(layer_prefix + "attention.query_key_value.bias",
+                                transformer.layers[i].attention.query_key_value.bias);
+        state_dict.emplace_back(layer_prefix + "attention.dense.weight", transformer.layers[i].attention.dense.weight);
+        state_dict.emplace_back(layer_prefix + "attention.dense.bias", transformer.layers[i].attention.dense.bias);
+        state_dict.emplace_back(layer_prefix + "post_attention_layernorm.weight",
+                                transformer.layers[i].post_attention_layernorm.weight);
+        state_dict.emplace_back(layer_prefix + "post_attention_layernorm.bias",
+                                transformer.layers[i].post_attention_layernorm.bias);
+        state_dict.emplace_back(layer_prefix + "mlp.dense_h_to_4h.weight",
+                                transformer.layers[i].mlp.dense_h_to_4h.weight);
+        state_dict.emplace_back(layer_prefix + "mlp.dense_h_to_4h.bias", transformer.layers[i].mlp.dense_h_to_4h.bias);
+        state_dict.emplace_back(layer_prefix + "mlp.dense_4h_to_h.weight",
+                                transformer.layers[i].mlp.dense_4h_to_h.weight);
+        state_dict.emplace_back(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
-    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);
+    state_dict.emplace_back("transformer.final_layernorm.weight", transformer.final_layernorm.weight);
+    state_dict.emplace_back("transformer.final_layernorm.bias", transformer.final_layernorm.bias);
+
+    for (auto &item : state_dict) {
+        const std::string &name = item.first;
+        ggml_tensor *tensor = item.second;
+        loader.read_tensor(name, tensor);
+        if (name != "transformer.word_embeddings.weight") {
+            tensor_to_device(tensor);
+        }
     }
-#endif
+
+    lm_head.weight->data = transformer.word_embeddings.weight->data; // tied weight
+    tensor_to_device(lm_head.weight);
 }
 
 ggml_tensor *ChatGLMForConditionalGeneration::forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past,
@@ -794,9 +931,9 @@ ggml_tensor *ChatGLMForConditionalGeneration::forward(ForwardContext *ctx, ggml_
         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));
+        tensor_assign_buffers(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);
+    ggml_tensor *lm_logits = lm_head.forward(ctx, transformer_outputs);
     return lm_logits;
 }
 
@@ -873,47 +1010,88 @@ ggml_tensor *GLM2SelfAttention::forward(ForwardContext *ctx, ggml_tensor *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]
+#ifdef GGML_USE_CUBLAS
+    query_layer = ggml_cont(ctx->gctx.get(), query_layer);
+    tensor_assign_buffers(query_layer);
+#endif
     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]
+    tensor_assign_buffers(query_layer);
+    query_layer =
+        ggml_cont(ctx->gctx.get(), ggml_permute(ctx->gctx.get(), query_layer, 0, 2, 1, 3)); // [heads, qlen, head_size]
+    tensor_assign_buffers(query_layer);
+    query_layer = ggml_reshape_3d(ctx->gctx.get(), query_layer, head_size, mqa_scale * qlen,
+                                  num_kv_heads); // [kv_heads, mqa_scale * qlen, head_size]
+    tensor_assign_buffers(query_layer);
 
     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]
+#ifdef GGML_USE_CUBLAS
+    key_layer = ggml_cont(ctx->gctx.get(), key_layer);
+    tensor_assign_buffers(key_layer);
+#endif
     key_layer = ggml_rope_inplace(ctx->gctx.get(), key_layer, n_past, rope_dim, 0, 0);
+    tensor_assign_buffers(key_layer);
     key_layer = ggml_permute(ctx->gctx.get(), key_layer, 0, 2, 1, 3); // [kv_heads, qlen, head_size]
+    tensor_assign_buffers(key_layer);
 
     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]
+    tensor_assign_buffers(value_layer);
 
     // 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]
+    tensor_assign_buffers(k_cache_view);
     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]
+    tensor_assign_buffers(v_cache_view);
     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]
+    key_layer =
+        ggml_view_3d(ctx->gctx.get(), k_cache, head_size, n_past + qlen, num_kv_heads, k_cache->nb[1], k_cache->nb[2],
+                     0); // [kv_heads, klen, head_size]
+    tensor_assign_buffers(key_layer);
+    value_layer =
+        ggml_view_3d(ctx->gctx.get(), v_cache, n_past + qlen, head_size, num_kv_heads, v_cache->nb[1], v_cache->nb[2],
+                     0); // [kv_heads, head_size, klen]
+    tensor_assign_buffers(value_layer);
+
+    // attention
+    ggml_tensor *attn_scores =
+        ggml_mul_mat(ctx->gctx.get(), key_layer, query_layer); // [kv_heads, mqa_scale * qlen, klen]
+    tensor_assign_buffers(attn_scores);
+    attn_scores =
+        ggml_scale_inplace(ctx->gctx.get(), attn_scores, ggml_new_f32(ctx->gctx.get(), 1.f / std::sqrt(head_size)));
+    tensor_assign_buffers(attn_scores);
+    if (n_past == 0) {
+        // build attention mask for context input
+        attn_scores = ggml_reshape_3d(ctx->gctx.get(), attn_scores, n_past + qlen, qlen,
+                                      num_attention_heads); // [heads, qlen, klen]
+        attn_scores = ggml_diag_mask_inf_inplace(ctx->gctx.get(), attn_scores, n_past);
+        tensor_assign_buffers(attn_scores);
+        attn_scores = ggml_reshape_3d(ctx->gctx.get(), attn_scores, n_past + qlen, mqa_scale * qlen,
+                                      num_kv_heads); // [kv_heads, mqa_scale * qlen, klen]
+    }
+    ggml_tensor *attn_probs = ggml_soft_max_inplace(ctx->gctx.get(), attn_scores); // [kv_heads, mqa_scale * qlen, klen]
+    tensor_assign_buffers(attn_probs);
+
+    ggml_tensor *context_layer =
+        ggml_mul_mat(ctx->gctx.get(), value_layer, attn_probs); // [kv_heads, mqa_scale * qlen, head_size]
+    tensor_assign_buffers(context_layer);
+    context_layer = ggml_reshape_3d(ctx->gctx.get(), context_layer, head_size, qlen,
+                                    num_attention_heads); // [heads, qlen, head_size]
+    context_layer = ggml_cont(ctx->gctx.get(),
+                              ggml_permute(ctx->gctx.get(), context_layer, 0, 2, 1, 3)); // [qlen, heads, head_size]
+    tensor_assign_buffers(context_layer);
+    context_layer = ggml_reshape_2d(ctx->gctx.get(), context_layer, hidden_size, qlen); // [qlen, hidden]
+    tensor_assign_buffers(context_layer);
 
     ggml_tensor *attn_output = dense.forward(ctx, context_layer);
     return attn_output;
@@ -924,28 +1102,37 @@ ggml_tensor *GLM2MLP::forward(ForwardContext *ctx, ggml_tensor *hidden_states) c
 
     // swiglu activation
     ggml_tensor *x0 = ggml_view_2d(ctx->gctx.get(), output, output->ne[0] / 2, output->ne[1], output->nb[1], 0);
+    x0 = ggml_cont(ctx->gctx.get(), x0);
+    tensor_assign_buffers(x0);
+    x0 = ggml_silu_inplace(ctx->gctx.get(), x0);
+    tensor_assign_buffers(x0);
+
     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);
+#ifdef GGML_USE_CUBLAS
+    x1 = ggml_cont(ctx->gctx.get(), x1);
+#endif
+    tensor_assign_buffers(x1);
+
+    output = ggml_mul_inplace(ctx->gctx.get(), x0, x1);
+    tensor_assign_buffers(output);
 
     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);
-
+    ggml_tensor *residual = hidden_states;
     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);
+    tensor_assign_buffers(hidden_states);
 
-    residual = ggml_dup(ctx->gctx.get(), hidden_states);
-    ggml_build_forward_expand(&ctx->gf, residual);
-
+    residual = hidden_states;
     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);
+    tensor_assign_buffers(hidden_states);
 
     return hidden_states;
 }
@@ -984,54 +1171,55 @@ ChatGLM2ForConditionalGeneration::ChatGLM2ForConditionalGeneration(const ChatGLM
     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);
+                                 config.num_attention_heads * config.num_kv_heads * 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);
+    for (auto &layer : transformer.layers) {
+        layer.attention.k_cache->data = kv_cache_ptr;
+        kv_cache_ptr += ggml_nbytes(layer.attention.k_cache);
+        tensor_to_device(layer.attention.k_cache);
+        layer.attention.v_cache->data = kv_cache_ptr;
+        kv_cache_ptr += ggml_nbytes(layer.attention.v_cache);
+        tensor_to_device(layer.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);
+    std::vector<std::pair<std::string, ggml_tensor *>> state_dict;
+    state_dict.reserve(3 + config.num_hidden_layers * 7);
+
+    state_dict.emplace_back("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);
+        state_dict.emplace_back(layer_prefix + "input_layernorm.weight", transformer.layers[i].input_layernorm.weight);
+        state_dict.emplace_back(layer_prefix + "self_attention.query_key_value.weight",
+                                transformer.layers[i].attention.query_key_value.weight);
+        state_dict.emplace_back(layer_prefix + "self_attention.query_key_value.bias",
+                                transformer.layers[i].attention.query_key_value.bias);
+        state_dict.emplace_back(layer_prefix + "self_attention.dense.weight",
+                                transformer.layers[i].attention.dense.weight);
+        state_dict.emplace_back(layer_prefix + "post_attention_layernorm.weight",
+                                transformer.layers[i].post_attention_layernorm.weight);
+        state_dict.emplace_back(layer_prefix + "mlp.dense_h_to_4h.weight",
+                                transformer.layers[i].mlp.dense_h_to_4h.weight);
+        state_dict.emplace_back(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);
+    state_dict.emplace_back("transformer.encoder.final_layernorm.weight", transformer.final_layernorm.weight);
+    state_dict.emplace_back("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);
+    for (auto &item : state_dict) {
+        const std::string &name = item.first;
+        ggml_tensor *tensor = item.second;
+        loader.read_tensor(name, tensor);
+        if (name != "transformer.embedding.word_embeddings.weight") {
+            tensor_to_device(tensor);
+        }
     }
-    transform_tensor_to_cuda(lm_head.weight);
-#endif
 }
 
 ggml_tensor *ChatGLM2ForConditionalGeneration::forward(ForwardContext *ctx, ggml_tensor *input_ids, int n_past,
@@ -1042,6 +1230,7 @@ ggml_tensor *ChatGLM2ForConditionalGeneration::forward(ForwardContext *ctx, ggml
         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));
+        tensor_assign_buffers(transformer_outputs);
     }
     ggml_tensor *lm_logits = lm_head.forward(ctx, transformer_outputs);
     return lm_logits;
diff --git a/chatglm.h b/chatglm.h
index 83375f7..03091f8 100644
--- a/chatglm.h
+++ b/chatglm.h
@@ -25,8 +25,20 @@ class LogMessageFatal {
     if (!(cond))                                                                                                       \
     CHATGLM_THROW << "check failed (" #cond ") "
 
+#define CHATGLM_CHECK_CUDA(call)                                                                                       \
+    do {                                                                                                               \
+        cudaError_t error = (call);                                                                                    \
+        CHATGLM_CHECK(error == cudaSuccess) << "CUDA error: " << cudaGetErrorString(error);                            \
+    } while (0)
+
 std::string to_string(ggml_tensor *tensor, bool with_data = true);
 
+void tensor_assign_buffers(ggml_tensor *tensor, bool scratch = true, bool force_inplace = false);
+
+void tensor_to_device(ggml_tensor *tensor);
+
+void tensor_to_cpu(ggml_tensor *tensor);
+
 struct BaseConfig {
     // common attributes
     ggml_type dtype;
@@ -145,14 +157,15 @@ class LayerNorm {
 
 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)) {}
+    RMSNorm() : weight(nullptr), inplace(true) {}
+    RMSNorm(InitContext *ctx, int normalized_shape, bool inplace = true)
+        : weight(ggml_new_tensor_1d(ctx->gctx.get(), GGML_TYPE_F32, normalized_shape)), inplace(inplace) {}
 
     ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *input) const;
 
   public:
     ggml_tensor *weight;
+    bool inplace;
 };
 
 class BaseStreamer {
@@ -200,10 +213,14 @@ class PerfStreamer : public BaseStreamer {
 
     int64_t num_prompt_tokens() const { return num_prompt_tokens_; }
     int64_t prompt_total_time_us() const { return prompt_us_ - start_us_; }
-    int64_t prompt_token_time_us() const { return prompt_total_time_us() / num_prompt_tokens(); }
+    int64_t prompt_token_time_us() const {
+        return num_prompt_tokens() ? prompt_total_time_us() / num_prompt_tokens() : 0;
+    }
     int64_t num_output_tokens() const { return num_output_tokens_; }
     int64_t output_total_time_us() const { return end_us_ - prompt_us_; }
-    int64_t output_token_time_us() const { return output_total_time_us() / num_output_tokens(); }
+    int64_t output_token_time_us() const {
+        return num_output_tokens() ? output_total_time_us() / num_output_tokens() : 0;
+    }
 
   private:
     int64_t start_us_;
@@ -270,8 +287,25 @@ enum ModelType {
     MODEL_TYPE_CHATGLM2 = 2,
 };
 
+int get_num_physical_cores();
+
 std::string to_string(ModelType model_type);
 
+struct TokenIdScore {
+    int id;
+    float score;
+
+    TokenIdScore() = default;
+    TokenIdScore(int id, float score) : id(id), score(score) {}
+
+    bool operator<(const TokenIdScore &other) const { return score < other.score; }
+    bool operator>(const TokenIdScore &other) const { return score > other.score; }
+
+    friend std::ostream &operator<<(std::ostream &os, const TokenIdScore &self) {
+        return os << "TokenIdScore(id=" << self.id << ", score=" << self.score << ")";
+    }
+};
+
 class BaseModelForConditionalGeneration {
   public:
     BaseModelForConditionalGeneration(ModelType model_type, BaseConfig config, size_t mem_size, size_t scratch_size)
@@ -291,18 +325,10 @@ class BaseModelForConditionalGeneration {
     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;
-
-        TokenIdScore() = default;
-        TokenIdScore(int id, float score) : id(id), score(score) {}
+    static void sampling_temperature(float *first, float *last, float temp);
+    static void sampling_top_k(TokenIdScore *first, TokenIdScore *kth, TokenIdScore *last);
+    static TokenIdScore *sampling_top_p(TokenIdScore *first, TokenIdScore *last, float top_p);
 
-        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:
@@ -361,13 +387,7 @@ 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)) {}
+    GLMSelfAttention(InitContext *ctx, int hidden_size, int num_attention_heads, int max_length);
 
     ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past, int n_ctx) const;
 
@@ -420,9 +440,10 @@ class ChatGLMForConditionalGeneration : public BaseModelForConditionalGeneration
 
   public:
     static constexpr size_t MEM_SIZE = 272ull * 1024 * 1024;
-    static constexpr size_t SCRATCH_SIZE = 128ull * 1024 * 1024;
+    static constexpr size_t SCRATCH_SIZE = 256ull * 1024 * 1024;
     ChatGLMConfig config;
     ChatGLMModel transformer;
+    Linear lm_head;
 
   private:
     // hold ggml_context & kv_cache
@@ -467,9 +488,9 @@ class GLM2SelfAttention {
         : 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,
+          k_cache(ggml_new_tensor_3d(ctx->gctx.get(), GGML_TYPE_F16, 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,
+          v_cache(ggml_new_tensor_3d(ctx->gctx.get(), GGML_TYPE_F16, max_length, hidden_size / num_attention_heads,
                                      num_kv_heads)) {}
 
     ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past) const;
@@ -501,8 +522,9 @@ class GLM2Block {
     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) {}
+        : input_layernorm(ctx, hidden_size, false),
+          attention(ctx, hidden_size, num_attention_heads, num_kv_heads, max_length),
+          post_attention_layernorm(ctx, hidden_size, false), mlp(ctx, hidden_size, intermediate_size) {}
 
     ggml_tensor *forward(ForwardContext *ctx, ggml_tensor *hidden_states, int n_past) const;
 
@@ -537,7 +559,7 @@ class ChatGLM2ForConditionalGeneration : public BaseModelForConditionalGeneratio
 
   public:
     static constexpr size_t MEM_SIZE = 512ull * 1024 * 1024;
-    static constexpr size_t SCRATCH_SIZE = 144ull * 1024 * 1024;
+    static constexpr size_t SCRATCH_SIZE = 256ull * 1024 * 1024;
 
     ChatGLM2Config config;
     ChatGLM2Model transformer;
diff --git a/third_party/ggml b/third_party/ggml
index ee1b372..8b36204 160000
--- a/third_party/ggml
+++ b/third_party/ggml
@@ -1,1 +1,1 @@
-Subproject commit ee1b3727e60403012dd2b57d35b60558f4db66d8
+Subproject commit 8b3620478d9b5becfdf3b91f845ee0004ef9b01c
diff --git a/third_party/pybind11 b/third_party/pybind11
index be97c5a..8b03ffa 160000
--- a/third_party/pybind11
+++ b/third_party/pybind11
@@ -1,1 +1,1 @@
-Subproject commit be97c5a98b4b252c524566f508b5c79410d118c6
+Subproject commit 8b03ffa7c06cd9c8a38297b1c8923695d1ff1b07
diff --git a/third_party/sentencepiece b/third_party/sentencepiece
index 7b694e4..635fe84 160000
--- a/third_party/sentencepiece
+++ b/third_party/sentencepiece
@@ -1,1 +1,1 @@
-Subproject commit 7b694e4bdb8d08be405f3370a6d147d9b8fcea76
+Subproject commit 635fe8423a249b6e081aacd290d8aef7476c6a28


Test patch:
diff --git a/chatglm_test.cpp b/chatglm_test.cpp
index 1d800d7..6b8c7f9 100644
--- a/chatglm_test.cpp
+++ b/chatglm_test.cpp
@@ -2,10 +2,25 @@
 #include <filesystem>
 #include <gtest/gtest.h>
 
+#ifdef GGML_USE_CUBLAS
+#include <cuda_runtime.h>
+#include <ggml-cuda.h>
+#endif
+
 namespace chatglm {
 
 namespace fs = std::filesystem;
 
+static inline int get_num_threads_for_benchmark() {
+#ifdef GGML_USE_CUBLAS
+    return 1;
+#endif
+    const char *chatglm_num_threads_env = getenv("CHATGLM_NUM_THREADS");
+    int num_threads_for_benchmark =
+        chatglm_num_threads_env ? std::stoi(chatglm_num_threads_env) : get_num_physical_cores();
+    return num_threads_for_benchmark;
+}
+
 static inline void expect_all_close(ggml_tensor *a, ggml_tensor *b, float atol = 1e-5) {
     ASSERT_EQ(a->type, b->type);
     ASSERT_EQ(a->type, GGML_TYPE_F32);
@@ -16,32 +31,195 @@ static inline void expect_all_close(ggml_tensor *a, ggml_tensor *b, float atol =
     }
 }
 
-static inline char *map_tensor_data(char *ptr, ggml_tensor *tensor) {
-    tensor->data = ptr;
+static inline char *read_tensor_data(char *ptr, ggml_tensor *tensor) {
+    memcpy(tensor->data, ptr, ggml_nbytes(tensor));
     return ptr + ggml_nbytes(tensor);
 }
 
+static inline float random() { return rand() / (float)RAND_MAX; }
+
+static inline void random_fill(ggml_tensor *tensor) {
+    std::vector<float> values(ggml_nelements(tensor));
+    for (float &v : values) {
+        v = random();
+    }
+    int64_t hist[16]{};
+
+    if (tensor->type == GGML_TYPE_F32) {
+        memcpy(tensor->data, values.data(), sizeof(float) * values.size());
+    } else if (tensor->type == GGML_TYPE_F16) {
+        ggml_fp32_to_fp16_row(values.data(), (ggml_fp16_t *)tensor->data, values.size());
+    } else if (tensor->type == GGML_TYPE_Q8_0) {
+        ggml_quantize_q8_0(values.data(), tensor->data, ggml_nelements(tensor), tensor->ne[0], hist);
+    } else if (tensor->type == GGML_TYPE_Q4_0) {
+        ggml_quantize_q4_0(values.data(), tensor->data, ggml_nelements(tensor), tensor->ne[0], hist);
+    } else {
+        CHATGLM_THROW << "unsupported dtype " << ggml_type_name(tensor->type);
+    }
+}
+
+// return elapsed time in milliseconds
+static inline float timeit(std::function<void()> fn, int warmup, int active) {
+    float elapsed_ms;
+    for (int i = 0; i < warmup; i++) {
+        fn();
+    }
+
+#ifdef GGML_USE_CUBLAS
+    cudaEvent_t start, stop;
+    CHATGLM_CHECK_CUDA(cudaEventCreate(&start));
+    CHATGLM_CHECK_CUDA(cudaEventCreate(&stop));
+    CHATGLM_CHECK_CUDA(cudaEventRecord(start));
+    for (int i = 0; i < active; i++) {
+        fn();
+    }
+    CHATGLM_CHECK_CUDA(cudaEventRecord(stop));
+    CHATGLM_CHECK_CUDA(cudaEventSynchronize(stop));
+    CHATGLM_CHECK_CUDA(cudaEventElapsedTime(&elapsed_ms, start, stop));
+    CHATGLM_CHECK_CUDA(cudaEventDestroy(start));
+    CHATGLM_CHECK_CUDA(cudaEventDestroy(stop));
+#else
+    int64_t start_us = ggml_time_us();
+    for (int i = 0; i < active; i++) {
+        fn();
+    }
+    int64_t end_us = ggml_time_us();
+    elapsed_ms = (end_us - start_us) / 1000.f;
+#endif
+
+    return elapsed_ms / active;
+}
+
+static bool equal(const std::vector<int> &a, const std::vector<int> &b) {
+    if (a.size() != b.size()) {
+        return false;
+    }
+    for (size_t i = 0; i < a.size(); i++) {
+        if (a[i] != b[i]) {
+            return false;
+        }
+    }
+    return true;
+}
+
+static std::vector<int> extract_sorted_ids(std::vector<TokenIdScore> &token_scores) {
+    std::vector<int> token_ids(token_scores.size());
+    for (size_t i = 0; i < token_scores.size(); i++) {
+        token_ids[i] = token_scores[i].id;
+    }
+    std::sort(token_ids.begin(), token_ids.end());
+    return token_ids;
+}
+
+TEST(Sampling, Temperature) {
+    constexpr float temp = 0.7;
+    std::vector<float> logits(64);
+    for (float &v : logits) {
+        v = random();
+    }
+    // reference
+    std::vector<float> target = logits;
+    for (auto &v : target) {
+        v /= temp;
+    }
+    // test
+    BaseModelForConditionalGeneration::sampling_temperature(logits.data(), logits.data() + logits.size(), temp);
+    // compare
+    for (size_t i = 0; i < logits.size(); i++) {
+        EXPECT_FLOAT_EQ(logits[i], target[i]);
+    }
+}
+
+TEST(Sampling, TopK) {
+    constexpr int top_k = 20;
+    std::vector<TokenIdScore> token_scores(64);
+    for (size_t i = 0; i < token_scores.size(); i++) {
+        token_scores[i] = TokenIdScore(i, random());
+    }
+
+    // reference
+    std::vector<TokenIdScore> target = token_scores;
+    std::sort(target.begin(), target.end(), std::greater<TokenIdScore>());
+    target.resize(top_k);
+
+    // test
+    BaseModelForConditionalGeneration::sampling_top_k(token_scores.data(), token_scores.data() + top_k,
+                                                      token_scores.data() + token_scores.size());
+    token_scores.resize(top_k);
+
+    // sort & compare
+    EXPECT_TRUE(equal(extract_sorted_ids(token_scores), extract_sorted_ids(target)));
+}
+
+static void reference_top_p(std::vector<TokenIdScore> &token_scores, float top_p) {
+    std::sort(token_scores.begin(), token_scores.end(), std::greater<TokenIdScore>());
+    BaseModelForConditionalGeneration::sampling_softmax_inplace(token_scores.data(),
+                                                                token_scores.data() + token_scores.size());
+    float cumsum = 0.f;
+    for (size_t i = 0; i < token_scores.size(); i++) {
+        cumsum += token_scores[i].score;
+        if (cumsum >= top_p) {
+            token_scores.resize(i + 1);
+            break;
+        }
+    }
+}
+
+TEST(Sampling, TopP) {
+    constexpr float top_p = 0.7;
+    for (int i = 0; i < 10; i++) {
+        std::vector<TokenIdScore> token_scores(1024);
+        for (size_t i = 0; i < token_scores.size(); i++) {
+            token_scores[i] = TokenIdScore(i, random());
+        }
+
+        // reference
+        std::vector<TokenIdScore> target = token_scores;
+        reference_top_p(target, top_p);
+        EXPECT_TRUE(!token_scores.empty());
+
+        // test
+        TokenIdScore *pos = BaseModelForConditionalGeneration::sampling_top_p(
+            token_scores.data(), token_scores.data() + token_scores.size(), top_p);
+        token_scores.resize(pos - token_scores.data());
+
+        // sort & compare
+        auto output_ids = extract_sorted_ids(token_scores);
+        auto target_ids = extract_sorted_ids(target);
+        EXPECT_TRUE(equal(output_ids, target_ids)) << "size " << output_ids.size() << " vs " << target_ids.size();
+    }
+}
+
 class ChatGLMTest : public ::testing::Test {
   protected:
     InitContext ictx;
     ForwardContext ctx;
     std::vector<char> scratch_buf;
+    int num_benchmark_threads_;
 
     void SetUp() override {
         ictx.dtype = GGML_TYPE_F32;
-        ictx.gctx = GGMLContext({1024 * 1024, nullptr, false});
+        ictx.gctx = GGMLContext({1024 * 1024 * 1024, nullptr, false});
 
         scratch_buf.resize(1024 * 1024);
 
-        ctx.gctx = GGMLContext({1024 * 1024, nullptr, false});
+        ctx.gctx = GGMLContext({1024 * 1024 * 1024, nullptr, false});
         ctx.scratch = {0, scratch_buf.size(), scratch_buf.data()};
 
         reset_cgraph();
+
+        num_benchmark_threads_ = get_num_threads_for_benchmark();
     }
 
-    void reset_cgraph() {
-        ctx.gf = {};
-        ctx.gf.n_threads = 1;
+    void reset_cgraph() { ctx.gf = {}; }
+
+    float perf_compute() {
+        auto fn = [this] { ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, num_benchmark_threads_); };
+        if (ggml_cpu_has_cublas()) {
+            return timeit(fn, 10, 100);
+        } else {
+            return timeit(fn, 1, 3);
+        }
     }
 };
 
@@ -62,7 +240,7 @@ TEST_F(ChatGLMTest, Embedding) {
     ggml_tensor *out = model.forward(&ctx, x);
 
     ggml_build_forward_expand(&ctx.gf, out);
-    ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+    ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, 1);
 
     expect_all_close(ref, out);
 }
@@ -73,28 +251,37 @@ TEST_F(ChatGLMTest, Linear) {
     char *ptr = mapped_file.data;
 
     ggml_tensor *w = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 32, 16);
-    ptr = map_tensor_data(ptr, w);
+    ptr = read_tensor_data(ptr, w);
     ggml_tensor *b = ggml_new_tensor_1d(ctx.gctx.get(), GGML_TYPE_F32, 16);
-    ptr = map_tensor_data(ptr, b);
+    ptr = read_tensor_data(ptr, b);
     ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 32, 2);
-    ptr = map_tensor_data(ptr, x);
+    ptr = read_tensor_data(ptr, x);
     ggml_tensor *ref = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 16, 2);
-    ptr = map_tensor_data(ptr, ref);
+    ptr = read_tensor_data(ptr, ref);
     ASSERT_EQ(ptr, mapped_file.data + mapped_file.size);
 
+    tensor_to_device(x);
+
     // fp32
     {
         ictx.dtype = GGML_TYPE_F32;
         Linear model(&ictx, 32, 16);
         model.weight->data = w->data;
         model.bias->data = b->data;
+        tensor_to_device(model.weight);
+        tensor_to_device(model.bias);
 
         ggml_tensor *out = model.forward(&ctx, x);
+        EXPECT_EQ(out->backend, x->backend);
+        out->backend = GGML_BACKEND_CPU;
 
         ggml_build_forward_expand(&ctx.gf, out);
-        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+        ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, 1);
 
         expect_all_close(ref, out);
+
+        tensor_to_cpu(model.weight);
+        tensor_to_cpu(model.bias);
     }
     // fp16
     {
@@ -104,14 +291,43 @@ TEST_F(ChatGLMTest, Linear) {
         Linear model(&ictx, 32, 16);
         ggml_fp32_to_fp16_row((float *)w->data, (ggml_fp16_t *)model.weight->data, ggml_nelements(model.weight));
         model.bias->data = b->data;
+        tensor_to_device(model.weight);
+        tensor_to_device(model.bias);
 
         ggml_tensor *out = model.forward(&ctx, x);
+        EXPECT_EQ(out->backend, x->backend);
+        out->backend = GGML_BACKEND_CPU;
 
         ggml_build_forward_expand(&ctx.gf, out);
-        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+        ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, 1);
 
         EXPECT_EQ(out->type, GGML_TYPE_F32);
         expect_all_close(ref, out, 5e-3);
+
+        tensor_to_cpu(model.weight);
+        tensor_to_cpu(model.bias);
+    }
+    tensor_to_cpu(x);
+}
+
+TEST_F(ChatGLMTest, BenchmarkLinear) {
+    constexpr int M = 64, N = 1024, K = 1024 * 3;
+    ictx.dtype = GGML_TYPE_F32;
+    Linear m(&ictx, K, N);
+    ggml_tensor *x = ggml_new_tensor_2d(ictx.gctx.get(), GGML_TYPE_F32, K, M);
+
+    std::vector<ggml_tensor *> all_tensors{m.weight, m.bias, x};
+    for (auto tensor : all_tensors) {
+        random_fill(tensor);
+        tensor_to_device(tensor);
+    }
+
+    ggml_tensor *y = m.forward(&ctx, x);
+    ggml_build_forward_expand(&ctx.gf, y);
+    std::cout << "[Benchmark] Linear " << ggml_type_name(ictx.dtype) << " time: " << perf_compute() << " ms\n";
+
+    for (auto tensor : all_tensors) {
+        tensor_to_cpu(tensor);
     }
 }
 
@@ -121,23 +337,50 @@ TEST_F(ChatGLMTest, LayerNorm) {
     char *ptr = mapped_file.data;
 
     LayerNorm model(&ictx, 64);
-    ptr = map_tensor_data(ptr, model.weight);
-    ptr = map_tensor_data(ptr, model.bias);
-
     ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 64, 3);
-    ptr = map_tensor_data(ptr, x);
-    x = ggml_dup(ctx.gctx.get(), x);
-
     ggml_tensor *ref = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 64, 3);
-    ptr = map_tensor_data(ptr, ref);
 
+    std::vector<ggml_tensor *> all_tensors{model.weight, model.bias, x, ref};
+    for (auto tensor : all_tensors) {
+        ptr = read_tensor_data(ptr, tensor);
+        tensor_to_device(tensor);
+    }
     ASSERT_EQ(ptr, mapped_file.data + mapped_file.size);
 
     ggml_tensor *out = model.forward(&ctx, x);
+    EXPECT_EQ(out->backend, x->backend);
+    out->backend = GGML_BACKEND_CPU;
+
     ggml_build_forward_expand(&ctx.gf, out);
-    ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+    ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, num_benchmark_threads_);
 
     expect_all_close(ref, out);
+
+    for (auto tensor : all_tensors) {
+        tensor_to_cpu(tensor);
+    }
+}
+
+TEST_F(ChatGLMTest, BenchmarkLayerNorm) {
+    constexpr int seq_len = 64;
+    constexpr int hidden = 1024;
+
+    LayerNorm m(&ictx, hidden);
+    ggml_tensor *x = ggml_new_tensor_2d(ictx.gctx.get(), GGML_TYPE_F32, hidden, seq_len);
+
+    std::vector<ggml_tensor *> all_tensors{m.weight, m.bias, x};
+    for (auto tensor : all_tensors) {
+        random_fill(tensor);
+        tensor_to_device(tensor);
+    }
+
+    ggml_tensor *y = m.forward(&ctx, x);
+    ggml_build_forward_expand(&ctx.gf, y);
+    std::cout << "[Benchmark] LayerNorm " << ggml_type_name(ictx.dtype) << " time: " << perf_compute() << " ms\n";
+
+    for (auto tensor : all_tensors) {
+        tensor_to_cpu(tensor);
+    }
 }
 
 TEST_F(ChatGLMTest, RMSNorm) {
@@ -146,23 +389,50 @@ TEST_F(ChatGLMTest, RMSNorm) {
     char *ptr = mapped_file.data;
 
     RMSNorm model(&ictx, 64);
-    ptr = map_tensor_data(ptr, model.weight);
-
     ggml_tensor *x = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 64, 3);
-    ptr = map_tensor_data(ptr, x);
-    x = ggml_dup(ctx.gctx.get(), x);
-
     ggml_tensor *ref = ggml_new_tensor_2d(ctx.gctx.get(), GGML_TYPE_F32, 64, 3);
-    ptr = map_tensor_data(ptr, ref);
 
+    std::vector<ggml_tensor *> all_tensors{model.weight, x, ref};
+    for (auto tensor : all_tensors) {
+        ptr = read_tensor_data(ptr, tensor);
+        tensor_to_device(tensor);
+    }
     ASSERT_EQ(ptr, mapped_file.data + mapped_file.size);
 
     ggml_tensor *out = model.forward(&ctx, x);
+    EXPECT_EQ(out->backend, x->backend);
+    out->backend = GGML_BACKEND_CPU;
 
     ggml_build_forward_expand(&ctx.gf, out);
-    ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+    ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, num_benchmark_threads_);
 
     expect_all_close(ref, out);
+
+    for (auto tensor : all_tensors) {
+        tensor_to_cpu(tensor);
+    }
+}
+
+TEST_F(ChatGLMTest, BenchmarkRMSNorm) {
+    constexpr int seq_len = 64;
+    constexpr int hidden = 1024;
+
+    RMSNorm m(&ictx, hidden);
+    ggml_tensor *x = ggml_new_tensor_2d(ictx.gctx.get(), GGML_TYPE_F32, hidden, seq_len);
+
+    std::vector<ggml_tensor *> all_tensors{m.weight, x};
+    for (auto tensor : all_tensors) {
+        random_fill(tensor);
+        tensor_to_device(tensor);
+    }
+
+    ggml_tensor *y = m.forward(&ctx, x);
+    ggml_build_forward_expand(&ctx.gf, y);
+    std::cout << "[Benchmark] RMSNorm " << ggml_type_name(ictx.dtype) << " time: " << perf_compute() << " ms\n";
+
+    for (auto tensor : all_tensors) {
+        tensor_to_cpu(tensor);
+    }
 }
 
 TEST_F(ChatGLMTest, GLMBlock) {
@@ -176,70 +446,139 @@ TEST_F(ChatGLMTest, GLMBlock) {
     constexpr int max_length = 16;
     constexpr int seq_len = 4;
     GLMBlock model(&ictx, hidden_size, num_attention_heads, num_hidden_layers, max_length);
-
-    ptr = map_tensor_data(ptr, model.input_layernorm.weight);
-    ptr = map_tensor_data(ptr, model.input_layernorm.bias);
-    ptr = map_tensor_data(ptr, model.attention.query_key_value.weight);
-    ptr = map_tensor_data(ptr, model.attention.query_key_value.bias);
-    ptr = map_tensor_data(ptr, model.attention.dense.weight);
-    ptr = map_tensor_data(ptr, model.attention.dense.bias);
-    ptr = map_tensor_data(ptr, model.post_attention_layernorm.weight);
-    ptr = map_tensor_data(ptr, model.post_attention_layernorm.bias);
-    ptr = map_tensor_data(ptr, model.mlp.dense_h_to_4h.weight);
-    ptr = map_tensor_data(ptr, model.mlp.dense_h_to_4h.bias);
-    ptr = map_tensor_data(ptr, model.mlp.dense_4h_to_h.weight);
-    ptr = map_tensor_data(ptr, model.mlp.dense_4h_to_h.bias);
+    tensor_to_device(model.attention.k_cache);
+    tensor_to_device(model.attention.v_cache);
 
     ggml_tensor *x1 = ggml_new_tensor_2d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size, seq_len);
-    ptr = map_tensor_data(ptr, x1);
-    x1 = ggml_dup(ictx.gctx.get(), x1);
-
     ggml_tensor *ref_y1 = ggml_new_tensor_2d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size, seq_len);
-    ptr = map_tensor_data(ptr, ref_y1);
-
     ggml_tensor *x2 = ggml_new_tensor_1d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size);
-    ptr = map_tensor_data(ptr, x2);
-    x2 = ggml_dup(ictx.gctx.get(), x2);
-
     ggml_tensor *ref_y2 = ggml_new_tensor_1d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size);
-    ptr = map_tensor_data(ptr, ref_y2);
-
     ggml_tensor *x3 = ggml_new_tensor_1d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size);
-    ptr = map_tensor_data(ptr, x3);
-    x3 = ggml_dup(ictx.gctx.get(), x3);
-
     ggml_tensor *ref_y3 = ggml_new_tensor_1d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size);
-    ptr = map_tensor_data(ptr, ref_y3);
+
+    std::vector<ggml_tensor *> all_tensors{model.input_layernorm.weight,
+                                           model.input_layernorm.bias,
+                                           model.attention.query_key_value.weight,
+                                           model.attention.query_key_value.bias,
+                                           model.attention.dense.weight,
+                                           model.attention.dense.bias,
+                                           model.post_attention_layernorm.weight,
+                                           model.post_attention_layernorm.bias,
+                                           model.mlp.dense_h_to_4h.weight,
+                                           model.mlp.dense_h_to_4h.bias,
+                                           model.mlp.dense_4h_to_h.weight,
+                                           model.mlp.dense_4h_to_h.bias,
+                                           x1,
+                                           ref_y1,
+                                           x2,
+                                           ref_y2,
+                                           x3,
+                                           ref_y3};
+
+    for (auto tensor : all_tensors) {
+        ptr = read_tensor_data(ptr, tensor);
+        tensor_to_device(tensor);
+    }
 
     ASSERT_EQ(ptr, mapped_file.data + mapped_file.size);
 
     // self attention
     {
         ggml_tensor *out_y1 = model.forward(&ctx, x1, 0, seq_len);
-
+        EXPECT_EQ(out_y1->backend, x1->backend);
+        out_y1->backend = GGML_BACKEND_CPU;
         ggml_build_forward_expand(&ctx.gf, out_y1);
-        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+        ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, 1);
 
         expect_all_close(ref_y1, out_y1, 5e-4);
     }
 
     // cross attention
+    reset_cgraph();
     {
         ggml_tensor *out_y2 = model.forward(&ctx, x2, seq_len, seq_len);
-
+        EXPECT_EQ(out_y2->backend, x2->backend);
+        out_y2->backend = GGML_BACKEND_CPU;
         ggml_build_forward_expand(&ctx.gf, out_y2);
-        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+        ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, 1);
 
         expect_all_close(ref_y2, out_y2, 5e-4);
     }
+    reset_cgraph();
     {
         ggml_tensor *out_y3 = model.forward(&ctx, x3, seq_len + 1, seq_len);
-
+        EXPECT_EQ(out_y3->backend, x3->backend);
+        out_y3->backend = GGML_BACKEND_CPU;
         ggml_build_forward_expand(&ctx.gf, out_y3);
-        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+        ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, 1);
 
         expect_all_close(ref_y3, out_y3, 5e-4);
     }
+
+    for (auto tensor : all_tensors) {
+        tensor_to_cpu(tensor);
+    }
+}
+
+TEST_F(ChatGLMTest, BenchmarkGLMBlock) {
+    constexpr int hidden_size = 4096;
+    constexpr int num_attention_heads = 32;
+    constexpr int num_hidden_layers = 28;
+    constexpr int max_length = 2048;
+    constexpr int seq_len = 64;
+
+    ggml_type dtypes[]{GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0};
+    for (const auto dtype : dtypes) {
+        SetUp();
+
+        ictx.dtype = dtype;
+        GLMBlock model(&ictx, hidden_size, num_attention_heads, num_hidden_layers, max_length);
+
+        ggml_tensor *self_attn_x = ggml_new_tensor_2d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size, seq_len);
+        ggml_tensor *cross_attn_x = ggml_new_tensor_1d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size);
+
+        std::vector<ggml_tensor *> all_tensors{model.input_layernorm.weight,
+                                               model.input_layernorm.bias,
+                                               model.attention.query_key_value.weight,
+                                               model.attention.query_key_value.bias,
+                                               model.attention.dense.weight,
+                                               model.attention.dense.bias,
+                                               model.post_attention_layernorm.weight,
+                                               model.post_attention_layernorm.bias,
+                                               model.mlp.dense_h_to_4h.weight,
+                                               model.mlp.dense_h_to_4h.bias,
+                                               model.mlp.dense_4h_to_h.weight,
+                                               model.mlp.dense_4h_to_h.bias,
+                                               self_attn_x,
+                                               cross_attn_x};
+
+        for (auto tensor : all_tensors) {
+            random_fill(tensor);
+            tensor_to_device(tensor);
+        }
+
+        // self attention
+        reset_cgraph();
+        {
+            ggml_tensor *self_attn_y = model.forward(&ctx, self_attn_x, 0, seq_len);
+            ggml_build_forward_expand(&ctx.gf, self_attn_y);
+            std::cout << "[Benchmark] GLMBlock " << ggml_type_name(dtype) << " self attn time: " << perf_compute()
+                      << " ms\n";
+        }
+
+        // cross attention
+        reset_cgraph();
+        {
+            ggml_tensor *cross_attn_y = model.forward(&ctx, cross_attn_x, seq_len, seq_len);
+            ggml_build_forward_expand(&ctx.gf, cross_attn_y);
+            std::cout << "[Benchmark] GLMBlock " << ggml_type_name(dtype) << " cross attn time: " << perf_compute()
+                      << " ms\n";
+        }
+
+        for (auto tensor : all_tensors) {
+            tensor_to_cpu(tensor);
+        }
+    }
 }
 
 TEST_F(ChatGLMTest, GLM2Block) {
@@ -255,64 +594,135 @@ TEST_F(ChatGLMTest, GLM2Block) {
     constexpr int max_length = 8;
 
     GLM2Block model(&ictx, hidden_size, num_attention_heads, num_kv_heads, ffn_hidden_size, max_length);
-    ptr = map_tensor_data(ptr, model.input_layernorm.weight);
-    ptr = map_tensor_data(ptr, model.attention.query_key_value.weight);
-    ptr = map_tensor_data(ptr, model.attention.query_key_value.bias);
-    ptr = map_tensor_data(ptr, model.attention.dense.weight);
-    ptr = map_tensor_data(ptr, model.post_attention_layernorm.weight);
-    ptr = map_tensor_data(ptr, model.mlp.dense_h_to_4h.weight);
-    ptr = map_tensor_data(ptr, model.mlp.dense_4h_to_h.weight);
+    tensor_to_device(model.attention.k_cache);
+    tensor_to_device(model.attention.v_cache);
 
     ggml_tensor *x1 = ggml_new_tensor_2d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size, seq_len);
-    ptr = map_tensor_data(ptr, x1);
-    x1 = ggml_dup(ictx.gctx.get(), x1);
-
     ggml_tensor *ref_y1 = ggml_new_tensor_2d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size, seq_len);
-    ptr = map_tensor_data(ptr, ref_y1);
-
     ggml_tensor *x2 = ggml_new_tensor_1d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size);
-    ptr = map_tensor_data(ptr, x2);
-    x2 = ggml_dup(ictx.gctx.get(), x2);
-
     ggml_tensor *ref_y2 = ggml_new_tensor_1d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size);
-    ptr = map_tensor_data(ptr, ref_y2);
-
     ggml_tensor *x3 = ggml_new_tensor_1d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size);
-    ptr = map_tensor_data(ptr, x3);
-    x3 = ggml_dup(ictx.gctx.get(), x3);
-
     ggml_tensor *ref_y3 = ggml_new_tensor_1d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size);
-    ptr = map_tensor_data(ptr, ref_y3);
 
+    std::vector<ggml_tensor *> all_tensors{model.input_layernorm.weight,
+                                           model.attention.query_key_value.weight,
+                                           model.attention.query_key_value.bias,
+                                           model.attention.dense.weight,
+                                           model.post_attention_layernorm.weight,
+                                           model.mlp.dense_h_to_4h.weight,
+                                           model.mlp.dense_4h_to_h.weight,
+                                           x1,
+                                           ref_y1,
+                                           x2,
+                                           ref_y2,
+                                           x3,
+                                           ref_y3};
+
+    for (auto tensor : all_tensors) {
+        ptr = read_tensor_data(ptr, tensor);
+        tensor_to_device(tensor);
+    }
     ASSERT_EQ(ptr, mapped_file.data + mapped_file.size);
 
     // self attention
+    reset_cgraph();
     {
         ggml_tensor *out_y1 = model.forward(&ctx, x1, 0);
+        EXPECT_EQ(out_y1->backend, x1->backend);
+        out_y1->backend = GGML_BACKEND_CPU;
         ggml_build_forward_expand(&ctx.gf, out_y1);
-        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+        ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, num_benchmark_threads_);
 
-        expect_all_close(ref_y1, out_y1, 5e-5);
+        expect_all_close(ref_y1, out_y1, 1e-4);
     }
 
     // cross attention
+    reset_cgraph();
     {
-        reset_cgraph();
-
         ggml_tensor *out_y2 = model.forward(&ctx, x2, seq_len);
+        EXPECT_EQ(out_y2->backend, x2->backend);
+        out_y2->backend = GGML_BACKEND_CPU;
         ggml_build_forward_expand(&ctx.gf, out_y2);
-        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+        ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, num_benchmark_threads_);
 
-        expect_all_close(ref_y2, out_y2, 5e-5);
+        expect_all_close(ref_y2, out_y2, 1e-4);
     }
+    reset_cgraph();
     {
-        reset_cgraph();
-
         ggml_tensor *out_y3 = model.forward(&ctx, x3, seq_len + 1);
+        EXPECT_EQ(out_y3->backend, x3->backend);
+        out_y3->backend = GGML_BACKEND_CPU;
         ggml_build_forward_expand(&ctx.gf, out_y3);
-        ggml_graph_compute(ctx.gctx.get(), &ctx.gf);
+        ggml_graph_compute_with_ctx(ctx.gctx.get(), &ctx.gf, num_benchmark_threads_);
+
+        expect_all_close(ref_y3, out_y3, 1e-4);
+    }
 
-        expect_all_close(ref_y3, out_y3, 5e-5);
+    for (auto tensor : all_tensors) {
+        tensor_to_cpu(tensor);
+    }
+    tensor_to_cpu(model.attention.k_cache);
+    tensor_to_cpu(model.attention.v_cache);
+}
+
+TEST_F(ChatGLMTest, BenchmarkGLM2Block) {
+    constexpr int seq_len = 64;
+    constexpr int hidden_size = 4096;
+    constexpr int num_attention_heads = 32;
+    constexpr int num_kv_heads = 2;
+    constexpr int ffn_hidden_size = 13696;
+    constexpr int max_length = 2048;
+
+    ggml_type dtypes[]{GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0};
+    for (const auto dtype : dtypes) {
+        SetUp();
+
+        ictx.dtype = dtype;
+        GLM2Block model(&ictx, hidden_size, num_attention_heads, num_kv_heads, ffn_hidden_size, max_length);
+        tensor_to_device(model.attention.k_cache);
+        tensor_to_device(model.attention.v_cache);
+
+        ggml_tensor *self_attn_x = ggml_new_tensor_2d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size, seq_len);
+        ggml_tensor *cross_attn_x = ggml_new_tensor_1d(ictx.gctx.get(), GGML_TYPE_F32, hidden_size);
+
+        std::vector<ggml_tensor *> all_tensors{model.input_layernorm.weight,
+                                               model.attention.query_key_value.weight,
+                                               model.attention.query_key_value.bias,
+                                               model.attention.dense.weight,
+                                               model.post_attention_layernorm.weight,
+                                               model.mlp.dense_h_to_4h.weight,
+                                               model.mlp.dense_4h_to_h.weight,
+                                               self_attn_x,
+                                               cross_attn_x};
+
+        for (auto tensor : all_tensors) {
+            random_fill(tensor);
+            tensor_to_device(tensor);
+        }
+
+        // self attention
+        reset_cgraph();
+        {
+            ggml_tensor *self_attn_y = model.forward(&ctx, self_attn_x, 0);
+            ggml_build_forward_expand(&ctx.gf, self_attn_y);
+            std::cout << "[Benchmark] GLM2Block " << ggml_type_name(dtype) << " self attn time: " << perf_compute()
+                      << " ms\n";
+        }
+
+        // cross attention
+        reset_cgraph();
+        {
+            ggml_tensor *cross_attn_y = model.forward(&ctx, cross_attn_x, seq_len);
+            ggml_build_forward_expand(&ctx.gf, cross_attn_y);
+            std::cout << "[Benchmark] GLM2Block " << ggml_type_name(dtype) << " cross attn time: " << perf_compute()
+                      << " ms\n";
+        }
+
+        for (auto tensor : all_tensors) {
+            tensor_to_cpu(tensor);
+        }
+        tensor_to_cpu(model.attention.k_cache);
+        tensor_to_cpu(model.attention.v_cache);
     }
 }
 
@@ -417,18 +827,6 @@ TEST_F(ChatGLMTest, quantize) {
     }
 }
 
-static bool equal(const std::vector<int> &a, const std::vector<int> &b) {
-    if (a.size() != b.size()) {
-        return false;
-    }
-    for (size_t i = 0; i < a.size(); i++) {
-        if (a[i] != b[i]) {
-            return false;
-        }
-    }
-    return true;
-}
-
 TEST(Pipeline, ChatGLM) {
     fs::path model_path = fs::path(__FILE__).parent_path() / "chatglm-ggml.bin";
     if (!fs::exists(model_path)) {
@@ -588,10 +986,7 @@ static void run_benchmark(const fs::path &model_path) {
 
     GenerationConfig gen_config;
     gen_config.do_sample = false;
-    char *num_threads_env = getenv("CHATGLM_NUM_THREADS");
-    if (num_threads_env) {
-        gen_config.num_threads = std::stoi(num_threads_env);
-    }
+    gen_config.num_threads = get_num_threads_for_benchmark();
 
     PerfStreamer streamer;
     start_ms = ggml_time_ms();
diff --git a/tests/test_convert.py b/tests/test_convert.py
index e8be3ad..3502d78 100644
--- a/tests/test_convert.py
+++ b/tests/test_convert.py
@@ -366,13 +366,13 @@ def make_data_glm2_block():
 
 
 def main():
-    sys.path.append(str(CHATGLM2_MODEL_PATH))
+    sys.path.append(str(CHATGLM_MODEL_PATH))
     torch.manual_seed(0)
     (HERE / "data").mkdir(parents=True, exist_ok=True)
     # make_data_linear()
-    make_data_layernorm()
+    # make_data_layernorm()
     # make_data_rms_norm()
-    # make_data_glm_block()
+    make_data_glm_block()
     # make_data_glm2_block()