diff --git a/drivers/arm/mm_monitor/mm_spe.c b/drivers/arm/mm_monitor/mm_spe.c
index 192f1f9c6adb8093366dac553ef8e9fb00e1fe43..73fa4acdc998bcb97c1f39e9f696b0e732c6886b 100644
--- a/drivers/arm/mm_monitor/mm_spe.c
+++ b/drivers/arm/mm_monitor/mm_spe.c
@@ -114,8 +114,8 @@ void mm_spe_buffer_free(void)
 			continue;
 		mm_spe_percpu_buffer_free(cpu);
 	}
-	spe_probe_status -= 1;
 	set_mem_sampling_state(false);
+	spe_probe_status -= 1;
 }
 EXPORT_SYMBOL_GPL(mm_spe_buffer_free);
 
@@ -366,11 +366,11 @@ void mm_spe_decoding(void)
 	arm_spe_decode_buf(spe_buf->cur, spe_buf->size);
 }
 
-struct mm_spe_buf *mm_spe_getbuf_addr(void)
+void *mm_spe_getbuf_addr(void)
 {
 	struct mm_spe_buf *spe_buf = this_cpu_ptr(&per_cpu_spe_buf);
 
-	return spe_buf;
+	return spe_buf->record_base;
 }
 
 int mm_spe_getnum_record(void)
diff --git a/include/linux/mem_sampling.h b/include/linux/mem_sampling.h
index 71d098fd3c05bf2b84bfbef98adbba4a2e64e704..b325c9955d31be888dafb103272bd6b530e10424 100644
--- a/include/linux/mem_sampling.h
+++ b/include/linux/mem_sampling.h
@@ -49,6 +49,7 @@ struct mem_sampling_record {
 	u64				context_id;
 	u64				boost_spe_addr[8];
 	u64				rem_addr;
+	u16				boost_spe_idx;
 	u16				source;
 };
 
@@ -57,7 +58,7 @@ struct mem_sampling_ops_struct {
 	void (*sampling_stop)(void);
 	void (*sampling_continue)(void);
 	void (*sampling_decoding)(void);
-	struct mm_spe_buf* (*mm_spe_getbuf_addr)(void);
+	void* (*mm_spe_getbuf_addr)(void);
 	int (*mm_spe_getnum_record)(void);
 
 };
@@ -83,7 +84,7 @@ void mm_spe_stop(void);
 void mm_spe_continue(void);
 void mm_spe_decoding(void);
 int mm_spe_getnum_record(void);
-struct mm_spe_buf *mm_spe_getbuf_addr(void);
+void *mm_spe_getbuf_addr(void);
 int mm_spe_enabled(void);
 void arm_spe_set_probe_status(int status);
 #else
@@ -93,7 +94,7 @@ static inline void mm_spe_decoding(void) { }
 static inline void arm_spe_set_probe_status(int status) { }
 static inline int mm_spe_start(void) { return 0; }
 static inline int mm_spe_getnum_record(void) { return 0; }
-static inline struct mm_spe_buf *mm_spe_getbuf_addr(void) { return NULL; }
+static inline void *mm_spe_getbuf_addr(void) { return NULL; }
 static inline int mm_spe_enabled(void) { return 0; }
 #endif /* CONFIG_ARM_SPE_MEM_SAMPLING */
 
diff --git a/include/trace/events/kmem.h b/include/trace/events/kmem.h
index 28b9679c474d05368d94f79987a28393dc58158a..3e78e6bd6e18e9ca5d8562fda43a6a2a7089d2c6 100644
--- a/include/trace/events/kmem.h
+++ b/include/trace/events/kmem.h
@@ -523,6 +523,27 @@ TRACE_EVENT(mm_mem_sampling_access_record,
 		__entry->cpuid, __entry->pid)
 );
 #endif /* CONFIG_NUMABALANCING_MEM_SAMPLING */
+
+#ifdef CONFIG_DAMON_MEM_SAMPLING
+TRACE_EVENT(mm_mem_sampling_damon_record,
+
+	TP_PROTO(u64 vaddr, int pid),
+
+	TP_ARGS(vaddr, pid),
+
+	TP_STRUCT__entry(
+		__field(u64, vaddr)
+		__field(int, pid)
+	),
+
+	TP_fast_assign(
+		__entry->vaddr = vaddr;
+		__entry->pid = pid;
+	),
+
+	TP_printk("vaddr=%llx pid=%d", __entry->vaddr, __entry->pid)
+);
+#endif /* CONFIG_DAMON_MEM_SAMPLING */
 #endif /* _TRACE_KMEM_H */
 
 /* This part must be outside protection */
diff --git a/mm/mem_sampling.c b/mm/mem_sampling.c
index 9ee68e15d1f67a3bbcf72cb5b603dc383fde5a40..74f95e4611fe1d8965573128b3f52f3928c94d5b 100644
--- a/mm/mem_sampling.c
+++ b/mm/mem_sampling.c
@@ -278,6 +278,9 @@ static void numa_balancing_mem_sampling_cb_unregister(void)
 }
 static void set_numabalancing_mem_sampling_state(bool enabled)
 {
+	if (!mem_sampling_ops.sampling_start || !mm_spe_enabled())
+		return;
+
 	if (enabled) {
 		numa_balancing_mem_sampling_cb_register();
 		static_branch_enable(&sched_numabalancing_mem_sampling);
@@ -316,6 +319,7 @@ static void damon_mem_sampling_record_cb(struct mem_sampling_record *record)
 	mmput(mm);
 
 	domon_record.vaddr = record->virt_addr;
+	trace_mm_mem_sampling_damon_record(record->virt_addr, (pid_t)record->context_id);
 
 	/* only the proc under monitor now has damon_fifo */
 	if (damon_fifo) {
@@ -341,6 +345,9 @@ static void damon_mem_sampling_record_cb_unregister(void)
 
 static void set_damon_mem_sampling_state(bool enabled)
 {
+	if (!mem_sampling_ops.sampling_start || !mm_spe_enabled())
+		return;
+
 	if (enabled) {
 		damon_mem_sampling_record_cb_register();
 		static_branch_enable(&mm_damon_mem_sampling);
@@ -411,14 +418,15 @@ static void __set_mem_sampling_state(bool enabled)
 
 void set_mem_sampling_state(bool enabled)
 {
+	if (!mem_sampling_ops.sampling_start || !mm_spe_enabled())
+		return;
+
 	if (mem_sampling_saved_state != MEM_SAMPLING_STATE_EMPTY) {
 		mem_sampling_saved_state = enabled ? MEM_SAMPLING_STATE_ENABLE :
 					    MEM_SAMPLING_STATE_DISABLE;
 		return;
 	}
 
-	if (!mem_sampling_ops.sampling_start || !mm_spe_enabled())
-		return;
 	if (enabled)
 		sysctl_mem_sampling_mode = MEM_SAMPLING_NORMAL;
 	else