From b5d6284190d83cc0be11a691e1045d88510b657a Mon Sep 17 00:00:00 2001
From: Concedo <39025047+LostRuins@users.noreply.github.com>
Date: Sun, 23 Apr 2023 00:07:33 +0800
Subject: [PATCH] increase initial buffer too

---
 otherarch/gpt2_v2.cpp | 2 +-
 otherarch/gptj_v2.cpp | 2 +-
 otherarch/neox.cpp    | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/otherarch/gpt2_v2.cpp b/otherarch/gpt2_v2.cpp
index 151e9dd4a..f86750656 100644
--- a/otherarch/gpt2_v2.cpp
+++ b/otherarch/gpt2_v2.cpp
@@ -371,8 +371,8 @@ bool gpt2_eval(
     const int n_vocab = hparams.n_vocab;
 
     //todo: there is a bug that causes the buffer to oom and I cannot figure it out, hack to increase size for now  
-    static size_t buf_size = 256u*1024*1024;
     const size_t extra_buf = 64u*1024*1024;
+    static size_t buf_size = 256u*1024*1024 + extra_buf;    
     static void * buf = malloc(buf_size);
 
     if (mem_per_token > 0 && mem_per_token*N*1.6 > buf_size) {
diff --git a/otherarch/gptj_v2.cpp b/otherarch/gptj_v2.cpp
index 7d548081c..a81053e62 100644
--- a/otherarch/gptj_v2.cpp
+++ b/otherarch/gptj_v2.cpp
@@ -382,8 +382,8 @@ bool gptj_eval(
     const int d_key = n_embd/n_head;
 
     //todo: there is a bug that causes the buffer to oom and I cannot figure it out, hack to increase size for now  
-    static size_t buf_size = 256u*1024*1024;
     const size_t extra_buf = 64u*1024*1024;
+    static size_t buf_size = 256u*1024*1024 + extra_buf;    
     static void * buf = malloc(buf_size);
 
     if (mem_per_token > 0 && mem_per_token*N*1.5 > buf_size) {
diff --git a/otherarch/neox.cpp b/otherarch/neox.cpp
index 5b9656aa5..51450c620 100644
--- a/otherarch/neox.cpp
+++ b/otherarch/neox.cpp
@@ -364,8 +364,8 @@ bool stablelm_eval(
     const int n_vocab = hparams.n_vocab;
     const int n_rot   = hparams.n_rot;
 
-    static size_t buf_size = 256u*1024*1024;
     const size_t extra_buf = 64u*1024*1024;
+    static size_t buf_size = 256u*1024*1024 + extra_buf;    
     static void * buf = malloc(buf_size);
 
     if (mem_per_token > 0 && mem_per_token*N*1.5 > buf_size) {