comment

lucidrains · Jul 27, 2022 · 665966e · 665966e
1 parent e429d05
commit 665966e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flash_attention_jax/cosine_sim_flash_attention.py b/flash_attention_jax/cosine_sim_flash_attention.py
@@ -12,7 +12,7 @@
 
 Q_CHUNK_SIZE = 1024
 K_CHUNK_SIZE = 1024
-COSINE_SIM_SCALE = 16
+COSINE_SIM_SCALE = 16 # this may need to be a function of log(sequence length), but 16 was sufficient for 2048 and 4096 in my tests
 
 # flash attention