diff --git a/drivers/gpu/drm/amd/display/amdgpu_dm/amdgpu_dm_helpers.c b/drivers/gpu/drm/amd/display/amdgpu_dm/amdgpu_dm_helpers.c
index c4c35f6844f468d9e7e2f1332cfbd61ab82ecc8f..f91d759793f004807b14387411f5dc40539e1343 100644
--- a/drivers/gpu/drm/amd/display/amdgpu_dm/amdgpu_dm_helpers.c
+++ b/drivers/gpu/drm/amd/display/amdgpu_dm/amdgpu_dm_helpers.c
@@ -204,15 +204,16 @@ void dm_helpers_dp_update_branch_info(
 {}
 
 static void dm_helpers_construct_old_payload(
-			struct dc_link *link,
-			int pbn_per_slot,
+			struct drm_dp_mst_topology_mgr *mgr,
+			struct drm_dp_mst_topology_state *mst_state,
 			struct drm_dp_mst_atomic_payload *new_payload,
 			struct drm_dp_mst_atomic_payload *old_payload)
 {
-	struct link_mst_stream_allocation_table current_link_table =
-									link->mst_stream_alloc_table;
-	struct link_mst_stream_allocation *dc_alloc;
-	int i;
+	struct drm_dp_mst_atomic_payload *pos;
+	int pbn_per_slot = mst_state->pbn_div;
+	u8 next_payload_vc_start = mgr->next_start_slot;
+	u8 payload_vc_start = new_payload->vc_start_slot;
+	u8 allocated_time_slots;
 
 	*old_payload = *new_payload;
 
@@ -221,20 +222,17 @@ static void dm_helpers_construct_old_payload(
 	 * struct drm_dp_mst_atomic_payload are don't care fields
 	 * while calling drm_dp_remove_payload_part2()
 	 */
-	for (i = 0; i < current_link_table.stream_count; i++) {
-		dc_alloc =
-			&current_link_table.stream_allocations[i];
-
-		if (dc_alloc->vcp_id == new_payload->vcpi) {
-			old_payload->time_slots = dc_alloc->slot_count;
-			old_payload->pbn = dc_alloc->slot_count * pbn_per_slot;
-			break;
-		}
+	list_for_each_entry(pos, &mst_state->payloads, next) {
+		if (pos != new_payload &&
+		    pos->vc_start_slot > payload_vc_start &&
+		    pos->vc_start_slot < next_payload_vc_start)
+			next_payload_vc_start = pos->vc_start_slot;
 	}
 
-	/* make sure there is an old payload*/
-	ASSERT(i != current_link_table.stream_count);
+	allocated_time_slots = next_payload_vc_start - payload_vc_start;
 
+	old_payload->time_slots = allocated_time_slots;
+	old_payload->pbn = allocated_time_slots * pbn_per_slot;
 }
 
 /*
@@ -271,8 +269,8 @@ bool dm_helpers_dp_mst_write_payload_allocation_table(
 		drm_dp_add_payload_part1(mst_mgr, mst_state, new_payload);
 	} else {
 		/* construct old payload by VCPI*/
-		dm_helpers_construct_old_payload(stream->link, mst_state->pbn_div,
-						new_payload, &old_payload);
+		dm_helpers_construct_old_payload(mst_mgr, mst_state,
+						 new_payload, &old_payload);
 		target_payload = &old_payload;
 
 		drm_dp_remove_payload_part1(mst_mgr, mst_state, new_payload);
@@ -364,7 +362,7 @@ bool dm_helpers_dp_mst_send_payload_allocation(
 	if (enable) {
 		ret = drm_dp_add_payload_part2(mst_mgr, mst_state->base.state, new_payload);
 	} else {
-		dm_helpers_construct_old_payload(stream->link, mst_state->pbn_div,
+		dm_helpers_construct_old_payload(mst_mgr, mst_state,
 						 new_payload, &old_payload);
 		drm_dp_remove_payload_part2(mst_mgr, mst_state, &old_payload, new_payload);
 	}