ADFA-5188 | Enable KV cache quantization and flash attention

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

@@ -26,6 +26,7 @@ android {
 				arguments += "-DLLAMA_BUILD_COMMON=ON"
 				arguments += "-DGGML_LLAMAFILE=OFF"
 				arguments += "-DCMAKE_BUILD_TYPE=Release"
+				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,27 @@ 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.
+ *
+ * @param requested the context asked for, in tokens
+ * @param trained_ctx what the model was trained for, or 0 when it does not say
+ * @return the context to configure, never above trained_ctx when the model declares one
+ */
+static int clamp_context(int requested, int trained_ctx) {
+    int clamped = requested > 0 ? requested : DEFAULT_N_CTX;
+    if (trained_ctx > 0 && clamped > trained_ctx) {
+        LOGi("context: n_ctx %d exceeds the model's trained %d; clamping", clamped, trained_ctx);
+        clamped = trained_ctx;
+    }
+    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,20 +407,48 @@ 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 what the model was trained for.
     const int trained_ctx = llama_model_n_ctx_train(model);
-    if (trained_ctx > 0 && requested_ctx > trained_ctx) {
-        LOGi("context: requested n_ctx %d exceeds the model's trained %d; clamping", requested_ctx, trained_ctx);
-        requested_ctx = trained_ctx;
+    const int requested_ctx = clamp_context(jn_ctx, trained_ctx);
+    // Sized by Kotlin against f16, the type the fallback below drops to; the two sizes differ
+    // because f16 costs nearly twice as much per cached token.
+    const int fallback_ctx = clamp_context(jfallback_n_ctx, trained_ctx);
+    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;
 
+    LOGi("Creating context: n_ctx = %d (model trained for %d), kv cache = %s", requested_ctx,
+         trained_ctx, quantize_kv ? "q8_0" : "f16");
+
     llama_context *context = llama_init_from_model(model, ctx_params);
+    bool quantized_in_use = quantize_kv;
+
+    if (!context) {
+        // The safety fallback: f16 with flash attention off is the one configuration nothing here
+        // can refuse — no block-size constraint on the cache, and no graph for AUTO to fail to
+        // place. It costs the attention speed-up on a model whose only problem was the cache type,
+        // which is the cheaper mistake to make. Kotlin already screens the head width, so getting
+        // here at all means the header and llama.cpp disagreed.
+        LOGe("Context creation failed; retrying at f16 with flash attention off and n_ctx %d",
+             fallback_ctx);
+        ctx_params.type_k = GGML_TYPE_F16;
+        ctx_params.type_v = GGML_TYPE_F16;
+        ctx_params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
+        ctx_params.n_ctx = fallback_ctx;
+        context = llama_init_from_model(model, ctx_params);
+        quantized_in_use = false;
+    }
 
     if (!context) {
         LOGe("context: llama_new_context_with_model() returned null");
@@ -411,9 +457,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), (int) jn_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,21 +246,31 @@ 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, and so that the size and the type
+     * cannot disagree: the context is created on the run loop, well after the caller picked them.
      *
      * @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; defaults to [nCtx], correct when [quantizeKv] is false
      */
-    suspend fun load(pathToModel: String, nCtx: Int) {
+    suspend fun load(
+        pathToModel: String,
+        nCtx: Int,
+        quantizeKv: Boolean = false,
+        fallbackNCtx: Int = nCtx,
+    ) {
         withContext(runLoop()) {
             when (threadLocalState.get()) {
                 is State.Idle -> {
                     val model = load_model(pathToModel)
                     if (model == 0L) throw IllegalStateException("load_model() failed")
 
-                    val context = new_context(model, nCtx)
+                    val context = new_context(model, nCtx, quantizeKv, fallbackNCtx)
                     if (context == 0L) throw IllegalStateException("new_context() failed")
 
                     val batch = new_batch(2048, 0, 1)

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

@@ -12,7 +12,9 @@ import com.itsaky.androidide.plugins.aiagentlocal.feedback.ModelNotConfiguredExc
 import com.itsaky.androidide.plugins.aiagentlocal.feedback.UserActionableLlmException
 import com.itsaky.androidide.plugins.aiagentlocal.feedback.UserFeedback
 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
@@ -314,11 +316,16 @@ class LocalLlmBackend(
             throw ModelLoadException(loadMessages.describe(shortfall), shortfall)
         }
 
-        val contextTokens = resolveContextSize(resolvedPath, availableBytes)
+        val contextSize = resolveContextSize(resolvedPath, availableBytes)
 
         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) {
@@ -334,27 +341,33 @@ 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. [ModelContextResolver] 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
-     * @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 suspend fun resolveContextSize(resolvedPath: String, availableBytes: Long): Int {
+    private suspend fun resolveContextSize(
+        resolvedPath: String,
+        availableBytes: Long,
+    ): ModelContextSize {
         val resolved = withContext(Dispatchers.IO) {
             ModelContextResolver.resolve(availableBytes.takeIf { it >= 0L }) {
                 File(resolvedPath).takeIf { it.isFile }?.inputStream()
             }
         }
         // Unconditional: a wrongly sized context otherwise just reads as the assistant forgetting.
         context.logger.info(
-            "Context size for $resolvedPath: ${resolved.contextTokens} tokens" +
+            "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 resolved.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,20 +34,37 @@ 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
+
     /**
      * @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
+     * @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?): Int {
+    fun choose(
+        header: GgufHeader?,
+        availableBytes: 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 || availableBytes == null) return DEFAULT_CONTEXT_TOKENS
         val modelTokens = header.contextLength?.takeIf { it > 0L } ?: return DEFAULT_CONTEXT_TOKENS
         // 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
 
         // Compute buffers come off the top; goes negative on a short device, which the floor absorbs.

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

@@ -0,0 +1,41 @@
+package com.itsaky.androidide.plugins.aiagentlocal.model
+
+/**
+ * How llama.cpp stores one element of the KV cache. The cache is the largest allocation a load
+ * makes and the one this plugin sizes, so what it costs per element lives only here — the native
+ * side is told a size per type rather than deriving one. Pure arithmetic, so it is unit-testable.
+ *
+ * @property bytesPerBlock what one block of [blockSize] elements occupies once stored
+ * @property blockSize elements per stored block; 1 for a type that is not quantized
+ */
+enum class KvCacheType(private val bytesPerBlock: Long, private val blockSize: Long) {
+
+    /** Two bytes per element, and llama.cpp's own default. Works for every model. */
+    F16(bytesPerBlock = 2L, blockSize = 1L),
+
+    /**
+     * 32 elements in 34 bytes — 32 quantized bytes plus one f16 scale — so a shade over half of
+     * [F16] for the same context. Usable only where [supports] holds, and only with flash
+     * attention, which llama.cpp requires for a quantized value cache. See ADFA-5188.
+     */
+    Q8_0(bytesPerBlock = 34L, blockSize = 32L);
+
+    /**
+     * @param elements cached elements, at most 2^43 for the shapes [ModelMemoryEstimator] admits
+     * @return what they occupy, exact whenever [elements] is a whole number of blocks
+     */
+    fun bytesFor(elements: Long): Long = elements * bytesPerBlock / blockSize
+
+    /**
+     * Whether a model's cached rows divide into whole blocks. llama.cpp refuses a quantized cache
+     * whose head width does not, so asking anyway costs a failed context creation and a retry.
+     *
+     * @param header the model's metadata, or null when it could not be read
+     * @return true when this type can hold that model's cache
+     */
+    fun supports(header: GgufHeader?): Boolean {
+        if (blockSize == 1L) return true
+        val widths = header?.let { ModelMemoryEstimator.headWidths(it) } ?: return false
+        return widths.first % blockSize == 0L && widths.second % blockSize == 0L
+    }
+}

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

@@ -3,13 +3,20 @@ package com.itsaky.androidide.plugins.aiagentlocal.model
 import java.io.InputStream
 
 /**
- * The context size one model load should get, and the header it was decided from.
+ * What one model load should get — context size and KV cache type — and the header behind them.
  *
  * @property contextTokens the context to load with; always a value [ContextSizePolicy] returned
+ * @property kvType the type the KV cache will be stored as, which is what that context was sized
+ *   against; the two are decided together or they describe different allocations
+ * @property fallbackContextTokens the context the same RAM buys under [KvCacheType.F16], for the
+ *   native fallback when llama.cpp refuses a quantized cache; equals [contextTokens] when [kvType]
+ *   is already [KvCacheType.F16]
  * @property header the model's parsed metadata, or null when it could not be read
  */
 internal data class ModelContextSize(
     val contextTokens: Int,
+    val kvType: KvCacheType,
+    val fallbackContextTokens: Int,
     val header: GgufHeader?,
 ) {
 
@@ -18,9 +25,10 @@ internal data class ModelContextSize(
 }
 
 /**
- * Decides how large a context a given model gets on this device: reads the model's GGUF header and
- * hands it to [ContextSizePolicy]. Owning both steps lets the load path and the pre-load memory
- * warning derive their numbers the same way, and stays Android-free to test. See ADFA-5187.
+ * Decides what one model load gets on this device — context size and KV cache type: reads the
+ * model's GGUF header and hands it to [ContextSizePolicy]. Owning every step lets the load path and
+ * the pre-load memory warning derive their numbers the same way, and stays Android-free to test.
+ * See ADFA-5187 and ADFA-5188.
  */
 internal object ModelContextResolver {
 
@@ -31,12 +39,18 @@ internal object ModelContextResolver {
      *
      * @param availableBytes free RAM in bytes, or null when it could not be read
      * @param openStream opens the model file, or returns null when there is nothing to open
-     * @return the context to load with, and the header behind it
+     * @return the context and cache type to load with, and the header behind them
      */
     fun resolve(availableBytes: Long?, openStream: () -> InputStream?): ModelContextSize {
         val header = GgufHeaderReader.read(openStream)
+        // A quantized cache buys about twice the context, so the type is picked before the size.
+        val kvType = ContextSizePolicy.chooseKvCache(header)
         return ModelContextSize(
-            contextTokens = ContextSizePolicy.choose(header, availableBytes),
+            contextTokens = ContextSizePolicy.choose(header, availableBytes, kvType),
+            kvType = kvType,
+            // Sized here rather than scaled natively, so the fallback obeys the one policy that
+            // knows the floor, the ceiling and the rounding.
+            fallbackContextTokens = ContextSizePolicy.choose(header, availableBytes, KvCacheType.F16),
             header = header,
         )
     }