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/mm/mem_sampling.c b/mm/mem_sampling.c
index 9ee68e15d1f67a3bbcf72cb5b603dc383fde5a40..fabf44ec20e364adb55f25bad96e72aee3368836 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);
@@ -341,6 +344,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 +417,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