Repository: NVIDIA/Fuser
Pull request #3953: DID loop split for reshape without pre-sharding reshape propagation
Author: login: Priya2698
Created: 2025-02-24T19:53:56Z | Merged: 2025-03-18T21:57:10Z

This PR:
1. Adds a WAR to identify reshapes as non-resharding if both the producer and consumer use either outer or inner split to produce the DIDx axis. This avoids having to propagate the reshape transform upwards.
2. Adds example for split reshape, merge reshape, and part of transformer forward block demonstrating propagation of shardings through reshape, permute and SDPA operations.

Will be rebased on #3875.


Diff:
diff --git a/csrc/multidevice/utils.cpp b/csrc/multidevice/utils.cpp
index 7b000ace95b..84e3ca27d4f 100644
--- a/csrc/multidevice/utils.cpp
+++ b/csrc/multidevice/utils.cpp
@@ -398,28 +398,90 @@ bool haveDifferentShardings(
       }
     }
 
-    auto is_mapped_in_id_model =
-        [](IterDomain* a, IterDomain* b, const IdModel& id_model) -> bool {
-      if (a == nullptr && b == nullptr) {
+    auto is_mapped_in_id_model = [producer, consumer, mapped_c_root_ids](
+                                     IterDomain* p_loop_id,
+                                     IterDomain* c_loop_id,
+                                     const IdModel& id_model) -> bool {
+      if (p_loop_id == nullptr && c_loop_id == nullptr) {
         return true;
       }
 
-      if (a == nullptr || b == nullptr) {
+      if (p_loop_id == nullptr || c_loop_id == nullptr) {
         return false;
       }
 
+      if (consumer->definition()->isA<ViewOp>()) {
+        // Consider the reshape: [h]-> [a, h/a] where `h` and `a` are sharded,
+        // IdModel currently does not map the two DID axes. Similarly for the
+        // corresponding merge reshape. This is a temporary workaround to check
+        // if the reshape is resharding when sharding is on reshaped ids.
+
+        // Get the root iterdomains in the consumer that produce reshape logical
+        // Ids.
+        std::unordered_set<IterDomain*> c_reshape_ids;
+        for (auto id : consumer->getLogicalDomain()) {
+          if (id->isRFactorProduct() && id->definition() &&
+              !id->definition()->isA<Resize>()) {
+            auto root_ids = getInputsInTargetDomain(
+                id, {mapped_c_root_ids.begin(), mapped_c_root_ids.end()});
+            for (auto root_id : root_ids) {
+              c_reshape_ids.insert(root_id);
+            }
+          }
+        }
+
+        // Get the logical iterdomains in the producer that are reshaped.
+        const std::unordered_map<IterDomain*, IterDomain*>& c2p =
+            PairwiseLogicalDomainMap(producer, consumer)
+                .mapConsumerToProducer();
+        std::unordered_set<IterDomain*> p_reshape_ids;
+        for (auto id : c_reshape_ids) {
+          auto p_id = c2p.find(id);
+          if (p_id != c2p.end()) {
+            p_reshape_ids.insert(p_id->second);
+          }
+        }
+
+        auto p_transforms = DependencyCheck::getAllExprsBetween(
+            {p_reshape_ids.begin(), p_reshape_ids.end()}, {p_loop_id});
+        auto c_transforms = DependencyCheck::getAllExprsBetween(
+            {c_reshape_ids.begin(), c_reshape_ids.end()}, {c_loop_id});
+
+        if (p_transforms.size() > 0 && c_transforms.size() > 0) {
+          // Both a and b come from reshaped ids.
+          NVF_ERROR(
+              p_loop_id->definition()->isA<Split>(),
+              p_loop_id,
+              " is not a Split: ",
+              p_loop_id->definition());
+          NVF_ERROR(
+              c_loop_id->definition()->isA<Split>(),
+              c_loop_id,
+              " is not a Split: ",
+              c_loop_id->definition());
+          // Get the split producing the sharded axis
+          auto* p_split = p_transforms.back()->as<Split>();
+          auto* c_split = c_transforms.back()->as<Split>();
+
+          // Reshape is not resharding if both the DID splits are either inner
+          // or outer splits. Note that we currently do not exercise inner
+          // splits in the multidevice tests.
+          return p_split->innerSplit() == c_split->innerSplit();
+        }
+      }
+
       // Going between bDIDx{1} and iDIDx{N} doesn't trigger resharding, but
       // would be flagged by ALMOSTEXACT as a false positive.
       if (id_model.idGraph(IdMappingMode::BROADCAST)
               .disjointValSets()
-              .strictAreMapped(a, b)) {
+              .strictAreMapped(p_loop_id, c_loop_id)) {
         return true;
       }
 
       // Check ALMOSTEXACT so iDIDx{N}*b{1} and iDIDx{N} are mapped.
       return id_model.idGraph(IdMappingMode::ALMOSTEXACT)
           .disjointValSets()
-          .strictAreMapped(a, b);
+          .strictAreMapped(p_loop_id, c_loop_id);
     };
 
     if (!is_mapped_in_id_model(p_loop_id, c_loop_id, id_model)) {
diff --git a/csrc/scheduler/utils.cpp b/csrc/scheduler/utils.cpp
index 124c0994277..ff5bb48777d 100644
--- a/csrc/scheduler/utils.cpp
+++ b/csrc/scheduler/utils.cpp
@@ -294,7 +294,8 @@ void parallelizeAllLike(
     int64_t pos,
     std::vector<TensorView*> selected_tvs,
     const std::unordered_set<ParallelType>& selected_parallel_types,
-    bool propagate_padding) {
+    bool propagate_padding,
+    bool parallelize_inputs) {
   FusionGuard fg(reference_tv->fusion());
 
   if (pos < 0) {
@@ -320,7 +321,7 @@ void parallelizeAllLike(
     selected_tvs = reference_tv->fusion()->allTvs();
   }
   for (auto tv : selected_tvs) {
-    if (tv->isFusionInput()) {
+    if (tv->isFusionInput() && !parallelize_inputs) {
       continue;
     }
     for (const auto i : c10::irange((int64_t)tv->getLoopDomain().size())) {
@@ -2285,6 +2286,16 @@ void propagateReshapeTransforms(Fusion* fusion, const ComputeAtMap& ca_map) {
     tv->reorder(old2new);
     //! Propagate current transformations on from_tv to all graphs
     transformPropagateToAllFrom(tv, (int64_t)old2new.size());
+
+    // Propgating the transforms will not replay the DIDx parallelization, so we
+    // need to do it manually here.
+    parallelizeAllLike(
+        tv,
+        /*pos=*/(int64_t)old2new.size(),
+        /*selected_tvs=*/{},
+        /*selected_parallel_types=*/{ParallelType::DIDx},
+        /*propagate_padding=*/false,
+        /*parallelize_inputs=*/true);
   }
 }
 
diff --git a/csrc/scheduler/utils.h b/csrc/scheduler/utils.h
index dc0ae305704..4e95d854c28 100644
--- a/csrc/scheduler/utils.h
+++ b/csrc/scheduler/utils.h
@@ -153,48 +153,60 @@ int64_t mergeNonReduction(TensorView* tv);
 // DAG. Empty `selected_tvs` means selecting all tensors in the fusion of
 // `reference_tv`. `selected_parallel_types` are the selected parallel types.
 // Empty `selected_parallel_types` means selecting all parallel types.
+// `parallelize_inputs` is a boolean flag that determines whether to parallelize
+// the inputs of the fusion. This is generally not required except in some cases
+// for DID parallelization. For eg: propagateReshapeTransforms using
+// TransformPropagator will replay the transforms onto the inputs of the fusion
+// but not the DID parallelization.
 void parallelizeAllLike(
     TensorView* reference_tv,
     int64_t pos = -1,
     std::vector<TensorView*> selected_tvs = {},
     const std::unordered_set<ParallelType>& selected_parallel_types = {},
-    bool propagate_padding = true);
+    bool propagate_padding = true,
+    bool parallelize_inputs = false);
 
 inline void parallelizeAllLike(
     TensorView* reference_tv,
     std::vector<TensorView*> selected_tvs,
     const std::unordered_set<ParallelType>& selected_parallel_types = {},
-    bool propagate_padding = true) {
+    bool propagate_padding = true,
+    bool parallelize_inputs = false) {
   parallelizeAllLike(
       reference_tv,
       -1,
       std::move(selected_tvs),
       selected_parallel_types,
-      propagate_padding);
+      propagate_padding,
+      parallelize_inputs);
 }
 
 inline void parallelizeAllLike(
     TensorView* reference_tv,
     std::initializer_list<TensorView*> selected_tvs,
     const std::unordered_set<ParallelType>& selected_parallel_types = {},
-    bool propagate_padding = true) {
+    bool propagate_padding = true,
+    bool parallelize_inputs = false) {
   parallelizeAllLike(
       reference_tv,
       std::vector<TensorView*>(selected_tvs),
       selected_parallel_types,
-      propagate_padding);
+      propagate_padding,
+      parallelize_inputs);
 }
 
 inline void parallelizeAllLike(
     TensorView* reference_tv,
     const std::unordered_set<ParallelType>& selected_parallel_types,
-    bool propagate_padding = true) {
+    bool propagate_padding = true,
+    bool parallelize_inputs = false) {
   parallelizeAllLike(
       reference_tv,
       -1,
       std::vector<TensorView*>{},
       selected_parallel_types,
-      propagate_padding);
+      propagate_padding,
+      parallelize_inputs);
 }
 
 // Common hyperparameters used in heuristic scheduler. These hyperparameters
diff --git a/tests/cpp/test_multidevice_sharding.cpp b/tests/cpp/test_multidevice_sharding.cpp
index 9940922af3f..94ef116bab2 100644
--- a/tests/cpp/test_multidevice_sharding.cpp
+++ b/tests/cpp/test_multidevice_sharding.cpp
@@ -745,7 +745,7 @@ TEST_F(MultiDeviceTest, ReorderDIDToFront) {
       __FILE__);
 }
 
-TEST_F(MultiDeviceTest, ShardedSplitReshapeIds) {
+TEST_F(MultiDeviceTest, TransformPropagatorSplitReshape) {
   auto fusion = std::make_unique<Fusion>();
   FusionGuard fg(fusion.get());
 
@@ -796,6 +796,51 @@ TEST_F(MultiDeviceTest, ShardedSplitReshapeIds) {
   FusionExecutorCache executor_cache(std::move(fusion));
   at::Tensor inp = at::randn({b, s, d * h * e}, tensor_options);
   at::Tensor sharded_inp = shardTensor(inp, tv0);
+
+  at::Tensor nvf_out =
+      executor_cache.runFusionWithInputs({sharded_inp})[0].as<at::Tensor>();
+
+  testValidate(
+      executor_cache.fusion(),
+      {nvf_out},
+      {sharded_inp},
+      {sharded_inp.view({b, s, h, e})},
+      __LINE__,
+      __FILE__);
+}
+
+TEST_F(MultiDeviceTest, LoopShardedSplitReshapeIds) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  const int d = communicator_->size();
+  const int64_t b = 2, s = 3, h = 8, e = 4;
+
+  TensorView* tv0 = makeContigConcreteTensor({b, s, d * h * e});
+  TensorView* tv1 = reshape(tv0, {b, s, d * h * e}, {b, s, d * h, e});
+
+  fusion->addInput(tv0);
+  fusion->addOutput(tv1);
+
+  auto mesh = DeviceMesh::createForNumDevices(d);
+
+  tv0->setDeviceMesh(mesh);
+  tv0->split(-1, d, /*inner_split=*/false);
+  tv0->axis(-2)->parallelize(ParallelType::DIDx);
+
+  tv1->setDeviceMesh(mesh);
+  tv1->split(-2, d, /*inner_split=*/false);
+  tv1->axis(-3)->parallelize(ParallelType::DIDx);
+
+  for (auto* tv : {tv0, tv1}) {
+    reorderDIDToFront(tv);
+    tv->setAllocationDomain(tv->getLoopDomain(), true);
+  }
+
+  FusionExecutorCache executor_cache(std::move(fusion));
+  at::Tensor inp = at::randn({b, s, d * h * e}, tensor_options);
+  at::Tensor sharded_inp = shardTensor(inp, -1, mesh);
+
   at::Tensor nvf_out =
       executor_cache.runFusionWithInputs({sharded_inp})[0].as<at::Tensor>();
   testValidate(
@@ -807,4 +852,203 @@ TEST_F(MultiDeviceTest, ShardedSplitReshapeIds) {
       __FILE__);
 }
 
+TEST_F(MultiDeviceTest, LoopShardedMergeReshapeIds) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  const int d = communicator_->size();
+  const int64_t b = 2, s = 3, h = 8, e = 4;
+
+  TensorView* tv0 = makeContigConcreteTensor({b, s, d * h, e});
+  TensorView* tv1 = reshape(tv0, {b, s, d * h, e}, {b, s, d * h * e});
+
+  fusion->addInput(tv0);
+  fusion->addOutput(tv1);
+
+  auto mesh = DeviceMesh::createForNumDevices(d);
+  tv0->setDeviceMesh(mesh);
+  tv0->split(-2, d, /*inner_split=*/false);
+  tv0->axis(-3)->parallelize(ParallelType::DIDx);
+
+  tv1->setDeviceMesh(mesh);
+  tv1->split(-1, d, /*inner_split=*/false);
+  tv1->axis(-2)->parallelize(ParallelType::DIDx);
+
+  for (auto* tv : {tv0, tv1}) {
+    reorderDIDToFront(tv);
+    tv->setAllocationDomain(tv->getLoopDomain(), true);
+  }
+
+  FusionExecutorCache executor_cache(std::move(fusion));
+  at::Tensor inp = at::randn({b, s, d * h, e}, tensor_options);
+  at::Tensor sharded_inp = shardTensor(inp, -2, mesh);
+  at::Tensor nvf_out =
+      executor_cache.runFusionWithInputs({sharded_inp})[0].as<at::Tensor>();
+  testValidate(
+      executor_cache.fusion(),
+      {nvf_out},
+      {sharded_inp},
+      {sharded_inp.view({b, s, h * e})},
+      __LINE__,
+      __FILE__);
+}
+
+namespace {
+// This is a simplified version of what we will eventually do in the
+// pre-segmentation pass
+void propagateShardings(Fusion* fusion, int64_t num_devices) {
+  for (Expr* expr : fusion->exprs()) {
+    if (expr->isA<SliceOp>()) {
+      NVF_THROW("SliceOp is not currently supported");
+    }
+
+    if (expr->isA<ViewOp>()) {
+      // TransformPropagator cannot be directly used.
+      // It raises an error for conflicting transformations from root domain to
+      // logical domain. Instead, we manually find the reshaped iterdomain and
+      // outer split DID. This might have to be extended further in the
+      // presegmentation pass.
+      // Note: For simplicity, this assumes that the sharding is on reshaped
+      // IDs. It is possible that the non-reshaped IDs are sharded, in which
+      // case we can use the TransformPropagator.
+      TensorView* reshaped_tv = expr->as<ViewOp>()->out();
+      auto transform_exprs = StmtSort::getExprsBetween(
+          {reshaped_tv->getMaybeRootDomain().begin(),
+           reshaped_tv->getMaybeRootDomain().end()},
+          {reshaped_tv->getLogicalDomain().begin(),
+           reshaped_tv->getLogicalDomain().end()});
+      NVF_CHECK(transform_exprs.size() == 1);
+      auto transform = transform_exprs[0];
+      NVF_CHECK(transform->isA<Split>() || transform->isA<Merge>());
+
+      // Get the reshaped ID (outer ID for split reshape).
+      // This is the ID that will be parallelized.
+      IterDomain* reshaped_id = transform->isA<Split>()
+          ? transform->as<Split>()->outer()
+          : transform->as<Merge>()->out();
+
+      auto reshaped_it = std::find(
+          reshaped_tv->getLoopDomain().begin(),
+          reshaped_tv->getLoopDomain().end(),
+          reshaped_id);
+      int64_t reshaped_axis =
+          std::distance(reshaped_tv->getLoopDomain().begin(), reshaped_it);
+
+      // Apply sharding to the reshaped tensor
+      reshaped_tv->split(reshaped_axis, num_devices, false);
+      reshaped_tv->axis(reshaped_axis)->parallelize(ParallelType::DIDx);
+      reorderDIDToFront(reshaped_tv);
+      continue;
+    }
+
+    // For other ops, propagate sharding from input to outputs
+    auto input_tv = expr->input(0)->as<TensorView>();
+    std::vector<TensorView*> output_tvs;
+    for (auto output : expr->outputs()) {
+      output_tvs.push_back(output->as<TensorView>());
+    }
+
+    TransformPropagator propagator(input_tv);
+
+    // Note: We will finally propagate from each input iteratively.
+    SetSelector selector(
+        std::unordered_set<TensorView*>(output_tvs.begin(), output_tvs.end()));
+    MaxLogicalDomainInfoSpanningTree(input_tv, &selector).traverse(&propagator);
+    scheduler_utils::parallelizeAllLike(
+        input_tv,
+        /*pos=*/-1,
+        /*selected_tv=*/output_tvs);
+  }
+}
+
+} // namespace
+
+TEST_F(MultiDeviceTest, TransformerFwd) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  const int d = communicator_->size();
+  const int64_t b = 2, s = 3, h = 8, e = 16;
+  auto mesh = DeviceMesh::createForNumDevices(d);
+
+  std::vector<int64_t> in_shape = {b, s, d * h * e};
+  std::vector<int64_t> out_shape = {b, s, d * h, e};
+
+  // The transformer block produces hq/hk/hv after slicing the MHA linear
+  // output.
+  TensorView* hq = makeConcreteTensor(in_shape, DataType::Half);
+  TensorView* hk = makeConcreteTensor(in_shape, DataType::Half);
+  TensorView* hv = makeConcreteTensor(in_shape, DataType::Half);
+
+  TensorView* q = reshape(hq, in_shape, out_shape);
+  TensorView* q_permuted = permute(q, {0, 2, 1, 3});
+  TensorView* k = reshape(hk, in_shape, out_shape);
+  TensorView* k_permuted = permute(k, {0, 2, 1, 3});
+  TensorView* v = reshape(hv, in_shape, out_shape);
+  TensorView* v_permuted = permute(v, {0, 2, 1, 3});
+
+  SdpfaFwdResult sdpa_out = sdpfa_fwd(
+      q_permuted,
+      k_permuted,
+      v_permuted,
+      /*dropout_p=*/IrBuilder::create<Val>(0.0),
+      /*is_causal=*/IrBuilder::create<Val>(false),
+      /*scale=*/nullptr);
+
+  TensorView* attn = sdpa_out.output;
+  TensorView* attn_permute = permute(attn, {0, 2, 1, 3});
+  TensorView* out = reshape(attn_permute, out_shape, in_shape);
+
+  fusion->addInput(hq);
+  fusion->addInput(hk);
+  fusion->addInput(hv);
+  fusion->addOutput(out);
+
+  // Shard input tensors
+  for (auto* tv : {hq, hk, hv}) {
+    tv->setDeviceMesh(mesh);
+    tv->split(-1, d, /*inner_split=*/false);
+    tv->axis(-2)->parallelize(ParallelType::DIDx);
+    reorderDIDToFront(tv);
+  }
+  propagateShardings(fusion.get(), d);
+
+  for (auto tv : fusion->allTvs()) {
+    tv->setAllocationDomain(tv->getLoopDomain(), true);
+  }
+
+  FusionExecutorCache executor_cache(std::move(fusion));
+  at::Tensor hq_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));
+  at::Tensor hk_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));
+  at::Tensor hv_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));
+
+  at::Tensor sharded_hq = shardTensor(hq_tensor, -1, mesh);
+  at::Tensor sharded_hk = shardTensor(hk_tensor, -1, mesh);
+  at::Tensor sharded_hv = shardTensor(hv_tensor, -1, mesh);
+
+  auto nvf_out =
+      executor_cache
+          .runFusionWithInputs({sharded_hq, sharded_hk, sharded_hv})[0]
+          .as<at::Tensor>();
+
+  double scale = 1.0 / std::sqrt(e);
+  auto reference_out = at::_scaled_dot_product_flash_attention(
+      hq_tensor.view(out_shape).transpose(1, 2),
+      hk_tensor.view(out_shape).transpose(1, 2),
+      hv_tensor.view(out_shape).transpose(1, 2),
+      /*dropout_p=*/0.0,
+      /*is_causal=*/false,
+      /*return_debug_mask=*/false,
+      scale);
+  at::Tensor ref_attn = shardTensor(
+      std::get<0>(reference_out).transpose(1, 2).view(in_shape), -1, mesh);
+
+  testValidate(
+      executor_cache.fusion(),
+      {nvf_out},
+      {sharded_hq, sharded_hk, sharded_hv},
+      {ref_attn},
+      __LINE__,
+      __FILE__);
+}
 } // namespace nvfuser


