diff --git a/src/Plugin/DisplayVariant/BlockDisplayVariant.php b/src/Plugin/DisplayVariant/BlockDisplayVariant.php
index d652104..fb37f7b 100644
--- a/src/Plugin/DisplayVariant/BlockDisplayVariant.php
+++ b/src/Plugin/DisplayVariant/BlockDisplayVariant.php
@@ -129,11 +129,7 @@ class BlockDisplayVariant extends VariantBase implements ContextAwareVariantInte
    */
   public function build() {
     $page = $this->executable->getPage();
-
     // Set default page cache keys that include the page and display.
-    // @todo Make the whole display a pre_render callback to fully benefit from
-    //   render caching and avoid to re-calculate contexts and access for all
-    //   blocks.
     $build['regions']['#cache']['keys'] = [
       'page_manager_page',
       // The page ID.
@@ -142,8 +138,17 @@ class BlockDisplayVariant extends VariantBase implements ContextAwareVariantInte
       // @todo should have an API for this?
       $this->configuration['uuid'],
     ];
+    $build['#pre_render'][] = [$this, 'buildRegions'];
+    return $build;
+  }
 
-    $cacheability = CacheableMetadata::createFromRenderArray($build['regions'])
+  /**
+   * #pre_render callback for building the regions.
+   */
+  public function buildRegions(array $build) {
+    $page = $this->executable->getPage();
+
+    $cacheability = CacheableMetadata::createFromRenderArray($build)
       ->addCacheableDependency($page);
 
     $contexts = $this->getContexts();
@@ -153,8 +158,8 @@ class BlockDisplayVariant extends VariantBase implements ContextAwareVariantInte
       }
 
       $region_name = Html::getClass("block-region-$region");
-      $build['regions'][$region]['#prefix'] = '<div class="' . $region_name . '">';
-      $build['regions'][$region]['#suffix'] = '</div>';
+      $build[$region]['#prefix'] = '<div class="' . $region_name . '">';
+      $build[$region]['#suffix'] = '</div>';
 
       /** @var $blocks \Drupal\Core\Block\BlockPluginInterface[] */
       $weight = 0;
@@ -197,13 +202,13 @@ class BlockDisplayVariant extends VariantBase implements ContextAwareVariantInte
         if (!empty($block_build['#configuration']['label'])) {
           $block_build['#configuration']['label'] = SafeMarkup::checkPlain($block_build['#configuration']['label']);
         }
-        $build['regions'][$region][$block_id] = $block_build;
+        $build[$region][$block_id] = $block_build;
       }
     }
 
     $build['#title'] = $this->renderPageTitle($this->configuration['page_title']);
 
-    $cacheability->applyTo($build['regions']);
+    $cacheability->applyTo($build);
 
     return $build;
   }
diff --git a/tests/src/Unit/BlockDisplayVariantTest.php b/tests/src/Unit/BlockDisplayVariantTest.php
index 1bb03c9..d546cb4 100644
--- a/tests/src/Unit/BlockDisplayVariantTest.php
+++ b/tests/src/Unit/BlockDisplayVariantTest.php
@@ -202,7 +202,7 @@ class BlockDisplayVariantTest extends UnitTestCase {
     // metadata of accessible blocks is merged to avoid cache redirects when
     // possible.
     $expected_cache_page = [
-      'keys' => ['page_manager_page', 'page_id', 'UUID'],
+      //'keys' => ['page_manager_page', 'page_id', 'UUID'],
       'contexts' => ['url', 'user', 'user.permissions'],
       'tags' => ['block_plugin1:block_plugin_id', 'page:page_id'],
       'max-age' => 3600,
@@ -211,16 +211,17 @@ class BlockDisplayVariantTest extends UnitTestCase {
     // Build the variant and ensure that pre_render is set only for the first
     // block.
     $build = $display_variant->build();
-    $this->assertSame([$display_variant, 'buildBlock'], $build['regions']['top']['block1']['#pre_render'][0]);
-    $this->assertTrue(empty($build['regions']['top']['block2']));
-    $this->assertSame($expected_cache_block1, $build['regions']['top']['block1']['#cache']);
-    $this->assertSame($expected_cache_page, $build['regions']['#cache']);
+    $build = $display_variant->buildRegions($build);
+    $this->assertSame([$display_variant, 'buildBlock'], $build['top']['block1']['#pre_render'][0]);
+    $this->assertTrue(empty($build['top']['block2']));
+    $this->assertSame($expected_cache_block1, $build['top']['block1']['#cache']);
+    $this->assertSame($expected_cache_page, $build['#cache']);
 
     // Ensure that building the block returns the correct markup.
     $block1->build()->willReturn([
       '#markup' => 'block1_build_value',
     ]);
-    $block1_build = $display_variant->buildBlock($build['regions']['top']['block1']);
+    $block1_build = $display_variant->buildBlock($build['top']['block1']);
     $this->assertSame(['#markup' => 'block1_build_value'], $block1_build['content']);
   }