*/
        WARN_ONCE(vram_needed && !adev,
                  "adev reference can't be null when vram is used");
-       if (adev)
+       if (adev) {
                adev->kfd.vram_used += vram_needed;
+               adev->kfd.vram_used_aligned += ALIGN(vram_needed, VRAM_AVAILABLITY_ALIGN);
+       }
        kfd_mem_limit.system_mem_used += system_mem_needed;
        kfd_mem_limit.ttm_mem_used += ttm_mem_needed;
 
        } else if (alloc_flag & KFD_IOC_ALLOC_MEM_FLAGS_VRAM) {
                WARN_ONCE(!adev,
                          "adev reference can't be null when alloc mem flags vram is set");
-               if (adev)
+               if (adev) {
                        adev->kfd.vram_used -= size;
+                       adev->kfd.vram_used_aligned -= ALIGN(size, VRAM_AVAILABLITY_ALIGN);
+               }
        } else if (alloc_flag & KFD_IOC_ALLOC_MEM_FLAGS_USERPTR) {
                kfd_mem_limit.system_mem_used -= size;
        } else if (!(alloc_flag &
        size_t available;
        spin_lock(&kfd_mem_limit.mem_limit_lock);
        available = adev->gmc.real_vram_size
-               - adev->kfd.vram_used
+               - adev->kfd.vram_used_aligned
                - atomic64_read(&adev->vram_pin_size)
                - reserved_for_pt;
        spin_unlock(&kfd_mem_limit.mem_limit_lock);