Reviews:
user: login: Priya2698
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-05T00:17:41Z
user: login: Priya2698
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-05T00:18:28Z
user: login: wujingyue
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-05T00:34:00Z
user: login: wujingyue
body: LGTM overall
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-05T00:51:43Z
user: login: Priya2698
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-05T01:37:04Z
user: login: Priya2698
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-05T23:26:35Z
user: login: Priya2698
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-06T00:29:21Z
user: login: Priya2698
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-07T20:15:33Z
user: login: wujingyue
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-11T22:56:04Z

Review comments:
diff_hunk: @@ -2234,17 +2235,38 @@ void propagateReshapeTransforms(Fusion* fusion, const ComputeAtMap& ca_map) {
     for (auto logical_id : tv->getLogicalDomain()) {
       if (terminating_reshape_dims.find(logical_id) !=
           terminating_reshape_dims.end()) {
+        // Check if logical ID is directly in the loop domain
         auto find_it = std::find(
             tv->getLoopDomain().begin(), tv->getLoopDomain().end(), logical_id);
+
+        // If not found directly and there is a sharded loop ID,
+        // check if the logical ID is the same as the producer of the DID split.
+        if (find_it == tv->getLoopDomain().end()) {
+          int64_t sharded_axis = getShardedLoopAxis(tv, ParallelType::DIDx);
+          if (sharded_axis != -1) {
+            // Get the split operation that created the DIDx dimension
+            auto split = dynamic_cast<Split*>(
+                tv->getLoopDomain().at(sharded_axis)->definition());
+            if (split != nullptr && split->in() == logical_id) {
+              old2new[sharded_axis] = (int64_t)old2new.size();
+              find_it = std::find(
+                  tv->getLoopDomain().begin(),
+                  tv->getLoopDomain().end(),
+                  split->inner());
+            }
+          }
+        }
+
path: csrc/scheduler/utils.cpp
user: login: Priya2698
body: This will be superseded by changes in PR #3875.
created_at: 2025-03-05T00:17:41Z
author_association: COLLABORATOR
diff_hunk: @@ -2257,6 +2279,15 @@ void propagateReshapeTransforms(Fusion* fusion, const ComputeAtMap& ca_map) {
     tv->reorder(old2new);
     //! Propagate current transformations on from_tv to all graphs
     transformPropagateToAllFrom(tv, (int64_t)old2new.size());
+    // Propgating the transforms will not replay the DIDx parallelization, so we
+    // need to do it manually here.
+    parallelizeAllLike(
+        tv,
+        /*pos=*/(int64_t)old2new.size(),
+        /*selected_tvs=*/{},
+        /*selected_parallel_types=*/{ParallelType::DIDx},
+        /*propagate_padding=*/false,
+        /*parallelize_inputs=*/true);
path: csrc/scheduler/utils.cpp
user: login: Priya2698
body: This will be the only new change in this PR for `propagateReshapeTransforms`.
created_at: 2025-03-05T00:18:28Z
author_association: COLLABORATOR
diff_hunk: @@ -745,4 +745,307 @@ TEST_F(MultiDeviceTest, ReorderDIDToFront) {
       __FILE__);
 }
 
+TEST_F(MultiDeviceTest, TransformPropagatorSplitReshape) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  const int d = communicator_->size();
+  const int64_t b = 2, s = 2, h = 4, e = 3;
+
+  TensorView* in = makeContigConcreteTensor(
+      {b, s, d * h * e}); // in: loop domain: {b, s, d*h*e}
+  TensorView* out = reshape(
+      in,
+      {b, s, d * h * e},
+      {b, s, d * h, e}); // out: loop domain: {b, s, d*h, e}
+
+  fusion->addInput(in);
+  fusion->addOutput(out);
+
+  auto mesh = DeviceMesh::createForNumDevices(d);
+
+  // Propagate transform from reshaped output to input.
+  // Without this propagation, the two DID axes on `in` and `out` will not be
+  // mapped in together in ID model. This causes scheduling to fail due to
+  // resharding.
+  TransformPropagator propagator_c2p(out);
+  MaxLogicalDomainInfoSpanningTree(out).traverse(&propagator_c2p);
+  // in: loop domain: {b, s, d*h, e} after transform propagation
path: tests/cpp/test_multidevice_sharding.cpp
user: login: wujingyue
body: Does propagateShardings in this file fail on split reshapes? 
created_at: 2025-03-05T00:34:00Z
author_association: COLLABORATOR
diff_hunk: @@ -399,7 +399,8 @@ bool haveDifferentShardings(
     }
 
     auto is_mapped_in_id_model =
-        [](IterDomain* a, IterDomain* b, const IdModel& id_model) -> bool {
+        [&consumer](
path: csrc/multidevice/utils.cpp
user: login: wujingyue
body: ```suggestion
        [consumer](
```
created_at: 2025-03-05T00:50:08Z
author_association: COLLABORATOR
diff_hunk: @@ -408,6 +409,33 @@ bool haveDifferentShardings(
         return false;
       }
 
+      if (consumer->definition()->isA<ViewOp>()) {
+        // Consider the reshape: [h]-> [a, h/a] where `h` and `a` are sharded,
+        // IdModel currently does not map the two DID axes. Similarly for the
+        // corresponding merge reshape. This is a temporary workaround to check
+        // if the reshape is resharding.
+        NVF_ERROR(
+            a->definition()->isA<Split>(),
+            a,
+            " is not a Split: ",
+            a->definition());
+        NVF_ERROR(
+            b->definition()->isA<Split>(),
+            b,
+            " is not a Split: ",
+            b->definition());
+        // Get the split producing the sharded axis
+        auto* a_split = a->definition()->as<Split>();
+        auto* b_split = b->definition()->as<Split>();
+        // Reshape is not resharding if both the DID splits are either inner or
+        // outer splits. Note that we currently do not exercise inner splits in
+        // the multidevice tests.
+        if (a_split->innerSplit() == b_split->innerSplit()) {
+          return true;
+        }
+        return false;
path: csrc/multidevice/utils.cpp
user: login: wujingyue
body: ```suggestion
        return a_split->innerSplit() == b_split->innerSplit();
```
created_at: 2025-03-05T00:50:29Z
author_association: COLLABORATOR
diff_hunk: @@ -745,4 +745,307 @@ TEST_F(MultiDeviceTest, ReorderDIDToFront) {
       __FILE__);
 }
 
+TEST_F(MultiDeviceTest, TransformPropagatorSplitReshape) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  const int d = communicator_->size();
+  const int64_t b = 2, s = 2, h = 4, e = 3;
+
+  TensorView* in = makeContigConcreteTensor(
+      {b, s, d * h * e}); // in: loop domain: {b, s, d*h*e}
+  TensorView* out = reshape(
+      in,
+      {b, s, d * h * e},
+      {b, s, d * h, e}); // out: loop domain: {b, s, d*h, e}
+
+  fusion->addInput(in);
+  fusion->addOutput(out);
+
+  auto mesh = DeviceMesh::createForNumDevices(d);
+
+  // Propagate transform from reshaped output to input.
+  // Without this propagation, the two DID axes on `in` and `out` will not be
+  // mapped in together in ID model. This causes scheduling to fail due to
+  // resharding.
+  TransformPropagator propagator_c2p(out);
+  MaxLogicalDomainInfoSpanningTree(out).traverse(&propagator_c2p);
+  // in: loop domain: {b, s, d*h, e} after transform propagation
path: tests/cpp/test_multidevice_sharding.cpp
user: login: Priya2698
body: I've kept the example where we propagate reshapes before loop split in the case we decide to go that route. Hence, keeping this example from PR #3875 .
created_at: 2025-03-05T01:37:04Z
author_association: COLLABORATOR
diff_hunk: @@ -408,6 +409,30 @@ bool haveDifferentShardings(
         return false;
       }
 
+      if (consumer->definition()->isA<ViewOp>()) {
+        // Consider the reshape: [h]-> [a, h/a] where `h` and `a` are sharded,
+        // IdModel currently does not map the two DID axes. Similarly for the
+        // corresponding merge reshape. This is a temporary workaround to check
+        // if the reshape is resharding.
+        NVF_ERROR(
path: csrc/multidevice/utils.cpp
user: login: Priya2698
body: This is incorrect since we do not have this split if the DID parallelization is on logical domain. I'll add a fix.
created_at: 2025-03-05T23:26:35Z
author_association: COLLABORATOR
diff_hunk: @@ -408,6 +409,30 @@ bool haveDifferentShardings(
         return false;
       }
 
+      if (consumer->definition()->isA<ViewOp>()) {
+        // Consider the reshape: [h]-> [a, h/a] where `h` and `a` are sharded,
+        // IdModel currently does not map the two DID axes. Similarly for the
+        // corresponding merge reshape. This is a temporary workaround to check
+        // if the reshape is resharding.
+        NVF_ERROR(
path: csrc/multidevice/utils.cpp
user: login: Priya2698
body: (TODO) I need to add 2 things here:
1. Check if the sharding is on the reshaped IDs for both producer and consumer, else use IdModel.
2. If the sharding is on reshaped IDs, it is outer split and on the outer reshaped ID.
created_at: 2025-03-06T00:29:21Z
author_association: COLLABORATOR
diff_hunk: @@ -399,27 +399,90 @@ bool haveDifferentShardings(
     }
 
     auto is_mapped_in_id_model =
-        [](IterDomain* a, IterDomain* b, const IdModel& id_model) -> bool {
-      if (a == nullptr && b == nullptr) {
+        [producer, consumer, mapped_p_logical_ids, mapped_c_root_ids](
+            IterDomain* p_loop_id,
+            IterDomain* c_loop_id,
+            const IdModel& id_model) -> bool {
+      if (p_loop_id == nullptr && c_loop_id == nullptr) {
         return true;
       }
 
-      if (a == nullptr || b == nullptr) {
+      if (p_loop_id == nullptr || c_loop_id == nullptr) {
         return false;
       }
 
+      if (consumer->definition()->isA<ViewOp>()) {
+        // Consider the reshape: [h]-> [a, h/a] where `h` and `a` are sharded,
+        // IdModel currently does not map the two DID axes. Similarly for the
+        // corresponding merge reshape. This is a temporary workaround to check
+        // if the reshape is resharding when sharding is on reshaped ids.
+
+        // Get the root iterdomains in the consumer that produce reshape logical
+        // Ids.
+        std::unordered_set<IterDomain*> c_reshape_ids;
+        for (auto id : consumer->getLogicalDomain()) {
+          if (id->isRFactorProduct() && id->definition() &&
+              !id->definition()->isA<Resize>()) {
+            auto root_ids = getInputsInTargetDomain(
+                id, {mapped_c_root_ids.begin(), mapped_c_root_ids.end()});
+            for (auto root_id : root_ids) {
+              c_reshape_ids.insert(root_id);
+            }
+          }
+        }
+
+        // Get the logical iterdomains in the producer that are reshaped.
+        const std::unordered_map<IterDomain*, IterDomain*>& c2p =
+            PairwiseLogicalDomainMap(producer, consumer)
+                .mapConsumerToProducer();
+        std::unordered_set<IterDomain*> p_reshape_ids;
+        for (auto id : c_reshape_ids) {
+          auto p_id = c2p.find(id);
+          if (p_id != c2p.end()) {
+            p_reshape_ids.insert(p_id->second);
+          }
+        }
+
+        auto p_transforms = DependencyCheck::getAllExprsBetween(
+            {p_reshape_ids.begin(), p_reshape_ids.end()}, {p_loop_id});
+        auto c_transforms = DependencyCheck::getAllExprsBetween(
+            {c_reshape_ids.begin(), c_reshape_ids.end()}, {c_loop_id});
+
+        if (p_transforms.size() > 0 && c_transforms.size() > 0) {
path: csrc/multidevice/utils.cpp
user: login: Priya2698
body: @wujingyue I am not checking for these transforms, if we are always on the outer since it will be replaced soon. Does that work in the interim while your PR #3482 lands?
created_at: 2025-03-07T20:15:33Z
author_association: COLLABORATOR
diff_hunk: @@ -158,43 +158,50 @@ void parallelizeAllLike(
     int64_t pos = -1,
     std::vector<TensorView*> selected_tvs = {},
     const std::unordered_set<ParallelType>& selected_parallel_types = {},
-    bool propagate_padding = true);
+    bool propagate_padding = true,
+    bool parallelize_inputs = false);
path: csrc/scheduler/utils.h
user: login: wujingyue
body: Can you code-comment this? Make sure you explain this is about fusion inputs or just tv's inputs. 
created_at: 2025-03-11T22:52:42Z
author_association: COLLABORATOR
diff_hunk: @@ -2285,6 +2286,15 @@ void propagateReshapeTransforms(Fusion* fusion, const ComputeAtMap& ca_map) {
     tv->reorder(old2new);
     //! Propagate current transformations on from_tv to all graphs
     transformPropagateToAllFrom(tv, (int64_t)old2new.size());
+    // Propgating the transforms will not replay the DIDx parallelization, so we
path: csrc/scheduler/utils.cpp
user: login: wujingyue
body: ```suggestion

    // Propgating the transforms will not replay the DIDx parallelization, so we
```
created_at: 2025-03-11T22:53:52Z
author_association: COLLABORATOR
diff_hunk: @@ -807,4 +809,246 @@ TEST_F(MultiDeviceTest, ShardedSplitReshapeIds) {
       __FILE__);
 }
 
+TEST_F(MultiDeviceTest, LoopShardedSplitReshapeIds) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  const int d = communicator_->size();
+  const int64_t b = 2, s = 3, h = 8, e = 4;
+
+  TensorView* tv0 = makeContigConcreteTensor({b, s, d * h * e});
+  TensorView* tv1 = reshape(tv0, {b, s, d * h * e}, {b, s, d * h, e});
+
+  fusion->addInput(tv0);
+  fusion->addOutput(tv1);
+
+  auto mesh = DeviceMesh::createForNumDevices(d);
+
+  tv0->setDeviceMesh(mesh);
+  tv0->split(-1, d, /*inner_split=*/false);
+  tv0->axis(-2)->parallelize(ParallelType::DIDx);
+
+  tv1->setDeviceMesh(mesh);
+  tv1->split(-2, d, /*inner_split=*/false);
+  tv1->axis(-3)->parallelize(ParallelType::DIDx);
+
+  for (auto* tv : {tv0, tv1}) {
+    reorderDIDToFront(tv);
+    tv->setAllocationDomain(tv->getLoopDomain(), true);
+  }
+
+  FusionExecutorCache executor_cache(std::move(fusion));
+  at::Tensor inp = at::randn({b, s, d * h * e}, tensor_options);
+  at::Tensor sharded_inp = shardTensor(inp, -1, mesh);
+
+  at::Tensor nvf_out =
+      executor_cache.runFusionWithInputs({sharded_inp})[0].as<at::Tensor>();
+  testValidate(
+      executor_cache.fusion(),
+      {nvf_out},
+      {sharded_inp},
+      {sharded_inp.view({b, s, h, e})},
+      __LINE__,
+      __FILE__);
+}
+
+TEST_F(MultiDeviceTest, LoopShardedMergeReshapeIds) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  const int d = communicator_->size();
+  const int64_t b = 2, s = 3, h = 8, e = 4;
+
+  TensorView* tv0 = makeContigConcreteTensor({b, s, d * h, e});
+  TensorView* tv1 = reshape(tv0, {b, s, d * h, e}, {b, s, d * h * e});
+
+  fusion->addInput(tv0);
+  fusion->addOutput(tv1);
+
+  auto mesh = DeviceMesh::createForNumDevices(d);
+  tv0->setDeviceMesh(mesh);
+  tv0->split(-2, d, /*inner_split=*/false);
+  tv0->axis(-3)->parallelize(ParallelType::DIDx);
+
+  tv1->setDeviceMesh(mesh);
+  tv1->split(-1, d, /*inner_split=*/false);
+  tv1->axis(-2)->parallelize(ParallelType::DIDx);
+
+  for (auto* tv : {tv0, tv1}) {
+    reorderDIDToFront(tv);
+    tv->setAllocationDomain(tv->getLoopDomain(), true);
+  }
+
+  FusionExecutorCache executor_cache(std::move(fusion));
+  at::Tensor inp = at::randn({b, s, d * h, e}, tensor_options);
+  at::Tensor sharded_inp = shardTensor(inp, -2, mesh);
+  at::Tensor nvf_out =
+      executor_cache.runFusionWithInputs({sharded_inp})[0].as<at::Tensor>();
+  testValidate(
+      executor_cache.fusion(),
+      {nvf_out},
+      {sharded_inp},
+      {sharded_inp.view({b, s, h * e})},
+      __LINE__,
+      __FILE__);
+}
+
+namespace {
+// This is a simplified version of what we will eventually do in the
+// pre-segmentation pass
+void propagateShardings(Fusion* fusion, int64_t num_devices) {
+  for (Expr* expr : fusion->exprs()) {
+    if (expr->isA<SliceOp>()) {
+      NVF_THROW("SliceOp is not currently supported");
+    }
+
+    if (expr->isA<ViewOp>()) {
+      // TransformPropagator cannot be directly used.
+      // It raises an error for conflicting transformations from root domain to
+      // logical domain. Instead, we manually find the reshaped iterdomain and
+      // outer split DID. This might have to be extended further in the
+      // presegmentation pass.
+      // Note: For simplicity, this assumes that the sharding is on reshaped
+      // IDs. It is possible that the non-reshaped IDs are sharded, in which
+      // case we can use the TransformPropagator.
+      TensorView* reshaped_tv = expr->as<ViewOp>()->out();
+      auto transform_exprs = StmtSort::getExprsBetween(
+          {reshaped_tv->getMaybeRootDomain().begin(),
+           reshaped_tv->getMaybeRootDomain().end()},
+          {reshaped_tv->getLogicalDomain().begin(),
+           reshaped_tv->getLogicalDomain().end()});
+      NVF_CHECK(transform_exprs.size() == 1);
+      auto transform = transform_exprs[0];
+      NVF_CHECK(transform->isA<Split>() || transform->isA<Merge>());
+
+      // Get the reshaped ID (outer ID for split reshape).
+      // This is the ID that will be parallelized.
+      IterDomain* reshaped_id = transform->isA<Split>()
+          ? transform->as<Split>()->outer()
+          : transform->as<Merge>()->out();
+
+      auto reshaped_it = std::find(
+          reshaped_tv->getLoopDomain().begin(),
+          reshaped_tv->getLoopDomain().end(),
+          reshaped_id);
+      int64_t reshaped_axis =
+          std::distance(reshaped_tv->getLoopDomain().begin(), reshaped_it);
+
+      // Apply sharding to the reshaped tensor
+      reshaped_tv->split(reshaped_axis, num_devices, false);
+      reshaped_tv->axis(reshaped_axis)->parallelize(ParallelType::DIDx);
+      reorderDIDToFront(reshaped_tv);
+      continue;
+    }
+
+    // For other ops, propagate sharding from input to outputs
+    auto input_tv = expr->input(0)->as<TensorView>();
+    std::vector<TensorView*> output_tvs;
+    for (auto output : expr->outputs()) {
+      output_tvs.push_back(output->as<TensorView>());
+    }
+
+    TransformPropagator propagator(input_tv);
+
+    // Note: We will finally propagate from each input iteratively.
+    SetSelector selector(
+        std::unordered_set<TensorView*>(output_tvs.begin(), output_tvs.end()));
+    MaxLogicalDomainInfoSpanningTree(input_tv, &selector).traverse(&propagator);
+    scheduler_utils::parallelizeAllLike(
+        input_tv,
+        /*pos=*/-1,
+        /*selected_tv=*/output_tvs);
+  }
+}
+
+} // namespace
+
+TEST_F(MultiDeviceTest, TransformerFwd) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  const int d = communicator_->size();
+  const int64_t b = 2, s = 3, h = 8, e = 16;
+  auto mesh = DeviceMesh::createForNumDevices(d);
+
+  std::vector<int64_t> in_shape = {b, s, d * h * e};
+  std::vector<int64_t> out_shape = {b, s, d * h, e};
+
+  // The transformer block produces hq/hk/hv after slicing the MHA linear
+  // output.
+  TensorView* hq = makeConcreteTensor(in_shape, DataType::Half);
+  TensorView* hk = makeConcreteTensor(in_shape, DataType::Half);
+  TensorView* hv = makeConcreteTensor(in_shape, DataType::Half);
+
+  TensorView* q = reshape(hq, in_shape, out_shape);
+  TensorView* q_permuted = permute(q, {0, 2, 1, 3});
+  TensorView* k = reshape(hk, in_shape, out_shape);
+  TensorView* k_permuted = permute(k, {0, 2, 1, 3});
+  TensorView* v = reshape(hv, in_shape, out_shape);
+  TensorView* v_permuted = permute(v, {0, 2, 1, 3});
+
+  SdpfaFwdResult sdpa_out = sdpfa_fwd(
+      q_permuted,
+      k_permuted,
+      v_permuted,
+      /*dropout_p=*/IrBuilder::create<Val>(0.0),
+      /*is_causal=*/IrBuilder::create<Val>(false),
+      /*scale=*/nullptr);
+
+  TensorView* attn = sdpa_out.output;
+  TensorView* attn_permute = permute(attn, {0, 2, 1, 3});
+  TensorView* out = reshape(attn_permute, out_shape, in_shape);
+
+  fusion->addInput(hq);
+  fusion->addInput(hk);
+  fusion->addInput(hv);
+  fusion->addOutput(out);
+
+  // Shard input tensors
+  for (auto* tv : {hq, hk, hv}) {
+    tv->setDeviceMesh(mesh);
+    tv->split(-1, d, /*inner_split=*/false);
+    tv->axis(-2)->parallelize(ParallelType::DIDx);
+    reorderDIDToFront(tv);
+  }
+  propagateShardings(fusion.get(), d);
+
+  for (auto tv : fusion->allTvs()) {
+    tv->setAllocationDomain(tv->getLoopDomain(), true);
+  }
+
+  FusionExecutorCache executor_cache(std::move(fusion));
+  at::Tensor hq_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));
+  at::Tensor hk_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));
+  at::Tensor hv_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));
+
+  at::Tensor sharded_hq = shardTensor(hq_tensor, -1, mesh);
+  at::Tensor sharded_hk = shardTensor(hk_tensor, -1, mesh);
+  at::Tensor sharded_hv = shardTensor(hv_tensor, -1, mesh);
+
+  at::Tensor nvf_out =
path: tests/cpp/test_multidevice_sharding.cpp
user: login: wujingyue
body: ```suggestion
  auto nvf_out =
```
the `.as` is clear about the type
created_at: 2025-03-11T22:55:29Z
author_association: COLLABORATOR
diff_hunk: @@ -807,4 +809,246 @@ TEST_F(MultiDeviceTest, ShardedSplitReshapeIds) {
       __FILE__);
 }
 
+TEST_F(MultiDeviceTest, LoopShardedSplitReshapeIds) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  const int d = communicator_->size();
+  const int64_t b = 2, s = 3, h = 8, e = 4;
+
+  TensorView* tv0 = makeContigConcreteTensor({b, s, d * h * e});
+  TensorView* tv1 = reshape(tv0, {b, s, d * h * e}, {b, s, d * h, e});
+
+  fusion->addInput(tv0);
+  fusion->addOutput(tv1);
+
+  auto mesh = DeviceMesh::createForNumDevices(d);
+
+  tv0->setDeviceMesh(mesh);
+  tv0->split(-1, d, /*inner_split=*/false);
+  tv0->axis(-2)->parallelize(ParallelType::DIDx);
+
+  tv1->setDeviceMesh(mesh);
+  tv1->split(-2, d, /*inner_split=*/false);
+  tv1->axis(-3)->parallelize(ParallelType::DIDx);
+
+  for (auto* tv : {tv0, tv1}) {
+    reorderDIDToFront(tv);
+    tv->setAllocationDomain(tv->getLoopDomain(), true);
+  }
+
+  FusionExecutorCache executor_cache(std::move(fusion));
+  at::Tensor inp = at::randn({b, s, d * h * e}, tensor_options);
+  at::Tensor sharded_inp = shardTensor(inp, -1, mesh);
+
+  at::Tensor nvf_out =
+      executor_cache.runFusionWithInputs({sharded_inp})[0].as<at::Tensor>();
+  testValidate(
+      executor_cache.fusion(),
+      {nvf_out},
+      {sharded_inp},
+      {sharded_inp.view({b, s, h, e})},
+      __LINE__,
+      __FILE__);
+}
+
+TEST_F(MultiDeviceTest, LoopShardedMergeReshapeIds) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  const int d = communicator_->size();
+  const int64_t b = 2, s = 3, h = 8, e = 4;
+
+  TensorView* tv0 = makeContigConcreteTensor({b, s, d * h, e});
+  TensorView* tv1 = reshape(tv0, {b, s, d * h, e}, {b, s, d * h * e});
+
+  fusion->addInput(tv0);
+  fusion->addOutput(tv1);
+
+  auto mesh = DeviceMesh::createForNumDevices(d);
+  tv0->setDeviceMesh(mesh);
+  tv0->split(-2, d, /*inner_split=*/false);
+  tv0->axis(-3)->parallelize(ParallelType::DIDx);
+
+  tv1->setDeviceMesh(mesh);
+  tv1->split(-1, d, /*inner_split=*/false);
+  tv1->axis(-2)->parallelize(ParallelType::DIDx);
+
+  for (auto* tv : {tv0, tv1}) {
+    reorderDIDToFront(tv);
+    tv->setAllocationDomain(tv->getLoopDomain(), true);
+  }
+
+  FusionExecutorCache executor_cache(std::move(fusion));
+  at::Tensor inp = at::randn({b, s, d * h, e}, tensor_options);
+  at::Tensor sharded_inp = shardTensor(inp, -2, mesh);
+  at::Tensor nvf_out =
+      executor_cache.runFusionWithInputs({sharded_inp})[0].as<at::Tensor>();
+  testValidate(
+      executor_cache.fusion(),
+      {nvf_out},
+      {sharded_inp},
+      {sharded_inp.view({b, s, h * e})},
+      __LINE__,
+      __FILE__);
+}
+
+namespace {
+// This is a simplified version of what we will eventually do in the
+// pre-segmentation pass
+void propagateShardings(Fusion* fusion, int64_t num_devices) {
+  for (Expr* expr : fusion->exprs()) {
+    if (expr->isA<SliceOp>()) {
+      NVF_THROW("SliceOp is not currently supported");
+    }
+
+    if (expr->isA<ViewOp>()) {
+      // TransformPropagator cannot be directly used.
+      // It raises an error for conflicting transformations from root domain to
+      // logical domain. Instead, we manually find the reshaped iterdomain and
+      // outer split DID. This might have to be extended further in the
+      // presegmentation pass.
+      // Note: For simplicity, this assumes that the sharding is on reshaped
+      // IDs. It is possible that the non-reshaped IDs are sharded, in which
+      // case we can use the TransformPropagator.
+      TensorView* reshaped_tv = expr->as<ViewOp>()->out();
+      auto transform_exprs = StmtSort::getExprsBetween(
+          {reshaped_tv->getMaybeRootDomain().begin(),
+           reshaped_tv->getMaybeRootDomain().end()},
+          {reshaped_tv->getLogicalDomain().begin(),
+           reshaped_tv->getLogicalDomain().end()});
+      NVF_CHECK(transform_exprs.size() == 1);
+      auto transform = transform_exprs[0];
+      NVF_CHECK(transform->isA<Split>() || transform->isA<Merge>());
+
+      // Get the reshaped ID (outer ID for split reshape).
+      // This is the ID that will be parallelized.
+      IterDomain* reshaped_id = transform->isA<Split>()
+          ? transform->as<Split>()->outer()
+          : transform->as<Merge>()->out();
+
+      auto reshaped_it = std::find(
+          reshaped_tv->getLoopDomain().begin(),
+          reshaped_tv->getLoopDomain().end(),
+          reshaped_id);
+      int64_t reshaped_axis =
+          std::distance(reshaped_tv->getLoopDomain().begin(), reshaped_it);
+
+      // Apply sharding to the reshaped tensor
+      reshaped_tv->split(reshaped_axis, num_devices, false);
+      reshaped_tv->axis(reshaped_axis)->parallelize(ParallelType::DIDx);
+      reorderDIDToFront(reshaped_tv);
+      continue;
+    }
+
+    // For other ops, propagate sharding from input to outputs
+    auto input_tv = expr->input(0)->as<TensorView>();
+    std::vector<TensorView*> output_tvs;
+    for (auto output : expr->outputs()) {
+      output_tvs.push_back(output->as<TensorView>());
+    }
+
+    TransformPropagator propagator(input_tv);
+
+    // Note: We will finally propagate from each input iteratively.
+    SetSelector selector(
+        std::unordered_set<TensorView*>(output_tvs.begin(), output_tvs.end()));
+    MaxLogicalDomainInfoSpanningTree(input_tv, &selector).traverse(&propagator);
+    scheduler_utils::parallelizeAllLike(
+        input_tv,
+        /*pos=*/-1,
+        /*selected_tv=*/output_tvs);
+  }
+}
+
+} // namespace
+
+TEST_F(MultiDeviceTest, TransformerFwd) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  const int d = communicator_->size();
+  const int64_t b = 2, s = 3, h = 8, e = 16;
+  auto mesh = DeviceMesh::createForNumDevices(d);
+
+  std::vector<int64_t> in_shape = {b, s, d * h * e};
+  std::vector<int64_t> out_shape = {b, s, d * h, e};
+
+  // The transformer block produces hq/hk/hv after slicing the MHA linear
+  // output.
+  TensorView* hq = makeConcreteTensor(in_shape, DataType::Half);
+  TensorView* hk = makeConcreteTensor(in_shape, DataType::Half);
+  TensorView* hv = makeConcreteTensor(in_shape, DataType::Half);
+
+  TensorView* q = reshape(hq, in_shape, out_shape);
+  TensorView* q_permuted = permute(q, {0, 2, 1, 3});
+  TensorView* k = reshape(hk, in_shape, out_shape);
+  TensorView* k_permuted = permute(k, {0, 2, 1, 3});
+  TensorView* v = reshape(hv, in_shape, out_shape);
+  TensorView* v_permuted = permute(v, {0, 2, 1, 3});
+
+  SdpfaFwdResult sdpa_out = sdpfa_fwd(
+      q_permuted,
+      k_permuted,
+      v_permuted,
+      /*dropout_p=*/IrBuilder::create<Val>(0.0),
+      /*is_causal=*/IrBuilder::create<Val>(false),
+      /*scale=*/nullptr);
+
+  TensorView* attn = sdpa_out.output;
+  TensorView* attn_permute = permute(attn, {0, 2, 1, 3});
+  TensorView* out = reshape(attn_permute, out_shape, in_shape);
+
+  fusion->addInput(hq);
+  fusion->addInput(hk);
+  fusion->addInput(hv);
+  fusion->addOutput(out);
+
+  // Shard input tensors
+  for (auto* tv : {hq, hk, hv}) {
+    tv->setDeviceMesh(mesh);
+    tv->split(-1, d, /*inner_split=*/false);
+    tv->axis(-2)->parallelize(ParallelType::DIDx);
+    reorderDIDToFront(tv);
+  }
+  propagateShardings(fusion.get(), d);
+
+  for (auto tv : fusion->allTvs()) {
+    tv->setAllocationDomain(tv->getLoopDomain(), true);
+  }
+
+  FusionExecutorCache executor_cache(std::move(fusion));
+  at::Tensor hq_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));
+  at::Tensor hk_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));
+  at::Tensor hv_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));
+
+  at::Tensor sharded_hq = shardTensor(hq_tensor, -1, mesh);
+  at::Tensor sharded_hk = shardTensor(hk_tensor, -1, mesh);
+  at::Tensor sharded_hv = shardTensor(hv_tensor, -1, mesh);
+
+  at::Tensor nvf_out =
+      executor_cache
+          .runFusionWithInputs({sharded_hq, sharded_hk, sharded_hv})[0]
+          .as<at::Tensor>();
+
+  double scale = 1.0 / std::sqrt(e);
+  auto reference_out = at::_scaled_dot_product_flash_attention(
+      hq_tensor.view(out_shape).transpose(1, 2),
+      hk_tensor.view(out_shape).transpose(1, 2),
+      hv_tensor.view(out_shape).transpose(1, 2),
+      /*dropout_p=*/0.0,
+      /*is_causal=*/false,
+      /*return_debug_mask=*/false,
+      /*scale=*/scale);
path: tests/cpp/test_multidevice_sharding.cpp
user: login: wujingyue
body: ```suggestion
      scale);
```
created_at: 2025-03-11T22:55:44Z
author_association: COLLABORATOR
diff_hunk: @@ -399,27 +399,90 @@ bool haveDifferentShardings(
     }
 
     auto is_mapped_in_id_model =
