chore: import upstream snapshot with attribution

cpp/serve/engine_actions/eagle_batch_draft.cc

@@ -0,0 +1,234 @@
+/*!
+ *  Copyright (c) 2023-2025 by Contributors
+ * \file serve/engine_actions/eagle_batch_draft.cc
+ */
+
+#include <numeric>
+
+#include "../config.h"
+#include "../model.h"
+#include "../sampler/sampler.h"
+#include "action.h"
+#include "action_commons.h"
+
+namespace mlc {
+namespace llm {
+namespace serve {
+
+/*!
+ * \brief The action that runs draft proposal 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 EagleBatchDraftActionObj : public EngineActionObj {
+ public:
+  explicit EagleBatchDraftActionObj(Array<Model> models, LogitProcessor logit_processor,
+                                    Sampler sampler, std::vector<ModelWorkspace> model_workspaces,
+                                    DraftTokenWorkspaceManager draft_token_workspace_manager,
+                                    EngineConfig engine_config,
+                                    Optional<EventTraceRecorder> trace_recorder)
+      : models_(std::move(models)),
+        logit_processor_(std::move(logit_processor)),
+        sampler_(std::move(sampler)),
+        model_workspaces_(std::move(model_workspaces)),
+        draft_token_workspace_manager_(std::move(draft_token_workspace_manager)),
+        engine_config_(std::move(engine_config)),
+        trace_recorder_(std::move(trace_recorder)) {}
+
+  Array<Request> Step(EngineState estate) final {
+    // - Only run spec decode when there are two models (llm+ssm) and >=1 running requests.
+    if (models_.size() != 2 || estate->running_queue.empty()) {
+      return {};
+    }
+
+    // Preempt request state entries when decode cannot apply.
+    std::vector<RequestStateEntry> running_rsentries = estate->GetRunningRequestStateEntries();
+    while (!CanDecode(running_rsentries.size())) {
+      if (estate->prefix_cache->TryFreeMemory()) continue;
+      RequestStateEntry preempted = PreemptLastRunningRequestStateEntry(
+          estate, models_, draft_token_workspace_manager_, trace_recorder_);
+      if (preempted.same_as(running_rsentries.back())) {
+        running_rsentries.pop_back();
+      }
+    }
+
+    auto tstart = std::chrono::high_resolution_clock::now();
+
+    int num_rsentries = running_rsentries.size();
+    TVM_FFI_ICHECK_GT(num_rsentries, 0)
+        << "There should be at least one request state entry that can run decode. "
+           "Possible failure reason: none of the prefill phase of the running requests is finished";
+    TVM_FFI_ICHECK_LE(num_rsentries, engine_config_->max_num_sequence)
+        << "The number of running requests exceeds the max number of sequence in EngineConfig. "
+           "Possible failure reason: the prefill action allows new sequence in regardless of the "
+           "max num sequence.";
+
+    Array<String> request_ids;
+    std::vector<int64_t> request_internal_ids;
+    Array<GenerationConfig> generation_cfg;
+    std::vector<RandomGenerator*> rngs;
+    std::vector<std::vector<int>> draft_token_indices;
+    request_ids.reserve(num_rsentries);
+    request_internal_ids.reserve(num_rsentries);
+    generation_cfg.reserve(num_rsentries);
+    draft_token_indices.reserve(num_rsentries);
+    for (const RequestStateEntry& rsentry : running_rsentries) {
+      request_ids.push_back(rsentry->request->id);
+      request_internal_ids.push_back(rsentry->mstates[0]->internal_id);
+      generation_cfg.push_back(rsentry->request->generation_cfg);
+      rngs.push_back(&rsentry->rng);
+    }
+
+    TVM_FFI_ICHECK_GT(estate->spec_draft_length, 0)
+        << "The speculative decoding draft length must be positive.";
+    // The first model doesn't get involved in draft proposal.
+    for (int model_id = 1; model_id < static_cast<int>(models_.size()); ++model_id) {
+      // Collect
+      // - the last committed token,
+      // - the request model state
+      // of each request.
+      std::vector<int> input_tokens;
+      Array<RequestModelState> mstates;
+      input_tokens.reserve(num_rsentries);
+      mstates.reserve(num_rsentries);
+      for (const RequestStateEntry& rsentry : running_rsentries) {
+        mstates.push_back(rsentry->mstates[model_id]);
+      }
+      // draft_length_ rounds of draft proposal.
+      ObjectRef hidden_states = model_workspaces_[model_id].hidden_states;
+      // Concat last hidden_states
+      draft_token_slots_.clear();
+      if (estate->spec_draft_length > 1) {
+        for (int i = 0; i < num_rsentries; ++i) {
+          draft_token_slots_.push_back(mstates[i]->draft_token_slots.back());
+        }
+        hidden_states = models_[model_id]->GatherHiddenStates(
+            model_workspaces_[0].draft_hidden_states_storage, draft_token_slots_, &hidden_states);
+      }
+      // The first draft token has been generated in prefill/verify stage
+      for (int draft_id = 1; draft_id < estate->spec_draft_length; ++draft_id) {
+        draft_token_indices.clear();
+        auto tdraft_start = std::chrono::high_resolution_clock::now();
+        // prepare new input tokens
+        input_tokens.clear();
+        for (int i = 0; i < num_rsentries; ++i) {
+          TVM_FFI_ICHECK(!mstates[i]->draft_output_tokens.empty());
+          input_tokens.push_back(mstates[i]->draft_output_tokens.back().GetTokenId());
+          draft_token_indices.emplace_back(
+              std::vector<int>{static_cast<int>(mstates[i]->draft_output_tokens.size() - 1)});
+        }
+
+        // - Compute embeddings.
+        RECORD_EVENT(trace_recorder_, request_ids, "start proposal embedding");
+        ObjectRef embeddings =
+            models_[model_id]->TokenEmbed({Shape{input_tokens.begin(), input_tokens.end()}});
+        RECORD_EVENT(trace_recorder_, request_ids, "finish proposal embedding");
+
+        // - Invoke model decode.
+        RECORD_EVENT(trace_recorder_, request_ids, "start proposal decode");
+        ObjectRef fused_embedding_hidden_states = models_[model_id]->FuseEmbedHidden(
+            embeddings, hidden_states, /*batch_size*/ num_rsentries, /*seq_len*/ 1);
+        hidden_states = models_[model_id]->BatchDecodeToLastHidden(fused_embedding_hidden_states,
+                                                                   request_internal_ids);
+        Tensor logits;
+        if (models_[model_id]->CanGetLogits()) {
+          logits = models_[model_id]->GetLogits(hidden_states);
+        } else {
+          // - Use base model's head.
+          logits = models_[0]->GetLogits(hidden_states);
+        }
+        RECORD_EVENT(trace_recorder_, request_ids, "finish proposal decode");
+        TVM_FFI_ICHECK_EQ(logits->ndim, 2);
+        TVM_FFI_ICHECK_EQ(logits->shape[0], num_rsentries);
+
+        // - Update logits.
+        logit_processor_->InplaceUpdateLogits(logits, generation_cfg, mstates, request_ids, nullptr,
+                                              &mstates, &draft_token_indices);
+
+        // - Compute probability distributions.
+        Tensor probs_on_device =
+            logit_processor_->ComputeProbsFromLogits(logits, generation_cfg, request_ids);
+
+        // - Commit the prefix cache changes from previous round of action.
+        // Note: we commit prefix cache changes here to overlap this commit with the GPU execution.
+        estate->prefix_cache->CommitSequenceExtention();
+
+        // - Sample tokens.
+        // Fill range [0, num_rsentries) into `sample_indices`.
+        std::vector<int> sample_indices(num_rsentries);
+        std::iota(sample_indices.begin(), sample_indices.end(), 0);
+        Tensor renormalized_probs = sampler_->BatchRenormalizeProbsByTopP(
+            probs_on_device, sample_indices, request_ids, generation_cfg);
+        std::vector<SampleResult> sample_results = sampler_->BatchSampleTokensWithProbAfterTopP(
+            renormalized_probs, sample_indices, request_ids, generation_cfg, rngs);
+        TVM_FFI_ICHECK_EQ(sample_results.size(), num_rsentries);
+
+        // - Add draft token to the state.
+        draft_token_workspace_manager_->AllocSlots(num_rsentries, &draft_token_slots_);
+        models_[model_id]->ScatterDraftProbs(probs_on_device, draft_token_slots_,
+                                             &model_workspaces_[0].draft_probs_storage);
+        // No need to save hidden states as they are not used by subsequent engine actions
+        for (int i = 0; i < num_rsentries; ++i) {
+          int64_t parent_idx = static_cast<int64_t>(mstates[i]->draft_output_tokens.size()) - 1;
+          mstates[i]->AddDraftToken(sample_results[i], draft_token_slots_[i], parent_idx);
+        }
+
+        auto tdraft_end = std::chrono::high_resolution_clock::now();
+        estate->metrics.UpdateDraftTimeByBatchSize(
+            num_rsentries, static_cast<double>((tdraft_end - tdraft_start).count()) / 1e9);
+      }
+    }
+
+    auto tend = std::chrono::high_resolution_clock::now();
+    estate->metrics.engine_decode_time_sum += static_cast<double>((tend - tstart).count()) / 1e9;
+
+    return {};
+  }
+
+ private:
+  /*! \brief Check if the input requests can be decoded under conditions. */
+  bool CanDecode(int num_rsentries) {
+    // The first model is not involved in draft proposal.
+    for (int model_id = 1; model_id < static_cast<int>(models_.size()); ++model_id) {
+      // Check if the model has enough available pages.
+      int num_available_pages = models_[model_id]->GetNumAvailablePages();
+      if (num_rsentries > num_available_pages) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /*! \brief The model to run draft generation in speculative decoding. */
+  Array<Model> models_;
+  /*! \brief The logit processor. */
+  LogitProcessor logit_processor_;
+  /*! \brief The sampler to sample new tokens. */
+  Sampler sampler_;
+  /*! \brief Workspace of each model. */
+  std::vector<ModelWorkspace> model_workspaces_;
+  /*! \brief The draft token workspace manager. */
+  DraftTokenWorkspaceManager draft_token_workspace_manager_;
+  /*! \brief The engine config. */
+  EngineConfig engine_config_;
+  /*! \brief Event trace recorder. */
+  Optional<EventTraceRecorder> trace_recorder_;
+  /*! \brief Temporary buffer to store the slots of the current draft tokens */
+  std::vector<int> draft_token_slots_;
+};
+
+EngineAction EngineAction::EagleBatchDraft(Array<Model> models, LogitProcessor logit_processor,
+                                           Sampler sampler,
+                                           std::vector<ModelWorkspace> model_workspaces,
+                                           DraftTokenWorkspaceManager draft_token_workspace_manager,
+                                           EngineConfig engine_config,
+                                           Optional<EventTraceRecorder> trace_recorder) {
+  return EngineAction(tvm::ffi::make_object<EagleBatchDraftActionObj>(
+      std::move(models), std::move(logit_processor), std::move(sampler),
+      std::move(model_workspaces), std::move(draft_token_workspace_manager),
+      std::move(engine_config), std::move(trace_recorder)));
+}
+
+}  // namespace serve
+}  // namespace llm
+}  // namespace mlc