Merge branch 'main' into fused_attn/graph_api_v1

Signed-off-by: cyanguwa <[email protected]>

Author: cyanguwa

tests/jax/test_distributed_fused_attn.py

@@ -197,6 +197,7 @@ def target_func(q, kv, mask):
             return jnp.mean(
                 cross_fused_attn(q,
                                  kv,
+                                 None,
                                  mask,
                                  None,
                                  attn_bias_type=attn_bias_type,

tests/jax/test_fused_attn.py

@@ -163,7 +163,7 @@ def customcall_cross_fused_attn(q, kv, q_token, kv_token, dropout_rng, **kwargs)
     # mask invert
     mask = (mask == 0)
 
-    return cross_fused_attn(q, kv, mask, dropout_rng, **kwargs)
+    return cross_fused_attn(q, kv, None, mask, dropout_rng, **kwargs)
 
 
 @pytest.mark.parametrize('b, s, h, d', SELF_CASES)