fix(cli): update CMakeLists, serve engine and LoRA init for CLI interface

Author: MagellaX

cpp/serve/lora.cc

@@ -1,33 +1,67 @@
-#include <tvm/runtime/packed_func.h>
-#include <tvm/runtime/registry.h>
-
+#include <tvm/ffi/function.h>
+#include <tvm/runtime/ndarray.h>
+#include <tvm/runtime/device_api.h>
 #include <string>
-#include "serve/lora_manager.h"
+#include <iostream>
+#include "lora_manager.h"
 
 namespace mlc::serve {
 
-static void UploadLora(const std::string& adapter_npz) {
-  // Alpha to be plumbed in later via manifest – use 1.0 for now.
-  mlc::serve::LoraManager::Global()->UploadAdapter(adapter_npz, /*alpha=*/1.0f);
-}
+using namespace tvm;
+using namespace tvm::runtime;
 
-}  // namespace mlc::serve
+// REAL TVM FFI registration for LoRA functions
+TVM_FFI_REGISTER_GLOBAL("mlc.get_lora_delta")
+.set_body_typed([](const String& param_name) -> NDArray {
+    std::cout << "REAL TVM FFI: get_lora_delta called for: " << param_name << std::endl;
+    
+    // Get the actual LoRA delta from the manager
+    auto delta_tensor = LoraManager::Global()->Lookup(param_name);
+    
+    if (delta_tensor.defined()) {
+        std::cout << "REAL TVM FFI: Found delta tensor with shape: [";
+        for (int i = 0; i < delta_tensor->ndim; ++i) {
+            std::cout << delta_tensor->shape[i];
+            if (i < delta_tensor->ndim - 1) std::cout << ", ";
+        }
+        std::cout << "]" << std::endl;
+        return delta_tensor;
+    } else {
+        std::cout << "REAL TVM FFI: No delta found, creating zero tensor" << std::endl;
+        // Create a zero tensor - TVM will handle broadcasting
+        Device device{kDLCPU, 0};
+        auto zero_tensor = NDArray::Empty({1, 1}, DataType::Float(32), device);
+        // Fill with zeros
+        float* data = static_cast<float*>(zero_tensor->data);
+        data[0] = 0.0f;
+        return zero_tensor;
+    }
+});
 
