ADFA-5188 | Enable KV cache quantization and flash attention

[jatezzz]

Description

This PR implements KV cache quantization (q8_0) and enables flash attention in the llama.cpp context parameters to drastically improve memory efficiency and generation speed.

• What: Configures the context parameters type_k and type_v to use GGML_TYPE_Q8_0 and enables flash attention (LLAMA_FLASH_ATTN_TYPE_AUTO).
• How: These settings are exposed through the configure-before-load pattern. If the native context creation fails (e.g., the model's head width is incompatible with the quantized block size), it gracefully falls back to the previous defaults: f16 KV cache and flash attention disabled.
• Why: Storing the KV cache as q8_0 halves its byte size, allowing for much longer conversations before context is dropped and drastically reducing mid-generation crashes on 4–6 GB RAM devices. Flash attention mitigates generation latency on longer contexts.

Details

Logs confirming n_ctx initialization, cache type used (q8_0 vs f16), and fallback activations.

Flash Attention Enabled

• Model: Qwen3 0.6B Q8_0
• Device RAM: ≥8 GB

Q8_0 KV cache

• Model: Llama 3.2 1B Instruct
• Device RAM: ≤3 GB
• Prompt: ~10K Tokens

Ticket

ADFA-5188

Observation

This implementation works in tandem with dynamic n_ctx sizing and should be validated alongside ADFA-5187, as the memory measurement relies on both features working concurrently.

#75 Needs to be merged first

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[hal-eisen-adfa]

Four things worth a look before this merges — details inline. In short: the native fallback allocates about the same number of bytes as the attempt that just failed, a failed new_context leaks the whole model, the KV budget doesn't account for the weights that get mmap'd right after, and there's a 16 KB page-size ABI change riding along unannounced.

[itsaky-adfa]

Re-reviewed at 46d69ec against the stacked base fix/ADFA-5187-dynamic-n-ctx. The KV quantization arithmetic is sound and I verified it end to end; the findings are concentrated in the native fallback ladder, which is the part nothing exercises.

Re-check of the previous round (4 findings from @hal-eisen-adfa, all four threads currently marked resolved):

• Fallback retries at essentially the same allocation size (llama-android.cpp) - addressed, but incompletely, and the gap is the one IMPORTANT finding in this round. Replied in the thread and reopened it; f570d09 added the 4096-floor retry, but its guard still fires an attempt that needs more memory than the one that just failed in the common case.
• Failed new_context leaks the whole model - fixed, verified at head, not on the reply: LLamaAndroid.kt:286-293 frees sampler, batch and context in reverse order and then free_model(model) before rethrowing, covering the new_batch and new_sampler branches too.
• KV budget ignores the weights mmap'd right after - fixed, verified at head: ContextSizePolicy.kt:82-84 now subtracts weightBytes before RUN_BUFFER_BYTES and the divisor, and choose returns the floor outright when modelSizeBytes is null.
• 16 KB page-size ABI change riding along unannounced - half fixed. The requested assertion landed and genuinely pins something: I extracted both AARs and origin/main has p_align = 0x1000 on five of six libraries (libomp.so was already 0x4000) against 0x4000 on all six here, so the new test fails against the pre-fix artifact rather than passing vacuously. The description half was not done; replied in that thread and reopened it.

Verification run for this round

• ../gradlew testDebugUnitTest at head: 98 tests, 0 failures, including ContextSizePolicyTest 27, KvCacheTypeTest 10, ModelContextResolverTest 7, ModelMemoryEstimatorTest 5, PrebuiltAarAbiTest 5.
• The committed AAR is genuinely rebuilt from the new .cpp, not a stale copy - libllama-android.so carries the new log strings (kv cache = %s, retrying at f16 with flash attention off, retrying at the n_ctx). Worth stating because a stale AAR would not have failed anything: JNI resolves new_context by name alone, so a 2-arg native function would have silently ignored the two new arguments.
• The new ELF offsets in PrebuiltAarAbiTest are correct for ELF64 (e_phoff 32, e_phentsize 54, e_phnum 56, p_align 48, magic 0x464C457F LE), and the probe window is ample - the program-header table here ends at byte 624 of 4096.
• The q8_0 arithmetic in ContextSizePolicyTest:158 checks out by hand: the same budget gives 5000 -> 4864 tokens at f16 and 9411 -> 9216 at q8_0.
• Overflow reasoning in kvBytesPerToken still holds after the refactor: at the declared ceilings, elements peak at 2^43, so the q8_0 intermediate elements * 34 reaches ~2^48 and the f16 product ~2^58 - both inside a Long, and f16 is indeed the dearer of the two as the doc claims.

Not a finding, recording it because it was raised and is wrong: a parallel pass flagged KvCacheType.supports deriving head width from embeddingLength / headCount as a SIGABRT risk when the guess is wrong. It is not. llama.cpp re-checks against its own hparams.n_embd_head_k/n_embd_head_v and returns nullptr (llama-context.cpp:2311-2327), which is the whole reason AUTO is used over ENABLED. A wrong Kotlin guess costs a failed context creation and a retry, exactly as the comment at :420 claims.

Device verification is still owed. Per the repo's CLAUDE.md, assemblePlugin succeeding is not verification for these plugins, and I have not run this on a device. Two things need it specifically: the 16 KB page-size change on a 16 KB-page arm64 device, and any rung of the fallback ladder at all.

Severity key: CRITICAL / IMPORTANT / MINOR / NITPICK, one per comment. This repo has no written approve/request-changes rule - CLAUDE.md covers conventions and verification but sets no merge gate, and there is no REVIEW.md or CONTRIBUTING.md - so the default applied: any confirmed IMPORTANT blocks.

[itsaky-adfa]

Requesting changes on one confirmed IMPORTANT, detailed in the reopened llama-android.cpp thread.

The 4096-floor retry that f570d09 added to answer the previous round does not shrink the allocation in the case the previous round was about. Its guard fires on flash_attn_type != DISABLED alone, so whenever the policy already returned its floor - and ContextSizePolicy.choose never returns below 4096, the same value as native DEFAULT_N_CTX - the retry re-requests the identical n_ctx and identical f16 cache with flash attention turned off, which adds the [n_kv x n_ubatch x n_head] KQ buffer the FA path never allocates (llama-graph.cpp:1287 vs :1329). On the 3 GB device from that thread it asks for more than the attempt that just OOMd.

To clear this: gate the retry on floor_ctx < current_ctx alone, and drop flash attention only where a quantized cache forced it.

Everything else is MINOR or NITPICK and none of it blocks. The KV quantization arithmetic is correct - I re-derived the per-token costs, the overflow ceilings and the two exact-value assertions by hand, and the 98 unit tests pass at head. The three other findings from the previous round are genuinely fixed at head, not just claimed fixed.

Two things still owed before QA, both noted inline: device verification of the 16 KB page-size change on a 16 KB-page device, and any exercise at all of the fallback ladder - no test and neither model in the description can reach it.

[itsaky-adfa]

Re-reviewed at 7d51871 against the stacked base fix/ADFA-5187-dynamic-n-ctx. The IMPORTANT from the last round is fixed, so nothing here blocks: two MINORs, one NITPICK, and two reopened threads that are both asks about the PR description rather than the code.

plugin-examples has no written approve/request-changes rule, so the default applied and MINOR does not block. One thing to flag: my earlier CHANGES_REQUESTED still stands on this PR and a COMMENT does not lift it. Say the word and I will approve, since the finding it was raised for is resolved.

Previous rounds, verified at head rather than from the replies

Round 1 (@hal-eisen-adfa):

• Fallback retried at essentially the same allocation size - FIXED. The floor rung is gated on floor_ctx < current_ctx alone now, so it fires whenever the current context is above 4096. Worked the roomy case through with the test header: 14080 q8_0 (368 MB) -> 7424 f16 (365 MB) -> 4096 f16 (201 MB). The ladder finally has a rung that shrinks.
• Failed new_context leaked the whole llama_model - FIXED. The catch frees sampler, batch and context, then free_model(model), then rethrows.
• KV budget took no account of the model weights - FIXED. ContextSizePolicy.choose subtracts weightBytes and RUN_BUFFER_BYTES before halving.
• 16 KB page-size ABI migration unpinned - PARTLY FIXED, thread reopened. The test is done and genuinely pins; the body line is not.

Round 2 (mine):

• Floor retry still asked for more memory than the attempt that just failed (IMPORTANT) - FIXED, as above.
• Fallback turned flash attention off - CODE FIXED, thread reopened for the description half only.
• Neither native rung is exercised by a test - NOT FIXED. Deferral accepted with a reason; details in the thread. Worth a follow-up ticket, not a blocker here.
• Pre-flight estimate was no longer an upper bound - FIXED. estimateForSelection hardcodes KvCacheType.F16, and givenASelectedModel_whenEstimating_thenTheCacheIsPricedAtTheFloorContext asserts against F16_PER_TOKEN, so a revert to chooseKvCache fails the suite.
• fallbackNCtx defaulted to nCtx - FIXED. The default is 0, which clamp_context maps to DEFAULT_N_CTX.
• Success line logged jn_ctx instead of the clamped ask - FIXED, logs requested_ctx.
• Unrelated churn in the leak-cleanup hunk - FIXED, reverted; the hunk is now the signature plus the new_context call.
• Redundant second choose call - FIXED, structural on kvType == F16 now.

What I checked independently

• The committed AAR really was rebuilt from the C++ at head. Every added format string (flash attention = %s, retrying at f16 with n_ctx %d, context: %s n_ctx %d exceeds) is in libllama-android.so at 7d51871 and absent at 46d69ec; every removed one is gone. Only that library changed - the other five are byte-identical, consistent with the llama.cpp submodule pointer not moving.
• The fallback design holds at the llama.cpp level. libllama.so carries both quantized V cache was requested, but this requires Flash Attention and %s: failed to initialize the context: %s, so a refusal reaches the retry as a null return rather than propagating out of llama_init_from_model.
• Nothing on the Kotlin side budgets prompts against the requested context. reportEffectiveContextSize reads the native n_ctx back and warns when it differs, and no other caller stores contextTokens - so a shortened f16 fallback cannot desync prompt limits.
• The new tests are not vacuous. For the test header q8_0 costs 26112 B/token against f16's 49152, and the resolver returns 14080 vs 7424, so fallbackContextTokens < contextTokens pins a real gap. assertNotEquals and atQ8 are both still live assertions.

Coverage of this review

No device run, so per this repo's own rule I am not calling the change verified: the native ladder has still never executed here, and both models in the description take the happy path (head_dim 128 and 64 both divide 32). The /code-review pass I kicked off resolved "76" against the wrong repository - it returned findings about deploy-to-firestore-app-distribution.yml and gradle.properties, which are not in this PR - so its output was discarded in full and every finding above is hand-derived from the diff at 7d51871.

This cannot merge before #75; the base is still that branch, and it is currently level with main.

[itsaky-adfa]

Approving - this lifts my earlier CHANGES_REQUESTED. The floor retry no longer fires an attempt larger than the one that just failed, and I verified the fix at 7d51871 rather than from the reply: the guard is floor_ctx < current_ctx alone, and the roomy case now steps 14080 q8_0 (368 MB) -> 7424 f16 (365 MB) -> 4096 f16 (201 MB), so the ladder finally has a rung that shrinks.

Nothing outstanding blocks the merge. Before it goes in, two things are just description edits and cost nothing:

• Add the 16 KB page-size line to the body (the reopened build.gradle.kts thread) - the text you drafted is fine as-is.
• Correct the body's "falls back to f16 KV cache and flash attention disabled", which this PR deliberately stopped doing, and the framing of it as enabling flash attention when llama_context_default_params() already set AUTO.

The two MINORs and the NITPICK on this round are safe to merge as they stand; the stale LocalLlmSettingsViewModel comment has a one-click suggestion. The untested native rungs stay a known residual - better as a follow-up ticket than a blocker here.

Note this still cannot merge before #75.