diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
index 5fe3680da9b8d6..f765dad0640671 100644
--- a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
@@ -645,6 +645,212 @@ int64_t GetShapeSizeBytes(const CostAnalysis* cost_analysis,
   return cost_analysis ? cost_analysis->GetShapeSizeBytes(shape) : 0;
 }
 
+void FindAliases(std::vector<AllocationValue>* allocation_values,
+                 bool has_async_pipelined_while_loops = false) {
+  absl::flat_hash_map<const HloInstruction*,
+                      std::vector<const AllocationValue*>>
+      values_by_defining_inst;
+  for (AllocationValue& value : *allocation_values) {
+    values_by_defining_inst[value.position().instruction].push_back(&value);
+  }
+  auto maybe_add_alias_with_instruction =
+      [&](const HloInstruction* instruction, AllocationValue::Use* use,
+          std::optional<ShapeIndex> expected_index = std::nullopt) {
+        auto aliased_values_it = values_by_defining_inst.find(instruction);
+        if (aliased_values_it != values_by_defining_inst.end()) {
+          for (const AllocationValue* aliased_value :
+               aliased_values_it->second) {
+            // Ensure that only matching tuple shape indexes are linked together
+            // as aliases when an expected shape index is specified (e.g. across
+            // while loop boundaries or async chain tuples).
+            if (expected_index.has_value() &&
+                !aliased_value->position().index.empty() &&
+                aliased_value->position().index != *expected_index) {
+              continue;
+            }
+            if (absl::c_find(use->aliases, aliased_value->position()) ==
+                use->aliases.end()) {
+              VLOG(3) << "Adding aliasing for use " << use->hlo_use.ToString()
+                      << " to " << aliased_value->ToShortString();
+              use->aliases.push_back(aliased_value->position());
+            }
+          }
+        }
+      };
+
+  for (AllocationValue& value : *allocation_values) {
+    for (AllocationValue::Use& use : value.uses()) {
+      // Find any aliases with the instruction itself (operand and output must
+      // alias), and any aliases with the parameters of called computations.
+      if (use.hlo_use.instruction->IsAsynchronous()) {
+        HloComputation* wrapped_computation =
+            use.hlo_use.instruction->async_wrapped_computation();
+        if (use.hlo_use.instruction->opcode() == HloOpcode::kAsyncStart) {
+          // Bound operand i of async-start is placed at tuple index {0, i} of
+          // async-start's output shape.
+          ShapeIndex expected_async_index = {0, use.hlo_use.operand_number};
+          expected_async_index.insert(expected_async_index.end(),
+                                      use.hlo_use.operand_index.begin(),
+                                      use.hlo_use.operand_index.end());
+          maybe_add_alias_with_instruction(use.hlo_use.instruction, &use,
+                                           expected_async_index);
+          if (use.hlo_use.instruction->async_wrapped_opcode() ==
+                  HloOpcode::kDynamicUpdateSlice &&
+              use.hlo_use.operand_number == 0) {
+            ShapeIndex expected_dus_index = {1};
+            expected_dus_index.insert(expected_dus_index.end(),
+                                      use.hlo_use.operand_index.begin(),
+                                      use.hlo_use.operand_index.end());
+            maybe_add_alias_with_instruction(use.hlo_use.instruction, &use,
+                                             expected_dus_index);
+          }
+          // Operands bound to async-start map directly to the initial
+          // parameters of the wrapped computation 1-to-1 by operand index.
+          if (use.hlo_use.operand_number <
+              wrapped_computation->num_parameters()) {
+            maybe_add_alias_with_instruction(
+                wrapped_computation->parameter_instruction(
+                    use.hlo_use.operand_number),
+                &use);
+          }
+        } else if (use.hlo_use.instruction->opcode() ==
+                   HloOpcode::kAsyncUpdate) {
+          // For async-update, operand 0 is the async context tuple (the result
+          // of async-start or a prior async-update), which forwards matching
+          // tuple indices and does not correspond to a computation parameter.
+          // Operands 1, 2, ... are late-bound arguments mapping to wrapped
+          // computation parameters offset by the number of operands already
+          // bound by previous instructions in the async chain, and placed at
+          // tuple index {0, param_idx} of async-update's output shape.
+          if (use.hlo_use.operand_number == 0) {
+            maybe_add_alias_with_instruction(use.hlo_use.instruction, &use,
+                                             use.hlo_use.operand_index);
+          } else if (use.hlo_use.operand_number >= 1) {
+            const HloInstruction* operand0 =
+                use.hlo_use.instruction->operand(0);
+            CHECK(operand0 != nullptr) << "Operand 0 is null for async update: "
+                                       << use.hlo_use.instruction->ToString();
+            CHECK(operand0->opcode() == HloOpcode::kAsyncStart ||
+                  operand0->opcode() == HloOpcode::kAsyncUpdate)
+                << "Unexpected operand for async update: "
+                << use.hlo_use.instruction->ToString();
+            std::vector<const HloInstruction*> previously_bound_operands =
+                hlo_instruction_utils::async::GetAsyncBoundOperands(
+                    Cast<HloAsyncInstruction>(operand0));
+            int previously_bound_operand_count =
+                previously_bound_operands.size();
+            int param_idx =
+                use.hlo_use.operand_number - 1 + previously_bound_operand_count;
+            ShapeIndex expected_async_index = {0, param_idx};
+            expected_async_index.insert(expected_async_index.end(),
+                                        use.hlo_use.operand_index.begin(),
+                                        use.hlo_use.operand_index.end());
+            maybe_add_alias_with_instruction(use.hlo_use.instruction, &use,
+                                             expected_async_index);
+            if (param_idx < wrapped_computation->num_parameters()) {
+              maybe_add_alias_with_instruction(
+                  wrapped_computation->parameter_instruction(param_idx), &use);
+            }
+          }
+        } else {
+          maybe_add_alias_with_instruction(use.hlo_use.instruction, &use);
+        }
+      } else if (use.hlo_use.instruction->opcode() == HloOpcode::kConditional) {
+        maybe_add_alias_with_instruction(use.hlo_use.instruction, &use);
+        // A conditional instruction has branch_count + 1 operands: operand 0 is
+        // the branch index/predicate (which does not correspond to any branch
+        // computation parameter), and each operand i >= 1 is the single
+        // argument forwarded to branch_computation(i - 1)'s parameter 0.
+        if (use.hlo_use.operand_number >= 1) {
+          const HloComputation* branch_computation =
+              use.hlo_use.instruction->branch_computation(
+                  use.hlo_use.operand_number - 1);
+          if (branch_computation->num_parameters() > 0) {
+            maybe_add_alias_with_instruction(
+                branch_computation->parameter_instruction(0), &use);
+          }
+        }
+      } else {
+        std::optional<ShapeIndex> expected_index = std::nullopt;
+        if (use.hlo_use.instruction->opcode() == HloOpcode::kWhile) {
+          expected_index = use.hlo_use.operand_index;
+        }
+        maybe_add_alias_with_instruction(use.hlo_use.instruction, &use,
+                                         expected_index);
+        // For synchronous instructions with called computations (such as kCall
+        // and kWhile), caller operand i maps 1-to-1 to parameter i of each
+        // called computation.
+        for (const HloComputation* called_computation :
+             use.hlo_use.instruction->called_computations()) {
+          if (use.hlo_use.operand_number <
+              called_computation->num_parameters()) {
+            maybe_add_alias_with_instruction(
+                called_computation->parameter_instruction(
+                    use.hlo_use.operand_number),
+                &use);
+          }
+        }
+      }
+
+      // Special case for kWhile: the root of the body computation must alias as
+      // well.
+      if (use.hlo_use.instruction->opcode() == HloOpcode::kWhile) {
+        HloPosition root_alias{
+            use.hlo_use.instruction->while_body()->root_instruction(),
+            use.hlo_use.operand_index};
+        VLOG(3) << "Adding while body root aliasing for use "
+                << use.hlo_use.ToString() << " to " << root_alias;
+        use.aliases.push_back(root_alias);
+      }
+
+      // Special case for conditionals - the output of a conditional op must
+      // alias with the branch computation outputs.
+      HloInstruction* conditional_instruction =
+          GetConditionalForBranchRoot(use.hlo_use.instruction);
+      if (conditional_instruction != nullptr &&
+          use.hlo_use.instruction->opcode() == HloOpcode::kTuple) {
+        // We only need to add a use alias if the branch root is a tuple,
+        // because a tuple is a use and any other instruction would be a
+        // definition or a position.
+        ShapeIndex index = use.hlo_use.operand_index;
+        index.push_front(use.hlo_use.operand_number);
+        HloPosition conditional_output_position{conditional_instruction, index};
+        VLOG(1) << "Add use alias for counditional output position "
+                << conditional_output_position.ToString() << " to use "
+                << use.hlo_use.ToString();
+        use.aliases.push_back(conditional_output_position);
+      }
+
+      // Special case for async pipelined while loops: an instruction
+      // producing a loop return operand that feeds into the body root tuple
+      // must alias with both the while body parameter(0) and the caller while
+      // instruction at the corresponding tuple index.
+      if (has_async_pipelined_while_loops) {
+        HloInstruction* while_instruction =
+            GetWhileForBodyRoot(use.hlo_use.instruction);
+        if (while_instruction != nullptr &&
+            IsAsyncPipelinedWhileLoop(while_instruction) &&
+            use.hlo_use.instruction->opcode() == HloOpcode::kTuple &&
+            while_instruction->while_body()->num_parameters() > 0) {
+          ShapeIndex index = use.hlo_use.operand_index;
+          index.push_front(use.hlo_use.operand_number);
+          HloInstruction* parameter_instruction =
+              while_instruction->while_body()->parameter_instruction(0);
+          HloPosition parameter_position{parameter_instruction, index};
+          HloPosition while_output_position{while_instruction, index};
+          VLOG(1)
+              << "Add use aliases for while body root position to parameter "
+              << parameter_position.ToString() << " and while output "
+              << while_output_position.ToString() << " for use "
+              << use.hlo_use.ToString();
+          use.aliases.push_back(parameter_position);
+          use.aliases.push_back(while_output_position);
+        }
+      }
+    }
+  }
+}
+
 }  // namespace
 
 // Helper to inspect the async wrapped opcode of a pipelined while loop
