[MLX][Spec] Add Gemma 4 MTP E2E and docs

docs_new/cookbook/autoregressive/Google/Gemma4.mdx

@@ -183,6 +183,12 @@ sglang serve --model-path google/gemma-4-26B-A4B-it \
 
 Each Gemma 4 variant ships with a paired `*-assistant` draft model for NEXTN multi-token prediction. Use the commands below to enable MTP for the corresponding target model. These match the configuration generated when you toggle **Speculative Decoding (MTP) → Enabled** in the [interactive selector](#3-1-basic-configuration).
 
+<Note>
+The commands below target GPUs. On Apple Silicon, use the experimental E2B-only
+[Gemma 4 MTP guide](/docs/hardware-platforms/apple_metal#experimental-gemma-4-frozen-kv-mtp)
+instead.
+</Note>
+
 ```bash Command
 # Gemma 4 E2B + MTP
 sglang serve \

docs_new/docs/advanced_features/speculative_decoding.mdx

@@ -395,6 +395,12 @@ print(response.choices[0].message.content)
 
 We support [MTP (Multi-Token Prediction)](https://arxiv.org/pdf/2404.19737) in SGLang by using speculative decoding. We use `XiaomiMiMo/MiMo-7B-RL` as an example here (for DeepSeek MTP usage, refer to [DeepSeek-V3.2 cookbook §4.2.3](/cookbook/autoregressive/DeepSeek/DeepSeek-V3_2#4-2-3-multi-token-prediction-eagle-speculative-decoding)).
 
+<Note>
+Gemma 4 E2B on Apple Silicon uses a separate experimental, one-token path.
+Follow the [Apple Silicon Gemma 4 MTP guide](/docs/hardware-platforms/apple_metal#experimental-gemma-4-frozen-kv-mtp)
+for its pinned models and native-cache constraints.
+</Note>
+
 ```bash Command
 python3 -m sglang.launch_server \
     --model XiaomiMiMo/MiMo-7B-RL \

docs_new/docs/hardware-platforms/apple_metal.mdx

@@ -87,6 +87,64 @@ The MLX backend supports two quantization paths on Apple Silicon:
    ```
    The MLX backend silently ignores `--quantization mlx_q4` when the model is already quantized in its HF config (path 1), so the same flag is safe to pass either way.
 
+## Experimental Gemma 4 Frozen-KV MTP
+
+The MLX backend has a correctness-first, one-token MTP prototype for Gemma 4
+E2B. It is experimental: batching, overlap scheduling, radix/prefix reuse,
+paged KV cache, and performance claims are out of scope.
+
+Install the pinned assistant provider:
+
+```bash
+uv pip install "mlx-vlm==0.5.0"
+```
+
+Launch the tested target and assistant revisions with the complete Level 1
+configuration:
+
+```bash
+SGLANG_USE_MLX=1 python -m sglang.launch_server \
+  --model-path mlx-community/gemma-4-e2b-it-4bit \
+  --revision 238767527555cb75a05732a84dff5d6ba0dd6809 \
+  --speculative-algorithm FROZEN_KV_MTP \
+  --speculative-draft-model-path mlx-community/gemma-4-E2B-it-assistant-bf16 \
+  --speculative-draft-model-revision a7770799b560135ebdbfae8b7f468947415003bc \
+  --speculative-num-steps 1 \
+  --speculative-eagle-topk 1 \
+  --speculative-num-draft-tokens 2 \
+  --disable-radix-cache \
+  --disable-overlap-schedule \
+  --chunked-prefill-size -1 \
+  --max-running-requests 1 \
+  --context-length 2048 \
+  --max-total-tokens 2048
+```
+
+Every request must set `temperature` to `0` explicitly. Do not add
+`--language-only`; this prototype is already text-only, while that flag selects
+an unsupported encoder-disaggregation mode.
+
+Confirm activation after a request with:
+
+```bash
+curl -s http://127.0.0.1:30000/server_info \
+  | jq '.internal_states[0].speculative_worker'
+```
+
+The state reports `implementation: "mlx_gemma4_frozen_kv_mtp"` and positive
+`proposed_tokens` and `verified_tokens`. After completion or
+`POST /flush_cache`, `active_request_count`, `native_request_count`, and
+`assistant_request_binding_count` return to zero; flushing retains the same
+`assistant_generation`.
+
+Keep the prototype constraints in place: the exact text-only E2B pair; greedy
+decoding with one proposal, one step, top-k 1, and verify width 2; one request
+on one Apple host; synchronous scheduling; radix cache and chunked prefill
+disabled; and context length no greater than 2,048. If `--max-total-tokens` is
+set, it must be at least the context length. Sampling penalties, logprobs,
+structured output, custom logits processors, sessions, LoRA, multimodal input,
+and distributed execution are rejected.
+
 ## Benchmarking with Requests
 
 `sglang.benchmark_one_batch` calls the synchronous prefill/decode methods directly without going through the scheduler and the overlap code path.

python/sglang/srt/hardware_backend/mlx/speculative_worker.py

@@ -133,6 +133,9 @@ def __init__(self, server_args: ServerArgs, gpu_id, ps, nccl_port, target_worker
         )
         self._proposer = MlxGemma4MTPProposer(self._assistant_runtime)
         self._active_rids: set[str] = set()
+        self._proposed_tokens = 0
+        self._verified_tokens = 0
+        self._accepted_draft_tokens = 0
 
     @property
     def draft_worker(self):
@@ -144,6 +147,22 @@ def carries_draft_hidden_states(self) -> bool:
     def get_draft_kv_pool(self):
         return None
 
+    def get_speculative_internal_state(self) -> dict[str, object]:
+        """Return JSON-safe activation and request-lifecycle counters."""
+
+        return {
+            "implementation": "mlx_gemma4_frozen_kv_mtp",
+            "proposed_tokens": self._proposed_tokens,
+            "verified_tokens": self._verified_tokens,
+            "accepted_draft_tokens": self._accepted_draft_tokens,
+            "assistant_generation": self._assistant_loader.generation,
+            "active_request_count": len(self._active_rids),
+            "native_request_count": len(self._native_runner._req_caches),
+            "assistant_request_binding_count": (
+                self._assistant_runtime.request_binding_count
+            ),
+        }
+
     def alloc_memory_pool(
         self,
         memory_pool_config=None,
@@ -207,6 +226,7 @@ def _propose_from_output(
         proposal = int(self._proposer.propose_one(request_id, seed, cache))
         if proposal < 0 or proposal >= self._target_adapter.vocab_size:
             raise ValueError("assistant proposal is outside the target vocabulary")
+        self._proposed_tokens += 1
         return proposal
 
     def _forward_prefill(self, batch: ScheduleBatch) -> GenerationBatchResult:
@@ -273,6 +293,8 @@ def _forward_decode(self, batch: ScheduleBatch) -> GenerationBatchResult:
                 cache=candidate_cache,
             )
             self._native_runner.verify_commit(pending, decision)
+            self._verified_tokens += 1
+            self._accepted_draft_tokens += int(decision.accepted_draft)
         except BaseException:
             if pending is not None:
                 self._native_runner.verify_abort(pending)

python/sglang/srt/managers/scheduler.py

@@ -3956,6 +3956,12 @@ def get_internal_state(self, recv_req: GetInternalStateReq):
                 / self.metrics_reporter.spec_total_num_forward_ct
             )
 
+        get_speculative_state = getattr(
+            self.draft_worker, "get_speculative_internal_state", None
+        )
+        if callable(get_speculative_state):
+            ret["speculative_worker"] = get_speculative_state()
+
         if RECORD_STEP_TIME:
             ret["step_time_dict"] = self.metrics_reporter.step_time_dict
 

test/registered/mlx/models_e2e/test_gemma4_mtp_mlx_correctness.py

@@ -0,0 +1,200 @@
+"""Pinned real-model parity gate for the experimental MLX Gemma 4 MTP path."""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+import platform
+import time
+import unittest
+from contextlib import contextmanager
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+    try_cached_model,
+)
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+register_mlx_ci(est_time=8, suite="stage-b-e2e-mlx")
+
+TARGET = "mlx-community/gemma-4-e2b-it-4bit"
+TARGET_REVISION = "238767527555cb75a05732a84dff5d6ba0dd6809"
+ASSISTANT = "mlx-community/gemma-4-E2B-it-assistant-bf16"
+ASSISTANT_REVISION = "a7770799b560135ebdbfae8b7f468947415003bc"
+CROSS_WINDOW_PROMPT = "blue green amber violet " * 140
+
+_HAS_RUNTIME = (
+    platform.system() == "Darwin"
+    and platform.machine() == "arm64"
+    and importlib.util.find_spec("mlx") is not None
+    and importlib.util.find_spec("mlx_vlm.speculative.drafters.gemma4_assistant")
+    is not None
+)
+
+
+@unittest.skipUnless(_HAS_RUNTIME, "requires Apple Silicon, MLX, and mlx-vlm")
+class TestGemma4MTPMlxCorrectness(CustomTestCase):
+    base_url = DEFAULT_URL_FOR_TEST
+
+    @classmethod
+    @contextmanager
+    def _server(cls, *, mtp: bool):
+        args = [
+            "--revision",
+            TARGET_REVISION,
+            "--served-model-name",
+            TARGET,
+            "--disable-radix-cache",
+            "--disable-overlap-schedule",
+            "--chunked-prefill-size",
+            "-1",
+            "--max-running-requests",
+            "1",
+            "--context-length",
+            "2048",
+            "--max-total-tokens",
+            "2048",
+            "--log-level",
+            "warning",
+        ]
+        if mtp:
+            args.extend(
+                [
+                    "--speculative-algorithm",
+                    "FROZEN_KV_MTP",
+                    "--speculative-draft-model-path",
+                    try_cached_model(ASSISTANT),
+                    "--speculative-draft-model-revision",
+                    ASSISTANT_REVISION,
+                    "--speculative-num-steps",
+                    "1",
+                    "--speculative-eagle-topk",
+                    "1",
+                    "--speculative-num-draft-tokens",
+                    "2",
+                ]
+            )
+        env = os.environ.copy()
+        env["SGLANG_USE_MLX"] = "1"
+        process = popen_launch_server(
+            try_cached_model(TARGET),
+            cls.base_url,
+            timeout=max(DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, 600),
+            other_args=args,
+            env=env,
+        )
+        try:
+            yield
+        finally:
+            kill_process_tree(process.pid)
+
+    @classmethod
+    def _generate(cls, max_new_tokens: int) -> dict:
+        response = requests.post(
+            f"{cls.base_url}/generate",
+            json={
+                "text": CROSS_WINDOW_PROMPT,
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": max_new_tokens,
+                    "ignore_eos": True,
+                },
+            },
+            timeout=180,
+        )
+        response.raise_for_status()
+        return response.json()
+
+    @staticmethod
+    def _signature(result: dict) -> tuple[list[int], str, str]:
+        return (
+            result["output_ids"],
+            result["text"],
+            result["meta_info"]["finish_reason"]["type"],
+        )
+
+    @classmethod
+    def _spec_state(cls) -> dict:
+        response = requests.get(f"{cls.base_url}/server_info", timeout=30)
+        response.raise_for_status()
+        states = response.json()["internal_states"]
+        assert len(states) == 1
+        return states[0]["speculative_worker"]
+
+    @classmethod
+    def _wait_for_clean_state(cls, timeout: float = 30) -> dict:
+        deadline = time.monotonic() + timeout
+        last = None
+        while time.monotonic() < deadline:
+            last = cls._spec_state()
+            if all(
+                last[key] == 0
+                for key in (
+                    "active_request_count",
+                    "native_request_count",
+                    "assistant_request_binding_count",
+                )
+            ):
+                return last
+            time.sleep(0.05)
+        raise AssertionError(f"MLX MTP request state did not drain: {last}")
+
+    def test_cross_window_exact_token_parity_and_flush(self):
+        with self._server(mtp=False):
+            target = self._generate(64)
+            target_signature = self._signature(target)
+            self.assertEqual(len(target_signature[0]), 64)
+            self.assertEqual(target_signature[2], "length")
+            self.assertGreater(target["meta_info"]["prompt_tokens"], 512)
+
+        with self._server(mtp=True):
+            baseline_state = self._wait_for_clean_state()
+            result = self._generate(64)
+            self.assertEqual(self._signature(result), target_signature)
+
+            meta = result["meta_info"]
+            proposed = int(meta["spec_num_proposed_drafts"])
+            accepted = int(meta["spec_num_correct_drafts"])
+            self.assertGreater(proposed, 0)
+            self.assertLessEqual(0, accepted)
+            self.assertLessEqual(accepted, proposed)
+
+            state = self._wait_for_clean_state()
+            self.assertEqual(state["implementation"], "mlx_gemma4_frozen_kv_mtp")
+            self.assertGreaterEqual(state["proposed_tokens"], state["verified_tokens"])
+            self.assertGreater(state["verified_tokens"], 0)
+            self.assertLessEqual(
+                state["accepted_draft_tokens"], state["verified_tokens"]
+            )
+            self.assertEqual(
+                state["verified_tokens"] - baseline_state["verified_tokens"],
+                proposed,
+            )
+            self.assertEqual(
+                state["accepted_draft_tokens"]
+                - baseline_state["accepted_draft_tokens"],
+                accepted,
+            )
+
+            generation = state["assistant_generation"]
+            flush = requests.post(f"{self.base_url}/flush_cache", timeout=30)
+            flush.raise_for_status()
+            self.assertTrue(flush.text.startswith("Cache flushed"))
+            self.assertEqual(
+                self._wait_for_clean_state()["assistant_generation"], generation
+            )
+
+            post_flush = self._generate(8)
+            self.assertEqual(post_flush["output_ids"], target_signature[0][:8])
+            self._wait_for_clean_state()
+
+
+if __name__ == "__main__":
+    unittest.main()