Repository: ggml-org/llama.cpp
Pull request: [CANN] Optimize CANN buffer pool memory management

### Overview

Multiple optional memory pools are provided for CANN, including VMM, priority queue-based, and traditional memory pools.
1. When the memory pool is available and GGML_CANN_DISABLE_VMM_POOL is not defined, the VMM pool is selected by default.
2. Otherwise, if GGML_CANN_ENABLE_BUF_PRIO_POOL is defined, the priority queue-based memory pool is used.
3. If neither condition is met, the default memory pool is used.

> This improvement was implemented by [Frank Mai](https://github.com/thxCode) at first, from his repo: [llama-box](https://github.com/gpustack/llama-box/blob/main/llama-box/patches/llama.cpp), and I applied the `cann.patch` to `llama.cpp`.

### Environment

- **OS:** ubuntu 22.04
- **NPU:** 910B3
- **CANN:** 8.0.RC3


### Benchmark

We use model **qwen2.5:0.5b-instruct-fp16.gguf** for our benchmark.

`./llama.cpp/build/bin/llama-cli -m ds/qwen2.5\:0.5b-instruct-fp16 -p "Building a website can be done in 10 steps:" -ngl 32`

**Use VMM pool**

If no environment variables are set, VMM pool is used by default.

```
Building a website can indeed be a complex process that involves several stages. Here’s a step-by-step breakdown of what you might consider in the 10-step process:

1. **Define Your Goals and Objectives**: Clearly define what you want to achieve with your website. Are you looking to sell products, offer services, or promote a brand? Knowing your goals will guide your website design and development.

2. **Choose a Web Development Platform**: Select the platform or framework you will use to build your website. Popular choices include WordPress, Wix, Squarespace, and Netlify.

3. **Create a Website Structure**: Decide on the layout, navigation, and content of your website. Decide on the hierarchy of your site to ensure a clear user flow.

4. **Design Your Website**: Use graphic design software to create a visually appealing design. Consider themes, color schemes, and typography to enhance the user experience.

5. **Choose a Theme or Template**: Select a template or theme that fits your design and your goals. This will help you create a cohesive and professional website.

6. **Develop the Website**: Start coding your website from scratch, or use a pre-built template. This involves coding the HTML, CSS, and JavaScript to create the user interface.

7. **Test Your Website**: Use tools like browser developer tools to test your website for compatibility, responsiveness, and security. This will help you identify and fix any issues before launch.

8. **Optimize Your Website**: Consider the SEO aspects of your website, including optimizing for search engines and ensuring your website is mobile-friendly.

9. **Publish Your Website**: Once your website is tested and optimized, you can publish it on a website hosting service like WordPress.com, Squarespace.com, or Wix.com.

10. **Maintain and Update Your Website**: Continuously monitor your website for updates and improvements based on user feedback and changing market demands.

Each step can be challenging and requires careful planning and attention to detail. It's important to start with a clear vision, choose a platform that suits your goals, and follow through with consistent quality control and maintenance.

>
llama_perf_sampler_print:    sampling time =     846.24 ms /   473 runs   (    1.79 ms per token,   558.94 tokens per second)
llama_perf_context_print:        load time =    4304.85 ms
llama_perf_context_print: prompt eval time =      22.32 ms /    41 tokens (    0.54 ms per token,  1837.25 tokens per second)
llama_perf_context_print:        eval time =    7556.07 ms /   431 runs   (   17.53 ms per token,    57.04 tokens per second)
llama_perf_context_print:       total time =  101961.78 ms /   472 tokens
```

**Use BUF_PRIO_POOL**  

If set `export GGML_CANN_DISABLE_VMM_POOL=1 && export GGML_CANN_ENABLE_BUF_PRIO_POOL=1`, priority queue-based memory pool will be enabled.

```
Building a website can be done in 10 steps:
assistant
Certainly! Building a website is a comprehensive process that typically involves several steps:

1. **Planning and Research**: Define the purpose and target audience of your website. Gather information on your competitors, the platforms on which they operate, and their features.

2. **Design**: Choose a design system or choose your own design. Decide on the color scheme, typography, layout, and overall design of your website.

3. **Development**: Start coding the website's structure and content. This involves HTML, CSS (for styling), and JavaScript (for interactivity).

4. **Deployment**: Decide on where you will store your files (e.g., on a server, in a cloud service, or locally). Choose a deployment method, such as static site generation (SSG) for static websites or server-side rendering (SSR) for dynamic websites.

5. **Testing**: Test your website across different browsers and devices to ensure compatibility and functionality. Use tools like browser testers or a developer’s testing suite.

6. **Launch**: Once testing is complete, deploy your website to the internet. Make sure your server can handle the traffic and that you have backups.

7. **Maintenance**: Regularly update your website with new content, improve its performance, and fix any bugs or issues that arise.

8. **SEO Optimization**: Enhance your website’s visibility in search engine results. Use relevant keywords, meta descriptions, and titles, and consider using off-page SEO strategies like backlinking.

9. **Security**: Ensure that your website is secure. This includes protecting user data, preventing hacking, and maintaining compliance with data privacy laws.

10. **Analytics and Feedback**: Use tools like Google Analytics to track website performance. Collect feedback from users and incorporate it into future updates.

Building a website is a collaborative process involving many stakeholders, so it's crucial to communicate effectively with them throughout the development and deployment phases.

>
llama_perf_sampler_print:    sampling time =     750.06 ms /   424 runs   (    1.77 ms per token,   565.29 tokens per second)
llama_perf_context_print:        load time =    4422.03 ms
llama_perf_context_print: prompt eval time =      30.95 ms /    41 tokens (    0.75 ms per token,  1324.93 tokens per second)
llama_perf_context_print:        eval time =    9070.38 ms /   382 runs   (   23.74 ms per token,    42.12 tokens per second)
llama_perf_context_print:       total time =   99563.38 ms /   423 tokens
```

**BUF_POOL**  

If set `export GGML_CANN_DISABLE_VMM_POOL=1`, traditional buffer pools will be enabled.

```
Building a website can indeed be a complex process, but I'll break down the 10 steps into simpler, more manageable parts:

1. **Define the Purpose and Audience**: Start by defining what your website is for and who your target audience is. This will help you determine the content and design elements you need.

2. **Choose a Platform**: Select the platform or framework you'll use to build your website. Popular options include WordPress, Joomla, WordPress with built-in themes, or custom development.

3. **Design the Website**: Use the platform to design your website. This step is crucial for ensuring that your website is visually appealing and user-friendly.

4. **Create Content**: Write content that is relevant and informative. This could include blog posts, infographics, videos, or other forms of media.

5. **Choose a Theme or Template**: Select a theme or template that matches your brand identity and helps guide user experience.

6. **Implement Search Engine Optimization (SEO)**: Use SEO tools and practices to improve your website's search engine ranking.

7. **Launch the Website**: Once everything is set up, launch your website and start promoting it to attract traffic.

8. **Maintain and Update the Site**: Regularly update your website with new content and improve it based on user feedback and analytics.

9. **SEO Optimization**: Continue to optimize your website for search engines to increase visibility and organic traffic.

10. **User Experience (UX)**: Focus on creating a seamless and user-friendly experience for visitors.

Each step builds upon the previous one, so don't get discouraged if you hit roadblocks. A good website is not just a physical space, but a living, breathing entity that provides value to its users.

This process, while complex, can be broken down into manageable steps, making it easier to start and maintain a website.

>
llama_perf_sampler_print:    sampling time =     719.61 ms /   416 runs   (    1.73 ms per token,   578.09 tokens per second)
llama_perf_context_print:        load time =    4279.12 ms
llama_perf_context_print: prompt eval time =      21.54 ms /    41 tokens (    0.53 ms per token,  1903.26 tokens per second)
llama_perf_context_print:        eval time =    6469.97 ms /   374 runs   (   17.30 ms per token,    57.81 tokens per second)
llama_perf_context_print:       total time =  266109.47 ms /   415 tokens
```



Changed files:
filename: ggml/src/ggml-cann/aclnn_ops.cpp
patch: @@ -1783,7 +1783,7 @@ void ggml_cann_get_rows(ggml_backend_cann_context& ctx, ggml_tensor* dst) {
                 src0->data, ACL_INT8, sizeof(int8_t), weight_ne, weight_nb,
                 GGML_MAX_DIMS + 1);
             aclTensor* acl_scale_tensor = ggml_cann_create_tensor(
-                src0->data, ACL_FLOAT16, sizeof(float16_t), scale_ne, scale_nb,
+                src0->data, ACL_FLOAT16, sizeof(uint16_t), scale_ne, scale_nb,
                 GGML_MAX_DIMS + 1, ACL_FORMAT_ND, scale_offset);
             aclTensor* dequant_tensor = ggml_cann_create_tensor(
                 dequant_buffer_allocator.get(), ACL_FLOAT, sizeof(float_t),

filename: ggml/src/ggml-cann/ggml-cann.cpp
patch: @@ -29,6 +29,8 @@
 #include <cstdio>
 #include <cstring>
 #include <mutex>
+#include <queue>
+#include <chrono>
 
 #include "ggml-impl.h"
 #include "ggml-backend-impl.h"
@@ -119,9 +121,10 @@ static ggml_cann_device_info ggml_cann_init() {
         prop.location.type = ACL_MEM_LOCATION_TYPE_DEVICE;
         prop.location.id = id;
         prop.reserve = 0;
-        ACL_CHECK(aclrtMemGetAllocationGranularity(
+        err = aclrtMemGetAllocationGranularity(
             &prop, ACL_RT_MEM_ALLOC_GRANULARITY_RECOMMENDED,
-            &info.devices[id].vmm_granularity));
+            &info.devices[id].vmm_granularity);
+        info.devices[id].vmm = err == ACL_SUCCESS;
 
         size_t free, total;
         ggml_backend_cann_get_device_memory(id, &free, &total);
@@ -148,11 +151,222 @@ const ggml_cann_device_info& ggml_cann_info() {
 
 //#define DEBUG_CANN_MALLOC
 /**
- * @brief A pool of CANN buffers(legacy).
+ * @brief A pool of CANN buffers(priority segment buffer).
  *
  * This class manages a pool of CANN buffers for a specific device.
  */
-struct ggml_cann_pool_leg : public ggml_cann_pool {
+struct ggml_cann_pool_buf_prio : public ggml_cann_pool {
+        /**
+         * @brief The maximum reuse margin for a buffer.
+         */
+        static const size_t max_reuse_margin = 1ull << 22;  // 4MB
+
+        /**
+         * @brief The minimum free margin for a buffer.
+         */
+        static const size_t min_free_margin = 1ull << 20;   // 1MB
+
+        /**
+         * @brief The alignment for buffer allocation.
+         */
+        static const size_t alignment = 128;
+
+        /**
+         * @brief The device ID associated with this buffer pool.
+         */
+        int device;
+
+        /**
+         * @brief Whether to disable clean during buffer allocation.
+         */
+        bool disable_clean = false;
+
+        /**
+         * @brief Structure representing a CANN buffer.
+         */
+        struct ggml_cann_buffer {
+            void* ptr = nullptr;  ///< Pointer to the buffer.
+            size_t size = 0;      ///< Size of the buffer.
+            std::chrono::steady_clock::time_point last_used;  ///< Last used time.
+
+            bool operator>(const ggml_cann_buffer& other) const {
+                return size > other.size;
+            }
+        };
+
+        /**
+         * @brief Array of CANN buffers in the pool.
+         */
+        std::unordered_map<void*, size_t> buffer_pool;
+        std::priority_queue<ggml_cann_buffer,
+                            std::vector<ggml_cann_buffer>,
+                            std::greater<>> free_buffers ;
+
+        /**
+         * @brief Total size of all buffers in the pool.
+         */
+        size_t pool_size = 0;
+
+        /**
+         * @brief Constructor to initialize the buffer pool for a specific device.
+         *
+         * @param device The device ID to associate with this buffer pool.
+         */
+        explicit ggml_cann_pool_buf_prio(int device) : device(device) {
+            disable_clean = getenv("GGML_CANN_DISABLE_BUF_POOL_CLEAN") != nullptr;
+        }
+
+        /**
+         * @brief Destructor to free all buffers in the pool.
+         */
+        ~ggml_cann_pool_buf_prio() {
+            ggml_cann_set_device(device);
+            for (auto& [b_ptr, b_size] : buffer_pool) {
+                aclrtFree(b_ptr);
+               pool_size -= b_size;
+            }
+            buffer_pool.clear();
+            GGML_ASSERT(pool_size == 0);
+        }
+
+        /**
+         * @brief Allocate a buffer of the given size.
+         *
+         * @param size The size of the buffer to allocate.
+         * @param actual_size A pointer to a variable to receive the actual size of
+         * the allocated buffer.
+         * @return A pointer to the allocated buffer.
+         */
+        void* alloc(size_t size, size_t* actual_size) override {
+            size = GGML_PAD(size, alignment);
+            if (size == 0) {
+                size = alignment;
+            }
+
+            void* ptr = nullptr;
+            auto now = std::chrono::steady_clock::now();
+
+            std::vector<ggml_cann_buffer> free_buffers_rest;
+            free_buffers_rest.reserve(free_buffers.size());
+            while (!free_buffers.empty()) {
+                auto b = free_buffers.top();
+                free_buffers.pop();
+
+                if (b.size >= size) {
+                    // reuse the buffer if the size is enough
+                    const size_t margin = b.size - size;
+                    if (margin <= max_reuse_margin) {
+                        *actual_size = b.size;
+                        ptr = b.ptr;
+    #ifdef DEBUG_CANN_MALLOC
+                        GGML_LOG_INFO(
+                            "cann pool[%d]: reused   %p, "
+                            "pool_size = %5u MB, "
+                            "size = %5u MB, "
+                            "margin = %5u MB\n",
+                            device, b.ptr,
+                            (uint32_t)(GGML_PAD(pool_size, 1048576) / 1048576),
+                            (uint32_t)(GGML_PAD(size, 1048576) / 1048576),
+                            (uint32_t)(GGML_PAD(margin, 1048576) / 1048576));
+    #endif
+                        break;
+                    }
+                }
+
+                bool should_clean = !disable_clean &&
+                                   b.size > min_free_margin &&
+                                   std::chrono::duration_cast<std::chrono::milliseconds>(now - b.last_used).count() > 100;
+                if (should_clean) {
+                    // free the buffer if the size is needed to be freed
+                    ACL_CHECK(aclrtFree(b.ptr));
+                    pool_size -= b.size;
+                    buffer_pool.erase(b.ptr);
+    #ifdef DEBUG_CANN_MALLOC
+                    GGML_LOG_INFO(
+                        "cann pool[%d]: clean    %p, "
+                        "pool_size = %5u MB, "
+                        "size = %5u MB\n",
+                        device, b.ptr,
+                        (uint32_t)(GGML_PAD(pool_size, 1048576) / 1048576),
+                        (uint32_t)(GGML_PAD(b.size, 1048576) / 1048576));
+    #endif
+                    continue;
+                }
+                free_buffers_rest.push_back(b);
+            }
+            for (ggml_cann_buffer &b : free_buffers_rest) {
+                free_buffers.push(std::move(b));
+            }
+
+    #ifdef DEBUG_CANN_MALLOC
+            GGML_LOG_INFO("cann pool[%d] free pool_size = %5u MB\n\n", device, (uint32_t)(GGML_PAD(pool_size, 1048576) / 1048576));
+    #endif
+            if (ptr != nullptr) {
+                return ptr;
+            }
+
+            // allocate a new buffer if no buffer can be reused
+            ggml_cann_set_device(device);
+            ACL_CHECK(aclrtMalloc(&ptr, size, ACL_MEM_MALLOC_HUGE_FIRST));
+            *actual_size = size;
+            pool_size += size;
+    #ifdef DEBUG_CANN_MALLOC
+            GGML_LOG_INFO(
+                "cann pool[%d]: allocate %p, "
+                "pool_size = %5u MB, "
+                "size = %5u MB\n",
+                device, ptr, (uint32_t)(GGML_PAD(pool_size, 1048576) / 1048576),
+                (uint32_t)(GGML_PAD(size, 1048576) / 1048576));
+    #endif
+            buffer_pool.emplace(ptr, size);
+            return ptr;
+        }
+
+        /**
+         * @brief Free a buffer and return it to the pool.
+         *
+         * @param ptr Pointer to the buffer to free.
+         * @param size Size of the buffer to free.
+         */
+        void free(void* ptr, size_t size) override {
+            auto it = buffer_pool.find(ptr);
+            if (it == buffer_pool.end()) {
+                GGML_ABORT("cann pool[%d]: buffer %p not found in pool\n", device, ptr);
+            }
+
+            auto now = std::chrono::steady_clock::now();
+            free_buffers.emplace(ggml_cann_buffer{ptr, it->second, now});
+    #ifdef DEBUG_CANN_MALLOC
+            GGML_LOG_INFO(
+                "cann pool[%d]: return   %p, "
+                "pool_size = %5u MB\n",
+                device, ptr,
+                (uint32_t)(GGML_PAD(pool_size, 1048576) / 1048576));
+    #endif
+        }
+    };
+
+/**
+ * @brief A pool of CANN buffers(segment buffer).
+ *
+ * This class manages a pool of CANN buffers for a specific device.
+ */
+struct ggml_cann_pool_buf : public ggml_cann_pool {
+    /**
+     * @brief The maximum reuse margin for a buffer.
+     */
+    static const size_t max_reuse_margin = 1ull << 22;  // 4MB
+
+    /**
+     * @brief The minimum free margin for a buffer.
+     */
+    static const size_t min_free_margin = 1ull << 20;   // 1MB
+
+    /**
+     * @brief The alignment for buffer allocation.
+     */
+    static const size_t alignment = 128;
+
     /**
      * @brief The maximum number of buffers in the pool.
      */
@@ -163,12 +377,19 @@ struct ggml_cann_pool_leg : public ggml_cann_pool {
      */
     int device;
 
+    /**
+     * @brief Whether to disable clean during buffer allocation.
+     */
+    bool disable_clean = false;
+
     /**
      * @brief Structure representing a CANN buffer.
      */
     struct ggml_cann_buffer {
         void* ptr = nullptr;  ///< Pointer to the buffer memory.
         size_t size = 0;      ///< Size of the buffer.
+        bool used = false;    ///< Whether the buffer is currently in use.
+        std::chrono::steady_clock::time_point last_used;  ///< Last used time.
     };
 
     /**
@@ -186,17 +407,19 @@ struct ggml_cann_pool_leg : public ggml_cann_pool {
      *
      * @param device The device ID to associate with this buffer pool.
      */
-    explicit ggml_cann_pool_leg(int device) : device(device) {}
+    explicit ggml_cann_pool_buf(int device) : device(device) {
+        disable_clean = getenv("GGML_CANN_DISABLE_BUF_POOL_CLEAN") != nullptr;
+    }
 
     /**
      * @brief Destructor to free all buffers in the pool.
      */
-    ~ggml_cann_pool_leg() {
+    ~ggml_cann_pool_buf() {
         ggml_cann_set_device(device);
         for (int i = 0; i < MAX_BUFFERS; ++i) {
             ggml_cann_buffer& b = buffer_pool[i];
             if (b.ptr != nullptr) {
-                ACL_CHECK(aclrtFree(b.ptr));
+                aclrtFree(b.ptr);
                 pool_size -= b.size;
             }
         }
@@ -212,63 +435,93 @@ struct ggml_cann_pool_leg : public ggml_cann_pool {
      * @return A pointer to the allocated buffer.
      */
     void* alloc(size_t size, size_t* actual_size) override {
-        const size_t alignment = 128;
         size = GGML_PAD(size, alignment);
         if (size == 0) {
             size = alignment;
         }
-#ifdef DEBUG_CANN_MALLOC
-        int nnz = 0;
-        size_t max_size = 0;
-#endif
-        size_t best_diff = 1ull << 36;
-        int ibest = -1;
-        for (int i = 0; i < MAX_BUFFERS; ++i) {
+
+        void* ptr = nullptr;
+        auto now = std::chrono::steady_clock::now();
+
+        int i = 0;
+        for (; i < MAX_BUFFERS; ++i) {
             ggml_cann_buffer& b = buffer_pool[i];
-            if (b.ptr != nullptr) {
+            if (b.ptr == nullptr) {
+                break;
+            }
+            if (b.used) {
+                continue;
+            }
+            if (b.size >= size) {
+                // reuse the buffer if the size is enough
+                const size_t margin = b.size - size;
+                if (margin <= max_reuse_margin) {
+                    *actual_size = b.size;
+                    b.used = true;
+                    ptr = b.ptr;
 #ifdef DEBUG_CANN_MALLOC
-                ++nnz;
-                if (b.size > max_size) max_size = b.size;
+                    GGML_LOG_INFO(
+                        "cann pool[%d]: reused   %p, "
+                        "pool_size = %5u MB, "
+                        "size = %5u MB, "
+                        "margin = %5u MB\n",
+                        device, b.ptr,
+                        (uint32_t)(GGML_PAD(pool_size, 1048576) / 1048576),
+                        (uint32_t)(GGML_PAD(size, 1048576) / 1048576),
+                        (uint32_t)(GGML_PAD(margin, 1048576) / 1048576));
 #endif
-                if (b.size >= size) {
-                    size_t diff = b.size - size;
-                    if (diff < best_diff) {
-                        best_diff = diff;
-                        ibest = i;
-                        if (!best_diff) {
-                            void* ptr = b.ptr;
-                            *actual_size = b.size;
-                            b.ptr = nullptr;
-                            b.size = 0;
-                            return ptr;
-                        }
-                    }
+                    break;
                 }
             }
+
+            bool should_clean = !disable_clean &&
+                                b.size > min_free_margin &&
+                                std::chrono::duration_cast<std::chrono::milliseconds>(now - b.last_used).count() > 100;
+            if (should_clean) {
+                // free the buffer if the size is needed to be freed
+                ACL_CHECK(aclrtFree(b.ptr));
+                pool_size -= b.size;
+#ifdef DEBUG_CANN_MALLOC
+                GGML_LOG_INFO(
+                    "cann pool[%d]: clean    %p, "
+                    "pool_size = %5u MB, "
+                    "size = %5u MB\n",
+                    device, b.ptr,
+                    (uint32_t)(GGML_PAD(pool_size, 1048576) / 1048576),
+                    (uint32_t)(GGML_PAD(b.size, 1048576) / 1048576));
+#endif
+                b.ptr = nullptr;
+            }
         }
-        if (ibest >= 0) {
-            ggml_cann_buffer& b = buffer_pool[ibest];
-            void* ptr = b.ptr;
-            *actual_size = b.size;
-            b.ptr = nullptr;
-            b.size = 0;
+        if (ptr != nullptr) {
             return ptr;
         }
-        void* ptr;
-        ggml_cann_set_device(device);
-        ACL_CHECK(
-            aclrtMalloc(&ptr, size, ACL_MEM_MALLOC_HUGE_FIRST));
-        *actual_size = size;
-        pool_size += size;
+
+        if (i < MAX_BUFFERS) {
+            // allocate a new buffer if no buffer can be reused
+            ggml_cann_buffer& b = buffer_pool[i];
+            ggml_cann_set_device(device);
+            ACL_CHECK(aclrtMalloc(&b.ptr, size, ACL_MEM_MALLOC_HUGE_FIRST));
+            pool_size += size;
+            *actual_size = size;
+            b.size = size;
+            b.used = true;
+            if (i >= MAX_BUFFERS - 8) {
+                GGML_LOG_WARN("cann pool[%d]: slots almost full\n", device);
+            }
 #ifdef DEBUG_CANN_MALLOC
-        GGML_LOG_INFO(
-            "%s[%d]: %d buffers, max_size = %u MB, pool_size = %u MB, "
-            "requested %u MB\n",
-            __func__, device, nnz, (uint32_t)(max_size / 1024 / 1024),
-            (uint32_t)(pool_size / 1024 / 1024),
-            (uint32_t)(size / 1024 / 1024));
+            GGML_LOG_INFO(
+                "cann pool[%d]: allocate %p, "
+                "pool_size = %5u MB, "
+                "size = %5u MB\n",
+                device, b.ptr,
+                (uint32_t)(GGML_PAD(pool_size, 1048576) / 1048576),
+                (uint32_t)(GGML_PAD(b.size, 1048576) / 1048576));
 #endif
-        return ptr;
+            return b.ptr;
+        }
+
+        GGML_ABORT("cann pool[%d]: slots full\n", device);
     }
 
     /**
@@ -280,16 +533,21 @@ struct ggml_cann_pool_leg : public ggml_cann_pool {
     void free(void* ptr, size_t size) override {
         for (int i = 0; i < MAX_BUFFERS; ++i) {
             ggml_cann_buffer& b = buffer_pool[i];
-            if (b.ptr == nullptr) {
-                b.ptr = ptr;
-                b.size = size;
-                return;
+            if (b.ptr != ptr) {
+                continue;
             }
+            b.used = false;
+            b.last_used = std::chrono::steady_clock::now();
+#ifdef DEBUG_CANN_MALLOC
+            GGML_LOG_INFO(
+                "cann pool[%d]: return   %p, "
+                "pool_size = %5u MB\n",
+                device, b.ptr,
+                (uint32_t)(GGML_PAD(pool_size, 1048576) / 1048576));
+#endif
+            return;
         }
-        // memory should always buffered. these memory may still needed by
-        // tasks in stream.
-        // TODO, fix me.
-        GGML_ABORT("Cann buffer pool full, increase MAX_CANN_BUFFERS\n");
+        GGML_ABORT("cann pool[%d]: slots full\n", device);
     }
 };
 
@@ -347,8 +605,7 @@ struct ggml_cann_pool_vmm : public ggml_cann_pool {
      * @param device The device ID to associate with this buffer pool.
      */
     explicit ggml_cann_pool_vmm(int device)
-        : device(device),
-          granularity(ggml_cann_info().devices[device].vmm_granularity) {
+    : device(device) {
         auto dev = ggml_cann_info().devices[device];
         granularity = dev.vmm_granularity;
         max_size = dev.total_vram;
@@ -471,7 +728,18 @@ struct ggml_cann_pool_vmm : public ggml_cann_pool {
  */
 std::unique_ptr<ggml_cann_pool> ggml_backend_cann_context::new_pool_for_device(
     int device) {
-    return std::unique_ptr<ggml_cann_pool>(new ggml_cann_pool_vmm(device));
+    bool disable_vmm = (getenv("GGML_CANN_DISABLE_VMM_POOL") != nullptr);
+    if (!disable_vmm && ggml_cann_info().devices[device].vmm) {
+        GGML_LOG_INFO("%s: device %d use vmm pool\n", __func__, device);
+        return std::unique_ptr<ggml_cann_pool>(new ggml_cann_pool_vmm(device));
+    }
+    bool enable_buf_prio = (getenv("GGML_CANN_ENABLE_BUF_PRIO_POOL") != nullptr);
+    if (enable_buf_prio) {
+        GGML_LOG_INFO("%s: device %d use buffer pool with priority queue\n", __func__, device);
+        return std::unique_ptr<ggml_cann_pool>(new ggml_cann_pool_buf_prio(device));
+    }
+    GGML_LOG_INFO("%s: device %d use buffer pool\n", __func__, device);
+    return std::unique_ptr<ggml_cann_pool>(new ggml_cann_pool_buf(device));
 }
 
 // cann buffer
@@ -1020,8 +1288,11 @@ ggml_backend_cann_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft,
 
     ggml_cann_set_device(buft_ctx->device);
 
-    size = std::max(size, (size_t)1);
-
+    const size_t alignment = 128;
+    size = GGML_PAD(size, alignment);
+    if (size == 0) {
+        size = alignment;
+    }
     void* dev_ptr;
     aclError err = aclrtMalloc(&dev_ptr, size, ACL_MEM_MALLOC_HUGE_FIRST);
     if (err != ACL_SUCCESS) {