diff --git a/ai-agent-local/libs/llama-api.jar b/ai-agent-local/libs/llama-api.jar index d25bb977..6b996697 100644 Binary files a/ai-agent-local/libs/llama-api.jar and b/ai-agent-local/libs/llama-api.jar differ diff --git a/ai-agent-local/libs/v8/llama-v8-release.aar b/ai-agent-local/libs/v8/llama-v8-release.aar index 06fb4742..a440446d 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/src/main/cpp/llama-android.cpp b/ai-agent-local/llama-impl/src/main/cpp/llama-android.cpp index 22e5cb5c..2698b07c 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 @@ -1,5 +1,6 @@ #include #include +#include #include #include #include @@ -12,7 +13,7 @@ #include "llama.h" #include "common.h" -#define TAG "llama-android.cpp" +#define TAG "AiAgentLocal.llama-android" #define LOGi(...) __android_log_print(ANDROID_LOG_INFO, TAG, __VA_ARGS__) #define LOGe(...) __android_log_print(ANDROID_LOG_ERROR, TAG, __VA_ARGS__) @@ -37,6 +38,38 @@ static std::string g_generated_text; static std::atomic g_stop_requested(false); static std::mutex g_globals_mutex; +/** + * Raises a Java exception, tolerating a FindClass that cannot resolve the name, since ThrowNew on a + * null jclass is undefined behaviour. Callers still own their resources: release them first, because + * only Release/Delete/Exception calls are legal once an exception is pending. + * + * @param env the calling thread's JNI environment + * @param class_name JNI name of the exception to raise, e.g. "java/lang/IllegalStateException" + * @param message the exception message + */ +static void throw_java(JNIEnv *env, const char *class_name, const char *message) { + jclass exception_class = env->FindClass(class_name); + if (!exception_class) { + LOGe("jni: cannot raise %s (\"%s\"): class not found", class_name, message); + return; + } + env->ThrowNew(exception_class, message); + env->DeleteLocalRef(exception_class); +} + +/** + * The token capacity a batch was allocated with, recorded by new_batch(). llama_batch itself only + * carries n_tokens (how full it is), not how large it is, so the map is the only record. + * + * @param batch a batch created by new_batch() + * @return its capacity in tokens, or 0 if it was not created here + */ +static size_t batch_capacity_of(llama_batch *batch) { + std::lock_guard lock(g_globals_mutex); + auto it = g_batch_n_tokens.find(batch); + return it == g_batch_n_tokens.end() ? 0 : (size_t) std::max(0, it->second); +} + bool is_valid_utf8(const char *string) { if (!string) { return true; @@ -82,10 +115,26 @@ static std::atomic g_n_threads_batch(-1); static std::atomic g_temperature(0.7f); static std::atomic g_top_p(0.9f); static std::atomic g_top_k(40); -static std::atomic g_n_ctx(4096); +/** + * Context used when the caller passes a non-positive one; mirrors ContextSizePolicy's floor. Only a + * guard against a bad argument — the size is chosen in Kotlin and passed to new_context per load. + */ +static constexpr int DEFAULT_N_CTX = 4096; static std::atomic g_kv_cache_reuse(true); static std::vector g_cached_tokens; +/** + * Drops both the KV cache and the record of what it held, after a prefill that did not complete. + * Leaving either behind would have the next turn reuse a prefix the cache no longer matches. + * + * @param context the context whose memory to clear + */ +static void forget_cached_prefix(llama_context *context) { + llama_memory_clear(llama_get_memory(context), true); + std::lock_guard lock(g_globals_mutex); + g_cached_tokens.clear(); +} + // Converts standard UTF-8 to UTF-16. NewStringUTF() is unusable here because it // expects modified UTF-8 (CESU-8), so 4-byte sequences such as emoji would mangle. // Invalid bytes become '?' so a truncated sequence cannot corrupt the remainder. @@ -191,15 +240,6 @@ Java_android_llama_cpp_LLamaAndroid_native_1configureSampling(JNIEnv *, jclass, g_top_k.store(validated_top_k); } -extern "C" -JNIEXPORT void JNICALL -Java_android_llama_cpp_LLamaAndroid_native_1configureContext(JNIEnv *, jclass, jint n_ctx) { - if (n_ctx <= 0) { - return; - } - g_n_ctx.store(n_ctx); -} - extern "C" JNIEXPORT void JNICALL Java_android_llama_cpp_LLamaAndroid_native_1configureKvCacheReuse(JNIEnv *, jclass, jboolean enabled) { @@ -305,14 +345,14 @@ Java_android_llama_cpp_LLamaAndroid_load_1model(JNIEnv *env, jobject, jstring fi llama_model_params model_params = llama_model_default_params(); auto path_to_model = env->GetStringUTFChars(filename, 0); - LOGi("Loading model from %s", path_to_model); + LOGi("model: loading from %s", path_to_model); auto model = llama_model_load_from_file(path_to_model, model_params); env->ReleaseStringUTFChars(filename, path_to_model); if (!model) { - LOGe("load_model() failed"); - env->ThrowNew(env->FindClass("java/lang/IllegalStateException"), "load_model() failed"); + LOGe("model: load_model() failed"); + throw_java(env, "java/lang/IllegalStateException", "load_model() failed"); return 0; } @@ -327,12 +367,12 @@ Java_android_llama_cpp_LLamaAndroid_free_1model(JNIEnv *, jobject, jlong model) extern "C" JNIEXPORT jlong JNICALL -Java_android_llama_cpp_LLamaAndroid_new_1context(JNIEnv *env, jobject, jlong jmodel) { +Java_android_llama_cpp_LLamaAndroid_new_1context(JNIEnv *env, jobject, jlong jmodel, jint jn_ctx) { auto model = reinterpret_cast(jmodel); if (!model) { - LOGe("new_context(): model cannot be null"); - env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"), "Model cannot be null"); + LOGe("context: model cannot be null"); + throw_java(env, "java/lang/IllegalArgumentException", "Model cannot be null"); return 0; } @@ -345,24 +385,36 @@ Java_android_llama_cpp_LLamaAndroid_new_1context(JNIEnv *env, jobject, jlong jmo if (n_threads_batch <= 0) { n_threads_batch = n_threads; } - LOGi("Using %d threads (batch=%d)", n_threads, n_threads_batch); + LOGi("context: using %d threads (batch=%d)", n_threads, n_threads_batch); llama_context_params ctx_params = llama_context_default_params(); - const int configured_ctx = g_n_ctx.load(); - ctx_params.n_ctx = configured_ctx > 0 ? configured_ctx : 4096; + 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; + } + + ctx_params.n_ctx = requested_ctx; ctx_params.n_threads = n_threads; ctx_params.n_threads_batch = n_threads_batch; llama_context *context = llama_init_from_model(model, ctx_params); if (!context) { - LOGe("llama_new_context_with_model() returned null)"); - env->ThrowNew(env->FindClass("java/lang/IllegalStateException"), - "llama_new_context_with_model() returned null)"); + LOGe("context: llama_new_context_with_model() returned null"); + throw_java(env, "java/lang/IllegalStateException", + "llama_new_context_with_model() returned null)"); 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)); + // A fresh context has an empty KV cache, so the prefix record must start empty too. { std::lock_guard lock(g_globals_mutex); @@ -419,12 +471,12 @@ Java_android_llama_cpp_LLamaAndroid_bench_1model( const int n_ctx = llama_n_ctx(context); - LOGi("n_ctx = %d", n_ctx); + LOGi("bench: n_ctx = %d", n_ctx); int i, j; int nri; for (nri = 0; nri < nr; nri++) { - LOGi("Benchmark prompt processing (pp)"); + LOGi("bench: prompt processing (pp)"); common_batch_clear(*batch); @@ -438,13 +490,13 @@ Java_android_llama_cpp_LLamaAndroid_bench_1model( const auto t_pp_start = ggml_time_us(); if (llama_decode(context, *batch) != 0) { - LOGi("llama_decode() failed during prompt processing"); + LOGi("bench: llama_decode() failed during prompt processing"); } const auto t_pp_end = ggml_time_us(); // bench text generation - LOGi("Benchmark text generation (tg)"); + LOGi("bench: text generation (tg)"); llama_memory_clear(llama_get_memory(context), false); const auto t_tg_start = ggml_time_us(); @@ -455,9 +507,9 @@ Java_android_llama_cpp_LLamaAndroid_bench_1model( common_batch_add(*batch, 0, i, {j}, true); } - LOGi("llama_decode() text generation: %d", i); + LOGi("bench: llama_decode() text generation: %d", i); if (llama_decode(context, *batch) != 0) { - LOGi("llama_decode() failed during text generation"); + LOGi("bench: llama_decode() failed during text generation"); } } @@ -477,7 +529,7 @@ Java_android_llama_cpp_LLamaAndroid_bench_1model( pp_std += speed_pp * speed_pp; tg_std += speed_tg * speed_tg; - LOGi("pp %f t/s, tg %f t/s", speed_pp, speed_tg); + LOGi("bench: pp %f t/s, tg %f t/s", speed_pp, speed_tg); } pp_avg /= double(nr); @@ -737,23 +789,23 @@ Java_android_llama_cpp_LLamaAndroid_completion_1init( int n_ctx = llama_n_ctx(context); size_t n_kv_req = tokens_list.size() + static_cast(n_len); - LOGi("n_len = %d, n_ctx = %d, n_kv_req = %zu", n_len, n_ctx, n_kv_req); + LOGi("prefill: n_len = %d, n_ctx = %d, n_kv_req = %zu", n_len, n_ctx, n_kv_req); if (n_kv_req > n_ctx) { - LOGe("error: n_kv_req > n_ctx, the required KV cache size is not big enough"); - env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"), - "Prompt is too long for the model's context size."); + LOGe("prefill: n_kv_req > n_ctx, the required KV cache size is not big enough"); + // Released before returning, as on every other exit from here: jtext is pinned until it is. + env->ReleaseStringUTFChars(jtext, text); + throw_java(env, "java/lang/IllegalArgumentException", + "Prompt is too long for the model's context size."); return 0; } g_prompt_tokens = static_cast(tokens_list.size()); for (auto id: tokens_list) { - LOGv("token: `%s`-> %d ", common_token_to_piece(context, id).c_str(), id); + LOGv("prefill: token `%s` -> %d", common_token_to_piece(context, id).c_str(), id); } - common_batch_clear(*batch); - // Reuse the longest common prefix with the cached sequence so the unchanged prefix (system prompt) isn't re-prefilled. size_t lcp = 0; { @@ -780,24 +832,54 @@ Java_android_llama_cpp_LLamaAndroid_completion_1init( llama_memory_seq_rm(mem, 0, (llama_pos) lcp, -1); } - { - std::lock_guard lock(g_globals_mutex); - g_cached_tokens.assign(tokens_list.begin(), tokens_list.end()); + // Sliced: the batch's fixed capacity can now sit far below n_ctx, and overrunning it wrecks the heap. + const size_t batch_capacity = batch_capacity_of(batch); + const size_t chunk_limit = std::min(batch_capacity, llama_n_batch(context)); + + if (chunk_limit == 0) { + // Not llama_n_batch(context): an untracked batch has an unknown allocation to overrun. + LOGe("prefill: batch was not created by new_batch(), so its capacity is unknown"); + forget_cached_prefix(context); + env->ReleaseStringUTFChars(jtext, text); + throw_java(env, "java/lang/IllegalStateException", + "Batch capacity is unknown."); + return 0; } - // Prefill only the divergent tail. - for (size_t i = lcp; i < tokens_list.size(); i++) { - common_batch_add(*batch, tokens_list[i], (llama_pos) i, {0}, false); - } + const size_t prefill_tokens = tokens_list.size() - lcp; + const size_t slices = (prefill_tokens + chunk_limit - 1) / chunk_limit; + // The only direct evidence the chunked path ran rather than the old single-batch prefill. + LOGi("prefill: %zu tokens (%zu reused from cache) in %zu slice(s) of at most %zu", + prefill_tokens, lcp, slices, chunk_limit); - if (batch->n_tokens > 0) { - // llama_decode will output logits only for the last token of the prompt - batch->logits[batch->n_tokens - 1] = true; - if (llama_decode(context, *batch) != 0) { - LOGe("llama_decode() failed"); + for (size_t start = lcp; start < tokens_list.size(); start += chunk_limit) { + const size_t end = std::min(start + chunk_limit, tokens_list.size()); + common_batch_clear(*batch); + for (size_t i = start; i < end; i++) { + common_batch_add(*batch, tokens_list[i], (llama_pos) i, {0}, false); + } + + // Only the last prompt token needs logits; earlier slices just populate the KV cache. + if (end == tokens_list.size() && batch->n_tokens > 0) { + batch->logits[batch->n_tokens - 1] = true; + } + + if (batch->n_tokens > 0 && llama_decode(context, *batch) != 0) { + LOGe("prefill: llama_decode() failed for tokens %zu..%zu", start, end); + forget_cached_prefix(context); + env->ReleaseStringUTFChars(jtext, text); + throw_java(env, "java/lang/IllegalStateException", + "Failed to process the prompt."); + return 0; } } + // Recorded only after every slice decoded, so the record matches what the KV cache holds. + { + std::lock_guard lock(g_globals_mutex); + g_cached_tokens.assign(tokens_list.begin(), tokens_list.end()); + } + env->ReleaseStringUTFChars(jtext, text); return g_prompt_tokens; @@ -871,7 +953,7 @@ Java_android_llama_cpp_LLamaAndroid_completion_1loop( if (!stop_str.empty() && generated_snapshot.length() >= stop_str.length()) { auto pos = generated_snapshot.find(stop_str); if (pos != std::string::npos) { - LOGi("Stop string matched: %s", stop_str.c_str()); + LOGi("generate: stop string matched: %s", stop_str.c_str()); size_t prefix_len = pos > prior_len ? pos - prior_len : 0; if (prefix_len > 0) { std::string prefix; @@ -928,7 +1010,7 @@ Java_android_llama_cpp_LLamaAndroid_completion_1loop( env->CallVoidMethod(intvar_ncur, la_int_var_inc); if (llama_decode(context, *batch) != 0) { - LOGe("llama_decode() returned null"); + LOGe("generate: llama_decode() returned null"); return nullptr; } 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 10f73297..ecf304cf 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 @@ -12,12 +12,19 @@ import java.util.concurrent.Executors import java.util.concurrent.atomic.AtomicBoolean import kotlin.concurrent.thread +/** + * Prefix on every logger name this module creates. Duplicated from LOG_PREFIX in the plugin's + * logging/LogTags.kt, which this module cannot import, and kept in step with it by hand: llama-impl + * only ever ships inside ai-agent-local's AAR, so a name without it points at no plugin. + */ +private const val LOG_PREFIX = "AiAgentLocal" + /** * Static library loader - ensures native library is loaded before any static methods are called. * This object's init block runs when the object is first accessed. */ private object NativeLibraryLoader { - private val log = LoggerFactory.getLogger("llama.cpp.loader") + private val log = LoggerFactory.getLogger("$LOG_PREFIX.NativeLibraryLoader") @Volatile private var loaded = false @@ -49,7 +56,7 @@ private object NativeLibraryLoader { class LLamaAndroid : ILlamaController { - private val log = LoggerFactory.getLogger(LLamaAndroid::class.java) + private val log = LoggerFactory.getLogger("$LOG_PREFIX.LLamaAndroid") init { // Ensure native library is loaded when any instance is created @@ -158,7 +165,7 @@ 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): Long + private external fun new_context(model: Long, nCtx: Int): Long private external fun free_context(context: Long) private external fun backend_init(numa: Boolean) private external fun backend_free() @@ -231,14 +238,24 @@ class LLamaAndroid : ILlamaController { } } - override suspend fun load(pathToModel: String) { + 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. + * + * @param pathToModel filesystem path to the `.gguf` model + * @param nCtx context size in tokens; anything non-positive means [DEFAULT_N_CTX] + */ + suspend fun load(pathToModel: String, nCtx: Int) { 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) + val context = new_context(model, nCtx) if (context == 0L) throw IllegalStateException("new_context() failed") val batch = new_batch(2048, 0, 1) @@ -339,7 +356,10 @@ class LLamaAndroid : ILlamaController { } companion object { - private val nativeLog = LoggerFactory.getLogger("llama.cpp") + private val nativeLog = LoggerFactory.getLogger("$LOG_PREFIX.llama.cpp") + + /** Context a [load] gets when the caller does not pick one; matches DEFAULT_N_CTX natively. */ + const val DEFAULT_N_CTX = 4096 // External native methods @JvmStatic @@ -348,9 +368,6 @@ class LLamaAndroid : ILlamaController { @JvmStatic private external fun native_configureSampling(temperature: Float, topP: Float, topK: Int) - @JvmStatic - private external fun native_configureContext(nCtx: Int) - @JvmStatic private external fun native_configureKvCacheReuse(enabled: Boolean) @@ -367,12 +384,6 @@ class LLamaAndroid : ILlamaController { native_configureSampling(temperature, topP, topK) } - @JvmStatic - fun configureContext(nCtx: Int) { - NativeLibraryLoader.ensureLoaded() - native_configureContext(nCtx) - } - @JvmStatic fun configureKvCacheReuse(enabled: Boolean) { NativeLibraryLoader.ensureLoaded() 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 cb11c9b2..f83a2616 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,6 +12,7 @@ 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.ModelContextResolver import com.itsaky.androidide.plugins.aiagentlocal.model.ModelLoadDiagnostics import com.itsaky.androidide.plugins.aiagentlocal.model.ModelLoadMessages import com.itsaky.androidide.plugins.aiagentlocal.preferences.LocalLlmPreferences @@ -33,6 +34,7 @@ import kotlinx.coroutines.ensureActive import kotlinx.coroutines.launch import kotlinx.coroutines.sync.Mutex import kotlinx.coroutines.sync.withLock +import kotlinx.coroutines.withContext /** * Local LLM backend using llama-impl for on-device inference. @@ -307,13 +309,16 @@ class LocalLlmBackend( } // Measured after the unload: availMem excludes the context and batch it just released. - ModelLoadDiagnostics.refuseBeforeLoad(availableMemoryBytes())?.let { shortfall -> + val availableBytes = availableMemoryBytes() + ModelLoadDiagnostics.refuseBeforeLoad(availableBytes)?.let { shortfall -> throw ModelLoadException(loadMessages.describe(shortfall), shortfall) } + val contextTokens = resolveContextSize(resolvedPath, availableBytes) + context.logger.info("Loading model: $resolvedPath") try { - llama.load(resolvedPath) + llama.load(resolvedPath, contextTokens) } catch (e: CancellationException) { throw e } catch (e: Exception) { @@ -328,6 +333,30 @@ class LocalLlmBackend( context.logger.info("Model loaded successfully") } + /** + * 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. + * + * @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 + */ + private suspend fun resolveContextSize(resolvedPath: String, availableBytes: Long): Int { + val resolved = withContext(Dispatchers.IO) { + ModelContextResolver.resolve(availableBytes.takeIf { it >= 0L }) { + File(resolvedPath).takeIf { it.isFile }?.inputStream() + } + } + // Unconditional: a wrongly sized context otherwise just reads as the assistant forgetting. + context.logger.info( + "Context size for $resolvedPath: ${resolved.contextTokens} tokens" + + " (model advertises ${resolved.advertisedTokens ?: "unknown"}," + + " ${if (availableBytes >= 0L) "$availableBytes bytes free" else "free RAM unknown"})" + ) + return resolved.contextTokens + } + /** * @return free RAM the OS reports, or -1 if unreadable (diagnosis then skips the low-memory case) */ 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 new file mode 100644 index 00000000..28789162 --- /dev/null +++ b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicy.kt @@ -0,0 +1,61 @@ +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. + */ +object ContextSizePolicy { + + /** + * The context every load got before this policy existed, and now both the fallback for any + * unreadable input and the floor. Below it the native prompt check starts rejecting + * conversations that fit today, so a smaller context costs working prompts rather than saving. + */ + const val DEFAULT_CONTEXT_TOKENS = 4096 + + /** + * Ceiling, whatever the model advertises and the device can afford. Four times the old fixed + * context and already past the point of diminishing returns, since prefill cost grows with the + * prompt. Models advertising 32k+ are capped here rather than taken at their word. + */ + const val MAX_CONTEXT_TOKENS = 16384 + + /** + * Contexts are rounded down to a multiple of this. Purely cosmetic — it keeps the chosen value + * and the llama.cpp context dump readable instead of reporting a number like 11417. + */ + private const val GRANULARITY_TOKENS = 256 + + /** + * The share of usable free RAM the KV cache may claim. The IDE and the app being edited draw on + * the same pool, and `availMem` is a snapshot taken before a load that then takes seconds, so + * half is left alone rather than sizing to the last free byte. + */ + private const val KV_BUDGET_DIVISOR = 2L + + /** + * @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 + * @return the context to configure, always between [DEFAULT_CONTEXT_TOKENS] and + * [MAX_CONTEXT_TOKENS] inclusive + */ + fun choose(header: GgufHeader?, availableBytes: Long?): 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 } + ?: return DEFAULT_CONTEXT_TOKENS + + // Compute buffers come off the top; goes negative on a short device, which the floor absorbs. + val budgetBytes = (availableBytes - ModelMemory.RUN_BUFFER_BYTES) / KV_BUDGET_DIVISOR + val affordableTokens = budgetBytes / perToken + + val ceiling = minOf(modelTokens, affordableTokens, MAX_CONTEXT_TOKENS.toLong()) + val rounded = (ceiling / GRANULARITY_TOKENS) * GRANULARITY_TOKENS + return rounded.coerceIn(DEFAULT_CONTEXT_TOKENS.toLong(), MAX_CONTEXT_TOKENS.toLong()).toInt() + } +} diff --git a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/GgufHeaderReader.kt b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/GgufHeaderReader.kt index 25693daa..985b5ee5 100644 --- a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/GgufHeaderReader.kt +++ b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/GgufHeaderReader.kt @@ -13,6 +13,8 @@ import java.io.InputStream * * @property architecture `general.architecture`; also the prefix every other key here is read under * @property blockCount transformer layers, `{arch}.block_count` + * @property contextLength the context the model was trained for, `{arch}.context_length`; the + * ceiling [ContextSizePolicy] sizes the KV cache against, and absent on files that omit it * @property embeddingLength model width, `{arch}.embedding_length` * @property headCount attention heads, `{arch}.attention.head_count` * @property headCountKv key/value heads under grouped-query attention, absent for plain MHA @@ -23,6 +25,7 @@ import java.io.InputStream data class GgufHeader( val architecture: String?, val blockCount: Long?, + val contextLength: Long? = null, val embeddingLength: Long?, val headCount: Long?, val headCountKv: Long?, @@ -59,6 +62,7 @@ internal object GgufHeaderReader { // Matched by suffix, then attributed to the "{arch}." prefix they carry — see [readHeader]. private const val SUFFIX_BLOCK_COUNT = ".block_count" + private const val SUFFIX_CONTEXT_LENGTH = ".context_length" private const val SUFFIX_EMBEDDING_LENGTH = ".embedding_length" private const val SUFFIX_HEAD_COUNT = ".attention.head_count" private const val SUFFIX_HEAD_COUNT_KV = ".attention.head_count_kv" @@ -138,6 +142,7 @@ internal object GgufHeaderReader { /** The shape values seen under one `{arch}.` prefix. A file may carry more than one. */ private class ArchShape { var blockCount: Long? = null + var contextLength: Long? = null var embeddingLength: Long? = null var headCount: Long? = null var headCountKv: Long? = null @@ -170,6 +175,9 @@ internal object GgufHeaderReader { key.endsWith(SUFFIX_BLOCK_COUNT) -> shapeFor(shapes, key, SUFFIX_BLOCK_COUNT).blockCount = readInteger(input, type, wide) + key.endsWith(SUFFIX_CONTEXT_LENGTH) -> + shapeFor(shapes, key, SUFFIX_CONTEXT_LENGTH).contextLength = readInteger(input, type, wide) + key.endsWith(SUFFIX_EMBEDDING_LENGTH) -> shapeFor(shapes, key, SUFFIX_EMBEDDING_LENGTH).embeddingLength = readInteger(input, type, wide) @@ -195,6 +203,7 @@ internal object GgufHeaderReader { return GgufHeader( architecture = architecture, blockCount = shape?.blockCount, + contextLength = shape?.contextLength, embeddingLength = shape?.embeddingLength, headCount = shape?.headCount, headCountKv = shape?.headCountKv, 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 new file mode 100644 index 00000000..b7ecfd0d --- /dev/null +++ b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolver.kt @@ -0,0 +1,43 @@ +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. + * + * @property contextTokens the context to load with; always a value [ContextSizePolicy] returned + * @property header the model's parsed metadata, or null when it could not be read + */ +internal data class ModelContextSize( + val contextTokens: Int, + val header: GgufHeader?, +) { + + /** The context the model claims to support, or null when the header did not say. */ + val advertisedTokens: Long? get() = header?.contextLength +} + +/** + * 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. + */ +internal object ModelContextResolver { + + /** + * Fails open by construction, with no error path of its own: [GgufHeaderReader.read] turns + * anything thrown while opening or parsing into a null header, and [ContextSizePolicy.choose] + * answers its default for one. Blocking — the header sits at the front of the model file. + * + * @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 + */ + fun resolve(availableBytes: Long?, openStream: () -> InputStream?): ModelContextSize { + val header = GgufHeaderReader.read(openStream) + return ModelContextSize( + contextTokens = ContextSizePolicy.choose(header, availableBytes), + 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 0f318c1e..a25e7c26 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,17 +23,11 @@ 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 and batch sizes below are ai-agent-local's, - * fixed on its native side: an estimate has to model the loader that will actually run. + * 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). */ object ModelMemoryEstimator { - /** - * The context every load gets, hard-coded as `ctx_params.n_ctx` in ai-agent-local's `llama-android.cpp`. - * The KV cache is sized from it, so keep the two in step. - */ - const val RUNTIME_CONTEXT_TOKENS = 4096L - /** Two bytes per cached element: f16, the default KV type. */ private const val KV_BYTES_PER_ELEMENT = 2L @@ -55,11 +49,17 @@ object ModelMemoryEstimator { /** * @param fileSizeBytes the model file's size, or null when it is unknown * @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. * @return the estimate, or null when there is nothing to base one on */ - fun estimate(fileSizeBytes: Long?, header: GgufHeader?): MemoryEstimate? { + fun estimate( + fileSizeBytes: Long?, + header: GgufHeader?, + contextTokens: Int, + ): MemoryEstimate? { if (fileSizeBytes == null || fileSizeBytes <= 0L) return null - val kvCacheBytes = header?.let(::kvCacheBytes) + val kvCacheBytes = header?.let { kvCacheBytes(it, contextTokens) } return if (kvCacheBytes != null) { MemoryEstimate(fileSizeBytes, kvCacheBytes + COMPUTE_BUFFER_BYTES, fromHeader = true) } else { @@ -73,17 +73,31 @@ object ModelMemoryEstimator { } /** - * KV cache size for a full context: one key and one value entry per kv head, per layer, per - * position. Null unless every value it needs is present and within its ceiling. + * 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? { + if (contextTokens <= 0) return null + val perToken = kvBytesPerToken(header) ?: return null + return perToken * contextTokens + } + + /** + * What one cached position costs: one key and one value entry per kv head, per layer. The + * factor [ContextSizePolicy] divides a RAM budget by, so sizing and estimate cannot drift apart. + * Stays under 2^44 within the ceilings below, so any context the policy returns fits a Long. + * + * @param header the model's metadata + * @return bytes of KV cache per token, or null if the header does not say enough */ - private fun kvCacheBytes(header: GgufHeader): Long? { + internal fun kvBytesPerToken(header: GgufHeader): 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 = 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 * RUNTIME_CONTEXT_TOKENS * kvHeads * (keyWidth + valueWidth) + return KV_BYTES_PER_ELEMENT * layers * kvHeads * (keyWidth + 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 f5f5cf89..3e9b1dba 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 @@ -14,7 +14,7 @@ import com.itsaky.androidide.plugins.aiagentlocal.logging.LOG_PREFIX import com.itsaky.androidide.plugins.aiagentlocal.model.ContentModelFileSource import com.itsaky.androidide.plugins.aiagentlocal.model.DeviceMemory import com.itsaky.androidide.plugins.aiagentlocal.model.GgufFileInspector -import com.itsaky.androidide.plugins.aiagentlocal.model.GgufHeaderReader +import com.itsaky.androidide.plugins.aiagentlocal.model.ModelContextResolver import com.itsaky.androidide.plugins.aiagentlocal.model.ModelFileInfo import com.itsaky.androidide.plugins.aiagentlocal.model.ModelFileSource import com.itsaky.androidide.plugins.aiagentlocal.model.ModelMemoryEstimator @@ -276,12 +276,17 @@ class LocalLlmSettingsViewModel( context: Context ): Boolean { val modelName = fileInfo.displayName + // Never cached: the user may have just closed apps to make room. + val availableBytes = deviceMemory.availableBytes() + // Resolved the same way the load will resolve it, so the warning describes the real allocation. + val resolved = ModelContextResolver.resolve(availableBytes) { + modelFiles.openStream(context, uriString) + } val estimate = ModelMemoryEstimator.estimate( fileSizeBytes = fileInfo.sizeBytes, - header = GgufHeaderReader.read { modelFiles.openStream(context, uriString) }, + header = resolved.header, + contextTokens = resolved.contextTokens, ) - // Read last and never cached: the user may have just closed apps to make room. - val availableBytes = deviceMemory.availableBytes() return when (val verdict = ModelMemoryGate.evaluate(estimate, availableBytes)) { ModelMemoryGate.Verdict.Safe -> true 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 new file mode 100644 index 00000000..c1f2098c --- /dev/null +++ b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicyTest.kt @@ -0,0 +1,139 @@ +package com.itsaky.androidide.plugins.aiagentlocal.model + +import com.itsaky.androidide.plugins.aiagentlocal.model.ContextSizePolicy.DEFAULT_CONTEXT_TOKENS +import com.itsaky.androidide.plugins.aiagentlocal.model.ContextSizePolicy.MAX_CONTEXT_TOKENS +import org.junit.Assert.assertEquals +import org.junit.Assert.assertTrue +import org.junit.Test + +class ContextSizePolicyTest { + + /** A small GQA model: 24 layers, 8 kv heads of 64, so 2 * 24 * 8 * 128 = 49_152 B/token. */ + private fun header( + contextLength: Long? = 32768L, + blockCount: Long? = 24L, + headCount: Long? = 16L, + headCountKv: Long? = 8L, + keyLength: Long? = 64L, + valueLength: Long? = 64L, + ) = GgufHeader( + architecture = "llama", + blockCount = blockCount, + contextLength = contextLength, + embeddingLength = 1024L, + headCount = headCount, + headCountKv = headCountKv, + keyLength = keyLength, + valueLength = valueLength, + ) + + private val bytesPerToken = 2L * 24L * 8L * (64L + 64L) + + /** Free RAM that affords exactly [tokens], undoing the reserve and the budget divisor. */ + private fun ramAffording(tokens: Long): Long = + tokens * bytesPerToken * 2L + ModelMemory.RUN_BUFFER_BYTES + + @Test + fun givenNoHeader_whenChoosing_thenFallsBackToDefault() { + assertEquals(DEFAULT_CONTEXT_TOKENS, ContextSizePolicy.choose(null, ramAffording(100_000L))) + } + + @Test + fun givenNoContextLengthInHeader_whenChoosing_thenFallsBackToDefault() { + val result = ContextSizePolicy.choose(header(contextLength = null), ramAffording(100_000L)) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenUnreadableMemory_whenChoosing_thenFallsBackToDefault() { + assertEquals(DEFAULT_CONTEXT_TOKENS, ContextSizePolicy.choose(header(), null)) + } + + @Test + fun givenHeaderMissingShapeValues_whenChoosing_thenFallsBackToDefault() { + // No block count and no way to derive one: the per-token cost is unknowable. + val result = ContextSizePolicy.choose(header(blockCount = null), ramAffording(100_000L)) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenLargeContextModelAndAmpleRam_whenChoosing_thenCapsAtMaximum() { + val result = ContextSizePolicy.choose(header(contextLength = 32768L), ramAffording(100_000L)) + assertEquals(MAX_CONTEXT_TOKENS, result) + } + + @Test + fun givenModelContextBelowFloor_whenChoosing_thenHoldsTheFloor() { + val result = ContextSizePolicy.choose(header(contextLength = 2048L), ramAffording(100_000L)) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenModelContextBetweenFloorAndMaximum_whenChoosing_thenUsesTheModelContext() { + val result = ContextSizePolicy.choose(header(contextLength = 8192L), ramAffording(100_000L)) + assertEquals(8192, result) + } + + @Test + fun givenRamBoundDevice_whenChoosing_thenReturnsRoundedAffordableContext() { + // Affords 10_000 tokens; expect it rounded down to a multiple of 256. + val result = ContextSizePolicy.choose(header(), ramAffording(10_000L)) + assertEquals(9984, result) + assertTrue("must stay under the model's own context", result < 32768) + } + + @Test + fun givenTightRam_whenChoosing_thenNeverGoesBelowTheFloor() { + val result = ContextSizePolicy.choose(header(), ramAffording(1_000L)) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenNoFreeMemory_whenChoosing_thenReturnsFloorRatherThanZero() { + assertEquals(DEFAULT_CONTEXT_TOKENS, ContextSizePolicy.choose(header(), 0L)) + } + + @Test + fun givenLessFreeRamThanTheComputeReserve_whenChoosing_thenReturnsFloor() { + // Budget goes negative here; the floor has to absorb it rather than a negative context. + val result = ContextSizePolicy.choose(header(), ModelMemory.RUN_BUFFER_BYTES / 2) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenAbsurdContextLength_whenChoosing_thenCapsAtMaximumWithoutOverflow() { + val result = ContextSizePolicy.choose(header(contextLength = Long.MAX_VALUE), ramAffording(100_000L)) + assertEquals(MAX_CONTEXT_TOKENS, result) + } + + @Test + fun givenNegativeContextLength_whenChoosing_thenFallsBackToDefault() { + val result = ContextSizePolicy.choose(header(contextLength = -1L), ramAffording(100_000L)) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenAbsurdShapeValues_whenChoosing_thenFallsBackToDefaultWithoutOverflow() { + val result = ContextSizePolicy.choose( + header(blockCount = Long.MAX_VALUE, keyLength = Long.MAX_VALUE), + ramAffording(100_000L), + ) + 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) + } + } + } +} 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 new file mode 100644 index 00000000..69647436 --- /dev/null +++ b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolverTest.kt @@ -0,0 +1,58 @@ +package com.itsaky.androidide.plugins.aiagentlocal.model + +import com.itsaky.androidide.plugins.aiagentlocal.model.ContextSizePolicy.DEFAULT_CONTEXT_TOKENS +import java.io.IOException +import java.io.InputStream +import org.junit.Assert.assertEquals +import org.junit.Assert.assertNull +import org.junit.Test + +/** + * Pins the fail-open contract: no way of failing to read a header may propagate out of [resolve], + * because the caller is a model load that should proceed at the default context instead of aborting. + */ +class ModelContextResolverTest { + + @Test + fun givenAnOpenerThatThrows_whenResolving_thenReturnsTheDefaultContext() { + val resolved = ModelContextResolver.resolve(Long.MAX_VALUE) { + throw IOException("permission denied") + } + assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) + assertNull(resolved.header) + } + + @Test + fun givenAnOpenerReturningNoStream_whenResolving_thenReturnsTheDefaultContext() { + val resolved = ModelContextResolver.resolve(Long.MAX_VALUE) { null } + assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) + assertNull(resolved.header) + } + + @Test + fun givenAStreamThatIsNotGguf_whenResolving_thenReturnsTheDefaultContext() { + val resolved = ModelContextResolver.resolve(Long.MAX_VALUE) { + "not a model file".byteInputStream() + } + assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) + assertNull(resolved.header) + } + + @Test + fun givenAStreamThatThrowsMidRead_whenResolving_thenReturnsTheDefaultContext() { + val resolved = ModelContextResolver.resolve(Long.MAX_VALUE) { ThrowingStream() } + assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) + assertNull(resolved.header) + } + + @Test + fun givenUnknownFreeMemory_whenResolving_thenReturnsTheDefaultContext() { + val resolved = ModelContextResolver.resolve(null) { "not a model file".byteInputStream() } + assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) + } + + /** 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") + } +}