-        [](IterDomain* a, IterDomain* b, const IdModel& id_model) -> bool {
-      if (a == nullptr && b == nullptr) {
+        [producer, consumer, mapped_p_logical_ids, mapped_c_root_ids](
+            IterDomain* p_loop_id,
+            IterDomain* c_loop_id,
+            const IdModel& id_model) -> bool {
+      if (p_loop_id == nullptr && c_loop_id == nullptr) {
         return true;
       }
 
-      if (a == nullptr || b == nullptr) {
+      if (p_loop_id == nullptr || c_loop_id == nullptr) {
         return false;
       }
 
+      if (consumer->definition()->isA<ViewOp>()) {
+        // Consider the reshape: [h]-> [a, h/a] where `h` and `a` are sharded,
+        // IdModel currently does not map the two DID axes. Similarly for the
+        // corresponding merge reshape. This is a temporary workaround to check
+        // if the reshape is resharding when sharding is on reshaped ids.
+
+        // Get the root iterdomains in the consumer that produce reshape logical
+        // Ids.
+        std::unordered_set<IterDomain*> c_reshape_ids;
+        for (auto id : consumer->getLogicalDomain()) {
+          if (id->isRFactorProduct() && id->definition() &&
+              !id->definition()->isA<Resize>()) {
+            auto root_ids = getInputsInTargetDomain(
+                id, {mapped_c_root_ids.begin(), mapped_c_root_ids.end()});
+            for (auto root_id : root_ids) {
+              c_reshape_ids.insert(root_id);
+            }
+          }
+        }
+
+        // Get the logical iterdomains in the producer that are reshaped.
+        const std::unordered_map<IterDomain*, IterDomain*>& c2p =
+            PairwiseLogicalDomainMap(producer, consumer)
+                .mapConsumerToProducer();
+        std::unordered_set<IterDomain*> p_reshape_ids;
+        for (auto id : c_reshape_ids) {
+          auto p_id = c2p.find(id);
+          if (p_id != c2p.end()) {
+            p_reshape_ids.insert(p_id->second);
+          }
+        }
+
+        auto p_transforms = DependencyCheck::getAllExprsBetween(
+            {p_reshape_ids.begin(), p_reshape_ids.end()}, {p_loop_id});
+        auto c_transforms = DependencyCheck::getAllExprsBetween(
+            {c_reshape_ids.begin(), c_reshape_ids.end()}, {c_loop_id});
+
+        if (p_transforms.size() > 0 && c_transforms.size() > 0) {
path: csrc/multidevice/utils.cpp
user: login: wujingyue
body: SGTM! 
created_at: 2025-03-11T22:56:01Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-02-24T19:54:42Z
body: <!-- llm action combined persistent comment, dont modify this line -->

*Review updated until commit https://github.com/NVIDIA/Fuser/commit/5111615e1ad70cc16cf6796d83bdd27a4dab687b*


### **Description**
- Added WAR to identify non-resharding reshapes in DID loops.

- Enhanced `parallelizeAllLike` with `parallelize_inputs` flag.

- Added tests for split, merge reshapes, and transformer forward block.


___



### **Changes walkthrough** 📝
<table><thead><tr><th></th><th align="left">Relevant files</th></tr></thead><tbody><tr><td><strong>Enhancement</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>utils.cpp</strong><dd><code>Add WAR for non-resharding reshapes in DID loops</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/multidevice/utils.cpp

<li>Added logic to identify non-resharding reshapes in DID loops.<br> <li> Updated <code>is_mapped_in_id_model</code> to handle reshapes without resharding.


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3953/files#diff-4739b2bc50e3c078541f1ea6665370ea6dd853d1fa0b4d1085b95d7c4791d96d">+68/-6</a>&nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>utils.cpp</strong><dd><code>Add `parallelize_inputs` flag to `parallelizeAllLike`</code>&nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/utils.cpp

<li>Added <code>parallelize_inputs</code> flag to <code>parallelizeAllLike</code>.<br> <li> Updated <code>propagateReshapeTransforms</code> to manually replay DID <br>parallelization.


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3953/files#diff-078826ec098dc190fe9aef215a8130a7f5d268dc36f9537b2bb9e6f255d589d6">+13/-2</a>&nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>utils.h</strong><dd><code>Update `parallelizeAllLike` with `parallelize_inputs`</code>&nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/utils.h

<li>Updated <code>parallelizeAllLike</code> function signature to include <br><code>parallelize_inputs</code>.


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3953/files#diff-ee71c807e8995e215fce4d916b5de8ee271959c631f37a4c368c03636bd1337c">+19/-7</a>&nbsp; &nbsp; </td>

</tr>
</table></td></tr><tr><td><strong>Tests</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>test_multidevice_sharding.cpp</strong><dd><code>Add tests for reshape propagation and transformer forward block</code></dd></summary>
<hr>

tests/cpp/test_multidevice_sharding.cpp

<li>Added <code>TransformPropagatorSplitReshape</code> test.<br> <li> Added <code>LoopShardedSplitReshapeIds</code> and <code>LoopShardedMergeReshapeIds</code> tests.<br> <li> Added <code>TransformerFwd</code> test for transformer forward block.


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3953/files#diff-c5175dc597d3e4dd05ea38630935938b69c33328284a9e3b689a60ccaaef5b7b">+245/-1</a>&nbsp; </td>

</tr>
</table></td></tr></tr></tbody></table>

---

### PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

<table>
<tr><td>🧪&nbsp;<strong>PR contains tests</strong></td></tr>
<tr><td>⚡&nbsp;<strong>Recommended focus areas for review</strong><br><br>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/3953/files#diff-4739b2bc50e3c078541f1ea6665370ea6dd853d1fa0b4d1085b95d7c4791d96dR401-R472'><strong>Performance Impact</strong></a>

The new logic for identifying reshapes as non-resharding may have performance implications. Ensure that this change does not introduce unnecessary computations or overhead.</summary>

```c++
auto is_mapped_in_id_model = [producer, consumer, mapped_c_root_ids](
                                 IterDomain* p_loop_id,
                                 IterDomain* c_loop_id,
                                 const IdModel& id_model) -> bool {
  if (p_loop_id == nullptr && c_loop_id == nullptr) {
    return true;
  }

  if (p_loop_id == nullptr || c_loop_id == nullptr) {
    return false;
  }

  if (consumer->definition()->isA<ViewOp>()) {
    // Consider the reshape: [h]-> [a, h/a] where `h` and `a` are sharded,
    // IdModel currently does not map the two DID axes. Similarly for the
    // corresponding merge reshape. This is a temporary workaround to check
    // if the reshape is resharding when sharding is on reshaped ids.

    // Get the root iterdomains in the consumer that produce reshape logical
    // Ids.
    std::unordered_set<IterDomain*> c_reshape_ids;
    for (auto id : consumer->getLogicalDomain()) {
      if (id->isRFactorProduct() && id->definition() &&
          !id->definition()->isA<Resize>()) {
        auto root_ids = getInputsInTargetDomain(
            id, {mapped_c_root_ids.begin(), mapped_c_root_ids.end()});
        for (auto root_id : root_ids) {
          c_reshape_ids.insert(root_id);
        }
      }
    }

    // Get the logical iterdomains in the producer that are reshaped.
    const std::unordered_map<IterDomain*, IterDomain*>& c2p =
        PairwiseLogicalDomainMap(producer, consumer)
            .mapConsumerToProducer();
    std::unordered_set<IterDomain*> p_reshape_ids;
    for (auto id : c_reshape_ids) {
      auto p_id = c2p.find(id);
      if (p_id != c2p.end()) {
        p_reshape_ids.insert(p_id->second);
      }
    }

    auto p_transforms = DependencyCheck::getAllExprsBetween(
        {p_reshape_ids.begin(), p_reshape_ids.end()}, {p_loop_id});
    auto c_transforms = DependencyCheck::getAllExprsBetween(
        {c_reshape_ids.begin(), c_reshape_ids.end()}, {c_loop_id});

    if (p_transforms.size() > 0 && c_transforms.size() > 0) {
      // Both a and b come from reshaped ids.
      NVF_ERROR(
          p_loop_id->definition()->isA<Split>(),
          p_loop_id,
          " is not a Split: ",
          p_loop_id->definition());
      NVF_ERROR(
          c_loop_id->definition()->isA<Split>(),
          c_loop_id,
          " is not a Split: ",
          c_loop_id->definition());
      // Get the split producing the sharded axis
      auto* p_split = p_transforms.back()->as<Split>();
      auto* c_split = c_transforms.back()->as<Split>();

      // Reshape is not resharding if both the DID splits are either inner
      // or outer splits. Note that we currently do not exercise inner
      // splits in the multidevice tests.
      return p_split->innerSplit() == c_split->innerSplit();
    }
  }

```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/3953/files#diff-078826ec098dc190fe9aef215a8130a7f5d268dc36f9537b2bb9e6f255d589d6R292-R324'><strong>Code Complexity</strong></a>

The addition of new parameters and logic in `parallelizeAllLike` increases the complexity of the function. Consider refactoring to improve readability and maintainability.</summary>

```c++
void parallelizeAllLike(
    TensorView* reference_tv,
    int64_t pos,
    std::vector<TensorView*> selected_tvs,
    const std::unordered_set<ParallelType>& selected_parallel_types,
    bool propagate_padding,
    bool parallelize_inputs) {
  FusionGuard fg(reference_tv->fusion());

  if (pos < 0) {
    pos += (int64_t)reference_tv->nDims() + 1;
  }
  NVF_CHECK(
      pos >= 0 && pos <= (int64_t)reference_tv->nDims(),
      "parallelizeAllLike called on an position outside valid range.");

  std::unordered_map<IterDomain*, IterDomain*> concrete_to_reference_map;

  auto ca_map = ComputeAtMap(FusionGuard::getCurFusion());

  const auto& reference_dom = reference_tv->getLoopDomain();
  for (auto it = reference_dom.begin(); it != reference_dom.begin() + pos;
       it++) {
    auto ca_id =
        ca_map.getConcreteMappedID(*it, IdMappingMode::PERMISSIVE_RESIZE);
    concrete_to_reference_map[ca_id] = *it;
  }

  if (selected_tvs.empty()) {
    selected_tvs = reference_tv->fusion()->allTvs();
  }
  for (auto tv : selected_tvs) {
    if (tv->isFusionInput() && !parallelize_inputs) {
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/3953/files#diff-c5175dc597d3e4dd05ea38630935938b69c33328284a9e3b689a60ccaaef5b7bR748-R1054'><strong>Test Coverage</strong></a>

The new tests cover a variety of scenarios, but ensure that edge cases and potential failure points are also tested.</summary>

```c++
TEST_F(MultiDeviceTest, TransformPropagatorSplitReshape) {
  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  const int d = communicator_->size();
  const int64_t b = 2, s = 2, h = 4, e = 3;

  TensorView* tv0 = makeContigConcreteTensor(
      {b, s, d * h * e}); // in: loop domain: {b, s, d*h*e}
  TensorView* tv1 = reshape(
      tv0,
      {b, s, d * h * e},
      {b, s, d * h, e}); // out: loop domain: {b, s, d*h, e}

  fusion->addInput(tv0);
  fusion->addOutput(tv1);

  auto mesh = DeviceMesh::createForNumDevices(d);

  // Propagate transform from reshaped output to input.
  // Without this propagation, the two DID axes on `in` and `out` will not be
  // mapped in together in ID model. This causes scheduling to fail due to
  // resharding.
  TransformPropagator propagator_c2p(tv1);
  MaxLogicalDomainInfoSpanningTree(tv1).traverse(&propagator_c2p);
  // in: loop domain: {b, s, d*h, e} after transform propagation

  // Loop split and parallelize input
  tv0->setDeviceMesh(mesh);
  tv0->split(-2, d, /*inner_split=*/false);
  tv0->axis(-3)->parallelize(ParallelType::DIDx);
  // in: loop domain: {b, s, DIDx{d}, h, e}

  // Propagate DID loop split to output
  TransformPropagator propagator_p2c(tv0);
  MaxLogicalDomainInfoSpanningTree(tv0).traverse(&propagator_p2c);
  // out: loop domain: {b, s, d, h, e} after transform propagation

  // Parallelize output
  scheduler_utils::parallelizeAllLike(
      tv0,
      /*pos=*/-1,
      /*selected_tv=*/{tv1});
  // out: loop domain: {b, s, DIDx{d}, h, e} after parallelization

  tv0->setAllocationDomain(tv0->getLoopDomain(), true);
  tv1->setAllocationDomain(tv1->getLoopDomain(), true);

  FusionExecutorCache executor_cache(std::move(fusion));
  at::Tensor inp = at::randn({b, s, d * h * e}, tensor_options);
  at::Tensor sharded_inp = shardTensor(inp, tv0);

  at::Tensor nvf_out =
      executor_cache.runFusionWithInputs({sharded_inp})[0].as<at::Tensor>();

  testValidate(
      executor_cache.fusion(),
      {nvf_out},
      {sharded_inp},
      {sharded_inp.view({b, s, h, e})},
      __LINE__,
      __FILE__);
}

TEST_F(MultiDeviceTest, LoopShardedSplitReshapeIds) {
  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  const int d = communicator_->size();
  const int64_t b = 2, s = 3, h = 8, e = 4;

  TensorView* tv0 = makeContigConcreteTensor({b, s, d * h * e});
  TensorView* tv1 = reshape(tv0, {b, s, d * h * e}, {b, s, d * h, e});

  fusion->addInput(tv0);
  fusion->addOutput(tv1);

  auto mesh = DeviceMesh::createForNumDevices(d);

  tv0->setDeviceMesh(mesh);
  tv0->split(-1, d, /*inner_split=*/false);
  tv0->axis(-2)->parallelize(ParallelType::DIDx);

  tv1->setDeviceMesh(mesh);
  tv1->split(-2, d, /*inner_split=*/false);
  tv1->axis(-3)->parallelize(ParallelType::DIDx);

  for (auto* tv : {tv0, tv1}) {
    reorderDIDToFront(tv);
    tv->setAllocationDomain(tv->getLoopDomain(), true);
  }

  FusionExecutorCache executor_cache(std::move(fusion));
  at::Tensor inp = at::randn({b, s, d * h * e}, tensor_options);
  at::Tensor sharded_inp = shardTensor(inp, -1, mesh);

  at::Tensor nvf_out =
      executor_cache.runFusionWithInputs({sharded_inp})[0].as<at::Tensor>();
  testValidate(
      executor_cache.fusion(),
      {nvf_out},
      {sharded_inp},
      {sharded_inp.view({b, s, h, e})},
      __LINE__,
      __FILE__);
}

TEST_F(MultiDeviceTest, LoopShardedMergeReshapeIds) {
  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  const int d = communicator_->size();
  const int64_t b = 2, s = 3, h = 8, e = 4;

  TensorView* tv0 = makeContigConcreteTensor({b, s, d * h, e});
  TensorView* tv1 = reshape(tv0, {b, s, d * h, e}, {b, s, d * h * e});

  fusion->addInput(tv0);
  fusion->addOutput(tv1);

  auto mesh = DeviceMesh::createForNumDevices(d);
  tv0->setDeviceMesh(mesh);
  tv0->split(-2, d, /*inner_split=*/false);
  tv0->axis(-3)->parallelize(ParallelType::DIDx);

  tv1->setDeviceMesh(mesh);
  tv1->split(-1, d, /*inner_split=*/false);
  tv1->axis(-2)->parallelize(ParallelType::DIDx);

  for (auto* tv : {tv0, tv1}) {
    reorderDIDToFront(tv);
    tv->setAllocationDomain(tv->getLoopDomain(), true);
  }

  FusionExecutorCache executor_cache(std::move(fusion));
  at::Tensor inp = at::randn({b, s, d * h, e}, tensor_options);
  at::Tensor sharded_inp = shardTensor(inp, -2, mesh);
  at::Tensor nvf_out =
      executor_cache.runFusionWithInputs({sharded_inp})[0].as<at::Tensor>();
  testValidate(
      executor_cache.fusion(),
      {nvf_out},
      {sharded_inp},
      {sharded_inp.view({b, s, h * e})},
      __LINE__,
      __FILE__);
}

namespace {
// This is a simplified version of what we will eventually do in the
// pre-segmentation pass
void propagateShardings(Fusion* fusion, int64_t num_devices) {
  for (Expr* expr : fusion->exprs()) {
    if (expr->isA<SliceOp>()) {
      NVF_THROW("SliceOp is not currently supported");
    }

    if (expr->isA<ViewOp>()) {
      // TransformPropagator cannot be directly used.
      // It raises an error for conflicting transformations from root domain to
      // logical domain. Instead, we manually find the reshaped iterdomain and
      // outer split DID. This might have to be extended further in the
      // presegmentation pass.
      // Note: For simplicity, this assumes that the sharding is on reshaped
      // IDs. It is possible that the non-reshaped IDs are sharded, in which
      // case we can use the TransformPropagator.
      TensorView* reshaped_tv = expr->as<ViewOp>()->out();
      auto transform_exprs = StmtSort::getExprsBetween(
          {reshaped_tv->getMaybeRootDomain().begin(),
           reshaped_tv->getMaybeRootDomain().end()},
          {reshaped_tv->getLogicalDomain().begin(),
           reshaped_tv->getLogicalDomain().end()});
      NVF_CHECK(transform_exprs.size() == 1);
      auto transform = transform_exprs[0];
      NVF_CHECK(transform->isA<Split>() || transform->isA<Merge>());

      // Get the reshaped ID (outer ID for split reshape).
      // This is the ID that will be parallelized.
      IterDomain* reshaped_id = transform->isA<Split>()
          ? transform->as<Split>()->outer()
          : transform->as<Merge>()->out();

      auto reshaped_it = std::find(
          reshaped_tv->getLoopDomain().begin(),
          reshaped_tv->getLoopDomain().end(),
          reshaped_id);
      int64_t reshaped_axis =
          std::distance(reshaped_tv->getLoopDomain().begin(), reshaped_it);

      // Apply sharding to the reshaped tensor
      reshaped_tv->split(reshaped_axis, num_devices, false);
      reshaped_tv->axis(reshaped_axis)->parallelize(ParallelType::DIDx);
      reorderDIDToFront(reshaped_tv);
      continue;
    }

    // For other ops, propagate sharding from input to outputs
    auto input_tv = expr->input(0)->as<TensorView>();
    std::vector<TensorView*> output_tvs;
    for (auto output : expr->outputs()) {
      output_tvs.push_back(output->as<TensorView>());
    }

    TransformPropagator propagator(input_tv);

    // Note: We will finally propagate from each input iteratively.
    SetSelector selector(
        std::unordered_set<TensorView*>(output_tvs.begin(), output_tvs.end()));
    MaxLogicalDomainInfoSpanningTree(input_tv, &selector).traverse(&propagator);
    scheduler_utils::parallelizeAllLike(
        input_tv,
        /*pos=*/-1,
        /*selected_tv=*/output_tvs);
  }
}

} // namespace

TEST_F(MultiDeviceTest, TransformerFwd) {
  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  const int d = communicator_->size();
  const int64_t b = 2, s = 3, h = 8, e = 16;
  auto mesh = DeviceMesh::createForNumDevices(d);

  std::vector<int64_t> in_shape = {b, s, d * h * e};
  std::vector<int64_t> out_shape = {b, s, d * h, e};

  // The transformer block produces hq/hk/hv after slicing the MHA linear
  // output.
  TensorView* hq = makeConcreteTensor(in_shape, DataType::Half);
  TensorView* hk = makeConcreteTensor(in_shape, DataType::Half);
  TensorView* hv = makeConcreteTensor(in_shape, DataType::Half);

  TensorView* q = reshape(hq, in_shape, out_shape);
  TensorView* q_permuted = permute(q, {0, 2, 1, 3});
  TensorView* k = reshape(hk, in_shape, out_shape);
  TensorView* k_permuted = permute(k, {0, 2, 1, 3});
  TensorView* v = reshape(hv, in_shape, out_shape);
  TensorView* v_permuted = permute(v, {0, 2, 1, 3});

  SdpfaFwdResult sdpa_out = sdpfa_fwd(
      q_permuted,
      k_permuted,
      v_permuted,
      /*dropout_p=*/IrBuilder::create<Val>(0.0),
      /*is_causal=*/IrBuilder::create<Val>(false),
      /*scale=*/nullptr);

  TensorView* attn = sdpa_out.output;
  TensorView* attn_permute = permute(attn, {0, 2, 1, 3});
  TensorView* out = reshape(attn_permute, out_shape, in_shape);

  fusion->addInput(hq);
  fusion->addInput(hk);
  fusion->addInput(hv);
  fusion->addOutput(out);

  // Shard input tensors
  for (auto* tv : {hq, hk, hv}) {
    tv->setDeviceMesh(mesh);
    tv->split(-1, d, /*inner_split=*/false);
    tv->axis(-2)->parallelize(ParallelType::DIDx);
    reorderDIDToFront(tv);
  }
  propagateShardings(fusion.get(), d);

  for (auto tv : fusion->allTvs()) {
    tv->setAllocationDomain(tv->getLoopDomain(), true);
  }

  FusionExecutorCache executor_cache(std::move(fusion));
  at::Tensor hq_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));
  at::Tensor hk_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));
  at::Tensor hv_tensor = at::randn({in_shape}, tensor_options.dtype(at::kHalf));

  at::Tensor sharded_hq = shardTensor(hq_tensor, -1, mesh);
  at::Tensor sharded_hk = shardTensor(hk_tensor, -1, mesh);
  at::Tensor sharded_hv = shardTensor(hv_tensor, -1, mesh);

  auto nvf_out =
      executor_cache
          .runFusionWithInputs({sharded_hq, sharded_hk, sharded_hv})[0]
          .as<at::Tensor>();

  double scale = 1.0 / std::sqrt(e);
  auto reference_out = at::_scaled_dot_product_flash_attention(
      hq_tensor.view(out_shape).transpose(1, 2),
      hk_tensor.view(out_shape).transpose(1, 2),
      hv_tensor.view(out_shape).transpose(1, 2),
      /*dropout_p=*/0.0,
      /*is_causal=*/false,
      /*return_debug_mask=*/false,
      scale);
  at::Tensor ref_attn = shardTensor(
      std::get<0>(reference_out).transpose(1, 2).view(in_shape), -1, mesh);

  testValidate(
      executor_cache.fusion(),
      {nvf_out},
      {sharded_hq, sharded_hk, sharded_hv},
      {ref_attn},
      __LINE__,
      __FILE__);
}
} // namespace nvfuser
```

</details>

</td></tr>
</table>

author_association: NONE
user: login: Priya2698
created_at: 2025-03-04T23:20:39Z
body: !test 
author_association: COLLABORATOR
user: login: Priya2698
created_at: 2025-03-05T00:15:32Z
body: !test
author_association: COLLABORATOR
user: login: Priya2698
created_at: 2025-03-05T00:19:47Z
body: @wujingyue some of the changes in this PR are duplicates and will be superseded by PR #3875. I have pointed those out above.
author_association: COLLABORATOR
user: login: Priya2698
created_at: 2025-03-05T01:37:48Z
body: !test
author_association: COLLABORATOR
user: login: Priya2698
created_at: 2025-03-07T01:19:41Z
body: !test
author_association: COLLABORATOR
user: login: Priya2698
created_at: 2025-03-07T20:12:38Z
body: !test
author_association: COLLABORATOR
user: login: Priya2698
created_at: 2025-03-18T18:10:45Z
body: !test
author_association: COLLABORATOR