Merge pull request #262 from google:mohit/fix_flash_flag

PiperOrigin-RevId: 583172562
google · Nov 16, 2023 · da44730 · da44730
2 parents 82ab361 + 4a47db8
commit da44730
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/MaxText/layers.py b/MaxText/layers.py
@@ -368,8 +368,8 @@ def wrap_flash_attention(query, key, value, decoder_segment_ids):
   def __call__(self,
                inputs_q: Array,
                inputs_kv: Array,
+               enable_flash_attention,
                decoder_segment_ids = None,
-               enable_flash_attention = True,
                inputs_positions:Optional[Array] = None,
                mask: Optional[Array] = None,
                bias: Optional[Array] = None,
@@ -1046,6 +1046,7 @@ def __call__(self,
         mesh = mesh)(
             lnx,
             lnx,
+            enable_flash_attention=cfg.enable_flash_attention,
             decoder_segment_ids=decoder_segment_ids,
             inputs_positions=decoder_positions,
             mask=decoder_mask,

diff --git a/MaxText/tests/attention_test.py b/MaxText/tests/attention_test.py
@@ -54,7 +54,7 @@ def setUp(self):
         config=self.cfg,
         mesh = self.mesh)
     self.variable = self.attention.init({'params': self.rng, 'aqt': self.rng}, jnp.ones((self.BS, self.MAX_TARGET_LENGTH, self.BASE_EMB_DIM)), 
-                            jnp.ones((self.BS, self.MAX_TARGET_LENGTH, self.BASE_EMB_DIM)), jnp.ones((self.BS, self.MAX_TARGET_LENGTH)))
+                            jnp.ones((self.BS, self.MAX_TARGET_LENGTH, self.BASE_EMB_DIM)), False)
 
 
   def get_decoder_mask(self):