diff --git a/renderdoc/driver/d3d12/d3d12_command_list4_wrap.cpp b/renderdoc/driver/d3d12/d3d12_command_list4_wrap.cpp
index 6bfea1045..fc9d4ba4c 100644
--- a/renderdoc/driver/d3d12/d3d12_command_list4_wrap.cpp
+++ b/renderdoc/driver/d3d12/d3d12_command_list4_wrap.cpp
@@ -847,8 +847,7 @@ bool WrappedID3D12GraphicsCommandList::PatchAccStructBlasAddress(
     ID3D12Resource *addressPairRes = m_pDevice->GetBLASAddressBufferResource();
     D3D12_GPU_VIRTUAL_ADDRESS addressPairResAddress = addressPairRes->GetGPUVirtualAddress();
 
-    // TODO: Update to gather the right count.
-    uint64_t addressCount = 0;
+    uint64_t addressCount = m_pDevice->GetBLASAddressCount();
 
     dxrCmd->SetPipelineState(patchInfo.m_pipeline);
     dxrCmd->SetComputeRootSignature(patchInfo.m_rootSignature);
@@ -930,7 +929,6 @@ bool WrappedID3D12GraphicsCommandList::Serialise_BuildRaytracingAccelerationStru
           }
         }
 
-        // AMD TODO: Find out do we need pre callback before build Acc struct call
         dxrCmd->BuildRaytracingAccelerationStructure(&AccStructDesc, NumPostbuildInfoDescs,
                                                      pPostbuildInfoDescs);
       }
diff --git a/renderdoc/driver/d3d12/d3d12_device.cpp b/renderdoc/driver/d3d12/d3d12_device.cpp
index 48f4fa749..abcb27d87 100644
--- a/renderdoc/driver/d3d12/d3d12_device.cpp
+++ b/renderdoc/driver/d3d12/d3d12_device.cpp
@@ -3122,6 +3122,8 @@ void WrappedID3D12Device::UploadBLASBufferAddresses()
     }
   }
 
+  m_blasAddressCount = (uint32_t)blasAddressPair.size();
+
   uint64_t requiredSize = blasAddressPair.size() * sizeof(BlasAddressPair);
 
   D3D12_RESOURCE_DESC addressBufferResDesc;
diff --git a/renderdoc/driver/d3d12/d3d12_device.h b/renderdoc/driver/d3d12/d3d12_device.h
index 1b7326ab5..927f64bd3 100644
--- a/renderdoc/driver/d3d12/d3d12_device.h
+++ b/renderdoc/driver/d3d12/d3d12_device.h
@@ -992,6 +992,12 @@ public:
   // pending at once then they will be split up and issued in multiple calls.
   static const UINT executeListsMaxSize = 50;
 
+  // GPU side pair-struct for remapping BLAS addresses just in time before builds
+  uint32_t m_blasAddressCount = 0;
+  ID3D12Resource *m_blasAddressBufferResource = NULL;
+  ID3D12Resource *m_blasAddressBufferUploadResource = NULL;
+  bool m_addressBufferUploaded = false;
+
   ID3D12GraphicsCommandListX *GetNewList();
   ID3D12GraphicsCommandListX *GetInitialStateList();
 
@@ -1260,9 +1266,7 @@ public:
 
   void UploadBLASBufferAddresses();
   ID3D12Resource *GetBLASAddressBufferResource() const { return m_blasAddressBufferResource; }
-  ID3D12Resource *m_blasAddressBufferResource = NULL;
-  ID3D12Resource *m_blasAddressBufferUploadResource = NULL;
-  bool m_addressBufferUploaded = false;
+  uint32_t GetBLASAddressCount() const { return m_blasAddressCount; }
 
   void ReleaseResource(ID3D12DeviceChild *pResource);