ADFA-5188 | Enable KV cache quantization and flash attention

ai-agent-local/llama-impl/build.gradle.kts

@@ -25,6 +25,9 @@ android {
 				arguments += "-DLLAMA_BUILD_COMMON=ON"
 				arguments += "-DGGML_LLAMAFILE=OFF"
 				arguments += "-DCMAKE_BUILD_TYPE=Release"
+				// 16 KB page alignment, required on newer arm64 devices; PrebuiltAarAbiTest asserts it
+				// on the committed AAR, since dropping this still builds and still loads on 4 KB.
+				arguments += "-DANDROID_SUPPORT_FLEXIBLE_PAGE_SIZES=ON"
 				cppFlags += listOf()
 				arguments += listOf()
 

ai-agent-local/llama-impl/src/main/cpp/llama-android.cpp

@@ -365,9 +365,31 @@ Java_android_llama_cpp_LLamaAndroid_free_1model(JNIEnv *, jobject, jlong model)
     llama_model_free(reinterpret_cast<llama_model *>(model));
 }
 
+/**
+ * Backstops a context size Kotlin chose: a misparsed header must not ask for more than the model
+ * was trained for, and a non-positive argument falls back to the default. Never clamps below
+ * DEFAULT_N_CTX, the context a 2048-trained model always got, so no prompt that fit regresses.
+ *
+ * @param requested the context asked for, in tokens
+ * @param trained_ctx what the model was trained for, or 0 when it does not say
+ * @param which which size is being clamped, so two clamps in one load read as two in the log
+ * @return the context to configure, never above trained_ctx unless that is below DEFAULT_N_CTX
+ */
+static int clamp_context(int requested, int trained_ctx, const char *which) {
+    int clamped = requested > 0 ? requested : DEFAULT_N_CTX;
+    const int ceiling = std::max(trained_ctx, DEFAULT_N_CTX);
+    if (trained_ctx > 0 && clamped > ceiling) {
+        LOGi("context: %s n_ctx %d exceeds the model's trained %d; clamping to %d", which, clamped,
+             trained_ctx, ceiling);
+        clamped = ceiling;
+    }
+    return clamped;
+}
+
 extern "C"
 JNIEXPORT jlong JNICALL
