From f4fc06b5bb4f6174e8f39599ba1fb5d8b054e48d Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 12 Jun 2026 09:37:44 -0300 Subject: [PATCH] Studio: persist speculative decoding preference across restart and model switch (#6169) * Studio: persist speculative decoding preference across restart and model switch * Studio: persist speculative preference on apply, not on edit * Fix/adjust speculative decoding persistence for PR #6169 * Fix speculative ngram alias for PR #6169 * Fix compare speculative preference for PR #6169 --------- Co-authored-by: wasimysaid --- .../src/features/chat/api/chat-adapter.ts | 16 +++ .../chat/hooks/use-chat-model-runtime.ts | 25 +++-- .../lib/apply-inference-status-to-store.ts | 9 +- .../src/features/chat/shared-composer.tsx | 10 ++ .../chat/stores/chat-runtime-store.ts | 97 ++++++++++++++++++- 5 files changed, 144 insertions(+), 13 deletions(-) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 5903bef107..824c8dcbd8 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -36,7 +36,10 @@ import { import { type PendingImageEditReference, type RagAutoInject, + resolveLoadedSpeculativeSettings, + resolveSpeculativeSettingsForLoad, resolveToolsEnabledOnLoad, + saveSpeculativeType, useChatRuntimeStore, } from "../stores/chat-runtime-store"; import { useExternalProvidersStore } from "../stores/external-providers-store"; @@ -1137,6 +1140,7 @@ async function autoLoadSmallestModel(): Promise<{ const store = useChatRuntimeStore.getState(); const hfToken = store.hfToken || null; const trustRemoteCode = store.params.trustRemoteCode ?? false; + const specSettings = resolveSpeculativeSettingsForLoad(); const toastId = toast("Loading a model…", { description: "Auto-selecting the smallest downloaded model.", duration: 5000, @@ -1201,7 +1205,10 @@ async function autoLoadSmallestModel(): Promise<{ is_lora: false, gguf_variant: variant.quant, trust_remote_code: trustRemoteCode, + speculative_type: specSettings.speculativeType, + spec_draft_n_max: specSettings.specDraftNMax, }); + saveSpeculativeType(specSettings.speculativeType); useChatRuntimeStore .getState() .setCheckpoint(repo.repo_id, variant.quant); @@ -1250,6 +1257,7 @@ async function autoLoadSmallestModel(): Promise<{ chatTemplateOverride: null, loadedChatTemplateOverride: null, loadedIsMultimodal: isMultimodalResponse(loadResp), + ...resolveLoadedSpeculativeSettings(loadResp), }); toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, { id: toastId, @@ -1290,7 +1298,10 @@ async function autoLoadSmallestModel(): Promise<{ is_lora: false, gguf_variant: null, trust_remote_code: trustRemoteCode, + speculative_type: specSettings.speculativeType, + spec_draft_n_max: specSettings.specDraftNMax, }); + saveSpeculativeType(specSettings.speculativeType); useChatRuntimeStore.getState().setCheckpoint(repo.repo_id); const store = useChatRuntimeStore.getState(); store.setModelRequiresTrustRemoteCode( @@ -1310,6 +1321,7 @@ async function autoLoadSmallestModel(): Promise<{ defaultChatTemplate: sfLoadResp.chat_template ?? null, chatTemplateOverride: null, loadedChatTemplateOverride: null, + ...resolveLoadedSpeculativeSettings(sfLoadResp), }); const sfModel: ChatModelSummary = { id: repo.repo_id, @@ -1372,7 +1384,10 @@ async function autoLoadSmallestModel(): Promise<{ is_lora: false, gguf_variant: "UD-Q4_K_XL", trust_remote_code: trustRemoteCode, + speculative_type: specSettings.speculativeType, + spec_draft_n_max: specSettings.specDraftNMax, }); + saveSpeculativeType(specSettings.speculativeType); useChatRuntimeStore .getState() .setCheckpoint("unsloth/Qwen3.5-4B-MTP-GGUF", "UD-Q4_K_XL"); @@ -1412,6 +1427,7 @@ async function autoLoadSmallestModel(): Promise<{ defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, loadedIsMultimodal: isMultimodalResponse(loadResp), + ...resolveLoadedSpeculativeSettings(loadResp), }); toast.success("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)", { id: toastId }); return { loaded: true, blockedByTrustRemoteCode: false }; diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 70cba21afc..e513d351c4 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -24,7 +24,9 @@ import { } from "../api/chat-api"; import { formatEta, formatRate } from "../utils/format-transfer"; import { + readPersistedSpeculativeType, resolveToolsEnabledOnLoad, + saveSpeculativeType, useChatRuntimeStore, } from "../stores/chat-runtime-store"; import { @@ -493,16 +495,16 @@ export function useChatModelRuntime() { } if (abortCtrl.signal.aborted) throw new Error("Cancelled"); - // Reset Speculative Decoding to Auto on model switch: spec - // strategy is per-model, so a sub-3B non-MTP GGUF's "Off" must - // not carry into a 27B MTP GGUF where Auto auto-promotes to - // draft-mtp. Clears the stale prior choice so the backend's - // platform-aware path runs by default; same for spec_draft_n_max - // (MTP-only). The user can still force a mode on the new model. + // On a model switch, fall back to the persisted standing + // preference rather than null so a per-session forced MTP mode + // can't follow the user onto a model without an MTP head. + // spec_draft_n_max is MTP-only and always resets. The loaded + // shadow is seeded too, preventing a transient dirty Apply state. if (currentCheckpoint && currentCheckpoint !== modelId) { + const persistedSpeculativeType = readPersistedSpeculativeType(); useChatRuntimeStore.setState({ - speculativeType: null, - loadedSpeculativeType: null, + speculativeType: persistedSpeculativeType, + loadedSpeculativeType: persistedSpeculativeType, specDraftNMax: null, loadedSpecDraftNMax: null, }); @@ -551,6 +553,11 @@ export function useChatModelRuntime() { // the model as active -- it's being unloaded. if (abortCtrl.signal.aborted) throw new Error("Cancelled"); + // The load applied this spec mode, so persist the user's standing + // preference now (the requested intent, not the resolved echo; + // saveSpeculativeType keeps only the universal auto/ngram/off). + saveSpeculativeType(speculativeType); + const currentParams = useChatRuntimeStore.getState().params; setParams( mergeBackendRecommendedInference({ @@ -711,6 +718,8 @@ export function useChatModelRuntime() { }); useChatRuntimeStore.setState({ activeNativePathToken: previousActiveNativePathToken ?? null, + loadedSpeculativeType: null, + loadedSpecDraftNMax: null, }); await refresh(); } catch { diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index 332bde6a3c..376329b0cc 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -24,13 +24,16 @@ export function normalizeSpeculativeType( if (!s) return null; if (s === "auto" || s === "default") return "auto"; if (s === "off") return "off"; - if (s === "ngram-simple") return "ngram-simple"; if (s === "mtp" || s === "draft-mtp") return "mtp"; - if (s === "ngram" || s === "ngram-mod") return "ngram"; + if (s === "ngram" || s === "ngram-mod" || s === "ngram-simple") { + return "ngram"; + } if (s === "mtp+ngram") return "mtp+ngram"; const parts = s.split(",").map((p) => p.trim()).filter(Boolean); const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp"); - const hasNgram = parts.some((p) => p === "ngram" || p === "ngram-mod"); + const hasNgram = parts.some( + (p) => p === "ngram" || p === "ngram-mod" || p === "ngram-simple", + ); if (hasMtp && hasNgram) return "mtp+ngram"; if (hasMtp) return "mtp"; if (hasNgram) return "ngram"; diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 9f190355de..a8fa2360a0 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -75,6 +75,9 @@ import { } from "./stores/plus-menu-prefs-store"; import { type ReasoningEffort, + resolveLoadedSpeculativeSettings, + resolveSpeculativeSettingsForLoad, + saveSpeculativeType, useChatRuntimeStore, } from "./stores/chat-runtime-store"; import { @@ -910,6 +913,9 @@ export function SharedComposer({ const effectiveChatTemplateOverride = chatTemplateOverride?.trim() ? chatTemplateOverride : null; + const specSettings = resolveSpeculativeSettingsForLoad({ + usePersistedPreference: true, + }); function modelDisplayName(id: string): string { const parts = id.split("/"); @@ -951,9 +957,12 @@ export function SharedComposer({ gguf_variant: sel.ggufVariant ?? null, trust_remote_code: trustRemoteCode, chat_template_override: effectiveChatTemplateOverride, + speculative_type: specSettings.speculativeType, + spec_draft_n_max: specSettings.specDraftNMax, // Honor the Tensor Parallelism toggle on compare loads too. tensor_parallel: currentStore.tensorParallel, }); + saveSpeculativeType(specSettings.speculativeType); const store = useChatRuntimeStore.getState(); store.setCheckpoint( resp.model, @@ -971,6 +980,7 @@ export function SharedComposer({ tensorParallel: resp.tensor_parallel ?? false, loadedTensorParallel: resp.tensor_parallel ?? false, loadedIsMultimodal: isMultimodalResponse(resp), + ...resolveLoadedSpeculativeSettings(resp), }); // Sync the models[] entry with the load response so attach/send gates // read fresh capabilities. /api/models/list can lag a model's actual diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 2d23777378..7aecc1fcd9 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -43,6 +43,12 @@ export const CHAT_RAG_TOP_K_KEY = "unsloth_chat_rag_top_k"; export const CHAT_RAG_AUTOINJECT_KEY = "unsloth_chat_rag_autoinject"; export const CHAT_RAG_AUTOINJECT_MIN_SCORE_KEY = "unsloth_chat_rag_autoinject_min_score"; +export const CHAT_SPECULATIVE_TYPE_KEY = "unsloth_chat_speculative_type"; + +// Persist only the model-agnostic intents (auto/ngram/off). MTP modes +// (mtp/mtp+ngram) and spec_draft_n_max stay session-only: a persisted MTP +// choice would silently no-op on models without an MTP head. Unknown -> auto. +const PERSISTED_SPEC_MODES = new Set(["auto", "ngram", "off"]); export type RagSource = | { type: "thread" } @@ -319,6 +325,71 @@ function saveString(key: string, value: string): void { } } +// Canonicalises any backend value onto the Speculative Decoding dropdown's +// modes ("auto"/"mtp"/"ngram"/"mtp+ngram"/"off"/null). Backend-only +// legacy aliases map to their closest UI mode. +export function normalizeSpeculativeType( + v: string | null | undefined, +): string | null { + if (v == null) return null; + const s = String(v).trim().toLowerCase(); + if (!s) return null; + if (s === "auto" || s === "default") return "auto"; + if (s === "off") return "off"; + if (s === "mtp" || s === "draft-mtp") return "mtp"; + if (s === "ngram" || s === "ngram-mod" || s === "ngram-simple") { + return "ngram"; + } + if (s === "mtp+ngram") return "mtp+ngram"; + // Comma-chained legacy values (e.g. from older backend echoes). + const parts = s.split(",").map((p) => p.trim()).filter(Boolean); + const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp"); + const hasNgram = parts.some( + (p) => p === "ngram" || p === "ngram-mod" || p === "ngram-simple", + ); + if (hasMtp && hasNgram) return "mtp+ngram"; + if (hasMtp) return "mtp"; + if (hasNgram) return "ngram"; + // Unknown -> safe fallback to Auto so the dropdown stays controlled. + return "auto"; +} + +export function resolveLoadedSpeculativeSettings(response: { + speculative_type?: string | null; + spec_draft_n_max?: number | null; +}): { + speculativeType: string | null; + loadedSpeculativeType: string | null; + specDraftNMax: number | null; + loadedSpecDraftNMax: number | null; +} { + const loadedSpeculativeType = normalizeSpeculativeType( + response.speculative_type, + ); + const loadedSpecDraftNMax = response.spec_draft_n_max ?? null; + return { + speculativeType: loadedSpeculativeType, + loadedSpeculativeType, + specDraftNMax: loadedSpecDraftNMax, + loadedSpecDraftNMax, + }; +} + +// The user's standing preference, sanitized to the universal set. +export function readPersistedSpeculativeType(): string { + const raw = loadString(CHAT_SPECULATIVE_TYPE_KEY, "auto"); + return PERSISTED_SPEC_MODES.has(raw) ? raw : "auto"; +} + +// MTP / null / unknown values are left unwritten so they stay session-only. +// Called from the load path so only an applied preference is persisted, not an +// unapplied dropdown edit the user might Reset or abandon before Apply. +export function saveSpeculativeType(value: string | null): void { + if (value && PERSISTED_SPEC_MODES.has(value)) { + saveString(CHAT_SPECULATIVE_TYPE_KEY, value); + } +} + function notifyHfTokenChanged(value: string): void { if (!canUseStorage()) return; try { @@ -809,7 +880,7 @@ export const useChatRuntimeStore = create((set, get) => ({ toolCallTimeout: 5, kvCacheDtype: null, loadedKvCacheDtype: null, - speculativeType: "auto", + speculativeType: readPersistedSpeculativeType(), loadedSpeculativeType: null, specFallbackReason: null, specDraftNMax: null, @@ -1024,7 +1095,7 @@ export const useChatRuntimeStore = create((set, get) => ({ toolStatus: null, kvCacheDtype: null, loadedKvCacheDtype: null, - speculativeType: "auto", + speculativeType: readPersistedSpeculativeType(), loadedSpeculativeType: null, specFallbackReason: null, specDraftNMax: null, @@ -1233,3 +1304,25 @@ export const useChatRuntimeStore = create((set, get) => ({ set({ pendingImageEditReference: null }), setContextUsage: (contextUsage) => set({ contextUsage }), })); + +export function resolveSpeculativeSettingsForLoad({ + usePersistedPreference = false, +}: { + usePersistedPreference?: boolean; +} = {}): { + speculativeType: string | null; + specDraftNMax: number | null; +} { + const state = useChatRuntimeStore.getState(); + const speculativeType = usePersistedPreference + ? readPersistedSpeculativeType() + : state.speculativeType ?? readPersistedSpeculativeType(); + return { + speculativeType, + specDraftNMax: + !usePersistedPreference && + (speculativeType === "mtp" || speculativeType === "mtp+ngram") + ? state.specDraftNMax + : null, + }; +}