cui-llama.rn 1.3.0 → 1.3.3

package/android/src/main/CMakeLists.txt

@@ -13,12 +13,17 @@ set(
     ${RNLLAMA_LIB_DIR}/llama-sampling.cpp
     ${RNLLAMA_LIB_DIR}/llama-vocab.cpp
     ${RNLLAMA_LIB_DIR}/log.cpp
+    
+    ${RNLLAMA_LIB_DIR}/amx/amx.cpp
+    ${RNLLAMA_LIB_DIR}/amx/mmq.cpp
 
     ${RNLLAMA_LIB_DIR}/ggml-aarch64.c
     ${RNLLAMA_LIB_DIR}/llama-grammar.cpp
     ${RNLLAMA_LIB_DIR}/llama-sampling.cpp
     ${RNLLAMA_LIB_DIR}/llama-vocab.cpp
     ${RNLLAMA_LIB_DIR}/log.cpp
+    ${RNLLAMA_LIB_DIR}/json.hpp
+    ${RNLLAMA_LIB_DIR}/json-schema-to-grammar.cpp
 
     ${RNLLAMA_LIB_DIR}/ggml-aarch64.c
     ${RNLLAMA_LIB_DIR}/ggml-alloc.c
@@ -53,7 +58,7 @@ function(build_library target_name cpu_flags)
 
     target_link_libraries(${target_name} ${LOG_LIB} android)
 
-    target_compile_options(${target_name} PRIVATE -pthread ${cpu_flags})
+    target_compile_options(${target_name} PRIVATE -pthread ${cpu_flags} -DLM_GGML_USE_CPU -DLM_GGML_USE_CPU_AARCH64)
 
     if (${CMAKE_BUILD_TYPE} STREQUAL "Debug")
         target_compile_options(${target_name} PRIVATE -DRNLLAMA_ANDROID_ENABLE_LOGGING)

package/android/src/main/jni.cpp

@@ -259,7 +259,7 @@ Java_com_rnllama_LlamaContext_initContext(
 
     const char *model_path_chars = env->GetStringUTFChars(model_path_str, nullptr);
     defaultParams.model = model_path_chars;
-
+    
     defaultParams.n_ctx = n_ctx;
     defaultParams.n_batch = n_batch;
 
@@ -281,7 +281,7 @@ Java_com_rnllama_LlamaContext_initContext(
     int default_n_threads = max_threads == 4 ? 2 : min(4, max_threads);
     defaultParams.cpuparams.n_threads = n_threads > 0 ? n_threads : default_n_threads;
 
-    defaultParams.n_gpu_layers = n_gpu_layers;
+    // defaultParams.n_gpu_layers = n_gpu_layers;
     defaultParams.flash_attn = flash_attn;
 
     const char *cache_type_k_chars = env->GetStringUTFChars(cache_type_k, nullptr);
@@ -558,7 +558,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
     //llama_reset_timings(llama->ctx);
 
     llama->params.prompt = env->GetStringUTFChars(prompt, nullptr);
-    llama->params.sparams.seed = (seed == -1) ? time(NULL) : seed;
+    llama->params.sampling.seed = (seed == -1) ? time(NULL) : seed;
 
     int max_threads = std::thread::hardware_concurrency();
     // Use 2 threads by default on 4-core devices, 4 threads on more cores
@@ -566,9 +566,9 @@ Java_com_rnllama_LlamaContext_doCompletion(
     llama->params.cpuparams.n_threads = n_threads > 0 ? n_threads : default_n_threads;
 
     llama->params.n_predict = n_predict;
-    llama->params.sparams.ignore_eos = ignore_eos;
+    llama->params.sampling.ignore_eos = ignore_eos;
 
-    auto & sparams = llama->params.sparams;
+    auto & sparams = llama->params.sampling;
     sparams.temp = temperature;
     sparams.penalty_last_n = penalty_last_n;
     sparams.penalty_repeat = penalty_repeat;
@@ -693,7 +693,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
             auto tokenResult = createWriteableMap(env);
             putString(env, tokenResult, "token", to_send.c_str());
 
-            if (llama->params.sparams.n_probs > 0) {
+            if (llama->params.sampling.n_probs > 0) {
               const std::vector<llama_token> to_send_toks = common_tokenize(llama->ctx, to_send, false);
               size_t probs_pos = std::min(sent_token_probs_index, llama->generated_token_probs.size());
               size_t probs_stop_pos = std::min(sent_token_probs_index + to_send_toks.size(), llama->generated_token_probs.size());

package/cpp/amx/amx.cpp

@@ -0,0 +1,196 @@
+#include "amx.h"
+#include "common.h"
+#include "mmq.h"
+#include "ggml-backend-impl.h"
+#include "ggml-backend.h"
+#include "ggml-impl.h"
+#include "ggml-cpu.h"
+
+#if defined(__gnu_linux__)
+#include <sys/syscall.h>
+#include <unistd.h>
+#endif
+
+#include <cstdlib>
+#include <cstring>
+#include <memory>
+
+#if defined(__AMX_INT8__) && defined(__AVX512VNNI__)
+
+// AMX buffer interface
+static void lm_ggml_backend_amx_buffer_free_buffer(lm_ggml_backend_buffer_t buffer) {
+    free(buffer->context);
+}
+
+static void * lm_ggml_backend_amx_buffer_get_base(lm_ggml_backend_buffer_t buffer) {
+    return (void *)(buffer->context);
+}
+
+static void lm_ggml_backend_amx_buffer_memset_tensor(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
+    memset((char *)tensor->data + offset, value, size);
+
+    LM_GGML_UNUSED(buffer);
+}
+
+static void lm_ggml_backend_amx_buffer_set_tensor(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
+    if (qtype_has_amx_kernels(tensor->type)) {
+        lm_ggml_backend_amx_convert_weight(tensor, data, offset, size);
+    } else {
+        memcpy((char *)tensor->data + offset, data, size);
+    }
+
+    LM_GGML_UNUSED(buffer);
+}
+
+static void lm_ggml_backend_amx_buffer_get_tensor(lm_ggml_backend_buffer_t buffer, const struct lm_ggml_tensor * tensor, void * data, size_t offset, size_t size) {
+    LM_GGML_ASSERT(!qtype_has_amx_kernels(tensor->type));
+    memcpy(data, (const char *)tensor->data + offset, size);
+
+    LM_GGML_UNUSED(buffer);
+}
+
+static bool lm_ggml_backend_amx_buffer_cpy_tensor(lm_ggml_backend_buffer_t buffer, const struct lm_ggml_tensor * src, struct lm_ggml_tensor * dst) {
+    if (lm_ggml_backend_buffer_is_host(src->buffer)) {
+        if (qtype_has_amx_kernels(src->type)) {
+            lm_ggml_backend_amx_convert_weight(dst, src->data, 0, lm_ggml_nbytes(dst));
+        } else {
+            memcpy(dst->data, src->data, lm_ggml_nbytes(src));
+        }
+        return true;
+    }
+    return false;
+
+    LM_GGML_UNUSED(buffer);
+}
+
+static void lm_ggml_backend_amx_buffer_clear(lm_ggml_backend_buffer_t buffer, uint8_t value) {
+    memset(buffer->context, value, buffer->size);
+}
+
+static lm_ggml_backend_buffer_i lm_ggml_backend_amx_buffer_interface = {
+    /* .free_buffer     = */ lm_ggml_backend_amx_buffer_free_buffer,
+    /* .get_base        = */ lm_ggml_backend_amx_buffer_get_base,
+    /* .init_tensor     = */ NULL, // no initialization required
+    /* .memset_tensor   = */ lm_ggml_backend_amx_buffer_memset_tensor,
+    /* .set_tensor      = */ lm_ggml_backend_amx_buffer_set_tensor,
+    /* .get_tensor      = */ lm_ggml_backend_amx_buffer_get_tensor,
+    /* .cpy_tensor      = */ lm_ggml_backend_amx_buffer_cpy_tensor,
+    /* .clear           = */ lm_ggml_backend_amx_buffer_clear,
+    /* .reset           = */ NULL,
+};
+
+static const char * lm_ggml_backend_amx_buffer_type_get_name(lm_ggml_backend_buffer_type_t buft) {
+    return "AMX";
+
+    LM_GGML_UNUSED(buft);
+}
+
+static lm_ggml_backend_buffer_t lm_ggml_backend_amx_buffer_type_alloc_buffer(lm_ggml_backend_buffer_type_t buft, size_t size) {
+    void * data = aligned_alloc(TENSOR_ALIGNMENT, size);
+    if (data == NULL) {
+        fprintf(stderr, "%s: failed to allocate buffer of size %zu\n", __func__, size);
+        return NULL;
+    }
+
+    return lm_ggml_backend_buffer_init(buft, lm_ggml_backend_amx_buffer_interface, data, size);
+}
+
+static size_t lm_ggml_backend_amx_buffer_type_get_alignment(lm_ggml_backend_buffer_type_t buft) {
+    return TENSOR_ALIGNMENT;
+
+    LM_GGML_UNUSED(buft);
+}
+
+static size_t lm_ggml_backend_amx_buffer_type_get_alloc_size(lm_ggml_backend_buffer_type_t buft, const lm_ggml_tensor* tensor) {
+    return lm_ggml_backend_amx_get_alloc_size(tensor);
+
+    LM_GGML_UNUSED(buft);
+}
+
+static bool lm_ggml_backend_amx_buffer_type_is_host(lm_ggml_backend_buffer_type_t buft) {
+    return false;
+
+    LM_GGML_UNUSED(buft);
+}
+
+#define ARCH_GET_XCOMP_PERM     0x1022
+#define ARCH_REQ_XCOMP_PERM     0x1023
+#define XFEATURE_XTILECFG       17
+#define XFEATURE_XTILEDATA      18
+
+static bool lm_ggml_amx_init() {
+#if defined(__gnu_linux__)
+    if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA)) {
+        fprintf(stderr, "AMX is not ready to be used!\n");
+        return false;
+    }
+    return true;
+#elif defined(_WIN32)
+    return true;
+#endif
+}
+lm_ggml_backend_buffer_type_t lm_ggml_backend_amx_buffer_type() {
+    static struct lm_ggml_backend_buffer_type lm_ggml_backend_buffer_type_amx = {
+        /* .iface = */ {
+            /* .get_name         = */ lm_ggml_backend_amx_buffer_type_get_name,
+            /* .alloc_buffer     = */ lm_ggml_backend_amx_buffer_type_alloc_buffer,
+            /* .get_alignment    = */ lm_ggml_backend_amx_buffer_type_get_alignment,
+            /* .get_max_size     = */ NULL, // defaults to SIZE_MAX
+            /* .get_alloc_size   = */ lm_ggml_backend_amx_buffer_type_get_alloc_size,
+            /* .is_host          = */ lm_ggml_backend_amx_buffer_type_is_host,
+        },
+        /* .device  = */ lm_ggml_backend_reg_dev_get(lm_ggml_backend_cpu_reg(), 0),
+        /* .context = */ NULL,
+    };
+
+    if (!lm_ggml_amx_init()) {
+        return NULL;
+    }
+
+    return &lm_ggml_backend_buffer_type_amx;
+}
+
+bool lm_ggml_backend_amx_buft_is_amx(lm_ggml_backend_buffer_type_t buft) {
+    return buft->iface.get_name == lm_ggml_backend_amx_buffer_type_get_name;
+}
+
+bool lm_ggml_backend_amx_device_supports_op(const struct lm_ggml_tensor * op) {
+    // handle only 2d gemm for now
+    auto is_contiguous_2d = [](const struct lm_ggml_tensor * t) {
+        return lm_ggml_is_contiguous(t) && t->ne[3] == 1 && t->ne[2] == 1;
+    };
+
+    switch (op->op) {
+        case LM_GGML_OP_NONE:
+        case LM_GGML_OP_RESHAPE:
+        case LM_GGML_OP_VIEW:
+        case LM_GGML_OP_PERMUTE:
+        case LM_GGML_OP_TRANSPOSE:
+            return true;
+
+        case LM_GGML_OP_MUL_MAT: {
+            const struct lm_ggml_tensor * src0 = op->src[0];
+            const struct lm_ggml_tensor * src1 = op->src[1];
+
+            const enum lm_ggml_type type = src0->type;
+            const int64_t ne0 = op->ne[0];
+
+            // amx kernels enables for Q4_0, Q4_1, Q8_0, F16
+            // Q4_K, Q5_K, Q6_K, IQ4_XS enabled for QK_K = 256
+            bool has_amx_kernels = qtype_has_amx_kernels(type) || (type == LM_GGML_TYPE_F16);
+
+            bool can_use_amx =
+                is_contiguous_2d(src0) &&       // src0 must be contiguous
+                is_contiguous_2d(src1) &&       // src1 must be contiguous
+                src1->type == LM_GGML_TYPE_F32 &&  // src1 must be float32
+                has_amx_kernels &&              // with amx kernel impls
+                ne0 % (TILE_N * 2) == 0;        // out_features is 32x
+
+            return can_use_amx;
+        }
+        default:
+            return false;
+    }
+}
+
+#endif // defined(__AMX_INT8__) && defined(__AVX512VNNI__)

package/cpp/amx/amx.h

@@ -0,0 +1,20 @@
+#include "ggml-backend.h"
+#include "ggml-cpu-impl.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#if defined(__AMX_INT8__) && defined(__AVX512VNNI__)
+
+lm_ggml_backend_buffer_type_t lm_ggml_backend_amx_buffer_type(void);
+bool lm_ggml_backend_amx_buft_is_amx(lm_ggml_backend_buffer_type_t buft);
+bool lm_ggml_backend_amx_device_supports_op(const struct lm_ggml_tensor * op);
+void lm_ggml_backend_amx_mul_mat(const struct lm_ggml_compute_params * params, struct lm_ggml_tensor * dst);
+size_t lm_ggml_backend_amx_desired_wsize(const struct lm_ggml_tensor * dst);
+
+#endif
+
+#ifdef __cplusplus
+}
+#endif

package/cpp/amx/common.h

@@ -0,0 +1,101 @@
+#pragma once
+
+#include "ggml.h"
+#include "ggml-cpu-impl.h"
+
+#include <algorithm>
+#include <memory>
+#include <type_traits>
+
+#if defined(_OPENMP)
+#include <omp.h>
+#endif
+
+#define TILE_M 16
+#define TILE_N 16
+#define TILE_K 32
+#define VNNI_BLK 4
+
+#define AMX_BLK_SIZE 32
+
+#define TMM0 0
+#define TMM1 1
+#define TMM2 2
+#define TMM3 3
+#define TMM4 4
+#define TMM5 5
+#define TMM6 6
+#define TMM7 7
+
+// parallel routines
+template <typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0>
+inline T div_up(T x, T y) { return (x + y - 1) / y; }
+
+template <typename T>
+inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
+#if 0
+    // onednn partition pattern
+    T& n_my = n_end;
+    if (nth <= 1 || n == 0) {
+        n_start = 0;
+        n_my = n;
+    } else {
+        T n1 = div_up(n, nth);
+        T n2 = n1 - 1;
+        T T1 = n - n2 * nth;
+        n_my = ith < T1 ? n1 : n2;
+        n_start = ith <= T1 ? ith*n1 : T1 * n1 + (ith - T1) * n2;
+    }
+    n_end += n_start;
+#else
+    // pytorch aten partition pattern
+    T n_my = div_up(n, nth);
+    n_start = ith * n_my;
+    n_end = std::min(n_start + n_my, n);
+#endif
+}
+
+template <typename func_t>
+inline void parallel_for(int nth, int n, const func_t& f) {
+#if defined(_OPENMP)
+#pragma omp parallel num_threads(nth)
+{
+    //int nth = omp_get_num_threads();
+    int ith = omp_get_thread_num();
+    int tbegin, tend;
+    balance211(n, nth, ith, tbegin, tend);
+    f(tbegin, tend);
+}
+#else
+    f(0, n);
+
+    LM_GGML_UNUSED(nth);
+#endif
+}
+
+template <typename func_t>
+inline void parallel_for_ggml(const lm_ggml_compute_params * params, int n, const func_t & f) {
+    int tbegin, tend;
+    balance211(n, params->nth, params->ith, tbegin, tend);
+    f(tbegin, tend);
+    lm_ggml_barrier(params->threadpool); // TODO: might not always be needed
+}
+
+// quantized types that have AMX support
+inline bool qtype_has_amx_kernels(const enum lm_ggml_type type) {
+    // TODO: fix padding for vnni format
+    return (type == LM_GGML_TYPE_Q4_0) ||
+        (type == LM_GGML_TYPE_Q4_1) ||
+        (type == LM_GGML_TYPE_Q8_0) ||
+        (type == LM_GGML_TYPE_Q4_K) ||
+        (type == LM_GGML_TYPE_Q5_K) ||
+        (type == LM_GGML_TYPE_Q6_K) ||
+        (type == LM_GGML_TYPE_IQ4_XS);
+}
+
+// ggml backend context
+struct lm_ggml_backend_amx_context {
+    int n_threads = LM_GGML_DEFAULT_N_THREADS;
+    std::unique_ptr<char[]> work_data;
+    size_t work_size = 0;
+};