fix: max_past default value must be -1, not 0 (#1348)

huggingface · Dec 15, 2023 · 37555cf · 37555cf
1 parent 9b78a6e
commit 37555cf
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py b/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py
@@ -149,7 +149,7 @@ def __init__(
     ):
         super().__init__()
         self.max_past = (
-            config.sliding_window if config.sliding_window is not None else 0
+            config.sliding_window if config.sliding_window is not None else -1
         )
         self.num_heads = config.num_attention_heads
         self.hidden_size = config.hidden_size

diff --git a/server/text_generation_server/models/custom_modeling/flash_mixtral_modeling.py b/server/text_generation_server/models/custom_modeling/flash_mixtral_modeling.py
@@ -204,7 +204,7 @@ def __init__(
     ):
         super().__init__()
         self.max_past = (
-            config.sliding_window if config.sliding_window is not None else 0
+            config.sliding_window if config.sliding_window is not None else -1
         )
         self.num_heads = config.num_attention_heads
         self.hidden_size = config.hidden_size

diff --git a/server/text_generation_server/utils/flash_attn.py b/server/text_generation_server/utils/flash_attn.py
@@ -72,6 +72,9 @@ def attention(
     softmax_scale,
     window_size_left=-1,
 ):
+    if window_size_left <= 0 and window_size_left != -1:
+        raise ValueError("`window_size_left` must be > 0 or -1")
+
     if HAS_FLASH_ATTN_V2_CUDA:
         return flash_attn_2_cuda.varlen_fwd(
             q,