From 8ccb4615701a6ef70dec2f9157e4ce41bf7aa1cd Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 14 Mar 2026 09:01:30 +0000
Subject: [PATCH] studio: group GGUF shards by variant in size-based fallback

The smallest-fitting-variant fallback now groups split GGUF shards
by their variant prefix and sums all shard sizes per variant.

For example, DeepSeek-V3.2 UD-Q4_K_XL has 9 shards totaling
379.8 GB. The previous code treated each shard as a separate
"variant" and would have incorrectly selected a single 50 GB shard
as fitting, ignoring the other 8 shards needed.

Tested with unsloth/DeepSeek-V3.2-GGUF (237 GGUF files, 27
variants from 150 GB to 1.25 TB). Correctly groups and sorts
all variants by total size.
---
 studio/backend/core/inference/llama_cpp.py | 42 ++++++++++++++++------
 1 file changed, 31 insertions(+), 11 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index ebcac15aff..a6f67eb6e4 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -287,10 +287,15 @@ class LlamaCppBackend:
         free_bytes: int,
         hf_token: Optional[str] = None,
     ) -> Optional[tuple[str, int]]:
-        """Find the smallest single-file GGUF variant that fits in free_bytes.
+        """Find the smallest GGUF variant (including all shards) that fits.
 
-        Returns (filename, size_bytes) or None if nothing fits.
+        Groups split shards by variant prefix and sums their sizes.
+        For example, UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total.
+
+        Returns (first_shard_filename, total_size_bytes) or None if nothing fits.
         """
+        import re
+
         try:
             from huggingface_hub import get_paths_info, list_repo_files
 
@@ -300,16 +305,31 @@ class LlamaCppBackend:
                 return None
 
             # Get sizes for all GGUF files
-            path_infos = list(get_paths_info(hf_repo, gguf_files, token = hf_token))
-            sized = [(p.path, p.size) for p in path_infos if p.size and p.size > 0]
-            if not sized:
-                return None
+            path_infos = list(
+                get_paths_info(hf_repo, gguf_files, token = hf_token)
+            )
+            size_map = {p.path: (p.size or 0) for p in path_infos}
 
-            # Sort by size ascending and pick the smallest that fits
-            sized.sort(key = lambda x: x[1])
-            for filename, size in sized:
-                if size <= free_bytes:
-                    return filename, size
+            # Group files by variant: shards share a prefix before -NNNNN-of-NNNNN
+            shard_pat = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
+            variants: dict[str, list[str]] = {}
+            for f in gguf_files:
+                m = shard_pat.match(f)
+                key = m.group(1) if m else f
+                variants.setdefault(key, []).append(f)
+
+            # Sum shard sizes per variant, track the first shard (for download)
+            variant_sizes: list[tuple[str, int, list[str]]] = []
+            for key, shard_files in variants.items():
+                total = sum(size_map.get(f, 0) for f in shard_files)
+                first = sorted(shard_files)[0]
+                variant_sizes.append((first, total, shard_files))
+
+            # Sort by total size ascending and pick the smallest that fits
+            variant_sizes.sort(key = lambda x: x[1])
+            for first_file, total_size, _ in variant_sizes:
+                if total_size > 0 and total_size <= free_bytes:
+                    return first_file, total_size
 
             return None
         except Exception: