add window_left option

predibase · Oct 25, 2024 · d880464 · d880464
1 parent 8540461
commit d880464
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/server/lorax_server/utils/flashinfer_attention.py b/server/lorax_server/utils/flashinfer_attention.py
@@ -108,6 +108,7 @@ def use_prefill_state(
     num_kv_heads: int,
     head_size: int,
     query_dtype: str = "float16",
+    window_left: int,
 ):
     """
     Context manager to set the active flashinfer prefill state to the given
@@ -124,6 +125,7 @@ def use_prefill_state(
             num_kv_heads=num_kv_heads,
             head_dim=head_size,
             q_data_type=query_dtype,
+            # window_left=window_left, TODO
         )
         yield
     finally: