From 50c43a3fa659683d4515ef3e8f4624c8471d7538 Mon Sep 17 00:00:00 2001 From: "T.J. Alumbaugh" Date: Fri, 13 Dec 2024 13:34:24 -0800 Subject: [PATCH] Enable DUS-based KV Cache update for OSS PiperOrigin-RevId: 705986742 --- ai_edge_torch/generative/layers/kv_cache.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ai_edge_torch/generative/layers/kv_cache.py b/ai_edge_torch/generative/layers/kv_cache.py index 50433a98..66f906c4 100644 --- a/ai_edge_torch/generative/layers/kv_cache.py +++ b/ai_edge_torch/generative/layers/kv_cache.py @@ -20,6 +20,7 @@ from ai_edge_torch import hlfb from ai_edge_torch.generative.layers import model_config +from ai_edge_torch.generative.utilities.dynamic_update_slice import dynamic_update_slice import torch import torch.utils._pytree as pytree @@ -159,8 +160,6 @@ def update( Returns: KVCacheEntry: The updated KVCache entry based on the passed inputs. """ - # Turn dynamic_update_slice updates off for now. - use_dus=False update_kv_cache = _update_kv_impl if use_dus else _update_kv_base_impl return update_kv_cache(cache, input_pos, k_slice, v_slice)