diff --git a/ramalama/amdkfd.py b/ramalama/amdkfd.py
index c69919256..24eca4712 100644
--- a/ramalama/amdkfd.py
+++ b/ramalama/amdkfd.py
@@ -2,6 +2,16 @@
 
 import glob
 
+# Heap types in memory properties
+#
+# Imported from /usr/include/linux/kfd_sysfs.h
+HEAP_TYPE_SYSTEM = 0
+HEAP_TYPE_FB_PUBLIC = 1
+HEAP_TYPE_FB_PRIVATE = 2
+HEAP_TYPE_GPU_GDS = 3
+HEAP_TYPE_GPU_LDS = 4
+HEAP_TYPE_GPU_SCRATCH = 5
+
 
 def parse_props(path):
     """Returns a dict corresponding to a KFD properties file"""
diff --git a/ramalama/common.py b/ramalama/common.py
index c5a2eb205..eb693270e 100644
--- a/ramalama/common.py
+++ b/ramalama/common.py
@@ -462,7 +462,7 @@ def check_rocm_amd():
             # See /usr/include/linux/kfd_sysfs.h for possible heap types
             #
             # Count public and private framebuffer memory as VRAM
-            if bank_props['heap_type'] in [1, 2]:
+            if bank_props['heap_type'] in [amdkfd.HEAP_TYPE_FB_PUBLIC, amdkfd.HEAP_TYPE_FB_PRIVATE]:
                 mem_bytes += int(bank_props['size_in_bytes'])
 
         if mem_bytes > MIN_VRAM_BYTES and mem_bytes > gpu_bytes: