diff --git a/kernel/sched/walt/boost.c b/kernel/sched/walt/boost.c
index 5ac7dce8160c..ded1ce3d0ac5 100644
--- a/kernel/sched/walt/boost.c
+++ b/kernel/sched/walt/boost.c
@@ -18,9 +18,6 @@
 static enum sched_boost_policy boost_policy_dt = SCHED_BOOST_NONE;
 static DEFINE_MUTEX(boost_mutex);
 
-struct task_group *task_group_topapp;
-struct task_group *task_group_foreground;
-
 void walt_init_tg(struct task_group *tg)
 {
 	struct walt_task_group *wtg;
diff --git a/kernel/sched/walt/sysctl.c b/kernel/sched/walt/sysctl.c
index 4816d702c790..e7738f397f09 100644
--- a/kernel/sched/walt/sysctl.c
+++ b/kernel/sched/walt/sysctl.c
@@ -39,7 +39,6 @@ unsigned int sysctl_sched_wake_up_idle[2];
 unsigned int sysctl_input_boost_ms;
 unsigned int sysctl_input_boost_freq[8];
 unsigned int sysctl_sched_boost_on_input;
-unsigned int sysctl_sched_init_stage;
 
 /* sysctl nodes accesed by other files */
 unsigned int __read_mostly sysctl_sched_coloc_downmigrate_ns;
@@ -61,53 +60,6 @@ unsigned int sysctl_walt_rtg_cfs_boost_prio = 99; /* disabled by default */
 /* range is [1 .. INT_MAX] */
 static int sysctl_task_read_pid = 1;
 
-static void init_tg_pointers(void)
-{
-	struct cgroup_subsys_state *css = &root_task_group.css;
-	struct cgroup_subsys_state *top_css = css;
-
-	/* ptrs are already initialized */
-	if (task_group_topapp)
-		return;
-
-	css_for_each_child(css, top_css) {
-		if (!strcmp(css->cgroup->kn->name, "top-app")) {
-			task_group_topapp = css_tg(css);
-			walt_init_topapp_tg(task_group_topapp);
-		} else if (!strcmp(css->cgroup->kn->name, "foreground")) {
-			task_group_foreground = css_tg(css);
-			walt_init_foreground_tg(task_group_foreground);
-		} else {
-			walt_init_tg(css_tg(css));
-		}
-	}
-}
-
-static int walt_init_stage_handler(struct ctl_table *table,
-				int write, void __user *buffer, size_t *lenp,
-				loff_t *ppos)
-{
-	int ret;
-	static DEFINE_MUTEX(mutex);
-	int old_value = sysctl_sched_init_stage;
-
-	mutex_lock(&mutex);
-
-	ret = proc_dointvec(table, write, buffer, lenp, ppos);
-
-	if (ret || !write)
-		goto unlock;
-
-	if (sysctl_sched_init_stage == 1 &&
-			old_value != sysctl_sched_init_stage) {
-		init_tg_pointers();
-	}
-
-unlock:
-	mutex_unlock(&mutex);
-	return ret;
-}
-
 static int walt_proc_group_thresholds_handler(struct ctl_table *table, int write,
 				       void __user *buffer, size_t *lenp,
 				       loff_t *ppos)
@@ -493,13 +445,6 @@ struct ctl_table input_boost_sysctls[] = {
 };
 
 struct ctl_table walt_table[] = {
-	{
-		.procname	= "sched_init_stage",
-		.data		= &sysctl_sched_init_stage,
-		.maxlen		= sizeof(unsigned int),
-		.mode		= 0644,
-		.proc_handler	= walt_init_stage_handler,
-	},
 	{
 		.procname	= "sched_user_hint",
 		.data		= &sysctl_sched_user_hint,
diff --git a/kernel/sched/walt/walt.c b/kernel/sched/walt/walt.c
index ab3f4a6d0b62..602eaead531c 100644
--- a/kernel/sched/walt/walt.c
+++ b/kernel/sched/walt/walt.c
@@ -2984,6 +2984,24 @@ static int sync_cgroup_colocation(struct task_struct *p, bool insert)
 	return __sched_set_group_id(p, grp_id);
 }
 
+static void walt_update_tg_pointer(struct cgroup_subsys_state *css)
+{
+	if (!strcmp(css->cgroup->kn->name, "top-app"))
+		walt_init_topapp_tg(css_tg(css));
+	else if (!strcmp(css->cgroup->kn->name, "foreground"))
+		walt_init_foreground_tg(css_tg(css));
+	else
+		walt_init_tg(css_tg(css));
+}
+
+static void android_rvh_cpu_cgroup_online(void *unused, struct cgroup_subsys_state *css)
+{
+	if (unlikely(walt_disabled))
+		return;
+
+	walt_update_tg_pointer(css);
+}
+
 static void android_rvh_cpu_cgroup_attach(void *unused,
 						struct cgroup_taskset *tset)
 {
@@ -2993,6 +3011,9 @@ static void android_rvh_cpu_cgroup_attach(void *unused,
 	struct task_group *tg;
 	struct walt_task_group *wtg;
 
+	if (unlikely(walt_disabled))
+		return;
+
 	cgroup_taskset_first(tset, &css);
 	if (!css)
 		return;
@@ -3981,6 +4002,7 @@ static void register_walt_hooks(void)
 	register_trace_android_rvh_schedule(android_rvh_schedule, NULL);
 	register_trace_android_rvh_resume_cpus(android_rvh_resume_cpus, NULL);
 	register_trace_android_rvh_cpu_cgroup_attach(android_rvh_cpu_cgroup_attach, NULL);
+	register_trace_android_rvh_cpu_cgroup_online(android_rvh_cpu_cgroup_online, NULL);
 	register_trace_android_rvh_update_cpus_allowed(android_rvh_update_cpus_allowed, NULL);
 	register_trace_android_rvh_sched_fork_init(android_rvh_sched_fork_init, NULL);
 	register_trace_android_rvh_ttwu_cond(android_rvh_ttwu_cond, NULL);
@@ -4039,6 +4061,17 @@ static int walt_init_stop_handler(void *data)
 	return 0;
 }
 
+static void walt_init_tg_pointers(void)
+{
+	struct cgroup_subsys_state *css = &root_task_group.css;
+	struct cgroup_subsys_state *top_css = css;
+
+	rcu_read_lock();
+	css_for_each_child(css, top_css)
+		walt_update_tg_pointer(css);
+	rcu_read_unlock();
+}
+
 static void walt_init(void)
 {
 	struct ctl_table_header *hdr;
@@ -4054,6 +4087,7 @@ static void walt_init(void)
 	BUG_ON(alloc_related_thread_groups());
 	walt_init_cycle_counter();
 	init_clusters();
+	walt_init_tg_pointers();
 
 	register_walt_hooks();
 	walt_fixup_init();
diff --git a/kernel/sched/walt/walt.h b/kernel/sched/walt/walt.h
index 1d87cdf8ef19..e94e50b4f507 100644
--- a/kernel/sched/walt/walt.h
+++ b/kernel/sched/walt/walt.h
@@ -282,8 +282,6 @@ extern void walt_update_group_thresholds(void);
 extern void sched_window_nr_ticks_change(void);
 extern unsigned long sched_user_hint_reset_time;
 extern struct irq_work walt_migration_irq_work;
-extern struct task_group *task_group_topapp;
-extern struct task_group *task_group_foreground;
 
 #define LIB_PATH_LENGTH 512
 extern unsigned int cpuinfo_max_freq_cached;