diff --git a/shared/source/command_stream/tbx_command_stream_receiver_hw.h b/shared/source/command_stream/tbx_command_stream_receiver_hw.h
index 049056842f..1b3320208e 100644
--- a/shared/source/command_stream/tbx_command_stream_receiver_hw.h
+++ b/shared/source/command_stream/tbx_command_stream_receiver_hw.h
@@ -31,6 +31,9 @@ class TbxCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFa
     uint32_t getMaskAndValueForPollForCompletion() const;
     bool getpollNotEqualValueForPollForCompletion() const;
     void flushSubmissionsAndDownloadAllocations(TaskCountType taskCount, bool skipAllocationsDownload);
+    MOCKABLE_VIRTUAL uint64_t getNonBlockingDownloadTimeoutMs() const {
+        return 2000; // 2s
+    }
 
   public:
     using CommandStreamReceiverSimulatedCommonHw<GfxFamily>::initAdditionalMMIO;
diff --git a/shared/source/command_stream/tbx_command_stream_receiver_hw.inl b/shared/source/command_stream/tbx_command_stream_receiver_hw.inl
index fed032000c..7f15b48c72 100644
--- a/shared/source/command_stream/tbx_command_stream_receiver_hw.inl
+++ b/shared/source/command_stream/tbx_command_stream_receiver_hw.inl
@@ -578,7 +578,6 @@ void TbxCommandStreamReceiverHw<GfxFamily>::downloadAllocationTbx(GraphicsAlloca
 template <typename GfxFamily>
 void TbxCommandStreamReceiverHw<GfxFamily>::downloadAllocations(bool blockingWait, TaskCountType taskCount) {
     volatile TagAddressType *pollAddress = this->getTagAddress();
-    constexpr uint64_t timeoutMs = 1000 * 2; // 2s
 
     auto waitTaskCount = std::min(taskCount, this->latestFlushedTaskCount.load());
 
@@ -593,7 +592,7 @@ void TbxCommandStreamReceiverHw<GfxFamily>::downloadAllocations(bool blockingWai
                 if (!blockingWait) {
                     // Additional delay to reach PC in case of Event wait
                     timeDiff = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::high_resolution_clock::now() - startTime).count();
-                    if (timeDiff > timeoutMs) {
+                    if (timeDiff > getNonBlockingDownloadTimeoutMs()) {
                         return;
                     }
                 }
diff --git a/shared/test/common/mocks/mock_tbx_csr.h b/shared/test/common/mocks/mock_tbx_csr.h
index fcef4d1317..1448990489 100644
--- a/shared/test/common/mocks/mock_tbx_csr.h
+++ b/shared/test/common/mocks/mock_tbx_csr.h
@@ -111,6 +111,11 @@ struct MockTbxCsrRegisterDownloadedAllocations : TbxCommandStreamReceiverHw<GfxF
         obtainUniqueOwnershipCalled++;
         return TbxCommandStreamReceiverHw<GfxFamily>::obtainUniqueOwnership();
     }
+
+    uint64_t getNonBlockingDownloadTimeoutMs() const override {
+        return 1;
+    }
+
     std::set<GraphicsAllocation *> downloadedAllocations;
     bool flushBatchedSubmissionsCalled = false;
     bool flushTagCalled = false;