cui-llama.rn 1.4.4 → 1.5.0

package/android/src/main/CMakeLists.txt

@@ -19,6 +19,10 @@ set(
     ${RNLLAMA_LIB_DIR}/ggml-alloc.c
     ${RNLLAMA_LIB_DIR}/ggml-backend.cpp
     ${RNLLAMA_LIB_DIR}/ggml-backend-reg.cpp
+    ${RNLLAMA_LIB_DIR}/ops.cpp
+    ${RNLLAMA_LIB_DIR}/unary-ops.cpp
+    ${RNLLAMA_LIB_DIR}/binary-ops.cpp
+    ${RNLLAMA_LIB_DIR}/vec.cpp
     ${RNLLAMA_LIB_DIR}/ggml-cpu.c
     ${RNLLAMA_LIB_DIR}/ggml-cpu.cpp
     ${RNLLAMA_LIB_DIR}/ggml-cpu-aarch64.cpp
@@ -46,16 +50,19 @@ set(
     ${RNLLAMA_LIB_DIR}/llama-model-loader.cpp
     ${RNLLAMA_LIB_DIR}/llama-mmap.cpp
     ${RNLLAMA_LIB_DIR}/llama-vocab.cpp
+    ${RNLLAMA_LIB_DIR}/llama-memory.cpp
+    ${RNLLAMA_LIB_DIR}/llama-io.cpp
+    ${RNLLAMA_LIB_DIR}/llama-graph.cpp
     ${RNLLAMA_LIB_DIR}/sampling.cpp
     ${RNLLAMA_LIB_DIR}/unicode-data.cpp
     ${RNLLAMA_LIB_DIR}/unicode.cpp
     ${RNLLAMA_LIB_DIR}/sgemm.cpp
     ${RNLLAMA_LIB_DIR}/common.cpp
     ${RNLLAMA_LIB_DIR}/chat.cpp
-    ${RNLLAMA_LIB_DIR}/chat-template.hpp
     ${RNLLAMA_LIB_DIR}/json-schema-to-grammar.cpp
-    ${RNLLAMA_LIB_DIR}/minja.hpp
     ${RNLLAMA_LIB_DIR}/json.hpp
+    ${RNLLAMA_LIB_DIR}/minja/minja.hpp
+    ${RNLLAMA_LIB_DIR}/minja/chat-template.hpp
     ${RNLLAMA_LIB_DIR}/rn-llama.cpp
     ${CMAKE_SOURCE_DIR}/jni-utils.h
     ${CMAKE_SOURCE_DIR}/jni.cpp

package/android/src/main/jni.cpp

@@ -11,6 +11,7 @@
 #include <unordered_map>
 #include "json-schema-to-grammar.h"
 #include "llama.h"
+#include "chat.h"
 #include "llama-impl.h"
 #include "ggml.h"
 #include "rn-llama.h"
@@ -263,7 +264,7 @@ Java_com_rnllama_LlamaContext_initContext(
     }
 
     const char *model_path_chars = env->GetStringUTFChars(model_path_str, nullptr);
-    defaultParams.model = model_path_chars;
+    defaultParams.model = { model_path_chars };
 
     const char *chat_template_chars = env->GetStringUTFChars(chat_template, nullptr);
     defaultParams.chat_template = chat_template_chars;
@@ -445,7 +446,7 @@ Java_com_rnllama_LlamaContext_loadModelDetails(
 
     auto default_caps = createWriteableMap(env);
 
-    auto default_tmpl = llama->templates.template_default.get();
+    auto default_tmpl = llama->templates.get()->template_default.get();
     auto default_tmpl_caps = default_tmpl->original_caps();
     putBoolean(env, default_caps, "tools", default_tmpl_caps.supports_tools);
     putBoolean(env, default_caps, "toolCalls", default_tmpl_caps.supports_tool_calls);
@@ -456,7 +457,7 @@ Java_com_rnllama_LlamaContext_loadModelDetails(
     putMap(env, minja, "defaultCaps", default_caps);
 
     putBoolean(env, minja, "toolUse", llama->validateModelChatTemplate(true, "tool_use"));
-    auto tool_use_tmpl = llama->templates.template_tool_use.get();
+    auto tool_use_tmpl = llama->templates.get()->template_tool_use.get();
     if (tool_use_tmpl != nullptr) {
       auto tool_use_caps = createWriteableMap(env);
       auto tool_use_tmpl_caps = tool_use_tmpl->original_caps();
@@ -510,15 +511,16 @@ Java_com_rnllama_LlamaContext_getFormattedChatWithJinja(
             parallel_tool_calls,
             tool_choice_chars
         );
-        putString(env, result, "prompt", formatted.prompt.get<std::string>().c_str());
+        putString(env, result, "prompt", formatted.prompt.c_str());
         putInt(env, result, "chat_format", static_cast<int>(formatted.format));
         putString(env, result, "grammar", formatted.grammar.c_str());
         putBoolean(env, result, "grammar_lazy", formatted.grammar_lazy);
         auto grammar_triggers = createWritableArray(env);
         for (const auto &trigger : formatted.grammar_triggers) {
             auto trigger_map = createWriteableMap(env);
-            putString(env, trigger_map, "word", trigger.word.c_str());
-            putBoolean(env, trigger_map, "at_start", trigger.at_start);
+            putInt(env, trigger_map, "type", trigger.type);
+            putString(env, trigger_map, "value", trigger.value.c_str());
+            putInt(env, trigger_map, "token", trigger.token);
             pushMap(env, grammar_triggers, trigger_map);
         }
         putArray(env, result, "grammar_triggers", grammar_triggers);
@@ -732,23 +734,53 @@ Java_com_rnllama_LlamaContext_doCompletion(
       sparams.grammar = grammar_chars;
     }
     sparams.grammar_lazy = grammar_lazy;
+
+    if (preserved_tokens != nullptr) {
+        int preserved_tokens_size = readablearray::size(env, preserved_tokens);
+        for (int i = 0; i < preserved_tokens_size; i++) {
+            jstring preserved_token = readablearray::getString(env, preserved_tokens, i);
+            auto ids = common_tokenize(llama->ctx, env->GetStringUTFChars(preserved_token, nullptr), /* add_special= */ false, /* parse_special= */ true);
+            if (ids.size() == 1) {
+                sparams.preserved_tokens.insert(ids[0]);
+            } else {
+                LOGI("[RNLlama] Not preserved because more than 1 token (wrong chat template override?): %s", env->GetStringUTFChars(preserved_token, nullptr));
+            }
+        }
+    }
+
     if (grammar_triggers != nullptr) {
         int grammar_triggers_size = readablearray::size(env, grammar_triggers);
         for (int i = 0; i < grammar_triggers_size; i++) {
-            common_grammar_trigger trigger;
             auto trigger_map = readablearray::getMap(env, grammar_triggers, i);
-            jstring trigger_word = readablemap::getString(env, trigger_map, "word", nullptr);
-            jboolean trigger_at_start = readablemap::getBool(env, trigger_map, "at_start", false);
-            trigger.word = env->GetStringUTFChars(trigger_word, nullptr);
-            trigger.at_start = trigger_at_start;
-
-            auto ids = common_tokenize(llama->ctx, trigger.word, /* add_special= */ false, /* parse_special= */ true);
-            if (ids.size() == 1) {
-                sparams.grammar_trigger_tokens.push_back(ids[0]);
-                sparams.preserved_tokens.insert(ids[0]);
-                continue;
+            const auto type = static_cast<common_grammar_trigger_type>(readablemap::getInt(env, trigger_map, "type", 0));
+            jstring trigger_word = readablemap::getString(env, trigger_map, "value", nullptr);
+            auto word = env->GetStringUTFChars(trigger_word, nullptr);
+
+            if (type == COMMON_GRAMMAR_TRIGGER_TYPE_WORD) {
+                auto ids = common_tokenize(llama->ctx, word, /* add_special= */ false, /* parse_special= */ true);
+                if (ids.size() == 1) {
+                    auto token = ids[0];
+                    if (std::find(sparams.preserved_tokens.begin(), sparams.preserved_tokens.end(), (llama_token) token) == sparams.preserved_tokens.end()) {
+                        throw std::runtime_error("Grammar trigger word should be marked as preserved token");
+                    }
+                    common_grammar_trigger trigger;
+                    trigger.type = COMMON_GRAMMAR_TRIGGER_TYPE_TOKEN;
+                    trigger.value = word;
+                    trigger.token = token;
+                    sparams.grammar_triggers.push_back(std::move(trigger));
+                } else {
+                    sparams.grammar_triggers.push_back({COMMON_GRAMMAR_TRIGGER_TYPE_WORD, word});
+                }
+            } else {
+                common_grammar_trigger trigger;
+                trigger.type = type;
+                trigger.value = word;
+                if (type == COMMON_GRAMMAR_TRIGGER_TYPE_TOKEN) {
+                    const auto token = (llama_token) readablemap::getInt(env, trigger_map, "token", 0);
+                    trigger.token = token;
+                }
+                sparams.grammar_triggers.push_back(std::move(trigger));
             }
-            sparams.grammar_trigger_words.push_back(trigger);
         }
     }
 
@@ -759,18 +791,6 @@ Java_com_rnllama_LlamaContext_doCompletion(
     }
     env->ReleaseStringUTFChars(json_schema, json_schema_chars);
 
-    if (preserved_tokens != nullptr) {
-        int preserved_tokens_size = readablearray::size(env, preserved_tokens);
-        for (int i = 0; i < preserved_tokens_size; i++) {
-            jstring preserved_token = readablearray::getString(env, preserved_tokens, i);
-            auto ids = common_tokenize(llama->ctx, env->GetStringUTFChars(preserved_token, nullptr), /* add_special= */ false, /* parse_special= */ true);
-            if (ids.size() == 1) {
-                sparams.preserved_tokens.insert(ids[0]);
-            } else {
-                LOGI("[RNLlama] Not preserved because more than 1 token (wrong chat template override?): %s", env->GetStringUTFChars(preserved_token, nullptr));
-            }
-        }
-    }
 
     const llama_model * model = llama_get_model(llama->ctx);
     const llama_vocab * vocab = llama_model_get_vocab(model);
@@ -902,7 +922,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
 
     auto toolCalls = createWritableArray(env);
     std::string reasoningContent = "";
-    std::string *content = nullptr;
+    std::string content;
     auto toolCallsSize = 0;
     if (!llama->is_interrupted) {
         try {
@@ -910,7 +930,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
             if (!message.reasoning_content.empty()) {
                 reasoningContent = message.reasoning_content;
             }
-            content = &message.content;
+            content = message.content;
             for (const auto &tc : message.tool_calls) {
                 auto toolCall = createWriteableMap(env);
                 putString(env, toolCall, "type", "function");
@@ -931,8 +951,8 @@ Java_com_rnllama_LlamaContext_doCompletion(
 
     auto result = createWriteableMap(env);
     putString(env, result, "text", llama->generated_text.c_str());
-    if (content) {
-        putString(env, result, "content", content->c_str());
+    if (!content.empty()) {
+        putString(env, result, "content", content.c_str());
     }
     if (!reasoningContent.empty()) {
         putString(env, result, "reasoning_content", reasoningContent.c_str());

package/cpp/binary-ops.cpp

@@ -0,0 +1,158 @@
+#include "binary-ops.h"
+
+#if defined(LM_GGML_USE_ACCELERATE)
+#include <Accelerate/Accelerate.h>
+
+using vDSP_fn_t = void (*)(const float *, vDSP_Stride, const float *, vDSP_Stride, float *, vDSP_Stride, vDSP_Length);
+#endif
+
+static inline float op_add(float a, float b) {
+    return a + b;
+}
+
+static inline float op_sub(float a, float b) {
+    return a - b;
+}
+
+static inline float op_mul(float a, float b) {
+    return a * b;
+}
+
+static inline float op_div(float a, float b) {
+    return a / b;
+}
+
+template <float (*op)(float, float), typename src0_t, typename src1_t, typename dst_t>
+static inline void vec_binary_op_contiguous(const int64_t n, dst_t * z, const src0_t * x, const src1_t * y) {
+    constexpr auto src0_to_f32 = type_conversion_table<src0_t>::to_f32;
+    constexpr auto src1_to_f32 = type_conversion_table<src1_t>::to_f32;
+    constexpr auto f32_to_dst  = type_conversion_table<dst_t >::from_f32;
+
+    for (int i = 0; i < n; i++) {
+        z[i] = f32_to_dst(op(src0_to_f32(x[i]), src1_to_f32(y[i])));
+    }
+}
+
+template <float (*op)(float, float), typename src0_t, typename src1_t, typename dst_t>
+static inline void vec_binary_op_non_contiguous(const int64_t n, const int64_t ne10, const int64_t nb10, dst_t * z, const src0_t * x, const src1_t * y) {
+    constexpr auto src0_to_f32 = type_conversion_table<src0_t>::to_f32;
+    constexpr auto src1_to_f32 = type_conversion_table<src1_t>::to_f32;
+    constexpr auto f32_to_dst  = type_conversion_table<dst_t >::from_f32;
+
+    for (int i = 0; i < n; i++) {
+        int i10 = i % ne10;
+        const src1_t * y_ptr = (const src1_t *)((const char *)y + i10*nb10);
+        z[i] = f32_to_dst(op(src0_to_f32(x[i]), src1_to_f32(*y_ptr)));
+    }
+}
+
+template <float (*op)(float, float), typename src0_t, typename src1_t, typename dst_t>
+static void apply_binary_op(const lm_ggml_compute_params * params, lm_ggml_tensor * dst) {
+    const lm_ggml_tensor * src0 = dst->src[0];
+    const lm_ggml_tensor * src1 = dst->src[1];
+
+    LM_GGML_ASSERT(lm_ggml_can_repeat(src1, src0) && lm_ggml_are_same_shape(src0, dst));
+
+    LM_GGML_TENSOR_BINARY_OP_LOCALS
+
+    LM_GGML_ASSERT( nb0 == sizeof(dst_t));
+    LM_GGML_ASSERT(nb00 == sizeof(src0_t));
+
+    const auto [ir0, ir1] = get_thread_range(params, src0);
+    const bool is_src1_contiguous = (nb10 == sizeof(src1_t));
+
+    if (!is_src1_contiguous) { // broadcast not implemented yet for non-contiguous
+        LM_GGML_ASSERT(lm_ggml_are_same_shape(src0, src1));
+    }
+
+#ifdef LM_GGML_USE_ACCELERATE
+    vDSP_fn_t vDSP_op = nullptr;
+    // TODO - avoid the f32-only check using type 'trait' lookup tables and row-based src-to-float conversion functions
+    if (src0->type == LM_GGML_TYPE_F32 && src1->type == LM_GGML_TYPE_F32 && dst->type == LM_GGML_TYPE_F32) {
+        if (op == op_add) {
+            vDSP_op = vDSP_vadd;
+        } else if (op == op_sub) {
+            vDSP_op = vDSP_vsub;
+        } else if (op == op_mul) {
+            vDSP_op = vDSP_vmul;
+        } else if (op == op_div) {
+            vDSP_op = vDSP_vdiv;
+        }
+    }
+#endif
+
+    for (int64_t ir = ir0; ir < ir1; ++ir) {
+        const int64_t i03 = ir/(ne02*ne01);
+        const int64_t i02 = (ir - i03*ne02*ne01)/ne01;
+        const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01);
+
+        const int64_t i13 = i03 % ne13;
+        const int64_t i12 = i02 % ne12;
+        const int64_t i11 = i01 % ne11;
+
+        dst_t        * dst_ptr  = (dst_t  *)       ((char *)       dst->data  + i03*nb3  + i02*nb2  + i01*nb1 );
+        const src0_t * src0_ptr = (const src0_t *) ((const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01);
+        const src1_t * src1_ptr = (const src1_t *) ((const char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11);
+
+        if (is_src1_contiguous) {
+            // src1 is broadcastable across src0 and dst in i1, i2, i3
+            const int64_t nr0 = ne00 / ne10;
+
+            for (int64_t r = 0; r < nr0; ++r) {
+#ifdef LM_GGML_USE_ACCELERATE
+                if constexpr (std::is_same_v<src0_t, float> && std::is_same_v<src1_t, float> && std::is_same_v<dst_t, float>) {
+                    if (vDSP_op != nullptr) {
+                        vDSP_op(src1_ptr, 1, src0_ptr + r*ne10, 1, dst_ptr + r*ne10, 1, ne10);
+                        continue;
+                    }
+                }
+#endif
+                vec_binary_op_contiguous<op>(ne10, dst_ptr + r*ne10, src0_ptr + r*ne10, src1_ptr);
+            }
+        } else {
+            vec_binary_op_non_contiguous<op>(ne0, ne10, nb10, dst_ptr, src0_ptr, src1_ptr);
+        }
+    }
+}
+
+// TODO: Use the 'traits' lookup table (for type conversion fns), instead of a mass of 'if' conditions with long templates
+template <float (*op)(float, float)>
+static void binary_op(const lm_ggml_compute_params * params, lm_ggml_tensor * dst) {
+    const lm_ggml_tensor * src0 = dst->src[0];
+    const lm_ggml_tensor * src1 = dst->src[1];
+
+    /*  */ if (src0->type == LM_GGML_TYPE_F32  && src1->type == LM_GGML_TYPE_F32  && dst->type == LM_GGML_TYPE_F32) { // all f32
+        apply_binary_op<op, float, float, float>(params, dst);
+    } else if (src0->type == LM_GGML_TYPE_F16  && src1->type == LM_GGML_TYPE_F16  && dst->type == LM_GGML_TYPE_F16) { // all f16
+        apply_binary_op<op, lm_ggml_fp16_t, lm_ggml_fp16_t, lm_ggml_fp16_t>(params, dst);
+    } else if (src0->type == LM_GGML_TYPE_BF16 && src1->type == LM_GGML_TYPE_BF16 && dst->type == LM_GGML_TYPE_BF16) { // all bf16
+        apply_binary_op<op, lm_ggml_bf16_t, lm_ggml_bf16_t, lm_ggml_bf16_t>(params, dst);
+    } else if (src0->type == LM_GGML_TYPE_BF16 && src1->type == LM_GGML_TYPE_F32  && dst->type == LM_GGML_TYPE_BF16) {
+        apply_binary_op<op, lm_ggml_bf16_t, float, lm_ggml_bf16_t>(params, dst);
+    } else if (src0->type == LM_GGML_TYPE_BF16 && src1->type == LM_GGML_TYPE_F32  && dst->type == LM_GGML_TYPE_F32) {
+        apply_binary_op<op, lm_ggml_bf16_t, float, float>(params, dst);
+    } else if (src0->type == LM_GGML_TYPE_F16  && src1->type == LM_GGML_TYPE_F32  && dst->type == LM_GGML_TYPE_F16) {
+        apply_binary_op<op, lm_ggml_fp16_t, float, lm_ggml_fp16_t>(params, dst);
+    } else if (src0->type == LM_GGML_TYPE_F16  && src1->type == LM_GGML_TYPE_F32  && dst->type == LM_GGML_TYPE_F32) {
+        apply_binary_op<op, lm_ggml_fp16_t, float, float>(params, dst);
+    } else {
+        LM_GGML_ABORT("%s: unsupported types: dst: %s, src0: %s, src1: %s\n", __func__,
+            lm_ggml_type_name(dst->type), lm_ggml_type_name(src0->type), lm_ggml_type_name(src1->type));
+    }
+}
+
+void lm_ggml_compute_forward_add_non_quantized(const lm_ggml_compute_params * params, lm_ggml_tensor * dst) {
+    binary_op<op_add>(params, dst);
+}
+
+void lm_ggml_compute_forward_sub(const lm_ggml_compute_params * params, lm_ggml_tensor * dst) {
+    binary_op<op_sub>(params, dst);
+}
+
+void lm_ggml_compute_forward_mul(const lm_ggml_compute_params * params, lm_ggml_tensor * dst) {
+    binary_op<op_mul>(params, dst);
+}
+
+void lm_ggml_compute_forward_div(const lm_ggml_compute_params * params, lm_ggml_tensor * dst) {
+    binary_op<op_div>(params, dst);
+}

package/cpp/binary-ops.h

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "cpu-common.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+void lm_ggml_compute_forward_add_non_quantized(const struct lm_ggml_compute_params * params, struct lm_ggml_tensor * dst);
+void lm_ggml_compute_forward_sub(const struct lm_ggml_compute_params * params, struct lm_ggml_tensor * dst);
+void lm_ggml_compute_forward_mul(const struct lm_ggml_compute_params * params, struct lm_ggml_tensor * dst);
+void lm_ggml_compute_forward_div(const struct lm_ggml_compute_params * params, struct lm_ggml_tensor * dst);
+
+#ifdef __cplusplus
+}
+#endif