Fix the performance regression with ragged attention on for llama2 7b model.

Author: wang2yn84

jetstream_pt/environment.py

@@ -124,9 +124,10 @@ class JetEngineEnvironmentData:
   # The ratio between query heads and kv heads
   n_reps: int = 0
 
+
 # pylint: disable-next=all
 class JetEngineEnvironment:
-# pylint: disable-next=all
+  # pylint: disable-next=all
   def __init__(self, data: JetEngineEnvironmentData):
     self._data = data
 

jetstream_pt/layers.py

@@ -438,7 +438,7 @@ def attend(xq, keys, values, local_mask=None):
             xq, (0, 0, 0, true_len - seqlen), "constant", 0
         )
 
-      if self.env.ragged_mha and seqlen == 1:
+      if self.env.ragged_mha and seqlen == 1 and keys.shape[-2] > 1:
         local_output, (local_max, local_denom) = torch_xla2.interop.call_jax(
             impl,
             xq,
@@ -589,7 +589,7 @@ def attend(xq, keys, values, k_scaler, v_scaler, local_mask=None):
         )
 
       # We are not using ragged attention for prefill yet.
-      if self.env.ragged_mha and seqlen == 1:
+      if self.env.ragged_mha and seqlen == 1 and keys.shape[-2] > 1:
         local_output, (local_max, local_denom) = torch_xla2.interop.call_jax(
             impl,
             xq,

tests/test_run_server.py

@@ -45,7 +45,7 @@ def reset_flags(self):
 
   def setup(self):
     """Setup."""
-     # pylint: disable-next=all
+    # pylint: disable-next=all
     from run_server import flags
 
     f = flags.FLAGS