Skip to content

issue/477 - feat: support partial cache hits for multimodal visual embeddings #484

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ma-hang wants to merge 1 commit into
base: main
Choose a base branch
from
Open

issue/477 - feat: support partial cache hits for multimodal visual embeddings #484

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions csrc/engine/infer_engine.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
Expand Down
5 changes: 5 additions & 0 deletions csrc/engine/rank_worker.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,13 @@ class RankWorker {
/// Image pixel values for multi-modal models.
std::optional<std::vector<infinicore::Tensor>> pixel_values;
/// Image placeholder bounds for MiniCPM-V style replacement.
/// Vector of tensors shape: [1, n_patch, 2].
std::optional<std::vector<infinicore::Tensor>> image_bound;
/// Source embedding bounds for partial multimodal prefill replacement.
/// Vector of tensors shape: [1, n_patch, 2].
std::optional<std::vector<infinicore::Tensor>> 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<std::vector<infinicore::Tensor>> tgt_sizes;
/// req_id for each pixel_values among a batch
std::optional<std::vector<size_t>> image_req_ids;
Expand Down
7 changes: 5 additions & 2 deletions csrc/models/infinilm_model.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::vector<infinicore::Tensor>> 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<std::vector<infinicore::Tensor>> image_bound;
/// Source embedding bounds for partial multimodal prefill replacement.
/// Vector of tensors shape: [1, n_patch, 2].
std::optional<std::vector<infinicore::Tensor>> 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<std::vector<infinicore::Tensor>> tgt_sizes;
/// Flattened [start, end) visual token ranges in the packed language sequence.
std::optional<std::vector<size_t>> visual_token_ranges;
Expand Down
77 changes: 68 additions & 9 deletions csrc/models/minicpmv/minicpmv_model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,23 +41,50 @@ MiniCPMVModel::MiniCPMVModel(std::shared_ptr<infinilm::config::ModelConfig> 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<const int64_t *>(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<const int64_t *>(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<size_t>(end) > out_slice->size(0) || static_cast<size_t>(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<size_t>(start), static_cast<size_t>(dst_len)}})
->copy_from(patch_embed->narrow({{0, static_cast<size_t>(src_start), static_cast<size_t>(dst_len)}}));
}
}

Expand All @@ -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<size_t>(req_start), static_cast<size_t>(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(
Expand Down
3 changes: 2 additions & 1 deletion csrc/models/minicpmv/minicpmv_model.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<infinilm::config::ModelConfig> config_;

Expand Down
128 changes: 108 additions & 20 deletions csrc/models/videonsa/videonsa_for_conditional_generation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,34 @@ std::shared_ptr<infinilm::config::ModelConfig> text_config_from(std::shared_ptr<
return std::make_shared<infinilm::config::ModelConfig>(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<const int64_t *>(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<int64_t>(spatial_merge_size) != 0
|| grid_w_i % static_cast<int64_t>(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<size_t>(grid_t_i);
const size_t llm_grid_h = static_cast<size_t>(grid_h_i / static_cast<int64_t>(spatial_merge_size));
const size_t llm_grid_w = static_cast<size_t>(grid_w_i / static_cast<int64_t>(spatial_merge_size));
total += grid_t * llm_grid_h * llm_grid_w;
}
return total;
}

} // namespace

VideoNSAForConditionalGeneration::VideoNSAForConditionalGeneration(std::shared_ptr<infinilm::config::ModelConfig> model_config,
Expand All @@ -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<const int64_t *>(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<size_t>(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<const int64_t *>(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<int64_t>(len)) {
throw std::runtime_error("VideoNSAForConditionalGeneration: image_bound and image_embed_bound length mismatch");
}
const size_t src_start = static_cast<size_t>(src_start_i);
const size_t src_len = len;
if (static_cast<size_t>(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<int32_t *>(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<infinicore::Tensor> pixel_tensors;
std::vector<infinicore::Tensor> grid_tensors;
Expand All @@ -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<size_t>();
}
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<const int64_t *>(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<size_t>(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<size_t>(req_start), static_cast<size_t>(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);
Expand Down
Loading
Loading