diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 1dd92fe181..a657431fdc 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -957,7 +957,7 @@ def test_forward(self):
 
         customcall_args = [
             # Put test data onto each GPU for distributed.
-            # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mas and
+            # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mask and
             # THD params once we support those features on CP.
             jax.device_put(self.cp_reorder_fn(self.q), self.qkvo_sharding),
             jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding),
@@ -1089,7 +1089,7 @@ def grad_func(
             self.dropout_rng,
         ]
         customcall_args = [
-            # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mas and
+            # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mask and
             # THD params once we support those features on CP.
             jax.device_put(self.cp_reorder_fn(self.q), self.qkvo_sharding),
             jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding),