issue/343 optimize siglip attention

PanZezhong1725 · PanZezhong1725 · commit cd3bb2ba90bb · 2026-06-05T06:18:00.000Z
diff --git a/csrc/layers/linear/fused_linear.cpp b/csrc/layers/linear/fused_linear.cpp
@@ -31,14 +31,14 @@ QKVParallelLinear::QKVParallelLinear(size_t hidden_size,
                                      const infinicore::Device &device,
                                      engine::distributed::RankInfo rank_info)
     : infinilm::nn::ColumnParallelLinear(
-          hidden_size,
-          calculate_out_feature_size(num_q_head, q_dim, num_k_head, k_dim, num_v_head, v_dim, rank_info),
-          quantization,
-          (q_bias || k_bias || v_bias),
-          dtype,
-          device,
-          rank_info.tp_rank,
-          rank_info.tp_size),
+        hidden_size,
+        calculate_out_feature_size(num_q_head, q_dim, num_k_head, k_dim, num_v_head, v_dim, rank_info),
+        quantization == nullptr ? std::make_shared<infinilm::quantization::NoneQuantization>() : quantization,
+        (q_bias || k_bias || v_bias),
+        dtype,
+        device,
+        rank_info.tp_rank,
+        rank_info.tp_size),
       q_dim_(q_dim),
       k_dim_(k_dim),
       v_dim_(v_dim),
@@ -120,7 +120,17 @@ GateUpParallelLinear::GateUpParallelLinear(size_t hidden_size, size_t intermedia
                                            std::shared_ptr<infinilm::quantization::BaseQuantization> quantization,
                                            const infinicore::DataType &dtype, const infinicore::Device &device,
                                            engine::distributed::RankInfo rank_info)
