perf: token-level multimodal mm token expansion.

.pre-commit-config.yaml

@@ -16,7 +16,8 @@ repos:
     args: [
       --toml, pyproject.toml,
       '--skip', 'build/**,third_party/**,*.svg,*.ptx.h',
-      '-L', 'CANN,cann,NNAL,nnal,ASCEND,ascend,tbe,copyin,nd,ND,bilt'
+      '-L', 'CANN,cann,NNAL,nnal,ASCEND,ascend,tbe,copyin,nd,ND,bilt',
+      '-L', 'MultiDimension'
     ]
     additional_dependencies:
       - tomli

xllm/core/common/metrics.cpp

@@ -41,6 +41,8 @@ DEFINE_COUNTER(request_handling_latency_seconds_completion,
 
 DEFINE_COUNTER(tokenization_latency_seconds,
                "Prompt tokenization latency in seconds");
+DEFINE_COUNTER(mm_prompt_expansion_latency_seconds,
+               "Multimodal prompt token expansion latency in seconds");
 DEFINE_COUNTER(chat_template_latency_seconds,
                "Chat template latency in seconds");
 

xllm/core/common/metrics.h

@@ -141,6 +141,7 @@ DECLARE_COUNTER(request_status_total_unimplemented);
 DECLARE_COUNTER(request_handling_latency_seconds_chat);
 DECLARE_COUNTER(request_handling_latency_seconds_completion);
 DECLARE_COUNTER(tokenization_latency_seconds);
+DECLARE_COUNTER(mm_prompt_expansion_latency_seconds);
 DECLARE_COUNTER(chat_template_latency_seconds);
 
 // latency of prefix cache operations in seconds

xllm/processors/glm4v_prompt_processor.cpp

@@ -15,10 +15,35 @@ limitations under the License.
 
 #include "processors/glm4v_prompt_processor.h"
 
+#include <glog/logging.h>
+
 #include <cstdint>
 #include <cstdio>
 
+#include "core/framework/tokenizer/tokenizer.h"
+
 namespace xllm {
+namespace {
+
+void append_encoded_tokens(const Tokenizer* tokenizer,
+                           std::vector<int32_t>& token_ids,
+                           const std::string& text) {
+  CHECK(tokenizer != nullptr) << "video expansion requires tokenizer";
+  std::vector<int32_t> encoded;
+  CHECK(tokenizer->encode(text, &encoded));
+  token_ids.insert(token_ids.end(), encoded.begin(), encoded.end());
+}
+
+int32_t resolve_video_placeholder_id(const Tokenizer* tokenizer,
+                                     const std::string& video_token) {
+  std::vector<int32_t> encoded;
+  CHECK(tokenizer->encode(video_token, &encoded));
+  CHECK_EQ(encoded.size(), 1U)
+      << "video placeholder must tokenize to a single id";
+  return encoded[0];
+}
+
+}  // namespace
 
 GLM4VPromptProcessor::GLM4VPromptProcessor(const ModelArgs& args) {
   merge_size_ = args.mm_image_merge_size();
@@ -29,7 +54,17 @@ GLM4VPromptProcessor::GLM4VPromptProcessor(const ModelArgs& args) {
   image_token_id_ = args.image_token_id();
 }
 
-void GLM4VPromptProcessor::process(std::string& prompt, const MMData& mm_data) {
+void GLM4VPromptProcessor::process(std::string& /*prompt*/,
+                                   const MMData& mm_data) {
+  // Token-level expansion is handled in expand_mm_tokens().
+  DLOG_IF(WARNING, !mm_data.empty())
+      << "GLM4VPromptProcessor::process is unused when token-level "
+         "expansion is enabled";
+}
+
+void GLM4VPromptProcessor::expand_mm_tokens(std::vector<int32_t>& token_ids,
+                                            MMData& mm_data,
+                                            const Tokenizer* tokenizer) {
   torch::Tensor image_grid_thw;
   if (auto res = mm_data.get<torch::Tensor>("image_grid_thw")) {
     image_grid_thw = res.value();
@@ -52,76 +87,63 @@ void GLM4VPromptProcessor::process(std::string& prompt, const MMData& mm_data) {
   }
 
   const int32_t merge_length = merge_size_ * merge_size_;
-  int32_t total_image_token = 0;
-  if (image_grid_thw.defined()) {
-    const int64_t count = image_grid_thw.sizes()[0];
-    for (int64_t idx = 0; idx < count; ++idx) {
-      total_image_token +=
-          image_grid_thw[idx].prod().item<int32_t>() / merge_length;
-    }
-  }
-
-  int32_t total_video_token = 0;
+  int32_t video_placeholder_id = 0;
   if (video_grid_thw.defined()) {
-    const int64_t count = video_grid_thw.sizes()[0];
-    for (int64_t idx = 0; idx < count; ++idx) {
-      total_video_token += video_grid_thw[idx].prod().item<int32_t>() /
-                           merge_length /
-                           video_grid_thw[idx][0].item<int32_t>();
-    }
+    CHECK(tokenizer != nullptr);
+    video_placeholder_id =
+        resolve_video_placeholder_id(tokenizer, video_token_);
   }
 
-  size_t total_token_len = total_image_token * image_token_.size() +
-                           total_video_token * image_token_.size();
-  std::string data;
-  data.reserve(prompt.size() + total_token_len);
+  std::vector<int32_t> expanded;
+  expanded.reserve(token_ids.size());
 
   int32_t image_index = 0;
   int32_t video_index = 0;
-  size_t begin = 0;
-  auto pair = find_vision_token(prompt, begin);
 
-  while (pair.second != std::string::npos) {
-    data.append(prompt, begin, pair.second - begin);
-    if (pair.first == TokenType::IMAGE) {
-      auto token_num =
+  for (size_t index = 0; index < token_ids.size(); ++index) {
+    if (token_ids[index] == image_start_token_id_ &&
+        index + 2 < token_ids.size() &&
+        token_ids[index + 2] == image_end_token_id_ &&
+        token_ids[index + 1] == image_token_id_ && image_grid_thw.defined() &&
+        image_index < image_grid_thw.size(0)) {
+      const int32_t token_num =
           image_grid_thw[image_index].prod().item<int32_t>() / merge_length;
-      while (token_num--) {
-        data.append(image_token_);
-      }
-
+      expanded.push_back(image_start_token_id_);
+      expanded.insert(
+          expanded.end(), static_cast<size_t>(token_num), image_token_id_);
+      expanded.push_back(image_end_token_id_);
       ++image_index;
-      begin = pair.second + image_token_.size();
-    } else if (pair.first == TokenType::VIDEO) {
-      auto num_frames = video_grid_thw[video_index][0].item<int32_t>();
-      auto timestamps = video_metadata[video_index].timestamps;
+      index += 2;
+      continue;
+    }
+
+    if (video_grid_thw.defined() && video_index < video_grid_thw.size(0) &&
+        token_ids[index] == video_placeholder_id) {
+      const int32_t num_frames = video_grid_thw[video_index][0].item<int32_t>();
+      const auto& timestamps = video_metadata[video_index].timestamps;
       CHECK(!timestamps.empty());
 
-      auto selected = build_timestamps(timestamps, num_frames);
-      auto token_num = video_grid_thw[video_index].prod().item<int32_t>() /
-                       merge_length / num_frames;
+      const auto selected =
+          build_timestamps(timestamps, static_cast<size_t>(num_frames));
+      const int32_t token_num =
+          video_grid_thw[video_index].prod().item<int32_t>() / merge_length /
+          num_frames;
       for (int32_t idx = 0; idx < num_frames; ++idx) {
-        data.append(begin_of_image_token_);
-        auto num = token_num;
-        while (num--) {
-          data.append(image_token_);
-        }
-        data.append(end_of_image_token_);
-        data.append(format_timestamp_str(selected[idx]));
+        expanded.push_back(image_start_token_id_);
+        expanded.insert(
+            expanded.end(), static_cast<size_t>(token_num), image_token_id_);
+        expanded.push_back(image_end_token_id_);
+        append_encoded_tokens(
+            tokenizer, expanded, format_timestamp_str(selected[idx]));
       }
-
       ++video_index;
-      begin = pair.second + video_token_.size();
-    } else {
-      LOG(FATAL) << "Unexpected token type encountered.";
+      continue;
     }
-    pair = find_vision_token(prompt, begin);
-  }
 
-  if (begin < prompt.size()) {
-    data.append(prompt, begin, std::string::npos);
+    expanded.push_back(token_ids[index]);
   }
-  prompt = std::move(data);
+
+  token_ids = std::move(expanded);
 }
 
 void GLM4VPromptProcessor::find_mm_spans(const std::vector<int32_t>& token_ids,
@@ -139,7 +161,7 @@ void GLM4VPromptProcessor::find_mm_spans(const std::vector<int32_t>& token_ids,
     auto token = token_ids[idx];
     if (token == video_start_token_id_) {
       is_video = true;
-      video_offset = idx + 1;
+      video_offset = static_cast<int32_t>(idx) + 1;
       video_mask.clear();
       continue;
     } else if (token == video_end_token_id_) {
@@ -165,7 +187,7 @@ void GLM4VPromptProcessor::find_mm_spans(const std::vector<int32_t>& token_ids,
       continue;
     }
     if (token == image_start_token_id_) {
-      image_span_offset = idx + 1;
+      image_span_offset = static_cast<int32_t>(idx) + 1;
     }
     if (token == image_token_id_) {
       ++image_span_length;

xllm/processors/glm4v_prompt_processor.h

@@ -38,6 +38,12 @@ class GLM4VPromptProcessor final : public PromptProcessor {
   void process(std::string& prompt, const MMData& mm_data) override;
   void find_mm_spans(const std::vector<int32_t>& token_ids,
                      MMData& mm_data) override;
+  bool uses_token_level_expansion() const override {
+    return !force_legacy_mm_expansion();
+  }
+  void expand_mm_tokens(std::vector<int32_t>& token_ids,
+                        MMData& mm_data,
+                        const Tokenizer* tokenizer = nullptr) override;
 
  private:
   std::pair<TokenType, size_t> find_vision_token(const std::string& prompt,

xllm/processors/multimodal_processor.cpp

@@ -47,6 +47,16 @@ bool MultimodalProcessorBase::tokenize(const std::string& prompt,
   return true;
 }
 
+void MultimodalProcessorBase::run_mm_token_expansion(
+    PromptProcessor* prompt_processor,
+    std::vector<int32_t>& token_ids,
+    MMData& mm_data) {
+  Timer timer;
+  prompt_processor->expand_mm_tokens(token_ids, mm_data, tokenizer_.get());
+  COUNTER_ADD(mm_prompt_expansion_latency_seconds, timer.elapsed_seconds());
+  prompt_processor->find_mm_spans(token_ids, mm_data);
+}
+
 void MultimodalProcessorBase::assign_mm_hash_keys(const MMInput& mm_input,
                                                   MMData& mm_data) const {
   const std::vector<MMInputItem>& input_items = mm_input.items();

xllm/processors/multimodal_processor.h

@@ -59,6 +59,10 @@ class MultimodalProcessorBase {
   bool tokenize(const std::string& prompt,
                 std::vector<int32_t>& token_ids) const;
 
+  void run_mm_token_expansion(PromptProcessor* prompt_processor,
+                              std::vector<int32_t>& token_ids,
+                              MMData& mm_data);
+
   void assign_mm_hash_keys(const MMInput& mm_input, MMData& mm_data) const;
 
   void pad_to_max_length(std::vector<int32_t>& token_ids) const;
@@ -94,6 +98,14 @@ class MultimodalProcessor final : public MultimodalProcessorBase {
   bool process_prompt(std::string& prompt,
                       MMData& mm_data,
                       std::vector<int32_t>& token_ids) override {
+    if (prompt_processor_->uses_token_level_expansion()) {
+      if (!tokenize(prompt, token_ids)) {
+        return false;
+      }
+      run_mm_token_expansion(prompt_processor_.get(), token_ids, mm_data);
+      return true;
+    }
+
     prompt_processor_->process(prompt, mm_data);
     if (!tokenize(prompt, token_ids)) {
       return false;

xllm/processors/prompt_processor.h

@@ -18,6 +18,7 @@ limitations under the License.
 #include <torch/torch.h>
 
 #include <cstdint>
+#include <cstdlib>
 #include <string>
 #include <vector>
 
@@ -27,13 +28,28 @@ limitations under the License.
 
 namespace xllm {
 
+class Tokenizer;
+
+inline bool force_legacy_mm_expansion() {
+  const char* env = std::getenv("XLLM_FORCE_LEGACY_MM_EXPANSION");
+  return env != nullptr && env[0] != '\0' && env[0] != '0';
+}
+
 class PromptProcessor {
  public:
   virtual ~PromptProcessor() = default;
 
   virtual void process(std::string& prompt, const MMData& mm_data) = 0;
   virtual void find_mm_spans(const std::vector<int32_t>& token_ids,
                              MMData& mm_data) = 0;
+
+  // Token-level expansion tokenizes the unexpanded prompt, then inserts N
+  // identical placeholder ids. Processors that return true implement
+  // expand_mm_tokens() and leave process() as a no-op.
+  virtual bool uses_token_level_expansion() const { return false; }
+  virtual void expand_mm_tokens(std::vector<int32_t>& token_ids,
+                                MMData& mm_data,
+                                const Tokenizer* tokenizer = nullptr) {}
 };
 
 }  // namespace xllm