Llama: support quantized KV cache export for Arm

[xingguo01]

Add a calibrated static int8 KV cache mode with per-channel qparams learned before export. Use standard tensor cache updates for TOSA, VGF, and Ethos-U. Wire model configuration and export validation, and cover cache replacement, calibration, and Arm lowering.

Load the standard quantized AOT library for static KVQ exports when portable QDQ boundaries require per-tensor out variants.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21574

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 1 Unclassified Failure

As of commit 7b6c114 with merge base 27d4962 ():

NEW FAILURES - The following jobs have failed:

• Apple / build-demo-ios / macos-job (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 65
• pull / test-arduino-library / run / arduino-library (gh)
  RuntimeError: Command docker exec -t a4fe2b4dae45730a408e230a0e449a207d7241a92e2dd1c66515b160167db481 /exec failed with exit code 1

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• MLX / test-mlx-muse-glimmer-build / test-mlx-muse-glimmer-build (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 2

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[zingo]

Hi @digantdesai and @rascani this might need your review :)

[digantdesai]

Nit:

Suggested change

	static_kv_cache_scale: float = 1.0 / 127.0
	static_quantize_kv_cache_scale: float = 1.0 / 127.0

[xingguo01]

Updated to static_quantize_kv_cache_scale for consistency with static_quantize_kv_cache.

[digantdesai]

can we use torchao affine quant util functions here?

[xingguo01]

I investigated torchao.quantize_affine. It exports as torch.ops.torchao.quantize_affine.default, which currently lacks the required ExecuTorch Arm/TOSA lowering and runtime support. I retained the decomposed implementation and documented the constraint.

[digantdesai]

Thanks. Wonder how well does this per channel quant performs for a random prompts to the LLM.. it would be good to include w/ and w/o static quant eval numbers.

[digantdesai]

Thanks. Good to share some w/ and w/o static quant eval for ramdom llm prompts.

[digantdesai]

LGTM, can't stamp it ATM, Github is having some issues..

[digantdesai]

LGTM, can't stamp it ATM, Github is having some issues..

[digantdesai]

LGTM, can't stamp it ATM, Github is having some issues..

[digantdesai]

curious to know w/ and w/o static quant for a random LLM. Thanks.

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 2m 19s —— View job

---

PR Review: Llama quantized KV cache export for Arm

• Read changed files
• Analyze correctness & design
• Post review feedback

Overall this is a clean, well-tested addition. The StaticQuantizedKVCache mirrors the existing QuantizedKVCache conventions, validation is layered, and the test coverage (unit replace/calibrate + TOSA/VGF/Ethos-U lowering) is solid. Below are the points worth addressing before merge.

---

Correctness

1. finalize_calibration only guards against all-zero, not per-channel zero — custom_kv_cache.py:311

if not torch.any(self.k_observed_max) or not torch.any(self.v_observed_max):
    raise RuntimeError("Static KV cache calibration observed no K/V values")

This fires only if every head-dim channel observed zero. If a single channel never sees a non-zero value, its scale becomes 0 → clamped to finfo.tiny, and any real value in that channel later quantizes to ±128→±127 (saturated). That's a graceful fallback, but silent. Consider a log.warning when any per-channel scale had to be clamped to tiny, so calibration coverage problems are visible.

2. finfo(...).tiny uses the observed-max dtype, not the scale-buffer dtype — custom_kv_cache.py:315-317

self.k_cache_scales.copy_(k_scales.clamp_min(torch.finfo(k_scales.dtype).tiny))

k_scales inherits k_observed_max's dtype (the model dtype, possibly fp16/bf16), while the clamp then gets copied into the fp32→(model dtype) scale buffer. For bf16 tiny ≈ 1.2e-38 but the eps is huge, so the clamp floor is effectively meaningless for bf16. It works out because your preserves_small_scales test only checks fp16/bf16 values well above tiny, but the intent (protect against underflow to 0) isn't actually enforced for bf16. Worth a comment noting this is a divide-by-zero guard, not an accuracy floor.

3. Full-cache dequant every decode step — custom_kv_cache.py _update_and_return_float_values

k_out = self._dequantize(self.k_cache, self.k_cache_scales, k_val.dtype)

Each update dequantizes the entire [B, S_max, H, D] cache, then overwrites the current positions with the exact float values. This is consistent with the existing QuantizedKVCache behavior, so not a blocker — but it's O(context) work per token and defeats part of the memory-bandwidth benefit of int8 storage on-device. Fine for a first landing; flagging for the eval discussion below.

Design / maintainability

4. Reuse torchao affine quant utils (echoing @digantdesai's inline comment at custom_kv_cache.py:377) — _quantize/_dequantize are hand-rolled round/clamp/cast. torchao's affine quant primitives would reduce the surface area and align with the rest of the quant stack. If there's a reason not to (e.g. avoiding the dep in this export path), a one-line comment would help.

5. from_float shape unpacking is order-dependent and easy to misread — custom_kv_cache.py

max_batch_size, n_heads, max_context_length, head_dim = kv_cache.k_cache.shape
if isinstance(kv_cache, CustomKVCache):
    max_batch_size, max_context_length, n_heads, head_dim = kv_cache.k_cache.shape

The first unpack is immediately overwritten for CustomKVCache. A branch (if isinstance(...): ... else: ...) reads more clearly and avoids the "why is n_heads/max_context_length swapped" double-take.

6. Validation is duplicated between ModelConfig.__post_init__ and _validate_static_kv_cache_args — this is not dead code (the from_args path mutates llm_config.model.* after construction, so __post_init__ never re-runs on those values), but the two copies can drift. A short comment on _validate_static_kv_cache_args explaining why it re-validates (post-construction mutation in from_args) would prevent a future reader from "simplifying" it away.

Nits

• Naming (echoing @digantdesai llm_config.py:210): the flag is static_quantize_kv_cache but the scale is static_kv_cache_scale. Consider static_quantize_kv_cache_scale for symmetry.
• export_llama_lib.py _validate_static_kv_cache_args: string concat leftovers — "...does not support " "model.local_global_attention" (two adjacent literals). Collapse to one string.

Eval numbers

@digantdesai asked several times (and it's the right ask): please include with/without static-quant eval numbers on a few random LLM prompts. Per-head-dim symmetric quant shared across heads (scales are [1,1,1,head_dim], reduced over batch/seq/heads) is a fairly coarse granularity — a perplexity or task-accuracy delta would validate that the coarse per-channel scheme holds up. This is the main open item before merge, since correctness of the plumbing is well-covered but the accuracy impact isn't yet demonstrated.

---

None of the above are hard blockers on the mechanics — the plumbing and lowering look correct. Items 1, 4, and the eval numbers are the ones I'd prioritize.
· llm-extension-arm-static-quantized-kv-cache

[digantdesai]

Try to address small things Claude finds.

[xingguo01]

@digantdesai thanks for your review. We evaluated SmolLM2 linear16a8w using nine greedy prompts to check quality by eye and 50 Wikitext prompts with 3,150 scored tokens for perplexity calculation on FVP CS-320. The context length is 64. Static KVQ produced coherent generations and perplexity 45.53, compared with the plain-KV linear16a8w baseline around 45. Though the current implementation of static KV cache is a functional but not performant path, and more work needs to be done to get it performant.

[xingguo01]

@claude here are my replies, all implemented in the latest PR.
Zero calibration channels
Addressed. Calibration now warns when any K/V channel remains zero and uses the smallest positive FP32 scale. Unit coverage was added.

Scale dtype
Addressed. Observed maxima are converted to the FP32 scale-buffer dtype before division and clamping. The comment clarifies that tiny prevents division by zero rather than defining an accuracy floor.

Full-cache dequantization
Agreed. This is intentional for the initial implementation because attention still consumes floating-point K/V. Avoiding O(context) dequantization requires backend-supported quantized attention and is follow-up work.

torchao affine utilities
Investigated, but the exported torchao affine operator does not currently have the required Arm/TOSA lowering and ExecuTorch runtime support. The decomposed implementation is therefore retained with an explanatory comment.

from_float() shape handling
Addressed with explicit CustomKVCache and standard KVCache branches.

Duplicated validation
Addressed with a comment explaining that from_args() mutates ModelConfig after post_init, requiring export-time revalidation.

Naming and string cleanup
Addressed. The scale was renamed to static_quantize_kv_cache_scale, and the adjacent string literals were collapsed.

Evaluation
Static KVQ achieved perplexity 45.53 over 50 Wikitext prompts and produced coherent output across nine greedy prompts. The plain-KV linear16a8w baseline was approximately 45.

[claude]

Claude finished @xingguo01's task in 1m 0s —— View job

---

Verification of addressed review items

I re-read the latest revision on llm-extension-arm-static-quantized-kv-cache. All the follow-ups check out:

• Zero calibration channels — custom_kv_cache.py:326-337 now computes per-channel scales, emits logging.warning when any K/V channel is all-zero, and floors with clamp_min(finfo.tiny). ✅
• Scale dtype — k_observed_max/v_observed_max are converted to the fp32 scale-buffer dtype before the divide-by-127 and clamp (custom_kv_cache.py:324-325), and the comment at line 331 correctly frames tiny as a divide-by-zero guard, not an accuracy floor. ✅
• from_float() shape handling — now explicit if isinstance(kv_cache, CustomKVCache): ... else: ... branches (custom_kv_cache.py:444-451), no more overwrite-then-unpack. ✅
• Duplicated validation — comment at export_llama_lib.py:1073 explains why it re-validates (from_args mutates ModelConfig after __post_init__). ✅
• Naming — renamed to static_quantize_kv_cache_scale consistently across llm_config.py (:210), export_llama_lib.py (arg :341, plumbing :948, :1902, :2083), and validation. ✅
• String-literal cleanup — export_llama_lib.py:1088 is now a single string. ✅

On the remaining items

• Full-cache dequant and torchao affine utils — both correctly deferred with explanatory comments (custom_kv_cache.py:364-365). Agreed these are follow-up work; the torchao op lacking Arm/TOSA lowering is a legitimate reason to keep the decomposed path.
• Eval numbers — thanks for running SmolLM2 linear16a8w: perplexity 45.53 (static KVQ) vs ~45 (plain-KV baseline) over 50 Wikitext prompts, plus coherent greedy generations. That ~0.5 delta at context length 64 with the coarse per-head-dim scheme is a reasonable result and satisfies @digantdesai's ask. Worth capturing this in the PR description or a comment near the feature so it's discoverable later.

Everything I flagged has been addressed or reasonably deferred. No further changes requested from my side — this LGTM on the mechanics. (Note: I can't submit a formal GitHub approval; @digantdesai's approval stands.)
· llm-extension-arm-static-quantized-kv-cache