-    : infinilm::nn::ColumnParallelLinear(hidden_size, intermediate_size * 2, quantization, gate_bias || up_bias, dtype, device, rank_info.tp_rank, rank_info.tp_size), gate_bias_(gate_bias), up_bias_(up_bias) {
+    : infinilm::nn::ColumnParallelLinear(
+        hidden_size,
+        intermediate_size * 2,
+        quantization == nullptr ? std::make_shared<infinilm::quantization::NoneQuantization>() : quantization,
+        gate_bias || up_bias,
+        dtype,
+        device,
+        rank_info.tp_rank,
+        rank_info.tp_size),
+      gate_bias_(gate_bias),
+      up_bias_(up_bias) {
     if (gate_bias_ != up_bias_) {
         throw std::runtime_error("Not supported yet: gate_bias and up_bias should be given at the same time");
     }
diff --git a/csrc/layers/linear/fused_linear.hpp b/csrc/layers/linear/fused_linear.hpp
@@ -1,7 +1,7 @@
 #pragma once
 #include "../../engine/distributed/communication_group.hpp"
-#include "linear.hpp"
 #include "../quantization/quantization.hpp"
+#include "linear.hpp"
 #include <functional>
 
 namespace infinilm::layers::linear {
@@ -13,15 +13,15 @@ class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear {
                                size_t q_dim, size_t k_dim, size_t v_dim,
                                size_t num_q_head, size_t num_k_head, size_t num_v_head,
                                bool q_bias, bool k_bias, bool v_bias,
-                               std::shared_ptr<infinilm::quantization::BaseQuantization> quantization,
+                               std::shared_ptr<infinilm::quantization::BaseQuantization> quantization = nullptr,
                                const infinicore::DataType &dtype = infinicore::DataType::F32,
                                const infinicore::Device &device = infinicore::Device(),
                                engine::distributed::RankInfo rank_info = engine::distributed::RankInfo());
 
     explicit QKVParallelLinear(size_t hidden_size,
                                size_t head_dim,
                                size_t num_q_head, size_t num_kv_head,
-                               std::shared_ptr<infinilm::quantization::BaseQuantization> quantization,
+                               std::shared_ptr<infinilm::quantization::BaseQuantization> quantization = nullptr,
                                bool bias = false,
                                const infinicore::DataType &dtype = infinicore::DataType::F32,
                                const infinicore::Device &device = infinicore::Device(),
@@ -32,7 +32,7 @@ class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear {
                       size_t num_q_head, size_t num_kv_head,
                       const std::string &q_name, const std::string &k_name, const std::string &v_name,
                       RegisterParamFn register_fn,
-                      std::shared_ptr<infinilm::quantization::BaseQuantization> quantization,
+                      std::shared_ptr<infinilm::quantization::BaseQuantization> quantization = nullptr,
                       bool bias = false,
                       const infinicore::DataType &dtype = infinicore::DataType::F32,
                       const infinicore::Device &device = infinicore::Device(),
@@ -84,21 +84,22 @@ class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear {
 
 class GateUpParallelLinear : public infinilm::nn::ColumnParallelLinear {
 public:
-    GateUpParallelLinear(size_t hidden_size, size_t intermediate_size, std::shared_ptr<infinilm::quantization::BaseQuantization> quantization,
+    GateUpParallelLinear(size_t hidden_size, size_t intermediate_size,
+                         std::shared_ptr<infinilm::quantization::BaseQuantization> quantization = nullptr,
                          bool bias = false,
                          const infinicore::DataType &dtype = infinicore::DataType::F32,
                          const infinicore::Device &device = infinicore::Device(),
                          engine::distributed::RankInfo rank_info = engine::distributed::RankInfo());
 
     GateUpParallelLinear(size_t hidden_size, size_t intermediate_size, bool gate_bias, bool up_bias,
-                         std::shared_ptr<infinilm::quantization::BaseQuantization> quantization,
+                         std::shared_ptr<infinilm::quantization::BaseQuantization> quantization = nullptr,
                          const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(),
                          engine::distributed::RankInfo rank_info = engine::distributed::RankInfo());
 
     GateUpParallelLinear(size_t hidden_size, size_t intermediate_size,
                          const std::string &gate_name, const std::string &up_name,
                          RegisterParamFn register_fn,
-                         std::shared_ptr<infinilm::quantization::BaseQuantization> quantization,
+                         std::shared_ptr<infinilm::quantization::BaseQuantization> quantization = nullptr,
                          bool bias = false,
                          const infinicore::DataType &dtype = infinicore::DataType::F32,
                          const infinicore::Device &device = infinicore::Device(),
diff --git a/csrc/layers/quantization/none_quantization.cpp b/csrc/layers/quantization/none_quantization.cpp
@@ -4,6 +4,8 @@
 
 namespace infinilm::quantization {
 
+NoneQuantization::NoneQuantization() : NoneQuantization(nlohmann::json()) {}
+
 std::vector<ParamDescriptor> NoneQuantization::get_param_layout(
     size_t in_features, size_t out_features,
     int split_dim, int tp_rank, int tp_size,
@@ -14,8 +16,7 @@ std::vector<ParamDescriptor> NoneQuantization::get_param_layout(
     std::vector<ParamDescriptor> descs;
     descs.push_back({"weight", {out_features, in_features}, dtype, split_dim, tp_rank, tp_size});
     if (bias) {
-        descs.push_back({"bias", {out_features}, dtype, split_dim >= 0 ? 0 : -1,
-                         split_dim >= 0 ? tp_rank : 0, split_dim >= 0 ? tp_size : 1});
+        descs.push_back({"bias", {out_features}, dtype, split_dim >= 0 ? 0 : -1, split_dim >= 0 ? tp_rank : 0, split_dim >= 0 ? tp_size : 1});
     }
     return descs;
 }
diff --git a/csrc/layers/quantization/none_quantization.hpp b/csrc/layers/quantization/none_quantization.hpp
@@ -6,7 +6,9 @@ namespace infinilm::quantization {
 class NoneQuantization : public BaseQuantization {
 public:
     explicit NoneQuantization(const nlohmann::json &quant_config)
-        : BaseQuantization(quant_config) {};
+        : BaseQuantization(quant_config){};
+
+    NoneQuantization();
 
     QuantScheme get_quant_scheme() const override {
         return QuantScheme::NONE;
diff --git a/csrc/models/minicpmv/resampler.hpp b/csrc/models/minicpmv/resampler.hpp
@@ -1,8 +1,8 @@
 #pragma once
 
 #include "../../config/model_config.hpp"
+#include "../../layers/linear/fused_linear.hpp"
 #include "infinicore/nn/layer_norm.hpp"
-#include "infinicore/nn/linear.hpp"
 #include "infinicore/nn/module.hpp"
 #include "infinicore/tensor.hpp"
 
@@ -30,7 +30,7 @@ class ResamplerAttention : public infinicore::nn::Module {
 
     INFINICORE_NN_PARAMETER(in_proj_weight);
     INFINICORE_NN_PARAMETER(in_proj_bias);
-    INFINICORE_NN_MODULE(infinicore::nn::Linear, out_proj);
+    INFINICORE_NN_MODULE(infinilm::nn::Linear, out_proj);
 };
 
 class Resampler : public infinicore::nn::Module {
@@ -59,7 +59,7 @@ class Resampler : public infinicore::nn::Module {
     INFINICORE_NN_PARAMETER(query);
     INFINICORE_NN_PARAMETER(proj);
     INFINICORE_NN_BUFFER(embedding_table);
-    INFINICORE_NN_MODULE(infinicore::nn::Linear, kv_proj);
+    INFINICORE_NN_MODULE(infinilm::nn::Linear, kv_proj);
     INFINICORE_NN_MODULE(ResamplerAttention, attn);
     INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, ln_q);
     INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, ln_kv);
diff --git a/csrc/models/minicpmv/siglip_vision.cpp b/csrc/models/minicpmv/siglip_vision.cpp
@@ -1,6 +1,8 @@
 #include "siglip_vision.hpp"
 
+#include "../../global_state/global_state.hpp"
 #include "infinicore/ops.hpp"
+#include "infinicore/ops/mha.hpp"
 
 #include <cmath>
 #include <cstring>
@@ -92,44 +94,52 @@ SiglipAttention::SiglipAttention(const nlohmann::json &config,
     if (embed_dim_ % num_heads_ != 0) {
         throw std::runtime_error("SiglipAttention: embed_dim must be divisible by num_heads");
     }
-    INFINICORE_NN_MODULE_INIT(q_proj, embed_dim_, embed_dim_, true, dtype, device);
-    INFINICORE_NN_MODULE_INIT(k_proj, embed_dim_, embed_dim_, true, dtype, device);
-    INFINICORE_NN_MODULE_INIT(v_proj, embed_dim_, embed_dim_, true, dtype, device);
+    qkv_proj_ = std::make_shared<infinilm::layers::linear::QKVParallelLinear>(
+        embed_dim_, head_dim_, num_heads_, num_heads_,
+        "q_proj", "k_proj", "v_proj", [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); },
+        nullptr, true, dtype, device);
+
     INFINICORE_NN_MODULE_INIT(out_proj, embed_dim_, embed_dim_, true, dtype, device);
+
+    attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend;
 }
 
 infinicore::Tensor SiglipAttention::forward(const infinicore::Tensor &hidden_states,
                                             const std::optional<infinicore::Tensor> &attention_mask) const {
-    (void)attention_mask;
     auto shape = hidden_states->shape();
     size_t batch_size = shape[0];
     size_t seq_len = shape[1];
 
-    auto q = q_proj_->forward(const_cast<infinicore::Tensor &>(hidden_states));
-    auto k = k_proj_->forward(const_cast<infinicore::Tensor &>(hidden_states));
-    auto v = v_proj_->forward(const_cast<infinicore::Tensor &>(hidden_states));
-
-    auto q_reshaped = q->view({batch_size, seq_len, num_heads_, head_dim_})->permute({0, 2, 1, 3})->contiguous();
-    auto k_reshaped = k->view({batch_size, seq_len, num_heads_, head_dim_})->permute({0, 2, 1, 3})->contiguous();
-    auto v_reshaped = v->view({batch_size, seq_len, num_heads_, head_dim_})->permute({0, 2, 1, 3})->contiguous();
-
-    auto q_flat = q_reshaped->view({batch_size * num_heads_, seq_len, head_dim_});
-    auto k_flat = k_reshaped->view({batch_size * num_heads_, seq_len, head_dim_});
-    auto v_flat = v_reshaped->view({batch_size * num_heads_, seq_len, head_dim_});
-
-    auto k_t = k_flat->permute({0, 2, 1});
-    auto attn_weights = infinicore::op::matmul(q_flat, k_t, scale_);
+    auto qkv = qkv_proj_->forward(const_cast<infinicore::Tensor &>(hidden_states))->view({batch_size, seq_len, num_heads_ * 3, head_dim_});
+    auto q = qkv->narrow({{2, 0, num_heads_}});
+    auto k = qkv->narrow({{2, num_heads_, num_heads_}});
+    auto v = qkv->narrow({{2, num_heads_ * 2, num_heads_}});
 
-    auto attn_view = attn_weights->view({batch_size * num_heads_, seq_len, seq_len});
-    infinicore::op::softmax_(attn_view, attn_view, -1);
-
-    auto attn_output = infinicore::op::matmul(attn_weights, v_flat);
-    auto out = attn_output->view({batch_size, num_heads_, seq_len, head_dim_})
-                   ->permute({0, 2, 1, 3})
-                   ->contiguous()
-                   ->view({batch_size, seq_len, embed_dim_});
-
-    return out_proj_->forward(out);
+    if (attention_backend_ == infinilm::backends::AttentionBackend::FLASH_ATTN) {
+        auto out = infinicore::op::mha(q, k, v, std::nullopt, scale_, false)->view({batch_size, seq_len, num_heads_ * head_dim_});
+        return out_proj_->forward(out);
+    } else {
+        auto q_reshaped = q->view({batch_size, seq_len, num_heads_, head_dim_})->permute({0, 2, 1, 3})->contiguous();
+        auto k_reshaped = k->view({batch_size, seq_len, num_heads_, head_dim_})->permute({0, 2, 1, 3})->contiguous();
+        auto v_reshaped = v->view({batch_size, seq_len, num_heads_, head_dim_})->permute({0, 2, 1, 3})->contiguous();
+
+        auto q_flat = q_reshaped->view({batch_size * num_heads_, seq_len, head_dim_});
+        auto k_flat = k_reshaped->view({batch_size * num_heads_, seq_len, head_dim_});
+        auto v_flat = v_reshaped->view({batch_size * num_heads_, seq_len, head_dim_});
+
+        auto k_t = k_flat->permute({0, 2, 1});
+        auto attn_weights = infinicore::op::matmul(q_flat, k_t, scale_);
+
+        auto attn_view = attn_weights->view({batch_size * num_heads_, seq_len, seq_len});
+        infinicore::op::softmax_(attn_view, attn_view, -1);
+
+        auto attn_output = infinicore::op::matmul(attn_weights, v_flat);
+        auto out = attn_output->view({batch_size, num_heads_, seq_len, head_dim_})
+                       ->permute({0, 2, 1, 3})
+                       ->contiguous()
+                       ->view({batch_size, seq_len, embed_dim_});
+        return out_proj_->forward(out);
+    }
 }
 
 SiglipMLP::SiglipMLP(const nlohmann::json &config,
diff --git a/csrc/models/minicpmv/siglip_vision.hpp b/csrc/models/minicpmv/siglip_vision.hpp
@@ -1,9 +1,10 @@
 #pragma once
 
+#include "../../backends/attention_backends.hpp"
 #include "../../config/model_config.hpp"
+#include "../../layers/linear/fused_linear.hpp"
 #include "infinicore/nn/embedding.hpp"
 #include "infinicore/nn/layer_norm.hpp"
-#include "infinicore/nn/linear.hpp"
 #include "infinicore/nn/module.hpp"
 #include "infinicore/tensor.hpp"
 #include <nlohmann/json.hpp>
@@ -61,11 +62,10 @@ class SiglipAttention : public infinicore::nn::Module {
     size_t num_heads_;
     size_t head_dim_;
     float scale_;
+    infinilm::backends::AttentionBackend attention_backend_;
 
-    INFINICORE_NN_MODULE(infinicore::nn::Linear, q_proj);
-    INFINICORE_NN_MODULE(infinicore::nn::Linear, k_proj);
-    INFINICORE_NN_MODULE(infinicore::nn::Linear, v_proj);
-    INFINICORE_NN_MODULE(infinicore::nn::Linear, out_proj);
+    INFINICORE_NN_MODULE(infinilm::layers::linear::QKVParallelLinear, qkv_proj);
+    INFINICORE_NN_MODULE(infinilm::nn::Linear, out_proj);
 };
 
 class SiglipMLP : public infinicore::nn::Module {
@@ -78,8 +78,8 @@ class SiglipMLP : public infinicore::nn::Module {
 
 private:
     std::string activation_;
-    INFINICORE_NN_MODULE(infinicore::nn::Linear, fc1);
-    INFINICORE_NN_MODULE(infinicore::nn::Linear, fc2);
+    INFINICORE_NN_MODULE(infinilm::nn::Linear, fc1);
+    INFINICORE_NN_MODULE(infinilm::nn::Linear, fc2);
 };
 
 class SiglipEncoderLayer : public infinicore::nn::Module {
diff --git a/test/service/request.py b/test/service/request.py
@@ -134,7 +134,7 @@ def main():
     )
 
     client = AsyncOpenAI(base_url=api_url, api_key="default")
-    asyncio.run(benchmark_user(client, messages, args.model_name))
+    asyncio.run(benchmark_user(client, messages, args.model))
 
 
 if __name__ == "__main__":

Original file line number	Diff line number	Diff line change
`@@ -134,7 +134,7 @@ def main():`
`134`	`134`	`)`
`135`	`135`
`136`	`136`	`client = AsyncOpenAI(base_url=api_url, api_key="default")`
`137`		`- asyncio.run(benchmark_user(client, messages, args.model_name))`
	`137`	`+ asyncio.run(benchmark_user(client, messages, args.model))`
`138`	`138`
`139`	`139`
`140`	`140`	`if __name__ == "__main__":`