feat(LlamaRankingContext): add an option to chunk large documents

llama/CMakeLists.txt

@@ -116,47 +116,21 @@ include_directories("./llama.cpp/common")
 # This is needed to use methods in "llama-grammar.h" and "unicode.h"
 target_include_directories(llama PUBLIC "./llama.cpp/src")
 
-unset(GPU_INFO_HEADERS)
-unset(GPU_INFO_SOURCES)
-unset(GPU_INFO_EXTRA_LIBS)
-
-if (GGML_VULKAN OR GGML_KOMPUTE)
-    find_package(Vulkan)
-    if (Vulkan_FOUND)
-        if (GGML_VULKAN)
-            message(STATUS "Using Vulkan for GPU info")
-        elseif (GGML_KOMPUTE)
-            message(STATUS "Using Vulkan for GPU info because Kompute is enabled")
-        endif()
-
-        list(APPEND GPU_INFO_HEADERS gpuInfo/vulkan-gpu-info.h)
-        list(APPEND GPU_INFO_SOURCES gpuInfo/vulkan-gpu-info.cpp)
-
-        add_compile_definitions(GPU_INFO_USE_VULKAN)
-
-        list(APPEND GPU_INFO_EXTRA_LIBS Vulkan::Vulkan)
-    else()
-        message(FATAL_ERROR "Vulkan was not found")
-    endif()
-endif()
-
-list(REMOVE_DUPLICATES GPU_INFO_HEADERS)
-list(REMOVE_DUPLICATES GPU_INFO_SOURCES)
-list(REMOVE_DUPLICATES GPU_INFO_EXTRA_LIBS)
-
 addVariantSuffix(llama "${NLC_VARIANT}")
 addVariantSuffix(ggml "${NLC_VARIANT}")
 
-file(GLOB SOURCE_FILES "addon/*.cpp" "addon/**/*.cpp" ${GPU_INFO_SOURCES})
+file(GLOB SOURCE_FILES "addon/*.cpp" "addon/**/*.cpp")
 
-add_library(${PROJECT_NAME} SHARED ${SOURCE_FILES} ${CMAKE_JS_SRC} ${GPU_INFO_HEADERS})
+add_library(${PROJECT_NAME} SHARED ${SOURCE_FILES} ${CMAKE_JS_SRC})
 set_target_properties(${PROJECT_NAME} PROPERTIES PREFIX "" SUFFIX ".node")
 target_link_libraries(${PROJECT_NAME} ${CMAKE_JS_LIB})
 target_link_libraries(${PROJECT_NAME} "llama")
 target_link_libraries(${PROJECT_NAME} "llama-common")
 
-if (DEFINED GPU_INFO_EXTRA_LIBS)
-    target_link_libraries(${PROJECT_NAME} ${GPU_INFO_EXTRA_LIBS})
+if (GGML_VULKAN)
+    find_package(Vulkan REQUIRED)
+    target_compile_definitions(${PROJECT_NAME} PRIVATE GPU_INFO_USE_VULKAN)
+    target_link_libraries(${PROJECT_NAME} Vulkan::Vulkan)
 endif()
 
 if(MSVC AND CMAKE_JS_NODELIB_DEF AND CMAKE_JS_NODELIB_TARGET)

llama/addon/AddonContext.cpp

@@ -574,17 +574,21 @@ Napi::Value AddonContext::InitBatch(const Napi::CallbackInfo& info) {
         return info.Env().Undefined();
     }
 
+    int32_t n_tokens = info[0].As<Napi::Number>().Int32Value();
+    if (n_tokens <= 0 || static_cast<uint32_t>(n_tokens) > context_params.n_batch) {
+        Napi::RangeError::New(info.Env(), "Invalid batch size").ThrowAsJavaScriptException();
+        return info.Env().Undefined();
+    }
+
     if (has_batch) {
         llama_batch_free(batch);
     }
 
-    int32_t n_tokens = info[0].As<Napi::Number>().Int32Value();
-
     batch = llama_batch_init(n_tokens, 0, 1);
     has_batch = true;
     batch_n_tokens = n_tokens;
 
