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,77 @@
+/*!
+ * Copyright (c) 2023-2025 by Contributors
+ * \file serve/draft_token_workspace_manager.cc
+ */
+
+#include "draft_token_workspace_manager.h"
+
+#include "model.h"
+
+namespace mlc {
+namespace llm {
+namespace serve {
+
+TVM_FFI_STATIC_INIT_BLOCK() { DraftTokenWorkspaceManagerObj::RegisterReflection(); }
+
+DraftTokenWorkspaceManagerObj::DraftTokenWorkspaceManagerObj(int max_num_tokens, int vocab_size,
+                                                             int hidden_size,
+                                                             DLDataType hidden_states_dtype,
+                                                             DLDevice device,
+                                                             const FunctionTable& ft)
+    : max_num_tokens_(max_num_tokens),
+      vocab_size_(vocab_size),
+      hidden_size_(hidden_size),
+      hidden_states_dtype_(hidden_states_dtype),
+      device_(device),
+      ft_(ft) {
+  free_slots_.resize(max_num_tokens);
+  std::iota(free_slots_.begin(), free_slots_.end(), 0);
+}
+
+void DraftTokenWorkspaceManagerObj::AllocSlots(int num_slots, std::vector<int>* result) {
+  TVM_FFI_ICHECK_LE(num_slots, free_slots_.size());
+  result->assign(free_slots_.rbegin(), free_slots_.rbegin() + num_slots);
+  free_slots_.resize(free_slots_.size() - num_slots);
+  for (int slot : (*result)) {
+    ref_count_[slot] = 1;
+  }
+}
+
+void DraftTokenWorkspaceManagerObj::AllocSlots(int num_slots,
+                                               const std::vector<int>& initial_ref_count,
+                                               std::vector<int>* result) {
+  TVM_FFI_ICHECK_LE(num_slots, free_slots_.size());
+  TVM_FFI_ICHECK_EQ(num_slots, initial_ref_count.size());
+  result->assign(free_slots_.rbegin(), free_slots_.rbegin() + num_slots);
+  free_slots_.resize(free_slots_.size() - num_slots);
+  for (int i = 0; i < num_slots; ++i) {
+    int slot = (*result)[i];
+    TVM_FFI_ICHECK(initial_ref_count[i] > 0);
+    ref_count_[slot] = initial_ref_count[i];
+  }
+}
+
+void DraftTokenWorkspaceManagerObj::FreeSlots(const std::vector<int>& slots) {
+  for (int slot : slots) {
+    if (--ref_count_.at(slot) == 0) {
+      free_slots_.push_back(slot);
+      ref_count_.erase(slot);
+    }
+  }
+}
+
+void DraftTokenWorkspaceManagerObj::AllocWorkspace(ModelWorkspace* workspace,
+                                                   bool require_hidden_states) {
+  workspace->draft_probs =
+      Tensor::Empty({max_num_tokens_, vocab_size_}, DLDataType{kDLFloat, 32, 1}, device_);
+  workspace->draft_probs_storage =
+      Tensor::Empty({max_num_tokens_, vocab_size_}, DLDataType{kDLFloat, 32, 1}, device_);
+  if (require_hidden_states) {
+    workspace->draft_hidden_states_storage = ft_.Empty(
+        {max_num_tokens_, hidden_size_}, hidden_states_dtype_, device_, /*worker0_only=*/false);
+  }
+}
+
+}  // namespace serve
+}  // namespace llm
+}  // namespace mlc