diff --git a/src/kosmickrisp/vulkan/kk_physical_device.c b/src/kosmickrisp/vulkan/kk_physical_device.c
index 91a32a2bbaa..7058a62874d 100644
--- a/src/kosmickrisp/vulkan/kk_physical_device.c
+++ b/src/kosmickrisp/vulkan/kk_physical_device.c
@@ -18,13 +18,17 @@
 
 #include "util/disk_cache.h"
 #include "util/mesa-blake3.h"
+#include "util/os_misc.h"
 #include "git_sha1.h"
 
 #include "vulkan/wsi/wsi_common.h"
 #include "vk_device.h"
 #include "vk_drm_syncobj.h"
+#include "vk_physical_device.h"
 #include "vk_shader_module.h"
 
+#define KK_HEAP_SIZE_PERCENT (0.75f)
+
 static uint32_t
 kk_get_vk_version()
 {
@@ -766,30 +770,6 @@ kk_physical_device_free_disk_cache(struct kk_physical_device *pdev)
 #endif
 }
 
-static uint64_t
-kk_get_sysmem_heap_size(void)
-{
-   uint64_t sysmem_size_B = 0;
-   if (!os_get_total_physical_memory(&sysmem_size_B))
-      return 0;
-
-   /* Use 3/4 of total size to avoid swapping */
-   return ROUND_DOWN_TO(sysmem_size_B * 3 / 4, 1 << 20);
-}
-
-static uint64_t
-kk_get_sysmem_heap_available(struct kk_physical_device *pdev)
-{
-   uint64_t sysmem_size_B = 0;
-   if (!os_get_available_system_memory(&sysmem_size_B)) {
-      vk_loge(VK_LOG_OBJS(pdev), "Failed to query available system memory");
-      return 0;
-   }
-
-   /* Use 3/4 of available to avoid swapping */
-   return ROUND_DOWN_TO(sysmem_size_B * 3 / 4, 1 << 20);
-}
-
 static void
 get_metal_limits(struct kk_physical_device *pdev)
 {
@@ -853,7 +833,7 @@ kk_enumerate_physical_devices(struct vk_instance *_instance)
 
    kk_physical_device_init_pipeline_cache(pdev);
 
-   uint64_t sysmem_size_B = kk_get_sysmem_heap_size();
+   uint64_t sysmem_size_B = os_get_gpu_heap_size(KK_HEAP_SIZE_PERCENT, NULL);
    if (sysmem_size_B == 0) {
       result = vk_errorf(instance, VK_ERROR_INITIALIZATION_FAILED,
                          "Failed to query total system memory");
@@ -864,7 +844,6 @@ kk_enumerate_physical_devices(struct vk_instance *_instance)
    pdev->mem_heaps[sysmem_heap_idx] = (struct kk_memory_heap){
       .size = sysmem_size_B,
       .flags = VK_MEMORY_HEAP_DEVICE_LOCAL_BIT,
-      .available = kk_get_sysmem_heap_available,
    };
 
    pdev->mem_types[pdev->mem_type_count++] = (VkMemoryType){
@@ -963,33 +942,12 @@ kk_GetPhysicalDeviceMemoryProperties2(
              */
             p->heapUsage[i] = used;
 
-            uint64_t available = heap->size;
-            if (heap->available)
-               available = heap->available(pdev);
-
-            /* From the Vulkan 1.3.278 spec:
-             *
-             *    "heapBudget is an array of VK_MAX_MEMORY_HEAPS VkDeviceSize
-             *    values in which memory budgets are returned, with one
-             *    element for each memory heap. A heap’s budget is a rough
-             *    estimate of how much memory the process can allocate from
-             *    that heap before allocations may fail or cause performance
-             *    degradation. The budget includes any currently allocated
-             *    device memory."
-             *
-             * and
-             *
-             *    "The heapBudget value must be less than or equal to
-             *    VkMemoryHeap::size for each heap."
-             *
-             * available (queried above) is the total amount free memory
-             * system-wide and does not include our allocations so we need
-             * to add that in.
+            /* Set the budget at 90% to avoid thrashing. Multiplying with
+             * KK_HEAP_SIZE_PERCENT to scale the budget the same way the heap
+             * was scaled.
              */
-            uint64_t budget = MIN2(available + used, heap->size);
-
-            /* Set the budget at 90% of available to avoid thrashing */
-            p->heapBudget[i] = ROUND_DOWN_TO(budget * 9 / 10, 1 << 20);
+            p->heapBudget[i] = vk_physical_device_heap_budget_from_system(
+               &pdev->vk, KK_HEAP_SIZE_PERCENT * 0.9f, heap->size, used);
          }
 
          /* From the Vulkan 1.3.278 spec:
diff --git a/src/kosmickrisp/vulkan/kk_physical_device.h b/src/kosmickrisp/vulkan/kk_physical_device.h
index 91c053a8310..1878a381ab5 100644
--- a/src/kosmickrisp/vulkan/kk_physical_device.h
+++ b/src/kosmickrisp/vulkan/kk_physical_device.h
@@ -32,7 +32,6 @@ struct kk_memory_heap {
    uint64_t size;
    uint64_t used;
    VkMemoryHeapFlags flags;
-   uint64_t (*available)(struct kk_physical_device *pdev);
 };
 
 struct kk_device_info {