add unique_count for dist_attn (#85)

Author: hanwen-sun

magi_attention/functional/dist_attn.py

@@ -324,6 +324,7 @@ def attn_fwd_partial(
                         **attn_arg.to_ffa_args(is_bwd=False),
                         merge_q_ranges=None,
                         qk_map=None,
+                        fwd_unique_count=None,
                         softmax_scale=q.shape[-1] ** -0.5,
                         deterministic=deterministic,
                         softcap=0.0,
@@ -394,6 +395,7 @@ def attn_bwd_partial(
                     **attn_arg.to_ffa_args(is_bwd=True),
                     merge_k_ranges=None,
                     bwd_kq_map=None,
+                    bwd_unique_count=None,
                     softmax_scale=q.shape[-1] ** -0.5,
                     deterministic=deterministic,
                     softcap=0.0,

magi_attention/functional/flex_flash_attn.py

@@ -130,7 +130,7 @@ def _flex_flash_attn_forward(
     attn_type_map,
     merge_q_ranges,
     qk_map,
-    unique_count,
+    fwd_unique_count,
     softmax_scale,
     softcap,
     disable_fwd_atomic_reduction,
@@ -157,7 +157,7 @@ def _flex_flash_attn_forward(
         attn_type_map,
         merge_q_ranges,
         qk_map,
-        unique_count,
+        fwd_unique_count,
         softmax_scale,
         softcap,
         disable_fwd_atomic_reduction,