Preserve requested GGUF GPU pools

Author: oobabooga

studio/backend/models/inference.py

@@ -485,7 +485,14 @@ class LoadResponse(BaseModel):
     )
     gpu_ids: Optional[List[int]] = Field(
         None,
-        description = "Physical GPU indices the model is pinned to, or None for automatic selection.",
+        description = "Effective GPU indices the model is using after fit-time narrowing, or None for automatic selection.",
+    )
+    requested_gpu_ids: Optional[List[int]] = Field(
+        None,
+        description = (
+            "GPU placement pool requested by the user before fit-time narrowing, "
+            "or None for automatic selection."
+        ),
     )
 
 
@@ -649,7 +656,14 @@ class InferenceStatusResponse(BaseModel):
     )
     gpu_ids: Optional[List[int]] = Field(
         None,
-        description = "Physical GPU indices the model is pinned to, or None for automatic selection.",
+        description = "Effective GPU indices the model is using after fit-time narrowing, or None for automatic selection.",
+    )
+    requested_gpu_ids: Optional[List[int]] = Field(
+        None,
+        description = (
+            "GPU placement pool requested by the user before fit-time narrowing, "
+            "or None for automatic selection."
+        ),
     )
     llama_cpp_supports_mtp: bool = Field(
         True,

studio/backend/routes/inference.py

@@ -4425,6 +4425,7 @@ async def _load_model_impl(
                     n_layers = llama_backend.n_layers,
                     n_moe_layers = llama_backend.n_moe_layers,
                     gpu_ids = llama_backend.gpu_ids,
+                    requested_gpu_ids = llama_backend.requested_gpu_ids,
                 )
         else:
             if (
@@ -4796,6 +4797,7 @@ async def _attempt_gguf_load(
                 n_layers = llama_backend.n_layers,
                 n_moe_layers = llama_backend.n_moe_layers,
                 gpu_ids = llama_backend.gpu_ids,
+                requested_gpu_ids = llama_backend.requested_gpu_ids,
             )
 
         # ── Standard path: load via Unsloth/transformers ──────────
@@ -5915,6 +5917,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
                 n_layers = llama_backend.n_layers,
                 n_moe_layers = llama_backend.n_moe_layers,
                 gpu_ids = llama_backend.gpu_ids,
+                requested_gpu_ids = llama_backend.requested_gpu_ids,
                 llama_cpp_supports_mtp = _supports_mtp,
                 spec_fallback_reason = llama_backend.spec_fallback_reason,
                 llama_cpp_prebuilt_stale = _stale,

studio/backend/tests/test_gpu_memory_mode.py

@@ -591,10 +591,23 @@ def test_load_request_accepts_gpu_ids():
 @pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse])
 def test_response_models_emit_gpu_ids(model_cls):
     if model_cls is LoadResponse:
-        obj = model_cls(status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1])
+        obj = model_cls(
+            status = "loaded",
+            model = "m",
+            display_name = "m",
+            inference = {},
+            gpu_ids = [1],
+            requested_gpu_ids = [1, 2],
+        )
     else:
-        obj = model_cls(gpu_ids = [1])
+        obj = model_cls(gpu_ids = [1], requested_gpu_ids = [1, 2])
     assert obj.model_dump()["gpu_ids"] == [1]
+    assert obj.model_dump()["requested_gpu_ids"] == [1, 2]
+
+
+def test_gguf_load_and_status_responses_include_requested_gpu_pool():
+    route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8")
+    assert route_src.count("requested_gpu_ids = llama_backend.requested_gpu_ids") == 3
 
 
 def test_gpu_ids_property_default_and_reset():

studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts

@@ -222,7 +222,9 @@ export function applyActiveModelStatusToStore(
     incomingGpuMode === "manual" ? (status.n_cpu_moe ?? null) : null;
   const incomingSplit =
     incomingGpuMode === "manual" ? (status.tensor_split ?? null) : null;
-  const incomingGpuIds = status.is_gguf ? (status.gpu_ids ?? null) : null;
+  const incomingGpuIds = status.is_gguf
+    ? (status.requested_gpu_ids ?? status.gpu_ids ?? null)
+    : null;
   const gpuStatusChanged =
     prevState.loadedGpuMemoryMode !== incomingGpuMode ||
     prevState.loadedGpuLayers !== incomingGpuLayers ||

studio/frontend/src/features/chat/stores/chat-runtime-store.ts

@@ -605,6 +605,7 @@ export function loadedGpuMemoryFields(resp: {
   n_layers?: number | null;
   n_moe_layers?: number;
   gpu_ids?: number[] | null;
+  requested_gpu_ids?: number[] | null;
 }) {
   // GPU-memory state is meaningful only for a GGUF chat load. A non-GGUF response
   // still carries gpu_memory_mode (its default "auto" is serialized), so gate on
@@ -631,7 +632,9 @@ export function loadedGpuMemoryFields(resp: {
     };
   }
   const mode = resp.gpu_memory_mode ?? "auto";
-  const gpuIds = resp.gpu_ids ?? null;
+  // Keep the user's placement pool editable across status/load hydration.
+  // gpu_ids remains the effective fitted subset for diagnostics.
+  const gpuIds = resp.requested_gpu_ids ?? resp.gpu_ids ?? null;
   // Layer/MoE/split knobs apply (and are reported) only in manual mode; in auto
   // the server ignores them, so don't seed the loaded baseline or the editable
   // knobs with values it never applied. In manual, the server reports gpu_layers
@@ -669,7 +672,7 @@ export function loadedGpuMemoryFields(resp: {
     ggufLayerCount: resp.n_layers ?? null,
     // MoE expert-layer count: the n_cpu_moe slider max, and 0 hides the slider.
     moeLayerCount: resp.n_moe_layers ?? null,
-    // The picker reflects what loaded (the request sent the user's pick).
+    // The picker reflects the requested placement pool, not a fitted subset.
     selectedGpuIds: gpuIds,
     loadedGpuIds: gpuIds,
     ...manualKnobs,

studio/frontend/src/features/chat/types/api.ts

@@ -188,7 +188,10 @@ export interface LoadModelResponse {
   n_layers?: number | null;
   /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */
   n_moe_layers?: number;
+  /** Effective GPU placement after fit-time narrowing. */
   gpu_ids?: number[] | null;
+  /** User-requested GPU placement pool before fit-time narrowing. */
+  requested_gpu_ids?: number[] | null;
 }
 
 export interface UnloadModelRequest {
@@ -240,7 +243,10 @@ export interface InferenceStatusResponse {
   /** n_ctx the active GGUF load was invoked with (0 = Auto); re-seeds a
    * Manual + Auto-layers context pin on hydration. Null for non-GGUF. */
   requested_context_length?: number | null;
+  /** Effective GPU placement after fit-time narrowing. */
   gpu_ids?: number[] | null;
+  /** User-requested GPU placement pool before fit-time narrowing. */
+  requested_gpu_ids?: number[] | null;
   n_layers?: number | null;
   /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */
   n_moe_layers?: number;

tests/studio/test_model_picker_contracts.py

@@ -151,6 +151,19 @@ def test_active_model_config_round_trips_gpu_fields():
     assert "export function gpuFieldsSignature" in shared
 
 
+def test_gpu_picker_round_trips_requested_pool_not_fitted_subset():
+    """A GGUF fit may narrow [0, 1] to [0], but load/status hydration must keep
+    [0, 1] as the editable pool so a later reload can grow back onto GPU 1."""
+    types = _read("features/chat/types/api.ts")
+    assert types.count("requested_gpu_ids?: number[] | null") >= 2
+
+    store = _read("features/chat/stores/chat-runtime-store.ts")
+    assert "resp.requested_gpu_ids ?? resp.gpu_ids ?? null" in store
+
+    status = _read("features/chat/lib/apply-inference-status-to-store.ts")
+    assert "status.requested_gpu_ids ?? status.gpu_ids ?? null" in status
+
+
 def test_compare_load_uses_each_models_gpu_config():
     src = _read("features/chat/shared-composer.tsx")
     assert "ownConfig.gpuMemoryMode ?? compareLoadKnobs.gpuMemoryMode" in src