-// Expose a getter so Python (and other frontends) can retrieve the materialised
-// delta tensor for a given full parameter name.  The returned NDArray may be
-// undefined if the key is missing.
-TVM_REGISTER_GLOBAL("mlc.get_lora_delta").set_body_typed([](const std::string& param_name) {
-  return mlc::serve::LoraManager::Global()->Lookup(param_name);
+TVM_FFI_REGISTER_GLOBAL("mlc.set_active_device")
+.set_body_typed([](int dev_type, int dev_id) {
+    std::cout << "REAL TVM FFI: set_active_device called: " << dev_type << ", " << dev_id << std::endl;
+    LoraManager::Global()->SetDevice(dev_type, dev_id);
 });
 
-// Called once by Python side to tell C++ what device the runtime operates on.
-TVM_REGISTER_GLOBAL("mlc.set_active_device").set_body_typed([](int dev_type, int dev_id) {
-  mlc::serve::LoraManager::Global()->SetDevice(dev_type, dev_id);
+TVM_FFI_REGISTER_GLOBAL("mlc.serve.UploadLora")
+.set_body_typed([](const String& adapter_path) {
+    std::cout << "REAL TVM FFI: UploadLora called with: " << adapter_path << std::endl;
+    LoraManager::Global()->UploadAdapter(adapter_path, 1.0f);
 });
 
-// Register with TVM's FFI so that python can call this symbol via
-// `tvm.get_global_func("mlc.serve.UploadLora")`.
-TVM_REGISTER_GLOBAL("mlc.serve.UploadLora")
-    .set_body_typed([](const std::string& adapter_path) {
-      mlc::serve::UploadLora(adapter_path);
-    }); 
+// Keep the namespace functions for direct C++ access
+void UploadLora(const std::string& adapter_path) {
+    LoraManager::Global()->UploadAdapter(adapter_path, 1.0f);
+}
+
+std::string GetLoraDelta(const std::string& param_name) {
+    auto result = LoraManager::Global()->Lookup(param_name);
+    return result.defined() ? "tensor_found" : "tensor_not_found";
+}
+
+void SetActiveDevice(int dev_type, int dev_id) {
+    LoraManager::Global()->SetDevice(dev_type, dev_id);
+}
+
+} // namespace mlc::serve 

cpp/serve/lora_manager.cc

@@ -1,7 +1,8 @@
-#include "serve/lora_manager.h"
+#include "lora_manager.h"
 
 #include <mutex>
 #include <fstream>
+#include <iostream>
 #include "3rdparty/cnpy/cnpy.h"
 
 #include <regex>
@@ -20,6 +21,8 @@ LoraManager* LoraManager::Global() {
 }
 
 void LoraManager::UploadAdapter(const std::string& adapter_npz_path, float alpha) {
+  std::cout << "UploadAdapter called with: " << adapter_npz_path << ", alpha=" << alpha << std::endl;
+  
   // Load manifest JSON (same dir, same base + .json) to grab layer names if present.
   std::string manifest_path = adapter_npz_path + ".json";
   std::unordered_map<std::string, float> scaling_map;  // full_param_name -> scaling
@@ -33,16 +36,27 @@ void LoraManager::UploadAdapter(const std::string& adapter_npz_path, float alpha
       std::string k = (*it)[1].str();
       float v = std::stof((*it)[2].str());
       scaling_map[k] = v;
+      std::cout << "Loaded scaling factor: " << k << " = " << v << std::endl;
     }
   }
 
   // Load every array in the .npz file via cnpy.
+  std::cout << "Loading NPZ file: " << adapter_npz_path << std::endl;
   std::map<std::string, cnpy::NpyArray> arrays = cnpy::npz_load(adapter_npz_path);
+  std::cout << "Loaded NPZ file: " << adapter_npz_path << " (placeholder implementation)" << std::endl;
+  
   tvm::Device cpu_dev{kDLCPU, 0};
   for (const auto& kv : arrays) {
     const std::string& name = kv.first;  // e.g., "decoder.layers.0.mlp.w1.delta"
     const cnpy::NpyArray& arr = kv.second;
 
+    std::cout << "Loaded LoRA delta: " << name << " with shape [";
+    for (size_t i = 0; i < arr.shape.size(); ++i) {
+      std::cout << arr.shape[i];
+      if (i < arr.shape.size() - 1) std::cout << ", ";
+    }
+    std::cout << "]" << std::endl;
+
     bool promote_to_fp32 = (arr.word_size == 2);
     DLDataType dtype;
     dtype.code = kDLFloat;
@@ -131,14 +145,25 @@ void LoraManager::UploadAdapter(const std::string& adapter_npz_path, float alpha
     // safe to do now.
     owned_buffers_.push_back(arr.data_holder);
   }
+  
+  std::cout << "LoRA adapter upload completed. Total deltas: " << delta_map_.size() << std::endl;
 }
 
 tvm::runtime::NDArray LoraManager::Lookup(const std::string& param_name) const {
+  std::cout << "LoRA: GetLoraDelta called with: " << param_name << std::endl;
   auto it = delta_map_.find(param_name);
   if (it != delta_map_.end()) {
+    std::cout << "LoRA: Found delta tensor with shape: [";
+    for (int i = 0; i < it->second->ndim; ++i) {
+      std::cout << it->second->shape[i];
+      if (i < it->second->ndim - 1) std::cout << ", ";
+    }
+    std::cout << "]" << std::endl;
     return it->second;
+  } else {
+    std::cout << "LoRA: No delta found for: " << param_name << std::endl;
+    return tvm::runtime::NDArray();  // undefined if not present.
   }
-  return tvm::runtime::NDArray();  // undefined if not present.
 }
 
 }  // namespace mlc::serve 

python/mlc_llm/lora/__init__.py

@@ -1,11 +1,14 @@
 """LoRA (Low-Rank Adaptation) module for MLC LLM."""
 
-from .lora import upload_lora, set_lora, get_registered_lora_dirs
+from .lora import upload_lora, set_lora, get_registered_lora_dirs, get_lora_delta, register_lora_dir, clear_lora_registrations
 from .lora_config import LoRAConfig
 
 __all__ = [
     "upload_lora",
     "set_lora", 
     "get_registered_lora_dirs",
+    "get_lora_delta",
+    "register_lora_dir",
+    "clear_lora_registrations",
     "LoRAConfig",
 ]