-    uint64_t newBatchMemorySize = calculateBatchMemorySize(n_tokens, llama_model_n_embd(model->model), context_params.n_batch);
+    uint64_t newBatchMemorySize = calculateBatchMemorySize(n_tokens, 0, 1);
     if (newBatchMemorySize > batchMemorySize) {
         adjustNapiExternalMemoryAdd(Env(), newBatchMemorySize - batchMemorySize);
         batchMemorySize = newBatchMemorySize;
@@ -618,12 +622,15 @@ Napi::Value AddonContext::AddToBatch(const Napi::CallbackInfo& info) {
 
     auto tokensLength = tokens.ElementLength();
     auto tokenLogitIndexesLength = tokenLogitIndexes.ElementLength();
-    GGML_ASSERT(batch.n_tokens + tokensLength <= batch_n_tokens);
+    if (tokensLength > static_cast<size_t>(batch_n_tokens - batch.n_tokens)) {
+        Napi::RangeError::New(info.Env(), "Tokens exceed the initialized batch size").ThrowAsJavaScriptException();
+        return info.Env().Undefined();
+    }
 
     Napi::Uint32Array resLogitIndexes = Napi::Uint32Array::New(info.Env(), tokenLogitIndexesLength);
 
     for (size_t i = 0, l = 0; i < tokensLength; i++) {
-        if (l < tokenLogitIndexesLength && l < tokenLogitIndexesLength && tokenLogitIndexes[l] == i) {
+        if (l < tokenLogitIndexesLength && tokenLogitIndexes[l] == i) {
             common_batch_add(batch, static_cast<llama_token>(tokens[i]), firstTokenContextIndex + i, { sequenceId }, true);
             resLogitIndexes[l] = batch.n_tokens - 1;
             l++;
@@ -724,26 +731,39 @@ Napi::Value AddonContext::GetEmbedding(const Napi::CallbackInfo& info) {
     }
 
     int32_t inputTokensLength = info[0].As<Napi::Number>().Int32Value();
-    int32_t maxVectorSize = (info.Length() > 1 && info[1].IsNumber()) ? info[1].As<Napi::Number>().Int32Value() : 0;
+    const double maxVectorSize = (info.Length() > 1 && info[1].IsNumber()) ? info[1].As<Napi::Number>().DoubleValue() : 0;
 
     if (inputTokensLength <= 0) {
         Napi::Error::New(info.Env(), "Invalid input tokens length").ThrowAsJavaScriptException();
         return info.Env().Undefined();
     }
 
-    const int n_embd = llama_model_n_embd(model->model);
+    if (!std::isfinite(maxVectorSize) || maxVectorSize < 0 || std::floor(maxVectorSize) != maxVectorSize) {
+        Napi::Error::New(info.Env(), "Invalid maximum embedding vector size").ThrowAsJavaScriptException();
+        return info.Env().Undefined();
+    }
+
     const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);
+    const int64_t n_embd = pooling_type == LLAMA_POOLING_TYPE_RANK
+        ? static_cast<int64_t>(llama_model_n_cls_out(model->model))
+        : llama_model_n_embd_out(model->model);
+
+    if (n_embd <= 0) {
+        Napi::Error::New(info.Env(), "Invalid embedding vector size").ThrowAsJavaScriptException();
+        return info.Env().Undefined();
+    }
+
     const auto* embeddings = pooling_type == LLAMA_POOLING_TYPE_NONE ? NULL : llama_get_embeddings_seq(ctx, 0);
     if (embeddings == NULL) {
-        embeddings = llama_get_embeddings_ith(ctx, inputTokensLength - 1);
+        embeddings = llama_get_embeddings_ith(ctx, -1);
     }
 
     if (embeddings == NULL) {
         Napi::Error::New(info.Env(), std::string("Failed to get embeddings for token ") + std::to_string(inputTokensLength - 1)).ThrowAsJavaScriptException();
         return info.Env().Undefined();
     }
 
-    size_t resultSize = maxVectorSize == 0 ? n_embd : std::min(n_embd, maxVectorSize);
+    const size_t resultSize = maxVectorSize == 0 ? n_embd : std::min<double>(n_embd, maxVectorSize);
     Napi::Float64Array result = Napi::Float64Array::New(info.Env(), resultSize);
     for (size_t i = 0; i < resultSize; i++) {
         result[i] = embeddings[i];
@@ -1059,10 +1079,10 @@ class RestoreCheckpointWorker : public Napi::AsyncWorker {
     public:
         AddonContext* context;
         AddonContextSequenceCheckpoint* checkpoint;
-        std::size_t maxPosIndex;
+        llama_pos maxPosIndex;
         bool restoreSuccess = false;
 
-        RestoreCheckpointWorker(const Napi::CallbackInfo& info, AddonContext* context, AddonContextSequenceCheckpoint* checkpoint, std::size_t maxPosIndex)
+        RestoreCheckpointWorker(const Napi::CallbackInfo& info, AddonContext* context, AddonContextSequenceCheckpoint* checkpoint, llama_pos maxPosIndex)
             : Napi::AsyncWorker(info.Env(), "RestoreCheckpointWorker"),
               context(context),
               checkpoint(checkpoint),
@@ -1085,7 +1105,22 @@ class RestoreCheckpointWorker : public Napi::AsyncWorker {
 
         void Execute() {
             try {
-                std::lock_guard<std::mutex> lock(checkpoint->dataMutex);
+                std::shared_lock<std::shared_mutex> lock(checkpoint->dataMutex);
+
+                if (checkpoint->disposed) {
+                    SetError("Checkpoint is disposed");
+                    return;
+                }
+
+                if (!checkpoint->initialized) {
+                    return;
+                }
+
+                if (checkpoint->maxPos < 0) {
+                    restoreSuccess = maxPosIndex == -1 &&
+                        llama_memory_seq_rm(llama_get_memory(context->ctx), checkpoint->sequenceId, 0, -1);
+                    return;
+                }
 
                 std::size_t dataSize = checkpoint->data.size();
                 std::size_t restoreSize = llama_state_seq_set_data_ext(context->ctx, checkpoint->data.data(), dataSize, checkpoint->sequenceId, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
@@ -1110,8 +1145,12 @@ class RestoreCheckpointWorker : public Napi::AsyncWorker {
 };
 
 Napi::Value AddonContext::RestoreCheckpoint(const Napi::CallbackInfo& info) {
+    if (disposed || !contextLoaded) {
+        throw Napi::Error::New(info.Env(), "Context is disposed or not loaded");
+    }
+
     AddonContextSequenceCheckpoint* checkpoint = Napi::ObjectWrap<AddonContextSequenceCheckpoint>::Unwrap(info[0].As<Napi::Object>());
-    std::size_t maxPosIndex = info[1].As<Napi::Number>().Int32Value();
+    const llama_pos maxPosIndex = static_cast<llama_pos>(info[1].As<Napi::Number>().Int32Value());
 
     RestoreCheckpointWorker* worker = new RestoreCheckpointWorker(info, this, checkpoint, maxPosIndex);
     worker->Queue();
@@ -1162,13 +1201,18 @@ AddonContextSequenceCheckpoint::~AddonContextSequenceCheckpoint() {
 
 class AddonContextSequenceCheckpointInitWorker : public Napi::AsyncWorker {
     public:
-    AddonContextSequenceCheckpoint* checkpoint;
+        AddonContextSequenceCheckpoint* checkpoint;
         AddonContext* context;
+        const llama_seq_id sequenceId;
+        llama_pos minPos = -1;
+        llama_pos maxPos = -1;
+        std::vector<uint8_t> data;
 
-        AddonContextSequenceCheckpointInitWorker(const Napi::CallbackInfo& info, AddonContextSequenceCheckpoint* checkpoint, AddonContext* context)
+        AddonContextSequenceCheckpointInitWorker(const Napi::CallbackInfo& info, AddonContextSequenceCheckpoint* checkpoint, AddonContext* context, llama_seq_id sequenceId)
             : Napi::AsyncWorker(info.Env(), "AddonContextSequenceCheckpointInitWorker"),
-            checkpoint(checkpoint),
+              checkpoint(checkpoint),
               context(context),
+              sequenceId(sequenceId),
               deferred(Napi::Promise::Deferred::New(info.Env())) {
             checkpoint->Ref();
             context->Ref();
@@ -1187,19 +1231,38 @@ class AddonContextSequenceCheckpointInitWorker : public Napi::AsyncWorker {
 
         void Execute() {
             try {
-                checkpoint->minPos = llama_memory_seq_pos_min(llama_get_memory(context->ctx), checkpoint->sequenceId);
-                checkpoint->maxPos = llama_memory_seq_pos_max(llama_get_memory(context->ctx), checkpoint->sequenceId);
-                const size_t checkpointSize = llama_state_seq_get_size_ext(context->ctx, checkpoint->sequenceId, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
-
-                checkpoint->data.resize(checkpointSize, 0);
-                llama_state_seq_get_data_ext(context->ctx, checkpoint->data.data(), checkpointSize, checkpoint->sequenceId, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+                minPos = llama_memory_seq_pos_min(llama_get_memory(context->ctx), sequenceId);
+                maxPos = llama_memory_seq_pos_max(llama_get_memory(context->ctx), sequenceId);
+                if (maxPos >= 0) {
+                    const size_t checkpointSize = llama_state_seq_get_size_ext(context->ctx, sequenceId, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+
+                    data.resize(checkpointSize, 0);
+                    if (checkpointSize == 0 || llama_state_seq_get_data_ext(context->ctx, data.data(), checkpointSize, sequenceId, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) != checkpointSize) {
+                        throw std::runtime_error("Failed to capture the complete checkpoint");
+                    }
+                }
             } catch (const std::exception& e) {
                 SetError(e.what());
             } catch(...) {
                 SetError("Unknown error when calling \"llama_state_seq_get_data_ext\"");
             }
         }
         void OnOK() {
+            if (checkpoint->disposed) {
+                deferred.Reject(Napi::Error::New(Env(), "Checkpoint is disposed").Value());
+                return;
+            }
+
+            {
+                std::unique_lock<std::shared_mutex> lock(checkpoint->dataMutex);
+                checkpoint->data.swap(data);
+                checkpoint->sequenceId = sequenceId;
+                checkpoint->minPos = minPos;
+                checkpoint->maxPos = maxPos;
+                checkpoint->initialized = true;
+            }
+            adjustNapiExternalMemorySubtract(Env(), data.size());
+            adjustNapiExternalMemoryAdd(Env(), checkpoint->data.size());
             deferred.Resolve(Env().Undefined());
         }
         void OnError(const Napi::Error& err) {
@@ -1209,9 +1272,12 @@ class AddonContextSequenceCheckpointInitWorker : public Napi::AsyncWorker {
 
 Napi::Value AddonContextSequenceCheckpoint::Init(const Napi::CallbackInfo& info) {
     AddonContext * context = Napi::ObjectWrap<AddonContext>::Unwrap(info[0].As<Napi::Object>());
-    sequenceId = info[1].As<Napi::Number>().Int32Value();
+    if (disposed || context->disposed || !context->contextLoaded) {
+        throw Napi::Error::New(info.Env(), "Checkpoint or context is disposed or not loaded");
+    }
+    const llama_seq_id requestedSequenceId = info[1].As<Napi::Number>().Int32Value();
 
-    AddonContextSequenceCheckpointInitWorker* worker = new AddonContextSequenceCheckpointInitWorker(info, this, context);
+    AddonContextSequenceCheckpointInitWorker* worker = new AddonContextSequenceCheckpointInitWorker(info, this, context, requestedSequenceId);
     worker->Queue();
     return worker->GetPromise();
 }
@@ -1222,9 +1288,12 @@ Napi::Value AddonContextSequenceCheckpoint::Dispose(const Napi::CallbackInfo& in
 }
 
 void AddonContextSequenceCheckpoint::dispose() {
-    std::lock_guard<std::mutex> lock(dataMutex);
-    data.clear();
-    data.resize(0);
+    std::unique_lock<std::shared_mutex> lock(dataMutex);
+    disposed = true;
+    adjustNapiExternalMemorySubtract(Env(), data.size());
+    std::vector<uint8_t>().swap(data);
+    minPos = -1;
+    maxPos = -1;
 }
 
 Napi::Value AddonContextSequenceCheckpoint::GetSize(const Napi::CallbackInfo& info) {

llama/addon/AddonContext.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <mutex>
+#include <shared_mutex>
 
 #include "llama.h"
 #include "napi.h"
@@ -11,8 +12,8 @@ class AddonContext : public Napi::ObjectWrap<AddonContext> {
     public:
         AddonModel* model;
         llama_context_params context_params;
-        llama_context* ctx;
-        llama_batch batch;
+        llama_context* ctx = nullptr;
+        llama_batch batch{};
         uint64_t batchMemorySize = 0;
         bool has_batch = false;
         int32_t batch_n_tokens = 0;
@@ -68,11 +69,13 @@ class AddonContext : public Napi::ObjectWrap<AddonContext> {
 
 class AddonContextSequenceCheckpoint : public Napi::ObjectWrap<AddonContextSequenceCheckpoint> {
     public:
-        std::mutex dataMutex;
+        std::shared_mutex dataMutex;
         std::vector<uint8_t> data;
         llama_seq_id sequenceId = 0;
-        std::size_t minPos = 0;
-        std::size_t maxPos = 0;
+        llama_pos minPos = -1;
+        llama_pos maxPos = -1;
+        bool initialized = false;
+        bool disposed = false;
 
         AddonContextSequenceCheckpoint(const Napi::CallbackInfo& info);
         ~AddonContextSequenceCheckpoint();

llama/addon/AddonGgufMetadata.cpp

@@ -43,7 +43,7 @@ static std::string getSplitPath(const std::string& splitPrefix, const uint16_t s
 
 AddonGgufMetadata::AddonGgufMetadata(const Napi::CallbackInfo& info)
     : Napi::ObjectWrap<AddonGgufMetadata>(info),
-      ggufMetadata(gguf_init_empty()) {
+      ggufMetadata(gguf_init_empty(), gguf_free) {
     if (ggufMetadata.get() == nullptr) {
         throw std::runtime_error("Failed to create an empty GGUF context");
     }
@@ -71,6 +71,7 @@ class AddonGgufMetadataInitWorker : public Napi::AsyncWorker {
         AddonGgufMetadata* addonGgufMetadata;
         std::vector<AddonGgufMetadataSource> sources;
         std::vector<Napi::Reference<Napi::Buffer<uint8_t>>> bufferRefs;
+        gguf_context_ptr ggufMetadata{gguf_init_empty()};
 
         AddonGgufMetadataInitWorker(const Napi::Env& env, AddonGgufMetadata* addonGgufMetadata)
             : Napi::AsyncWorker(env, "AddonGgufMetadataInitWorker"),
@@ -91,8 +92,6 @@ class AddonGgufMetadataInitWorker : public Napi::AsyncWorker {
 
         void Execute() {
             try {
-                gguf_context_ptr& ggufMetadata = addonGgufMetadata->ggufMetadata;
-
                 auto loadMetadataSource = [](const AddonGgufMetadataSource& itemSource, ggml_context_ptr& tensorContextGuard) {
                     struct ggml_context* tensorContext = nullptr;
                     struct gguf_init_params ggufParams = {
@@ -214,6 +213,11 @@ class AddonGgufMetadataInitWorker : public Napi::AsyncWorker {
             }
         }
         void OnOK() {
+            if (addonGgufMetadata->disposed) {
+                deferred.Reject(Napi::Error::New(Env(), "Metadata was disposed during initialization").Value());
+                return;
+            }
+            addonGgufMetadata->ggufMetadata = std::move(ggufMetadata);
             deferred.Resolve(Env().Undefined());
         }
         void OnError(const Napi::Error& err) {

llama/addon/AddonGgufMetadata.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <cstddef>
+#include <memory>
 
 #include "ggml-cpp.h"
 #include "napi.h"
@@ -34,7 +35,7 @@ struct AddonGgufMetadataSource {
 
 class AddonGgufMetadata : public Napi::ObjectWrap<AddonGgufMetadata> {
     public:
-        gguf_context_ptr ggufMetadata;
+        std::shared_ptr<gguf_context> ggufMetadata;
         bool disposed = false;
 
         AddonGgufMetadata(const Napi::CallbackInfo& info);