diff --git a/shared/source/memory_manager/gfx_partition.cpp b/shared/source/memory_manager/gfx_partition.cpp
index 524f378fbc..13aba39425 100644
--- a/shared/source/memory_manager/gfx_partition.cpp
+++ b/shared/source/memory_manager/gfx_partition.cpp
@@ -269,7 +269,7 @@ bool GfxPartition::init(uint64_t gpuAddressSpace, size_t cpuAddressRangeSizeToRe
             gfxBase = 0ull;
             heapInit(HeapIndex::heapSvm, 0ull, 0ull);
         } else {
-            if (!initAdditionalRange(cpuVirtualAddressSize, gpuAddressSpace, gfxBase, gfxTop, rootDeviceIndex, numRootDevices, systemMemorySize)) {
+            if (!initAdditionalRange(cpuVirtualAddressSize, gpuAddressSpace, gfxBase, gfxTop, rootDeviceIndex, systemMemorySize)) {
                 return false;
             }
         }
@@ -318,7 +318,7 @@ bool GfxPartition::init(uint64_t gpuAddressSpace, size_t cpuAddressRangeSizeToRe
     return true;
 }
 
-bool GfxPartition::initAdditionalRange(uint32_t cpuVirtualAddressSize, uint64_t gpuAddressSpace, uint64_t &gfxBase, uint64_t &gfxTop, uint32_t rootDeviceIndex, size_t numRootDevices, uint64_t systemMemorySize) {
+bool GfxPartition::initAdditionalRange(uint32_t cpuVirtualAddressSize, uint64_t gpuAddressSpace, uint64_t &gfxBase, uint64_t &gfxTop, uint32_t rootDeviceIndex, uint64_t systemMemorySize) {
     /*
      * 57-bit Full Range SVM gfx layout:
      *
@@ -381,8 +381,7 @@ bool GfxPartition::initAdditionalRange(uint32_t cpuVirtualAddressSize, uint64_t
 
     // Init HEAP_EXTENDED only for 57 bit GPU
     if (gpuAddressSpace == maxNBitValue(57)) {
-        // Split HEAP_EXTENDED among root devices (like HEAP_STANDARD64K)
-        auto heapExtendedSize = alignDown((maxNBitValue(48) + 1) / numRootDevices, GfxPartition::heapGranularity);
+        auto heapExtendedSize = alignDown((maxNBitValue(48) + 1), GfxPartition::heapGranularity);
         heapInit(HeapIndex::heapExtended, maxNBitValue(57 - 1) + 1 + rootDeviceIndex * heapExtendedSize, heapExtendedSize);
     }
 
diff --git a/shared/source/memory_manager/gfx_partition.h b/shared/source/memory_manager/gfx_partition.h
index ff06a603ea..6a96b6a321 100644
--- a/shared/source/memory_manager/gfx_partition.h
+++ b/shared/source/memory_manager/gfx_partition.h
@@ -104,7 +104,7 @@ class GfxPartition {
     static const std::array<HeapIndex, 8> heapNonSvmNames;
 
   protected:
-    bool initAdditionalRange(uint32_t cpuAddressWidth, uint64_t gpuAddressSpace, uint64_t &gfxBase, uint64_t &gfxTop, uint32_t rootDeviceIndex, size_t numRootDevices, uint64_t systemMemorySize);
+    bool initAdditionalRange(uint32_t cpuAddressWidth, uint64_t gpuAddressSpace, uint64_t &gfxBase, uint64_t &gfxTop, uint32_t rootDeviceIndex, uint64_t systemMemorySize);
 
     class Heap {
       public:
diff --git a/shared/test/unit_test/memory_manager/gfx_partition_tests.cpp b/shared/test/unit_test/memory_manager/gfx_partition_tests.cpp
index 01bf5e5c0d..52c34e0940 100644
--- a/shared/test/unit_test/memory_manager/gfx_partition_tests.cpp
+++ b/shared/test/unit_test/memory_manager/gfx_partition_tests.cpp
@@ -908,7 +908,7 @@ TEST(GfxPartitionTest, givenGpuAddressSpaceIs57BitAndSeveralRootDevicesThenHeapE
         EXPECT_TRUE(gfxPartition.init(maxNBitValue(57), reservedCpuAddressRangeSize, rootDeviceIndex, numRootDevices, false, 0u, gfxTop));
 
         auto heapExtendedTotalSize = maxNBitValue(48) + 1;
-        auto heapExtendedSize = alignDown(heapExtendedTotalSize / numRootDevices, GfxPartition::heapGranularity);
+        auto heapExtendedSize = alignDown(heapExtendedTotalSize, GfxPartition::heapGranularity);
 
         EXPECT_EQ(heapExtendedSize, gfxPartition.getHeapSize(HeapIndex::heapExtended));
         EXPECT_EQ(maxNBitValue(56) + 1 + rootDeviceIndex * heapExtendedSize, gfxPartition.getHeapBase(HeapIndex::heapExtended));
@@ -922,7 +922,7 @@ TEST(GfxPartitionTest, givenGpuAddressSpaceIs57BitAndSeveralRootDevicesThenHeapE
         EXPECT_TRUE(gfxPartition.init(maxNBitValue(57), reservedCpuAddressRangeSize, rootDeviceIndex, numRootDevices, false, 0u, gfxTop));
 
         auto heapExtendedTotalSize = maxNBitValue(48) + 1;
-        auto heapExtendedSize = alignDown(heapExtendedTotalSize / numRootDevices, GfxPartition::heapGranularity);
+        auto heapExtendedSize = alignDown(heapExtendedTotalSize, GfxPartition::heapGranularity);
 
         EXPECT_EQ(heapExtendedSize, gfxPartition.getHeapSize(HeapIndex::heapExtended));
         EXPECT_EQ(maxNBitValue(56) + 1 + rootDeviceIndex * heapExtendedSize, gfxPartition.getHeapBase(HeapIndex::heapExtended));