Update on "[ExecuTorch][Llama] Change runner to enable chunked prefill"

This diff adds code to chunk prompt longer than max_seq_len to enable prefill of larger context

Differential Revision: [D71833061](https://our.internmc.facebook.com/intern/diff/D71833061/)

[ghstack-poisoned]

Author: kimishpatel

examples/models/llama/runner/runner.cpp

@@ -141,7 +141,8 @@ Error Runner::load() {
   text_prefiller_ = std::make_unique<llm::TextPrefiller>(
       text_decoder_runner_.get(),
       metadata_.at(kUseKVCache),
-      metadata_.at(kEnableDynamicShape));
+      metadata_.at(kEnableDynamicShape),
+      metadata_.at(kMaxSeqLen));
 
   text_token_generator_ = std::make_unique<llm::TextTokenGenerator>(
       tokenizer_.get(),
@@ -242,24 +243,9 @@ Error Runner::generate(
     wrapped_callback(prompt);
   }
   int64_t pos = 0;
-  uint64_t cur_token;
-  int max_seq_len = metadata_.at(kMaxSeqLen) -
-      1; // -1 because for some reason tracing results in this upperbound
-  int num_tokens_to_process = 0;
-  while (num_tokens_to_process < num_prompt_tokens) {
-    auto num_tokens_to_prefill_with =
-        std::min(num_prompt_tokens - num_tokens_to_process, max_seq_len);
-    std::vector<uint64_t> prompt_tokens_to_process(num_tokens_to_prefill_with);
-    std::copy(
-        prompt_tokens.begin() + num_tokens_to_process,
-        prompt_tokens.begin() + num_tokens_to_process +
-            num_tokens_to_prefill_with,
-        prompt_tokens_to_process.begin());
-    auto prefill_res = text_prefiller_->prefill(prompt_tokens_to_process, pos);
-    ET_CHECK_OK_OR_RETURN_ERROR(prefill_res.error());
-    cur_token = prefill_res.get();
-    num_tokens_to_process += num_tokens_to_prefill_with;
-  }
+  auto prefill_res = text_prefiller_->prefill(prompt_tokens, pos);
+  ET_CHECK_OK_OR_RETURN_ERROR(prefill_res.error());
+  uint64_t cur_token = prefill_res.get();
   stats_.first_token_ms = llm::time_in_ms();
   stats_.prompt_eval_end_ms = llm::time_in_ms();
 

extension/llm/runner/text_prefiller.cpp

@@ -10,6 +10,7 @@
 // LLM.
 
 #include <executorch/extension/llm/runner/text_prefiller.h>
+#include <algorithm>
 
 namespace executorch {
 namespace extension {
@@ -18,10 +19,12 @@ namespace llm {
 TextPrefiller::TextPrefiller(
     TextDecoderRunner* text_decoder_runner,
     bool use_kv_cache,
-    bool enable_parallel_prefill)
+    bool enable_parallel_prefill,
+    int64_t max_seq_len)
     : text_decoder_runner_(text_decoder_runner),
       use_kv_cache_(use_kv_cache),
-      enable_parallel_prefill_(enable_parallel_prefill) {}
+      enable_parallel_prefill_(enable_parallel_prefill),
+      max_seq_len_(max_seq_len > 0 ? max_seq_len - 1 : 127) {} // -1 because for some reason tracing results in this upperbound
 
 ::executorch::runtime::Result<uint64_t> TextPrefiller::prefill(
     std::vector<uint64_t>& prompt_tokens,
@@ -30,6 +33,43 @@ ::executorch::runtime::Result<uint64_t> TextPrefiller::prefill(
   if (!text_decoder_runner_->is_method_loaded()) {
     ET_CHECK_OK_OR_RETURN_ERROR(text_decoder_runner_->load());
   }
+  
+  // Check if we need to chunk the prompt tokens
+  int32_t num_prompt_tokens = prompt_tokens.size();
+  
+  // If prompt tokens exceed max_seq_len_, we need to chunk them
+  if (num_prompt_tokens > max_seq_len_) {
+    uint64_t cur_token = 0;
+    int num_tokens_to_process = 0;
+    
+    while (num_tokens_to_process < num_prompt_tokens) {
+      auto num_tokens_to_prefill_with =
+          std::min<int>(num_prompt_tokens - num_tokens_to_process, max_seq_len_);
+      
+      std::vector<uint64_t> prompt_tokens_to_process(num_tokens_to_prefill_with);
+      std::copy(
+          prompt_tokens.begin() + num_tokens_to_process,
+          prompt_tokens.begin() + num_tokens_to_process + num_tokens_to_prefill_with,
+          prompt_tokens_to_process.begin());
+      
+      // Process this chunk
+      auto chunk_result = prefillChunk(prompt_tokens_to_process, start_pos);
+      ET_CHECK_OK_OR_RETURN_ERROR(chunk_result.error());
+      cur_token = chunk_result.get();
+      
+      num_tokens_to_process += num_tokens_to_prefill_with;
+    }
+    
+    return cur_token;
+  } else {
+    // If prompt tokens don't exceed max_seq_len_, process them directly
+    return prefillChunk(prompt_tokens, start_pos);
+  }
+}
+
+::executorch::runtime::Result<uint64_t> TextPrefiller::prefillChunk(
+    std::vector<uint64_t>& prompt_tokens,
+    int64_t& start_pos) {
   // enable_parallel_prefill_ maybe set even when not using kv cache
   // When kv cache is not used, start pos is ignored
   int32_t num_prompt_tokens = prompt_tokens.size();

extension/llm/runner/text_prefiller.h

@@ -22,7 +22,8 @@ class ET_EXPERIMENTAL TextPrefiller {
   TextPrefiller(
       TextDecoderRunner* text_decoder_runner,
       bool use_kv_cache_,
-      bool enable_parallel_prefill);
+      bool enable_parallel_prefill,
+      int64_t max_seq_len = 128);
   /**
    * Prefill an LLM Module with the given text input.
    * @param prompt_tokens The text prompt tokens to the LLM Module. Encoded by
@@ -35,10 +36,21 @@ class ET_EXPERIMENTAL TextPrefiller {
       std::vector<uint64_t>& prompt_tokens,
       int64_t& start_pos);
 
+  /**
+   * Helper method to prefill a chunk of tokens.
+   * @param prompt_tokens The chunk of text prompt tokens to process.
+   * @param start_pos The starting position in KV cache of the input in the LLM Module.
+   * @return The next token of the LLM Module after prefilling this chunk.
+   */
+  ::executorch::runtime::Result<uint64_t> prefillChunk(
+      std::vector<uint64_t>& prompt_tokens,
+      int64_t& start_pos);
+
  private:
   TextDecoderRunner* text_decoder_runner_;
   bool use_kv_cache_;
   bool enable_parallel_prefill_;
+  int64_t max_seq_len_;
 };
 
 } // namespace llm