diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp
index 72bb644b66..b9751f68b3 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp
@@ -574,6 +574,21 @@ hsa_status_t KfdDriver::SetTrapHandler(uint32_t node_id, const void* base, uint6
   return HSA_STATUS_SUCCESS;
 }
 
+hsa_status_t KfdDriver::AllocateScratchMemory(uint32_t node_id, uint64_t size, void** mem) const {
+  assert(mem);
+  assert(size > 0);
+
+  HsaMemFlags flags = {};
+  flags.ui32.Scratch = 1;
+  flags.ui32.HostAccess = 1;
+
+  void* ptr = AllocateKfdMemory(flags, node_id, size);
+  if (ptr == nullptr) return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+
+  *mem = ptr;
+  return HSA_STATUS_SUCCESS;
+}
+
 hsa_status_t KfdDriver::GetDeviceHandle(uint32_t node_id, void** device_handle) const {
   assert(device_handle);
 
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
index fd7fa21a04..76c4842d5f 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
@@ -877,6 +877,10 @@ hsa_status_t XdnaDriver::SetTrapHandler(uint32_t node_id, const void* base, uint
   return HSA_STATUS_ERROR;
 }
 
+hsa_status_t XdnaDriver::AllocateScratchMemory(uint32_t node_id, uint64_t size, void** mem) const {
+  return HSA_STATUS_ERROR;
+}
+
 hsa_status_t XdnaDriver::GetDeviceHandle(uint32_t node_id, void** device_handle) const {
   return HSA_STATUS_ERROR;
 }
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_kfd_driver.h b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_kfd_driver.h
index 4c56b8b9fe..fc0ee82cba 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_kfd_driver.h
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_kfd_driver.h
@@ -128,6 +128,7 @@ public:
   hsa_status_t GetClockCounters(uint32_t node_id, HsaClockCounters* clock_counter) const override;
   hsa_status_t GetTileConfig(uint32_t node_id, HsaGpuTileConfig* config) const override;
   hsa_status_t GetWallclockFrequency(uint32_t node_id, uint64_t* frequency) const override;
+  hsa_status_t AllocateScratchMemory(uint32_t node_id, uint64_t size, void** mem) const override;
 
   hsa_status_t OpenSMI(uint32_t node_id, int* fd) const override;
 
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_xdna_driver.h b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_xdna_driver.h
index 8642c8d462..0aee5bdfcb 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_xdna_driver.h
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_xdna_driver.h
@@ -242,6 +242,7 @@ public:
   hsa_status_t GetClockCounters(uint32_t node_id, HsaClockCounters* clock_counter) const override;
   hsa_status_t GetTileConfig(uint32_t node_id, HsaGpuTileConfig* config) const override;
   hsa_status_t GetWallclockFrequency(uint32_t node_id, uint64_t* frequency) const override;
+  hsa_status_t AllocateScratchMemory(uint32_t node_id, uint64_t size, void** mem) const override;
 
   hsa_status_t IsModelEnabled(bool* enable) const override;
 
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/driver.h b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/driver.h
index a74670b31a..8df7e73686 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/driver.h
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/driver.h
@@ -301,6 +301,13 @@ public:
   /// code.
   virtual hsa_status_t GetWallclockFrequency(uint32_t node_id, uint64_t* frequency) const = 0;
 
+  /// @brief Allocates scratch memory for the agent.
+  /// @param[in] node_id Node ID of the agent
+  /// @param[in] size Size of the scratch memory
+  /// @param[out] mem Pointer to the scratch memory
+  /// @return HSA_STATUS_SUCCESS if scratch memory allocated successfully.
+  virtual hsa_status_t AllocateScratchMemory(uint32_t node_id, uint64_t size, void** mem) const = 0;
+
   /// Unique identifier for supported kernel-mode drivers.
   const DriverType kernel_driver_type_;
 
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
index 446ec50ae6..9d687df62f 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
@@ -491,11 +491,6 @@ void GpuAgent::InitRegionList() {
 }
 
 void GpuAgent::InitScratchPool() {
-  HsaMemFlags flags;
-  flags.Value = 0;
-  flags.ui32.Scratch = 1;
-  flags.ui32.HostAccess = 1;
-
   scratch_per_thread_ =
       core::Runtime::runtime_singleton_->flag().scratch_mem_size();
   if (scratch_per_thread_ == 0)
@@ -516,14 +511,13 @@ void GpuAgent::InitScratchPool() {
 #endif
 
   void* scratch_base = nullptr;
-  HSAKMT_STATUS err =
-      HSAKMT_CALL(hsaKmtAllocMemory(node_id(), max_scratch_len, flags, &scratch_base));
-  assert(err == HSAKMT_STATUS_SUCCESS && "hsaKmtAllocMemory(Scratch) failed");
+  hsa_status_t err = driver().AllocateScratchMemory(node_id(), max_scratch_len, &scratch_base);
+  assert(err == HSA_STATUS_SUCCESS && "AllocateScratchMemory failed");
   assert(IsMultipleOf(scratch_base, 0x1000) &&
          "Scratch base is not page aligned!");
 
   scratch_pool_. ~SmallHeap();
-  if (HSAKMT_STATUS_SUCCESS == err) {
+  if (HSA_STATUS_SUCCESS == err) {
     new (&scratch_pool_) SmallHeap(scratch_base, max_scratch_len);
   } else {
     new (&scratch_pool_) SmallHeap();