diff --git a/ai-agent-local/libs/v8/llama-v8-release.aar b/ai-agent-local/libs/v8/llama-v8-release.aar index 9a4088a..0241922 100644 Binary files a/ai-agent-local/libs/v8/llama-v8-release.aar and b/ai-agent-local/libs/v8/llama-v8-release.aar differ diff --git a/ai-agent-local/llama-impl/build.gradle.kts b/ai-agent-local/llama-impl/build.gradle.kts index cc519ae..3332461 100644 --- a/ai-agent-local/llama-impl/build.gradle.kts +++ b/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() diff --git a/ai-agent-local/llama-impl/src/main/cpp/llama-android.cpp b/ai-agent-local/llama-impl/src/main/cpp/llama-android.cpp index 54697c0..e734740 100644 --- a/ai-agent-local/llama-impl/src/main/cpp/llama-android.cpp +++ b/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(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(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. { diff --git a/ai-agent-local/llama-impl/src/main/java/android/llama/cpp/LLamaAndroid.kt b/ai-agent-local/llama-impl/src/main/java/android/llama/cpp/LLamaAndroid.kt index 59b537b..fe86505 100644 --- a/ai-agent-local/llama-impl/src/main/java/android/llama/cpp/LLamaAndroid.kt +++ b/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,29 +246,36 @@ 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. - * - * 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. + * 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") + // 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) + context = new_context(model, nCtx, quantizeKv, fallbackNCtx) if (context == 0L) throw IllegalStateException("new_context() failed") batch = new_batch(2048, 0, 1) @@ -271,13 +283,13 @@ class LLamaAndroid : ILlamaController { sampler = new_sampler() if (sampler == 0L) throw IllegalStateException("new_sampler() failed") - } catch (failure: Throwable) { - // Reverse of the allocation order, and the model last: it owns the rest. + } 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 failure + throw e } log.info("Loaded model {}", pathToModel) diff --git a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/backend/LocalLlmBackend.kt b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/backend/LocalLlmBackend.kt index 0e00767..458e0cd 100644 --- a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/backend/LocalLlmBackend.kt +++ b/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 @@ -328,11 +330,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) { @@ -345,13 +352,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 */ @@ -375,22 +382,24 @@ 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 }, @@ -398,11 +407,12 @@ class LocalLlmBackend( ) // 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 } /** diff --git a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicy.kt b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicy.kt index 165f3c5..5aeaac7 100644 --- a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicy.kt +++ b/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. diff --git a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/KvCacheType.kt b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/KvCacheType.kt new file mode 100644 index 0000000..3784186 --- /dev/null +++ b/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 + } +} diff --git a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolver.kt b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolver.kt index b87704d..e183b3c 100644 --- a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolver.kt +++ b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolver.kt @@ -1,13 +1,20 @@ package com.itsaky.androidide.plugins.aiagentlocal.model /** - * 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?, ) { @@ -16,9 +23,9 @@ internal data class ModelContextSize( } /** - * Decides how large a context a given model gets on this device, from the header someone else - * already read. Pure, so the load path can take its free-RAM reading after an unload and still - * price the same header the embedding-model guard used. See ADFA-5187. + * Decides what one model load gets on this device — context size and KV cache type — from the header + * someone else already read. Pure, so the load path can take its free-RAM reading after an unload + * and still price the same header the embedding-model guard used. See ADFA-5187 and ADFA-5188. */ internal object ModelContextResolver { @@ -31,14 +38,24 @@ internal object ModelContextResolver { * @param availableBytes free RAM in bytes, or null when it could not be read * @param modelSizeBytes the model file's size in bytes, or null when it could not be read; the * weights are charged against free RAM before the KV cache gets a budget - * @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( header: GgufHeader?, availableBytes: Long?, modelSizeBytes: Long?, - ): ModelContextSize = ModelContextSize( - contextTokens = ContextSizePolicy.choose(header, availableBytes, modelSizeBytes), - header = header, - ) + ): ModelContextSize { + // 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, modelSizeBytes, 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, modelSizeBytes, KvCacheType.F16, + ), + header = header, + ) + } } diff --git a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelMemoryEstimator.kt b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelMemoryEstimator.kt index 0a3d31f..22c9c52 100644 --- a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelMemoryEstimator.kt +++ b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelMemoryEstimator.kt @@ -24,15 +24,12 @@ data class MemoryEstimate( /** * Estimates the memory a `.gguf` model needs, from its size and its declared shape. Pure and - * Android-free, so the arithmetic is unit-testable. The context it measures at is the caller's: the - * load path passes [ContextSizePolicy.choose]'s answer, the pre-flight warning the floor, since a - * figure derived from free RAM cannot then be judged against it (ADFA-5187). + * Android-free, so the arithmetic is unit-testable. The context and cache type it measures at are + * the caller's: the load path passes what [ModelContextResolver] resolved, the pre-flight warning + * the floor, since a figure derived from free RAM cannot then be judged against it (ADFA-5187/5188). */ object ModelMemoryEstimator { - /** Two bytes per cached element: f16, the default KV type. */ - private const val KV_BYTES_PER_ELEMENT = 2L - /** Graph and compute buffers every load allocates; see [ModelMemory.RUN_BUFFER_BYTES]. */ private const val COMPUTE_BUFFER_BYTES = ModelMemory.RUN_BUFFER_BYTES @@ -53,15 +50,18 @@ object ModelMemoryEstimator { * @param header the model's metadata, or null when it could not be read * @param contextTokens the context to price the cache at; required, because a default here * would silently describe an allocation nobody makes + * @param kvType the cache type the load will be given; required for the same reason, and from + * the same [ModelContextResolver] answer, since it halves what a cached token costs * @return the estimate, or null when there is nothing to base one on */ fun estimate( fileSizeBytes: Long?, header: GgufHeader?, contextTokens: Int, + kvType: KvCacheType, ): MemoryEstimate? { if (fileSizeBytes == null || fileSizeBytes <= 0L) return null - val kvCacheBytes = header?.let { kvCacheBytes(it, contextTokens) } + val kvCacheBytes = header?.let { kvCacheBytes(it, contextTokens, kvType) } return if (kvCacheBytes != null) { MemoryEstimate(fileSizeBytes, kvCacheBytes + COMPUTE_BUFFER_BYTES, fromHeader = true) } else { @@ -78,9 +78,9 @@ object ModelMemoryEstimator { * KV cache size for a full context of [contextTokens]. Null unless every value it needs is * present and within its ceiling, or the context is not positive. */ - private fun kvCacheBytes(header: GgufHeader, contextTokens: Int): Long? { + private fun kvCacheBytes(header: GgufHeader, contextTokens: Int, kvType: KvCacheType): Long? { if (contextTokens <= 0) return null - val perToken = kvBytesPerToken(header) ?: return null + val perToken = kvBytesPerToken(header, kvType) ?: return null return perToken * contextTokens } @@ -90,16 +90,29 @@ object ModelMemoryEstimator { * Stays under 2^44 within the ceilings below, so any context the policy returns fits a Long. * * @param header the model's metadata + * @param kvType the type the cache will be stored as * @return bytes of KV cache per token, or null if the header does not say enough */ - internal fun kvBytesPerToken(header: GgufHeader): Long? { + internal fun kvBytesPerToken(header: GgufHeader, kvType: KvCacheType = KvCacheType.F16): Long? { val layers = header.blockCount?.within(MAX_LAYERS) ?: return null val heads = header.headCount?.within(MAX_HEADS) ?: return null // Grouped-query attention caches only the kv heads; absent means one per head (plain MHA). val kvHeads = (header.headCountKv ?: heads).within(MAX_HEADS) ?: return null + val (keyWidth, valueWidth) = headWidths(header) ?: return null + return kvType.bytesFor(layers * kvHeads * (keyWidth + valueWidth)) + } + + /** + * The per-head key and value widths, each either declared or derived. Also what decides whether + * a quantized cache is possible at all, since it needs both to divide into whole blocks. + * + * @param header the model's metadata + * @return key width to value width, or null if the header does not say enough + */ + internal fun headWidths(header: GgufHeader): Pair? { val keyWidth = header.keyLength?.within(MAX_WIDTH) ?: defaultHeadWidth(header) ?: return null val valueWidth = header.valueLength?.within(MAX_WIDTH) ?: defaultHeadWidth(header) ?: return null - return KV_BYTES_PER_ELEMENT * layers * kvHeads * (keyWidth + valueWidth) + return keyWidth to valueWidth } /** The value when it is positive and no larger than [ceiling]; null when it is neither. */ diff --git a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/settings/LocalLlmSettingsViewModel.kt b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/settings/LocalLlmSettingsViewModel.kt index c91ed04..504b6c4 100644 --- a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/settings/LocalLlmSettingsViewModel.kt +++ b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/settings/LocalLlmSettingsViewModel.kt @@ -286,6 +286,8 @@ class LocalLlmSettingsViewModel( fileSizeBytes = fileInfo.sizeBytes, header = header, contextTokens = ContextSizePolicy.DEFAULT_CONTEXT_TOKENS, + // The type does not depend on free RAM, so the load will pick this same one. + kvType = ContextSizePolicy.chooseKvCache(header), ) // Read last and never cached: the header parse above is blocking I/O over the model file, // and the user may have just closed apps to make room. diff --git a/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicyTest.kt b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicyTest.kt index 4616357..e6bff5c 100644 --- a/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicyTest.kt +++ b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicyTest.kt @@ -41,12 +41,13 @@ class ContextSizePolicyTest { private fun ramAffording(tokens: Long, weightBytes: Long = modelSize): Long = tokens * bytesPerToken * 2L + weightBytes + ModelMemory.RUN_BUFFER_BYTES - /** [ContextSizePolicy.choose] with the model size defaulted, which most cases do not vary. */ + /** [ContextSizePolicy.choose] with the size and cache type defaulted, which most cases fix. */ private fun choose( header: GgufHeader?, availableBytes: Long?, modelSizeBytes: Long? = modelSize, - ): Int = ContextSizePolicy.choose(header, availableBytes, modelSizeBytes) + kvType: KvCacheType = KvCacheType.F16, + ): Int = ContextSizePolicy.choose(header, availableBytes, modelSizeBytes, kvType) @Test fun givenNoHeader_whenChoosing_thenFallsBackToDefault() { @@ -136,6 +137,41 @@ class ContextSizePolicyTest { assertEquals(DEFAULT_CONTEXT_TOKENS, result) } + @Test + fun givenAQuantizableModel_whenChoosingTheCacheType_thenPicksQ8_0() { + assertEquals(KvCacheType.Q8_0, ContextSizePolicy.chooseKvCache(header())) + } + + @Test + fun givenAHeadWidthQ8_0CannotHold_whenChoosingTheCacheType_thenFallsBackToF16() { + assertEquals(KvCacheType.F16, ContextSizePolicy.chooseKvCache(header(keyLength = 80L))) + } + + @Test + fun givenNoHeader_whenChoosingTheCacheType_thenFallsBackToF16() { + assertEquals(KvCacheType.F16, ContextSizePolicy.chooseKvCache(null)) + } + + @Test + fun givenTheSameRam_whenChoosingUnderQ8_0_thenAffordsNearlyTwiceTheContext() { + // The RAM that buys 5_000 f16 tokens buys 9_411 q8_0 ones, rounded down to whole blocks. + val ram = ramAffording(5_000L) + assertEquals(4864, choose(header(), ram, kvType = KvCacheType.F16)) + assertEquals(9216, choose(header(), ram, kvType = KvCacheType.Q8_0)) + } + + @Test + fun givenAModelContextBelowWhatQ8_0Affords_whenChoosing_thenTheModelStillCaps() { + val result = choose(header(contextLength = 8192L), ramAffording(5_000L), kvType = KvCacheType.Q8_0) + assertEquals(8192, result) + } + + @Test + fun givenTightRamUnderQ8_0_whenChoosing_thenNeverGoesBelowTheFloor() { + val result = choose(header(), ramAffording(1_000L), kvType = KvCacheType.Q8_0) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + @Test fun givenUnreadableModelSize_whenChoosing_thenFallsBackToDefault() { val result = choose(header(), ramAffording(100_000L), modelSizeBytes = null) @@ -193,12 +229,14 @@ class ContextSizePolicyTest { for (context in contexts) { for (memory in memories) { for (size in sizes) { - val result = choose(header(contextLength = context), memory, size) - assertTrue( - "context=$context memory=$memory size=$size gave $result", - result in DEFAULT_CONTEXT_TOKENS..MAX_CONTEXT_TOKENS, - ) - assertEquals("must be a whole number of 256-token blocks", 0, result % 256) + for (kvType in KvCacheType.entries) { + val result = choose(header(contextLength = context), memory, size, kvType) + assertTrue( + "context=$context memory=$memory size=$size kv=$kvType gave $result", + result in DEFAULT_CONTEXT_TOKENS..MAX_CONTEXT_TOKENS, + ) + assertEquals("must be a whole number of 256-token blocks", 0, result % 256) + } } } } diff --git a/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/KvCacheTypeTest.kt b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/KvCacheTypeTest.kt new file mode 100644 index 0000000..6d2ad9e --- /dev/null +++ b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/KvCacheTypeTest.kt @@ -0,0 +1,87 @@ +package com.itsaky.androidide.plugins.aiagentlocal.model + +import org.junit.Assert.assertEquals +import org.junit.Assert.assertFalse +import org.junit.Assert.assertTrue +import org.junit.Test + +class KvCacheTypeTest { + + /** Widths default to 64, a multiple of the q8_0 block, as most models' heads are. */ + private fun header( + embeddingLength: Long? = 1024L, + headCount: Long? = 16L, + keyLength: Long? = 64L, + valueLength: Long? = 64L, + ) = GgufHeader( + architecture = "llama", + blockCount = 24L, + contextLength = 8192L, + embeddingLength = embeddingLength, + headCount = headCount, + headCountKv = 8L, + keyLength = keyLength, + valueLength = valueLength, + ) + + @Test + fun givenF16_whenSizingElements_thenChargesTwoBytesEach() { + assertEquals(64L, KvCacheType.F16.bytesFor(32L)) + assertEquals(2L, KvCacheType.F16.bytesFor(1L)) + } + + @Test + fun givenQ8_0_whenSizingOneBlock_thenChargesTheBlockPlusItsScale() { + assertEquals(34L, KvCacheType.Q8_0.bytesFor(32L)) + } + + @Test + fun givenQ8_0_whenSizingATypicalToken_thenCostsJustOverHalfOfF16() { + val elements = 24L * 8L * (64L + 64L) + assertEquals(26112L, KvCacheType.Q8_0.bytesFor(elements)) + assertEquals(49152L, KvCacheType.F16.bytesFor(elements)) + } + + @Test + fun givenNoHeader_whenAskingF16_thenStillSupported() { + // f16 has no shape constraint, so an unreadable header cannot rule it out. + assertTrue(KvCacheType.F16.supports(null)) + } + + @Test + fun givenNoHeader_whenAskingQ8_0_thenNotSupported() { + assertFalse(KvCacheType.Q8_0.supports(null)) + } + + @Test + fun givenDeclaredWidthsInWholeBlocks_whenAskingQ8_0_thenSupported() { + assertTrue(KvCacheType.Q8_0.supports(header())) + } + + @Test + fun givenKeyWidthNotInWholeBlocks_whenAskingQ8_0_thenNotSupported() { + assertFalse(KvCacheType.Q8_0.supports(header(keyLength = 80L))) + } + + @Test + fun givenValueWidthNotInWholeBlocks_whenAskingQ8_0_thenNotSupported() { + assertFalse(KvCacheType.Q8_0.supports(header(valueLength = 48L))) + } + + @Test + fun givenUndeclaredWidths_whenAskingQ8_0_thenJudgesTheDerivedWidth() { + // 1024 / 16 = 64, a whole number of blocks; 1200 / 16 = 75 is not. + assertTrue(KvCacheType.Q8_0.supports(header(keyLength = null, valueLength = null))) + assertFalse( + KvCacheType.Q8_0.supports( + header(embeddingLength = 1200L, keyLength = null, valueLength = null) + ) + ) + } + + @Test + fun givenHeaderWithoutShapeValues_whenAskingQ8_0_thenNotSupported() { + val result = KvCacheType.Q8_0.supports(header(embeddingLength = null, keyLength = null, valueLength = null)) + assertFalse(result) + } +} diff --git a/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolverTest.kt b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolverTest.kt index b1ae2fd..277b8d2 100644 --- a/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolverTest.kt +++ b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolverTest.kt @@ -52,6 +52,14 @@ class ModelContextResolverTest { assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) } + @Test + fun givenAnUnreadableHeader_whenResolving_thenTheFallbackSizeIsTheDefaultToo() { + // No header means f16, so the native fallback has nothing shorter to drop to. + val resolved = resolve { null } + assertEquals(KvCacheType.F16, resolved.kvType) + assertEquals(resolved.contextTokens, resolved.fallbackContextTokens) + } + @Test fun givenUnknownModelSize_whenResolving_thenReturnsTheDefaultContext() { val resolved = ModelContextResolver.resolve(header(), Long.MAX_VALUE, null) diff --git a/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/packaging/PrebuiltAarAbiTest.kt b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/packaging/PrebuiltAarAbiTest.kt index 686cedb..9c53976 100644 --- a/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/packaging/PrebuiltAarAbiTest.kt +++ b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/packaging/PrebuiltAarAbiTest.kt @@ -1,6 +1,8 @@ package com.itsaky.androidide.plugins.aiagentlocal.packaging import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder import java.util.zip.ZipFile import org.junit.Assert.assertEquals import org.junit.Assert.assertTrue @@ -66,11 +68,73 @@ class PrebuiltAarAbiTest { assertTrue("libllama-android.so is missing or empty (size=$size)", size > 0) } + @Test + fun givenTheCommittedAar_whenInspected_thenEveryLibraryIsAlignedForLargePages() { + // Nothing else holds -DANDROID_SUPPORT_FLEXIBLE_PAGE_SIZES: dropped, the .so still builds, + // still loads on a 4 KB device, and fails only on a 16 KB one nobody tests on. + val offenders = jniLibs.keys + .filter { it.startsWith("jni/$expectedAbi/") } + .associateWith { loadSegmentAlignments(it) } + .filterValues { aligns -> aligns.isEmpty() || aligns.any { it < MIN_SEGMENT_ALIGNMENT } } + assertEquals( + "Native libraries are not built for $MIN_SEGMENT_ALIGNMENT-byte pages" + + " (PT_LOAD p_align per library; empty means the ELF could not be read): $offenders", + emptyMap>(), + offenders, + ) + } + + /** + * The PT_LOAD segment alignments of one ELF64 shared library in the AAR. + * + * @param entryName zip entry path of the `.so` + * @return one alignment per PT_LOAD segment; empty when the entry is not a readable ELF64, which + * the caller treats as a failure rather than a pass + */ + private fun loadSegmentAlignments(entryName: String): List { + // Only the ELF header and the program-header table are read, and both sit at the front. + val head = ZipFile(aar).use { zip -> + val entry = zip.getEntry(entryName) ?: return emptyList() + zip.getInputStream(entry).use { it.readNBytes(ELF_PROBE_BYTES) } + } + if (head.size < ELF64_HEADER_BYTES) return emptyList() + val elf = ByteBuffer.wrap(head).order(ByteOrder.LITTLE_ENDIAN) + if (elf.getInt(0) != ELF_MAGIC_LE || head[EI_CLASS].toInt() != ELF_CLASS_64) return emptyList() + + val tableOffset = elf.getLong(E_PHOFF) + val entrySize = elf.getShort(E_PHENTSIZE).toInt() + val entryCount = elf.getShort(E_PHNUM).toInt() + return (0 until entryCount).mapNotNull { index -> + val at = tableOffset + index.toLong() * entrySize + // A table reaching past the probe is an unfamiliar layout, so report nothing read. + if (at < 0L || at + entrySize > head.size) return emptyList() + val offset = at.toInt() + if (elf.getInt(offset) != PT_LOAD) null else elf.getLong(offset + P_ALIGN) + } + } + private companion object { const val DEFAULT_ABI = "arm64-v8a" const val AAR_RELATIVE_PATH = "libs/v8/llama-v8-release.aar" + /** The 16 KB page size arm64 Android may use; segments must align to it or the loader rejects. */ + const val MIN_SEGMENT_ALIGNMENT = 16384L + + // ELF64 offsets and values, from the spec; the program-header table follows the 64-byte header. + const val ELF_PROBE_BYTES = 4096 + const val ELF64_HEADER_BYTES = 64 + + /** `\x7fELF` read as a little-endian int, the byte order every Android ABI uses. */ + const val ELF_MAGIC_LE = 0x464C457F + const val EI_CLASS = 4 + const val ELF_CLASS_64 = 2 + const val E_PHOFF = 32 + const val E_PHENTSIZE = 54 + const val E_PHNUM = 56 + const val PT_LOAD = 1 + const val P_ALIGN = 48 + /** The libraries llama.cpp has to produce for the wrapper to load at all. */ val REQUIRED_LIBS = setOf( "libggml-base.so",