Sync managed_alloc.cu: selective prefetch (<2 GiB to GPU)

Previously the Dockerfile was copying the old version that had
cudaMemPrefetchAsync completely removed, which caused cublasGemmEx
crashes because model weights stayed in LPDDR and cuBLAS couldn't
access them. This version prefetches allocations <2 GiB (model
weights) to GPU but skips large allocations (KV cache).

vllm/managed_alloc.cu

@@ -7,9 +7,11 @@
 //   2. cudaMemAdviseSetPreferredLocation(GPU) → driver prefers keeping pages on GPU
 //   3. cudaMemAdviseSetAccessedBy(CPU) → CPU can access over C2C NVLink without
 //      triggering page migration back to system RAM (critical: prevents OOM)
-//   4. cudaMemPrefetchAsync(GPU) → actively migrates pages to GPU immediately,
-//      so subsequent writes go to HBM/EGM, not system RAM (prevents OOM on
-//      systems where EGM carved out most of system memory)
+//   4. Selective prefetching — small allocations (model weights, <2 GiB)
+//      are prefetched to GPU so cuBLAS/cuDNN kernels can access them
+//      directly from HBM. Large allocations (KV cache blocks) stay in
+//      managed memory and page-fault on demand, since they're too large
+//      to fit in HBM and attention ops can tolerate page faults.
 #include <cuda_runtime.h>
 #include <stdio.h>
 
@@ -56,16 +58,16 @@ void* managed_malloc(size_t size, int device, cudaStream_t stream) {
   cpu_loc.id = cudaCpuDeviceId;
   cudaMemAdvise(ptr, size, cudaMemAdviseSetAccessedBy, cpu_loc);
 
-  // Prefetch to GPU immediately. This actively migrates the virtual
-  // pages to the GPU side so that subsequent writes (e.g., model weight
-  // loading) go directly to HBM/EGM instead of pinning system RAM.
-  // Without this, the first write to each page faults into system RAM,
-  // which causes OOM when the OS only has ~102 GiB after EGM carveout.
-  //
-  // CUDA 13+ signature: cudaMemPrefetchAsync(ptr, size, cudaMemLocation, flags)
-  // Note: no stream parameter in the cudaMemLocation overload.
-  // flags=0 means no special flags.
-  if (size > 0) {
+  // Selective prefetch: migrate pages to GPU for small allocations only.
+  // Model weights (individual tensors) are typically <2 GiB and MUST be
+  // on GPU for cuBLAS GEMM operations — GPU compute kernels cannot
+  // page-fault into managed memory during execution.
+  // KV cache blocks are large and numerous; prefetching them all fills
+  // HBM and causes subsequent allocations to fail.
+  // The 2 GiB threshold separates "compute data" from "cache data".
+  const size_t PREFETCH_THRESHOLD = 2ULL * 1024 * 1024 * 1024; // 2 GiB
+
+  if (size > 0 && size < PREFETCH_THRESHOLD) {
     err = cudaMemPrefetchAsync(ptr, size, gpu_loc, 0);
     if (err != cudaSuccess) {
       // Non-fatal: prefetch failure shouldn't prevent allocation.