diff --git a/ai-agent-local/libs/v8/llama-v8-release.aar b/ai-agent-local/libs/v8/llama-v8-release.aar index 1fc67318..9a4088a2 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..54697c0f 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,39 @@ 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 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; + } + + 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 +474,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 +493,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 +510,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 +532,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 +792,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 +835,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 +956,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 +1013,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..59b537b6 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,21 +238,47 @@ 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. + * + * 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. + * + * @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) - if (context == 0L) throw IllegalStateException("new_context() failed") - - val batch = new_batch(2048, 0, 1) - if (batch == 0L) throw IllegalStateException("new_batch() failed") - - val sampler = new_sampler() - if (sampler == 0L) throw IllegalStateException("new_sampler() failed") + var context = 0L + var batch = 0L + var sampler = 0L + try { + context = new_context(model, nCtx) + if (context == 0L) throw IllegalStateException("new_context() failed") + + batch = new_batch(2048, 0, 1) + if (batch == 0L) throw IllegalStateException("new_batch() failed") + + sampler = new_sampler() + if (sampler == 0L) throw IllegalStateException("new_sampler() failed") + } catch (failure: Throwable) { + // Reverse of the allocation order, and the model last: it owns the rest. + if (sampler != 0L) free_sampler(sampler) + if (batch != 0L) free_batch(batch) + if (context != 0L) free_context(context) + free_model(model) + throw failure + } log.info("Loaded model {}", pathToModel) threadLocalState.set(State.Loaded(model, context, batch, sampler)) @@ -339,7 +372,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 +384,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 +400,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..0e007673 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 @@ -11,7 +11,10 @@ import com.itsaky.androidide.plugins.aiagentlocal.feedback.ModelLoadException import com.itsaky.androidide.plugins.aiagentlocal.feedback.ModelNotConfiguredException import com.itsaky.androidide.plugins.aiagentlocal.feedback.UserActionableLlmException 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.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 +36,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. @@ -287,11 +291,23 @@ class LocalLlmBackend( return // Already loaded } + // One parse of the metadata block per load, feeding both the guard below and the context + // sizing after the unload: it sits at the front of a multi-GB file, and a model switch + // used to walk it twice. + val modelFile = File(resolvedPath).takeIf { it.isFile } + val openModel = { modelFile?.inputStream() } + val (header, modelSizeBytes) = withContext(Dispatchers.IO) { + GgufHeaderReader.read(openModel) to modelFile?.length()?.takeIf { it > 0L } + } + // Guard the chat path against encoder-only embedding models. Running causal generation on // one aborts natively (SIGABRT) and takes the IDE down. Classify BEFORE unloading any // working chat model, so a wrong selection never tears down a good one. See ADFA-4388. - val kind = GgufModelInspector.classify(resolvedPath) - if (kind.isEmbeddingOnly) { + // The overload rescans for the architecture alone, and only if the parse above gave up. + val modelKind = withContext(Dispatchers.IO) { + GgufModelInspector.classify(header, openModel) + } + if (modelKind.isEmbeddingOnly) { throw IncompatibleModelException( "The selected model is an embedding model and can't be used for chat. " + "Choose a chat model in AI Settings." @@ -307,13 +323,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, header, modelSizeBytes) + context.logger.info("Loading model: $resolvedPath") try { - llama.load(resolvedPath) + llama.load(resolvedPath, contextTokens) } catch (e: CancellationException) { throw e } catch (e: Exception) { @@ -326,6 +345,64 @@ class LocalLlmBackend( modelLoaded = true currentModelPath = resolvedPath context.logger.info("Model loaded successfully") + reportEffectiveContextSize(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. + * + * @param requestedTokens the context [resolveContextSize] asked for + */ + private suspend fun reportEffectiveContextSize(requestedTokens: Int) { + val actual = try { + llama.getContextSize() + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { + context.logger.warn("Could not read the created context size: ${e.message}") + return + } + if (actual == requestedTokens) { + context.logger.info("Context size in effect: $actual tokens") + } else { + context.logger.warn( + "Context size in effect: $actual tokens, not the $requestedTokens requested;" + + " prompt-length limits follow the smaller number" + ) + } + } + + /** + * 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 + * @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 + */ + private fun resolveContextSize( + resolvedPath: String, + availableBytes: Long, + header: GgufHeader?, + modelSizeBytes: Long?, + ): Int { + val resolved = ModelContextResolver.resolve( + header = header, + availableBytes = availableBytes.takeIf { it >= 0L }, + modelSizeBytes = modelSizeBytes, + ) + // Unconditional: a wrongly sized context otherwise just reads as the assistant forgetting. + context.logger.info( + "Context size for $resolvedPath: ${resolved.contextTokens} tokens" + + " (model advertises ${resolved.advertisedTokens ?: "unknown"}," + + " ${if (availableBytes >= 0L) "$availableBytes bytes free" else "free RAM unknown"})" + ) + return resolved.contextTokens } /** 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..165f3c52 --- /dev/null +++ b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicy.kt @@ -0,0 +1,72 @@ +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 what the weights and compute buffers leave that 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 sized to the last free byte. + */ + private const val KV_BUDGET_DIVISOR = 2L + + /** + * 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 + * free RAM gets the floor — deliberate, since the alternative is sizing a cache it must fight. + * + * @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 + * @return the context to configure, always between [DEFAULT_CONTEXT_TOKENS] and + * [MAX_CONTEXT_TOKENS] inclusive + */ + fun choose(header: GgufHeader?, availableBytes: Long?, modelSizeBytes: Long?): 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. + val freeBytes = availableBytes?.takeIf { it >= 0L } ?: return DEFAULT_CONTEXT_TOKENS + val weightBytes = modelSizeBytes?.takeIf { it > 0L } ?: 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 + + // Each clamped at zero: an unclamped Long underflows on an absurd size and wraps positive. + val afterWeights = (freeBytes - weightBytes).coerceAtLeast(0L) + val spareBytes = (afterWeights - ModelMemory.RUN_BUFFER_BYTES).coerceAtLeast(0L) + val budgetBytes = spareBytes / 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..c23e0e07 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 @@ -11,6 +11,9 @@ import java.io.InputStream * The GGUF metadata the memory estimate needs. Every field is nullable because a file may omit any * key; the estimator then falls back to a heuristic instead of guessing. * + * New properties are appended rather than inserted: every one of them is `Long?`, so a positional + * construction that shifted would still compile and silently bind the wrong value to each name. + * * @property architecture `general.architecture`; also the prefix every other key here is read under * @property blockCount transformer layers, `{arch}.block_count` * @property embeddingLength model width, `{arch}.embedding_length` @@ -19,6 +22,8 @@ import java.io.InputStream * @property keyLength per-head key width, `{arch}.attention.key_length`; absent means it is the * model width divided by the head count, which is only the default and not always the truth * @property valueLength per-head value width, `{arch}.attention.value_length`; as [keyLength] + * @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 */ data class GgufHeader( val architecture: String?, @@ -28,13 +33,14 @@ data class GgufHeader( val headCountKv: Long?, val keyLength: Long? = null, val valueLength: Long? = null, + val contextLength: Long? = null, ) /** * Reads the metadata block at the front of a `.gguf` file — the shape values the KV-cache estimate - * needs. Never reads the weights, and fails closed to null: an unreadable header must mean "no - * estimate", never a wrong one. Parses the same metadata block as [GgufModelInspector], which - * answers a different question — chat model or embedding model. + * needs, plus the architecture [GgufModelInspector] classifies. Never reads the weights, and fails + * closed to null: an unreadable header must mean "no estimate", never a wrong one. The one parser + * for this block — a well-formed file is walked once, [readArchitecture] retries only after a null. */ internal object GgufHeaderReader { @@ -59,6 +65,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" @@ -101,6 +108,23 @@ internal object GgufHeaderReader { null } + /** + * The architecture and nothing else, for the crash guard when [read] already returned null. + * Stops at the first `general.architecture` — conventionally the first entry — so nothing later + * in the block can defeat it. Blocking I/O, bounded by [MAX_METADATA_BYTES]. + * + * @param openStream opens the candidate model, or returns null when it can't be opened + * @return the declared architecture, or null if the parse never reached it + */ + fun readArchitecture(openStream: () -> InputStream?): String? = try { + openStream()?.use { stream -> + val budgeted = BudgetedInputStream(stream, MAX_METADATA_BYTES) + readArchitectureOnly(DataInputStream(BufferedInputStream(budgeted, 1 shl 16))) + } + } catch (_: Throwable) { + null + } + /** * Aborts the parse once [limit] bytes have been consumed, so no declared count can make the * read run on past the metadata. Throwing is deliberate: [read] treats it like any other parse @@ -138,6 +162,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 +195,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) @@ -200,9 +228,34 @@ internal object GgufHeaderReader { headCountKv = shape?.headCountKv, keyLength = shape?.keyLength, valueLength = shape?.valueLength, + contextLength = shape?.contextLength, ) } + /** + * Deliberately laxer than [readHeader]: no entry-count ceiling, and it returns before reading + * whatever follows the architecture, so the constructs that make a full parse fail cannot make + * an embedding model look chat-capable. See ADFA-4388. + */ + private fun readArchitectureOnly(input: DataInputStream): String? { + if (readU32(input) != GGUF_MAGIC) return null + + val wide = readU32(input) >= 2 + readCount(input, wide) // tensor count, unused here + val entryCount = readCount(input, wide) + if (entryCount < 0L) return null + + var index = 0L + while (index < entryCount) { + val key = readString(input, wide) + val type = readU32(input) + if (key == KEY_ARCHITECTURE && type == T_STRING) return readString(input, wide) + skipValue(input, type, wide) + index++ + } + return null + } + /** The shape values for the architecture [key] belongs to, created on first sight. */ private fun shapeFor(shapes: HashMap, key: String, suffix: String): ArchShape = shapes.getOrPut(key.dropLast(suffix.length)) { ArchShape() } diff --git a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/GgufModelInspector.kt b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/GgufModelInspector.kt index 0e830a73..bd9f75c7 100644 --- a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/GgufModelInspector.kt +++ b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/GgufModelInspector.kt @@ -4,40 +4,25 @@ import java.io.BufferedInputStream import java.io.DataInputStream import java.io.File import java.io.FileInputStream +import java.io.InputStream /** - * Minimal GGUF header reader, just enough to tell whether a `.gguf` file is a chat/generation - * model or an embedding (encoder-only) model. + * Tells whether a `.gguf` file is a chat/generation model or an embedding (encoder-only) model, + * from the architecture [GgufHeaderReader] parsed out of its metadata block. * * WHY: the native chat path runs a causal `llama_decode`. Handed an encoder-only model (BERT * family, e.g. all-MiniLM), llama.cpp hits a `GGML_ASSERT` and calls `abort()` — a SIGABRT that * no Kotlin `try/catch` can intercept, taking the whole IDE process down. We classify the file * up front so the backend can refuse chat gracefully instead of crashing. See ADFA-4388. * - * This reads only the GGUF metadata header (`general.architecture` is almost always the first - * key), skipping over values without loading the model, so it's cheap. It deliberately - * **fails open**: any parse error or missing architecture is reported as [ModelKind.UNKNOWN] and - * treated as chat-capable, so a genuine chat model is never wrongly blocked by a header quirk. + * It deliberately **fails open**: an unreadable header or a missing architecture is reported as + * [ModelKind.UNKNOWN] and treated as chat-capable, so a genuine chat model is never wrongly + * blocked by a header quirk. */ object GgufModelInspector { private const val GGUF_MAGIC = GgufFormat.MAGIC_LE_INT - // GGUF metadata value types. - private const val T_UINT8 = 0 - private const val T_INT8 = 1 - private const val T_UINT16 = 2 - private const val T_INT16 = 3 - private const val T_UINT32 = 4 - private const val T_INT32 = 5 - private const val T_FLOAT32 = 6 - private const val T_BOOL = 7 - private const val T_STRING = 8 - private const val T_ARRAY = 9 - private const val T_UINT64 = 10 - private const val T_INT64 = 11 - private const val T_FLOAT64 = 12 - /** * Architectures that are encoder-only embedding models and cannot do causal generation. * The `contains("bert")` catch below covers the whole BERT family (bert, nomic-bert, @@ -62,105 +47,36 @@ object GgufModelInspector { false } - /** Reads [modelPath]'s GGUF header and classifies it. Never throws. */ - fun classify(modelPath: String): Result { - val arch = try { - readArchitecture(File(modelPath)) - } catch (_: Exception) { - null - } ?: return Result(ModelKind.UNKNOWN, null) + /** + * Classifies the header [GgufHeaderReader] read, so a caller that also needs the rest of it — + * the context sizing does — pays for one metadata parse rather than two. Never throws. + * + * @param header the model's metadata, or null when it could not be read + */ + fun classify(header: GgufHeader?): Result = classifyArchitecture(header?.architecture) + + /** + * As [classify], but a header with no architecture retries via + * [GgufHeaderReader.readArchitecture]. A full parse rejects far more files than an + * architecture-only scan, and each rejection is a model this guard would wave through (ADFA-4388). + * + * @param header the model's metadata, or null when it could not be read + * @param openStream reopens the same model for the fallback scan; blocking I/O + */ + fun classify(header: GgufHeader?, openStream: () -> InputStream?): Result = + header?.architecture?.let(::classifyArchitecture) + ?: classifyArchitecture(GgufHeaderReader.readArchitecture(openStream)) + private fun classifyArchitecture(architecture: String?): Result { + val arch = architecture ?: return Result(ModelKind.UNKNOWN, null) val a = arch.lowercase() val isEmbedding = a.contains("bert") || a in EMBEDDING_ARCHS return Result(if (isEmbedding) ModelKind.EMBEDDING else ModelKind.CHAT, arch) } - private fun readArchitecture(file: File): String? { - DataInputStream(BufferedInputStream(FileInputStream(file), 1 shl 16)).use { input -> - val magic = readU32(input) - if (magic != GGUF_MAGIC) return null - - val version = readU32(input) - // v1 used 32-bit counts/lengths; v2+ use 64-bit. - val wide = version >= 2 - - // tensor_count, then metadata_kv_count. - readCount(input, wide) - val kvCount = readCount(input, wide) - - for (i in 0 until kvCount) { - val key = readString(input, wide) - val valueType = readU32(input) - if (key == "general.architecture" && valueType == T_STRING) { - return readRawString(input, wide) - } - skipValue(input, valueType, wide) - } - } - return null - } - - private fun skipValue(input: DataInputStream, type: Int, wide: Boolean) { - when (type) { - T_UINT8, T_INT8, T_BOOL -> skipFully(input, 1) - T_UINT16, T_INT16 -> skipFully(input, 2) - T_UINT32, T_INT32, T_FLOAT32 -> skipFully(input, 4) - T_UINT64, T_INT64, T_FLOAT64 -> skipFully(input, 8) - T_STRING -> skipFully(input, readCount(input, wide)) - T_ARRAY -> { - val elemType = readU32(input) - val n = readCount(input, wide) - // GGUF arrays are never nested, so elements are scalars or strings. - repeat(n.toInt().coerceAtLeast(0)) { skipValue(input, elemType, wide) } - } - else -> throw IllegalStateException("Unknown GGUF value type: $type") - } - } - - // --- little-endian primitives --- - private fun readU32(input: DataInputStream): Int { val b0 = input.read(); val b1 = input.read(); val b2 = input.read(); val b3 = input.read() if (b3 < 0) throw java.io.EOFException() return (b0 and 0xFF) or ((b1 and 0xFF) shl 8) or ((b2 and 0xFF) shl 16) or ((b3 and 0xFF) shl 24) } - - private fun readU64(input: DataInputStream): Long { - var v = 0L - for (i in 0 until 8) { - val b = input.read() - if (b < 0) throw java.io.EOFException() - v = v or ((b.toLong() and 0xFF) shl (8 * i)) - } - return v - } - - /** A length/count field: 64-bit on GGUF v2+, 32-bit on v1. */ - private fun readCount(input: DataInputStream, wide: Boolean): Long = - if (wide) readU64(input) else readU32(input).toLong() and 0xFFFFFFFFL - - private fun readString(input: DataInputStream, wide: Boolean): String = readRawString(input, wide) - - private fun readRawString(input: DataInputStream, wide: Boolean): String { - val len = readCount(input, wide) - // Keys/arch strings are tiny; guard against a corrupt huge length. - if (len < 0 || len > 1 shl 20) throw IllegalStateException("Unreasonable GGUF string length: $len") - val bytes = ByteArray(len.toInt()) - input.readFully(bytes) - return String(bytes, Charsets.UTF_8) - } - - private fun skipFully(input: DataInputStream, n: Long) { - var remaining = n - while (remaining > 0) { - val skipped = input.skip(remaining) - if (skipped > 0) { - remaining -= skipped - } else { - // skip() can return 0 near buffer boundaries; fall back to a read. - if (input.read() < 0) throw java.io.EOFException() - remaining -= 1 - } - } - } } 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..b87704d1 --- /dev/null +++ b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolver.kt @@ -0,0 +1,44 @@ +package com.itsaky.androidide.plugins.aiagentlocal.model + +/** + * 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, 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. + */ +internal object ModelContextResolver { + + /** + * Fails open by construction, with no error path of its own: a null header is what + * [GgufHeaderReader.read] returns for anything it could not open or parse, and + * [ContextSizePolicy.choose] answers its default for one. + * + * @param header the model's parsed metadata, or null when it could not be read + * @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 + */ + fun resolve( + header: GgufHeader?, + availableBytes: Long?, + modelSizeBytes: Long?, + ): ModelContextSize = ModelContextSize( + contextTokens = ContextSizePolicy.choose(header, availableBytes, modelSizeBytes), + header = header, + ) +} diff --git a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelLoadDiagnostics.kt b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelLoadDiagnostics.kt index e1edb83e..2e0c663b 100644 --- a/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelLoadDiagnostics.kt +++ b/ai-agent-local/src/main/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelLoadDiagnostics.kt @@ -38,9 +38,9 @@ object ModelLoadDiagnostics { } /** - * Weights are mmap'd, so the file need not fit in free RAM — only the KV cache and compute - * buffers must be resident. The memory check therefore tests a conservative headroom, because - * overestimating it would blame a corrupt model on memory and send users chasing smaller files. + * Diagnoses a load that already failed, so it tests a conservative headroom rather than the + * file size: overestimating would blame a corrupt model on memory. [ContextSizePolicy] charges + * the mmap'd weights instead, because it sizes the cache before the load pages them in. * * @param modelPath resolved filesystem path the native loader was handed * @param availableMemoryBytes free RAM reported by the OS, or negative if unknown 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..0a3d31f0 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 @@ -5,8 +5,9 @@ import com.itsaky.androidide.plugins.aiagentlocal.model.GgufHeader /** * What a model will cost in memory, split the way it behaves at runtime. * - * @property loadBytes the weights. mmap'd, so they need not *fit*: when they don't, the device - * thrashes page cache instead of failing fast, which is the "ten minutes, then an error" report. + * @property loadBytes the weights. mmap'd, so a shortfall thrashes page cache instead of failing + * fast — the "ten minutes, then an error" report. Not a reason to treat them as free: + * [ContextSizePolicy] charges them, since it sizes the cache before the load pages them in. * @property runBytes KV cache and compute buffers. Ordinary allocations, so this part must fit. * @property fromHeader true when [runBytes] came from the model's own shape values rather than the * size-based fallback; diagnostics only. @@ -23,17 +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 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: 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). */ 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 +51,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 to price the cache at; required, because a default here + * would silently describe an allocation nobody makes * @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 +75,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..c91ed04d 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 @@ -12,6 +12,7 @@ import com.itsaky.androidide.plugins.aiagentlocal.R import com.itsaky.androidide.plugins.aiagentlocal.format.ByteSize import com.itsaky.androidide.plugins.aiagentlocal.logging.LOG_PREFIX import com.itsaky.androidide.plugins.aiagentlocal.model.ContentModelFileSource +import com.itsaky.androidide.plugins.aiagentlocal.model.ContextSizePolicy import com.itsaky.androidide.plugins.aiagentlocal.model.DeviceMemory import com.itsaky.androidide.plugins.aiagentlocal.model.GgufFileInspector import com.itsaky.androidide.plugins.aiagentlocal.model.GgufHeaderReader @@ -267,6 +268,9 @@ class LocalLlmSettingsViewModel( /** * Checks the model against free RAM and, when it looks too large, asks the user whether to go * ahead. Fails OPEN: an unreadable size or header means no warning rather than a wrong one. + * Prices the KV cache at the floor context, never at one derived from the free RAM it is then + * compared against — that would make a larger granted context the thing that trips the warning, + * and would move "needs X to run" between two selections of the same model. * * @return true to continue with this model */ @@ -276,11 +280,15 @@ class LocalLlmSettingsViewModel( context: Context ): Boolean { val modelName = fileInfo.displayName + // The floor is the least the load can use, so also the least this model can cost. + val header = GgufHeaderReader.read { modelFiles.openStream(context, uriString) } val estimate = ModelMemoryEstimator.estimate( fileSizeBytes = fileInfo.sizeBytes, - header = GgufHeaderReader.read { modelFiles.openStream(context, uriString) }, + header = header, + contextTokens = ContextSizePolicy.DEFAULT_CONTEXT_TOKENS, ) - // Read last and never cached: the user may have just closed apps to make room. + // 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. val availableBytes = deviceMemory.availableBytes() 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 new file mode 100644 index 00000000..4616357a --- /dev/null +++ b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ContextSizePolicyTest.kt @@ -0,0 +1,206 @@ +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) + + /** A model small enough that the weights term never decides these cases on its own. */ + private val modelSize = 64L * 1024 * 1024 + + /** + * Free RAM that affords exactly [tokens], undoing the weights, the reserve and the divisor. + * + * @param tokens the context the returned figure should just cover + * @param weightBytes the model size the caller will pass alongside it + */ + 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. */ + private fun choose( + header: GgufHeader?, + availableBytes: Long?, + modelSizeBytes: Long? = modelSize, + ): Int = ContextSizePolicy.choose(header, availableBytes, modelSizeBytes) + + @Test + fun givenNoHeader_whenChoosing_thenFallsBackToDefault() { + assertEquals(DEFAULT_CONTEXT_TOKENS, choose(null, ramAffording(100_000L))) + } + + @Test + fun givenNoContextLengthInHeader_whenChoosing_thenFallsBackToDefault() { + val result = choose(header(contextLength = null), ramAffording(100_000L)) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenUnreadableMemory_whenChoosing_thenFallsBackToDefault() { + assertEquals(DEFAULT_CONTEXT_TOKENS, 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 = choose(header(blockCount = null), ramAffording(100_000L)) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenLargeContextModelAndAmpleRam_whenChoosing_thenCapsAtMaximum() { + val result = choose(header(contextLength = 32768L), ramAffording(100_000L)) + assertEquals(MAX_CONTEXT_TOKENS, result) + } + + @Test + fun givenModelContextBelowFloor_whenChoosing_thenHoldsTheFloor() { + val result = choose(header(contextLength = 2048L), ramAffording(100_000L)) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenModelContextBetweenFloorAndMaximum_whenChoosing_thenUsesTheModelContext() { + val result = 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 = 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 = choose(header(), ramAffording(1_000L)) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenNoFreeMemory_whenChoosing_thenReturnsFloorRatherThanZero() { + assertEquals(DEFAULT_CONTEXT_TOKENS, 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 = choose(header(), ModelMemory.RUN_BUFFER_BYTES / 2) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenAbsurdContextLength_whenChoosing_thenCapsAtMaximumWithoutOverflow() { + val result = choose(header(contextLength = Long.MAX_VALUE), ramAffording(100_000L)) + assertEquals(MAX_CONTEXT_TOKENS, result) + } + + @Test + fun givenNegativeContextLength_whenChoosing_thenFallsBackToDefault() { + val result = choose(header(contextLength = -1L), ramAffording(100_000L)) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenAbsurdShapeValues_whenChoosing_thenFallsBackToDefaultWithoutOverflow() { + val result = choose( + header(blockCount = Long.MAX_VALUE, keyLength = Long.MAX_VALUE), + ramAffording(100_000L), + ) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenUnreadableModelSize_whenChoosing_thenFallsBackToDefault() { + val result = choose(header(), ramAffording(100_000L), modelSizeBytes = null) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenModelWeightsFillingMostOfFreeRam_whenChoosing_thenHoldsTheFloor() { + // 4.4 GB of weights in 3.5 GB free: nothing is left to spend, whatever the model advertises. + val result = choose(header(), availableBytes = 3_500L * 1024 * 1024, modelSizeBytes = 4_400L * 1024 * 1024) + assertEquals(DEFAULT_CONTEXT_TOKENS, result) + } + + @Test + fun givenAModelLargerThanFreeRam_whenChoosing_thenHoldsTheFloorRatherThanWrapping() { + // Unclamped, both subtractions underflow here and the wrapped budget picks the ceiling. + assertEquals(DEFAULT_CONTEXT_TOKENS, choose(header(), 0L, Long.MAX_VALUE)) + assertEquals(DEFAULT_CONTEXT_TOKENS, choose(header(), 1L, Long.MAX_VALUE)) + assertEquals(DEFAULT_CONTEXT_TOKENS, choose(header(), ramAffording(100_000L), Long.MAX_VALUE)) + } + + @Test + fun givenNegativeFreeMemory_whenChoosing_thenFallsBackToDefault() { + assertEquals(DEFAULT_CONTEXT_TOKENS, choose(header(), -1L)) + assertEquals(DEFAULT_CONTEXT_TOKENS, choose(header(), Long.MIN_VALUE)) + } + + @Test + fun givenGrowingModelSizes_whenChoosing_thenTheContextNeverGrows() { + // Monotonicity is what a bounds-only assertion misses: a wrap reads as more RAM, not less. + val availableBytes = ramAffording(50_000L, weightBytes = 0L) + val sizes = listOf(1L, 1L shl 20, 1L shl 30, 4L shl 30, Long.MAX_VALUE / 2, Long.MAX_VALUE) + var previous = MAX_CONTEXT_TOKENS + for (size in sizes) { + val result = choose(header(), availableBytes, size) + assertTrue("size=$size gave $result, up from $previous", result <= previous) + previous = result + } + } + + @Test + fun givenTwoModelSizesAndTheSameFreeRam_whenChoosing_thenTheLargerModelGetsLessContext() { + val availableBytes = ramAffording(12_000L, weightBytes = 0L) + val small = choose(header(), availableBytes, modelSizeBytes = 128L * 1024 * 1024) + val large = choose(header(), availableBytes, modelSizeBytes = 1_024L * 1024 * 1024) + assertTrue("weights must reduce the KV budget, got $small then $large", large < small) + } + + @Test + fun givenAnyInputs_whenChoosing_thenResultStaysWithinTheDeclaredBounds() { + val contexts = listOf(null, -1L, 0L, 512L, 4096L, 8192L, 32768L, Long.MAX_VALUE) + val memories = + listOf(null, Long.MIN_VALUE, -1L, 0L, 1L, ModelMemory.RUN_BUFFER_BYTES, ramAffording(50_000L), Long.MAX_VALUE) + val sizes = listOf(null, -1L, 0L, 1L, modelSize, Long.MAX_VALUE) + 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) + } + } + } + } +} diff --git a/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/GgufModelInspectorTest.kt b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/GgufModelInspectorTest.kt new file mode 100644 index 00000000..2b1312aa --- /dev/null +++ b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/GgufModelInspectorTest.kt @@ -0,0 +1,113 @@ +package com.itsaky.androidide.plugins.aiagentlocal.model + +import java.io.ByteArrayOutputStream +import java.io.InputStream +import org.junit.Assert.assertEquals +import org.junit.Assert.assertNull +import org.junit.Test + +/** + * Pins the encoder-only guard against the files a full header parse gives up on. Classifying an + * embedding model as [GgufModelInspector.ModelKind.UNKNOWN] lets it reach a causal `llama_decode`, + * which aborts the whole IDE process — the crash ADFA-4388 added the guard to prevent. + */ +class GgufModelInspectorTest { + + @Test + fun givenAWellFormedEmbeddingModel_whenClassifying_thenReportsEmbedding() { + val bytes = gguf(architectureEntry(EMBEDDING_ARCH)) + + assertEquals(EMBEDDING_ARCH, GgufHeaderReader.read { bytes.inputStream() }?.architecture) + assertEquals(GgufModelInspector.ModelKind.EMBEDDING, classify(bytes).kind) + } + + @Test + fun givenAnUnknownValueTypeAfterTheArchitecture_whenClassifying_thenStillReportsEmbedding() { + val bytes = gguf(architectureEntry(EMBEDDING_ARCH), unknownTypeEntry("quirk")) + + assertNull(GgufHeaderReader.read { bytes.inputStream() }) + assertEquals(GgufModelInspector.ModelKind.EMBEDDING, classify(bytes).kind) + } + + @Test + fun givenMoreEntriesThanTheParserAccepts_whenClassifying_thenStillReportsEmbedding() { + val bytes = gguf(architectureEntry(EMBEDDING_ARCH), declaredEntryCount = 5000L) + + assertNull(GgufHeaderReader.read { bytes.inputStream() }) + assertEquals(GgufModelInspector.ModelKind.EMBEDDING, classify(bytes).kind) + } + + @Test + fun givenAnUnparseableChatModel_whenClassifying_thenReportsChat() { + val bytes = gguf(architectureEntry("llama"), unknownTypeEntry("quirk")) + + assertEquals(GgufModelInspector.ModelKind.CHAT, classify(bytes).kind) + } + + @Test + fun givenNoArchitectureAtAll_whenClassifying_thenFailsOpenAsUnknown() { + val bytes = gguf(unknownTypeEntry("quirk")) + + assertEquals(GgufModelInspector.ModelKind.UNKNOWN, classify(bytes).kind) + } + + @Test + fun givenAnOpenerThatReturnsNoStream_whenClassifying_thenFailsOpenAsUnknown() { + val result = GgufModelInspector.classify(null) { null } + + assertEquals(GgufModelInspector.ModelKind.UNKNOWN, result.kind) + } + + /** Classifies the way the load path does: one full parse, then the architecture-only retry. */ + private fun classify(bytes: ByteArray): GgufModelInspector.Result { + val openStream: () -> InputStream? = { bytes.inputStream() } + return GgufModelInspector.classify(GgufHeaderReader.read(openStream), openStream) + } +} + +private const val EMBEDDING_ARCH = "nomic-bert" +private const val T_STRING = 8 +private const val UNKNOWN_VALUE_TYPE = 99 +private const val GGUF_VERSION = 3 + +/** + * @param entries the metadata key/value pairs, already encoded, in order + * @param declaredEntryCount the count to write into the header, defaulting to the truth; a larger + * one is how a corrupt file trips the parser's entry ceiling + */ +private fun gguf(vararg entries: ByteArray, declaredEntryCount: Long? = null): ByteArray = + ByteArrayOutputStream().apply { + write("GGUF".toByteArray(Charsets.US_ASCII)) + writeU32(GGUF_VERSION) + writeU64(0L) // tensor count + writeU64(declaredEntryCount ?: entries.size.toLong()) + entries.forEach { write(it) } + }.toByteArray() + +private fun architectureEntry(value: String): ByteArray = + ByteArrayOutputStream().apply { + writeString("general.architecture") + writeU32(T_STRING) + writeString(value) + }.toByteArray() + +/** A value type no parser can skip past, so it ends any parse that reaches it. */ +private fun unknownTypeEntry(key: String): ByteArray = + ByteArrayOutputStream().apply { + writeString(key) + writeU32(UNKNOWN_VALUE_TYPE) + }.toByteArray() + +private fun ByteArrayOutputStream.writeU32(value: Int) { + for (shift in 0 until 32 step 8) write((value ushr shift) and 0xFF) +} + +private fun ByteArrayOutputStream.writeU64(value: Long) { + for (shift in 0 until 64 step 8) write(((value ushr shift) and 0xFF).toInt()) +} + +private fun ByteArrayOutputStream.writeString(value: String) { + val bytes = value.toByteArray(Charsets.UTF_8) + writeU64(bytes.size.toLong()) + write(bytes) +} 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..b1ae2fd1 --- /dev/null +++ b/ai-agent-local/src/test/kotlin/com/itsaky/androidide/plugins/aiagentlocal/model/ModelContextResolverTest.kt @@ -0,0 +1,85 @@ +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 of the read-then-resolve pair the load path runs: no way of failing + * to read a header may propagate out of it, 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 = resolve { + throw IOException("permission denied") + } + assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) + assertNull(resolved.header) + } + + @Test + fun givenAnOpenerReturningNoStream_whenResolving_thenReturnsTheDefaultContext() { + val resolved = resolve { null } + assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) + assertNull(resolved.header) + } + + @Test + fun givenAStreamThatIsNotGguf_whenResolving_thenReturnsTheDefaultContext() { + val resolved = resolve { + "not a model file".byteInputStream() + } + assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) + assertNull(resolved.header) + } + + @Test + fun givenAStreamThatThrowsMidRead_whenResolving_thenReturnsTheDefaultContext() { + val resolved = resolve { ThrowingStream() } + assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) + assertNull(resolved.header) + } + + @Test + fun givenUnknownFreeMemory_whenResolving_thenReturnsTheDefaultContext() { + val resolved = ModelContextResolver.resolve(header(), null, MODEL_SIZE) + assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) + } + + @Test + fun givenUnknownModelSize_whenResolving_thenReturnsTheDefaultContext() { + val resolved = ModelContextResolver.resolve(header(), Long.MAX_VALUE, null) + assertEquals(DEFAULT_CONTEXT_TOKENS, resolved.contextTokens) + } + + /** Reads then resolves the way the load path does, with room to spare so only the header decides. */ + private fun resolve(openStream: () -> InputStream?) = + ModelContextResolver.resolve(GgufHeaderReader.read(openStream), Long.MAX_VALUE, MODEL_SIZE) + + /** A header that would earn more than the floor, so a null one is what the assertions catch. */ + private fun header() = GgufHeader( + architecture = "llama", + blockCount = 24L, + embeddingLength = 1024L, + headCount = 16L, + headCountKv = 8L, + keyLength = 64L, + valueLength = 64L, + contextLength = 32768L, + ) + + /** 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") + } + + private companion object { + const val MODEL_SIZE = 64L * 1024 * 1024 + } +}