clean AttentionData to dataclass

wnma3mz · Feb 6, 2025 · 68e67b6 · 68e67b6
1 parent 05235aa
commit 68e67b6
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 11 deletions.
diff --git a/tllm/commons/cache.py b/tllm/commons/cache.py
@@ -295,14 +295,8 @@ def update_tinygrad(self, key_states, value_states, uuid_list, layer_idx):
         return key_lst[0].cat(*key_lst[1:], dim=0), value_lst[0].cat(*value_lst[1:], dim=0)
 
 
+@dataclass
 class AttentionData:
-    def __init__(self, uuid_list: List[str], request_cache: RequestsCache, attn_mask: MIX_TENSOR) -> None:
-        self.uuid_list = uuid_list
-        self.request_cache = request_cache
-        self.attn_mask = attn_mask
-
-    def get_decoder_cache(self, uuid: str) -> DecoderCache:
-        return self.request_cache.get_decoder_cache(uuid)
-
-    def get_kv_len(self, uuid: str) -> int:
-        return self.request_cache.get_kv_len(uuid)
+    uuid_list: List[str]
+    request_cache: RequestsCache
+    attn_mask: MIX_TENSOR
diff --git a/tllm/commons/cache_manager.py b/tllm/commons/cache_manager.py
@@ -19,7 +19,7 @@ def contains(self, key) -> bool:
 
     def update_cache(self, seq_input):
         for uuid in seq_input.uuid_list:
-            self.cache.set(uuid, self.attn_data.get_decoder_cache(uuid))
+            self.cache.set(uuid, self.attn_data.request_cache.get_decoder_cache(uuid))
             self.cache.check_alive()
 
         if self.request_cache is not None: