cui-llama.rn 1.2.6 → 1.3.3

package/android/src/main/jni.cpp

@@ -4,13 +4,15 @@
 #include <android/log.h>
 #include <cstdlib>
 #include <ctime>
+#include <ctime>
 #include <sys/sysinfo.h>
 #include <string>
 #include <thread>
 #include <unordered_map>
 #include "llama.h"
-#include "rn-llama.hpp"
+#include "llama-impl.h"
 #include "ggml.h"
+#include "rn-llama.hpp"
 
 #define UNUSED(x) (void)(x)
 #define TAG "RNLLAMA_ANDROID_JNI"
@@ -22,6 +24,13 @@ static inline int min(int a, int b) {
     return (a < b) ? a : b;
 }
 
+static void log_callback(lm_ggml_log_level level, const char * fmt, void * data) {
+    if (level == LM_GGML_LOG_LEVEL_ERROR)     __android_log_print(ANDROID_LOG_ERROR, TAG, fmt, data);
+    else if (level == LM_GGML_LOG_LEVEL_INFO) __android_log_print(ANDROID_LOG_INFO, TAG, fmt, data);
+    else if (level == LM_GGML_LOG_LEVEL_WARN) __android_log_print(ANDROID_LOG_WARN, TAG, fmt, data);
+    else __android_log_print(ANDROID_LOG_DEFAULT, TAG, fmt, data);
+}
+
 extern "C" {
 
 // Method to create WritableMap
@@ -107,6 +116,15 @@ static inline void pushDouble(JNIEnv *env, jobject arr, double value) {
     env->CallVoidMethod(arr, pushDoubleMethod, value);
 }
 
+// Method to push string into WritableArray
+static inline void pushString(JNIEnv *env, jobject arr, const char *value) {
+    jclass mapClass = env->FindClass("com/facebook/react/bridge/WritableArray");
+    jmethodID pushStringMethod = env->GetMethodID(mapClass, "pushString", "(Ljava/lang/String;)V");
+
+    jstring jValue = env->NewStringUTF(value);
+    env->CallVoidMethod(arr, pushStringMethod, jValue);
+}
+
 // Method to push WritableMap into WritableArray
 static inline void pushMap(JNIEnv *env, jobject arr, jobject value) {
     jclass mapClass = env->FindClass("com/facebook/react/bridge/WritableArray");
@@ -125,6 +143,77 @@ static inline void putArray(JNIEnv *env, jobject map, const char *key, jobject v
     env->CallVoidMethod(map, putArrayMethod, jKey, value);
 }
 
+JNIEXPORT jobject JNICALL
+Java_com_rnllama_LlamaContext_modelInfo(
+    JNIEnv *env,
+    jobject thiz,
+    jstring model_path_str,
+    jobjectArray skip
+) {
+    UNUSED(thiz);
+
+    const char *model_path_chars = env->GetStringUTFChars(model_path_str, nullptr);
+
+    std::vector<std::string> skip_vec;
+    int skip_len = env->GetArrayLength(skip);
+    for (int i = 0; i < skip_len; i++) {
+        jstring skip_str = (jstring) env->GetObjectArrayElement(skip, i);
+        const char *skip_chars = env->GetStringUTFChars(skip_str, nullptr);
+        skip_vec.push_back(skip_chars);
+        env->ReleaseStringUTFChars(skip_str, skip_chars);
+    }
+
+    struct lm_gguf_init_params params = {
+        /*.no_alloc = */ false,
+        /*.ctx      = */ NULL,
+    };
+    struct lm_gguf_context * ctx = lm_gguf_init_from_file(model_path_chars, params);
+
+    if (!ctx) {
+        LOGI("%s: failed to load '%s'\n", __func__, model_path_chars);
+        return nullptr;
+    }
+
+    auto info = createWriteableMap(env);
+    putInt(env, info, "version", lm_gguf_get_version(ctx));
+    putInt(env, info, "alignment", lm_gguf_get_alignment(ctx));
+    putInt(env, info, "data_offset", lm_gguf_get_data_offset(ctx));
+    {
+        const int n_kv = lm_gguf_get_n_kv(ctx);
+
+        for (int i = 0; i < n_kv; ++i) {
+            const char * key = lm_gguf_get_key(ctx, i);
+
+            bool skipped = false;
+            if (skip_len > 0) {
+                for (int j = 0; j < skip_len; j++) {
+                    if (skip_vec[j] == key) {
+                        skipped = true;
+                        break;
+                    }
+                }
+            }
+
+            if (skipped) {
+                continue;
+            }
+
+            const std::string value = rnllama::lm_gguf_kv_to_str(ctx, i);
+            putString(env, info, key, value.c_str());
+        }
+    }
+
+    env->ReleaseStringUTFChars(model_path_str, model_path_chars);
+    lm_gguf_free(ctx);
+
+    return reinterpret_cast<jobject>(info);
+}
+
+struct callback_context {
+    JNIEnv *env;
+    rnllama::llama_rn_context *llama;
+    jobject callback;
+};
 
 std::unordered_map<long, rnllama::llama_rn_context *> context_map;
 
@@ -141,10 +230,14 @@ Java_com_rnllama_LlamaContext_initContext(
     jobject thiz,
     jstring model_path_str,
     jboolean embedding,
+    jint embd_normalize,
     jint n_ctx,
     jint n_batch,
     jint n_threads,
     jint n_gpu_layers, // TODO: Support this
+    jboolean flash_attn,
+    jstring cache_type_k,
+    jstring cache_type_v,
     jboolean use_mlock,
     jboolean use_mmap,
     jboolean vocab_only,
@@ -152,7 +245,8 @@ Java_com_rnllama_LlamaContext_initContext(
     jfloat lora_scaled,
     jfloat rope_freq_base,
     jfloat rope_freq_scale,
-    jobject javaLlamaContext
+    jint pooling_type,
+    jobject load_progress_callback
 ) {
     UNUSED(thiz);
 
@@ -165,65 +259,110 @@ Java_com_rnllama_LlamaContext_initContext(
 
     const char *model_path_chars = env->GetStringUTFChars(model_path_str, nullptr);
     defaultParams.model = model_path_chars;
-
-    defaultParams.embedding = embedding;
-
+    
     defaultParams.n_ctx = n_ctx;
     defaultParams.n_batch = n_batch;
 
+    if (pooling_type != -1) {
+        defaultParams.pooling_type = static_cast<enum llama_pooling_type>(pooling_type);
+    }
+
+    defaultParams.embedding = embedding;
+    if (embd_normalize != -1) {
+        defaultParams.embd_normalize = embd_normalize;
+    }
+    if (embedding) {
+        // For non-causal models, batch size must be equal to ubatch size
+        defaultParams.n_ubatch = defaultParams.n_batch;
+    }
+
     int max_threads = std::thread::hardware_concurrency();
     // Use 2 threads by default on 4-core devices, 4 threads on more cores
     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);
+    const char *cache_type_v_chars = env->GetStringUTFChars(cache_type_v, nullptr);
+    defaultParams.cache_type_k = cache_type_k_chars;
+    defaultParams.cache_type_v = cache_type_v_chars;
+
     defaultParams.use_mlock = use_mlock;
     defaultParams.use_mmap = use_mmap;
 
     const char *lora_chars = env->GetStringUTFChars(lora_str, nullptr);
     if (lora_chars != nullptr && lora_chars[0] != '\0') {
         defaultParams.lora_adapters.push_back({lora_chars, lora_scaled});
-        defaultParams.use_mmap = false;
     }
 
     defaultParams.rope_freq_base = rope_freq_base;
     defaultParams.rope_freq_scale = rope_freq_scale;
 
-    // progress callback when loading
-    jclass llamaContextClass = env->GetObjectClass(javaLlamaContext);
-    jmethodID sendProgressMethod = env->GetMethodID(llamaContextClass, "emitModelProgressUpdate", "(I)V");
-
-    CallbackContext callbackctx = {env, javaLlamaContext, sendProgressMethod, 0};
-
-    defaultParams.progress_callback_user_data = &callbackctx;
-    defaultParams.progress_callback = [](float progress, void * ctx) {
-        unsigned percentage = (unsigned) (100 * progress);
-        CallbackContext * cbctx = static_cast<CallbackContext*>(ctx);
-        // reduce call frequency by only calling method when value changes
-        if (percentage <= cbctx->current) return true; 
-        cbctx->current = percentage;
-        cbctx->env->CallVoidMethod(cbctx->thiz, cbctx->sendProgressMethod, percentage);
-        return true;
-    };
-    
-
     auto llama = new rnllama::llama_rn_context();
+    llama->is_load_interrupted = false;
+    llama->loading_progress = 0;
+
+    if (load_progress_callback != nullptr) {
+        defaultParams.progress_callback = [](float progress, void * user_data) {
+            callback_context *cb_ctx = (callback_context *)user_data;
+            JNIEnv *env = cb_ctx->env;
+            auto llama = cb_ctx->llama;
+            jobject callback = cb_ctx->callback;
+            int percentage = (int) (100 * progress);
+            if (percentage > llama->loading_progress) {
+                llama->loading_progress = percentage;
+                jclass callback_class = env->GetObjectClass(callback);
+                jmethodID onLoadProgress = env->GetMethodID(callback_class, "onLoadProgress", "(I)V");
+                env->CallVoidMethod(callback, onLoadProgress, percentage);
+            }
+            return !llama->is_load_interrupted;
+        };
+
+        callback_context *cb_ctx = new callback_context;
+        cb_ctx->env = env;
+        cb_ctx->llama = llama;
+        cb_ctx->callback = env->NewGlobalRef(load_progress_callback);
+        defaultParams.progress_callback_user_data = cb_ctx;
+    }
+
     bool is_model_loaded = llama->loadModel(defaultParams);
 
+    env->ReleaseStringUTFChars(model_path_str, model_path_chars);
+    env->ReleaseStringUTFChars(lora_str, lora_chars);
+    env->ReleaseStringUTFChars(cache_type_k, cache_type_k_chars);
+    env->ReleaseStringUTFChars(cache_type_v, cache_type_v_chars);
+
     LOGI("[RNLlama] is_model_loaded %s", (is_model_loaded ? "true" : "false"));
     if (is_model_loaded) {
+      if (embedding && llama_model_has_encoder(llama->model) && llama_model_has_decoder(llama->model)) {
+        LOGI("[RNLlama] computing embeddings in encoder-decoder models is not supported");
+        llama_free(llama->ctx);
+        return -1;
+      }
       context_map[(long) llama->ctx] = llama;
     } else {
       llama_free(llama->ctx);
     }
 
-    env->ReleaseStringUTFChars(model_path_str, model_path_chars);
-    env->ReleaseStringUTFChars(lora_str, lora_chars);
-
     return reinterpret_cast<jlong>(llama->ctx);
 }
 
+
+JNIEXPORT void JNICALL
+Java_com_rnllama_LlamaContext_interruptLoad(
+    JNIEnv *env,
+    jobject thiz,
+    jlong context_ptr
+) {
+    UNUSED(thiz);
+    auto llama = context_map[(long) context_ptr];
+    if (llama) {
+        llama->is_load_interrupted = true;
+    }
+}
+
 JNIEXPORT jobject JNICALL
 Java_com_rnllama_LlamaContext_loadModelDetails(
     JNIEnv *env,
@@ -397,13 +536,18 @@ Java_com_rnllama_LlamaContext_doCompletion(
     jint top_k,
     jfloat top_p,
     jfloat min_p,
-    jfloat xtc_t,
-    jfloat xtc_p,
+    jfloat xtc_threshold,
+    jfloat xtc_probability,
     jfloat typical_p,
     jint seed,
     jobjectArray stop,
     jboolean ignore_eos,
     jobjectArray logit_bias,
+    jfloat   dry_multiplier,
+    jfloat   dry_base,
+    jint dry_allowed_length,    
+    jint dry_penalty_last_n,
+    jobjectArray dry_sequence_breakers,
     jobject partial_completion_callback
 ) {
     UNUSED(thiz);
@@ -414,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
@@ -422,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;
@@ -440,14 +584,34 @@ Java_com_rnllama_LlamaContext_doCompletion(
     sparams.typ_p = typical_p;
     sparams.n_probs = n_probs;
     sparams.grammar = env->GetStringUTFChars(grammar, nullptr);
-    sparams.xtc_threshold = xtc_t;
-    sparams.xtc_probability = xtc_p;
+    sparams.xtc_threshold = xtc_threshold;
+    sparams.xtc_probability = xtc_probability;
+    sparams.dry_multiplier = dry_multiplier;
+    sparams.dry_base = dry_base;
+    sparams.dry_allowed_length = dry_allowed_length;
+    sparams.dry_penalty_last_n = dry_penalty_last_n;
 
     sparams.logit_bias.clear();
     if (ignore_eos) {
         sparams.logit_bias[llama_token_eos(llama->model)].bias = -INFINITY;
     }
 
+    // dry break seq
+
+    jint size = env->GetArrayLength(dry_sequence_breakers);
+    std::vector<std::string> dry_sequence_breakers_vector;
+
+    for (jint i = 0; i < size; i++) {
+        jstring javaString = (jstring)env->GetObjectArrayElement(dry_sequence_breakers, i);
+        const char *nativeString = env->GetStringUTFChars(javaString, 0);
+        dry_sequence_breakers_vector.push_back(std::string(nativeString));
+        env->ReleaseStringUTFChars(javaString, nativeString);
+        env->DeleteLocalRef(javaString);
+    }
+
+    sparams.dry_sequence_breakers = dry_sequence_breakers_vector;
+
+    // logit bias
     const int n_vocab = llama_n_vocab(llama_get_model(llama->ctx));
     jsize logit_bias_len = env->GetArrayLength(logit_bias);
 
@@ -529,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());
@@ -651,16 +815,27 @@ Java_com_rnllama_LlamaContext_isEmbeddingEnabled(
 
 JNIEXPORT jobject JNICALL
 Java_com_rnllama_LlamaContext_embedding(
-        JNIEnv *env, jobject thiz, jlong context_ptr, jstring text) {
+        JNIEnv *env, jobject thiz,
+        jlong context_ptr,
+        jstring text,
+        jint embd_normalize
+) {
     UNUSED(thiz);
     auto llama = context_map[(long) context_ptr];
 
+    common_params embdParams;
+    embdParams.embedding = true;
+    embdParams.embd_normalize = llama->params.embd_normalize;
+    if (embd_normalize != -1) {
+      embdParams.embd_normalize = embd_normalize;
+    }
+
     const char *text_chars = env->GetStringUTFChars(text, nullptr);
 
     llama->rewind();
 
     llama_perf_context_reset(llama->ctx);
-    
+
     llama->params.prompt = text_chars;
 
     llama->params.n_predict = 0;
@@ -675,7 +850,7 @@ Java_com_rnllama_LlamaContext_embedding(
     llama->loadPrompt();
     llama->doCompletion();
 
-    std::vector<float> embedding = llama->getEmbedding();
+    std::vector<float> embedding = llama->getEmbedding(embdParams);
 
     auto embeddings = createWritableArray(env);
     for (const auto &val : embedding) {
@@ -683,6 +858,12 @@ Java_com_rnllama_LlamaContext_embedding(
     }
     putArray(env, result, "embedding", embeddings);
 
+    auto promptTokens = createWritableArray(env);
+    for (const auto &tok : llama->embd) {
+      pushString(env, promptTokens, common_token_to_piece(llama->ctx, tok).c_str());
+    }
+    putArray(env, result, "prompt_tokens", promptTokens);
+
     env->ReleaseStringUTFChars(text, text_chars);
     return result;
 }
@@ -722,4 +903,11 @@ Java_com_rnllama_LlamaContext_freeContext(
     context_map.erase((long) llama->ctx);
 }
 
+JNIEXPORT void JNICALL
+Java_com_rnllama_LlamaContext_logToAndroid(JNIEnv *env, jobject thiz) {
+    UNUSED(env);
+    UNUSED(thiz);
+    llama_log_set(log_callback, NULL);
+}
+
 } // extern "C"

package/android/src/newarch/java/com/rnllama/RNLlamaModule.java

@@ -39,8 +39,13 @@ public class RNLlamaModule extends NativeRNLlamaSpec {
   }
 
   @ReactMethod
-  public void initContext(final ReadableMap params, final Promise promise) {
-    rnllama.initContext(params, promise);
+  public void modelInfo(final String model, final ReadableArray skip, final Promise promise) {
+    rnllama.modelInfo(model, skip, promise);
+  }
+
+  @ReactMethod
+  public void initContext(double id, final ReadableMap params, final Promise promise) {
+    rnllama.initContext(id, params, promise);
   }
 
   @ReactMethod
@@ -89,8 +94,8 @@ public class RNLlamaModule extends NativeRNLlamaSpec {
   }
 
   @ReactMethod
-  public void embedding(double id, final String text, final Promise promise) {
-    rnllama.embedding(id, text, promise);
+  public void embedding(double id, final String text, final ReadableMap params, final Promise promise) {
+    rnllama.embedding(id, text, params, promise);
   }
 
   @ReactMethod

package/android/src/oldarch/java/com/rnllama/RNLlamaModule.java

@@ -40,8 +40,13 @@ public class RNLlamaModule extends ReactContextBaseJavaModule {
   }
 
   @ReactMethod
-  public void initContext(final ReadableMap params, final Promise promise) {
-    rnllama.initContext(params, promise);
+  public void modelInfo(final String model, final ReadableArray skip, final Promise promise) {
+    rnllama.modelInfo(model, skip, promise);
+  }
+
+  @ReactMethod
+  public void initContext(double id, final ReadableMap params, final Promise promise) {
+    rnllama.initContext(id, params, promise);
   }
 
   @ReactMethod
@@ -90,8 +95,8 @@ public class RNLlamaModule extends ReactContextBaseJavaModule {
   }
 
   @ReactMethod
-  public void embedding(double id, final String text, final Promise promise) {
-    rnllama.embedding(id, text, promise);
+  public void embedding(double id, final String text, final ReadableMap params, final Promise promise) {
+    rnllama.embedding(id, text, params, promise);
   }
 
   @ReactMethod

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__)