Option to Disable kv_cache for Generative Models in ai_edge_torch

[Hardik-Choraria]

Description of the bug:

While the current ai_edge_torch.generative pipeline supports KV cache layout and update strategies (INPLACE, PREPEND_LEFT), there’s no explicit high-level option to disable the use of KV cache entirely during model export.

Actual vs expected behavior:

No response

Any other information you'd like to share?

No response

[gaikwadrahul8]

Hi, @Hardik-Choraria
You're right that the ai_edge_torch.generative pipeline doesn't have a direct high-level option to disable KV cache during model export. If I'm not wrong this is because the pipeline is built assuming the models will use KV cache for efficient on-device performance. lack of a high-level switch to disable KV cache. Generative API is designed with the assumption to use a KV cache. Disabling KV cache will require re-authoring the model and removing the KV cache mechanism. Please refer these official documentation [Ref-1], [Ref-2] and this file contains the conversion logic. The convert_to_tflite function processes the PyTorch model and sets up the signatures for prefill and decode which are designed to utilize KV cache. There's no flag to disable this behavior

If I have missed something here please let me know. Thank you for your cooperation and understanding

[github-actions]

Marking this issue as stale since it has been open for 7 days with no activity. This issue will be closed if no further activity occurs.

[hayhan]

To provide a high level option for generating a no kv cached model is helpful for the low mem devices.

[xwin]

Disabling KV cache is also helpful when using models with NPU accelerators. The accelerators do not support KV cache directly but can support LLM models. Please consider adding option to disable KV cache.