From a15e27e4c397b21d9d2e52ad8ceacf37ec8a2103 Mon Sep 17 00:00:00 2001 From: MaYuhang <2902139028@qq.com> Date: Mon, 6 Jul 2026 09:38:35 +0000 Subject: [PATCH] issue/477 - feat: support partial cache hits for multimodal visual embeddings --- csrc/engine/infer_engine.cpp | 1 + csrc/engine/rank_worker.hpp | 5 + csrc/models/infinilm_model.hpp | 7 +- csrc/models/minicpmv/minicpmv_model.cpp | 77 +++- csrc/models/minicpmv/minicpmv_model.hpp | 3 +- .../videonsa_for_conditional_generation.cpp | 128 +++++- .../videonsa_for_conditional_generation.hpp | 3 +- csrc/pybind11/engine/engine.hpp | 4 + python/infinilm/infer_engine.py | 5 + python/infinilm/llm/cache_manager.py | 121 +++--- python/infinilm/llm/llm.py | 14 +- python/infinilm/llm/request.py | 24 +- python/infinilm/llm/scheduler.py | 20 +- python/infinilm/multimodal/features.py | 68 ++++ .../processors/basic_llm_processor.py | 19 +- .../infinilm/processors/minicpmv_processor.py | 376 ++++++++++++------ python/infinilm/processors/processor.py | 16 +- .../infinilm/processors/videonsa_processor.py | 359 ++++++++++++----- 18 files changed, 940 insertions(+), 310 deletions(-) create mode 100644 python/infinilm/multimodal/features.py diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index 37c63bde2..11143c518 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -173,6 +173,7 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { to_device(mamba_final_state_indices), to_device_vec(pixel_values), to_device_vec(image_bound), + to_device_vec(image_embed_bound), to_device_vec(tgt_sizes), visual_token_ranges, }; diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index 5ca0dbcec..45c098a6d 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -59,8 +59,13 @@ class RankWorker { /// Image pixel values for multi-modal models. std::optional> pixel_values; /// Image placeholder bounds for MiniCPM-V style replacement. + /// Vector of tensors shape: [1, n_patch, 2]. std::optional> image_bound; + /// Source embedding bounds for partial multimodal prefill replacement. + /// Vector of tensors shape: [1, n_patch, 2]. + std::optional> image_embed_bound; /// Target patch sizes for each image (MiniCPM-V). + /// Vector of tensors shape is model-specific (MiniCPM-V: [n_patch, 2], VideoNSA: [n_media, 3]). std::optional> tgt_sizes; /// req_id for each pixel_values among a batch std::optional> image_req_ids; diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index de6b5f9e3..7fdcb16eb 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -42,10 +42,13 @@ class InfinilmModel : public infinicore::nn::Module { /// Vector of tensors. Shape is model-specific (e.g. LLaVA: [batch, 3, H, W], MiniCPM-V: [n_patch, 3, filter_H, H * W / filter_H]). std::optional> pixel_values; /// Image placeholder bounds for MiniCPM-V style replacement. - /// Vector of tensors shape: [n_patch, 2]. + /// Vector of tensors shape: [1, n_patch, 2]. std::optional> image_bound; + /// Source embedding bounds for partial multimodal prefill replacement. + /// Vector of tensors shape: [1, n_patch, 2]. + std::optional> image_embed_bound; /// Target patch sizes for each image (MiniCPM-V). - /// Vector of tensors shape: [n_path, 2] if pre-flattened. + /// Vector of tensors shape is model-specific (MiniCPM-V: [n_patch, 2], VideoNSA: [n_media, 3]). std::optional> tgt_sizes; /// Flattened [start, end) visual token ranges in the packed language sequence. std::optional> visual_token_ranges; diff --git a/csrc/models/minicpmv/minicpmv_model.cpp b/csrc/models/minicpmv/minicpmv_model.cpp index bbe1d227a..681a79530 100644 --- a/csrc/models/minicpmv/minicpmv_model.cpp +++ b/csrc/models/minicpmv/minicpmv_model.cpp @@ -41,23 +41,50 @@ MiniCPMVModel::MiniCPMVModel(std::shared_ptr mode void MiniCPMVModel::replace_embeddings(infinicore::Tensor inputs_embeds, const infinicore::Tensor &vision_hidden, - const infinicore::Tensor &image_bound) const { + const infinicore::Tensor &image_bound, + const infinicore::Tensor &image_embed_bound) const { auto bounds_cpu = image_bound->to(infinicore::Device::cpu()); + auto embed_bounds_cpu = image_embed_bound->to(infinicore::Device::cpu()); auto batch_size = inputs_embeds->size(0); ASSERT_EQ(batch_size, 1); ASSERT_EQ(bounds_cpu->size(0), 1); + ASSERT_EQ(bounds_cpu->size(2), 2); + ASSERT_EQ(embed_bounds_cpu->size(0), 1); + ASSERT_EQ(embed_bounds_cpu->size(1), bounds_cpu->size(1)); + ASSERT_EQ(embed_bounds_cpu->size(2), 2); auto out_slice = inputs_embeds->squeeze(0); auto bound_slice = bounds_cpu->squeeze(0); + auto embed_bound_slice = embed_bounds_cpu->squeeze(0); auto vision_len = vision_hidden->size(0); + ASSERT_EQ(vision_len, bound_slice->size(0)); for (size_t patch = 0; patch < vision_len; ++patch) { auto patch_embed = vision_hidden->narrow({{0, patch, 1}})->squeeze(0); auto bound = bound_slice->narrow({{0, patch, 1}}); auto bound_ptr = reinterpret_cast(bound->data()); - auto start = bound_ptr[0]; - auto end = bound_ptr[1]; + const int64_t start = bound_ptr[0]; + const int64_t end = bound_ptr[1]; + if (start < 0 || end < start) { + throw std::runtime_error("MiniCPMVModel: invalid image_bound"); + } + const int64_t dst_len = end - start; + + auto embed_bound = embed_bound_slice->narrow({{0, patch, 1}}); + auto embed_bound_ptr = reinterpret_cast(embed_bound->data()); + const int64_t src_start = embed_bound_ptr[0]; + const int64_t src_end = embed_bound_ptr[1]; + if (src_start < 0 || src_end < src_start) { + throw std::runtime_error("MiniCPMVModel: invalid image_embed_bound"); + } + if (dst_len != src_end - src_start) { + throw std::runtime_error("MiniCPMVModel: image_bound and image_embed_bound length mismatch"); + } + if (static_cast(end) > out_slice->size(0) || static_cast(src_end) > patch_embed->size(0)) { + throw std::runtime_error("MiniCPMVModel: multimodal embedding bounds are out of range"); + } - out_slice->narrow({{0, size_t(start), size_t(end - start)}})->copy_from(patch_embed); + out_slice->narrow({{0, static_cast(start), static_cast(dst_len)}}) + ->copy_from(patch_embed->narrow({{0, static_cast(src_start), static_cast(dst_len)}})); } } @@ -74,18 +101,50 @@ InfinilmModel::Output MiniCPMVModel::forward(const InfinilmModel::Input &input) if (input.pixel_values->size() != input.image_bound->size() || input.pixel_values->size() != input.tgt_sizes->size()) { throw std::runtime_error("MiniCPMVModel: pixel_values, image_bound and tgt_sizes must have the same number of elements"); } + const auto &mm_metadata = global_state::get_forward_context().mm_metadata; + if (!mm_metadata.image_req_ids.has_value()) { + throw std::runtime_error("MiniCPMVModel: image_req_ids must be provided with pixel_values"); + } + const auto &image_req_ids = mm_metadata.image_req_ids.value(); + if (input.pixel_values->size() != image_req_ids.size()) { + throw std::runtime_error("MiniCPMVModel: multimodal tensor lists must match image_req_ids"); + } + if (!input.image_embed_bound.has_value()) { + throw std::runtime_error("MiniCPMVModel: image_embed_bound must be provided with pixel_values"); + } + if (input.image_embed_bound->size() != image_req_ids.size()) { + throw std::runtime_error("MiniCPMVModel: image_embed_bound must match image_req_ids"); + } auto inputs_embeds = llm_->model().embed_tokens(input_ids); // inputs_embeds concat tokens from all requests, while images are processed per request // slice inputs_embeds using request offsets to get the embedding of each request + if (!input.input_offsets.has_value()) { + throw std::runtime_error("MiniCPMVModel: input_offsets is required with pixel_values"); + } infinicore::Tensor input_offsets_cpu = input.input_offsets.value()->to(infinicore::Device::cpu()); int32_t *offsets = (int32_t *)(input_offsets_cpu->data()); - for (size_t i : global_state::get_forward_context().mm_metadata.image_req_ids.value()) { - auto pixel_values = input.pixel_values.value().at(i); - auto vision_embedding = vpm_->forward(pixel_values, input.tgt_sizes.value().at(i)); - auto vision_hidden = resampler_->forward(vision_embedding, input.tgt_sizes.value().at(i)); - replace_embeddings(inputs_embeds->narrow({{1, size_t(offsets[i]), size_t(offsets[i + 1] - offsets[i])}}), vision_hidden, input.image_bound.value().at(i)); + const size_t num_offsets = input_offsets_cpu->size(0); + for (size_t media_idx = 0; media_idx < image_req_ids.size(); ++media_idx) { + const size_t req_id = image_req_ids[media_idx]; + if (num_offsets < 2 || req_id >= num_offsets - 1) { + throw std::runtime_error("MiniCPMVModel: image_req_ids is out of input_offsets range"); + } + const int32_t req_start = offsets[req_id]; + const int32_t req_end = offsets[req_id + 1]; + if (req_start < 0 || req_end < req_start) { + throw std::runtime_error("MiniCPMVModel: invalid input_offsets"); + } + auto pixel_values = input.pixel_values.value().at(media_idx); + auto tgt_sizes = input.tgt_sizes.value().at(media_idx); + auto vision_embedding = vpm_->forward(pixel_values, tgt_sizes); + auto vision_hidden = resampler_->forward(vision_embedding, tgt_sizes); + replace_embeddings( + inputs_embeds->narrow({{1, static_cast(req_start), static_cast(req_end - req_start)}}), + vision_hidden, + input.image_bound.value().at(media_idx), + input.image_embed_bound.value().at(media_idx)); } auto hidden_states = llm_->model().forward_embeds( diff --git a/csrc/models/minicpmv/minicpmv_model.hpp b/csrc/models/minicpmv/minicpmv_model.hpp index 98bfb6944..dad0dafb6 100644 --- a/csrc/models/minicpmv/minicpmv_model.hpp +++ b/csrc/models/minicpmv/minicpmv_model.hpp @@ -25,7 +25,8 @@ class MiniCPMVModel : public InfinilmModel { private: void replace_embeddings(infinicore::Tensor inputs_embeds, const infinicore::Tensor &vision_hidden, - const infinicore::Tensor &image_bound) const; + const infinicore::Tensor &image_bound, + const infinicore::Tensor &image_embed_bound) const; std::shared_ptr config_; diff --git a/csrc/models/videonsa/videonsa_for_conditional_generation.cpp b/csrc/models/videonsa/videonsa_for_conditional_generation.cpp index 54fbad26b..f092bc4cc 100644 --- a/csrc/models/videonsa/videonsa_for_conditional_generation.cpp +++ b/csrc/models/videonsa/videonsa_for_conditional_generation.cpp @@ -25,6 +25,34 @@ std::shared_ptr text_config_from(std::shared_ptr< return std::make_shared(text_config_json); } +size_t visual_length_from_grid(const infinicore::Tensor &grid_tensor, + size_t spatial_merge_size) { + if (spatial_merge_size == 0) { + throw std::runtime_error("VideoNSAForConditionalGeneration: spatial_merge_size must be positive"); + } + auto grid_cpu = grid_tensor->to(infinicore::Device::cpu()); + auto rows = grid_cpu->size(0); + auto grid_ptr = reinterpret_cast(grid_cpu->data()); + size_t total = 0; + for (size_t i = 0; i < rows; ++i) { + const int64_t grid_t_i = grid_ptr[i * 3]; + const int64_t grid_h_i = grid_ptr[i * 3 + 1]; + const int64_t grid_w_i = grid_ptr[i * 3 + 2]; + if (grid_t_i <= 0 || grid_h_i <= 0 || grid_w_i <= 0) { + throw std::runtime_error("VideoNSAForConditionalGeneration: invalid grid_thw"); + } + if (grid_h_i % static_cast(spatial_merge_size) != 0 + || grid_w_i % static_cast(spatial_merge_size) != 0) { + throw std::runtime_error("VideoNSAForConditionalGeneration: grid_thw is not divisible by spatial_merge_size"); + } + const size_t grid_t = static_cast(grid_t_i); + const size_t llm_grid_h = static_cast(grid_h_i / static_cast(spatial_merge_size)); + const size_t llm_grid_w = static_cast(grid_w_i / static_cast(spatial_merge_size)); + total += grid_t * llm_grid_h * llm_grid_w; + } + return total; +} + } // namespace VideoNSAForConditionalGeneration::VideoNSAForConditionalGeneration(std::shared_ptr model_config, @@ -42,41 +70,85 @@ VideoNSAForConditionalGeneration::VideoNSAForConditionalGeneration(std::shared_p void VideoNSAForConditionalGeneration::replace_embeddings(infinicore::Tensor inputs_embeds, const infinicore::Tensor &vision_hidden, - const infinicore::Tensor &image_bound) const { + const infinicore::Tensor &image_bound, + const infinicore::Tensor &image_embed_bound) const { auto bounds_cpu = image_bound->to(infinicore::Device::cpu()); + auto embed_bounds_cpu = image_embed_bound->to(infinicore::Device::cpu()); + ASSERT_EQ(inputs_embeds->size(0), 1); + ASSERT_EQ(bounds_cpu->size(0), 1); + ASSERT_EQ(bounds_cpu->size(2), 2); + ASSERT_EQ(embed_bounds_cpu->size(0), 1); + ASSERT_EQ(embed_bounds_cpu->size(1), bounds_cpu->size(1)); + ASSERT_EQ(embed_bounds_cpu->size(2), 2); auto out_slice = inputs_embeds->squeeze(0); auto bound_slice = bounds_cpu->squeeze(0); + auto embed_bound_slice = embed_bounds_cpu->squeeze(0); auto bound_count = bound_slice->size(0); - size_t vision_offset = 0; for (size_t i = 0; i < bound_count; ++i) { auto bound = bound_slice->narrow({{0, i, 1}}); auto bound_ptr = reinterpret_cast(bound->data()); - auto start = bound_ptr[0]; - auto end = bound_ptr[1]; - if (end <= start) { + const int64_t start = bound_ptr[0]; + const int64_t end = bound_ptr[1]; + if (start < 0 || end < start) { + throw std::runtime_error("VideoNSAForConditionalGeneration: invalid image_bound"); + } + if (start == end) { continue; } const size_t len = static_cast(end - start); - auto patch_embed = vision_hidden->narrow({{0, vision_offset, len}}); + + auto embed_bound = embed_bound_slice->narrow({{0, i, 1}}); + auto embed_bound_ptr = reinterpret_cast(embed_bound->data()); + const int64_t src_start_i = embed_bound_ptr[0]; + const int64_t src_end_i = embed_bound_ptr[1]; + if (src_start_i < 0 || src_end_i < src_start_i) { + throw std::runtime_error("VideoNSAForConditionalGeneration: invalid image_embed_bound"); + } + if (src_end_i - src_start_i != static_cast(len)) { + throw std::runtime_error("VideoNSAForConditionalGeneration: image_bound and image_embed_bound length mismatch"); + } + const size_t src_start = static_cast(src_start_i); + const size_t src_len = len; + if (static_cast(end) > out_slice->size(0) + || src_start > vision_hidden->size(0) + || src_len > vision_hidden->size(0) - src_start) { + throw std::runtime_error("VideoNSAForConditionalGeneration: multimodal embedding bounds are out of range"); + } + auto patch_embed = vision_hidden->narrow({{0, src_start, src_len}}); out_slice->narrow({{0, size_t(start), len}})->copy_from(patch_embed); - vision_offset += len; } } infinilm::InfinilmModel::Output VideoNSAForConditionalGeneration::forward(const infinilm::InfinilmModel::Input &input) const { if (input.pixel_values.has_value() && input.pixel_values.value().size() > 0) { + if (!input.input_ids.has_value()) { + throw std::runtime_error("VideoNSAForConditionalGeneration: input_ids is required"); + } if (!input.image_bound.has_value() || !input.tgt_sizes.has_value()) { throw std::runtime_error("VideoNSAForConditionalGeneration: image_bound and tgt_sizes must be provided with pixel_values"); } + if (!input.input_offsets.has_value()) { + throw std::runtime_error("VideoNSAForConditionalGeneration: input_offsets is required with pixel_values"); + } auto input_ids = input.input_ids.value(); auto inputs_embeds = model_->embed_tokens(input_ids); auto input_offsets_cpu = input.input_offsets.value()->to(infinicore::Device::cpu()); int32_t *offsets = reinterpret_cast(input_offsets_cpu->data()); - const auto &image_req_ids = global_state::get_forward_context().mm_metadata.image_req_ids.value(); + const auto &mm_metadata = global_state::get_forward_context().mm_metadata; + if (!mm_metadata.image_req_ids.has_value()) { + throw std::runtime_error("VideoNSAForConditionalGeneration: image_req_ids must be provided with pixel_values"); + } + const auto &image_req_ids = mm_metadata.image_req_ids.value(); if (input.pixel_values->size() != image_req_ids.size() || input.image_bound->size() != image_req_ids.size() || input.tgt_sizes->size() != image_req_ids.size()) { throw std::runtime_error("VideoNSAForConditionalGeneration: multimodal tensor lists must match image_req_ids"); } + if (!input.image_embed_bound.has_value()) { + throw std::runtime_error("VideoNSAForConditionalGeneration: image_embed_bound must be provided with pixel_values"); + } + if (input.image_embed_bound->size() != image_req_ids.size()) { + throw std::runtime_error("VideoNSAForConditionalGeneration: image_embed_bound must match image_req_ids"); + } std::vector pixel_tensors; std::vector grid_tensors; @@ -91,25 +163,41 @@ infinilm::InfinilmModel::Output VideoNSAForConditionalGeneration::forward(const auto batched_vision_hidden = visual_->forward(batched_pixels, batched_grids); size_t vision_offset = 0; + size_t spatial_merge_size = 2; + const auto &vision_config = model_config_->get_config_json()["vision_config"]; + if (vision_config.contains("spatial_merge_size")) { + spatial_merge_size = vision_config["spatial_merge_size"].get(); + } + const size_t num_offsets = input_offsets_cpu->size(0); for (size_t media_idx = 0; media_idx < image_req_ids.size(); ++media_idx) { const size_t req_id = image_req_ids[media_idx]; - auto bounds_cpu = input.image_bound.value().at(media_idx)->to(infinicore::Device::cpu())->squeeze(0); - auto bound_count = bounds_cpu->size(0); - auto bounds = reinterpret_cast(bounds_cpu->data()); - size_t vision_len = 0; - for (size_t i = 0; i < bound_count; ++i) { - auto start = bounds[i * 2]; - auto end = bounds[i * 2 + 1]; - if (end > start) { - vision_len += static_cast(end - start); - } + if (num_offsets < 2 || req_id >= num_offsets - 1) { + throw std::runtime_error("VideoNSAForConditionalGeneration: image_req_ids is out of input_offsets range"); + } + const int32_t req_start = offsets[req_id]; + const int32_t req_end = offsets[req_id + 1]; + if (req_start < 0 || req_end < req_start) { + throw std::runtime_error("VideoNSAForConditionalGeneration: invalid input_offsets"); + } + auto grid_tensor = input.tgt_sizes.value().at(media_idx); + const size_t vision_len = visual_length_from_grid(grid_tensor, spatial_merge_size); + if (vision_offset > batched_vision_hidden->size(0) + || vision_len > batched_vision_hidden->size(0) - vision_offset) { + throw std::runtime_error("VideoNSAForConditionalGeneration: visual hidden size does not match tgt_sizes"); } auto vision_hidden = batched_vision_hidden->narrow({{0, vision_offset, vision_len}}); - auto req_embeds = inputs_embeds->narrow({{1, size_t(offsets[req_id]), size_t(offsets[req_id + 1] - offsets[req_id])}}); - replace_embeddings(req_embeds, vision_hidden, input.image_bound.value().at(media_idx)); + auto req_embeds = inputs_embeds->narrow({{1, static_cast(req_start), static_cast(req_end - req_start)}}); + replace_embeddings( + req_embeds, + vision_hidden, + input.image_bound.value().at(media_idx), + input.image_embed_bound.value().at(media_idx)); vision_offset += vision_len; } + if (vision_offset != batched_vision_hidden->size(0)) { + throw std::runtime_error("VideoNSAForConditionalGeneration: unused visual hidden states after replacement"); + } auto hidden_states = model_->forward_embeds(inputs_embeds, input.position_ids.value()); auto logits = lm_head_->forward(hidden_states); diff --git a/csrc/models/videonsa/videonsa_for_conditional_generation.hpp b/csrc/models/videonsa/videonsa_for_conditional_generation.hpp index a31f5db87..a753ca46d 100644 --- a/csrc/models/videonsa/videonsa_for_conditional_generation.hpp +++ b/csrc/models/videonsa/videonsa_for_conditional_generation.hpp @@ -21,7 +21,8 @@ class VideoNSAForConditionalGeneration : public InfinilmModel { protected: void replace_embeddings(infinicore::Tensor inputs_embeds, const infinicore::Tensor &vision_hidden, - const infinicore::Tensor &image_bound) const; + const infinicore::Tensor &image_bound, + const infinicore::Tensor &image_embed_bound) const; INFINICORE_NN_MODULE(VideoNSATextModel, model); INFINICORE_NN_MODULE(VideoNSAVisionModel, visual); diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index 4c93c737e..70d2fee0c 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -144,6 +144,7 @@ inline void bind_infer_engine(py::module &m) { std::optional mamba_final_state_indices, std::optional> pixel_values, std::optional> image_bound, + std::optional> image_embed_bound, std::optional> tgt_sizes, std::optional> image_req_ids, std::optional> visual_token_ranges, @@ -161,6 +162,7 @@ inline void bind_infer_engine(py::module &m) { std::move(mamba_final_state_indices), std::move(pixel_values), std::move(image_bound), + std::move(image_embed_bound), std::move(tgt_sizes), std::move(image_req_ids), std::move(visual_token_ranges), @@ -209,6 +211,7 @@ inline void bind_infer_engine(py::module &m) { py::arg("mamba_final_state_indices") = std::nullopt, py::arg("pixel_values") = std::nullopt, py::arg("image_bound") = std::nullopt, + py::arg("image_embed_bound") = std::nullopt, py::arg("tgt_sizes") = std::nullopt, py::arg("image_req_ids") = std::nullopt, py::arg("visual_token_ranges") = std::nullopt) @@ -224,6 +227,7 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("mamba_final_state_indices", &InferEngine::Input::mamba_final_state_indices) .def_readwrite("pixel_values", &InferEngine::Input::pixel_values) .def_readwrite("image_bound", &InferEngine::Input::image_bound) + .def_readwrite("image_embed_bound", &InferEngine::Input::image_embed_bound) .def_readwrite("tgt_sizes", &InferEngine::Input::tgt_sizes) .def_readwrite("image_req_ids", &InferEngine::Input::image_req_ids) .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges) diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 3049d05d8..ad12daf33 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -210,6 +210,7 @@ def forward( mamba_final_state_indices=None, pixel_values=None, image_bound=None, + image_embed_bound=None, tgt_sizes=None, image_req_ids=None, visual_token_ranges=None, @@ -261,6 +262,7 @@ def convert_tensor_list(tensor_list_): pixel_values = convert_tensor_list(pixel_values) image_bound = convert_tensor_list(image_bound) + image_embed_bound = convert_tensor_list(image_embed_bound) tgt_sizes = convert_tensor_list(tgt_sizes) return infinicore.Tensor( @@ -279,6 +281,7 @@ def convert_tensor_list(tensor_list_): mamba_final_state_indices=mamba_final_state_indices, pixel_values=pixel_values, image_bound=image_bound, + image_embed_bound=image_embed_bound, tgt_sizes=tgt_sizes, image_req_ids=image_req_ids, visual_token_ranges=visual_token_ranges, @@ -300,6 +303,7 @@ def generate( *, pixel_values=None, image_bound=None, + image_embed_bound=None, tgt_sizes=None, _measure_and_log_time=False, ): @@ -439,6 +443,7 @@ def generate( mamba_init_state_indices=mamba_init_state_indices, mamba_final_state_indices=mamba_final_state_indices, image_bound=image_bound if iter == 0 else None, + image_embed_bound=image_embed_bound if iter == 0 else None, tgt_sizes=tgt_sizes if iter == 0 else None, temperature=generation_config.temperature, top_k=generation_config.top_k, diff --git a/python/infinilm/llm/cache_manager.py b/python/infinilm/llm/cache_manager.py index b10d1ce3d..cba1629fb 100644 --- a/python/infinilm/llm/cache_manager.py +++ b/python/infinilm/llm/cache_manager.py @@ -2,12 +2,17 @@ KV Cache Manager - Paged Attention block-based cache allocation and management. """ +import json from collections import deque -from typing import Dict, List, Set +from typing import Any, Dict, List, Set -import numpy as np import xxhash +from infinilm.multimodal.features import ( + MMFeature, + iter_mm_parts, +) + class Block: """KV Cache Block with reference counting and hash-based reuse support.""" @@ -88,18 +93,50 @@ def compute_hash( cls, token_ids: List[int], prefix_hash: int = -1, - mm_data_identifiers: List[str] = None, + extra_keys: tuple[Any, ...] | list[Any] | None = None, ) -> int: - """Compute hash for token sequence with optional prefix chaining.""" + """Compute xxhash64 for token IDs, prefix hash, and optional extra keys.""" h = xxhash.xxh64() - if prefix_hash != -1: - h.update(prefix_hash.to_bytes(8, "little")) - h.update(np.array(token_ids, dtype=np.int32).tobytes()) - if mm_data_identifiers is not None: - for identifier in mm_data_identifiers: - h.update(identifier.encode("utf-8")) + payload = { + "prefix_hash": int(prefix_hash), + "token_ids": [int(t) for t in token_ids], + "extra_keys": extra_keys or (), + } + h.update( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode( + "utf-8" + ) + ) return h.intdigest() + @staticmethod + def _get_mm_extra_keys_for_block( + mm_features: List[MMFeature], + block_start: int, + block_end: int, + ) -> tuple[Any, ...]: + extra_keys = [] + for feature, part in iter_mm_parts(mm_features): + if not part.overlaps(block_start, block_end): + continue + overlap_start = max(part.token_start, block_start) + overlap_end = min(part.token_end, block_end) + src_start = part.embed_start + (overlap_start - part.token_start) + src_end = part.embed_start + (overlap_end - part.token_start) + feature_key = feature.mm_hash or feature.identifier + extra_keys.append( + ( + "mm_part", + feature.modality, + feature_key, + part.part_index, + overlap_start - block_start, + src_start, + src_end - src_start, + ) + ) + return tuple(extra_keys) + def __init__(self, num_blocks: int, block_size: int): assert num_blocks > 0 and block_size > 0, ( "num_blocks and block_size must be positive" @@ -183,7 +220,7 @@ def get_total_usable_blocks(self) -> int: def get_computed_blocks( self, token_ids: List[int], - mm_token_index_mappings: List[dict] = None, + mm_features: List[MMFeature] | None = None, ) -> tuple[List[int], int, List[dict]]: """Find locally cached prefix blocks for the given token sequence. @@ -191,7 +228,7 @@ def get_computed_blocks( Args: token_ids: Input token sequence. - mm_token_index_mappings: List of multimodal token index mappings. + mm_features: Normalized multimodal features for prefix caching. Returns: A tuple of (cached_block_table, num_local_cached_tokens, blocks_blueprint): - cached_block_table: List of matched block IDs (each with ref_count incremented). @@ -202,15 +239,12 @@ def get_computed_blocks( num_blocks = (num_tokens + self.block_size - 1) // self.block_size num_full_blocks = num_tokens // self.block_size remain_tokens = num_tokens % self.block_size - mm_token_index_mappings = mm_token_index_mappings or [] - num_mm_inputs = len(mm_token_index_mappings) + mm_features = mm_features or [] # Variables cached_block_table = [] prefix_hash = -1 cache_miss = False - mm_start_counter = 0 - mm_caching_queue = deque() blocks_blueprint = [] # [{"prefix_hash": int or -1 if not a full block, "block_id": int or -1 if not cached}, ...] max_blocks_to_reuse = num_full_blocks @@ -219,23 +253,12 @@ def get_computed_blocks( end_idx = min(start_idx + self.block_size, num_tokens) block_tokens = token_ids[start_idx:end_idx] - # Process multimodal token index mappings for this block - mm_data_identifiers = [] - while ( - mm_start_counter < num_mm_inputs - and mm_token_index_mappings[mm_start_counter]["start_index"] < end_idx - and mm_token_index_mappings[mm_start_counter]["start_index"] - >= start_idx - ): - # for all mm_data whose start_index is within this block's token range, add its identifier to the list - mm_data_identifiers.append( - mm_token_index_mappings[mm_start_counter]["identifier"] - ) - mm_caching_queue.append(mm_start_counter) - mm_start_counter += 1 + extra_keys = self._get_mm_extra_keys_for_block( + mm_features, start_idx, end_idx + ) prefix_hash = ( - self.compute_hash(block_tokens, prefix_hash, mm_data_identifiers) + self.compute_hash(block_tokens, prefix_hash, extra_keys) if len(block_tokens) == self.block_size else -1 ) @@ -258,29 +281,10 @@ def get_computed_blocks( max_blocks_to_reuse = min(max_blocks_to_reuse, block_idx) cache_miss = True - if not cache_miss: - # pop fully cached mm_data - while ( - mm_caching_queue - and mm_token_index_mappings[mm_caching_queue[0]]["end_index"] - < end_idx - ): - mm_caching_queue.popleft() - blocks_blueprint.append( {"prefix_hash": prefix_hash, "block_id": cached_block_id} ) - # If there is one incomplete mm_data, tailing blocks need to fall back until all included mm_data are complete - if mm_caching_queue: - incomplete_mm = mm_token_index_mappings[mm_caching_queue.popleft()] - incomplete_mm_start = incomplete_mm[ - "start_index" - ] # Fall back until this index is no longer included in the block - max_blocks_to_reuse = min( - max_blocks_to_reuse, incomplete_mm_start // self.block_size - ) - num_local_cached_tokens = max_blocks_to_reuse * self.block_size for block_id in range(max_blocks_to_reuse): @@ -297,6 +301,7 @@ def allocate_slots( num_computed_tokens: int = 0, cached_block_table: List[int] = None, blocks_blueprint: List[dict] = None, + mm_features: List[MMFeature] = None, delay_cache_blocks: bool = False, ) -> tuple[List[int], List[int]] | None: """Allocate KV cache slots for a request (PD-disaggregation aware). @@ -353,7 +358,10 @@ def allocate_slots( if blocks_blueprint is not None and block_idx < len(blocks_blueprint): block_hash = blocks_blueprint[block_idx]["prefix_hash"] if block_hash == -1: - block_hash = self.compute_hash(block_tokens, prefix_hash) + extra_keys = self._get_mm_extra_keys_for_block( + mm_features or [], start_tok, end_tok + ) + block_hash = self.compute_hash(block_tokens, prefix_hash, extra_keys) prefix_hash = block_hash block = self._allocate_full_block() block.update(block_hash, block_tokens) @@ -377,7 +385,11 @@ def allocate_slots( return block_table, slot_mapping def append_slot( - self, block_table: List[int], num_tokens: int, total_token_ids: List[int] = None + self, + block_table: List[int], + num_tokens: int, + total_token_ids: List[int] = None, + mm_features: List[MMFeature] = None, ) -> tuple[List[int], int]: """Append slot for decode phase (generate one new token). @@ -410,7 +422,10 @@ def append_slot( else: prefix_hash = -1 - current_hash = self.compute_hash(block_tokens, prefix_hash) + extra_keys = self._get_mm_extra_keys_for_block( + mm_features or [], block_start_idx, block_end_idx + ) + current_hash = self.compute_hash(block_tokens, prefix_hash, extra_keys) last_block.update(current_hash, block_tokens) self.hash_to_block_id[current_hash] = last_block_id diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 82a4e443b..0959b00e4 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -411,7 +411,7 @@ def generate( for content in contents: request_id = f"cmpl-{uuid.uuid4().hex}" processed_inputs = None - mm_index_mappings = None + mm_features = None if apply_chat_template: prompt = self.engine.apply_chat_template( content, add_generation_prompt=True @@ -428,8 +428,9 @@ def generate( ) prompt_token_ids = processed_inputs.get("input_ids").flatten().tolist() - mm_index_mappings = self.engine.processor.get_mm_token_index_list( + mm_features = self.engine.processor.get_mm_features( prompt_token_ids, + processed_inputs=processed_inputs, image_ids=mm_inputs["image_urls"], video_ids=mm_inputs["video_urls"], audio_ids=mm_inputs["audio_urls"], @@ -443,7 +444,7 @@ def generate( prompt=prompt, prompt_token_ids=prompt_token_ids, processed_inputs=processed_inputs, - mm_token_index_mappings=mm_index_mappings, + mm_features=mm_features, sampling_params=sampling_params, eos_token_ids=self.engine.eos_token_ids, ) @@ -743,7 +744,7 @@ def add_request( if request_id is None: request_id = f"cmpl-{uuid.uuid4().hex}" - mm_index_mappings = None + mm_features = None processed_inputs = None if prompt_token_ids is not None: @@ -774,8 +775,9 @@ def add_request( ) prompt_token_ids = processed_inputs.get("input_ids").flatten().tolist() - mm_index_mappings = self.engine.processor.get_mm_token_index_list( + mm_features = self.engine.processor.get_mm_features( prompt_token_ids, + processed_inputs=processed_inputs, image_ids=mm_inputs["image_urls"], video_ids=mm_inputs["video_urls"], audio_ids=mm_inputs["audio_urls"], @@ -792,7 +794,7 @@ def add_request( prompt=prompt, prompt_token_ids=prompt_token_ids, processed_inputs=processed_inputs, - mm_token_index_mappings=mm_index_mappings, + mm_features=mm_features, sampling_params=sampling_params, eos_token_ids=self.engine.eos_token_ids, request_data=request_data, diff --git a/python/infinilm/llm/request.py b/python/infinilm/llm/request.py index fee1e90cd..b45557f95 100644 --- a/python/infinilm/llm/request.py +++ b/python/infinilm/llm/request.py @@ -12,6 +12,7 @@ import janus from infinilm.llm.sampling_params import SamplingParams +from infinilm.multimodal.features import MMFeature logger = logging.getLogger(__name__) @@ -120,7 +121,7 @@ def __init__( prompt: Optional[str] = None, prompt_token_ids: Optional[List[int]] = None, processed_inputs: Optional[dict] = None, - mm_token_index_mappings: Optional[List[dict]] = None, + mm_features: Optional[List[MMFeature]] = None, sampling_params: Optional[SamplingParams] = None, eos_token_ids: Optional[List[int]] = None, arrival_time: Optional[float] = None, @@ -138,7 +139,9 @@ def __init__( ) self.prompt_length: int = len(self.prompt_token_ids) self.processed_inputs: Optional[dict] = processed_inputs - self.mm_token_index_mappings: Optional[List[dict]] = mm_token_index_mappings + self.mm_features: List[MMFeature] = ( + mm_features if mm_features is not None else [] + ) self.priority: int = 0 # Sampling & stopping criteria @@ -161,6 +164,8 @@ def __init__( ) self.num_computed_tokens: int = 0 # Total tokens computed (local + remote) self.num_blocks: int = 0 + self.scheduled_prefill_start: int = 0 + self.scheduled_prefill_end: int = self.prompt_length # Mamba cache management. None means no mamba cache row is currently owned. self.mamba_cache_index: Optional[int] = None @@ -208,8 +213,19 @@ def get_num_blocks_required(self, block_size: int) -> int: def get_max_tokens(self) -> Optional[int]: return self.sampling_params.max_tokens - def get_mm_token_index_mappings(self) -> Optional[List[dict]]: - return self.mm_token_index_mappings + def get_mm_features(self) -> List[MMFeature]: + return self.mm_features + + def set_scheduled_prefill_window(self, start: int, end: int) -> None: + assert 0 <= start <= end <= self.prompt_length, ( + f"Invalid prefill window [{start}, {end}) for prompt length " + f"{self.prompt_length}" + ) + self.scheduled_prefill_start = start + self.scheduled_prefill_end = end + + def get_scheduled_prefill_window(self) -> tuple[int, int]: + return self.scheduled_prefill_start, self.scheduled_prefill_end def is_finished(self) -> bool: return self.status in [ diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index 32b5fd937..bed962896 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -128,7 +128,7 @@ def schedule(self) -> Optional[SchedulerOutput]: num_local_computed_tokens, blocks_blueprint, ) = self.cache_manager.get_computed_blocks( - req_tokens, req.get_mm_token_index_mappings() + req_tokens, mm_features=req.get_mm_features(), ) if self.connector is not None: ext_tokens, load_kv_async = ( @@ -147,6 +147,9 @@ def schedule(self) -> Optional[SchedulerOutput]: if load_kv_async: num_computed_tokens -= 1 num_new_tokens = req.get_prompt_length() - num_computed_tokens + req.set_scheduled_prefill_window( + num_computed_tokens, req.get_prompt_length() + ) # Early token budget check: skip can_accept_request and allocate_slots # for requests that would exceed the per-schedule token budget. @@ -185,6 +188,7 @@ def schedule(self) -> Optional[SchedulerOutput]: num_computed_tokens=num_computed_tokens, cached_block_table=cached_block_table, blocks_blueprint=blocks_blueprint, + mm_features=req.get_mm_features(), delay_cache_blocks=load_kv_async, ) @@ -224,9 +228,11 @@ def schedule(self) -> Optional[SchedulerOutput]: ) else: load_kv_async = False - num_tokens_this_step = ( - req.get_prompt_length() - req.num_local_cached_tokens + req.set_scheduled_prefill_window( + req.num_computed_tokens, req.get_prompt_length() ) + compute_start, compute_end = req.get_scheduled_prefill_window() + num_tokens_this_step = compute_end - compute_start if self._exceeds_token_budget( current_num_batched_tokens, num_tokens_this_step, @@ -249,7 +255,8 @@ def schedule(self) -> Optional[SchedulerOutput]: current_prefill_extra_blocks += self._get_prefill_extra_blocks(req) scheduled_requests.append(req) - num_tokens_this_step = req.get_prompt_length() - req.num_local_cached_tokens + compute_start, compute_end = req.get_scheduled_prefill_window() + num_tokens_this_step = compute_end - compute_start current_num_batched_tokens += num_tokens_this_step req.status = RequestStatus.RUNNING @@ -284,7 +291,10 @@ def schedule(self) -> Optional[SchedulerOutput]: # Decode phase: allocate slot for newly generated token try: req.block_table, new_slot = self.cache_manager.append_slot( - req.block_table, req.get_total_length(), req.get_all_token_ids() + req.block_table, + req.get_total_length(), + req.get_all_token_ids(), + req.get_mm_features(), ) req.slot_mapping = [new_slot] req.num_blocks = len(req.block_table) diff --git a/python/infinilm/multimodal/features.py b/python/infinilm/multimodal/features.py new file mode 100644 index 000000000..883a9eddb --- /dev/null +++ b/python/infinilm/multimodal/features.py @@ -0,0 +1,68 @@ +"""Lightweight multimodal feature metadata used by prefix caching.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Iterable, Iterator + + +@dataclass(frozen=True) +class MMPlaceholderRange: + """Placeholder token range for a multimodal input.""" + + offset: int + length: int + is_embed: Any | None = None + + @property + def end(self) -> int: + return self.offset + self.length + + def overlap(self, start: int, end: int) -> bool: + return end > self.offset and start < self.end + + def overlaps(self, start: int, end: int) -> bool: + return self.overlap(start, end) + + +@dataclass(frozen=True) +class MMFeaturePart: + """Model-specific part of a multimodal feature.""" + + token_start: int + token_end: int + embed_start: int + embed_end: int + item_index: int + part_index: int + data_index: int | None = None + + @property + def token_length(self) -> int: + return self.token_end - self.token_start + + def overlap(self, start: int, end: int) -> bool: + return end > self.token_start and start < self.token_end + + def overlaps(self, start: int, end: int) -> bool: + return self.overlap(start, end) + + +@dataclass(frozen=True) +class MMFeature: + """A cache-addressable multimodal input feature.""" + + modality: str + identifier: str + position: MMPlaceholderRange + mm_hash: str | None = None + data: Any | None = None + parts: tuple[MMFeaturePart, ...] = field(default_factory=tuple) + + +def iter_mm_parts( + mm_features: Iterable[MMFeature], +) -> Iterator[tuple[MMFeature, MMFeaturePart]]: + for feature in mm_features: + for part in feature.parts: + yield feature, part diff --git a/python/infinilm/processors/basic_llm_processor.py b/python/infinilm/processors/basic_llm_processor.py index a6fbc33ac..70bae9801 100644 --- a/python/infinilm/processors/basic_llm_processor.py +++ b/python/infinilm/processors/basic_llm_processor.py @@ -207,7 +207,8 @@ def _build_model_input_from_batch_scheduler_output( if scheduler_output.is_prefill: # Prefill phase req_tokens = req.get_input_tokens() - tokens_to_compute = req_tokens[num_cached:] + compute_start, compute_end = req.get_scheduled_prefill_window() + tokens_to_compute = req_tokens[compute_start:compute_end] tokens.extend(tokens_to_compute) compute_len = len(tokens_to_compute) @@ -218,8 +219,8 @@ def _build_model_input_from_batch_scheduler_output( seq_offsets.append(current_offset) slot_mapping.extend(req.slot_mapping) - cached_lens.append(num_cached) - position_ids.extend(range(num_cached, num_cached + compute_len)) + cached_lens.append(compute_start) + position_ids.extend(range(compute_start, compute_start + compute_len)) else: # Decode phase @@ -271,3 +272,15 @@ def get_mm_token_index_list( self, prompt_token_ids, image_ids=None, video_ids=None, audio_ids=None, **kwargs ): return [] + + @override + def get_mm_features( + self, + prompt_token_ids, + processed_inputs=None, + image_ids=None, + video_ids=None, + audio_ids=None, + **kwargs, + ): + return [] diff --git a/python/infinilm/processors/minicpmv_processor.py b/python/infinilm/processors/minicpmv_processor.py index e4f879c9c..821a78a22 100644 --- a/python/infinilm/processors/minicpmv_processor.py +++ b/python/infinilm/processors/minicpmv_processor.py @@ -1,6 +1,13 @@ from transformers import AutoConfig, AutoProcessor from typing_extensions import override +from infinilm.multimodal.features import ( + MMFeature, + MMFeaturePart, + MMPlaceholderRange, + iter_mm_parts, +) + from .processor import InfinilmProcessor, register_processor @@ -96,6 +103,158 @@ def apply_chat_template( **kwargs, ) + @staticmethod + def _append_mm_data(mm_data: dict, key: str, value): + mm_data.setdefault(key, []).append(value) + + @staticmethod + def _flatten_pixel_values(pixel_values): + flat = [] + if pixel_values is None: + return flat + if not isinstance(pixel_values, (list, tuple)): + return [pixel_values] + for item in pixel_values: + if isinstance(item, (list, tuple)): + flat.extend(item) + else: + flat.append(item) + return flat + + @staticmethod + def _flatten_tensor_rows(values): + import torch + + flat = [] + + def visit(value): + if value is None: + return + if isinstance(value, torch.Tensor): + if value.ndim <= 1: + flat.append(value) + else: + flat.extend(row for row in value.reshape(-1, value.shape[-1])) + return + if isinstance(value, (list, tuple)): + for item in value: + visit(item) + return + raise TypeError(f"Unsupported multimodal tensor row type: {type(value)}") + + visit(values) + return flat + + def _append_request_mm_data( + self, + mm_data: dict, + req, + req_batch_index: int, + compute_start: int, + compute_end: int, + ) -> None: + if req.processed_inputs is None or "pixel_values" not in req.processed_inputs: + return + + import torch + import infinicore + + flat_pixel_values = self._flatten_pixel_values( + req.processed_inputs.get("pixel_values") + ) + flat_tgt_sizes = self._flatten_tensor_rows( + req.processed_inputs.get("tgt_sizes") + ) + + selected = [] + for feature, part in iter_mm_parts(req.get_mm_features()): + if feature.modality != "image": + continue + copy_start = max(part.token_start, compute_start) + copy_end = min(part.token_end, compute_end) + if copy_start >= copy_end: + continue + + data_index = part.data_index if part.data_index is not None else part.part_index + if data_index >= len(flat_pixel_values): + raise IndexError( + f"MM part data_index={data_index} exceeds pixel_values " + f"count {len(flat_pixel_values)}" + ) + if data_index >= len(flat_tgt_sizes): + raise IndexError( + f"MM part data_index={data_index} exceeds tgt_sizes " + f"count {len(flat_tgt_sizes)}" + ) + + dst_start = copy_start - compute_start + dst_end = copy_end - compute_start + src_start = part.embed_start + (copy_start - part.token_start) + src_end = part.embed_start + (copy_end - part.token_start) + if (dst_end - dst_start) != (src_end - src_start): + raise ValueError( + "MiniCPM-V multimodal source and destination lengths differ: " + f"dst=[{dst_start},{dst_end}), src=[{src_start},{src_end})" + ) + + selected.append((data_index, dst_start, dst_end, src_start, src_end)) + + if not selected: + return + + selected_pixel_values = [ + flat_pixel_values[data_index] + .flatten(end_dim=1) + .permute(1, 0) + .to(self.pixel_values_dtype) + for data_index, _, _, _, _ in selected + ] + pixel_values_tensor = torch.nn.utils.rnn.pad_sequence( + selected_pixel_values, batch_first=True, padding_value=0.0 + ) + batch_size, seq_len, _ = pixel_values_tensor.shape + pixel_values_tensor = ( + pixel_values_tensor.permute(0, 2, 1) + .reshape(batch_size, 3, -1, seq_len) + .contiguous() + ) + + selected_tgt_sizes = [ + flat_tgt_sizes[data_index].reshape(-1) + for data_index, _, _, _, _ in selected + ] + for tgt_size in selected_tgt_sizes: + if tgt_size.numel() != 2: + raise ValueError( + "MiniCPM-V tgt_sizes row must contain exactly two values, " + f"got {tgt_size.numel()}" + ) + tgt_sizes_tensor = torch.vstack(selected_tgt_sizes).to(torch.int64) + image_bound_tensor = torch.tensor( + [[dst_start, dst_end] for _, dst_start, dst_end, _, _ in selected], + dtype=torch.int64, + ).unsqueeze(0) + image_embed_bound_tensor = torch.tensor( + [[src_start, src_end] for _, _, _, src_start, src_end in selected], + dtype=torch.int64, + ).unsqueeze(0) + + self._append_mm_data( + mm_data, "pixel_values", infinicore.from_torch(pixel_values_tensor) + ) + self._append_mm_data( + mm_data, "tgt_sizes", infinicore.from_torch(tgt_sizes_tensor) + ) + self._append_mm_data( + mm_data, "image_bound", infinicore.from_torch(image_bound_tensor) + ) + self._append_mm_data( + mm_data, + "image_embed_bound", + infinicore.from_torch(image_embed_bound_tensor), + ) + self._append_mm_data(mm_data, "image_req_ids", req_batch_index) + @override def build_model_inputs( self, @@ -129,11 +288,11 @@ def build_model_inputs( current_offset = 0 for req_id, req in enumerate(scheduler_output.scheduled_requests): - num_cached = req.num_local_cached_tokens if scheduler_output.is_prefill: # Prefill phase req_tokens = req.get_input_tokens() - tokens_to_compute = req_tokens[num_cached:] + compute_start, compute_end = req.get_scheduled_prefill_window() + tokens_to_compute = req_tokens[compute_start:compute_end] tokens.extend(tokens_to_compute) compute_len = len(tokens_to_compute) @@ -144,87 +303,11 @@ def build_model_inputs( seq_offsets.append(current_offset) slot_mapping.extend(req.slot_mapping) - cached_lens.append(num_cached) - position_ids.extend(range(num_cached, num_cached + compute_len)) - - if ( - req.processed_inputs is not None - and "pixel_values" in req.processed_inputs - ): - import torch - - num_cached_patch = ( - (req.processed_inputs["image_bound"][0][:, 1] <= num_cached) - .sum() - .item() - ) - - # if all patches are already cached, skip processing multimodal inputs and return text-only inputs for this request - if num_cached_patch < len(req.processed_inputs["pixel_values"]): - # 1. pixel_values - all_pixel_values = [] - pixel_values = req.processed_inputs["pixel_values"] - tgt_sizes = req.processed_inputs["tgt_sizes"] - image_bound = req.processed_inputs["image_bound"] - for pv in pixel_values: - all_pixel_values.extend( - [ - t.flatten(end_dim=1) - .permute(1, 0) - .to(self.pixel_values_dtype) - for i, t in enumerate(pv) - if i >= num_cached_patch - ] - ) - - pixel_values_tensor = torch.nn.utils.rnn.pad_sequence( - all_pixel_values, batch_first=True, padding_value=0.0 - ) - B, L, _ = pixel_values_tensor.shape - pixel_values_tensor = ( - pixel_values_tensor.permute(0, 2, 1) - .reshape(B, 3, -1, L) - .contiguous() - ) - pixel_values_infini = infinicore.from_torch(pixel_values_tensor) - - # 2. tgt_sizes - all_tgt_sizes = [ - tgt_size - for i, tgt_size in enumerate(tgt_sizes) - if isinstance(tgt_size, torch.Tensor) - and i >= num_cached_patch - ] - - tgt_sizes_tensor = torch.vstack(all_tgt_sizes).to(torch.int64) - - tgt_sizes_infini = infinicore.from_torch(tgt_sizes_tensor) - - # 3. image_bound - batch_size = len(image_bound) - max_ranges = max(len(b) for b in image_bound) - - bound = torch.zeros( - (batch_size, max_ranges, 2), dtype=torch.int64 - ) - - for i, bnd in enumerate(image_bound): - bnd = bnd[num_cached_patch:, :] - if len(bnd) > 0: - bound[i, : len(bnd), :] = bnd - - image_bound_infini = infinicore.from_torch(bound) - - def append_mm_data(mm_data__: dict, key__: str, value__): - if mm_data__.get(key__) is None: - mm_data[key__] = [value__] - else: - mm_data[key__].append(value__) - - append_mm_data(mm_data, "pixel_values", pixel_values_infini) - append_mm_data(mm_data, "tgt_sizes", tgt_sizes_infini) - append_mm_data(mm_data, "image_bound", image_bound_infini) - append_mm_data(mm_data, "image_req_ids", req_id) + cached_lens.append(compute_start) + position_ids.extend(range(compute_start, compute_start + compute_len)) + self._append_request_mm_data( + mm_data, req, req_id, compute_start, compute_end + ) else: # Decode phase @@ -241,7 +324,7 @@ def append_mm_data(mm_data__: dict, key__: str, value__): seq_offsets.append(current_offset) slot_mapping.extend(req.slot_mapping) - cached_lens.append(num_cached) + cached_lens.append(req.num_local_cached_tokens) position_ids.append(seq_len - 1) # Pad block_table to same length @@ -277,50 +360,111 @@ def get_tokenizer(self): def get_mm_token_index_list( self, prompt_token_ids, image_ids=None, video_ids=None, **kwargs ): + image_ids = image_ids or [] image_idx = -1 - patch_start = [] - patch_end = [] + data_index = 0 + patch_start = None + part_index = 0 mm_token_index_list = [] for i, token_id in enumerate(prompt_token_ids): if token_id == self.tokenizer.im_id_start_id: - assert len(patch_start) == len(patch_end), ( - "Invalid prompt format: image start token found before previous image end token is closed" - ) - # deal with previous image patches - if patch_start: - for start, end in zip(patch_start, patch_end): - mm_token_index_list.append( - { - "start_index": start, - "end_index": end, - "identifier": image_ids[image_idx], - } - ) - # reset patch start and end for next image - patch_start = [] - patch_end = [] - - # increment image index for next image image_idx += 1 - patch_start.append(i + 1) - elif token_id == self.tokenizer.slice_start_id: - patch_start.append(i + 1) - elif ( - token_id == self.tokenizer.im_id_end_id - or token_id == self.tokenizer.slice_end_id + part_index = 0 + elif token_id in ( + self.tokenizer.im_start_id, + self.tokenizer.slice_start_id, ): - patch_end.append(i - 1) - - if patch_start: - for start, end in zip(patch_start, patch_end): + if image_idx < 0: + image_idx = 0 + assert patch_start is None, ( + "Invalid prompt format: nested MiniCPM-V image placeholder" + ) + patch_start = i + 1 + elif token_id in ( + self.tokenizer.im_end_id, + self.tokenizer.slice_end_id, + ): + assert patch_start is not None, ( + "Invalid prompt format: MiniCPM-V image placeholder end found " + "before start" + ) mm_token_index_list.append( { - "start_index": start, - "end_index": end, + "start_index": patch_start, + "end_index": i - 1, "identifier": image_ids[image_idx], + "modality": "image", + "item_index": image_idx, + "part_index": part_index, + "data_index": data_index, } ) + data_index += 1 + part_index += 1 + patch_start = None + + assert patch_start is None, "Invalid prompt format: unclosed image placeholder" assert image_idx + 1 == len(image_ids), ( "The number of image tokens does not match the number of images data provided" ) return mm_token_index_list + + @override + def get_mm_features( + self, + prompt_token_ids, + processed_inputs=None, + image_ids=None, + video_ids=None, + audio_ids=None, + **kwargs, + ): + mappings = self.get_mm_token_index_list( + prompt_token_ids, + image_ids=image_ids, + video_ids=video_ids, + audio_ids=audio_ids, + **kwargs, + ) + image_bound_rows = [] + if processed_inputs is not None and "image_bound" in processed_inputs: + image_bound_rows = self._flatten_tensor_rows(processed_inputs["image_bound"]) + if len(image_bound_rows) != len(mappings): + raise ValueError( + "MiniCPM-V image_bound rows do not match placeholder ranges: " + f"{len(image_bound_rows)} != {len(mappings)}" + ) + + features = [] + for data_index, mapping in enumerate(mappings): + start = int(mapping["start_index"]) + end = int(mapping["end_index"]) + 1 + if image_bound_rows: + bound = image_bound_rows[data_index] + bound_start = int(bound[0].item()) + bound_end = int(bound[1].item()) + if (bound_start, bound_end) != (start, end): + raise ValueError( + "MiniCPM-V token placeholder range does not match " + f"image_bound row {data_index}: token=[{start},{end}), " + f"image_bound=[{bound_start},{bound_end})" + ) + + part = MMFeaturePart( + token_start=start, + token_end=end, + embed_start=0, + embed_end=end - start, + item_index=int(mapping["item_index"]), + part_index=int(mapping["part_index"]), + data_index=int(mapping.get("data_index", data_index)), + ) + features.append( + MMFeature( + modality="image", + identifier=str(mapping["identifier"]), + position=MMPlaceholderRange(offset=start, length=end - start), + parts=(part,), + ) + ) + return features diff --git a/python/infinilm/processors/processor.py b/python/infinilm/processors/processor.py index a2952bc1e..7ba08ba5a 100644 --- a/python/infinilm/processors/processor.py +++ b/python/infinilm/processors/processor.py @@ -37,11 +37,25 @@ def get_mm_token_index_list( self, prompt_token_ids, image_ids=None, video_ids=None, audio_ids=None, **kwargs ): """ - Get the list of starting token index and identifier mapping for multimodal inputs, sorted by index. + Legacy/helper API for placeholder range parsing. + + New scheduling and prefix-cache code should use get_mm_features(). Return: [{"start_index": , "identifier": }, ...] """ raise NotImplementedError("get_mm_token_index_list is not implemented yet") + def get_mm_features( + self, + prompt_token_ids, + processed_inputs=None, + image_ids=None, + video_ids=None, + audio_ids=None, + **kwargs, + ): + """Return normalized multimodal features for scheduling and prefix caching.""" + raise NotImplementedError("get_mm_features is not implemented yet") + # Global registry mapping model_type strings to their Processor classes _PROCESSOR_REGISTRY = {} diff --git a/python/infinilm/processors/videonsa_processor.py b/python/infinilm/processors/videonsa_processor.py index 781bf3141..5f04bf30e 100644 --- a/python/infinilm/processors/videonsa_processor.py +++ b/python/infinilm/processors/videonsa_processor.py @@ -4,6 +4,13 @@ from transformers import AutoConfig, AutoProcessor from typing_extensions import override +from infinilm.multimodal.features import ( + MMFeature, + MMFeaturePart, + MMPlaceholderRange, + iter_mm_parts, +) + from .processor import InfinilmProcessor, register_processor @@ -226,6 +233,127 @@ def _video_second_per_grid(self, processed_inputs, video_idx): except Exception: return 1.0 + @staticmethod + def _append_mm_data(mm_data: dict, key: str, value): + mm_data.setdefault(key, []).append(value) + + def _media_payloads(self, processed_inputs): + payloads = {"image": [], "video": []} + if processed_inputs is None: + return payloads + + pixel_values = processed_inputs.get("pixel_values") + image_grids = self._grid_list(processed_inputs, "image_grid_thw") + if pixel_values is not None: + offset = 0 + for grid in image_grids: + grid_tensor = torch.as_tensor(grid, dtype=torch.int64) + patch_count = int(grid_tensor.prod().item()) + payloads["image"].append( + (pixel_values[offset : offset + patch_count], grid_tensor) + ) + offset += patch_count + + pixel_values_videos = processed_inputs.get("pixel_values_videos") + video_grids = self._grid_list(processed_inputs, "video_grid_thw") + if pixel_values_videos is not None: + offset = 0 + for grid in video_grids: + grid_tensor = torch.as_tensor(grid, dtype=torch.int64) + patch_count = int(grid_tensor.prod().item()) + payloads["video"].append( + (pixel_values_videos[offset : offset + patch_count], grid_tensor) + ) + offset += patch_count + + return payloads + + def _append_request_mm_data( + self, + mm_data: dict, + req, + req_batch_index: int, + compute_start: int, + compute_end: int, + packed_request_start: int, + ) -> None: + if req.processed_inputs is None or "image_bound" not in req.processed_inputs: + return + + import infinicore + + payloads = self._media_payloads(req.processed_inputs) + selected = [] + + for feature, part in iter_mm_parts(req.get_mm_features()): + if feature.modality not in payloads: + continue + + copy_start = max(part.token_start, compute_start) + copy_end = min(part.token_end, compute_end) + if copy_start >= copy_end: + continue + + item_index = part.item_index + if item_index >= len(payloads[feature.modality]): + raise IndexError( + f"VideoNSA {feature.modality} item_index={item_index} exceeds " + f"processed input count {len(payloads[feature.modality])}" + ) + + dst_start = copy_start - compute_start + dst_end = copy_end - compute_start + src_start = part.embed_start + (copy_start - part.token_start) + src_end = part.embed_start + (copy_end - part.token_start) + selected.append( + ( + feature.modality, + payloads[feature.modality][item_index], + dst_start, + dst_end, + src_start, + src_end, + ) + ) + + if not selected: + return + + for modality, payload, dst_start, dst_end, src_start, src_end in selected: + pixel_tensor = payload[0].to(self.pixel_values_dtype).contiguous() + grid_tensor = payload[1].reshape(1, 3).to(torch.int64) + image_bound_tensor = torch.tensor( + [[[dst_start, dst_end]]], + dtype=torch.int64, + ) + image_embed_bound_tensor = torch.tensor( + [[[src_start, src_end]]], + dtype=torch.int64, + ) + + self._append_mm_data( + mm_data, "pixel_values", infinicore.from_torch(pixel_tensor) + ) + self._append_mm_data( + mm_data, "tgt_sizes", infinicore.from_torch(grid_tensor) + ) + self._append_mm_data( + mm_data, "image_bound", infinicore.from_torch(image_bound_tensor) + ) + self._append_mm_data( + mm_data, + "image_embed_bound", + infinicore.from_torch(image_embed_bound_tensor), + ) + self._append_mm_data(mm_data, "image_req_ids", req_batch_index) + + mm_data.setdefault("visual_token_ranges", []).extend( + [ + packed_request_start + int(dst_start), + packed_request_start + int(dst_end), + ] + ) + def _append_text_positions(self, axes, length, start_pos): if length <= 0: return start_pos @@ -374,88 +502,32 @@ def build_model_inputs( num_cached = req.num_local_cached_tokens if scheduler_output.is_prefill: req_tokens = req.get_input_tokens() - tokens_to_compute = req_tokens[num_cached:] + compute_start, compute_end = req.get_scheduled_prefill_window() + tokens_to_compute = req_tokens[compute_start:compute_end] tokens.extend(tokens_to_compute) compute_len = len(tokens_to_compute) seq_len = len(req_tokens) seq_lens.append(seq_len) + packed_request_start = current_offset current_offset += compute_len seq_offsets.append(current_offset) - cached_lens.append(num_cached) + cached_lens.append(compute_start) prompt_positions = self._prompt_mrope_positions( req_tokens, req.processed_inputs ) for axis in range(3): position_axes[axis].extend( - prompt_positions[axis][num_cached : num_cached + compute_len] + prompt_positions[axis][compute_start:compute_end] ) - if ( - req.processed_inputs is not None - and "image_bound" in req.processed_inputs - ): - image_bound = req.processed_inputs["image_bound"] - bounds = image_bound[0] - bounds = bounds[bounds[:, 1] > bounds[:, 0]] - num_cached_media = (bounds[:, 1] <= num_cached).sum().item() - if num_cached_media < len(bounds): - grids = [] - tensors = [] - offset = 0 - pixel_values = req.processed_inputs.get("pixel_values") - image_grid = req.processed_inputs.get("image_grid_thw") - if pixel_values is not None: - for media_idx, grid in enumerate(image_grid): - patch_count = int(grid.prod().item()) - if media_idx >= num_cached_media: - grids.append(grid) - tensors.append( - pixel_values[offset : offset + patch_count] - ) - offset += patch_count - pixel_values_videos = req.processed_inputs.get( - "pixel_values_videos" - ) - video_grid = req.processed_inputs.get("video_grid_thw") - if pixel_values_videos is not None: - offset = 0 - base = 0 if image_grid is None else len(image_grid) - for media_idx, grid in enumerate(video_grid): - patch_count = int(grid.prod().item()) - if base + media_idx >= num_cached_media: - grids.append(grid) - tensors.append( - pixel_values_videos[ - offset : offset + patch_count - ] - ) - offset += patch_count - if tensors: - pixel_tensor = torch.cat(tensors, dim=0).to( - self.pixel_values_dtype - ) - grid_tensor = torch.stack(grids).to(torch.int64) - bound = ( - bounds[num_cached_media:].unsqueeze(0).to(torch.int64) - ) - mm_data.setdefault("pixel_values", []).append( - infinicore.from_torch(pixel_tensor) - ) - mm_data.setdefault("tgt_sizes", []).append( - infinicore.from_torch(grid_tensor) - ) - mm_data.setdefault("image_bound", []).append( - infinicore.from_torch(bound) - ) - mm_data.setdefault("image_req_ids", []).append(req_id) - if pixel_values_videos is not None: - for start, end in bound.squeeze(0).tolist(): - if end > start: - packed_start = seq_offsets[-2] + int(start) - packed_end = seq_offsets[-2] + int(end) - mm_data.setdefault( - "visual_token_ranges", [] - ).extend([packed_start, packed_end]) + self._append_request_mm_data( + mm_data, + req, + req_id, + compute_start, + compute_end, + packed_request_start, + ) else: seq_len = req.get_total_length() last_token = ( @@ -516,40 +588,149 @@ def get_tokenizer(self): def get_mm_token_index_list( self, prompt_token_ids, image_ids=None, video_ids=None, **kwargs ): - ids = list(image_ids or []) + list(video_ids or []) + image_ids = list(image_ids or []) + video_ids = list(video_ids or []) out = [] - media_idx = 0 + image_idx = 0 + video_idx = 0 + data_index = 0 start = None current = None + current_modality = None + current_item_index = None + + def append_current(end): + nonlocal data_index + if start is None: + return + if current_modality == "image": + identifier = image_ids[current_item_index] + else: + identifier = video_ids[current_item_index] + out.append( + { + "start_index": start, + "end_index": end, + "end": end, + "identifier": identifier, + "modality": current_modality, + "item_index": current_item_index, + "part_index": 0, + "data_index": data_index, + } + ) + data_index += 1 + for i, token_id in enumerate(prompt_token_ids): if token_id in (self.image_token_id, self.video_token_id): + modality = "image" if token_id == self.image_token_id else "video" if start is None: start = i current = token_id + current_modality = modality + if modality == "image": + if image_idx >= len(image_ids): + raise ValueError( + "The number of image tokens exceeds image data provided" + ) + current_item_index = image_idx + image_idx += 1 + else: + if video_idx >= len(video_ids): + raise ValueError( + "The number of video tokens exceeds video data provided" + ) + current_item_index = video_idx + video_idx += 1 elif current != token_id: - out.append( - { - "start_index": start, - "end_index": i, - "identifier": ids[media_idx], - } - ) - media_idx += 1 + append_current(i) start = i current = token_id + current_modality = modality + if modality == "image": + if image_idx >= len(image_ids): + raise ValueError( + "The number of image tokens exceeds image data provided" + ) + current_item_index = image_idx + image_idx += 1 + else: + if video_idx >= len(video_ids): + raise ValueError( + "The number of video tokens exceeds video data provided" + ) + current_item_index = video_idx + video_idx += 1 elif start is not None: - out.append( - {"start_index": start, "end_index": i, "identifier": ids[media_idx]} - ) - media_idx += 1 + append_current(i) start = None current = None + current_modality = None + current_item_index = None if start is not None: - out.append( - { - "start_index": start, - "end_index": len(prompt_token_ids), - "identifier": ids[media_idx], - } + append_current(len(prompt_token_ids)) + if image_idx != len(image_ids): + raise ValueError( + "The number of image tokens does not match the number of images data provided" + ) + if video_idx != len(video_ids): + raise ValueError( + "The number of video tokens does not match the number of videos data provided" ) return out + + @override + def get_mm_features( + self, + prompt_token_ids, + processed_inputs=None, + image_ids=None, + video_ids=None, + audio_ids=None, + **kwargs, + ): + mappings = self.get_mm_token_index_list( + prompt_token_ids, + image_ids=image_ids, + video_ids=video_ids, + audio_ids=audio_ids, + **kwargs, + ) + bounds = self._valid_media_bounds(processed_inputs) + if bounds and len(bounds) != len(mappings): + raise ValueError( + "VideoNSA image_bound rows do not match placeholder ranges: " + f"{len(bounds)} != {len(mappings)}" + ) + + features = [] + for data_index, mapping in enumerate(mappings): + start = int(mapping["start_index"]) + end = int(mapping["end"]) + if bounds: + bound_start, bound_end = bounds[data_index] + if (bound_start, bound_end) != (start, end): + raise ValueError( + "VideoNSA token placeholder range does not match " + f"image_bound row {data_index}: token=[{start},{end}), " + f"image_bound=[{bound_start},{bound_end})" + ) + + part = MMFeaturePart( + token_start=start, + token_end=end, + embed_start=0, + embed_end=end - start, + item_index=int(mapping["item_index"]), + part_index=int(mapping["part_index"]), + data_index=int(mapping.get("data_index", data_index)), + ) + features.append( + MMFeature( + modality=str(mapping["modality"]), + identifier=str(mapping["identifier"]), + position=MMPlaceholderRange(offset=start, length=end - start), + parts=(part,), + ) + ) + return features