chore: import upstream snapshot with attribution

2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
@@ -0,0 +1,240 @@
+/*!
+ *  Copyright (c) 2023-2025 by Contributors
+ * \file serve/engine_actions/batch_verify.cc
+ */
+
+#include <tvm/support/cuda/nvtx.h>
+
+#include <cmath>
+#include <exception>
+
+#include "../config.h"
+#include "../model.h"
+#include "../sampler/sampler.h"
+#include "action.h"
+#include "action_commons.h"
+
+namespace mlc {
+namespace llm {
+namespace serve {
+
+using tvm::support::NVTXScopedRange;
+
+/*!
+ * \brief The action that runs verification for requests in the
+ * `running_queue` of engine state. Preempt low-priority requests
+ * accordingly when it is impossible to decode all the running requests.
+ */
+class BatchJumpForwardActionObj : public EngineActionObj {
+ public:
+  explicit BatchJumpForwardActionObj(Array<Model> models, Tokenizer tokenizer,
+                                     Optional<EventTraceRecorder> trace_recorder)
+      : models_(std::move(models)),
+        tokenizer_(tokenizer),
+        trace_recorder_(std::move(trace_recorder)) {}
+
+  Array<Request> Step(EngineState estate) final {
+    // - Do not run decode when there are multiple models or no running requests.
+    if (models_.size() > 1 || estate->running_queue.empty()) {
+      return {};
+    }
+
+    // Preempt request state entries when jump-forward decoding cannot apply.
+    std::vector<RequestStateEntry> running_rsentries;
+    {
+      NVTXScopedRange nvtx_scope("BatchJumpForward getting requests");
+      running_rsentries = estate->GetRunningRequestStateEntries();
+      while (!CheckMemForJumpForward(running_rsentries.size())) {
+        if (estate->prefix_cache->TryFreeMemory()) continue;
+        RequestStateEntry preempted =
+            PreemptLastRunningRequestStateEntry(estate, models_, std::nullopt, trace_recorder_);
+        if (preempted.same_as(running_rsentries.back())) {
+          running_rsentries.pop_back();
+        }
+      }
+    }
+
+    if (running_rsentries.empty()) {
+      return {};
+    }
+
+    auto tstart = std::chrono::high_resolution_clock::now();
+
+    for (auto rsentry : running_rsentries) {
+      if (!CanJumpForward(rsentry)) {
+        continue;
+      }
+
+      auto mstate = rsentry->mstates[0];
+      auto jump_forward_str = mstate->grammar_matcher->FindJumpForwardString();
+
+      if (jump_forward_str.empty()) {
+        continue;
+      }
+
+      auto [rollback_cnt, new_tokens, new_string] =
+          RetokenizeWithNewString(mstate, jump_forward_str, MAX_ROLLBACK_TOKENS_);
+
+      HandleRollback(rsentry, mstate, rollback_cnt, new_tokens, new_string);
+
+      // Commit new tokens (kv cache is handled in the next decode)
+      for (auto token_id : new_tokens) {
+        mstate->CommitToken({{token_id, 1.0}, {}});
+      }
+
+      mstate->require_retokenization_in_next_decode = true;
+
+      // Update metrics
+      rsentry->rstate->metrics.jump_forward_tokens +=
+          std::max(static_cast<int>(new_tokens.size()) - rollback_cnt, 0);
+
+      rsentry->rstate->metrics.completion_tokens +=
+          static_cast<int>(new_tokens.size()) - rollback_cnt;
+    }
+
+    auto tend = std::chrono::high_resolution_clock::now();
+    estate->metrics.engine_jump_forward_time_sum +=
+        static_cast<double>((tend - tstart).count()) / 1e9;
+
+    return {};
+  }
+
+ private:
+  /*! \brief Check if jump-forward decoding can be executed without exceeding the memory limit. */
+  bool CheckMemForJumpForward(int num_rsentries) {
+    static constexpr int MAX_AVG_JUMPFORWARD_PAGES_PER_REQUEST = 10;
+    int num_available_pages = models_[0]->GetNumAvailablePages();
+    return num_rsentries * MAX_AVG_JUMPFORWARD_PAGES_PER_REQUEST <= num_available_pages;
+  }
+
+  /*! \brief Check if the jump-forward can be executed. When logprobs is requested, or the
+   * grammar state matcher is not defined, jump-forward is not executed. */
+  bool CanJumpForward(const RequestStateEntry& rsentry) {
+    if (rsentry->request->generation_cfg->debug_config.grammar_execution_mode !=
+        GrammarExecutionMode::kJumpForward) {
+      return false;
+    }
+    if (rsentry->request->generation_cfg->logprobs) {
+      return false;
+    }
+    if (!rsentry->mstates[0]->grammar_matcher) {
+      return false;
+    }
+    return true;
+  }
+
+  /*!
+   * \brief Retokenize the input string with a new string.
+   * \param mstate The model state.
+   * \param new_string The new string to append.
+   * \param max_rollback_tokens The maximum number of tokens to rollback.
+   * \return The number of tokens to rollback, the new tokens and a delta string of output (equal to
+   * new_string if no cutoff happens; shorter than new_string if cutoff happens).
+   */
+  std::tuple<int, std::vector<int32_t>, std::string> RetokenizeWithNewString(
+      RequestModelState mstate, const std::string& new_string, int max_rollback_tokens) {
+    // Step 1. Get past tokens
+    // past_tokens = mstate[-max_rollback_tokens:]
+    // past_string = detokenize(past_tokens)
+    const auto& token_table = tokenizer_->PostProcessedTokenTable();
+    std::vector<int32_t> past_tokens;
+    std::string past_string;
+    auto past_begin_it = mstate->committed_tokens.size() >= max_rollback_tokens
+                             ? mstate->committed_tokens.end() - max_rollback_tokens
+                             : mstate->committed_tokens.begin();
+    for (auto it = past_begin_it; it != mstate->committed_tokens.end(); ++it) {
+      past_tokens.push_back(it->GetTokenId());
+      past_string += token_table[it->GetTokenId()];
+    }
+
+    // Step 2. Retokenize
+    // Compare tokenize(past_string + new_string) and past_tokens
+    auto new_tokens = tokenizer_->EncodeNoPrependSpace(past_string + new_string);
+    auto delta_string = new_string;
+
+    // Pop last token if it is a prefix of another token. That's because such tokens will often
+    // be rolled back in the next decode, which disturbs the distribution, so we will avoid
+    // generating them.
+    if (tokenizer_->GetPrefixTokenMask()[new_tokens.back()]) {
+      auto last_token = token_table[new_tokens.back()];
+      if (last_token.length() >= new_string.length()) {
+        return {0, {}, ""};
+      }
+
+      delta_string = delta_string.substr(0, delta_string.length() - last_token.length());
+      new_tokens.pop_back();
+    }
+
+    int first_differ_idx = past_tokens.size();
+    for (int i = 0; i < static_cast<int>(past_tokens.size()); ++i) {
+      if (i == static_cast<int>(new_tokens.size()) || past_tokens[i] != new_tokens[i]) {
+        first_differ_idx = i;
+        break;
+      }
+    }
+
+    return {past_tokens.size() - first_differ_idx,
+            std::vector<int32_t>(new_tokens.begin() + first_differ_idx, new_tokens.end()),
+            delta_string};
+  }
+
+  /*!
+   * \brief Handle rollback for the stream output, the model state and the kv cache.
+   * \param rsentry The request state entry.
+   * \param mstate The model state.
+   * \param rollback_cnt The number of tokens to rollback.
+   * \param new_tokens The new tokens. Useful for the stream output.
+   * \param new_string The delta string of output. Useful for the stream output.
+   */
+  void HandleRollback(const RequestStateEntry& rsentry, RequestModelState mstate, int rollback_cnt,
+                      const std::vector<int32_t>& new_tokens, const std::string& new_string) {
+    // 1. Handle rollback for the stream output
+    if (rollback_cnt >
+        static_cast<int>(mstate->committed_tokens.size()) - rsentry->next_callback_token_pos) {
+      const auto& token_table = tokenizer_->PostProcessedTokenTable();
+      for (auto i = rsentry->next_callback_token_pos; i < mstate->committed_tokens.size(); ++i) {
+        auto token_id = mstate->committed_tokens[i].GetTokenId();
+        rsentry->extra_prefix_string += token_table[token_id];
+      }
+      rsentry->extra_prefix_string += new_string;
+      rsentry->next_callback_token_pos = static_cast<int>(mstate->committed_tokens.size()) -
+                                         rollback_cnt + static_cast<int>(new_tokens.size());
+    }
+
+    // 2. Handle rollback for the model state
+    if (rollback_cnt > 0) {
+      mstate->RollbackTokens(rollback_cnt);
+    }
+
+    // 3. Handle rollback for the kv cache
+    if (rollback_cnt > mstate->num_tokens_for_next_decode) {
+      models_[0]->PopNFromKVCache(mstate->internal_id,
+                                  rollback_cnt - mstate->num_tokens_for_next_decode);
+      mstate->num_tokens_for_next_decode = 0;
+    } else {
+      mstate->num_tokens_for_next_decode -= rollback_cnt;
+    }
+  }
+
+  /*!
+   * \brief The model to run jump-forward decoding. When there are multiple
+   * models, the `Step` function of the created action will not take effect.
+   */
+  Array<Model> models_;
+  /*! \brief Tokenizer for retokenization. */
+  Tokenizer tokenizer_;
+  /*! \brief Event trace recorder. */
+  Optional<EventTraceRecorder> trace_recorder_;
+  /*! \brief The maximum number of tokens to rollback. */
+  const int MAX_ROLLBACK_TOKENS_ = 10;
+};
+
+EngineAction EngineAction::BatchJumpForward(Array<Model> models, Tokenizer tokenizer,
+                                            Optional<EventTraceRecorder> trace_recorder) {
+  return EngineAction(tvm::ffi::make_object<BatchJumpForwardActionObj>(
+      std::move(models), std::move(tokenizer), std::move(trace_recorder)));
+}
+
+}  // namespace serve
+}  // namespace llm
+}  // namespace mlc