diff --git a/kernel/sched/walt/walt.h b/kernel/sched/walt/walt.h
index 32f570b603cc..5792c86de8f2 100644
--- a/kernel/sched/walt/walt.h
+++ b/kernel/sched/walt/walt.h
@@ -719,11 +719,13 @@ static inline bool task_fits_capacity(struct task_struct *p,
 					int cpu)
 {
 	unsigned int margin;
+	struct walt_rq *src_wrq = (struct walt_rq *) cpu_rq(task_cpu(p))->android_vendor_data1;
+	struct walt_rq *dst_wrq = (struct walt_rq *) cpu_rq(cpu)->android_vendor_data1;
 
 	/*
 	 * Derive upmigration/downmigrate margin wrt the src/dest CPU.
 	 */
-	if (capacity_orig_of(task_cpu(p)) > capacity_orig_of(cpu))
+	if (src_wrq->cluster->id > dst_wrq->cluster->id)
 		margin = sched_capacity_margin_down[cpu];
 	else
 		margin = sched_capacity_margin_up[task_cpu(p)];
diff --git a/kernel/sched/walt/walt_cfs.c b/kernel/sched/walt/walt_cfs.c
index 82f54d187d0a..9cdbe5c59ad5 100644
--- a/kernel/sched/walt/walt_cfs.c
+++ b/kernel/sched/walt/walt_cfs.c
@@ -64,8 +64,10 @@ bias_to_this_cpu(struct task_struct *p, int cpu, int start_cpu)
 {
 	bool base_test = cpumask_test_cpu(cpu, &p->cpus_mask) &&
 						cpu_active(cpu);
-	bool start_cap_test = (capacity_orig_of(cpu) >=
-					capacity_orig_of(start_cpu));
+
+	struct walt_rq *wrq = (struct walt_rq *) cpu_rq(cpu)->android_vendor_data1;
+	struct walt_rq *start_wrq = (struct walt_rq *) cpu_rq(start_cpu)->android_vendor_data1;
+	bool start_cap_test = (wrq->cluster->id >= start_wrq->cluster->id);
 
 	return base_test && start_cap_test;
 }
@@ -246,9 +248,12 @@ static void walt_find_best_target(struct sched_domain *sd,
 	bool rtg_high_prio_task = task_rtg_high_prio(p);
 	cpumask_t visit_cpus;
 	struct walt_task_struct *wts = (struct walt_task_struct *) p->android_vendor_data1;
+	struct walt_rq *prev_wrq = (struct walt_rq *) cpu_rq(prev_cpu)->android_vendor_data1;
+	struct walt_rq *start_wrq;
 
 	/* Find start CPU based on boost value */
 	start_cpu = fbt_env->start_cpu;
+	start_wrq = (struct walt_rq *) cpu_rq(start_cpu)->android_vendor_data1;
 
 	/*
 	 * For higher capacity worth I/O tasks, stop the search
@@ -265,7 +270,7 @@ static void walt_find_best_target(struct sched_domain *sd,
 	}
 
 	/* fast path for prev_cpu */
-	if (((capacity_orig_of(prev_cpu) == capacity_orig_of(start_cpu)) ||
+	if (((prev_wrq->cluster->id == start_wrq->cluster->id) ||
 				asym_cap_siblings(prev_cpu, start_cpu)) &&
 				cpu_active(prev_cpu) &&
 				available_idle_cpu(prev_cpu) &&
@@ -703,10 +708,12 @@ static inline bool select_cpu_same_energy(int cpu, int best_cpu, int prev_cpu)
 {
 	bool new_cpu_is_idle = available_idle_cpu(cpu);
 	bool best_cpu_is_idle = available_idle_cpu(best_cpu);
+	struct walt_rq *wrq = (struct walt_rq *) cpu_rq(cpu)->android_vendor_data1;
+	struct walt_rq *best_wrq = (struct walt_rq *) cpu_rq(best_cpu)->android_vendor_data1;
 
-	if (capacity_orig_of(best_cpu) < capacity_orig_of(cpu))
+	if (best_wrq->cluster->id < wrq->cluster->id)
 		return false;
-	if (capacity_orig_of(cpu) < capacity_orig_of(best_cpu))
+	if (wrq->cluster->id < best_wrq->cluster->id)
 		return true;
 
 	if (best_cpu_is_idle && walt_get_idle_exit_latency(cpu_rq(best_cpu)) <= 1)
@@ -759,6 +766,8 @@ int walt_find_energy_efficient_cpu(struct task_struct *p, int prev_cpu,
 	int first_cpu;
 	bool energy_eval_needed = true;
 	struct compute_energy_output output;
+	struct walt_rq *prev_wrq = (struct walt_rq *) cpu_rq(prev_cpu)->android_vendor_data1;
+	struct walt_rq *start_wrq;
 
 	if (walt_is_many_wakeup(sibling_count_hint) && prev_cpu != cpu &&
 			cpumask_test_cpu(prev_cpu, &p->cpus_mask))
@@ -770,6 +779,7 @@ int walt_find_energy_efficient_cpu(struct task_struct *p, int prev_cpu,
 	walt_get_indicies(p, &order_index, &end_index, task_boost, uclamp_boost,
 								&energy_eval_needed);
 	start_cpu = cpumask_first(&cpu_array[order_index][0]);
+	start_wrq = (struct walt_rq *) cpu_rq(start_cpu)->android_vendor_data1;
 
 	is_rtg = task_in_related_thread_group(p);
 	curr_is_rtg = task_in_related_thread_group(cpu_rq(cpu)->curr);
@@ -895,7 +905,7 @@ int walt_find_energy_efficient_cpu(struct task_struct *p, int prev_cpu,
 	    walt_get_idle_exit_latency(cpu_rq(best_energy_cpu)) <= 1) &&
 	    (prev_energy != ULONG_MAX) && (best_energy_cpu != prev_cpu) &&
 	    ((prev_energy - best_energy) <= prev_energy >> 5) &&
-	    (capacity_orig_of(prev_cpu) <= capacity_orig_of(start_cpu)))
+	    (prev_wrq->cluster->id <= start_wrq->cluster->id))
 		best_energy_cpu = prev_cpu;
 
 unlock:
diff --git a/kernel/sched/walt/walt_lb.c b/kernel/sched/walt/walt_lb.c
index 7e361c89e9f7..09cbd735e009 100644
--- a/kernel/sched/walt/walt_lb.c
+++ b/kernel/sched/walt/walt_lb.c
@@ -266,11 +266,13 @@ static inline bool need_active_lb(struct task_struct *p, int dst_cpu,
 				  int src_cpu)
 {
 	struct walt_task_struct *wts = (struct walt_task_struct *) p->android_vendor_data1;
+	struct walt_rq *src_wrq = (struct walt_rq *) cpu_rq(src_cpu)->android_vendor_data1;
+	struct walt_rq *dst_wrq = (struct walt_rq *) cpu_rq(dst_cpu)->android_vendor_data1;
 
 	if (cpu_rq(src_cpu)->active_balance)
 		return false;
 
-	if (capacity_orig_of(dst_cpu) <= capacity_orig_of(src_cpu))
+	if (dst_wrq->cluster->id <= src_wrq->cluster->id)
 		return false;
 
 	if (!wts->misfit)
@@ -286,12 +288,13 @@ static int walt_lb_pull_tasks(int dst_cpu, int src_cpu)
 	unsigned long flags;
 	struct task_struct *pulled_task = NULL, *p;
 	bool active_balance = false, to_lower;
-	struct walt_rq *wrq = (struct walt_rq *) src_rq->android_vendor_data1;
+	struct walt_rq *src_wrq = (struct walt_rq *) src_rq->android_vendor_data1;
+	struct walt_rq *dst_wrq = (struct walt_rq *) cpu_rq(dst_cpu)->android_vendor_data1;
 	struct walt_task_struct *wts;
 
 	BUG_ON(src_cpu == dst_cpu);
 
-	to_lower = capacity_orig_of(dst_cpu) < capacity_orig_of(src_cpu);
+	to_lower = dst_wrq->cluster->id < src_wrq->cluster->id;
 
 	raw_spin_lock_irqsave(&src_rq->__lock, flags);
 
@@ -339,7 +342,7 @@ static int walt_lb_pull_tasks(int dst_cpu, int src_cpu)
 				src_rq->active_balance = 1;
 				src_rq->push_cpu = dst_cpu;
 				get_task_struct(p);
-				wrq->push_task = p;
+				src_wrq->push_task = p;
 				mark_reserved(dst_cpu);
 
 				/* lock must be dropped before waking the stopper */
@@ -577,13 +580,15 @@ static int walt_lb_find_busiest_cpu(int dst_cpu, const cpumask_t *src_mask, int
 {
 	int fsrc_cpu = cpumask_first(src_mask);
 	int busiest_cpu;
+	struct walt_rq *fsrc_wrq = (struct walt_rq *) cpu_rq(fsrc_cpu)->android_vendor_data1;
+	struct walt_rq *dst_wrq = (struct walt_rq *) cpu_rq(dst_cpu)->android_vendor_data1;
 
-	if (capacity_orig_of(dst_cpu) == capacity_orig_of(fsrc_cpu))
+	if (dst_wrq->cluster->id == fsrc_wrq->cluster->id)
 		busiest_cpu = walt_lb_find_busiest_similar_cap_cpu(dst_cpu,
 								src_mask, has_misfit);
-	else if (capacity_orig_of(dst_cpu) > capacity_orig_of(fsrc_cpu))
+	else if (dst_wrq->cluster->id > fsrc_wrq->cluster->id)
 		busiest_cpu = walt_lb_find_busiest_from_lower_cap_cpu(dst_cpu,
-								      src_mask, has_misfit);
+								src_mask, has_misfit);
 	else
 		busiest_cpu = walt_lb_find_busiest_from_higher_cap_cpu(dst_cpu,
 								       src_mask, has_misfit);
@@ -597,7 +602,8 @@ void walt_lb_tick(struct rq *rq)
 	int prev_cpu = rq->cpu, new_cpu, ret;
 	struct task_struct *p = rq->curr;
 	unsigned long flags;
-	struct walt_rq *wrq = (struct walt_rq *) rq->android_vendor_data1;
+	struct walt_rq *prev_wrq = (struct walt_rq *) rq->android_vendor_data1;
+	struct walt_rq *new_wrq;
 	struct walt_task_struct *wts = (struct walt_task_struct *) p->android_vendor_data1;
 
 	raw_spin_lock(&rq->__lock);
@@ -627,9 +633,11 @@ void walt_lb_tick(struct rq *rq)
 	new_cpu = walt_find_energy_efficient_cpu(p, prev_cpu, 0, 1);
 	rcu_read_unlock();
 
+	new_wrq = (struct walt_rq *) cpu_rq(new_cpu)->android_vendor_data1;
+
 	/* prevent active task migration to busy or same/lower capacity CPU */
 	if (new_cpu < 0 || !available_idle_cpu(new_cpu) ||
-		capacity_orig_of(new_cpu) <= capacity_orig_of(prev_cpu))
+		new_wrq->cluster->id <= prev_wrq->cluster->id)
 		goto out_unlock;
 
 	raw_spin_lock(&rq->__lock);
@@ -640,7 +648,7 @@ void walt_lb_tick(struct rq *rq)
 	rq->active_balance = 1;
 	rq->push_cpu = new_cpu;
 	get_task_struct(p);
-	wrq->push_task = p;
+	prev_wrq->push_task = p;
 	raw_spin_unlock(&rq->__lock);
 
 	mark_reserved(new_cpu);
@@ -1011,10 +1019,12 @@ static void walt_can_migrate_task(void *unused, struct task_struct *p,
 				  int dst_cpu, int *can_migrate)
 {
 	bool to_lower;
+	struct walt_rq *dst_wrq = (struct walt_rq *) cpu_rq(dst_cpu)->android_vendor_data1;
+	struct walt_rq *task_wrq = (struct walt_rq *) cpu_rq(task_cpu(p))->android_vendor_data1;
 
 	if (unlikely(walt_disabled))
 		return;
-	to_lower = capacity_orig_of(dst_cpu) < capacity_orig_of(task_cpu(p));
+	to_lower = dst_wrq->cluster->id < task_wrq->cluster->id;
 
 	if (_walt_can_migrate_task(p, dst_cpu, to_lower, true))
 		return;