From 7b6c11445750733fb9eb54e2d3e087a229cdc4bf Mon Sep 17 00:00:00 2001 From: Xingguo Li Date: Tue, 14 Jul 2026 15:00:29 +0100 Subject: [PATCH] Llama: support quantized KV cache export for Arm 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, export validation, 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. AI-assisted-by: Codex Signed-off-by: Xingguo Li Change-Id: Ica8dcc059df3bbdc94b426ee972a8a09f5efbff0 --- .../arm/test/modules/test_static_cache.py | 164 +++++++++++ backends/arm/test/setup_testing.sh | 1 + .../models/llama/config/test_llm_config.py | 39 +++ examples/models/llama/export_llama_lib.py | 219 +++++++++++++- .../source_transformation/custom_kv_cache.py | 272 ++++++++++++++++++ .../llama/tests/test_replace_kv_cache.py | 160 +++++++++++ extension/llm/export/config/llm_config.py | 43 +++ 7 files changed, 897 insertions(+), 1 deletion(-) diff --git a/backends/arm/test/modules/test_static_cache.py b/backends/arm/test/modules/test_static_cache.py index 0950a9a4ebf..86649f1e589 100644 --- a/backends/arm/test/modules/test_static_cache.py +++ b/backends/arm/test/modules/test_static_cache.py @@ -41,12 +41,114 @@ InputKind.USER_INPUT: 3, } +EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS = { + InputKind.BUFFER: 4, + InputKind.USER_INPUT: 3, +} + EXPECTED_OUTPUT_COUNTS = { OutputKind.BUFFER_MUTATION: 2, OutputKind.USER_OUTPUT: 2, } +DYNAMIC_KVQ_OPS = [ + "torch.ops.quantized_decomposed.choose_qparams_per_token_asymmetric.default", + "torch.ops.quantized_decomposed.quantize_per_token.default", + "torch.ops.quantized_decomposed.dequantize_per_token.default", + "torch.ops.llama.update_cache.default", + "torch.ops.llama.update_cache_with_indices.default", +] + + +def _reject_dynamic_kvq_ops(pipeline): + pipeline.add_stage_after( + "export", pipeline.tester.check_not, DYNAMIC_KVQ_OPS, suffix="dynamic_kvq_ops" + ) + + +@torch.no_grad() +class StaticQuantizedCacheModule(torch.nn.Module): + key_cache: torch.Tensor + value_cache: torch.Tensor + key_scale: torch.Tensor + value_scale: torch.Tensor + + def __init__( + self, + config: LlamaConfig, + max_cache_len: int = 10, + scale: float = 1.0 / 127.0, + ) -> None: + super().__init__() + + self.config = config + hidden_size = self.config.hidden_size + num_attention_heads = self.config.num_attention_heads + assert hidden_size is not None and num_attention_heads is not None + + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + self.head_dim = self.hidden_size // self.num_attention_heads + cache_shape = (1, self.num_attention_heads, max_cache_len, self.head_dim) + scale_shape = (1, 1, 1, self.head_dim) + + self.register_buffer("key_cache", torch.zeros(cache_shape, dtype=torch.int8)) + self.register_buffer("value_cache", torch.zeros(cache_shape, dtype=torch.int8)) + self.register_buffer( + "key_scale", torch.full(scale_shape, scale, dtype=torch.float32) + ) + self.register_buffer( + "value_scale", torch.full(scale_shape, scale, dtype=torch.float32) + ) + + # PT2E activation quantization does not create persistent int8 mutable buffers. + def _quantize(self, value: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + return torch.clamp(torch.round(value / scale), -128, 127).to(torch.int8) + + def forward( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + cache_position: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + key_q = self._quantize(key_states, self.key_scale) + value_q = self._quantize(value_states, self.value_scale) + + self.key_cache[:, :, cache_position] = key_q + self.value_cache[:, :, cache_position] = value_q + + key = self.key_cache.to(torch.float32) * self.key_scale + value = self.value_cache.to(torch.float32) * self.value_scale + key[:, :, cache_position] = key_states + value[:, :, cache_position] = value_states + + return key.clone(), value.clone() + + def get_inputs(self) -> input_t: + key_states = torch.randn( + ( + 1, + self.num_attention_heads, + 1, + self.head_dim, + ), + dtype=torch.float32, + ) + value_states = torch.randn( + ( + 1, + self.num_attention_heads, + 1, + self.head_dim, + ), + dtype=torch.float32, + ) + cache_position = torch.tensor([1], dtype=torch.int64) + + return key_states, value_states, cache_position + + @torch.no_grad() class StaticCacheModule(torch.nn.Module): def __init__( @@ -226,3 +328,65 @@ def test_static_cache_vgf_quant(test_data): ) pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) pipeline.run() + + +@common.parametrize("test_data", test_configs) +def test_static_quantized_cache_tosa_INT(test_data): + module = StaticQuantizedCacheModule(test_data).eval() + pipeline = TosaPipelineINT[input_t]( + module, module.get_inputs(), aten_op=[], exir_op=[], fold_quantize=False + ) + _reject_dynamic_kvq_ops(pipeline) + pipeline.change_args( + "check_count.exir", + {"torch.ops.higher_order.executorch_call_delegate": 2}, + ) + pipeline.count_program_io_kinds( + EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS + ) + pipeline.run() + + +@common.parametrize( + "test_data", + test_configs, + xfails={ + config: "Incorrect numerical behavior: MLBEDSW-11589" for config in test_configs + }, +) +def test_static_quantized_cache_u85_INT(test_data): + module = StaticQuantizedCacheModule(test_data).eval() + pipeline = EthosU85PipelineINT[input_t]( + module, module.get_inputs(), aten_ops=[], fold_quantize=False + ) + _reject_dynamic_kvq_ops(pipeline) + pipeline.change_args( + "check_count.exir", + {"torch.ops.higher_order.executorch_call_delegate": 2}, + ) + pipeline.tester.use_portable_ops = True + pipeline.count_program_io_kinds( + EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +@common.parametrize("test_data", test_configs) +def test_static_quantized_cache_vgf_quant(test_data): + module = StaticQuantizedCacheModule(test_data).eval() + pipeline = VgfPipeline[input_t]( + module, + module.get_inputs(), + aten_op=[], + exir_op=[], + quantize=True, + fold_quantize=False, + tosa_spec="TOSA-1.0+INT", + n_expected_delegates=2, + ) + _reject_dynamic_kvq_ops(pipeline) + pipeline.count_program_io_kinds( + EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS + ) + pipeline.run() diff --git a/backends/arm/test/setup_testing.sh b/backends/arm/test/setup_testing.sh index 7127bd029d4..12961daad45 100755 --- a/backends/arm/test/setup_testing.sh +++ b/backends/arm/test/setup_testing.sh @@ -84,6 +84,7 @@ ops_list_u85=( aten::bmm.out aten::scalar_tensor.out aten::index.Tensor_out + aten::copy_ aten::where.self_out dim_order_ops::_to_dim_order_copy.out "${ops_list_quantized_decomposed[@]}" diff --git a/examples/models/llama/config/test_llm_config.py b/examples/models/llama/config/test_llm_config.py index c5823eb3097..4bbd56cdb66 100644 --- a/examples/models/llama/config/test_llm_config.py +++ b/examples/models/llama/config/test_llm_config.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -37,6 +38,44 @@ def test_quantize_kv_without_kv(self): with self.assertRaises(ValueError): ModelConfig(quantize_kv_cache=True) + def test_static_quantize_kv_without_kv(self): + with self.assertRaises(ValueError): + ModelConfig(static_quantize_kv_cache=True) + + def test_static_quantize_kv_rejects_dynamic_shape(self): + with self.assertRaises(ValueError): + ModelConfig(use_kv_cache=True, static_quantize_kv_cache=True) + + def test_static_quantize_kv_rejects_dynamic_quantized_kv(self): + with self.assertRaises(ValueError): + ModelConfig( + use_kv_cache=True, + enable_dynamic_shape=False, + quantize_kv_cache=True, + static_quantize_kv_cache=True, + ) + + def test_static_quantize_kv_rejects_specialized_cache_modes(self): + with self.assertRaises(ValueError): + ModelConfig( + use_kv_cache=True, + enable_dynamic_shape=False, + static_quantize_kv_cache=True, + local_global_attention=[16], + ) + with self.assertRaises(ValueError): + ModelConfig( + use_kv_cache=True, + enable_dynamic_shape=False, + static_quantize_kv_cache=True, + use_attention_sink="4,2048", + ) + + def test_static_quantize_kv_rejects_non_finite_scale(self): + for scale in (float("nan"), float("inf"), float("-inf")): + with self.subTest(scale=scale), self.assertRaises(ValueError): + ModelConfig(static_quantize_kv_cache_scale=scale) + def test_local_global_attention_without_kv(self): with self.assertRaises(ValueError): ModelConfig(local_global_attention="[16]", use_kv_cache=False) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index a15268e9751..28b55b3bc58 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -13,6 +13,7 @@ import copy import json import logging +import math import re import shlex from functools import partial @@ -65,9 +66,12 @@ ) from .source_transformation.attention import replace_attention_to_attention_sha from .source_transformation.custom_kv_cache import ( + enable_static_kv_cache_calibration, + finalize_static_kv_cache_calibration, replace_kv_cache_with_custom_kv_cache, replace_kv_cache_with_quantized_kv_cache, replace_kv_cache_with_ring_kv_cache, + replace_kv_cache_with_static_quantized_kv_cache, ) from .source_transformation.quantize import ( get_quant_embedding_transform, @@ -327,6 +331,18 @@ def build_args_parser() -> argparse.ArgumentParser: action="store_true", help="Whether or not to export a model using int8 per token quantized kv cache", ) + parser.add_argument( + "--static_quantize_kv_cache", + default=False, + action="store_true", + help="Whether or not to export a model using static-qparams int8 KV cache storage", + ) + parser.add_argument( + "--static_quantize_kv_cache_scale", + type=float, + default=1.0 / 127.0, + help="Fixed symmetric per-head-dim scale for static quantized KV cache", + ) parser.add_argument( "--num_sharding", type=int, @@ -727,6 +743,115 @@ def export_llama( # noqa: C901 return filename +def _static_kvq_out_variants_registered() -> bool: + try: + _ = torch.ops.quantized_decomposed.quantize_per_tensor.out + _ = torch.ops.quantized_decomposed.dequantize_per_tensor.out + return True + except AttributeError: + return False + + +def _ensure_static_kvq_out_variants(llm_config: LlmConfig) -> None: + if _static_kvq_out_variants_registered(): + return + + if llm_config.export.so_library: + libraries = [llm_config.export.so_library] + else: + import glob + import os + + import executorch + from executorch.extension.pybindings import portable_lib # noqa # usort: skip + + libraries = sorted( + { + os.path.realpath(library) + for package_path in executorch.__path__ + for library in glob.glob( + f"{package_path}/**/*quantized_ops_aot_lib.*", recursive=True + ) + } + ) + if len(libraries) != 1: + discovered = ", ".join(libraries) if libraries else "none" + raise RuntimeError( + "Static quantized KV cache export could not select a packaged " + "quantized ops out-variant library. " + f"Discovered: {discovered}. Set export.so_library explicitly." + ) + + library = libraries[0] + logging.info("Loading quantized ops library for static KV cache: %s", library) + torch.ops.load_library(library) + if not _static_kvq_out_variants_registered(): + raise RuntimeError( + f"{library} does not register the required quantized_decomposed " + "per-tensor out variants." + ) + + +def _calibrate_static_kv_cache( + edge_manager: LLMEdgeManager, llm_config: LlmConfig +) -> None: + tasks = llm_config.quantization.calibration_tasks + if not tasks: + logging.warning( + "No calibration task provided for static KV cache; using the configured " + "fixed scale %s", + llm_config.model.static_quantize_kv_cache_scale, + ) + return + if llm_config.quantization.calibration_seq_length is None: + raise ValueError( + "Static KV cache calibration requires quantization.calibration_seq_length" + ) + if llm_config.base.tokenizer_path is None: + raise ValueError("Static KV cache calibration requires base.tokenizer_path") + + from executorch.examples.models.llama.eval_llama_lib import GraphModuleEvalWrapper + from lm_eval.evaluator import simple_evaluate + from pytorch_tokenizers import get_tokenizer + + caches = enable_static_kv_cache_calibration(edge_manager.model) + tokenizer = get_tokenizer(llm_config.base.tokenizer_path) + eval_wrapper = GraphModuleEvalWrapper( + model=edge_manager.model, + tokenizer=tokenizer, + max_seq_length=llm_config.quantization.calibration_seq_length, + use_kv_cache=True, + generate_full_logits=edge_manager.generate_full_logits, + enable_dynamic_shape=False, + device="cpu", + ) + logging.info( + "Calibrating static KV cache with tasks=%s, limit=%s, seq_length=%s", + tasks, + llm_config.quantization.calibration_limit, + llm_config.quantization.calibration_seq_length, + ) + with torch.no_grad(): + simple_evaluate( + model=eval_wrapper, + tasks=tasks, + limit=llm_config.quantization.calibration_limit, + ) + finalize_static_kv_cache_calibration(edge_manager.model) + + k_scales = torch.cat([cache.k_cache_scales.flatten() for cache in caches]) + v_scales = torch.cat([cache.v_cache_scales.flatten() for cache in caches]) + logging.info( + "Calibrated %d static KV caches: K scale range [%g, %g], " + "V scale range [%g, %g]", + len(caches), + k_scales.min().item(), + k_scales.max().item(), + v_scales.min().item(), + v_scales.max().item(), + ) + + def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager: """ Helper function for export_llama. Loads the model from checkpoint and params, @@ -777,6 +902,23 @@ def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager: # dtype_override afterward. IntxUnpackedToInt8Tensor.to() properly # propagates the dtype change to scale/zero_point/output dtype. logging.info(f"Checkpoint dtype: {edge_manager.model.checkpoint_dtype}") + if llm_config.backend.vgf.enabled and llm_config.model.static_quantize_kv_cache: + pt2e_quantize = ( + llm_config.quantization.pt2e_quantize.value + if llm_config.quantization.pt2e_quantize + else None + ) + if ( + pt2e_quantize != "vgf_16a8w" + or llm_config.backend.vgf.quantize_scope.value != "linear" + ): + raise ValueError( + "VGF static quantized KV cache is only supported with the " + "linear16a8w demo path. Set " + "quantization.pt2e_quantize=vgf_16a8w and " + "backend.vgf.quantize_scope=linear." + ) + edge_manager = edge_manager.set_output_dir(output_dir_path).source_transform( _get_source_transforms( dtype_override=dtype_override, @@ -802,8 +944,13 @@ def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager: or bool(llm_config.model.use_attention_sink), use_sdpa_with_kv_cache=llm_config.model.use_sdpa_with_kv_cache, quantize_kv_cache=llm_config.model.quantize_kv_cache, + static_quantize_kv_cache=llm_config.model.static_quantize_kv_cache, + static_quantize_kv_cache_scale=llm_config.model.static_quantize_kv_cache_scale, use_kv_cache=llm_config.model.use_kv_cache, qnn=llm_config.backend.qnn.enabled, + vgf=llm_config.backend.vgf.enabled, + tosa=llm_config.backend.tosa.enabled, + ethosu=llm_config.backend.ethosu.enabled, use_qnn_sha=llm_config.backend.qnn.use_sha, optimized_rotation_path=llm_config.backend.qnn.optimized_rotation_path, mps=llm_config.backend.mps.enabled, @@ -922,7 +1069,40 @@ def _qmode_type(value): ) +def _validate_static_kv_cache_args(llm_config): + # from_args mutates ModelConfig after __post_init__, so validate again here. + if not llm_config.model.static_quantize_kv_cache: + return + if not llm_config.model.use_kv_cache: + raise ValueError("static_quantize_kv_cache requires model.use_kv_cache=True") + if llm_config.model.quantize_kv_cache: + raise ValueError( + "Cannot enable both quantize_kv_cache and static_quantize_kv_cache" + ) + if llm_config.model.enable_dynamic_shape: + raise ValueError( + "static_quantize_kv_cache requires model.enable_dynamic_shape=False" + ) + if llm_config.model.local_global_attention: + raise ValueError( + "static_quantize_kv_cache does not support model.local_global_attention" + ) + if llm_config.model.use_attention_sink: + raise ValueError( + "static_quantize_kv_cache does not support model.use_attention_sink" + ) + if ( + not math.isfinite(llm_config.model.static_quantize_kv_cache_scale) + or llm_config.model.static_quantize_kv_cache_scale <= 0 + ): + raise ValueError( + "model.static_quantize_kv_cache_scale must be finite and positive" + ) + + def _validate_args(llm_config): + _validate_static_kv_cache_args(llm_config) + if llm_config.export.max_context_length < llm_config.export.max_seq_length: raise ValueError( f"max_context_length {llm_config.export.max_context_length} must be >= max_seq_len {llm_config.export.max_seq_length}. max_context_length impacts kv cache size that is used to remember history, while max_seq_length refers to user prompt length. Please use --max_context_length to specify context length." @@ -1438,6 +1618,8 @@ def _export_llama_multimethod(llm_config: LlmConfig) -> LLMEdgeManager: def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager: # noqa: C901 _validate_args(llm_config) + if llm_config.model.static_quantize_kv_cache: + _ensure_static_kvq_out_variants(llm_config) # Check for multimethod export if llm_config.multimethod.enabled: @@ -1450,6 +1632,12 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager: # noqa: C901 additional_passes = [] if llm_config.base.model_class.value in TORCHTUNE_DEFINED_MODELS: additional_passes = [InitializedMutableBufferPass(["kv_cache_pos"])] + if llm_config.model.use_kv_cache and ( + llm_config.backend.tosa.enabled + or llm_config.backend.vgf.enabled + or llm_config.backend.ethosu.enabled + ): + additional_passes.append(InitializedMutableBufferPass(["k_cache", "v_cache"])) # For attention sink models, cache_positions must be initialized to -1 # (sentinel for "empty slot"). Without this pass, ExecuTorch only serializes @@ -1460,6 +1648,8 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager: # noqa: C901 # export_to_edge builder_manager = _prepare_for_llama_export(llm_config) + if llm_config.model.static_quantize_kv_cache: + _calibrate_static_kv_cache(builder_manager, llm_config) if ( llm_config.backend.tosa.enabled or llm_config.backend.vgf.enabled @@ -1708,8 +1898,13 @@ def _get_source_transforms( # noqa use_custom_sdpa_with_attention_mask: bool = False, use_sdpa_with_kv_cache: bool = False, quantize_kv_cache: bool = False, + static_quantize_kv_cache: bool = False, + static_quantize_kv_cache_scale: float = 1.0 / 127.0, use_kv_cache: bool = False, qnn: bool = False, + vgf: bool = False, + tosa: bool = False, + ethosu: bool = False, use_qnn_sha: bool = False, optimized_rotation_path: Optional[str] = None, mps: bool = False, @@ -1746,8 +1941,12 @@ def _get_source_transforms( # noqa use_custom_sdpa_with_attention_mask: Whether to use custom SDPA with attention mask. use_sdpa_with_kv_cache: Whether to use SDPA with KV cache. quantize_kv_cache: Whether to quantize KV cache. + static_quantize_kv_cache: Whether to use static-qparams int8 KV cache storage. use_kv_cache: Whether to use KV cache. qnn: Whether to use QNN. + vgf: Whether to use VGF. + tosa: Whether to use TOSA. + ethosu: Whether to use Ethos-U. use_qnn_sha: Whether to use QNN SHA. optimized_rotation_path: Path to optimized rotation. mps: Whether to use MPS. @@ -1864,10 +2063,28 @@ def _get_source_transforms( # noqa if quantize_kv_cache: assert use_kv_cache, "quantize_kv_cache requires use_kv_cache=True" + if static_quantize_kv_cache: + raise ValueError( + "Cannot enable both quantize_kv_cache and static_quantize_kv_cache" + ) transforms.append(replace_kv_cache_with_quantized_kv_cache) - # Right now transforms.append(replace_sdpa_with_quantized_sdpa) + if static_quantize_kv_cache: + assert use_kv_cache, "static_quantize_kv_cache requires use_kv_cache=True" + if use_sdpa_with_kv_cache: + raise ValueError( + "Static quantized KV cache uses standard tensor cache updates; " + "disable model.use_sdpa_with_kv_cache." + ) + transforms.append( + partial( + replace_kv_cache_with_static_quantized_kv_cache, + scale=static_quantize_kv_cache_scale, + use_custom_update_cache_op=not (vgf or tosa or ethosu), + ) + ) + if use_kv_cache: if qnn: from executorch.backends.qualcomm.utils.utils import ( diff --git a/examples/models/llama/source_transformation/custom_kv_cache.py b/examples/models/llama/source_transformation/custom_kv_cache.py index b9a190a678c..dbaac9accf4 100644 --- a/examples/models/llama/source_transformation/custom_kv_cache.py +++ b/examples/models/llama/source_transformation/custom_kv_cache.py @@ -1,10 +1,12 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import logging +import math from enum import Enum from typing import Optional, Tuple @@ -253,6 +255,211 @@ def from_float( ) +class StaticQuantizedKVCache(nn.Module): + def __init__( + self, + max_batch_size, + max_context_length, + n_heads, + head_dim, + scale: float = 1.0 / 127.0, + use_custom_update_cache_op: bool = True, + return_float_values: bool = True, + dtype: torch.dtype = torch.float32, + ): + super().__init__() + + if not math.isfinite(scale) or scale <= 0: + raise ValueError("Static KV cache scale must be finite and positive") + + self.use_custom_update_cache_op = use_custom_update_cache_op + self.quantized_cache_dtype = torch.int8 + self.return_float_values = return_float_values + self.max_context_length = max_context_length + self.calibration_enabled = False + cache_shape = (max_batch_size, max_context_length, n_heads, head_dim) + scale_shape = (1, 1, 1, head_dim) + self.register_buffer( + "k_cache", + torch.zeros(cache_shape, dtype=self.quantized_cache_dtype), + persistent=False, + ) + self.register_buffer( + "v_cache", + torch.zeros(cache_shape, dtype=self.quantized_cache_dtype), + persistent=False, + ) + self.register_buffer( + "k_cache_scales", torch.full(scale_shape, scale, dtype=torch.float32) + ) + self.register_buffer( + "v_cache_scales", torch.full(scale_shape, scale, dtype=torch.float32) + ) + self.register_buffer("k_calibration_cache", None, persistent=False) + self.register_buffer("v_calibration_cache", None, persistent=False) + self.register_buffer( + "k_observed_max", + torch.zeros(scale_shape, dtype=dtype), + persistent=False, + ) + self.register_buffer( + "v_observed_max", + torch.zeros(scale_shape, dtype=dtype), + persistent=False, + ) + + def enable_calibration(self): + self.calibration_enabled = True + self.k_calibration_cache = self.k_observed_max.new_zeros(self.k_cache.shape) + self.v_calibration_cache = self.v_observed_max.new_zeros(self.v_cache.shape) + self.k_observed_max.zero_() + self.v_observed_max.zero_() + + def finalize_calibration(self): + if not self.calibration_enabled: + raise RuntimeError("Static KV cache calibration is not enabled") + 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") + + k_scales = self.k_observed_max.to(self.k_cache_scales.dtype) / 127.0 + v_scales = self.v_observed_max.to(self.v_cache_scales.dtype) / 127.0 + if torch.any(k_scales == 0) or torch.any(v_scales == 0): + logging.warning( + "Static KV cache calibration observed an all-zero K/V channel; " + "using the smallest positive scale for that channel." + ) + # This floor prevents division by zero; it is not an accuracy threshold. + self.k_cache_scales.copy_( + k_scales.clamp_min(torch.finfo(self.k_cache_scales.dtype).tiny) + ) + self.v_cache_scales.copy_( + v_scales.clamp_min(torch.finfo(self.v_cache_scales.dtype).tiny) + ) + self.calibration_enabled = False + self.k_calibration_cache = None + self.v_calibration_cache = None + self.k_cache.zero_() + self.v_cache.zero_() + + def _observe_and_update(self, input_pos, k_val, v_val): + if self.k_calibration_cache is None or self.v_calibration_cache is None: + raise RuntimeError("Static KV cache calibration is not enabled") + self.k_observed_max.copy_( + torch.maximum( + self.k_observed_max, + k_val.detach().abs().amax(dim=(0, 1, 2), keepdim=True), + ) + ) + self.v_observed_max.copy_( + torch.maximum( + self.v_observed_max, + v_val.detach().abs().amax(dim=(0, 1, 2), keepdim=True), + ) + ) + self.k_calibration_cache[:, input_pos] = k_val + self.v_calibration_cache[:, input_pos] = v_val + return self.k_calibration_cache, self.v_calibration_cache + + def _quantize(self, value, scales): + # torchao affine custom ops do not yet have the required Arm/TOSA + # lowering and ExecuTorch out-variant runtime support. + qmin = torch.iinfo(self.quantized_cache_dtype).min + qmax = torch.iinfo(self.quantized_cache_dtype).max + return torch.clamp(torch.round(value / scales), qmin, qmax).to( + self.quantized_cache_dtype + ) + + def _dequantize(self, value, scales, dtype): + return value.to(dtype) * scales.to(dtype) + + def _update_cache(self, value, cache, input_pos, indices=None): + start_pos = input_pos[0].item() + if self.use_custom_update_cache_op: + if indices is not None: + _ = torch.ops.llama.update_cache_with_indices( + value, cache, start_pos, indices + ) + else: + _ = torch.ops.llama.update_cache(value, cache, start_pos) + else: + assert indices is None, "Indices not supported for this path" + cache[:, input_pos] = value + + def _quantize_and_update(self, input_pos, k_val, v_val, indices=None): + quantized_k_val = self._quantize(k_val, self.k_cache_scales) + quantized_v_val = self._quantize(v_val, self.v_cache_scales) + + self._update_cache(quantized_k_val, self.k_cache, input_pos, indices) + self._update_cache(quantized_v_val, self.v_cache, input_pos, indices) + + def _update_and_return_float_values(self, input_pos, k_val, v_val, indices=None): + self._quantize_and_update(input_pos, k_val, v_val, indices) + + k_out = self._dequantize(self.k_cache, self.k_cache_scales, k_val.dtype) + v_out = self._dequantize(self.v_cache, self.v_cache_scales, v_val.dtype) + + self._update_cache(k_val, k_out, input_pos, indices) + self._update_cache(v_val, v_out, input_pos, indices) + + return k_out, v_out + + def _update_and_return_quantized_values( + self, input_pos, k_val, v_val, indices=None + ): + self._quantize_and_update(input_pos, k_val, v_val, indices) + + return self.k_cache, self.v_cache + + def update(self, input_pos, k_val, v_val, indices=None): + """ + k_val, v_val: [B, H, S, D] + return: [B, H, S, D] + Storage is [B, S, H, D], with static per-head-dim qparams. + """ + + k_val = k_val.transpose(1, 2) + v_val = v_val.transpose(1, 2) + + if self.calibration_enabled: + if indices is not None: + raise ValueError("Static KV calibration does not support indices") + k_out, v_out = self._observe_and_update(input_pos, k_val, v_val) + elif self.return_float_values: + k_out, v_out = self._update_and_return_float_values( + input_pos, k_val, v_val, indices + ) + else: + k_out, v_out = self._update_and_return_quantized_values( + input_pos, k_val, v_val, indices + ) + return k_out.transpose(1, 2), v_out.transpose(1, 2) + + @classmethod + def from_float( + cls, + kv_cache, + scale: float = 1.0 / 127.0, + use_custom_update_cache_op: bool = True, + ): + if isinstance(kv_cache, CustomKVCache): + max_batch_size, max_context_length, n_heads, head_dim = ( + kv_cache.k_cache.shape + ) + else: + max_batch_size, n_heads, max_context_length, head_dim = ( + kv_cache.k_cache.shape + ) + return cls( + max_batch_size, + max_context_length, + n_heads, + head_dim, + scale=scale, + use_custom_update_cache_op=use_custom_update_cache_op, + dtype=kv_cache.k_cache.dtype, + ) + + def replace_kv_cache_with_quantized_kv_cache(module): try: op = torch.ops.quantized_decomposed.quantize_per_token.out @@ -304,6 +511,71 @@ def _replace_kv_cache_with_quantized_kv_cache(module): return module +def replace_kv_cache_with_static_quantized_kv_cache( + module, scale: float = 1.0 / 127.0, use_custom_update_cache_op: bool = True +): + if use_custom_update_cache_op: + from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 + + logging.info( + "Replacing KVCache with StaticQuantizedKVCache. This modifies the model " + "in place. use_custom_update_cache_op=%s, scale=%s", + use_custom_update_cache_op, + scale, + ) + return _replace_kv_cache_with_static_quantized_kv_cache( + module, scale, use_custom_update_cache_op + ) + + +def _replace_kv_cache_with_static_quantized_kv_cache( + module, scale: float, use_custom_update_cache_op: bool +): + for name, child in module.named_children(): + if isinstance(child, KVCache) or isinstance(child, CustomKVCache): + if type(child) not in (KVCache, CustomKVCache): + raise ValueError( + "Static quantized KV cache does not support specialized " + f"cache type {type(child).__name__}" + ) + setattr( + module, + name, + StaticQuantizedKVCache.from_float( + child, + scale=scale, + use_custom_update_cache_op=use_custom_update_cache_op, + ), + ) + else: + _replace_kv_cache_with_static_quantized_kv_cache( + child, scale, use_custom_update_cache_op + ) + return module + + +def enable_static_kv_cache_calibration(module): + caches = [ + child for child in module.modules() if isinstance(child, StaticQuantizedKVCache) + ] + if not caches: + raise ValueError("No StaticQuantizedKVCache modules found for calibration") + for cache in caches: + cache.enable_calibration() + return caches + + +def finalize_static_kv_cache_calibration(module): + caches = [ + child for child in module.modules() if isinstance(child, StaticQuantizedKVCache) + ] + if not caches: + raise ValueError("No StaticQuantizedKVCache modules found for calibration") + for cache in caches: + cache.finalize_calibration() + return caches + + class CustomKVCache(nn.Module): def __init__( self, diff --git a/examples/models/llama/tests/test_replace_kv_cache.py b/examples/models/llama/tests/test_replace_kv_cache.py index 8d7171633b2..383fe6d6882 100644 --- a/examples/models/llama/tests/test_replace_kv_cache.py +++ b/examples/models/llama/tests/test_replace_kv_cache.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -7,6 +8,7 @@ import unittest from typing import List +import torch import torch.nn as nn from executorch.examples.models.llama.attention import ( @@ -25,6 +27,8 @@ replace_kv_cache_with_custom_kv_cache, replace_kv_cache_with_quantized_kv_cache, replace_kv_cache_with_ring_kv_cache, + replace_kv_cache_with_static_quantized_kv_cache, + StaticQuantizedKVCache, ) @@ -135,6 +139,162 @@ def test_replace_quantized_kv_cache_with_quantized_ring_kv_cache(self): # Verify that QuantizedKVCache has been replaced with QuantizedRingKVCache self.assertIsInstance(model.layers[0].attention.kv_cache, QuantizedRingKVCache) + def test_replace_static_quantized_kv_cache(self): + """Test replacing KVCache with static-qparams int8 KV storage.""" + attention = self._create_attention_with_kv_cache() + model = self._create_mock_model([attention]) + + replace_kv_cache_with_static_quantized_kv_cache( + model, scale=0.25, use_custom_update_cache_op=False + ) + + cache = model.layers[0].attention.kv_cache + self.assertIsInstance(cache, StaticQuantizedKVCache) + self.assertFalse(cache.use_custom_update_cache_op) + self.assertEqual(cache.k_cache.dtype, cache.quantized_cache_dtype) + self.assertEqual(cache.k_cache_scales.shape[-1], self.head_dim) + self.assertIsNone(cache.k_calibration_cache) + self.assertIsNone(cache.v_calibration_cache) + + def test_calibrate_static_quantized_kv_cache(self): + attention = self._create_attention_with_kv_cache() + model = self._create_mock_model([attention]) + replace_kv_cache_with_static_quantized_kv_cache( + model, scale=0.25, use_custom_update_cache_op=False + ) + cache = model.layers[0].attention.kv_cache + self.assertIsNone(cache.k_calibration_cache) + self.assertIsNone(cache.v_calibration_cache) + cache.enable_calibration() + self.assertEqual(cache.k_calibration_cache.shape, cache.k_cache.shape) + self.assertEqual(cache.v_calibration_cache.shape, cache.v_cache.shape) + + input_pos = torch.tensor([0, 1]) + shape = (self.batch_size, self.n_kv_heads, 2, self.head_dim) + k_val = torch.arange(torch.tensor(shape).prod(), dtype=torch.float32).reshape( + shape + ) + v_val = -2.0 * k_val + k_out, v_out = cache.update(input_pos, k_val, v_val) + + torch.testing.assert_close(k_out[:, :, input_pos], k_val) + torch.testing.assert_close(v_out[:, :, input_pos], v_val) + expected_k_scales = k_val.abs().amax(dim=(0, 1, 2)) / 127.0 + expected_v_scales = v_val.abs().amax(dim=(0, 1, 2)) / 127.0 + + cache.finalize_calibration() + + self.assertFalse(cache.calibration_enabled) + self.assertIsNone(cache.k_calibration_cache) + self.assertIsNone(cache.v_calibration_cache) + torch.testing.assert_close(cache.k_cache_scales.flatten(), expected_k_scales) + torch.testing.assert_close(cache.v_cache_scales.flatten(), expected_v_scales) + self.assertEqual(torch.count_nonzero(cache.k_cache), 0) + self.assertEqual(torch.count_nonzero(cache.v_cache), 0) + + def test_static_quantized_kv_cache_warns_for_zero_channel(self): + attention = self._create_attention_with_kv_cache() + model = self._create_mock_model([attention]) + replace_kv_cache_with_static_quantized_kv_cache( + model, scale=0.25, use_custom_update_cache_op=False + ) + cache = model.layers[0].attention.kv_cache + cache.enable_calibration() + + input_pos = torch.tensor([0]) + shape = (self.batch_size, self.n_kv_heads, 1, self.head_dim) + k_val = torch.ones(shape, dtype=torch.float32) + v_val = torch.ones(shape, dtype=torch.float32) + k_val[..., 0] = 0 + + cache.update(input_pos, k_val, v_val) + with self.assertLogs(level="WARNING") as logs: + cache.finalize_calibration() + + self.assertIn("all-zero K/V channel", logs.output[0]) + self.assertEqual( + cache.k_cache_scales[..., 0].item(), torch.finfo(torch.float32).tiny + ) + + def test_static_quantized_kv_cache_preserves_small_scales(self): + for dtype in (torch.float16, torch.bfloat16): + with self.subTest(dtype=dtype): + attention = self._create_attention_with_kv_cache() + model = self._create_mock_model([attention]) + replace_kv_cache_with_static_quantized_kv_cache( + model, scale=0.25, use_custom_update_cache_op=False + ) + model.to(dtype=dtype) + + cache = model.layers[0].attention.kv_cache + cache.enable_calibration() + input_pos = torch.tensor([0]) + shape = (self.batch_size, self.n_kv_heads, 1, self.head_dim) + k_val = torch.full(shape, 0.0625, dtype=dtype) + v_val = torch.full(shape, -0.03125, dtype=dtype) + cache.update(input_pos, k_val, v_val) + cache.finalize_calibration() + + expected_k_scale = (k_val.abs().amax() / 127.0).to( + cache.k_cache_scales.dtype + ) + expected_v_scale = (v_val.abs().amax() / 127.0).to( + cache.v_cache_scales.dtype + ) + torch.testing.assert_close( + cache.k_cache_scales, + torch.full_like(cache.k_cache_scales, expected_k_scale), + ) + torch.testing.assert_close( + cache.v_cache_scales, + torch.full_like(cache.v_cache_scales, expected_v_scale), + ) + self.assertLess( + cache.k_cache_scales.max().item(), torch.finfo(dtype).eps + ) + + def test_static_quantized_kv_cache_preserves_model_dtype(self): + for dtype in (torch.float16, torch.bfloat16): + for cast_before_replacement in (True, False): + with self.subTest( + dtype=dtype, cast_before_replacement=cast_before_replacement + ): + attention = self._create_attention_with_kv_cache() + model = self._create_mock_model([attention]) + if cast_before_replacement: + model.to(dtype=dtype) + replace_kv_cache_with_static_quantized_kv_cache( + model, scale=0.25, use_custom_update_cache_op=False + ) + if not cast_before_replacement: + model.to(dtype=dtype) + + cache = model.layers[0].attention.kv_cache + input_pos = torch.tensor([0, 1]) + shape = (self.batch_size, self.n_kv_heads, 2, self.head_dim) + k_val = torch.randn(shape, dtype=dtype) + v_val = torch.randn(shape, dtype=dtype) + k_out, v_out = cache.update(input_pos, k_val, v_val) + + self.assertEqual(k_out.dtype, dtype) + self.assertEqual(v_out.dtype, dtype) + + def test_static_quantized_kv_cache_rejects_specialized_cache(self): + attention = self._create_attention_with_kv_cache() + attention.kv_cache = RingKVCache( + self.batch_size, + self.max_context_len, + self.n_kv_heads, + self.head_dim, + self.enable_dynamic_shape, + ) + model = self._create_mock_model([attention]) + + with self.assertRaisesRegex(ValueError, "RingKVCache"): + replace_kv_cache_with_static_quantized_kv_cache( + model, use_custom_update_cache_op=False + ) + def test_multiple_layers_with_different_window_sizes(self): """Test replacing KV caches in multiple layers with different window sizes.""" # Create a model with multiple layers diff --git a/extension/llm/export/config/llm_config.py b/extension/llm/export/config/llm_config.py index f1f93ba28c6..8e04f359935 100644 --- a/extension/llm/export/config/llm_config.py +++ b/extension/llm/export/config/llm_config.py @@ -19,6 +19,7 @@ """ import argparse +import math import re from dataclasses import dataclass, field from enum import Enum @@ -188,6 +189,7 @@ class ModelConfig: input_prune_map: Path to the output pruning token mapping file (token_map.json). use_kv_cache: Whether to use KV cache. quantize_kv_cache: Whether to perform int8 per token quantization on the KV cache. + static_quantize_kv_cache: Whether to use static-qparams int8 KV cache storage. local_global_attention: List of integers specifying local and global attention pattern. e.g., [0, 16, 0, 16] to specify that every other layer is sliding window of 16. [0, 16, 32] pattern specifies 2nd and 3rd layers have sliding windows of 16 and 32. @@ -204,6 +206,8 @@ class ModelConfig: input_prune_map: Optional[str] = None use_kv_cache: bool = False quantize_kv_cache: bool = False + static_quantize_kv_cache: bool = False + static_quantize_kv_cache_scale: float = 1.0 / 127.0 local_global_attention: Optional[List[int]] = None # Replace eager MOEFeedForward modules with the # `llama::quantized_moe_ffn` portable-runtime custom op. @@ -217,6 +221,39 @@ def __post_init__(self): "Cannot quantize the KV cache (quantize_kv_cache) without enabling the KV cache (use_kv_cache)" ) + if self.static_quantize_kv_cache and not self.use_kv_cache: + raise ValueError( + "Cannot statically quantize the KV cache (static_quantize_kv_cache) without enabling the KV cache (use_kv_cache)" + ) + + if self.quantize_kv_cache and self.static_quantize_kv_cache: + raise ValueError( + "Cannot enable both dynamic quantized KV cache (quantize_kv_cache) and static quantized KV cache (static_quantize_kv_cache)" + ) + + if self.static_quantize_kv_cache and self.enable_dynamic_shape: + raise ValueError( + "static_quantize_kv_cache requires static export shapes (enable_dynamic_shape=False)" + ) + + if self.static_quantize_kv_cache and self.local_global_attention: + raise ValueError( + "static_quantize_kv_cache does not support local_global_attention" + ) + + if self.static_quantize_kv_cache and self.use_attention_sink: + raise ValueError( + "static_quantize_kv_cache does not support use_attention_sink" + ) + + if ( + not math.isfinite(self.static_quantize_kv_cache_scale) + or self.static_quantize_kv_cache_scale <= 0 + ): + raise ValueError( + "static_quantize_kv_cache_scale must be finite and positive" + ) + if self.local_global_attention and not self.use_kv_cache: raise ValueError( "Cannot use local_global_attention without enabling the KV cache (use_kv_cache)" @@ -729,6 +766,12 @@ def from_args(cls, args: argparse.Namespace) -> "LlmConfig": # noqa: C901 llm_config.model.use_kv_cache = args.use_kv_cache if hasattr(args, "quantize_kv_cache"): llm_config.model.quantize_kv_cache = args.quantize_kv_cache + if hasattr(args, "static_quantize_kv_cache"): + llm_config.model.static_quantize_kv_cache = args.static_quantize_kv_cache + if hasattr(args, "static_quantize_kv_cache_scale"): + llm_config.model.static_quantize_kv_cache_scale = ( + args.static_quantize_kv_cache_scale + ) if hasattr(args, "local_global_attention"): llm_config.model.local_global_attention = args.local_global_attention if hasattr(args, "use_moe_quantized_op"):