GPU index-space invalidation: apply only to real storage restores

Codex round on the fromPersisted marker: it was set on call sites that also
carry fresh, current-space picks, and the cold-cache revalidation raced the
load path.

- applyPerModelConfigToRuntime now takes { fromPersisted } (default false =
  live). Only a config actually restored from storage passes true.
- chat-page load-on-selection passes fromPersisted only for the
  rememberedConfigFor fallback, not an explicit selection.config (fresh Run
  settings / picker edit). Prevents a freshly selected Vulkan ordinal from being
  cleared and loaded on all GPUs.
- shared-composer compare pane gates fromPersisted on (!config && ownRemembered),
  so an explicit sel.config pane pick stays live.
- Cold-cache race: instead of leaking the raw persisted ids into the store
  (where the load path snapshots and sends them before the async revalidation
  runs), park the pick at null (all GPUs, safe) and fill in the validated pick
  once the cache warms. The load snapshot can no longer send a wrong-space pin.

Contract test updated + a regression guard for the storage-only invalidation.

studio/frontend/src/features/chat/chat-page.tsx

@@ -2323,6 +2323,9 @@ export function ChatPage({
       });
       const hasAppliedConfig = applyModelLoadConfigToRuntime(
         selection.config ?? rememberedConfigFor(selection),
+        // Only the rememberedConfigFor fallback is a storage restore; an
+        // explicit selection.config is a fresh pick in the current index space.
+        { fromPersisted: !selection.config },
       );
       await selectModel({
         ...selection,
@@ -2993,6 +2996,7 @@ export function ChatPage({
           });
           const hasAppliedConfig = applyModelLoadConfigToRuntime(
             rememberedConfigFor(selection),
+            { fromPersisted: true },
           );
           await selectModelRef.current({
             ...selection,

studio/frontend/src/features/chat/shared-composer.tsx

@@ -1058,7 +1058,10 @@ export function SharedComposer({
         const effectiveSelectedGpuIds =
           ownConfig.selectedGpuIds !== undefined
             ? reconcilePersistedGpuIds(ownConfig.selectedGpuIds, {
-                fromPersisted: true,
+                // Only a resolveInitialConfig storage hit is a real restore; an
+                // explicit sel.config is a fresh current-space pane pick, so it
+                // must stay live and never take the index-space invalidation.
+                fromPersisted: !config && ownRemembered,
               })
             : compareLoadKnobs.selectedGpuIds;
         // A pane's context comes from its own config only: a saved pin, or null

studio/frontend/src/features/hub/hub-page.tsx

@@ -1189,7 +1189,10 @@ export function ModelsPage() {
       const previousConfig = currentRuntimePerModelConfig({
         includeMaxSeqLength: true,
       });
-      const hasAppliedConfig = applyModelLoadConfigToRuntime(rememberedConfig);
+      const hasAppliedConfig = applyModelLoadConfigToRuntime(rememberedConfig, {
+        // rememberedConfig is the resolveInitialConfig storage hit (or null).
+        fromPersisted: true,
+      });
       void selectModel({
         id: runId,
         ggufVariant: opts.ggufVariant,

studio/frontend/src/features/model-picker/model-config/apply-per-model-config.ts

@@ -24,7 +24,18 @@ function cleanTemplate(value: string | null | undefined): string | null {
   return value?.trim() ? value : null;
 }
 
-export function applyPerModelConfigToRuntime(config: PerModelConfig): void {
+export function applyPerModelConfigToRuntime(
+  config: PerModelConfig,
+  // ``fromPersisted`` must be true only when ``config`` was actually restored
+  // from storage (a remembered per-model config). The GPU index-space
+  // invalidation (physical CUDA/ROCm ids vs ggml Vulkan ordinals, after a
+  // backend swap) applies only to such picks. A freshly edited config -- Run
+  // settings, a compare-pane selector, a cancel-restore of the in-session
+  // config -- is already in the current index space, so it stays live (false)
+  // and its pick is never cleared.
+  opts: { fromPersisted?: boolean } = {},
+): void {
+  const fromPersisted = opts.fromPersisted ?? false;
   // Fall back to the standing default when the model has no saved
   // maxSeqLength. maxSeqLength is the only per-model field carried on
   // params (the rest are reset below), so without this a model with no
@@ -36,6 +47,23 @@ export function applyPerModelConfigToRuntime(config: PerModelConfig): void {
   if (maxSeqLength !== store.params.maxSeqLength) {
     store.setParams({ ...store.params, maxSeqLength });
   }
+  // reconcilePersistedGpuIds can only decide the index-space question once the
+  // GPU cache is warm. This runs synchronously on model selection, which can
+  // happen before any GPU hook has fetched /api/system. For a persisted pick on
+  // a cold cache we cannot yet tell whether the saved ids belong to the current
+  // backend's index space, so parking the raw ids in the store would let the
+  // load path snapshot and send them as-is (pinning the wrong card after a
+  // swap). Park null instead (all GPUs -- safe) and fill in the validated pick
+  // once the cache warms. Live picks are current-space, so they pass through.
+  const coldCache = cachedPinnableGpuIndices() === null;
+  const parkColdPersistedPick =
+    fromPersisted && coldCache && config.selectedGpuIds != null;
+  const reconciledGpuIds =
+    config.selectedGpuIds === undefined
+      ? null
+      : parkColdPersistedPick
+        ? null
+        : reconcilePersistedGpuIds(config.selectedGpuIds, { fromPersisted });
   useChatRuntimeStore.setState({
     customContextLength: config.customContextLength ?? null,
     kvCacheDtype: config.kvCacheDtype ?? null,
@@ -54,33 +82,22 @@ export function applyPerModelConfigToRuntime(config: PerModelConfig): void {
     gpuLayers: config.gpuLayers ?? GPU_LAYERS_AUTO,
     nCpuMoe: config.nCpuMoe ?? 0,
     splitRatio: null,
-    selectedGpuIds:
-      config.selectedGpuIds !== undefined
-        ? reconcilePersistedGpuIds(config.selectedGpuIds, {
-            fromPersisted: true,
-          })
-        : null,
+    selectedGpuIds: reconciledGpuIds,
   });
 
-  // reconcilePersistedGpuIds can only clear a wrong-space pick (physical
-  // CUDA/ROCm ids vs ggml Vulkan ordinals, after a llama.cpp backend swap) once
-  // the GPU cache is warm. This runs synchronously on model selection, which can
-  // happen before any GPU hook has fetched /api/system -- on a cold cache the
-  // reconcile passes the saved ids through unvalidated, and if it then launders
-  // into the store the later load path reconciles it as a live pick
-  // (fromPersisted false) and never clears it, pinning the wrong card. So when
-  // the pick was applied cold, warm the cache and re-reconcile, overwriting only
-  // if the user has not replaced the pick in the meantime.
-  const persistedIds = config.selectedGpuIds;
-  if (persistedIds != null && cachedPinnableGpuIndices() === null) {
-    const applied = useChatRuntimeStore.getState().selectedGpuIds;
+  if (parkColdPersistedPick) {
+    // parkColdPersistedPick already established selectedGpuIds != null.
+    const persistedIds = config.selectedGpuIds as number[];
     void ensureGpuDeviceCache().then(() => {
-      const store = useChatRuntimeStore.getState();
-      if (store.selectedGpuIds !== applied) return; // user changed it since
+      // Only fill in the parked pick if the user has not chosen one meanwhile
+      // (reconciledGpuIds is null here, so the store must still read null).
+      if (useChatRuntimeStore.getState().selectedGpuIds !== reconciledGpuIds) {
+        return;
+      }
       const revalidated = reconcilePersistedGpuIds(persistedIds, {
         fromPersisted: true,
       });
-      if (revalidated !== store.selectedGpuIds) {
+      if (revalidated !== reconciledGpuIds) {
         useChatRuntimeStore.setState({ selectedGpuIds: revalidated });
       }
     });
@@ -89,9 +106,10 @@ export function applyPerModelConfigToRuntime(config: PerModelConfig): void {
 
 export function applyModelLoadConfigToRuntime(
   config: PerModelConfig | null | undefined,
+  opts: { fromPersisted?: boolean } = {},
 ): boolean {
   const hasConfig = config != null;
-  applyPerModelConfigToRuntime(config ?? DEFAULT_PER_MODEL_CONFIG);
+  applyPerModelConfigToRuntime(config ?? DEFAULT_PER_MODEL_CONFIG, opts);
   return hasConfig;
 }
 

tests/studio/test_model_picker_contracts.py

@@ -157,11 +157,11 @@ def test_compare_load_uses_each_models_gpu_config():
     assert "ownConfig.gpuLayers ?? compareLoadKnobs.gpuLayers" in src
     assert "ownConfig.nCpuMoe ?? compareLoadKnobs.nCpuMoe" in src
     assert "if (ownConfig.selectedGpuIds != null)" in src
-    # A per-model config pick is a persisted restore, so it must reconcile with
-    # fromPersisted: true (drops a pick saved under the other llama.cpp backend's
-    # GPU index space after a Vulkan/ROCm swap).
+    # A compare pane reconciles its own config's GPU pick, but the index-space
+    # invalidation (fromPersisted) must apply only to a real storage restore --
+    # never to an explicit sel.config, which is a fresh current-space pane pick.
     assert "reconcilePersistedGpuIds(ownConfig.selectedGpuIds, {" in src
-    assert "fromPersisted: true" in src
+    assert "fromPersisted: !config && ownRemembered" in src
     for field in (
         "gpu_memory_mode: effectiveGpuMemoryMode",
         "gpu_layers: effectiveGpuLayers",
@@ -171,6 +171,24 @@ def test_compare_load_uses_each_models_gpu_config():
         assert field in src
 
 
+def test_gpu_index_space_invalidation_only_for_storage_restores():
+    """The GPU index-space invalidation (physical CUDA/ROCm ids vs ggml Vulkan
+    ordinals after a backend swap) must apply only to picks actually restored
+    from storage. A freshly edited config -- an explicit selection.config, a
+    compare-pane sel.config -- is already current-space and must stay live, or
+    selecting a Vulkan ordinal and loading clears it and loads on all GPUs."""
+    # chat-page load-on-selection: fromPersisted only for the rememberedConfigFor
+    # fallback, never for an explicit selection.config.
+    chat = _read("features/chat/chat-page.tsx")
+    assert "fromPersisted: !selection.config" in chat
+
+    # apply-per-model-config: defaults to live, and on a cold cache it parks a
+    # persisted pick at null rather than leaking the raw ids to the load path.
+    apply = _read("features/model-picker/model-config/apply-per-model-config.ts")
+    assert "opts: { fromPersisted?: boolean } = {}" in apply
+    assert "const parkColdPersistedPick" in apply
+
+
 def test_active_native_gguf_metadata_uses_path_token():
     src = _read("features/model-picker/components/model-config-page.tsx")
     assert "(isActiveModel ? activeNativePathToken : null)" in src