[dtensor] fix scaled dot product flash attention sharding

XilunWu · XilunWu · commit 4439ac0e909c · 2025-02-27T16:31:26.000-08:00
ghstack-source-id: 408ec85 Pull Request resolved: #148125
diff --git a/test/distributed/tensor/test_attention.py b/test/distributed/tensor/test_attention.py
@@ -42,7 +42,6 @@
 if PLATFORM_SUPPORTS_MEM_EFF_ATTENTION:
     backends.append(SDPBackend.EFFICIENT_ATTENTION)
 
-
 rotater_enum_to_str = {
     _RotateMethod.ALL_GATHER: "allgather",
     _RotateMethod.ALL_TO_ALL: "alltoall",
@@ -360,6 +359,9 @@ def _test_ring_attention_custom_transformer(self, rotater: _RotateMethod) -> Non
             self.device_type,
             torch.arange(0, self.world_size),
         )
+        # early init DTensor RNG tracker to avoid broadcast be captuured in comm_mode
+        torch.distributed.tensor._random.manual_seed(10, device_mesh)
+
         dtype = torch.bfloat16
         bs = 2
         args = ModelArgs()
diff --git a/torch/distributed/tensor/_ops/_matrix_ops.py b/torch/distributed/tensor/_ops/_matrix_ops.py
@@ -252,8 +252,8 @@ def scaled_dot_product_flash_attention_strategy(
         None,  # cum_seq_k
         None,  # max_q
         None,  # max_k
-        None,  # philox_seed
-        None,  # philox_offset
+        Replicate(),  # rng_state
+        None,  # unused
         Replicate(),
         Replicate(),
         Replicate(),
@@ -279,8 +279,8 @@ def scaled_dot_product_flash_attention_strategy(
         None,  # cum_seq_k
         None,  # max_q
         None,  # max_k
-        None,  # philox_seed
-        None,  # philox_offset
+        Replicate(),  # rng_state
+        None,  # unused
         debug_attn_mask_sharding,
         qkv_sharding,
         qkv_sharding,
@@ -297,8 +297,8 @@ def scaled_dot_product_flash_attention_strategy(
             None,  # cum_seq_k
             None,  # max_q
             None,  # max_k
-            None,  # philox_seed
-            None,  # philox_offset
+            Replicate(),  # rng_state
+            None,  # unused
             Shard(2),  # debugattn
             Shard(2),  # q
             Shard(2),  # k
