diff --git a/runtime/command_stream/command_stream_receiver.cpp b/runtime/command_stream/command_stream_receiver.cpp
index 79ac28e8b6..3178b552ea 100644
--- a/runtime/command_stream/command_stream_receiver.cpp
+++ b/runtime/command_stream/command_stream_receiver.cpp
@@ -179,6 +179,11 @@ void CommandStreamReceiver::cleanupResources() {
         tagAllocation = nullptr;
         tagAddress = nullptr;
     }
+
+    if (preemptionAllocation) {
+        getMemoryManager()->freeGraphicsMemory(preemptionAllocation);
+        preemptionAllocation = nullptr;
+    }
 }
 
 bool CommandStreamReceiver::waitForCompletionWithTimeout(bool enableTimeout, int64_t timeoutMicroseconds, uint32_t taskCountToWait) {
@@ -364,6 +369,15 @@ bool CommandStreamReceiver::initializeTagAllocation() {
     return true;
 }
 
+bool CommandStreamReceiver::createPreemptionAllocation() {
+    auto hwInfo = executionEnvironment.getHardwareInfo();
+    AllocationProperties properties{true, hwInfo->capabilityTable.requiredPreemptionSurfaceSize, GraphicsAllocation::AllocationType::PREEMPTION, false};
+    properties.flags.uncacheable = hwInfo->workaroundTable.waCSRUncachable;
+    properties.alignment = 256 * MemoryConstants::kiloByte;
+    this->preemptionAllocation = getMemoryManager()->allocateGraphicsMemoryWithProperties(properties);
+    return this->preemptionAllocation != nullptr;
+}
+
 std::unique_lock<CommandStreamReceiver::MutexType> CommandStreamReceiver::obtainUniqueOwnership() {
     return std::unique_lock<CommandStreamReceiver::MutexType>(this->ownershipMutex);
 }
diff --git a/runtime/command_stream/command_stream_receiver.h b/runtime/command_stream/command_stream_receiver.h
index 5a4347137b..a0dd90bac9 100644
--- a/runtime/command_stream/command_stream_receiver.h
+++ b/runtime/command_stream/command_stream_receiver.h
@@ -120,9 +120,7 @@ class CommandStreamReceiver {
     GraphicsAllocation *getScratchAllocation();
     GraphicsAllocation *getDebugSurfaceAllocation() const { return debugSurface; }
     GraphicsAllocation *allocateDebugSurface(size_t size);
-
-    void setPreemptionCsrAllocation(GraphicsAllocation *allocation) { preemptionCsrAllocation = allocation; }
-
+    GraphicsAllocation *getPreemptionAllocation() const { return preemptionAllocation; }
     void requestThreadArbitrationPolicy(uint32_t requiredPolicy) { this->requiredThreadArbitrationPolicy = requiredPolicy; }
     void requestStallingPipeControlOnNextFlush() { stallingPipeControlOnNextFlushRequired = true; }
 
@@ -147,6 +145,7 @@ class CommandStreamReceiver {
     void setExperimentalCmdBuffer(std::unique_ptr<ExperimentalCommandBuffer> &&cmdBuffer);
 
     bool initializeTagAllocation();
+    MOCKABLE_VIRTUAL bool createPreemptionAllocation();
     MOCKABLE_VIRTUAL std::unique_lock<MutexType> obtainUniqueOwnership();
 
     bool peekTimestampPacketWriteEnabled() const { return timestampPacketWriteEnabled; }
@@ -207,7 +206,7 @@ class CommandStreamReceiver {
     volatile uint32_t *tagAddress = nullptr;
 
     GraphicsAllocation *tagAllocation = nullptr;
-    GraphicsAllocation *preemptionCsrAllocation = nullptr;
+    GraphicsAllocation *preemptionAllocation = nullptr;
     GraphicsAllocation *debugSurface = nullptr;
     OSInterface *osInterface = nullptr;
 
diff --git a/runtime/command_stream/command_stream_receiver_hw_base.inl b/runtime/command_stream/command_stream_receiver_hw_base.inl
index c42569a633..5ee6c6a747 100644
--- a/runtime/command_stream/command_stream_receiver_hw_base.inl
+++ b/runtime/command_stream/command_stream_receiver_hw_base.inl
@@ -366,8 +366,8 @@ CompletionStamp CommandStreamReceiverHw<GfxFamily>::flushTask(
 
     this->makeResident(*tagAllocation);
 
-    if (preemptionCsrAllocation)
-        makeResident(*preemptionCsrAllocation);
+    if (preemptionAllocation)
+        makeResident(*preemptionAllocation);
 
     if (dispatchFlags.preemptionMode == PreemptionMode::MidThread || device.isSourceLevelDebuggerActive()) {
         auto sipType = SipKernel::getSipKernelType(device.getHardwareInfo().platform.eRenderCoreFamily, device.isSourceLevelDebuggerActive());
@@ -636,7 +636,7 @@ inline void CommandStreamReceiverHw<GfxFamily>::waitForTaskCountWithKmdNotifyFal
 
 template <typename GfxFamily>
 inline void CommandStreamReceiverHw<GfxFamily>::programPreemption(LinearStream &csr, Device &device, DispatchFlags &dispatchFlags) {
-    PreemptionHelper::programCmdStream<GfxFamily>(csr, dispatchFlags.preemptionMode, this->lastPreemptionMode, preemptionCsrAllocation, device);
+    PreemptionHelper::programCmdStream<GfxFamily>(csr, dispatchFlags.preemptionMode, this->lastPreemptionMode, preemptionAllocation, device);
     this->lastPreemptionMode = dispatchFlags.preemptionMode;
 }
 
@@ -656,7 +656,7 @@ inline void CommandStreamReceiverHw<GfxFamily>::programStateSip(LinearStream &cm
 template <typename GfxFamily>
 inline void CommandStreamReceiverHw<GfxFamily>::programPreamble(LinearStream &csr, Device &device, DispatchFlags &dispatchFlags, uint32_t &newL3Config) {
     if (!this->isPreambleSent) {
-        PreambleHelper<GfxFamily>::programPreamble(&csr, device, newL3Config, this->requiredThreadArbitrationPolicy, this->preemptionCsrAllocation);
+        PreambleHelper<GfxFamily>::programPreamble(&csr, device, newL3Config, this->requiredThreadArbitrationPolicy, this->preemptionAllocation);
         this->isPreambleSent = true;
         this->lastSentL3Config = newL3Config;
         this->lastSentThreadArbitrationPolicy = this->requiredThreadArbitrationPolicy;
diff --git a/runtime/device/device.cpp b/runtime/device/device.cpp
index c6a04f4c98..9ef9df3dbe 100644
--- a/runtime/device/device.cpp
+++ b/runtime/device/device.cpp
@@ -96,13 +96,9 @@ Device::~Device() {
         executionEnvironment->sourceLevelDebugger->notifyDeviceDestruction();
     }
 
-    if (preemptionAllocation) {
-        executionEnvironment->memoryManager->freeGraphicsMemory(preemptionAllocation);
-        preemptionAllocation = nullptr;
-    }
     executionEnvironment->memoryManager->waitForDeletions();
-
     alignedFree(this->slmWindowStartAddress);
+
     executionEnvironment->decRefInternal();
 }
 
@@ -112,7 +108,6 @@ bool Device::createDeviceImpl() {
     if (!createEngines()) {
         return false;
     }
-
     executionEnvironment->memoryManager->setDefaultEngineIndex(defaultEngineIndex);
 
     auto osInterface = executionEnvironment->osInterface.get();
@@ -143,17 +138,9 @@ bool Device::createDeviceImpl() {
 
     executionEnvironment->memoryManager->setForce32BitAllocations(getDeviceInfo().force32BitAddressess);
 
-    if (preemptionMode == PreemptionMode::MidThread || isSourceLevelDebuggerActive()) {
-        preemptionAllocation = executionEnvironment->memoryManager->allocateGraphicsMemoryWithProperties(getAllocationPropertiesForPreemption());
-        if (!preemptionAllocation) {
-            return false;
-        }
-    }
-
-    for (auto &engine : engines) {
-        auto csr = engine.commandStreamReceiver;
-        csr->setPreemptionCsrAllocation(preemptionAllocation);
-        if (DebugManager.flags.EnableExperimentalCommandBuffer.get() > 0) {
+    if (DebugManager.flags.EnableExperimentalCommandBuffer.get() > 0) {
+        for (auto &engine : engines) {
+            auto csr = engine.commandStreamReceiver;
             csr->setExperimentalCmdBuffer(std::make_unique<ExperimentalCommandBuffer>(csr, getDeviceInfo().profilingTimerResolution));
         }
     }
@@ -161,13 +148,8 @@ bool Device::createDeviceImpl() {
     return true;
 }
 
-AllocationProperties Device::getAllocationPropertiesForPreemption() const {
-    AllocationProperties properties{true, getHardwareInfo().capabilityTable.requiredPreemptionSurfaceSize, GraphicsAllocation::AllocationType::PREEMPTION, false};
-    properties.flags.uncacheable = getWaTable()->waCSRUncachable;
-    properties.alignment = 256 * MemoryConstants::kiloByte;
-    return properties;
-}
 bool Device::createEngines() {
+
     auto &hwInfo = getHardwareInfo();
     auto defaultEngineType = getChosenEngineType(hwInfo);
     auto &gpgpuEngines = HwHelper::get(hwInfo.platform.eRenderCoreFamily).getGpgpuEngineInstances();
@@ -189,11 +171,14 @@ bool Device::createEngines() {
         if (!commandStreamReceiver->initializeTagAllocation()) {
             return false;
         }
-
         if (gpgpuEngines[deviceCsrIndex] == defaultEngineType && !lowPriority) {
             defaultEngineIndex = deviceCsrIndex;
         }
 
+        if ((preemptionMode == PreemptionMode::MidThread || isSourceLevelDebuggerActive()) && !commandStreamReceiver->createPreemptionAllocation()) {
+            return false;
+        }
+
         engines.push_back({commandStreamReceiver, osContext});
     }
     return true;
diff --git a/runtime/device/device.h b/runtime/device/device.h
index 7a77fc8062..b776d7c68b 100644
--- a/runtime/device/device.h
+++ b/runtime/device/device.h
@@ -104,12 +104,11 @@ class Device : public BaseObject<_cl_device_id> {
     PerformanceCounters *getPerformanceCounters() { return performanceCounters.get(); }
     static decltype(&PerformanceCounters::create) createPerformanceCountersFunc;
     PreemptionMode getPreemptionMode() const { return preemptionMode; }
-    GraphicsAllocation *getPreemptionAllocation() const { return preemptionAllocation; }
     MOCKABLE_VIRTUAL const WhitelistedRegisters &getWhitelistedRegisters() { return getHardwareInfo().capabilityTable.whitelistedRegisters; }
     std::vector<unsigned int> simultaneousInterops;
     std::string deviceExtensions;
     std::string name;
-    bool isSourceLevelDebuggerActive() const;
+    MOCKABLE_VIRTUAL bool isSourceLevelDebuggerActive() const;
     SourceLevelDebugger *getSourceLevelDebugger() { return executionEnvironment->sourceLevelDebugger.get(); }
     ExecutionEnvironment *getExecutionEnvironment() const { return executionEnvironment; }
     const HardwareCapabilities &getHardwareCapabilities() const { return hardwareCapabilities; }
@@ -137,14 +136,10 @@ class Device : public BaseObject<_cl_device_id> {
     MOCKABLE_VIRTUAL void initializeCaps();
     void setupFp64Flags();
     void appendOSExtensions(std::string &deviceExtensions);
-    AllocationProperties getAllocationPropertiesForPreemption() const;
-
     unsigned int enabledClVersion = 0u;
 
     HardwareCapabilities hardwareCapabilities = {};
     DeviceInfo deviceInfo;
-
-    GraphicsAllocation *preemptionAllocation = nullptr;
     std::unique_ptr<OSTime> osTime;
     std::unique_ptr<DriverInfo> driverInfo;
     std::unique_ptr<PerformanceCounters> performanceCounters;
diff --git a/runtime/memory_manager/os_agnostic_memory_manager.h b/runtime/memory_manager/os_agnostic_memory_manager.h
index 2f6d6e2f96..a3abb0d1ba 100644
--- a/runtime/memory_manager/os_agnostic_memory_manager.h
+++ b/runtime/memory_manager/os_agnostic_memory_manager.h
@@ -88,4 +88,5 @@ class OsAgnosticMemoryManager : public MemoryManager {
     unsigned long long counter = 0;
     bool fakeBigAllocations = false;
 };
+
 } // namespace NEO
diff --git a/unit_tests/command_queue/enqueue_command_without_kernel_tests.cpp b/unit_tests/command_queue/enqueue_command_without_kernel_tests.cpp
index 1d039c2d3a..f950b84d8b 100644
--- a/unit_tests/command_queue/enqueue_command_without_kernel_tests.cpp
+++ b/unit_tests/command_queue/enqueue_command_without_kernel_tests.cpp
@@ -54,10 +54,7 @@ HWTEST_F(EnqueueHandlerTest, whenEnqueueCommandWithoutKernelThenPassCorrectDispa
     mockCsr->initializeTagAllocation();
     auto oldCsr = mockCmdQ->engine->commandStreamReceiver;
     mockCmdQ->engine->commandStreamReceiver = mockCsr.get();
-
-    AllocationProperties properties(1, GraphicsAllocation::AllocationType::PREEMPTION);
-    auto preemptionAllocation = executionEnvironment->memoryManager->allocateGraphicsMemoryWithProperties(properties);
-    mockCsr->setPreemptionCsrAllocation(preemptionAllocation);
+    mockCsr->createPreemptionAllocation();
 
     auto blocking = true;
     TimestampPacketContainer previousTimestampPacketNodes;
@@ -69,7 +66,6 @@ HWTEST_F(EnqueueHandlerTest, whenEnqueueCommandWithoutKernelThenPassCorrectDispa
     EXPECT_EQ(mockCmdQ->isMultiEngineQueue(), mockCsr->passedDispatchFlags.multiEngineQueue);
     EXPECT_EQ(pDevice->getPreemptionMode(), mockCsr->passedDispatchFlags.preemptionMode);
     mockCmdQ->engine->commandStreamReceiver = oldCsr;
-    executionEnvironment->memoryManager->freeGraphicsMemory(preemptionAllocation);
 }
 
 HWTEST_F(EnqueueHandlerTest, GivenCommandStreamWithoutKernelAndZeroSurfacesWhenEnqueuedHandlerThenUsedSizeEqualZero) {
diff --git a/unit_tests/command_stream/command_stream_receiver_flush_task_1_tests.cpp b/unit_tests/command_stream/command_stream_receiver_flush_task_1_tests.cpp
index 02f3d71b68..dc6d8d4a0b 100644
--- a/unit_tests/command_stream/command_stream_receiver_flush_task_1_tests.cpp
+++ b/unit_tests/command_stream/command_stream_receiver_flush_task_1_tests.cpp
@@ -781,6 +781,7 @@ struct CommandStreamReceiverHwLog : public UltCommandStreamReceiver<FamilyType>
 HWTEST_F(CommandStreamReceiverFlushTaskTests, flushTaskWithBothCSCallsFlushOnce) {
     CommandStreamReceiverHwLog<FamilyType> commandStreamReceiver(*pDevice->executionEnvironment);
     commandStreamReceiver.initializeTagAllocation();
+    commandStreamReceiver.createPreemptionAllocation();
     commandStream.getSpace(sizeof(typename FamilyType::MI_NOOP));
 
     commandStreamReceiver.setupContext(*pDevice->getDefaultEngine().osContext);
diff --git a/unit_tests/command_stream/command_stream_receiver_flush_task_3_tests.cpp b/unit_tests/command_stream/command_stream_receiver_flush_task_3_tests.cpp
index a65329c5b5..1d0af2100e 100644
--- a/unit_tests/command_stream/command_stream_receiver_flush_task_3_tests.cpp
+++ b/unit_tests/command_stream/command_stream_receiver_flush_task_3_tests.cpp
@@ -669,8 +669,12 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, givenCsrInBatchingModeWhenTotalRes
     auto mockCsr = new MockCsrHw2<FamilyType>(*executionEnvironment);
     executionEnvironment->commandStreamReceivers.resize(1);
     executionEnvironment->commandStreamReceivers[0].push_back(std::unique_ptr<CommandStreamReceiver>(mockCsr));
+
+    if (pDevice->getPreemptionMode() == PreemptionMode::MidThread || pDevice->isSourceLevelDebuggerActive()) {
+        mockCsr->createPreemptionAllocation();
+    }
+
     mockCsr->initializeTagAllocation();
-    mockCsr->setPreemptionCsrAllocation(pDevice->getPreemptionAllocation());
     mockCsr->overrideDispatchPolicy(DispatchMode::BatchedDispatch);
     mockCsr->setupContext(*pDevice->getDefaultEngine().osContext);
 
diff --git a/unit_tests/event/event_tests.cpp b/unit_tests/event/event_tests.cpp
index b19556d1a9..7264f05159 100644
--- a/unit_tests/event/event_tests.cpp
+++ b/unit_tests/event/event_tests.cpp
@@ -958,6 +958,7 @@ HWTEST_F(InternalsEventTest, GivenBufferWithoutZeroCopyOnCommandMapOrUnmapFlushe
     MockNonZeroCopyBuff buffer(executionStamp);
     MockCsr<FamilyType> csr(executionStamp, *pDevice->executionEnvironment);
     csr.setTagAllocation(pDevice->getDefaultEngine().commandStreamReceiver->getTagAllocation());
+    csr.createPreemptionAllocation();
     csr.setupContext(*pDevice->getDefaultEngine().osContext);
 
     MemObjSizeArray size = {{4, 1, 1}};
diff --git a/unit_tests/gen9/test_preemption_gen9.cpp b/unit_tests/gen9/test_preemption_gen9.cpp
index 4ad0c371f0..f0202f113c 100644
--- a/unit_tests/gen9/test_preemption_gen9.cpp
+++ b/unit_tests/gen9/test_preemption_gen9.cpp
@@ -94,7 +94,7 @@ GEN9TEST_F(Gen9ThreadGroupPreemptionEnqueueKernelTest, givenSecondEnqueueWithThe
     auto &csr = pDevice->getUltCommandStreamReceiver<FamilyType>();
     csr.getMemoryManager()->setForce32BitAllocations(false);
     csr.setMediaVFEStateDirty(false);
-    auto csrSurface = csr.getPreemptionCsrAllocation();
+    auto csrSurface = csr.getPreemptionAllocation();
     EXPECT_EQ(nullptr, csrSurface);
     size_t off[3] = {0, 0, 0};
     size_t gws[3] = {1, 1, 1};
@@ -128,7 +128,7 @@ GEN9TEST_F(Gen9ThreadGroupPreemptionEnqueueKernelTest, givenSecondEnqueueWithThe
     auto &csr = pDevice->getUltCommandStreamReceiver<FamilyType>();
     csr.getMemoryManager()->setForce32BitAllocations(false);
     csr.setMediaVFEStateDirty(false);
-    auto csrSurface = csr.getPreemptionCsrAllocation();
+    auto csrSurface = csr.getPreemptionAllocation();
     EXPECT_EQ(nullptr, csrSurface);
     HardwareParse hwCsrParser;
     HardwareParse hwCmdQParser;
@@ -265,7 +265,7 @@ GEN9TEST_F(Gen9MidThreadPreemptionEnqueueKernelTest, givenSecondEnqueueWithTheSa
     auto &csr = pDevice->getUltCommandStreamReceiver<FamilyType>();
     csr.getMemoryManager()->setForce32BitAllocations(false);
     csr.setMediaVFEStateDirty(false);
-    auto csrSurface = csr.getPreemptionCsrAllocation();
+    auto csrSurface = csr.getPreemptionAllocation();
     ASSERT_NE(nullptr, csrSurface);
     HardwareParse hwCsrParser;
     HardwareParse hwCmdQParser;
@@ -345,7 +345,7 @@ GEN9TEST_F(Gen9MidThreadPreemptionEnqueueKernelTest, givenSecondEnqueueWithTheSa
     auto &csr = pDevice->getUltCommandStreamReceiver<FamilyType>();
     csr.getMemoryManager()->setForce32BitAllocations(false);
     csr.setMediaVFEStateDirty(false);
-    auto csrSurface = csr.getPreemptionCsrAllocation();
+    auto csrSurface = csr.getPreemptionAllocation();
     ASSERT_NE(nullptr, csrSurface);
     HardwareParse hwCsrParser;
     HardwareParse hwCmdQParser;
diff --git a/unit_tests/libult/ult_command_stream_receiver.h b/unit_tests/libult/ult_command_stream_receiver.h
index 289444ac65..b21230281c 100644
--- a/unit_tests/libult/ult_command_stream_receiver.h
+++ b/unit_tests/libult/ult_command_stream_receiver.h
@@ -70,18 +70,10 @@ class UltCommandStreamReceiver : public CommandStreamReceiverHw<GfxFamily>, publ
     using BaseClass::CommandStreamReceiver::waitForTaskCountAndCleanAllocationList;
 
     virtual ~UltCommandStreamReceiver() override {
-        if (tempPreemptionLocation) {
-            this->setPreemptionCsrAllocation(nullptr);
-        }
     }
 
     UltCommandStreamReceiver(ExecutionEnvironment &executionEnvironment) : BaseClass(executionEnvironment), recursiveLockCounter(0) {
-        if (executionEnvironment.getHardwareInfo()->capabilityTable.defaultPreemptionMode == PreemptionMode::MidThread) {
-            tempPreemptionLocation = std::make_unique<GraphicsAllocation>(GraphicsAllocation::AllocationType::UNKNOWN, nullptr, 0, 0, 0, MemoryPool::MemoryNull, false);
-            this->preemptionCsrAllocation = tempPreemptionLocation.get();
-        }
     }
-
     static CommandStreamReceiver *create(bool withAubDump, ExecutionEnvironment &executionEnvironment) {
         return new UltCommandStreamReceiver<GfxFamily>(executionEnvironment);
     }
@@ -114,6 +106,7 @@ class UltCommandStreamReceiver : public CommandStreamReceiverHw<GfxFamily>, publ
     size_t getPreferredTagPoolSize() const override {
         return BaseClass::getPreferredTagPoolSize() + 1;
     }
+    void setPreemptionAllocation(GraphicsAllocation *allocation) { this->preemptionAllocation = allocation; }
 
     bool waitForCompletionWithTimeout(bool enableTimeout, int64_t timeoutMicroseconds, uint32_t taskCountToWait) override {
         latestWaitForCompletionWithTimeoutTaskCount.store(taskCountToWait);
@@ -121,7 +114,7 @@ class UltCommandStreamReceiver : public CommandStreamReceiverHw<GfxFamily>, publ
     }
 
     void overrideCsrSizeReqFlags(CsrSizeRequestFlags &flags) { this->csrSizeRequestFlags = flags; }
-    GraphicsAllocation *getPreemptionCsrAllocation() const { return this->preemptionCsrAllocation; }
+    GraphicsAllocation *getPreemptionAllocation() const { return this->preemptionAllocation; }
 
     void makeResident(GraphicsAllocation &gfxAllocation) override {
         if (storeMakeResidentAllocations) {
@@ -187,9 +180,5 @@ class UltCommandStreamReceiver : public CommandStreamReceiverHw<GfxFamily>, publ
     uint32_t latestSentTaskCountValueDuringFlush = 0;
     uint32_t blitBufferCalled = 0;
     std::atomic<uint32_t> latestWaitForCompletionWithTimeoutTaskCount{0};
-
-  protected:
-    std::unique_ptr<GraphicsAllocation> tempPreemptionLocation;
 };
-
 } // namespace NEO
diff --git a/unit_tests/mocks/mock_device.cpp b/unit_tests/mocks/mock_device.cpp
index b6e83a255a..1b7f6fba3f 100644
--- a/unit_tests/mocks/mock_device.cpp
+++ b/unit_tests/mocks/mock_device.cpp
@@ -62,7 +62,10 @@ void MockDevice::resetCommandStreamReceiver(CommandStreamReceiver *newCsr) {
 void MockDevice::resetCommandStreamReceiver(CommandStreamReceiver *newCsr, uint32_t engineIndex) {
     executionEnvironment->commandStreamReceivers[getDeviceIndex()][engineIndex].reset(newCsr);
     executionEnvironment->commandStreamReceivers[getDeviceIndex()][engineIndex]->initializeTagAllocation();
-    executionEnvironment->commandStreamReceivers[getDeviceIndex()][engineIndex]->setPreemptionCsrAllocation(preemptionAllocation);
+
+    if (preemptionMode == PreemptionMode::MidThread || isSourceLevelDebuggerActive()) {
+        executionEnvironment->commandStreamReceivers[getDeviceIndex()][engineIndex]->createPreemptionAllocation();
+    }
     this->engines[engineIndex].commandStreamReceiver = newCsr;
 
     auto osContext = this->engines[engineIndex].osContext;
diff --git a/unit_tests/mocks/mock_device.h b/unit_tests/mocks/mock_device.h
index cbf5c52c04..ad6c7a17f3 100644
--- a/unit_tests/mocks/mock_device.h
+++ b/unit_tests/mocks/mock_device.h
@@ -101,14 +101,6 @@ class MockDevice : public Device {
         return createWithExecutionEnvironment<T>(pHwInfo, executionEnvironment, 0u);
     }
 
-    void allocatePreemptionAllocationIfNotPresent() {
-        if (this->preemptionAllocation == nullptr) {
-            if (preemptionMode == PreemptionMode::MidThread || isSourceLevelDebuggerActive()) {
-                this->preemptionAllocation = executionEnvironment->memoryManager->allocateGraphicsMemoryWithProperties(getAllocationPropertiesForPreemption());
-                this->engines[defaultEngineIndex].commandStreamReceiver->setPreemptionCsrAllocation(preemptionAllocation);
-            }
-        }
-    }
     std::unique_ptr<MemoryManager> mockMemoryManager;
 
   private:
diff --git a/unit_tests/os_interface/linux/drm_command_stream_tests.cpp b/unit_tests/os_interface/linux/drm_command_stream_tests.cpp
index c5766430d4..82eaff88b3 100644
--- a/unit_tests/os_interface/linux/drm_command_stream_tests.cpp
+++ b/unit_tests/os_interface/linux/drm_command_stream_tests.cpp
@@ -771,7 +771,7 @@ class DrmCommandStreamBatchingTests : public Test<DrmCommandStreamEnhancedFixtur
             GlobalMockSipProgram::sipProgram->resetAllocation(device->getMemoryManager()->allocateGraphicsMemoryWithProperties(MockAllocationProperties{MemoryConstants::pageSize}));
         }
         tagAllocation = static_cast<DrmAllocation *>(device->getDefaultEngine().commandStreamReceiver->getTagAllocation());
-        preemptionAllocation = static_cast<DrmAllocation *>(device->getPreemptionAllocation());
+        preemptionAllocation = static_cast<DrmAllocation *>(device->getDefaultEngine().commandStreamReceiver->getPreemptionAllocation());
     }
     void TearDown() override {
         if (PreemptionHelper::getDefaultPreemptionMode(*platformDevices[0]) == PreemptionMode::MidThread) {
@@ -828,7 +828,6 @@ TEST_F(DrmCommandStreamBatchingTests, givenCsrWhenDispatchPolicyIsSetToBatchingT
     tCsr->makeResident(*dummyAllocation);
 
     tCsr->setTagAllocation(tagAllocation);
-    tCsr->setPreemptionCsrAllocation(preemptionAllocation);
     DispatchFlags dispatchFlags;
     dispatchFlags.preemptionMode = PreemptionHelper::getDefaultPreemptionMode(device->getHardwareInfo());
     tCsr->flushTask(cs, 0u, cs, cs, cs, 0u, dispatchFlags, *device);
@@ -880,7 +879,6 @@ TEST_F(DrmCommandStreamBatchingTests, givenRecordedCommandBufferWhenItIsSubmitte
     IndirectHeap cs(commandBuffer);
 
     tCsr->setTagAllocation(tagAllocation);
-    tCsr->setPreemptionCsrAllocation(preemptionAllocation);
     auto &submittedCommandBuffer = tCsr->getCS(1024);
     //use some bytes
     submittedCommandBuffer.getSpace(4);
@@ -933,7 +931,6 @@ TEST_F(DrmCommandStreamBatchingTests, givenRecordedCommandBufferWhenItIsSubmitte
     EXPECT_EQ(ioctlExecCnt + ioctlUserPtrCnt, this->mock->ioctl_cnt.total);
 
     mm->freeGraphicsMemory(commandBuffer);
-    tCsr->setPreemptionCsrAllocation(nullptr);
 }
 
 typedef Test<DrmCommandStreamEnhancedFixture> DrmCommandStreamLeaksTest;
diff --git a/unit_tests/os_interface/windows/device_command_stream_tests.cpp b/unit_tests/os_interface/windows/device_command_stream_tests.cpp
index 8c7b808e26..2b13bed110 100644
--- a/unit_tests/os_interface/windows/device_command_stream_tests.cpp
+++ b/unit_tests/os_interface/windows/device_command_stream_tests.cpp
@@ -52,20 +52,21 @@ class WddmCommandStreamFixture {
     DeviceCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME> *csr;
     MockWddmMemoryManager *memoryManager = nullptr;
     WddmMock *wddm = nullptr;
+
     DebugManagerStateRestore stateRestore;
 
     virtual void SetUp() {
         HardwareInfo *hwInfo = nullptr;
         DebugManager.flags.CsrDispatchMode.set(static_cast<uint32_t>(DispatchMode::ImmediateDispatch));
         auto executionEnvironment = getExecutionEnvironmentImpl(hwInfo);
-        wddm = static_cast<WddmMock *>(executionEnvironment->osInterface->get()->getWddm());
-
-        csr = new WddmCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME>(*executionEnvironment);
 
         memoryManager = new MockWddmMemoryManager(*executionEnvironment);
         executionEnvironment->memoryManager.reset(memoryManager);
+        wddm = static_cast<WddmMock *>(executionEnvironment->osInterface->get()->getWddm());
 
+        csr = new WddmCommandStreamReceiver<DEFAULT_TEST_FAMILY_NAME>(*executionEnvironment);
         device.reset(MockDevice::create<MockDevice>(executionEnvironment, 0u));
+
         device->resetCommandStreamReceiver(csr);
         ASSERT_NE(nullptr, device);
     }
@@ -129,15 +130,9 @@ class WddmCommandStreamWithMockGdiFixture {
         device->resetCommandStreamReceiver(this->csr);
         ASSERT_NE(nullptr, device);
         this->csr->overrideRecorededCommandBuffer(*device);
-        if (device->getPreemptionMode() == PreemptionMode::MidThread) {
-            preemptionAllocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{MemoryConstants::pageSize});
-        }
     }
 
     virtual void TearDown() {
-        if (preemptionAllocation) {
-            memoryManager->freeGraphicsMemory(preemptionAllocation);
-        }
         wddm = nullptr;
     }
 };
@@ -735,7 +730,6 @@ HWTEST_F(WddmCommandStreamMockGdiTest, givenRecordedCommandBufferWhenItIsSubmitt
     auto sshAlloc = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{MemoryConstants::pageSize});
 
     auto tagAllocation = csr->getTagAllocation();
-    csr->setPreemptionCsrAllocation(preemptionAllocation);
 
     LinearStream cs(commandBuffer);
     IndirectHeap dsh(dshAlloc);
diff --git a/unit_tests/preamble/preamble_tests.cpp b/unit_tests/preamble/preamble_tests.cpp
index a99f6520e1..dcf610a696 100644
--- a/unit_tests/preamble/preamble_tests.cpp
+++ b/unit_tests/preamble/preamble_tests.cpp
@@ -149,13 +149,12 @@ HWTEST_F(PreambleTest, givenKernelDebuggingActiveWhenPreambleIsProgrammedThenPro
     auto miLoadRegImmCountWithoutDebugging = cmdList.size();
 
     mockDevice->setSourceLevelDebuggerActive(true);
-    mockDevice->allocatePreemptionAllocationIfNotPresent();
+    auto preemptionAllocation = mockDevice->getCommandStreamReceiver().getPreemptionAllocation();
 
     StackVec<char, 8192> preambleBuffer2(8192);
     preambleStream.replaceBuffer(&*preambleBuffer2.begin(), preambleBuffer2.size());
     PreambleHelper<FamilyType>::programPreamble(&preambleStream, *mockDevice, 0U,
-                                                ThreadArbitrationPolicy::RoundRobin, mockDevice->getPreemptionAllocation());
-
+                                                ThreadArbitrationPolicy::RoundRobin, preemptionAllocation);
     HardwareParse hwParser2;
     hwParser2.parseCommands<FamilyType>(preambleStream);
     cmdList = hwParser2.getCommandsList<MI_LOAD_REGISTER_IMM>();
diff --git a/unit_tests/preemption/preemption_tests.cpp b/unit_tests/preemption/preemption_tests.cpp
index 210f484353..b40b06fde4 100644
--- a/unit_tests/preemption/preemption_tests.cpp
+++ b/unit_tests/preemption/preemption_tests.cpp
@@ -439,6 +439,39 @@ HWTEST_P(PreemptionTest, whenInNonMidThreadModeThenCsrBaseAddressIsNotProgrammed
     EXPECT_EQ(0u, cmdStream.getUsed());
 }
 
+HWTEST_P(PreemptionTest, whenFailToCreatePreemptionAllocationThenFailToCreateDevice) {
+
+    class MockUltCsr : public UltCommandStreamReceiver<FamilyType> {
+
+      public:
+        MockUltCsr(ExecutionEnvironment &executionEnvironment) : UltCommandStreamReceiver<FamilyType>(executionEnvironment) {
+        }
+
+        bool createPreemptionAllocation() override {
+            return false;
+        }
+    };
+
+    class MockDeviceReturnedDebuggerActive : public MockDevice {
+      public:
+        MockDeviceReturnedDebuggerActive(ExecutionEnvironment *executionEnvironment, uint32_t deviceIndex)
+            : MockDevice(executionEnvironment, deviceIndex) {}
+        bool isSourceLevelDebuggerActive() const override {
+            return true;
+        }
+    };
+
+    ExecutionEnvironment *executionEnvironment = platformImpl->peekExecutionEnvironment();
+
+    platformImpl->peekExecutionEnvironment()->commandStreamReceivers.resize(1);
+    platformImpl->peekExecutionEnvironment()->commandStreamReceivers[0].resize(2);
+    executionEnvironment->commandStreamReceivers[0][1].reset(new MockUltCsr(*executionEnvironment));
+    executionEnvironment->commandStreamReceivers[0][0].reset(new MockUltCsr(*executionEnvironment));
+
+    std::unique_ptr<MockDevice> mockDevice(MockDevice::create<MockDeviceReturnedDebuggerActive>(executionEnvironment, 0));
+    EXPECT_EQ(nullptr, mockDevice);
+}
+
 INSTANTIATE_TEST_CASE_P(
     NonMidThread,
     PreemptionTest,
@@ -452,21 +485,22 @@ HWTEST_F(MidThreadPreemptionTests, createCsrSurfaceNoWa) {
     ASSERT_NE(nullptr, mockDevice.get());
 
     auto &csr = mockDevice->getUltCommandStreamReceiver<FamilyType>();
-    MemoryAllocation *csrSurface = static_cast<MemoryAllocation *>(csr.getPreemptionCsrAllocation());
+    MemoryAllocation *csrSurface = static_cast<MemoryAllocation *>(csr.getPreemptionAllocation());
     ASSERT_NE(nullptr, csrSurface);
     EXPECT_FALSE(csrSurface->uncacheable);
 
-    GraphicsAllocation *devCsrSurface = mockDevice->getPreemptionAllocation();
+    GraphicsAllocation *devCsrSurface = csr.getPreemptionAllocation();
     EXPECT_EQ(csrSurface, devCsrSurface);
 }
 
 HWTEST_F(MidThreadPreemptionTests, givenMidThreadPreemptionWhenFailingOnCsrSurfaceAllocationThenFailToCreateDevice) {
+
     class FailingMemoryManager : public OsAgnosticMemoryManager {
       public:
         FailingMemoryManager(ExecutionEnvironment &executionEnvironment) : OsAgnosticMemoryManager(executionEnvironment) {}
 
         GraphicsAllocation *allocateGraphicsMemoryWithAlignment(const AllocationData &allocationData) override {
-            if (++allocateGraphicsMemoryCount > HwHelper::get(platformDevices[0]->platform.eRenderCoreFamily).getGpgpuEngineInstances().size()) {
+            if (++allocateGraphicsMemoryCount > HwHelper::get(platformDevices[0]->platform.eRenderCoreFamily).getGpgpuEngineInstances().size() - 1) {
                 return nullptr;
             }
             return OsAgnosticMemoryManager::allocateGraphicsMemoryWithAlignment(allocationData);
@@ -489,11 +523,11 @@ HWTEST_F(MidThreadPreemptionTests, createCsrSurfaceWa) {
     ASSERT_NE(nullptr, mockDevice.get());
 
     auto &csr = mockDevice->getUltCommandStreamReceiver<FamilyType>();
-    MemoryAllocation *csrSurface = static_cast<MemoryAllocation *>(csr.getPreemptionCsrAllocation());
+    MemoryAllocation *csrSurface = static_cast<MemoryAllocation *>(csr.getPreemptionAllocation());
     ASSERT_NE(nullptr, csrSurface);
     EXPECT_TRUE(csrSurface->uncacheable);
 
-    GraphicsAllocation *devCsrSurface = mockDevice->getPreemptionAllocation();
+    GraphicsAllocation *devCsrSurface = csr.getPreemptionAllocation();
     EXPECT_EQ(csrSurface, devCsrSurface);
 }