Unverified Commit 19ba16aa authored by Yuanhang Sun's avatar Yuanhang Sun Committed by GitHub
Browse files

[Fix]: add missing device attribute to ChunkCache (#11493)

parent a2b3d9b9
...@@ -27,6 +27,10 @@ class ChunkCache(BasePrefixCache): ...@@ -27,6 +27,10 @@ class ChunkCache(BasePrefixCache):
self.req_to_token_pool = req_to_token_pool self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.page_size = page_size self.page_size = page_size
if self.token_to_kv_pool_allocator:
self.device = self.token_to_kv_pool_allocator.device
else:
self.device = torch.device("cpu")
# NOTE (csy): this is to determine if a cache has prefix matching feature. # NOTE (csy): this is to determine if a cache has prefix matching feature.
# Chunk cache always return True to indicate no prefix matching. # Chunk cache always return True to indicate no prefix matching.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment