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,30 @@ 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
+ * @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) {
+ 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: n_ctx %d exceeds the model's trained %d; clamping to %d", 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 +410,68 @@ 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);
+ // 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;
+
+ // 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) {
+ // 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;
+ }
+
+ // The only retry that shrinks the allocation, back to the context every load got before this was
+ // sized per device; fallback_ctx cannot, since f16 costs what q8_0 bought the extra tokens with.
+ // The guard skips an attempt that would re-request exactly what just failed.
+ // 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 ||
+ ctx_params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED)) {
+ LOGe("Context creation failed; retrying at the n_ctx %d floor with f16 and flash attention off",
+ floor_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 = 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 +480,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,28 +246,51 @@ 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)
- if (context == 0L) throw IllegalStateException("new_context() failed")
-
- val batch = new_batch(2048, 0, 1)
- if (batch == 0L) throw IllegalStateException("new_batch() failed")
-
- val sampler = new_sampler()
- if (sampler == 0L) throw IllegalStateException("new_sampler() failed")
+ // Only State.Loaded holds these, so a later step failing leaks them unless freed here.
+ var context = 0L
+ var batch = 0L
+ var sampler = 0L
+ try {
+ context = new_context(model, nCtx, quantizeKv, fallbackNCtx)
+ if (context == 0L) throw IllegalStateException("new_context() failed")
+
+ batch = new_batch(2048, 0, 1)
+ if (batch == 0L) throw IllegalStateException("new_batch() failed")
+
+ sampler = new_sampler()
+ if (sampler == 0L) throw IllegalStateException("new_sampler() failed")
+ } catch (e: Throwable) {
+ // Model last: the context borrows from it, so it has to outlive the context.
+ if (sampler != 0L) free_sampler(sampler)
+ if (batch != 0L) free_batch(batch)
+ if (context != 0L) free_context(context)
+ free_model(model)
+ throw e
+ }
 
  log.info("Loaded model {}", pathToModel)
  threadLocalState.set(State.Loaded(model, context, batch, sampler))

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

@@ -14,7 +14,9 @@ import com.itsaky.androidide.plugins.aiagentlocal.feedback.UserFeedback
 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
@@ -324,11 +326,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) {
@@ -341,13 +348,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
  */
@@ -371,34 +378,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. [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
  * @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 {
+ ): 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: ${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,15 +34,34 @@ 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; 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.
@@ -52,7 +71,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
 
  // Weights first, then the compute buffers, each clamped at zero rather than left to run