issue/343 optimize minicpmv resampler

PanZezhong1725 · PanZezhong1725 · commit 429471ec1270 · 2026-05-25T08:01:18.000Z
diff --git a/csrc/models/minicpmv/minicpmv_model.cpp b/csrc/models/minicpmv/minicpmv_model.cpp
@@ -33,6 +33,8 @@ MiniCPMVModel::MiniCPMVModel(std::shared_ptr<infinilm::config::ModelConfig> mode
                               embed_dim,
                               num_heads,
                               vision_cfg.value("hidden_size", 768),
+                              vision_cfg.value("image_size", 224),
+                              vision_cfg.value("patch_size", 16),
                               dtype,
                               device);
 }
diff --git a/csrc/models/minicpmv/resampler.cpp b/csrc/models/minicpmv/resampler.cpp
@@ -121,15 +121,20 @@ Resampler::Resampler(size_t num_queries,
                      size_t embed_dim,
                      size_t num_heads,
                      size_t kv_dim,
+                     size_t image_size,
+                     size_t patch_size,
                      const infinicore::DataType &dtype,
                      const infinicore::Device &device)
     : num_queries_(num_queries),
       embed_dim_(embed_dim),
       num_heads_(num_heads),
       kv_dim_(kv_dim),
+      image_size_(image_size),
+      patch_size_(patch_size),
       use_kv_proj_(kv_dim != embed_dim) {
     INFINICORE_NN_PARAMETER_INIT(query, ({num_queries_, embed_dim_}, dtype, device));
     INFINICORE_NN_PARAMETER_INIT(proj, ({embed_dim_, embed_dim_}, dtype, device));
+
     INFINICORE_NN_MODULE_INIT(attn, embed_dim_, num_heads_, dtype, device);
     INFINICORE_NN_MODULE_INIT(ln_q, embed_dim_, 1e-6, dtype, device);
     INFINICORE_NN_MODULE_INIT(ln_kv, embed_dim_, 1e-6, dtype, device);
@@ -138,6 +143,15 @@ Resampler::Resampler(size_t num_queries,
     if (use_kv_proj_) {
         INFINICORE_NN_MODULE_INIT(kv_proj, kv_dim_, embed_dim_, false, dtype, device);
     }
+
+    // Initialize full 2d embeddings with max size, calculate on cpu and copy to gpu
+    size_t num_patches = image_size_ / patch_size_;
+    INFINICORE_NN_BUFFER_INIT(embedding_table, ({num_patches, num_patches, embed_dim_}, dtype, device_));
+    std::vector<float> buf(num_patches * num_patches * embed_dim_);
+    compute_2d_sincos_pos_embed(buf.data(), embed_dim_, num_patches, num_patches);
+    auto embedding_table_cpu = infinicore::Tensor::zeros({num_patches, num_patches, embed_dim_}, dtype, infinicore::Device::cpu());
+    write_pos_embed(embedding_table_cpu->data(), embedding_table_cpu->dtype(), buf.data(), num_patches * num_patches * embed_dim_);
+    embedding_table_->copy_from(embedding_table_cpu);
 }
 
 infinicore::Tensor Resampler::forward(const infinicore::Tensor &x,
@@ -152,32 +166,22 @@ infinicore::Tensor Resampler::forward(const infinicore::Tensor &x,
     kv = ln_kv_->forward(kv);
 
     // Build positional embeddings on CPU
-    std::vector<int64_t> tgt_sizes_host;
-
     auto tgt_cpu = tgt_sizes->to(infinicore::Device::cpu());
-    auto n = tgt_cpu->numel();
-    tgt_sizes_host.resize(n);
-    std::memcpy(tgt_sizes_host.data(), tgt_cpu->data(), n * sizeof(int64_t));
+    int64_t *tgt_sizes_ptr = (int64_t *)(tgt_cpu->data());
 
-    auto pos_cpu = infinicore::Tensor::zeros({batch_size, seq_len, embed_dim_}, kv->dtype(), infinicore::Device::cpu());
-    auto *pos_ptr = reinterpret_cast<std::byte *>(pos_cpu->data());
-    const size_t elem_size = pos_cpu->element_size();
+    auto pos_embeddings = infinicore::Tensor::zeros(kv->shape(), kv->dtype(), kv->device());
 
     for (size_t b = 0; b < batch_size; ++b) {
-        size_t tgt_h = 1;
-        size_t tgt_w = seq_len;
-        if (!tgt_sizes_host.empty()) {
-            tgt_h = static_cast<size_t>(tgt_sizes_host[b * 2]);
-            tgt_w = static_cast<size_t>(tgt_sizes_host[b * 2 + 1]);
-        }
-        const size_t patch_len = tgt_h * tgt_w;
-        std::vector<float> buf(patch_len * embed_dim_);
-        compute_2d_sincos_pos_embed(buf.data(), embed_dim_, tgt_h, tgt_w);
-        write_pos_embed(pos_ptr + b * seq_len * embed_dim_ * elem_size, pos_cpu->dtype(), buf.data(), patch_len * embed_dim_);
+
+        auto tgt_h = static_cast<size_t>(tgt_sizes_ptr[b * 2]);
+        auto tgt_w = static_cast<size_t>(tgt_sizes_ptr[b * 2 + 1]);
+
+        auto src_embeddings = embedding_table_->narrow({{0, 0, tgt_h}, {1, 0, tgt_w}});
+        auto tgt_embeddings = pos_embeddings->narrow({{0, b, 1}, {1, 0, tgt_h * tgt_w}})->view({tgt_h, tgt_w, embed_dim_});
+        tgt_embeddings->copy_from(src_embeddings);
     }
 
-    auto pos = pos_cpu->to(kv->device());
-    auto kv_with_pos = infinicore::op::add(kv, pos);
+    auto kv_with_pos = infinicore::op::add(kv, pos_embeddings);
 
     auto q = ln_q_->forward(query_);
     if (q->shape().size() == 2) {
diff --git a/csrc/models/minicpmv/resampler.hpp b/csrc/models/minicpmv/resampler.hpp
@@ -39,6 +39,8 @@ class Resampler : public infinicore::nn::Module {
               size_t embed_dim,
               size_t num_heads,
               size_t kv_dim,
+              size_t image_size,
+              size_t patch_size,
               const infinicore::DataType &dtype,
               const infinicore::Device &device);
 
@@ -50,10 +52,13 @@ class Resampler : public infinicore::nn::Module {
     size_t embed_dim_;
     size_t num_heads_;
     size_t kv_dim_;
+    size_t image_size_;
+    size_t patch_size_;
     bool use_kv_proj_;
 
     INFINICORE_NN_PARAMETER(query);
     INFINICORE_NN_PARAMETER(proj);
+    INFINICORE_NN_BUFFER(embedding_table);
     INFINICORE_NN_MODULE(infinicore::nn::Linear, kv_proj);
     INFINICORE_NN_MODULE(ResamplerAttention, attn);
     INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, ln_q);
diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp
@@ -30,51 +30,6 @@ namespace infinilm::engine {
 
 inline void bind_infer_engine(py::module &m) {
     py::class_<InferEngine, std::shared_ptr<InferEngine>> infer_engine(m, "InferEngine");
-    infer_engine
-        .def(py::init([](
-                          const InfinilmModel::Config &cfg,
-                          const distributed::DistConfig &dist,
-                          infinicore::Device::Type dev,
-                          std::shared_ptr<const infinilm::cache::CacheConfig> cache_cfg,
-                          bool enable_graph_compiling,
-                          const std::string &attention_backend) {
-                 return std::make_shared<InferEngine>(
-                     cfg,
-                     dist,
-                     dev,
-                     cache_cfg ? cache_cfg.get() : nullptr,
-                     enable_graph_compiling,
-                     infinilm::backends::parse_attention_backend(attention_backend));
-             }),
-             py::arg("config"),
-             py::arg("distributed_config") = distributed::DistConfig(),
-             py::arg("device_type") = infinicore::context::getDevice().getType(),
-             py::arg("cache_config") = py::none(),
-             py::arg("enable_graph_compiling") = false,
-             py::arg("attention_backend") = "default")
-        .def("load_param", &InferEngine::load_param,
-             py::arg("name"), py::arg("param"),
-             "Load a parameter tensor into all workers (each worker picks its shard)")
-        .def("state_dict", [](InferEngine &self) {
-            py::list state_dict_tp_all;
-            for (const auto &state_dict_tp : self.state_dict()) {
-                py::dict result;
-                for (const auto &[name, param] : state_dict_tp) {
-                    result[py::cast(name)] = infinicore::Tensor(param);
-                }
-                state_dict_tp_all.append(result);
-            }
-            return state_dict_tp_all;
-        })
-        .def("process_weights_after_loading", &InferEngine::process_weights_after_loading, "Process the weights after loading on all workers (e.g., for quantization)")
-        .def(
-            "forward", [](InferEngine &self, const InferEngine::Input &input) -> InferEngine::Output { return self.forward(input); }, "Run inference on all ranks with arbitrary arguments")
-        .def(
-            "reset_cache", [](InferEngine &self, std::shared_ptr<cache::CacheConfig> cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none())
-        .def("get_cache_config", [](const InferEngine &self) -> std::shared_ptr<cache::CacheConfig> {
-            auto cfg = self.get_cache_config();
-            return cfg ? std::shared_ptr<cache::CacheConfig>(cfg->unique_copy()) : nullptr; })
-        .def("__repr__", [](const InferEngine &self) { return "<InferEngine: " + std::string(self.get_dist_config()) + ">"; });
 
     infer_engine
         .def(py::init([](
@@ -116,8 +71,10 @@ inline void bind_infer_engine(py::module &m) {
             return state_dict_tp_all;
         })
         .def("process_weights_after_loading", &InferEngine::process_weights_after_loading, "Process the weights after loading on all workers (e.g., for quantization)")
-        .def("forward", [](InferEngine &self, const InferEngine::Input &input) -> InferEngine::Output { return self.forward(input); }, "Run inference on all ranks with arbitrary arguments")
-        .def("reset_cache", [](InferEngine &self, std::shared_ptr<cache::CacheConfig> cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none())
+        .def(
+            "forward", [](InferEngine &self, const InferEngine::Input &input) -> InferEngine::Output { return self.forward(input); }, "Run inference on all ranks with arbitrary arguments")
+        .def(
+            "reset_cache", [](InferEngine &self, std::shared_ptr<cache::CacheConfig> cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none())
         .def("get_cache_config", [](const InferEngine &self) -> std::shared_ptr<cache::CacheConfig> {
             auto cfg = self.get_cache_config();
             return cfg ? std::shared_ptr<cache::CacheConfig>(cfg->unique_copy()) : nullptr; })
diff --git a/test/service/request.py b/test/service/request.py
@@ -25,7 +25,24 @@ def get_args():
         "--port", type=int, default=8000, help="Infer server port, default 8000"
     )
     parser.add_argument(
-        "--host", default="127.0.0.1", help="Infer server url, default 127.0.0.1"
+        "--host",
+        type=str,
+        default="127.0.0.1",
+        help="Infer server url, default 127.0.0.1",
+    )
+
+    parser.add_argument(
+        "--api-url",
+        type=str,
+        default=None,
+        help="Full service url, if given host and port will be ignored",
+    )
+
+    parser.add_argument(
+        "--model-name",
+        type=str,
+        default="default",
+        help="Name of the model being served, needed by vllm",
     )
 
     return parser.parse_args()
@@ -56,12 +73,12 @@ def build_messages(content_args, system_prompt):
     return messages
 
 
-async def benchmark_user(client, messages):
+async def benchmark_user(client, messages, model_name):
     try:
         print(f"  ❓ 提问: {messages}")
         start_time = time.time()
         stream = await client.chat.completions.create(
-            model="default",
+            model=model_name,
             messages=messages,
             stream=True,
         )
@@ -110,8 +127,14 @@ def main():
     if not args.content:
         args.content = ["text:山东最高的山是？"]
     messages = build_messages(args.content, args.system)
-    client = AsyncOpenAI(base_url=f"http://{args.host}:{args.port}", api_key="default")
-    asyncio.run(benchmark_user(client, messages))
+    api_url = (
+        f"http://{args.api_url}"
+        if args.api_url is not None
+        else f"http://{args.host}:{args.port}"
+    )
+
+    client = AsyncOpenAI(base_url=api_url, api_key="default")
+    asyncio.run(benchmark_user(client, messages, args.model_name))
 
 
 if __name__ == "__main__":

Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,8 @@ MiniCPMVModel::MiniCPMVModel(std::shared_ptr<infinilm::config::ModelConfig> mode`
`33`	`33`	`embed_dim,`
`34`	`34`	`num_heads,`
`35`	`35`	`vision_cfg.value("hidden_size", 768),`
	`36`	`+ vision_cfg.value("image_size", 224),`
	`37`	`+ vision_cfg.value("patch_size", 16),`
`36`	`38`	`dtype,`
`37`	`39`	`device);`
`38`	`40`	`}`