@@ -1107,145 +1313,10 @@ void MsaAlgorithm::CreateAllocationValues(
   }
 }
 
-void MsaAlgorithm::FindAliases(
-    std::vector<AllocationValue>* allocation_values) const {
-  absl::flat_hash_map<const HloInstruction*,
-                      std::vector<const AllocationValue*>>
-      values_by_defining_inst;
-  for (AllocationValue& value : *allocation_values) {
-    values_by_defining_inst[value.position().instruction].push_back(&value);
-  }
-  auto maybe_add_alias_with_instruction = [&](const HloInstruction* instruction,
-                                              AllocationValue::Use* use) {
-    auto aliased_values_it = values_by_defining_inst.find(instruction);
-    if (aliased_values_it != values_by_defining_inst.end()) {
-      for (const AllocationValue* aliased_value : aliased_values_it->second) {
-        // When aliasing while loop boundaries, ensure that only matching tuple
-        // shape indexes are linked together as aliases.
-        if (use->hlo_use.instruction->opcode() == HloOpcode::kWhile &&
-            !aliased_value->position().index.empty() &&
-            aliased_value->position().index != use->hlo_use.operand_index) {
-          continue;
-        }
-        if (absl::c_find(use->aliases, aliased_value->position()) ==
-            use->aliases.end()) {
-          VLOG(3) << "Adding aliasing for use " << use->hlo_use.ToString()
-                  << " to " << aliased_value->ToShortString();
-          use->aliases.push_back(aliased_value->position());
-        }
-      }
-    }
-  };
-
-  for (AllocationValue& value : *allocation_values) {
-    for (AllocationValue::Use& use : value.uses()) {
-      // Find any aliases with the instruction itself (operand and output must
-      // alias).
-      maybe_add_alias_with_instruction(use.hlo_use.instruction, &use);
-
-      // Find any aliases with the parameters of called computations.
-      if (use.hlo_use.instruction->IsAsynchronous()) {
-        HloComputation* wrapped_computation =
-            use.hlo_use.instruction->async_wrapped_computation();
-        if (use.hlo_use.instruction->opcode() == HloOpcode::kAsyncStart &&
-            use.hlo_use.instruction->operand_count() > 0) {
-          // Operands bound with async-start map directly to the initial
-          // parameters of the wrapped computation.
-          for (int i = 0; i < use.hlo_use.instruction->operand_count(); ++i) {
-            maybe_add_alias_with_instruction(
-                wrapped_computation->parameter_instruction(i), &use);
-          }
-        } else if (use.hlo_use.instruction->opcode() ==
-                       HloOpcode::kAsyncUpdate &&
-                   use.hlo_use.instruction->operand_count() > 1) {
-          // Operands bound with async-update map to parameters of the
-          // wrapped computation offset by the number of operands
-          // that were bound by previous instructions.
-          const HloInstruction* operand0 = use.hlo_use.instruction->operand(0);
-          CHECK(operand0 != nullptr) << "Operand 0 is null for async update: "
-                                     << use.hlo_use.instruction->ToString();
-          CHECK(operand0->opcode() == HloOpcode::kAsyncStart ||
-                operand0->opcode() == HloOpcode::kAsyncUpdate)
-              << "Unexpected operand for async update: "
-              << use.hlo_use.instruction->ToString();
-          std::vector<const HloInstruction*> previously_bound_operands =
-              hlo_instruction_utils::async::GetAsyncBoundOperands(
-                  Cast<HloAsyncInstruction>(operand0));
-          int previously_bound_operand_count = previously_bound_operands.size();
-          for (int i = 1; i < use.hlo_use.instruction->operand_count(); ++i) {
-            maybe_add_alias_with_instruction(
-                wrapped_computation->parameter_instruction(
-                    i - 1 + previously_bound_operand_count),
-                &use);
-          }
-        }
-      } else {
-        for (const HloComputation* called_computation :
-             use.hlo_use.instruction->called_computations()) {
-          for (const HloInstruction* parameter_instruction :
-               called_computation->parameter_instructions()) {
-            maybe_add_alias_with_instruction(parameter_instruction, &use);
-          }
-        }
-      }
-
-      // Special case for kWhile: the root of the body computation must alias as
-      // well.
-      if (use.hlo_use.instruction->opcode() == HloOpcode::kWhile) {
-        HloPosition root_alias{
-            use.hlo_use.instruction->while_body()->root_instruction(),
-            use.hlo_use.operand_index};
-        VLOG(3) << "Adding while body root aliasing for use "
-                << use.hlo_use.ToString() << " to " << root_alias;
-        use.aliases.push_back(root_alias);
-      }
-
-      // Special case for conditionals - the output of a conditional op must
-      // alias with the branch computation outputs.
-      HloInstruction* conditional_instruction =
-          GetConditionalForBranchRoot(use.hlo_use.instruction);
-      if (conditional_instruction != nullptr &&
-          use.hlo_use.instruction->opcode() == HloOpcode::kTuple) {
-        // We only need to add a use alias if the branch root is a tuple,
-        // because a tuple is a use and any other instruction would be a
-        // definition or a position.
-        ShapeIndex index = use.hlo_use.operand_index;
-        index.push_front(use.hlo_use.operand_number);
-        HloPosition conditional_output_position{conditional_instruction, index};
-        VLOG(1) << "Add use alias for counditional output position "
-                << conditional_output_position.ToString() << " to use "
-                << use.hlo_use.ToString();
-        use.aliases.push_back(conditional_output_position);
-      }
-
-      // Special case for async pipelined while loops: an instruction
-      // producing a loop return operand that feeds into the body root tuple
-      // must alias with both the while body parameter(0) and the caller while
-      // instruction at the corresponding tuple index.
-      if (has_async_pipelined_while_loops_) {
-        HloInstruction* while_instruction =
-            GetWhileForBodyRoot(use.hlo_use.instruction);
-        if (while_instruction != nullptr &&
-            IsAsyncPipelinedWhileLoop(while_instruction) &&
-            use.hlo_use.instruction->opcode() == HloOpcode::kTuple &&
-            while_instruction->while_body()->num_parameters() > 0) {
-          ShapeIndex index = use.hlo_use.operand_index;
-          index.push_front(use.hlo_use.operand_number);
-          HloInstruction* parameter_instruction =
-              while_instruction->while_body()->parameter_instruction(0);
-          HloPosition parameter_position{parameter_instruction, index};
-          HloPosition while_output_position{while_instruction, index};
-          VLOG(1)
-              << "Add use aliases for while body root position to parameter "
-              << parameter_position.ToString() << " and while output "
-              << while_output_position.ToString() << " for use "
-              << use.hlo_use.ToString();
-          use.aliases.push_back(parameter_position);
-          use.aliases.push_back(while_output_position);
-        }
-      }
-    }
-  }
+void MsaAlgorithm::FindAliasesForTesting(
+    std::vector<AllocationValue>* allocation_values,
+    bool has_async_pipelined_while_loops) {
+  FindAliases(allocation_values, has_async_pipelined_while_loops);
 }
 
 void MsaAlgorithm::ExtendScopedAlternateMemoryAllocations() {
@@ -5466,7 +5537,7 @@ void MsaAlgorithm::CreateAllocationValuesFromColocatedIntervals(
     }
   }
 
-  FindAliases(&new_allocation_values);
+  FindAliases(&new_allocation_values, has_async_pipelined_while_loops_);
   absl::c_move(new_allocation_values, std::back_inserter(allocation_values));
 }
 
diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.h b/third_party/xla/xla/service/memory_space_assignment/algorithm.h
index b982fb98513e37..9a45cfe3a6e1a3 100644
--- a/third_party/xla/xla/service/memory_space_assignment/algorithm.h
+++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.h
@@ -380,6 +380,11 @@ class MsaAlgorithm : public GlobalDecreasingSizeBestFitHeap<HloValue> {
 
   absl::StatusOr<HeapSimulator::Result<HloValue>> Finish() override;
 
+  // Testing wrapper for FindAliases.
+  static void FindAliasesForTesting(
+      std::vector<AllocationValue>* allocation_values,
+      bool has_async_pipelined_while_loops = false);
+
   // Block prefetching is an MSA feature that allows processing all prefetches
   // in one pass within a block of memory space in the alternate memory. This
   // guarantees FIFO ordering of all prefetches and allows for more aggressive
@@ -479,10 +484,6 @@ class MsaAlgorithm : public GlobalDecreasingSizeBestFitHeap<HloValue> {
       absl::Span<const MsaBufferInterval* const> colocated_intervals,
       std::vector<AllocationValue>& allocation_values);
 
-  // Go through all the uses in the AllocationValues and find the aliasing
-  // positions.
-  void FindAliases(std::vector<AllocationValue>* allocation_values) const;
-
   AllocationSequence* allocations() { return allocations_; }
   const Options& options() const { return options_; }
   const HloAliasAnalysis& alias_analysis() { return alias_analysis_; }
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc
index baf254a34e3cdc..92230bc9ec8858 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc
@@ -18683,6 +18683,205 @@ ENTRY entry {
             kDefaultMemorySpace);
 }
 
