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,100 @@
+/*!
+ * \file builtin.cc
+ * \brief Multi-GPU builtin functions in MLC LLM.
+ */
+#ifndef MLC_SINGLE_GPU_ONLY
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/shape.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/optional.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/runtime/disco/builtin.h>
+#include <tvm/runtime/disco/disco_worker.h>
+#include <tvm/runtime/tensor.h>
+#include <tvm/runtime/vm/vm.h>
+
+namespace mlc {
+namespace llm {
+namespace multi_gpu {
+
+using namespace tvm::runtime;
+using tvm::ffi::Array;
+using tvm::ffi::ObjectRef;
+using tvm::ffi::Optional;
+using tvm::ffi::Shape;
+
+ObjectRef DispatchFunctionByGroup(tvm::ffi::AnyView vm_arg,
+                                  Array<Array<ObjectRef>> funcs_and_args) {
+  using namespace vm;
+  VirtualMachine* vm = VirtualMachine::GetContextPtr(vm_arg);
+  DiscoWorker* worker = DiscoWorker::ThreadLocal();
+  int world_size = worker->num_workers;
+  int group_size = worker->num_workers / worker->num_groups;
+  int num_group = world_size / group_size;
+  TVM_FFI_ICHECK_EQ(funcs_and_args.size(), num_group)
+      << "Number of groups mismatches. There are " << num_group
+      << " groups while the function/arg array has " << funcs_and_args.size() << " elements.";
+
+  int group_id = worker->worker_id / group_size;
+  TVM_FFI_ICHECK(!funcs_and_args[group_id].empty())
+      << "No function is provided for group " << group_id;
+  VMClosure func = funcs_and_args[group_id][0].as_or_throw<VMClosure>();
+
+  int num_args = static_cast<int>(funcs_and_args[group_id].size()) - 1;
+  std::vector<tvm::ffi::AnyView> packed_args(num_args);
+  for (int i = 0; i < num_args; ++i) {
+    // NOTE: Need explicily define `arg` so that the argument does not
+    // have type code kTVMObjectRValueRefArg.
+    packed_args[i] = funcs_and_args[group_id][1 + i];
+  }
+
+  tvm::ffi::Any rv;
+  vm->InvokeClosurePacked(funcs_and_args[group_id][0].as_or_throw<VMClosure>(),
+                          tvm::ffi::PackedArgs(packed_args.data(), packed_args.size()), &rv);
+  return rv.cast<ObjectRef>();
+}
+
+ObjectRef SendFromLastGroupToWorker0(Tensor send, Optional<Tensor> recv, Shape shape,
+                                     DLDataType dtype) {
+  DiscoWorker* worker = DiscoWorker::ThreadLocal();
+  int worker_id = worker->worker_id;
+  int world_size = worker->num_workers;
+  int group_size = worker->num_workers / worker->num_groups;
+  TVM_FFI_ICHECK_NE(world_size, group_size) << "Cannot perform when there is only one group.";
+  int sender_id = world_size - group_size;
+  if (worker_id == 0) {
+    TVM_FFI_ICHECK(recv.has_value()) << "The receive Tensor is undefined for worker 0.";
+    Tensor recv_arr = recv.value().CreateView(shape, dtype);
+    RecvFromWorker(recv_arr, sender_id);
+    return recv_arr;
+  } else if (worker_id == sender_id) {
+    TVM_FFI_ICHECK_EQ(send->dtype, dtype)
+        << "The src Tensor has mismatched dtype than the expected dtype.";
+    TVM_FFI_ICHECK_EQ(send->ndim, shape.size())
+        << "The src Tensor has mismatched shape than the expected shape.";
+    for (int i = 0; i < send->ndim; ++i) {
+      TVM_FFI_ICHECK_EQ(send->shape[i], shape[i])
+          << "The src Tensor has mismatched shape than the expected shape.";
+    }
+    SendToWorker(send, /*receiver_id=*/0);
+    return recv.value_or(Tensor(nullptr));
+  }
+
+  // We only process for worker 0 and the first worker of the last group.
+  // For other workers, we return the input object.
+  return recv.value_or(Tensor(nullptr));
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+  namespace refl = tvm::ffi::reflection;
+  refl::GlobalDef()
+      .def("mlc.multi_gpu.DispatchFunctionByGroup", DispatchFunctionByGroup)
+      .def("mlc.multi_gpu.SendFromLastGroupToWorker0", SendFromLastGroupToWorker0);
+}
+
+}  // namespace multi_gpu
+}  // namespace llm
+}  // namespace mlc
+
+#endif  // MLC_SINGLE_GPU_ONLY