diff --git a/ai-agent-local/libs/v8/llama-v8-release.aar b/ai-agent-local/libs/v8/llama-v8-release.aar index a440446..fe7ab43 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 7bb87ea..dcc70a8 100644 --- a/ai-agent-local/llama-impl/build.gradle.kts +++ b/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() 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 2698b07..f23bbed 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,27 @@ 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. + * + * @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(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. { 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 ecf304c..b550700 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,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) 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 f83a261..d9c332d 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 @@ -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,15 +341,20 @@ 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() @@ -350,11 +362,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 2878916..e7e90be 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,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. 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 b7ecfd0..e9f8ac0 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 @@ -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, ) } 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 a25e7c2..64558e2 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 @@ -23,14 +23,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 — - * pass [ContextSizePolicy.choose]'s answer, or it describes an allocation nothing makes (ADFA-5187). + * Android-free, so the arithmetic is unit-testable. The context and cache type it measures at are + * the caller's — pass what [ModelContextResolver] resolved, or it describes an allocation nothing + * makes. See ADFA-5187 and ADFA-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 @@ -51,15 +49,18 @@ object ModelMemoryEstimator { * @param header the model's metadata, or null when it could not be read * @param contextTokens the context the load will be given; required, because a default here * would silently describe an allocation nobody makes. [ModelContextResolver] supplies it. + * @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 { @@ -76,9 +77,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 } @@ -88,16 +89,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 3e9b1db..78b1ca4 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,7 @@ class LocalLlmSettingsViewModel( fileSizeBytes = fileInfo.sizeBytes, header = resolved.header, contextTokens = resolved.contextTokens, + kvType = resolved.kvType, ) return when (val verdict = ModelMemoryGate.evaluate(estimate, availableBytes)) { 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 c1f2098..63639cd 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 @@ -121,18 +121,55 @@ 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, ContextSizePolicy.choose(header(), ram, KvCacheType.F16)) + assertEquals(9216, ContextSizePolicy.choose(header(), ram, KvCacheType.Q8_0)) + } + + @Test + fun givenAModelContextBelowWhatQ8_0Affords_whenChoosing_thenTheModelStillCaps() { + val result = ContextSizePolicy.choose(header(contextLength = 8192L), ramAffording(5_000L), KvCacheType.Q8_0) + assertEquals(8192, result) + } + + @Test + fun givenTightRamUnderQ8_0_whenChoosing_thenNeverGoesBelowTheFloor() { + val result = ContextSizePolicy.choose(header(), ramAffording(1_000L), KvCacheType.Q8_0) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + @Test fun givenAnyInputs_whenChoosing_thenResultStaysWithinTheDeclaredBounds() { val contexts = listOf(null, -1L, 0L, 512L, 4096L, 8192L, 32768L, Long.MAX_VALUE) val memories = listOf(null, 0L, 1L, ModelMemory.RUN_BUFFER_BYTES, ramAffording(50_000L), Long.MAX_VALUE) for (context in contexts) { for (memory in memories) { - val result = ContextSizePolicy.choose(header(contextLength = context), memory) - assertTrue( - "context=$context memory=$memory 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 = ContextSizePolicy.choose(header(contextLength = context), memory, kvType) + assertTrue( + "context=$context memory=$memory 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 6964743..b9b0147 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 @@ -51,6 +51,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 = ModelContextResolver.resolve(Long.MAX_VALUE) { null } + assertEquals(KvCacheType.F16, resolved.kvType) + assertEquals(resolved.contextTokens, resolved.fallbackContextTokens) + } + /** Opens fine and then fails, which is the case a null-check on the opener would not cover. */ private class ThrowingStream : InputStream() { override fun read(): Int = throw IOException("device is gone")