Upload cuda_backend.py

cuda_backend.py

@@ -935,6 +935,16 @@ def _apply_dc_nudge_cuda(
     return codes_out.reshape(n, QK_K)
 
 
+# VRAM-aware CUDA batching for the Q2_K factored pipeline.  The candidate
+# kernels materialize several live (B, K, 256) fp32 tensors (K = n_cand_d *
+# n_cand_m), so the per-block footprint scales with K (~6 fp32 temporaries of
+# 256 weights each).  The floor is the historical 4096-block batch (known to
+# fit even on small cards); on large GPUs the batch grows with the free VRAM.
+Q2_K_CUDA_FLOOR_BLOCKS = 4096
+Q2_K_CUDA_VRAM_FRACTION = 0.45
+Q2_K_CUDA_FP32_TEMPS_PER_CANDIDATE = 6
+
+
 def quantize_cuda(
     weights: np.ndarray,
     importance: np.ndarray | None,
@@ -981,7 +991,22 @@ def quantize_cuda(
     clamp_val = seed_clamp_value(x, cfg)
     dev = torch.device(device)
     gpu_name = cuda_device_name(device)
-    batch_blocks = 4096
+    batch_blocks = Q2_K_CUDA_FLOOR_BLOCKS
+    try:
+        # VRAM-aware default: scale the batch to free memory at start-up,
+        # never below the floor.  Footprint per block: K live (B, K, 256)
+        # fp32 temporaries drive the peak (see _factor_candidates_cuda).
+        k_eff = int(cfg.n_cand_d) * int(cfg.n_cand_m)
+        bytes_per_block = (
+            Q2_K_CUDA_FP32_TEMPS_PER_CANDIDATE * k_eff * QK_K * 4
+        )
+        free_bytes, _ = torch.cuda.mem_get_info(dev)
+        batch_blocks = max(
+            Q2_K_CUDA_FLOOR_BLOCKS,
+            int(Q2_K_CUDA_VRAM_FRACTION * free_bytes) // max(bytes_per_block, 1),
+        )
+    except Exception:
+        batch_blocks = Q2_K_CUDA_FLOOR_BLOCKS
 
     candidate_errors = np.empty((n_blocks, TOTAL_CANDIDATES), dtype=np.float32)
     scales_seed = np.empty((n_blocks, NS), dtype=np.float32)