Fix --attn_cache_tokens default (#392)

10 months ago · b6b3ae964f
parent d49d9ad0cf
commit b6b3ae964f
1 changed files with 1 additions and 1 deletions
--- a/src/petals/server/server.py
+++ b/src/petals/server/server.py
@ -186,7 +186,7 @@ class Server:

        # For attention cache in GPU or RAM
        if attn_cache_tokens is None:
-            attn_cache_tokens = 32768 if is_multiquery_attn else 2048
+            attn_cache_tokens = 32768 if is_multiquery_attn else 8192
        cache_values_per_block = 2 * self.block_config.hidden_size * attn_cache_tokens
        cache_values_per_block //= self.block_config.num_key_value_groups
        self._cache_bytes_per_block = cache_values_per_block * torch.finfo(self.torch_dtype).bits // 8