feat: speculative decoding with paired drafter checkpoints (#5241)

Author: qinxuye

doc/source/user_guide/backends.rst

@@ -7,6 +7,8 @@ Backends
 Xinference supports multiple backends for different models. After the user specifies the model,
 xinference will automatically select the appropriate backend.
 
+.. _llama_cpp_backend:
+
 llama.cpp
 =========
 
@@ -196,3 +198,120 @@ MLX
 `MLX <https://github.com/ml-explore/mlx-examples/tree/main/llms>`_ provides efficient runtime
 to run LLM on Apple silicon. It's recommended to use for Mac users when running on Apple silicon
 if the model has MLX format support.
+
+
+.. _speculative_decoding:
+
+Speculative decoding
+====================
+Some models ship a small paired drafter checkpoint that predicts several tokens
+ahead, which the target model then verifies in one pass. Output is unchanged,
+decoding gets faster. Gemma 4 calls this multi-token prediction (MTP) and
+publishes a ``*-it-assistant`` drafter for every variant.
+
+Pass ``--enable_mtp true`` at launch to download the drafter declared by the
+model spec and run it alongside the target model:
+
+.. code-block:: bash
+
+   xinference launch --model-name gemma-4 --model-engine vllm \
+     --model-format pytorch --size-in-billions 12 --quantization none \
+     --enable_mtp true
+
+In the Web UI the same options live under *Advanced Configuration → Speculative
+Decoding*, which only appears for a format/size that actually ships a drafter.
+
+Optional parameters:
+
+* ``--num_speculative_tokens <n>``: how many tokens the drafter proposes per
+  round, including the bonus token. Left unset, MLX reads it from the drafter,
+  which runs at the depth it was trained for (``4`` for Gemma 4). Gemma 4 on
+  llama.cpp, vLLM, and SGLang follows the same model-size recipe: ``2`` for
+  E2B, ``4`` for E4B and 26B-A4B, and the lower end (``4``) of the recommended
+  ``4-8`` range for 12B and 31B. Other llama.cpp models keep xllamacpp's own
+  default.
+* ``--draft_quantization <quantization>``: which drafter conversion to use, when
+  the spec declares more than one — the MLX build of Gemma 4 12B publishes
+  eight. Defaults to the first declared, which is the least quantized one: a
+  drafter is small and quantizing it costs acceptance rate, so pairing a
+  quantized target with an unquantized drafter is the recommended setup.
+* ``--draft_model_path <path>``: use a local drafter instead of the one declared
+  by the spec. Implies ``--enable_mtp true``.
+
+The drafter must match the target model family and size — it shares the target's
+KV cache, so a mismatched checkpoint is rejected rather than silently degrading.
+
+Engine support:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 18 32 50
+
+   * - Engine
+     - Requirement
+     - Notes
+   * - :ref:`vLLM <vllm_backend>`
+     - ``vllm>=0.22.0``, ``transformers>=5.8.0``
+     - Translated into ``speculative_config`` with ``method: mtp``. An
+       explicitly provided ``speculative_config`` is left untouched. Older
+       vLLM treats the drafter as a generic draft model, while older
+       Transformers does not recognize ``gemma4_assistant``; either case is
+       rejected before engine initialization. Virtual-environment launches
+       also synchronize ``flashinfer-cubin`` with ``flashinfer-python`` before
+       starting vLLM, repairing stale environments that contain mismatched
+       FlashInfer packages.
+   * - :ref:`SGLang <sglang_backend>`
+     - ``sglang==0.5.13.post1``, ``transformers==5.8.1``
+     - Translated into ``--speculative-algorithm NEXTN`` with the matching
+       ``speculative_num_steps`` / ``speculative_num_draft_tokens``. An
+       explicitly provided ``speculative_algorithm`` is left untouched.
+   * - :ref:`MLX <mlx_backend>`
+     - ``mlx-vlm>=0.5.0`` (``>=0.6.1`` for Gemma 4 12B)
+     - Served by the MLX vision engine, the one that runs multimodal models
+       such as Gemma 4. The drafter is validated against the target when the
+       model loads.
+   * - :ref:`llama.cpp <llama_cpp_backend>`
+     - ``xllamacpp>=2026.6.9713``
+     - Translated into the ``draft-mtp`` speculative implementation. The
+       drafter is a single gguf published inside the target's own repository,
+       so its quantizations are its own (``BF16``, ``F16``, ``Q8_0`` for Gemma
+       4) and independent of the target's. Earlier llama.cpp builds do not
+       know the ``gemma4-assistant`` architecture and cannot load it.
+
+Not supported by the Transformers engine: it runs its own continuous-batching
+loop rather than ``generate()``, so there is nowhere to attach a drafter.
+
+When it pays off
+----------------
+Speculation is only worth it when a drafting step is cheap *relative to* a
+target decoding step. The drafter is a small dense model whose cost does not
+change with the target, so the ratio is what decides the outcome — and a
+mixture-of-experts target can land on the wrong side of it, because only its
+activated slice is read per token:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 34 22 22 22
+
+   * - gemma-4 MLX, 4bit, M5 Pro
+     - Without a drafter
+     - With MTP
+     - Accepted per round
+   * - 31B (dense, 18.4 GB read per token)
+     - 14.9 tok/s
+     - **31.0 tok/s** (2.1x)
+     - 2.08 of 4
+   * - 26B-A4B (MoE, ~2.2 GB read per token)
+     - 73.2 tok/s
+     - 65.9 tok/s (0.9x)
+     - 1.40 of 4
+
+The 0.83 GB drafter costs about 5% of a 31B decoding step but 39% of a
+26B-A4B one, so on the MoE the three drafting steps of a round already exceed
+one plain decoding step — and the round has to win that back from a lower
+acceptance rate. The MoE is still the faster model here in absolute terms; it
+simply has no headroom left for speculation.
+
+So measure before leaving it on, and note that the verification step itself is
+not the problem: a four-token forward costs only ~40% more than a single-token
+one on either model.

frontend/.env.development

@@ -0,0 +1 @@
+NEXT_PUBLIC_API_URL=http://127.0.0.1:9997

frontend/.gitignore

@@ -34,6 +34,10 @@ yarn-error.log*
 
 # env files (can opt-in for committing if needed)
 .env*
+# ...except the dev defaults: without NEXT_PUBLIC_API_URL the browser
+# talks to the Next dev server, whose rewrite proxy buffers SSE and
+# silently kills streaming responses.
+!.env.development
 
 # vercel
 .vercel

frontend/README.md

@@ -22,15 +22,21 @@ npm run dev
 
 Open http://127.0.0.1:3999.
 
-In dev mode, the frontend proxies API requests to `http://127.0.0.1:9997`. To
-use a different backend endpoint:
+`.env.development` points the browser straight at `http://127.0.0.1:9997`.
+Keep it that way unless you know you want the proxy: with `NEXT_PUBLIC_API_URL`
+unset the browser talks to the Next dev server instead, and its rewrite proxy
+buffers responses — streaming endpoints (chat completions and friends) then
+deliver the whole answer in one go, which looks like the backend lost its
+streaming support.
+
+To use a different backend endpoint:
 
 ```bash
-XINFERENCE_API_URL=http://127.0.0.1:6735 npm run dev
+NEXT_PUBLIC_API_URL=http://127.0.0.1:6735 npm run dev
 ```
 
-For direct browser requests to a non-default backend, use
-`NEXT_PUBLIC_API_URL` instead.
+`XINFERENCE_API_URL` sets the proxy target instead, for the same-origin setup
+where the browser goes through the dev server.
 
 ## Build
 

frontend/src/components/pages/launch-model/launch-dialog/advanced-config.tsx

@@ -5,15 +5,22 @@ import { Settings } from 'lucide-react';
 
 import { CollapsiblePanel } from '@/components/ui/collapsible';
 import { FormField } from '@/components/ui/form-field';
+import { Input } from '@/components/ui/input';
 import { RadioGroup } from '@/components/ui/radio-group';
+import { Select } from '@/components/ui/select';
+import { Switch } from '@/components/ui/switch';
+import { InfoTooltip } from '@/components/ui/tooltip';
 import {
+  GEMMA_4_SPECULATIVE_TOKENS_BY_SIZE,
   KWARGS_OPTIONS_FOR_ENGINES,
   QUANTIZATION_OPTIONS,
+  SPECULATIVE_TOKENS_DEFAULT_BY_ENGINE,
   VIRTUAL_ENV_OPTIONS,
 } from '@/constants/launch';
 import { useI18n } from '@/contexts/i18n-context';
 import { useWatch } from '@/hooks/use-form';
 import type { FormInstance } from '@/types/form';
+import type { Option } from '@/types/common';
 import type { RequestModelType } from '../types';
 import { toOptionValue } from '../utils';
 import CommonFormList from './common-form-list';
@@ -22,13 +29,27 @@ import { ModelType } from '@/constants';
 interface AdvancedConfigProps {
   form: FormInstance;
   modelType: RequestModelType;
+  modelName?: string;
+  /** whether the selected spec ships a drafter for speculative decoding */
+  hasDrafter?: boolean;
+  /** available drafter conversions, empty when there is nothing to pick */
+  draftQuantizationOptions?: Option<string>[];
 }
 
 interface ConfigSectionProps {
   title: ReactNode;
   children: ReactNode;
 }
 
+function FieldLabel({ label, tip }: { label: string; tip: string }) {
+  return (
+    <span className="shrink-0 flex items-center gap-1 text-sm font-medium">
+      {label}
+      <InfoTooltip content={tip} />
+    </span>
+  );
+}
+
 function ConfigSection({ title, children }: ConfigSectionProps) {
   return (
     <section className="space-y-2">
@@ -40,13 +61,24 @@ function ConfigSection({ title, children }: ConfigSectionProps) {
   );
 }
 
-const AdvancedConfig: FC<AdvancedConfigProps> = ({ form, modelType }) => {
+const AdvancedConfig: FC<AdvancedConfigProps> = ({
+  form,
+  modelType,
+  modelName,
+  hasDrafter = false,
+  draftQuantizationOptions = [],
+}) => {
   const { t } = useI18n();
   const modelEngineValue = toOptionValue(useWatch('model_engine', form));
+  const modelSizeValue = toOptionValue(useWatch('model_size_in_billions', form));
+  const enableMtpValue = useWatch('enable_mtp', form);
 
-  const kwargsOptionsForEngine = modelEngineValue
-    ? KWARGS_OPTIONS_FOR_ENGINES[modelEngineValue.toLowerCase()]
-    : undefined;
+  const engineKey = modelEngineValue.toLowerCase();
+  const kwargsOptionsForEngine = engineKey ? KWARGS_OPTIONS_FOR_ENGINES[engineKey] : undefined;
+  const speculativeTokensDefault =
+    ['vllm', 'sglang', 'llama.cpp'].includes(engineKey) && modelName === 'gemma-4'
+      ? GEMMA_4_SPECULATIVE_TOKENS_BY_SIZE[modelSizeValue]
+      : SPECULATIVE_TOKENS_DEFAULT_BY_ENGINE[engineKey];
   const showLora = [ModelType.LLM, ModelType.Image, ModelType.Video].includes(modelType);
   const showLoraKwargs = [ModelType.Image, ModelType.Video].includes(modelType);
   return (
@@ -126,6 +158,62 @@ const AdvancedConfig: FC<AdvancedConfigProps> = ({ form, modelType }) => {
           />
         </div>
       </ConfigSection>
+
+      {hasDrafter && (
+        <ConfigSection title={t('launchModel.speculativeDecoding')}>
+          <div className="px-2 py-3.5 flex items-center justify-between">
+            <FieldLabel label={t('launchModel.enableMtp')} tip={t('launchModel.enableMtpTip')} />
+            <FormField name="enable_mtp" valuePropName="checked">
+              <Switch />
+            </FormField>
+          </div>
+
+          {enableMtpValue && draftQuantizationOptions.length > 0 && (
+            <div className="px-2 py-3.5 flex items-center justify-between gap-4">
+              <FieldLabel
+                label={t('launchModel.draftModelQuantization')}
+                tip={t('launchModel.draftModelQuantizationTip')}
+              />
+              <FormField
+                className="w-56"
+                name="draft_quantization"
+                placeholder={t('launchModel.draftModelQuantizationPlaceholder')}
+              >
+                <Select options={draftQuantizationOptions} />
+              </FormField>
+            </div>
+          )}
+
+          {enableMtpValue && (
+            <div className="px-2 py-3.5 flex items-center justify-between gap-4">
+              <FieldLabel
+                label={t('launchModel.numSpeculativeTokens')}
+                tip={t('launchModel.numSpeculativeTokensTip')}
+              />
+              <FormField
+                className="w-56"
+                name="num_speculative_tokens"
+                placeholder={
+                  speculativeTokensDefault
+                    ? t('launchModel.numSpeculativeTokensPlaceholderValue', {
+                        value: speculativeTokensDefault,
+                      })
+                    : t('launchModel.numSpeculativeTokensPlaceholder')
+                }
+                rules={[
+                  {
+                    pattern: /^[1-9]\d*$/,
+                    message: t('launchModel.enterIntegerGreaterThanZero'),
+                  },
+                ]}
+                normalize={(v) => (v === '' ? undefined : Number(v))}
+              >
+                <Input type="number" min={1} />
+              </FormField>
+            </div>
+          )}
+        </ConfigSection>
+      )}
     </CollapsiblePanel>
   );
 };

frontend/src/components/pages/launch-model/launch-dialog/launch-dialog.tsx

@@ -243,6 +243,51 @@ export default function LaunchDialog({
     t,
   ]);
 
+  // The spec currently selected in the form, used to tell whether this
+  // format/size/quantization ships a drafter for speculative decoding.
+  const selectedSpec = useMemo(
+    () =>
+      model?.modelSpecs?.find(
+        (spec) =>
+          spec.model_format === modelFormatValue &&
+          toOptionValue(spec.model_size_in_billions) === modelSizeInBillionsKey &&
+          (spec.quantization === quantizationValue ||
+            (Array.isArray(spec.quantizations) && spec.quantizations.includes(quantizationValue)))
+      ),
+    [model?.modelSpecs, modelFormatValue, modelSizeInBillionsKey, quantizationValue]
+  );
+  // Drafter conversions available for the selected spec, each flagged with
+  // whether it is already downloaded. Empty when the spec declares a drafter
+  // without alternatives, in which case there is nothing to pick.
+  // A drafter needs both a spec that ships one and an engine that can run it:
+  // gemma-4's pytorch spec has a drafter but Transformers cannot use it, and
+  // showing the switch there only moves the failure to the launch button.
+  const engineSupportsDrafter = useMemo(() => {
+    const entries = modelEngineMap[modelEngineValue];
+    if (!Array.isArray(entries)) return false;
+    const matching = entries.filter(
+      (entry) =>
+        entry.model_format === modelFormatValue &&
+        toOptionValue(entry.model_size_in_billions) === modelSizeInBillionsKey
+    );
+    return (matching.length ? matching : entries).some((entry) => entry.support_draft_model);
+  }, [modelEngineMap, modelEngineValue, modelFormatValue, modelSizeInBillionsKey]);
+  const hasDrafter =
+    engineSupportsDrafter &&
+    Boolean(selectedSpec?.draft_model_id || selectedSpec?.draft_model_file_name_template);
+  const draftQuantizationOptions = useMemo(() => {
+    const quantizations = (selectedSpec?.draft_quantizations as string[] | undefined) || [];
+    const cacheStatus = selectedSpec?.draft_cache_status;
+
+    return quantizations.map((quantization, index) => ({
+      label: quantization,
+      value: quantization,
+      suffix: (Array.isArray(cacheStatus) ? cacheStatus[index] : cacheStatus)
+        ? t('launchModel.cached')
+        : undefined,
+    }));
+  }, [selectedSpec, t]);
+
   const multimodalProjectorOptions = useMemo(
     () =>
       Array.from(
@@ -524,7 +569,15 @@ export default function LaunchDialog({
         name: 'collapsibleConfig',
         type: 'custom',
         colSpan: 2,
-        content: <CollapsibleConfig form={form} modelType={modelType} />,
+        content: (
+          <CollapsibleConfig
+            form={form}
+            modelType={modelType}
+            modelName={model?.model_name}
+            hasDrafter={hasDrafter}
+            draftQuantizationOptions={draftQuantizationOptions}
+          />
+        ),
       },
     ],
     [ModelType.Embedding]: [