Initialize v_cache to avoid NaNs (#12)

2023-08-23 14:23:59 -07:00 · 2023-08-23 14:23:59 -07:00 · 072f267cc3
parent 2fda8fe812
commit 072f267cc3
1 changed files with 2 additions and 1 deletions
--- a/server/text_generation_server/models/flash_causal_lm.py
+++ b/server/text_generation_server/models/flash_causal_lm.py
@ -53,8 +53,9 @@ class CacheManager:
                    dtype=dtype,
                    device=device,
                ),
-                torch.empty(
+                torch.full(
                    (num_blocks, num_heads, head_size, self.block_size),
+                    self.v_cache_initial_value,
                    dtype=dtype,
                    device=device,
                ),