nvidia: use vzalloc for oversized page table allocations

[guzebing1612-dev]

Summary

This PR avoids a failure in nvos_create_alloc() when registering very large
system-memory ranges.

The issue was found in a CXL-DRAM use case where a large CXL-DRAM backed host
memory region is registered with the NVIDIA driver so that it can be accessed by
the GPU. In the tested workload, vLLM + LMCache uses CXL-DRAM as host-side KV
cache storage, and a 512 GiB CXL-DRAM backed range needs to be registered with
cudaHostRegister().

nvos_create_alloc() currently allocates at->page_table as one large
kvzalloc() object. For very large registered host-memory ranges, the computed
page table metadata size can exceed INT_MAX. In that case, kvzalloc() fails
before falling back to vmalloc, and the registration fails.

This change uses vzalloc() directly only when pt_size > INT_MAX, while
preserving the existing kvzalloc() path for smaller allocations.

Background

The target deployment uses CXL-DRAM as an extension of host memory for GPU
workloads. The specific use case is to make a large CXL-DRAM backed host memory
range available to the GPU through CUDA host memory registration.

In this setup:

vLLM + LMCache
  -> KV cache is offloaded to host memory
  -> host memory is backed by CXL-DRAM
  -> cudaHostRegister() is used so the GPU can access the registered host memory

The failure was observed when trying to register a 512 GiB CXL-DRAM backed host
memory range for GPU use.

Problem

nvos_create_alloc() allocates at->page_table as a single metadata array:

NvU64 pt_size = num_pages * sizeof(nvidia_pte_t);

at->page_table = kvzalloc(pt_size, NV_GFP_KERNEL);

For very large registered system-memory ranges, pt_size can exceed the
kvmalloc() / kvzalloc() size guard.

For example, a 512 GiB registered range with 4 KiB base pages requires:

num_pages = 512 GiB / 4 KiB = 134217728
sizeof(nvidia_pte_t) = 16
pt_size = 134217728 * 16 = 2147483648 bytes

2147483648 bytes is exactly 2 GiB and is larger than INT_MAX. In this case,
kvzalloc() returns NULL before falling back to vmalloc, causing
nvos_create_alloc() to fail.

In CUDA applications this is visible as a large cudaHostRegister() failure,
because the CUDA user-space API eventually enters the NVIDIA kernel module path
which registers and pins the user pages.

Root cause

The failure is caused by a single oversized kernel allocation for CPU-side page
metadata:

registered system-memory range
  -> base-page count
  -> nvos_create_alloc()
  -> pt_size = num_pages * sizeof(nvidia_pte_t)
  -> kvzalloc(pt_size, NV_GFP_KERNEL)
  -> NULL when pt_size > INT_MAX

The page table metadata is CPU-side driver metadata. It does not require
physical contiguity. The existing cleanup path already uses kvfree(), which is
valid for both kvzalloc()-backed and vzalloc()-backed allocations.

Testing scope

This has been tested with the NVIDIA open kernel module flavor only.

Reproduction and validation environment:

GPU: NVIDIA H100 PCIe
OS: Ubuntu 24.04.2 LTS
Kernel: Linux 6.17.0-35-generic
Driver version: 590.48.01
Kernel module flavor: open
modinfo -F license nvidia: Dual MIT/GPL
CUDA version reported by nvidia-smi: 13.1
Use case: CXL-DRAM backed host memory registered for GPU access

I do not have access to the proprietary kernel module package of the same
version, so I have not verified whether the same issue reproduces with the
proprietary kernel module flavor.

Runtime evidence before this patch

For the 512 GiB failure case, bpftrace showed that the NVIDIA page table
metadata size reaches 2048 MiB:

num_pages=134217728
metadata_bytes=2147483648
metadata_mib=2048
nvos_create_alloc() returns NULL

A small kernel-module probe also reproduced the allocator boundary:

kvzalloc(2047 MiB) succeeds
kvzalloc(2048 MiB) fails

This matches the observed 512 GiB failure.

Validation after this patch

The driver was rebuilt with this change and the patched open kernel module was
loaded on the same test system.

After the patch:

512 GiB CXL-DRAM backed host memory registration:
  cudaHostRegister() succeeds

This verifies that the 512 GiB registration no longer fails at the
nvos_create_alloc() page table allocation path.

[CLAassistant]

All committers have signed the CLA.