From 83d4a42cb7653be81eaef681a05076d0ae2ae9df Mon Sep 17 00:00:00 2001 From: Scott Todd Date: Wed, 4 Dec 2024 10:12:52 -0800 Subject: [PATCH] Implement `__lt__` as required for `heapq.heapify`. --- .../llm/components/kvcache/trie_attention_cache.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/shortfin/python/shortfin_apps/llm/components/kvcache/trie_attention_cache.py b/shortfin/python/shortfin_apps/llm/components/kvcache/trie_attention_cache.py index fbb008005..3993e2444 100644 --- a/shortfin/python/shortfin_apps/llm/components/kvcache/trie_attention_cache.py +++ b/shortfin/python/shortfin_apps/llm/components/kvcache/trie_attention_cache.py @@ -90,6 +90,10 @@ def __eq__(self, other: object) -> bool: """Nodes are equal only if they are the same object.""" return self is other + def __lt__(self, other): + """Sort nodes by their memory address.""" + return id(self) < id(other) + class TriePagedAttentionCacheAllocation(PageAllocation): """Represents a page allocation in the trie-based cache.