+// Tests that FindAliases correctly maps operands to parameters for async-start,
+// async-update, call, and conditional instructions:
+// 1. Operands bound with async-start map 1-to-1 to initial parameters.
+// 2. Operand 0 of async-update (the async context token) does not alias
+//    parameters.
+// 3. Operands bound with async-update map 1-to-1 to subsequent parameters,
+//    offset by the number of previously bound operands.
+// 4. Operands of synchronous call map 1-to-1 to callee parameters by operand
+//    index, and operands 1..N of conditional map to branch_computation(i - 1)
+//    parameter 0.
+TEST_F(MemorySpaceAssignmentTest, FindAliasesAsyncStartAndAsyncUpdate) {
+  absl::string_view hlo_string = R"hlo(
+HloModule module, is_scheduled=true
+
+async_computation {
+  p0 = f32[4]{0} parameter(0)
+  p1 = f32[4]{0} parameter(1)
+  p2 = f32[4]{0} parameter(2)
+  p3 = f32[4]{0} parameter(3)
+  ROOT tuple = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) tuple(p0, p1, p2, p3)
+}
+
+call_computation {
+  cp0 = f32[4]{0} parameter(0)
+  cp1 = f32[4]{0} parameter(1)
+  ROOT call_tuple = (f32[4]{0}, f32[4]{0}) tuple(cp0, cp1)
+}
+
+branch0_computation {
+  ROOT bp0 = f32[4]{0} parameter(0)
+}
+
+branch1_computation {
+  ROOT bp1 = f32[4]{0} parameter(0)
+}
+
+ENTRY entry {
+  param0 = f32[4]{0} parameter(0)
+  param1 = f32[4]{0} parameter(1)
+  param2 = f32[4]{0} parameter(2)
+  param3 = f32[4]{0} parameter(3)
+  param_pred = pred[] parameter(4)
+  async-start = ((f32[4]{0}, f32[4]{0}), (), s32[]) async-start(param0, param1), calls=async_computation
+  async-update = ((f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}), (), s32[]) async-update(async-start, param2, param3), calls=async_computation
+  async-done = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) async-done(async-update), calls=async_computation
+  sync-call = (f32[4]{0}, f32[4]{0}) call(param0, param1), to_apply=call_computation
+  cond = f32[4]{0} conditional(param_pred, param2, param3), true_computation=branch0_computation, false_computation=branch1_computation
+  ROOT root = ((f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}), (f32[4]{0}, f32[4]{0}), f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) tuple(async-done, sync-call, cond, param0, param1, param2, param3)
+}
+  )hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  HloInstruction* param0 = FindInstruction(module.get(), "param0");
+  HloInstruction* param1 = FindInstruction(module.get(), "param1");
+  HloInstruction* param2 = FindInstruction(module.get(), "param2");
+  HloInstruction* param3 = FindInstruction(module.get(), "param3");
+  HloInstruction* pred = FindInstruction(module.get(), "param_pred");
+  HloInstruction* async_start = FindInstruction(module.get(), "async-start");
+  HloInstruction* async_update = FindInstruction(module.get(), "async-update");
+  HloInstruction* async_done = FindInstruction(module.get(), "async-done");
+  HloInstruction* sync_call = FindInstruction(module.get(), "sync-call");
+  HloInstruction* cond = FindInstruction(module.get(), "cond");
+
+  HloComputation* async_comp =
+      module->GetComputationWithName("async_computation");
+  HloInstruction* p0 = async_comp->parameter_instruction(0);
+  HloInstruction* p1 = async_comp->parameter_instruction(1);
+  HloInstruction* p2 = async_comp->parameter_instruction(2);
+  HloInstruction* p3 = async_comp->parameter_instruction(3);
+
+  HloComputation* call_comp =
+      module->GetComputationWithName("call_computation");
+  HloInstruction* cp0 = call_comp->parameter_instruction(0);
+  HloInstruction* cp1 = call_comp->parameter_instruction(1);
+
+  HloComputation* branch0_comp =
+      module->GetComputationWithName("branch0_computation");
+  HloInstruction* bp0 = branch0_comp->parameter_instruction(0);
+  HloComputation* branch1_comp =
+      module->GetComputationWithName("branch1_computation");
+  HloInstruction* bp1 = branch1_comp->parameter_instruction(0);
+
+  std::vector<AllocationValue> allocation_values;
+  // Allocation values for parameters of the async, call, and branch
+  // computations.
+  allocation_values.emplace_back(nullptr, HloPosition{p0, {}}, 16);
+  allocation_values.emplace_back(nullptr, HloPosition{p1, {}}, 16);
+  allocation_values.emplace_back(nullptr, HloPosition{p2, {}}, 16);
+  allocation_values.emplace_back(nullptr, HloPosition{p3, {}}, 16);
+  allocation_values.emplace_back(nullptr, HloPosition{cp0, {}}, 16);
+  allocation_values.emplace_back(nullptr, HloPosition{cp1, {}}, 16);
+  allocation_values.emplace_back(nullptr, HloPosition{bp0, {}}, 16);
+  allocation_values.emplace_back(nullptr, HloPosition{bp1, {}}, 16);
+
+  // Allocation values for caller operands passed to async-start (operands 0 and
+  // 1) and sync-call (operands 0 and 1).
+  AllocationValue val_param0(nullptr, HloPosition{param0, {}}, 16);
+  val_param0.uses().push_back(AllocationValue::Use{HloUse{async_start, 0}, 1});
+  val_param0.uses().push_back(AllocationValue::Use{HloUse{sync_call, 0}, 4});
+  allocation_values.push_back(std::move(val_param0));
+
+  AllocationValue val_param1(nullptr, HloPosition{param1, {}}, 16);
+  val_param1.uses().push_back(AllocationValue::Use{HloUse{async_start, 1}, 1});
+  val_param1.uses().push_back(AllocationValue::Use{HloUse{sync_call, 1}, 4});
+  allocation_values.push_back(std::move(val_param1));
+
+  // Allocation values for async-start bound operand elements ({0, 0} and
+  // {0, 1}) consumed as operand 0 of async-update.
+  AllocationValue val_async_start_0(nullptr, HloPosition{async_start, {0, 0}},
+                                    16);
+  val_async_start_0.uses().push_back(
+      AllocationValue::Use{HloUse{async_update, 0, {0, 0}}, 2});
+  allocation_values.push_back(std::move(val_async_start_0));
+
+  AllocationValue val_async_start_1(nullptr, HloPosition{async_start, {0, 1}},
+                                    16);
+  val_async_start_1.uses().push_back(
+      AllocationValue::Use{HloUse{async_update, 0, {0, 1}}, 2});
+  allocation_values.push_back(std::move(val_async_start_1));
+
+  // Allocation values for caller operands passed to async-update (operands 1
+  // and 2) and conditional (operands 1 and 2).
+  AllocationValue val_param2(nullptr, HloPosition{param2, {}}, 16);
+  val_param2.uses().push_back(AllocationValue::Use{HloUse{async_update, 1}, 2});
+  val_param2.uses().push_back(AllocationValue::Use{HloUse{cond, 1}, 5});
+  allocation_values.push_back(std::move(val_param2));
+
+  AllocationValue val_param3(nullptr, HloPosition{param3, {}}, 16);
+  val_param3.uses().push_back(AllocationValue::Use{HloUse{async_update, 2}, 2});
+  val_param3.uses().push_back(AllocationValue::Use{HloUse{cond, 2}, 5});
+  allocation_values.push_back(std::move(val_param3));
+
+  AllocationValue val_pred(nullptr, HloPosition{pred, {}}, 1);
+  val_pred.uses().push_back(AllocationValue::Use{HloUse{cond, 0}, 5});
+  allocation_values.push_back(std::move(val_pred));
+
+  // Allocation values for async-update bound operand elements ({0, 0}, {0, 1},
+  // {0, 2}, and {0, 3}) consumed by async-done.
+  for (int64_t i = 0; i < 4; ++i) {
+    AllocationValue val_async_update(nullptr, HloPosition{async_update, {0, i}},
+                                     16);
+    val_async_update.uses().push_back(
+        AllocationValue::Use{HloUse{async_done, 0, {0, i}}, 3});
+    allocation_values.push_back(std::move(val_async_update));
+  }
+
+  MsaAlgorithm::FindAliasesForTesting(&allocation_values);
+
+  auto find_val = [&](const HloInstruction* inst,
+                      ShapeIndex index = {}) -> const AllocationValue& {
+    for (const auto& v : allocation_values) {
+      if (v.position().instruction == inst && v.position().index == index) {
+        return v;
+      }
+    }
+    LOG(FATAL) << "Not found: " << inst->name() << " " << index.ToString();
+  };
+
+  // 1. Verify 1-to-1 mapping for async-start operands:
+  // param0 (operand 0) aliases {async_start, {0, 0}} and callee parameter p0.
+  EXPECT_THAT(find_val(param0).uses()[0].aliases,
+              ::testing::ElementsAre(HloPosition{async_start, {0, 0}},
+                                     HloPosition{p0, {}}));
+  // param1 (operand 1) aliases {async_start, {0, 1}} and callee parameter p1.
+  EXPECT_THAT(find_val(param1).uses()[0].aliases,
+              ::testing::ElementsAre(HloPosition{async_start, {0, 1}},
+                                     HloPosition{p1, {}}));
+
+  // 2. Verify operand 0 of async-update (the async context tuple) aliases
+  // matching elements in async_update, and does not alias any callee parameter:
+  EXPECT_THAT(find_val(async_start, {0, 0}).uses()[0].aliases,
+              ::testing::ElementsAre(HloPosition{async_update, {0, 0}}));
+  EXPECT_THAT(find_val(async_start, {0, 1}).uses()[0].aliases,
+              ::testing::ElementsAre(HloPosition{async_update, {0, 1}}));
+
+  // 3. Verify late-bound parameter offset calculation for async-update
+  // operands: param2 (operand 1 of async-update) aliases {async_update, {0, 2}}
+  // and callee parameter p2 (1 - 1 + 2 = 2).
+  EXPECT_THAT(find_val(param2).uses()[0].aliases,
+              ::testing::ElementsAre(HloPosition{async_update, {0, 2}},
+                                     HloPosition{p2, {}}));
+  // param3 (operand 2 of async-update) aliases {async_update, {0, 3}} and
+  // callee parameter p3 (2 - 1 + 2 = 3).
+  EXPECT_THAT(find_val(param3).uses()[0].aliases,
+              ::testing::ElementsAre(HloPosition{async_update, {0, 3}},
+                                     HloPosition{p3, {}}));
+
+  // 4. Verify 1-to-1 mapping for synchronous call and conditional operands:
+  EXPECT_THAT(find_val(param0).uses()[1].aliases,
+              ::testing::ElementsAre(HloPosition{cp0, {}}));
+  EXPECT_THAT(find_val(param1).uses()[1].aliases,
+              ::testing::ElementsAre(HloPosition{cp1, {}}));
+  EXPECT_THAT(find_val(pred).uses()[0].aliases, ::testing::IsEmpty());
+  EXPECT_THAT(find_val(param2).uses()[1].aliases,
+              ::testing::ElementsAre(HloPosition{bp0, {}}));
+  EXPECT_THAT(find_val(param3).uses()[1].aliases,
+              ::testing::ElementsAre(HloPosition{bp1, {}}));
+}
+
 // Tests a case where some operands bound in async-start/update are allowed
 // in alternate memory while others are restricted to default memory.
 // TODO(b/538345137): Re-enable this test once b/538345137 is fixed.