diff --git a/kernel/sched/walt/trace.h b/kernel/sched/walt/trace.h
index 0a53034c6917..d382b5fd018f 100644
--- a/kernel/sched/walt/trace.h
+++ b/kernel/sched/walt/trace.h
@@ -1237,6 +1237,31 @@ TRACE_EVENT(sched_overutilized,
 	 __entry->overutilized, __entry->span)
 );
 
+TRACE_EVENT(sched_cgroup_attach,
+
+	TP_PROTO(struct task_struct *p, unsigned int grp_id, int ret),
+
+	TP_ARGS(p, grp_id, ret),
+
+	TP_STRUCT__entry(
+		__array(char,		comm, TASK_COMM_LEN)
+		__field(pid_t,		pid)
+		__field(unsigned int,	grp_id)
+		__field(int,		ret)
+	),
+
+	TP_fast_assign(
+		memcpy(__entry->comm, p->comm, TASK_COMM_LEN);
+		__entry->pid	= p->pid;
+		__entry->grp_id = grp_id;
+		__entry->ret = ret;
+	),
+
+	TP_printk("comm=%s pid=%d grp_id=%u ret=%d",
+			__entry->comm, __entry->pid,
+			__entry->grp_id, __entry->ret)
+
+);
 #endif /* _TRACE_WALT_H */
 
 #undef TRACE_INCLUDE_PATH
diff --git a/kernel/sched/walt/walt.c b/kernel/sched/walt/walt.c
index 4f28f497e776..dd925a739681 100644
--- a/kernel/sched/walt/walt.c
+++ b/kernel/sched/walt/walt.c
@@ -3059,13 +3059,6 @@ static int create_default_coloc_group(void)
 	return 0;
 }
 
-static int sync_cgroup_colocation(struct task_struct *p, bool insert)
-{
-	unsigned int grp_id = insert ? DEFAULT_CGROUP_COLOC_ID : 0;
-
-	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"))
@@ -3089,9 +3082,10 @@ static void android_rvh_cpu_cgroup_attach(void *unused,
 {
 	struct task_struct *task;
 	struct cgroup_subsys_state *css;
-	bool colocate;
 	struct task_group *tg;
 	struct walt_task_group *wtg;
+	unsigned int grp_id;
+	int ret;
 
 	if (unlikely(walt_disabled))
 		return;
@@ -3102,10 +3096,12 @@ static void android_rvh_cpu_cgroup_attach(void *unused,
 
 	tg = container_of(css, struct task_group, css);
 	wtg = (struct walt_task_group *) tg->android_vendor_data1;
-	colocate = wtg->colocate;
 
-	cgroup_taskset_for_each(task, css, tset)
-		sync_cgroup_colocation(task, colocate);
+	cgroup_taskset_for_each(task, css, tset) {
+		grp_id = wtg->colocate ? DEFAULT_CGROUP_COLOC_ID : 0;
+		ret = __sched_set_group_id(task, grp_id);
+		trace_sched_cgroup_attach(task, grp_id, ret);
+	}
 }
 
 static bool is_cluster_hosting_top_app(struct walt_sched_cluster *cluster)