issue/407 - Release the GIL; preallocated workspace

Author: wangpengcheng

csrc/engine/infer_engine.cpp

@@ -14,15 +14,16 @@ InferEngine::InferEngine(
     const cache::CacheConfig *cache_config,
     bool enable_graph_compiling,
     backends::AttentionBackend attention_backend,
-    std::optional<infinicore::DataType> kv_cache_dtype) // Changed parameter
+    std::optional<infinicore::DataType> kv_cache_dtype, // Changed parameter
+    size_t max_num_batched_tokens)
     : communication_group_(distributed_config, device_type), attention_backend_(attention_backend) {
     if (cache_config != nullptr) {
         cache_config_ = cache_config->unique_copy();
     }
 
     // Load model config if model_path is provided, model_path must be valid, and config.json exists
     this->model_config_ = infinilm::config::ConfigFactory::createConfig(config_str);
-    auto infinilm_config = std::make_shared<infinilm::global_state::InfinilmConfig>(attention_backend, this->model_config_);
+    auto infinilm_config = std::make_shared<infinilm::global_state::InfinilmConfig>(attention_backend, this->model_config_, max_num_batched_tokens);
 
     // Only support offline int8 kv cache quantization in this version
     if (kv_cache_dtype.has_value()) {

csrc/engine/infer_engine.hpp

@@ -27,7 +27,8 @@ class InferEngine {
         const cache::CacheConfig *cache_config = nullptr,
         bool enable_graph_compiling = false,
         backends::AttentionBackend attention_backend = backends::AttentionBackend::Default,
-        std::optional<infinicore::DataType> kv_cache_dtype = std::nullopt);
+        std::optional<infinicore::DataType> kv_cache_dtype = std::nullopt,
+        size_t max_num_batched_tokens = 2048);
 
     // Load a parameter to all workers (each can extract its shard inside RankWorker)
     void load_param(const std::string &name, const infinicore::Tensor &param);

csrc/global_state/forward_context.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "../models/infinilm_model.hpp"
+#include <unordered_map>
 
 namespace infinilm::global_state {
 
@@ -43,6 +44,9 @@ struct AttentionMetadata {
 struct ForwardContext {
     AttentionMetadata attn_metadata;
     std::vector<infinicore::Tensor> kv_cache_vec;
+
+    // preallocated workspace for some modules
+    std::unordered_map<std::string, infinicore::Tensor> preallocated_workspace;
 };
 
 void initialize_forward_context(ForwardContext &forward_context);

csrc/global_state/infinilm_config.hpp

@@ -14,13 +14,19 @@ struct InfinilmConfig {
 public:
     InfinilmConfig() = default;
     InfinilmConfig(const infinilm::backends::AttentionBackend &backend,
-                   const std::shared_ptr<infinilm::config::ModelConfig> &model_config)
+                   const std::shared_ptr<infinilm::config::ModelConfig> &model_config,
+                   size_t max_num_batched_tokens)
         : attention_backend(backend),
-          model_config(model_config) {}
+          model_config(model_config),
+          max_num_batched_tokens(max_num_batched_tokens) {
+        const size_t max_position_embeddings = model_config->get<size_t>("max_position_embeddings");
+        ASSERT(max_num_batched_tokens >= 512 && max_num_batched_tokens <= max_position_embeddings);
+    }
 
 public:
     infinilm::backends::AttentionBackend attention_backend;
     std::shared_ptr<infinilm::config::ModelConfig> model_config;
+    size_t max_num_batched_tokens = 0;
 };
 
 /**

csrc/layers/attention/attention.cpp

@@ -1,17 +1,20 @@
 #include "attention.hpp"
+#include "../../global_state/global_state.hpp"
 #include "../../utils.hpp"
 #include "../rotary_embedding/rotary_embedding.hpp"
+#include <string>
 
 namespace infinilm::layers::attention {
 
 Attention::Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config,
                      size_t layer_idx,
-                     const infinicore::Device &device) {
+                     const infinicore::Device &device)
+    : device_(device),
+      dtype_(model_config->get_dtype()) {
     layer_idx_ = layer_idx;
     hidden_size_ = model_config->get<size_t>("hidden_size");
     head_dim_ = model_config->get<size_t>("head_dim");
 
-    const auto &dtype{model_config->get_dtype()};
     size_t total_num_heads = model_config->get<size_t>("num_attention_heads");
     size_t total_num_kv_heads = model_config->get<size_t>("num_key_value_heads");
     bool use_bias = model_config->get_or<bool>("attention_bias", true);
@@ -31,18 +34,21 @@ Attention::Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config
     qkv_proj_ = std::make_shared<layers::linear::QKVParallelLinear>(
         hidden_size_, head_dim_, total_num_heads, total_num_kv_heads,
         "q_proj", "k_proj", "v_proj", register_fn,
-        quantization_method, use_bias, dtype, device, rank_info);
+        quantization_method, use_bias, dtype_, device_, rank_info);
     o_proj_ = this->register_module<layers::linear::RowParallelLinear>(
         "o_proj", total_num_heads * head_dim_, hidden_size_, quantization_method,
-        use_output_bias, dtype, device, tp_rank, tp_size, rank_info.comm);
+        use_output_bias, dtype_, device_, tp_rank, tp_size, rank_info.comm);
 
-    rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device);
+    rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device_);
 
     float scaling = 1.0f / std::sqrt(static_cast<float>(head_dim_));
     attn_ = std::make_shared<AttentionLayer>(num_attention_heads_, head_dim_, scaling, num_key_value_heads_, layer_idx_,
-                                             kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_);
+                                             kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_, device_);
 
-    init_kv_cache_quant_params(register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_);
+    init_kv_cache_quant_params(register_fn, device_, kv_cache_k_scale_, kv_cache_v_scale_);
+
+    rank_qkv_output_size_ = qkv_proj_->out_features() / static_cast<size_t>(tp_size);
+    this->_initialize_preallocated_workspace();
 }
 
 infinicore::Tensor Attention::forward(const infinicore::Tensor &positions,
@@ -62,7 +68,8 @@ infinicore::Tensor Attention::forward_static_(const infinicore::Tensor &position
     size_t seq_len = shape[1];
 
     // 1. Project Q, K, V
-    auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable);
+    auto qkv_output = max_qkv_output_->narrow({{0, 0, batch_size * seq_len}})->view({batch_size, seq_len, rank_qkv_output_size_});
+    auto [q, k, v] = qkv_proj_->forward_split_(qkv_output, hidden_states_mutable);
 
     // 2. Reshape for multi-head attention
     auto q_reshaped = q->view({batch_size, seq_len, num_attention_heads_, head_dim_});
@@ -90,8 +97,9 @@ infinicore::Tensor Attention::forward_static_(const infinicore::Tensor &position
     auto attn_output = attn_->forward(q_rope, k_reshaped, v_reshaped);
 
     // 7. Project output
-    auto output = o_proj_->forward(attn_output);
-    return output;
+    auto o_output = max_o_output_->narrow({{0, 0, batch_size * seq_len}})->view({batch_size, seq_len, hidden_size_});
+    o_proj_->forward_(o_output, attn_output);
+    return o_output;
 }
 
 infinicore::Tensor Attention::forward_paged_(const infinicore::Tensor &position_ids,
@@ -106,7 +114,8 @@ infinicore::Tensor Attention::forward_paged_(const infinicore::Tensor &position_
     ASSERT_EQ(batch_size, 1);
 
     // 1. Project Q, K, V
-    auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable);
+    auto qkv_output = max_qkv_output_->narrow({{0, 0, seq_len}})->view({1, seq_len, rank_qkv_output_size_});
+    auto [q, k, v] = qkv_proj_->forward_split_(qkv_output, hidden_states_mutable);
 
     // 2. Reshape for multi-head attention
     auto q_reshaped = q->view({seq_len, num_attention_heads_, head_dim_});
@@ -133,8 +142,35 @@ infinicore::Tensor Attention::forward_paged_(const infinicore::Tensor &position_
     auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped);
 
     // 6. Project output
-    auto output = o_proj_->forward(attn_output);
-    return output;
+    auto o_output = max_o_output_->narrow({{0, 0, seq_len}})->view({1, seq_len, hidden_size_});
+    o_proj_->forward_(o_output, attn_output);
+    return o_output;
+}
+
+void Attention::_initialize_preallocated_workspace() {
+    const auto &infinilm_config = infinilm::global_state::get_infinilm_config();
+    auto &preallocated_workspace = infinilm::global_state::get_forward_context().preallocated_workspace;
+    const size_t max_num_batched_tokens = infinilm_config.max_num_batched_tokens;
+
+    const std::string attention_cache_key = std::string("Attention_max_num_batched_tokens_")
+                                          + std::to_string(max_num_batched_tokens) + "_rank_qkv_output_size_"
+                                          + std::to_string(rank_qkv_output_size_) + "_hidden_size_"
+                                          + std::to_string(hidden_size_) + "_dtype_"
+                                          + infinicore::toString(dtype_) + "_device_"
+                                          + device_.toString();
+
+    size_t max_output_size = std::max(rank_qkv_output_size_, hidden_size_);
+    if (preallocated_workspace.find(attention_cache_key) == preallocated_workspace.end()) {
+        auto attention_buffer = infinicore::Tensor::empty({max_num_batched_tokens * max_output_size}, dtype_, device_);
+        preallocated_workspace[attention_cache_key] = attention_buffer;
+    }
+
+    auto attention_buffer = preallocated_workspace.at(attention_cache_key);
+    const auto attention_buffer_shape = attention_buffer->shape();
+    ASSERT(attention_buffer_shape[0] == max_num_batched_tokens * max_output_size);
+
+    max_qkv_output_ = attention_buffer->narrow({{0, 0, max_num_batched_tokens * rank_qkv_output_size_}})->view({max_num_batched_tokens, rank_qkv_output_size_});
+    max_o_output_ = attention_buffer->narrow({{0, 0, max_num_batched_tokens * hidden_size_}})->view({max_num_batched_tokens, hidden_size_});
 }
 
 void init_kv_cache_quant_params(std::function<void(const std::string &, infinicore::nn::Parameter)> register_fn,

csrc/layers/attention/attention.hpp

@@ -5,6 +5,8 @@
 #include "../../global_state/global_state.hpp"
 #include "../linear/linear.hpp"
 #include "backends/attention_layer.hpp"
+#include "infinicore/device.hpp"
+#include "infinicore/dtype.hpp"
 #include "infinicore/nn/module.hpp"
 #include "infinicore/nn/rope.hpp"
 #include "infinicore/tensor.hpp"
@@ -37,6 +39,8 @@ class Attention : public infinicore::nn::Module {
     infinicore::Tensor forward_paged_(const infinicore::Tensor &positions,
                                       const infinicore::Tensor &hidden_states) const;
 
+    void _initialize_preallocated_workspace();
+
 protected:
     std::shared_ptr<infinilm::layers::linear::QKVParallelLinear> qkv_proj_;
     std::shared_ptr<infinilm::layers::linear::RowParallelLinear> o_proj_;
@@ -49,13 +53,22 @@ class Attention : public infinicore::nn::Module {
     size_t num_key_value_heads_;
     size_t hidden_size_;
     size_t head_dim_;
+    infinicore::Device device_;
+    infinicore::DataType dtype_;
 
     // For off-line kv cache quantization
     INFINICORE_NN_PARAMETER(kv_cache_k_scale);
     INFINICORE_NN_PARAMETER(kv_cache_v_scale);
+
+private:
+    size_t rank_qkv_output_size_;
+
+    // preallocated workspace for Attention
+    infinicore::Tensor max_qkv_output_;
+    infinicore::Tensor max_o_output_;
 };
 void init_kv_cache_quant_params(std::function<void(const std::string &, infinicore::nn::Parameter)> register_fn,
-                              const infinicore::Device &device,
-                              infinicore::nn::Parameter &kv_cache_k_scale,
-                              infinicore::nn::Parameter &kv_cache_v_scale);
+                                const infinicore::Device &device,
+                                infinicore::nn::Parameter &kv_cache_k_scale,
+                                infinicore::nn::Parameter &kv_cache_v_scale);
 } // namespace infinilm::layers::attention

csrc/layers/attention/backends/attention_layer.cpp

@@ -9,16 +9,17 @@ AttentionLayer::AttentionLayer(size_t num_heads,
                                size_t layer_idx,
                                infinicore::Tensor k_scale,
                                infinicore::Tensor v_scale,
-                               ::infinilm::backends::AttentionBackend attn_backend) : k_scale_(k_scale), v_scale_(v_scale), layer_idx_(layer_idx), attn_backend_(attn_backend) {
+                               ::infinilm::backends::AttentionBackend attn_backend,
+                               const infinicore::Device &device) : k_scale_(k_scale), v_scale_(v_scale), layer_idx_(layer_idx), attn_backend_(attn_backend) {
     switch (attn_backend) {
     case ::infinilm::backends::AttentionBackend::STATIC_ATTN:
-        attn_backend_impl_ = std::make_shared<backends::StaticAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx);
+        attn_backend_impl_ = std::make_shared<backends::StaticAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx, device);
         break;
     case ::infinilm::backends::AttentionBackend::PAGED_ATTN:
-        attn_backend_impl_ = std::make_shared<backends::PagedAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx);
+        attn_backend_impl_ = std::make_shared<backends::PagedAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx, device);
         break;
     case ::infinilm::backends::AttentionBackend::FLASH_ATTN:
-        attn_backend_impl_ = std::make_shared<backends::FlashAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx);
+        attn_backend_impl_ = std::make_shared<backends::FlashAttentionImpl>(num_heads, head_size, scale, num_kv_heads, layer_idx, device);
         break;
     default:
         throw std::runtime_error("infinilm::layers::attention::AttentionLayer: unsupported attention backend");

csrc/layers/attention/backends/attention_layer.hpp

@@ -31,7 +31,8 @@ class AttentionLayer {
                    size_t layer_idx,
                    infinicore::Tensor k_scale,
                    infinicore::Tensor v_scale,
-                   ::infinilm::backends::AttentionBackend attention_backend);
+                   ::infinilm::backends::AttentionBackend attention_backend,
+                   const infinicore::Device &device);
 
     infinicore::Tensor forward(infinicore::Tensor &query,
                                infinicore::Tensor &key,

csrc/layers/attention/backends/flash_attn.cpp

@@ -1,29 +1,38 @@
 #include "flash_attn.hpp"
 
+#include "../../../global_state/global_state.hpp"
 #include "../../../utils.hpp"
 #include "infinicore/ops.hpp"
 #include "infinicore/ops/mha_kvcache.hpp"
 #include "infinicore/ops/mha_varlen.hpp"
+#include <string>
 
 namespace infinilm::layers::attention::backends {
 
 FlashAttentionImpl::FlashAttentionImpl(size_t num_heads,
                                        size_t head_size,
                                        float scale,
                                        size_t num_kv_heads,
-                                       size_t layer_idx)
+                                       size_t layer_idx,
+                                       const infinicore::Device &device)
     : num_heads_(num_heads),
       head_size_(head_size),
       scale_(scale),
       num_kv_heads_(num_kv_heads),
       layer_idx_(layer_idx),
-      head_dim_(head_size) {
+      head_dim_(head_size),
+      device_(device) {
 
     const infinilm::global_state::InfinilmConfig &infinilm_config = infinilm::global_state::get_infinilm_config();
     if (!infinilm_config.model_config) {
         throw std::runtime_error("infinilm::layers::attention::backends::FlashAttentionImpl: model_config is null");
     }
-    max_position_embeddings_ = infinilm_config.model_config->get<size_t>("max_position_embeddings");
+
+    const auto &model_config = infinilm_config.model_config;
+    dtype_ = model_config->get_dtype();
+    max_position_embeddings_ = model_config->get<size_t>("max_position_embeddings");
+
+    this->_initialize_preallocated_workspace();
 }
 
 infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer,
@@ -48,8 +57,9 @@ infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer,
     bool is_prefill = (seq_len != total_sequence_lengths.value()->shape()[0]);
 
     // 2. Compute attention
-    infinicore::Tensor attn_output = infinicore::Tensor::empty({seq_len, num_heads_, head_dim_}, query->dtype(), query->device());
+    infinicore::Tensor attn_output;
     if (is_prefill) {
+        attn_output = max_attn_output_->narrow({{0, 0, seq_len}});
         infinicore::op::mha_varlen_(
             attn_output,
             query,
@@ -99,4 +109,27 @@ std::tuple<infinicore::Tensor, infinicore::Tensor> FlashAttentionImpl::do_kv_cac
     return {k_cache_layer, v_cache_layer};
 }
 
+void FlashAttentionImpl::_initialize_preallocated_workspace() {
+    const auto &infinilm_config = infinilm::global_state::get_infinilm_config();
+    auto &preallocated_workspace = infinilm::global_state::get_forward_context().preallocated_workspace;
+    const size_t max_num_batched_tokens = infinilm_config.max_num_batched_tokens;
+
+    const std::string cache_key = std::string("FlashAttentionImpl_max_num_batched_tokens_")
+                                + std::to_string(max_num_batched_tokens) + "_num_heads_"
+                                + std::to_string(num_heads_) + "_head_dim_"
+                                + std::to_string(head_dim_) + "_dtype_"
+                                + infinicore::toString(dtype_) + "_device_"
+                                + device_.toString();
+
+    if (preallocated_workspace.find(cache_key) == preallocated_workspace.end()) {
+        auto flash_attention_impl_buffer = infinicore::Tensor::empty({max_num_batched_tokens, num_heads_, head_dim_}, dtype_, device_);
+        preallocated_workspace[cache_key] = flash_attention_impl_buffer;
+    }
+
+    auto flash_attention_impl_buffer = preallocated_workspace.at(cache_key);
+    const auto buffer_shape = flash_attention_impl_buffer->shape();
+    ASSERT(buffer_shape[0] == max_num_batched_tokens && buffer_shape[1] == num_heads_ && buffer_shape[2] == head_dim_);
+
+    max_attn_output_ = flash_attention_impl_buffer;
+}
 } // namespace infinilm::layers::attention::backends

csrc/layers/attention/backends/flash_attn.hpp

@@ -1,6 +1,8 @@
 #pragma once
 
 #include "../../../global_state/global_state.hpp"
+#include "infinicore/device.hpp"
+#include "infinicore/dtype.hpp"
 #include "infinicore/tensor.hpp"
 #include <tuple>
 
@@ -16,7 +18,8 @@ class FlashAttentionImpl {
                        size_t head_size,
                        float scale,
                        size_t num_kv_heads,
-                       size_t layer_idx);
+                       size_t layer_idx,
+                       const infinicore::Device &device);
 
     /**
      * @brief Forward pass with FlashAttention.
@@ -43,12 +46,20 @@ class FlashAttentionImpl {
                                                                           const infinicore::Tensor slot_mapping) const;
 
 private:
+    void _initialize_preallocated_workspace();
+
     size_t num_heads_;
     size_t head_size_;
     float scale_;
     size_t num_kv_heads_;
     size_t layer_idx_;
     size_t head_dim_; // Note: head_dim equals to head_size
     size_t max_position_embeddings_;
+    infinicore::Device device_;
+    infinicore::DataType dtype_;
+
+    // preallocated workspace for FlashAttentionImpl
+    infinicore::Tensor max_attn_output_;
 };
+
 } // namespace infinilm::layers::attention::backends