-Java_android_llama_cpp_LLamaAndroid_new_1context(JNIEnv *env, jobject, jlong jmodel, jint jn_ctx) {
+Java_android_llama_cpp_LLamaAndroid_new_1context(JNIEnv *env, jobject, jlong jmodel, jint jn_ctx,
+                                                 jboolean jquantize_kv, jint jfallback_n_ctx) {
     auto model = reinterpret_cast<llama_model *>(jmodel);
 
     if (!model) {
@@ -389,23 +411,62 @@ Java_android_llama_cpp_LLamaAndroid_new_1context(JNIEnv *env, jobject, jlong jmo
 
     llama_context_params ctx_params = llama_context_default_params();
 
-    int requested_ctx = jn_ctx > 0 ? jn_ctx : DEFAULT_N_CTX;
-
-    // Backstop on Kotlin's number: a misparsed header must not exceed the trained context. Floored
-    // at DEFAULT_N_CTX, the context a 2048-trained model always got, so no prompt that fit regresses.
     const int trained_ctx = llama_model_n_ctx_train(model);
-    const int clamp_ctx = std::max(trained_ctx, DEFAULT_N_CTX);
-    if (trained_ctx > 0 && requested_ctx > clamp_ctx) {
-        LOGi("context: requested n_ctx %d exceeds the model's trained %d; clamping to %d",
-             requested_ctx, trained_ctx, clamp_ctx);
-        requested_ctx = clamp_ctx;
+    const int requested_ctx = clamp_context(jn_ctx, trained_ctx, "requested");
+    const bool quantize_kv = jquantize_kv == JNI_TRUE;
+
+    // AUTO rather than ENABLED: it is AUTO that makes llama.cpp validate a quantized cache against
+    // the model's head width and refuse it by returning null. ENABLED skips that check and aborts
+    // inside ggml instead, taking the IDE down with it.
+    ctx_params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO;
+    if (quantize_kv) {
+        // A quantized V cache is only defined with flash attention, which AUTO may still refuse.
+        ctx_params.type_k = GGML_TYPE_Q8_0;
+        ctx_params.type_v = GGML_TYPE_Q8_0;
     }
 
     ctx_params.n_ctx = requested_ctx;
     ctx_params.n_threads = n_threads;
     ctx_params.n_threads_batch = n_threads_batch;
 
+    // The request, not the outcome: whether AUTO resolved to on is llama.cpp's own line to log,
+    // under the AiAgentLocal.llama.cpp tag, since no public getter exposes cparams.flash_attn.
+    LOGi("Creating context: n_ctx = %d (model trained for %d), kv cache = %s, flash attention = %s",
+         requested_ctx, trained_ctx, quantize_kv ? "q8_0" : "f16",
+         llama_flash_attn_type_name(ctx_params.flash_attn_type));
+
     llama_context *context = llama_init_from_model(model, ctx_params);
+    bool quantized_in_use = quantize_kv;
+
+    // Two unrelated failures land here and want opposite retries: a refused quantized cache is not
+    // a shortage and keeps its long context, while a shortage is answered only by fewer bytes.
+    if (!context && quantize_kv) {
+        // Only the cache type moves: AUTO was the default before this PR and downgrades itself on
+        // an f16 cache, so disabling it would make a refused model slower than it is on main.
+        // Sized by Kotlin against f16, and clamped here so a load that never retries logs no clamp.
+        const int fallback_ctx = clamp_context(jfallback_n_ctx, trained_ctx, "fallback");
+        LOGe("Context creation failed; retrying at f16 with n_ctx %d", fallback_ctx);
+        ctx_params.type_k = GGML_TYPE_F16;
+        ctx_params.type_v = GGML_TYPE_F16;
+        ctx_params.n_ctx = fallback_ctx;
+        context = llama_init_from_model(model, ctx_params);
+        quantized_in_use = false;
+    }
+
+    // The only retry that shrinks the allocation, back to the context every load got before this
+    // was sized per device; the f16 retry above cannot, costing per token what it saved in tokens.
+    // Guarded on size alone: at the same n_ctx nothing left to change makes the attempt cheaper.
+    // n_ctx is unsigned; every value compared here is a clamped positive.
+    const int current_ctx = (int) ctx_params.n_ctx;
+    const int floor_ctx = std::min(current_ctx, DEFAULT_N_CTX);
+    if (!context && floor_ctx < current_ctx) {
+        LOGe("Context creation failed; retrying at the n_ctx %d floor with f16", floor_ctx);
+        ctx_params.type_k = GGML_TYPE_F16;
+        ctx_params.type_v = GGML_TYPE_F16;
+        ctx_params.n_ctx = floor_ctx;
+        context = llama_init_from_model(model, ctx_params);
+        quantized_in_use = false;
+    }
 
     if (!context) {
         LOGe("context: llama_new_context_with_model() returned null");
@@ -414,9 +475,11 @@ Java_android_llama_cpp_LLamaAndroid_new_1context(JNIEnv *env, jobject, jlong jmo
         return 0;
     }
 
-    // n_ctx now varies per model and device, so a wrong size is invisible in a report without this.
-    LOGi("context: created with n_ctx = %u (requested %d, model trained for %d), n_batch = %u",
-         llama_n_ctx(context), requested_ctx, trained_ctx, llama_n_batch(context));
+    // n_ctx and the cache type now vary per model and device, so a wrong one is invisible in a
+    // report without this.
+    LOGi("Context created: n_ctx = %u (requested %d, model trained for %d), n_batch = %u, kv cache = %s",
+         llama_n_ctx(context), requested_ctx, trained_ctx, llama_n_batch(context),
+         quantized_in_use ? "q8_0" : "f16");
 
     // A fresh context has an empty KV cache, so the prefix record must start empty too.
     {

ai-agent-local/llama-impl/src/main/java/android/llama/cpp/LLamaAndroid.kt

@@ -165,7 +165,12 @@ class LLamaAndroid : ILlamaController {
     private external fun log_to_android()
     private external fun load_model(filename: String): Long
     private external fun free_model(model: Long)
-    private external fun new_context(model: Long, nCtx: Int): Long
+    private external fun new_context(
+        model: Long,
+        nCtx: Int,
+        quantizeKv: Boolean,
+        fallbackNCtx: Int,
+    ): Long
     private external fun free_context(context: Long)
     private external fun backend_init(numa: Boolean)
     private external fun backend_free()
@@ -241,18 +246,30 @@ class LLamaAndroid : ILlamaController {
     override suspend fun load(pathToModel: String) = load(pathToModel, DEFAULT_N_CTX)
 
     /**
-     * Loads a model and gives its context [nCtx] tokens. The size is an argument rather than
-     * process-global state so that it cannot be overwritten between being chosen and being used:
-     * the context is created on the run loop, well after the caller picked the number.
+     * Loads a model and gives its context [nCtx] tokens, stored as q8_0 when [quantizeKv] asks for
+     * it. Every part of the shape is an argument rather than process-global state so that none of it
+     * can be overwritten between being chosen and being used: the context is created on the run
+     * loop, well after the caller picked them.
      *
      * A partial load frees what it allocated before rethrowing: [threadLocalState] stays `Idle`, so
      * nothing else can reach those handles, and a retry would otherwise mmap another model on top
      * of the leaked one for the process lifetime.
      *
      * @param pathToModel filesystem path to the `.gguf` model
-     * @param nCtx context size in tokens; anything non-positive means [DEFAULT_N_CTX]
+     * @param nCtx context size in tokens, sized for [quantizeKv]; non-positive means [DEFAULT_N_CTX]
+     * @param quantizeKv true to store the KV cache as q8_0, roughly half the bytes of f16; the
+     *   native side may still refuse it, in which case the load falls back to f16 at [fallbackNCtx]
+     * @param fallbackNCtx context size for that f16 fallback, which the caller sizes against f16's
+     *   own per-token cost; non-positive means [DEFAULT_N_CTX], since [nCtx] would be the wrong
+     *   default under [quantizeKv] — an f16 cache at a size q8_0 paid for asks for nearly twice the
+     *   bytes the attempt that just failed did
      */
-    suspend fun load(pathToModel: String, nCtx: Int) {
+    suspend fun load(
+        pathToModel: String,
+        nCtx: Int,
+        quantizeKv: Boolean = false,
+        fallbackNCtx: Int = 0,
+    ) {
         withContext(runLoop()) {
             when (threadLocalState.get()) {
                 is State.Idle -> {
@@ -263,7 +280,7 @@ class LLamaAndroid : ILlamaController {
                     var batch = 0L
                     var sampler = 0L
                     try {
-                        context = new_context(model, nCtx)
+                        context = new_context(model, nCtx, quantizeKv, fallbackNCtx)
                         if (context == 0L) throw IllegalStateException("new_context() failed")
 
                         batch = new_batch(2048, 0, 1)

ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/backend/LocalLlmBackend.kt

@@ -11,10 +11,12 @@ import com.itsaky.androidide.plugins.aiagentlocal.feedback.ModelLoadException
 import com.itsaky.androidide.plugins.aiagentlocal.feedback.ModelNotConfiguredException
 import com.itsaky.androidide.plugins.aiagentlocal.feedback.UserActionableLlmException
 import com.itsaky.androidide.plugins.aiagentlocal.feedback.UserFeedback
-import com.itsaky.androidide.plugins.aiagentlocal.model.ContextSizePolicy
 import com.itsaky.androidide.plugins.aiagentlocal.model.GgufHeader
 import com.itsaky.androidide.plugins.aiagentlocal.model.GgufHeaderReader
 import com.itsaky.androidide.plugins.aiagentlocal.model.GgufModelInspector
+import com.itsaky.androidide.plugins.aiagentlocal.model.KvCacheType
+import com.itsaky.androidide.plugins.aiagentlocal.model.ModelContextResolver
+import com.itsaky.androidide.plugins.aiagentlocal.model.ModelContextSize
 import com.itsaky.androidide.plugins.aiagentlocal.model.ModelLoadDiagnostics
 import com.itsaky.androidide.plugins.aiagentlocal.model.ModelLoadMessages
 import com.itsaky.androidide.plugins.aiagentlocal.preferences.LocalLlmPreferences
@@ -329,11 +331,16 @@ class LocalLlmBackend(
             throw ModelLoadException(loadMessages.describe(shortfall), shortfall)
         }
 
-        val contextTokens = resolveContextSize(resolvedPath, availableBytes, header, modelSizeBytes)
+        val contextSize = resolveContextSize(resolvedPath, availableBytes, header, modelSizeBytes)
 
         context.logger.info("Loading model: $resolvedPath")
         try {
-            llama.load(resolvedPath, contextTokens)
+            llama.load(
+                pathToModel = resolvedPath,
+                nCtx = contextSize.contextTokens,
+                quantizeKv = contextSize.kvType == KvCacheType.Q8_0,
+                fallbackNCtx = contextSize.fallbackContextTokens,
+            )
         } catch (e: CancellationException) {
             throw e
         } catch (e: Exception) {
@@ -346,13 +353,13 @@ class LocalLlmBackend(
         modelLoaded = true
         currentModelPath = resolvedPath
         context.logger.info("Model loaded successfully")
-        reportEffectiveContextSize(contextTokens)
+        reportEffectiveContextSize(contextSize.contextTokens)
     }
 
     /**
      * Logs the context the native side actually created. It can be smaller than what was asked for
-     * — `new_context` clamps a request above what the model was trained for — and without this the
-     * only visible number is the request, so a prompt rejected as too long looks like it fit.
+     * — clamped to the trained context, or dropped to the shorter f16 fallback when a quantized
+     * cache was refused — and a prompt rejected as too long otherwise looks like it fit.
      *
      * @param requestedTokens the context [resolveContextSize] asked for
      */
@@ -376,34 +383,37 @@ class LocalLlmBackend(
     }
 
     /**
-     * Sizes the KV cache for this model on this device. Must run after any unload, so the freed
-     * context is counted as available, and the answer is passed to [LLamaAndroid.load] rather than
-     * stored anywhere. [ContextSizePolicy.choose] fails open, so this has no failure of its own.
+     * Sizes the KV cache for this model on this device and picks the type it is stored as. Must run
+     * after any unload, so the freed context is counted as available. Answers rather than applies:
+     * every part of the shape is an argument to [LLamaAndroid.load], so nothing can drift between
+     * being chosen here and being used natively. [ModelContextResolver] fails open, so this has no
+     * failure of its own.
      *
      * @param resolvedPath filesystem path to the model, already resolved from any content URI
      * @param availableBytes free RAM as [availableMemoryBytes] reports it, negative if unknown
      * @param header the model's metadata as read once by [ensureModelLoaded], null if unreadable
      * @param modelSizeBytes the model file's size, null if unreadable
-     * @return the context size in tokens to load the model with
+     * @return the context size, cache type and f16 fallback size to load the model with
      */
     private fun resolveContextSize(
         resolvedPath: String,
         availableBytes: Long,
         header: GgufHeader?,
         modelSizeBytes: Long?,
-    ): Int {
-        val contextTokens = ContextSizePolicy.choose(
+    ): ModelContextSize {
+        val resolved = ModelContextResolver.resolve(
             header = header,
             availableBytes = availableBytes.takeIf { it >= 0L },
             modelSizeBytes = modelSizeBytes,
         )
         // Unconditional: a wrongly sized context otherwise just reads as the assistant forgetting.
         context.logger.info(
-            "Context size for $resolvedPath: $contextTokens tokens" +
-                " (model advertises ${header?.contextLength ?: "unknown"}," +
+            "Context size for $resolvedPath: ${resolved.contextTokens} tokens," +
+                " ${resolved.kvType} KV cache" +
+                " (model advertises ${resolved.advertisedTokens ?: "unknown"}," +
                 " ${if (availableBytes >= 0L) "$availableBytes bytes free" else "free RAM unknown"})"
         )
-        return contextTokens
+        return resolved
     }
 
     /**

ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicy.kt

@@ -1,9 +1,9 @@
 package com.itsaky.androidide.plugins.aiagentlocal.model
 
 /**
- * Picks the context size (`n_ctx`) one model load gets, from what the model advertises and what the
- * device can spare — the KV cache scales linearly with it and is the largest knob we control. Pure
- * and Android-free, so every boundary is unit-testable off-device. See ADFA-5187.
+ * Picks what one model load gets: the context size (`n_ctx`) and the type the KV cache is stored as,
+ * from what the model advertises and what the device can spare. Pure and Android-free, so every
+ * boundary is unit-testable off-device. See ADFA-5187 and ADFA-5188.
  */
 object ContextSizePolicy {
 
@@ -34,6 +34,17 @@ object ContextSizePolicy {
      */
     private const val KV_BUDGET_DIVISOR = 2L
 
+    /**
+     * The cache type a load should ask for. Quantized wherever the model allows it: it halves the
+     * bytes one cached token costs, which is what lets [choose] return a longer context on the same
+     * device. Falls back to [KvCacheType.F16] rather than risking a refused context. See ADFA-5188.
+     *
+     * @param header the model's GGUF metadata, or null when it could not be read
+     * @return the type to configure natively, and to size the context against
+     */
+    fun chooseKvCache(header: GgufHeader?): KvCacheType =
+        if (KvCacheType.Q8_0.supports(header)) KvCacheType.Q8_0 else KvCacheType.F16
+
     /**
      * The weights are charged against free RAM even though they are mmap'd: this reading is taken
      * before a load that then pages them in from the same pool. So a model whose file approaches
@@ -42,11 +53,19 @@ object ContextSizePolicy {
      * @param header the model's GGUF metadata, or null when it could not be read
      * @param availableBytes free RAM right now, or null when it could not be read; a negative
      *   reading is treated as unreadable too
-     * @param modelSizeBytes the model file's size, or null when it could not be read
+     * @param modelSizeBytes the model file's size, or null when it could not be read; the weights
+     *   are charged against free RAM before the cache gets a budget
+     * @param kvType the cache type this load will ask for, from [chooseKvCache]; the budget buys
+     *   about twice the context under [KvCacheType.Q8_0], so the two have to be decided together
      * @return the context to configure, always between [DEFAULT_CONTEXT_TOKENS] and
      *   [MAX_CONTEXT_TOKENS] inclusive
      */
-    fun choose(header: GgufHeader?, availableBytes: Long?, modelSizeBytes: Long?): Int {
+    fun choose(
+        header: GgufHeader?,
+        availableBytes: Long?,
+        modelSizeBytes: Long?,
+        kvType: KvCacheType = KvCacheType.F16,
+    ): Int {
         // Each null is a distinct "we don't know"; all of them mean the same fallback.
         if (header == null) return DEFAULT_CONTEXT_TOKENS
         // A negative reading is not free RAM this can reason about, so treat it as unreadable.
@@ -56,7 +75,7 @@ object ContextSizePolicy {
         // Nothing to weigh below the floor, and no reason to price a cache we would not shrink.
         if (modelTokens <= DEFAULT_CONTEXT_TOKENS) return DEFAULT_CONTEXT_TOKENS
 
-        val perToken = ModelMemoryEstimator.kvBytesPerToken(header)?.takeIf { it > 0L }
+        val perToken = ModelMemoryEstimator.kvBytesPerToken(header, kvType)?.takeIf { it > 0L }
             ?: return DEFAULT_CONTEXT_TOKENS
 
         // Each clamped at zero: an unclamped Long underflows on an absurd size and wraps positive.