ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: persist speculative decoding preference across restart and model switch (#6169)

Browse source 
* 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 <wasimysdev@gmail.com>
This commit is contained in:
oobabooga 2026-06-12 09:37:44 -03:00 • committed by GitHub
commit f4fc06b5bb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 144 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

16
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -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 };

25
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -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 {

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

@ -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";

10
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -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

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

@ -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<ChatRuntimeStore>((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<ChatRuntimeStore>((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<ChatRuntimeStore>((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,
};
}