455.23.04: Page allocation failure in kernel module at random points

I am attaching a patch here which I think is the right way to handle the BUG, if you need my signed-off-by, please reach out or just add it

--- nvidia-modeset/nvidia-modeset-linux.c.org	2020-11-23 20:46:12.817979880 +1100
+++ nvidia-modeset/nvidia-modeset-linux.c	2020-11-24 10:50:31.474395155 +1100
@@ -21,6 +21,7 @@
 #include <linux/file.h>
 #include <linux/list.h>
 #include <linux/rwsem.h>
+#include <linux/mm.h>
 
 #include "nvstatus.h"
 
@@ -169,33 +170,19 @@ static inline void nvkms_write_unlock_pm
  * are called while nvkms_lock is held.
  *************************************************************************/
 
-/* Don't use kmalloc for allocations larger than 128k */
-#define KMALLOC_LIMIT (128 * 1024)
-
+/*
+ * Let the system decide when to switch between kmalloc and vmalloc
+ */
 void* NVKMS_API_CALL nvkms_alloc(size_t size, NvBool zero)
 {
-    void *p;
-
-    if (size <= KMALLOC_LIMIT) {
-        p = kmalloc(size, GFP_KERNEL);
-    } else {
-        p = vmalloc(size);
-    }
-
-    if (zero && (p != NULL)) {
-        memset(p, 0, size);
-    }
-
-    return p;
+    if (zero)
+        return kvzalloc(size, GFP_KERNEL);
+    return kvmalloc(size, GFP_KERNEL);
 }
 
 void NVKMS_API_CALL nvkms_free(void *ptr, size_t size)
 {
-    if (size <= KMALLOC_LIMIT) {
-        kfree(ptr);
-    } else {
-        vfree(ptr);
-    }
+    return kvfree(ptr);
 }
 
 void* NVKMS_API_CALL nvkms_memset(void *ptr, NvU8 c, size_t size)