feat(cpu): add Qwen3.5 0.8B multi-image support

Author: Aharrypotter

examples/qwen3_5/README.md

@@ -129,6 +129,17 @@ mllm-qwen3-5-runner \
   --prompt "Describe the image." \
   --max_new_tokens 32
 
+# Qwen3.5-0.8B multi-image (image order is preserved)
+mllm-qwen3-5-runner \
+  --model_path /path/to/qwen3.5-0.8b-multimodal-w4a32-kai.mllm \
+  --model_version v2 \
+  --tokenizer_path /path/to/Qwen3.5-0.8B/tokenizer.json \
+  --config_path examples/qwen3_5/config_0.8B_multimodal_w4a32_kai.json \
+  --image_path /path/to/first.jpg \
+  --image_path /path/to/second.jpg \
+  --prompt "Compare the first and second images." \
+  --max_new_tokens 32
+
 # Qwen3.5-4B text-only
 mllm-qwen3-5-runner \
   --model_path /path/to/qwen3.5-4b-w4a32-kai.mllm \
@@ -139,6 +150,8 @@ mllm-qwen3-5-runner \
   --max_new_tokens 32
 ```
 
-Omit `--prompt` for the interactive loop. The same optional image is used for
-each independent prompt in that process. Omit `--image_path` to run text-only
-inference with either the text-only or multimodal 0.8B model.
+Omit `--prompt` for the interactive loop. Repeat `--image_path` to attach
+multiple still images in order; the same ordered image list is used for each
+independent prompt in that process. Omit `--image_path` to run text-only
+inference with either the text-only or multimodal 0.8B model. Video input is
+not supported.

examples/qwen3_5/main.cpp

@@ -1,4 +1,5 @@
 #include <fmt/core.h>
+#include <algorithm>
 #include <chrono>
 #include <cstdio>
 #include <filesystem>
@@ -30,8 +31,9 @@ MLLM_MAIN({
   auto& config_path = Argparse::add<std::string>("-c|--config_path").help("Config path").required(true);
   auto& prompt = Argparse::add<std::string>("-p|--prompt").help("Run one prompt non-interactively").required(false);
   auto& prompt_file = Argparse::add<std::string>("--prompt_file").help("Read one benchmark prompt from a file").required(false);
-  auto& image_path =
-      Argparse::add<std::string>("-i|--image_path").help("Optional still image for single-image inference").required(false);
+  auto& image_paths = Argparse::add<std::vector<std::string>>("-i|--image_path")
+                          .help("Optional still image; repeat in prompt order for multi-image inference")
+                          .required(false);
   auto& max_new_tokens = Argparse::add<int>("-g|--max_new_tokens").help("Maximum generated tokens per prompt").required(false);
   auto& print_token_ids = Argparse::add<bool>("--print_token_ids").help("Print generated token IDs to stderr").required(false);
   auto& benchmark_warmup =
@@ -76,6 +78,7 @@ MLLM_MAIN({
     }
 
     auto cfg = mllm::models::qwen3_5::Qwen3_5Config(config_path.get());
+    const auto configured_image_paths = image_paths.get();
     int generation_limit = max_new_tokens.isSet() ? max_new_tokens.get() : 64;
     if (generation_limit <= 0 || generation_limit > cfg.max_cache_length) {
       throw std::invalid_argument("max_new_tokens must be between 1 and max_cache_length");
@@ -85,11 +88,15 @@ MLLM_MAIN({
                                 || benchmark_warmup.isSet() || require_device_telemetry.isSet();
     if (prompt.isSet() && prompt_file.isSet()) { throw std::invalid_argument("prompt and prompt_file are mutually exclusive"); }
     if (prompt.isSet() && prompt.get().empty()) { throw std::invalid_argument("prompt must not be empty"); }
-    if (image_path.isSet() && image_path.get().empty()) { throw std::invalid_argument("image_path must not be empty"); }
-    if (image_path.isSet() && !cfg.vision_enabled) {
+    if (image_paths.isSet() && std::any_of(configured_image_paths.begin(), configured_image_paths.end(), [](const auto& path) {
+          return path.empty();
+        })) {
+      throw std::invalid_argument("image_path must not be empty");
+    }
+    if (image_paths.isSet() && !cfg.vision_enabled) {
       throw std::invalid_argument("image_path requires a config with vision_config");
     }
-    if (image_path.isSet() && benchmark_mode) { throw std::invalid_argument("image_path is not supported in benchmark mode"); }
+    if (image_paths.isSet() && benchmark_mode) { throw std::invalid_argument("image_path is not supported in benchmark mode"); }
     if (benchmark_mode) {
       if (!prompt_file.isSet() || !benchmark_samples.isSet() || !benchmark_jsonl.isSet() || !benchmark_variant.isSet()
           || !benchmark_source_sha.isSet() || !expected_prompt_tokens.isSet()) {
@@ -237,8 +244,7 @@ MLLM_MAIN({
           // KV cache and every GDN recurrent/conv state must start empty.
           model.resetState();
           fmt::print("Processing...\n");
-          auto inputs =
-              tokenizer.convertMessage({.prompt = prompt_text, .image_path = image_path.isSet() ? image_path.get() : ""});
+          auto inputs = tokenizer.convertMessage({.prompt = prompt_text, .image_paths = configured_image_paths});
           const auto prompt_length = inputs.at("sequence").shape()[1];
           if (prompt_length + generation_limit - 1 > cfg.max_cache_length) {
             throw std::invalid_argument(fmt::format("prompt token count ({}) plus max_new_tokens ({}) exceeds "

mllm/models/qwen3_5/image_preprocessor_qwen3_5.hpp

@@ -9,6 +9,7 @@
 #include <stdexcept>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "mllm/core/Tensor.hpp"
 #include "mllm/preprocessor/visual/Image.hpp"
@@ -68,6 +69,35 @@ class Qwen3_5ImagePreprocessor {
     return flattenNormalizedPatches(image.tensor());
   }
 
+  std::pair<Tensor, Tensor> operator()(const std::vector<std::string>& image_paths) const {
+    if (image_paths.empty()) { throw std::invalid_argument("Qwen3.5 image path list must not be empty"); }
+    std::vector<std::pair<Tensor, Tensor>> processed_images;
+    processed_images.reserve(image_paths.size());
+    int32_t total_patches = 0;
+    int32_t patch_features = -1;
+    for (const auto& image_path : image_paths) {
+      auto processed = (*this)(image_path);
+      if (patch_features < 0) patch_features = processed.first.shape()[1];
+      if (processed.first.shape().size() != 2 || processed.first.shape()[1] != patch_features) {
+        throw std::runtime_error("Qwen3.5 preprocessed images have incompatible patch features");
+      }
+      total_patches += processed.first.shape()[0];
+      processed_images.push_back(std::move(processed));
+    }
+
+    auto patches = Tensor::empty({total_patches, patch_features}, kFloat32, kCPU).alloc();
+    auto grids = Tensor::empty({static_cast<int32_t>(image_paths.size()), 3}, kInt32, kCPU).alloc();
+    int64_t patch_offset = 0;
+    for (size_t image_index = 0; image_index < processed_images.size(); ++image_index) {
+      const auto& [image_patches, image_grid] = processed_images[image_index];
+      std::copy(image_patches.ptr<float>(), image_patches.ptr<float>() + image_patches.numel(),
+                patches.ptr<float>() + patch_offset);
+      std::copy(image_grid.ptr<int32_t>(), image_grid.ptr<int32_t>() + 3, grids.ptr<int32_t>() + image_index * 3);
+      patch_offset += image_patches.numel();
+    }
+    return {patches, grids};
+  }
+
   // Exposed for focused tests and reference-oracle comparisons. Input is an
   // already-resized float32 RGB tensor in [H,W,3] with values in [0,255].
   [[nodiscard]] std::pair<Tensor, Tensor> flattenNormalizedPatches(const Tensor& image_hwc) const {

mllm/models/qwen3_5/modeling_qwen3_5.hpp

@@ -634,8 +634,8 @@ class Qwen3_5ForCausalLM : public ARGeneration, public nn::Module {
       position_ids = input.at("position_ids");
       if (seq_len == 1) { position_ids = advanceQwen3_5PositionIds(position_ids); }
     } else if (has_pixel_values) {
-      position_ids = makeQwen3_5SingleImagePositionIds(input.at("mm_token_type_ids"), input.at("image_grid_thw"),
-                                                       vision_spatial_merge_size_);
+      position_ids =
+          makeQwen3_5ImagePositionIds(input.at("mm_token_type_ids"), input.at("image_grid_thw"), vision_spatial_merge_size_);
     } else {
       position_ids = Tensor::empty({batch_size, seq_len}, kInt64, kCPU).alloc();
       auto position_ids_ptr = position_ids.ptr<int64_t>();
@@ -659,29 +659,25 @@ class Qwen3_5ForCausalLM : public ARGeneration, public nn::Module {
       }
       const auto* input_ids = sequence.ptr<int64_t>();
       const auto* types = token_types.ptr<int32_t>();
-      int32_t image_begin = -1;
       int32_t image_count = 0;
       for (int32_t s = 0; s < seq_len; ++s) {
         if (input_ids[s] == video_token_id_ || types[s] == 2) {
-          throw std::invalid_argument("Qwen3.5 CPU single-image support does not accept video tokens");
+          throw std::invalid_argument("Qwen3.5 CPU image support does not accept video tokens");
         }
         if ((input_ids[s] == image_token_id_) != (types[s] == 1)) {
           throw std::invalid_argument("Qwen3.5 image token IDs and modality token types disagree");
         }
-        if (types[s] == 1) {
-          if (image_begin < 0) image_begin = s;
-          ++image_count;
-        }
+        if (types[s] == 1) { ++image_count; }
       }
 
       auto input_embeddings = llm.embed(sequence);
       auto image_embeddings = llm.encodeImage(input.at("pixel_values"), input.at("image_grid_thw"));
-      if (image_begin < 0 || image_embeddings.shape().size() != 2 || image_embeddings.shape()[0] != image_count
+      if (image_count <= 0 || image_embeddings.shape().size() != 2 || image_embeddings.shape()[0] != image_count
           || input_embeddings.shape().size() != 3 || image_embeddings.shape()[1] != input_embeddings.shape()[2]
           || image_embeddings.dtype() != input_embeddings.dtype()) {
         throw std::invalid_argument("Qwen3.5 image features and image placeholders do not match");
       }
-      image_embeddings.copy2(input_embeddings[{kAll, {image_begin, image_begin + image_count}, kAll}]);
+      injectQwen3_5ImageEmbeddings(input_embeddings, image_embeddings, token_types);
       sequence = llm.forwardEmbeddings(input_embeddings, llm_embedding_sin, llm_embedding_cos, AnyValue(&kv_cache_));
     } else {
       sequence = llm(sequence, llm_embedding_sin, llm_embedding_cos, AnyValue(&kv_cache_))[0];

mllm/models/qwen3_5/modeling_qwen3_5_vision.hpp

@@ -327,20 +327,48 @@ class Qwen3_5VisionModel final : public nn::Module {
   std::vector<Tensor> forward(const std::vector<Tensor>& inputs, const std::vector<AnyValue>& args) override {
     const auto& pixel_values = inputs[0];
     const auto& grid_thw = inputs[1];
+    if (pixel_values.shape().size() != 2 || grid_thw.dtype() != kInt32 || grid_thw.device() != kCPU
+        || grid_thw.shape().size() != 2 || grid_thw.shape()[0] <= 0 || grid_thw.shape()[1] != 3) {
+      throw std::invalid_argument("Qwen3.5 vision model requires patches and one grid row per image");
+    }
+    const auto* grids = grid_thw.ptr<int32_t>();
+    int32_t patch_offset = 0;
+    std::vector<Tensor> image_outputs;
+    image_outputs.reserve(grid_thw.shape()[0]);
+    for (int32_t image_index = 0; image_index < grid_thw.shape()[0]; ++image_index) {
+      const auto* grid = grids + image_index * 3;
+      if (grid[0] <= 0 || grid[1] <= 0 || grid[2] <= 0) {
+        throw std::invalid_argument("Qwen3.5 vision image grids must be positive");
+      }
+      const int32_t patch_count = grid[0] * grid[1] * grid[2];
+      if (patch_count > pixel_values.shape()[0] - patch_offset) {
+        throw std::invalid_argument("Qwen3.5 image grids exceed the supplied patch rows");
+      }
+      auto image_pixels = pixel_values[{{patch_offset, patch_offset + patch_count}, kAll}].contiguous();
+      auto image_grid = grid_thw[{{image_index, image_index + 1}, kAll}].contiguous();
+      image_outputs.push_back(forwardImage(image_pixels, image_grid));
+      patch_offset += patch_count;
+    }
+    if (patch_offset != pixel_values.shape()[0]) {
+      throw std::invalid_argument("Qwen3.5 supplied patch rows exceed the image grids");
+    }
+    return {image_outputs.size() == 1 ? image_outputs[0] : nn::functional::concat(image_outputs, 0)};
+  }
+
+ private:
+  Tensor forwardImage(const Tensor& pixel_values, const Tensor& grid_thw) {
     auto hidden_states = patch_embed_(pixel_values)[0];
     auto position_embedding = makeQwen3_5VisionBilinearPositionEmbedding(pos_embed_.weight(), grid_thw, spatial_merge_size_);
     if (hidden_states.shape() != position_embedding.shape()) {
       throw std::invalid_argument("Qwen3.5 patch embeddings do not match the image grid");
     }
     hidden_states = hidden_states + position_embedding;
-
     auto position_ids = makeQwen3_5VisionPositionIds(grid_thw, spatial_merge_size_);
     auto [sin, cos] = makeQwen3_5VisionRotaryEmbedding(position_ids, hidden_size_ / num_heads_);
     for (auto& block : blocks_.list()) { hidden_states = block(hidden_states, sin, cos)[0]; }
-    return {merger_(hidden_states)[0]};
+    return merger_(hidden_states)[0];
   }
 
- private:
   int32_t hidden_size_ = 768;
   int32_t num_heads_ = 12;
   int32_t spatial_merge_size_ = 2;