Update attention_unfused required memory size

Author: andrewchernyh

csrc/transformer/inference/csrc/pt_binding.cpp

@@ -354,10 +354,7 @@ void attention_unfused(T* prev_key_cont,
     float layer_scale = alibi.sizes().size() > 1 ? std::max(1, layer_id) : 1.0;
     float alpha = norm_factor * norm_factor / layer_scale;
     float gemm_beta = 0.0;
-    // Always use the tail workspace
-    T* scratch = (T*)Context::Instance().GetWorkSpace();
-    T *workspace = scratch + ((Context::Instance().get_workspace_size() / sizeof(T)) -
-                            bsz * heads * seq_len * soft_len);
+    T* workspace = (T*)Context::Instance().GetAttentionUnfusedWorkspace();
 
     cublasSetStream(Context::Instance().GetCublasHandle(), Context::Instance().GetCurrentStream());
     cublas_strided_batched_gemm(Context::Instance().GetCublasHandle(),

csrc/transformer/inference/includes/inference_context.h

@@ -55,7 +55,8 @@ class Context {
           _curr_offset(0),
           _stream(0),
           _free_memory_size(0),
-          _num_tokens(1)
+          _num_tokens(1),
+          _attention_unfused_workspace_offset(0)
     {
         if (cublasCreate(&_cublasHandle) != CUBLAS_STATUS_SUCCESS) {
             auto message = std::string("Fail to create cublas handle.");
@@ -101,7 +102,7 @@ class Context {
         if (!_free_memory_size) { cudaMemGetInfo(&_free_memory_size, &total_size); }
 
         size_t activation_size = 16 * hidden_dim * batch_size;
-        size_t temp_size = batch_size * num_heads * prompt_len * prompt_len * elem_size;
+        size_t temp_size = batch_size * num_heads * prompt_len * prompt_len * elem_size / mp_size;
         size_t cache_size = num_layers * batch_size * (hidden_dim / mp_size) * 2;
         size_t minimal_requirements = temp_size + (_free_memory_size > GIGABYTE ? 500 : 100) * MEGABYTE;
         if (_free_memory_size < minimal_requirements) {
@@ -139,13 +140,15 @@ class Context {
             throw std::runtime_error("Workspace is null.");
         }
         _workSpaceSize = workSpaceSize;
+        _attention_unfused_workspace_offset = workSpaceSize - temp_size;
     }
     inline size_t GetMaxTokenLenght() const { return _max_seq_len; }
 
     cudaEvent_t GetCompEvent(int id) { return id == 1 ? _comp1_event : _comp2_event; }
 
     size_t get_workspace_size() const { return _workSpaceSize; }
     void* GetWorkSpace() { return _workspace; }
+    void* GetAttentionUnfusedWorkspace() { return _workspace + _attention_unfused_workspace_offset; }
 
     inline unsigned new_token(unsigned layer_id)
     {
@@ -211,6 +214,8 @@ class Context {
     cudaEvent_t _comm_event;
 
     void* _workspace;
+    // offset from _workspace for uttention unfused memory
+    size_t _attention_unfused_workspace_offset;
     uint64_t _seed;
     uint64_t _curr_offset;