[MLX][Spec] Add optional Gemma 4 target feedback

python/sglang/benchmark/one_batch.py

@@ -584,6 +584,7 @@ def __init__(self, model_runner, server_args):
         # standalone bench mode where no scheduler is present.
         init_kwargs = dict(
             model_path=server_args.model_path,
+            revision=server_args.revision,
             trust_remote_code=server_args.trust_remote_code,
             disable_radix_cache=True,
             mem_fraction_static=server_args.mem_fraction_static,

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

@@ -0,0 +1,92 @@
+"""Optional Gemma 4 target outputs needed by the MLX MTP proposer."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import mlx.core as mx
+
+
+@dataclass(frozen=True)
+class MlxTargetForwardOutput:
+    logits: mx.array
+    hidden_states: mx.array | None = None
+
+
+@dataclass(frozen=True)
+class MlxTargetSeed:
+    token_id: int
+    hidden_state: mx.array
+    token_embedding: mx.array
+
+
+def _extract_logits(model_output: Any) -> mx.array:
+    return model_output[0] if isinstance(model_output, tuple) else model_output
+
+
+class Gemma4TargetAdapter:
+    """Capture Gemma 4 backbone output without changing normal target math."""
+
+    def __init__(self, model: Any):
+        self.model = model
+        self.causal_model = getattr(model, "language_model", model)
+        self.backbone = getattr(self.causal_model, "model", None)
+        if self.backbone is None or not hasattr(self.backbone, "embed_tokens"):
+            raise TypeError(
+                "Gemma4TargetAdapter requires an mlx-lm Gemma 4 text model "
+                "or conditional-generation wrapper"
+            )
+
+    @property
+    def vocab_size(self) -> int:
+        return int(self.backbone.config.vocab_size)
+
+    def input_embeddings(self, input_ids: mx.array) -> mx.array:
+        embed_scale = float(getattr(self.backbone, "embed_scale", 1.0))
+        return self.backbone.embed_tokens(input_ids) * embed_scale
+
+    def forward(
+        self,
+        input_ids: mx.array,
+        *,
+        cache: list[Any] | None = None,
+        collect_hidden_states: bool = False,
+    ) -> MlxTargetForwardOutput:
+        if not collect_hidden_states:
+            return MlxTargetForwardOutput(
+                logits=_extract_logits(self.model(input_ids, cache=cache))
+            )
+
+        hidden = self.backbone(input_ids, cache=cache)
+        if bool(getattr(self.causal_model, "tie_word_embeddings", False)):
+            logits = self.backbone.embed_tokens.as_linear(hidden)
+        else:
+            logits = self.causal_model.lm_head(hidden)
+
+        softcap = getattr(self.causal_model, "final_logit_softcapping", None)
+        if softcap is not None:
+            logits = mx.tanh(logits / softcap) * softcap
+        return MlxTargetForwardOutput(logits=logits, hidden_states=hidden)
+
+    def make_seed(
+        self,
+        output: MlxTargetForwardOutput,
+        *,
+        hidden_row_index: int,
+        emitted_token_id: int,
+    ) -> MlxTargetSeed:
+        hidden = output.hidden_states
+        if hidden is None or hidden.ndim != 3:
+            raise ValueError("target output has no batch-major hidden states")
+        if hidden_row_index < 0 or hidden_row_index >= hidden.shape[1]:
+            raise IndexError("target hidden row is outside the query width")
+        if emitted_token_id < 0 or emitted_token_id >= self.vocab_size:
+            raise ValueError("assistant seed token is outside the target vocabulary")
+
+        token_ids = mx.array([[emitted_token_id]], dtype=mx.int32)
+        return MlxTargetSeed(
+            token_id=int(emitted_token_id),
+            hidden_state=hidden[:, hidden_row_index : hidden_row_index + 1, :],
+            token_embedding=self.input_embeddings(token_ids),
+        )

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

@@ -70,6 +70,7 @@ class MlxPendingPrefill:
     full_token_ids: list[int]
     req_pool_idx: int
     synced_offset: int
+    target_output: Any | None = None
 
 
 @dataclass
@@ -104,6 +105,17 @@ class MlxPendingDecode:
     caches: list[list[Any]]
 
 
+@dataclass
+class MlxPendingVerify:
+    """Isolated target verification awaiting proposal and commit."""
+
+    req_id: str
+    query_token_ids: tuple[int, int]
+    transaction: Any
+    speculative_cache: list[Any]
+    target_output: Any
+
+
 _MLX_QUANTIZATION_PRESETS: dict[str, tuple[int, int]] = {
     # name -> (bits, group_size). group_size=64 matches the mlx-community convention.
     "mlx_q4": (4, 64),
@@ -150,13 +162,15 @@ class MlxModelRunner:
     def __init__(
         self,
         model_path: str,
+        revision: str | None = None,
         trust_remote_code: bool = False,
         disable_radix_cache: bool = False,
         pool_size: int | None = None,
         mem_fraction_static: float = 0.8,
         quantization: str | None = None,
     ):
         self.model_path = model_path
+        self.revision = revision
         self.trust_remote_code = trust_remote_code
         self.model = None
         self.disable_radix_cache = disable_radix_cache
@@ -175,6 +189,7 @@ def __init__(
         self._native_cache_fallback = (
             _mlx_model_type(self.model) in _NATIVE_CACHE_FALLBACK_MODEL_TYPES
         )
+        self._target_adapter = None
         if self._native_cache_fallback and not self.disable_radix_cache:
             raise NotImplementedError(
                 "Gemma 4 on the MLX backend currently requires "
@@ -236,6 +251,19 @@ def _extract_logits(model_output):
             return model_output[0]
         return model_output
 
+    def _get_target_adapter(self):
+        """Construct the optional Gemma 4 feedback path only when MTP uses it."""
+
+        if not self.native_cache_fallback:
+            raise RuntimeError("Gemma 4 target feedback requires native caches")
+        if self._target_adapter is None:
+            from sglang.srt.hardware_backend.mlx.model_adapter import (
+                Gemma4TargetAdapter,
+            )
+
+            self._target_adapter = Gemma4TargetAdapter(self.model)
+        return self._target_adapter
+
     def _new_cache_skeleton(self) -> list[Any]:
         """Create a model-shaped cache list before attention cache wiring."""
         if self.native_cache_fallback:
@@ -468,6 +496,7 @@ def _load_model(self):
             self.model_path,
             tokenizer_config={"trust_remote_code": self.trust_remote_code},
             return_config=True,
+            revision=self.revision,
         )
         self.model, _tokenizer, config = loaded
 
@@ -724,6 +753,46 @@ def prefill(
         self._eval_with_cache(pending.lazy_token, pending.cache)
         return self.prefill_finalize(pending)
 
+    def prefill_for_mtp(
+        self,
+        req_id: str,
+        new_token_ids: list[int],
+        full_token_ids: list[int],
+        prefix_slot_ids: list[int],
+        new_slot_ids: list[int],
+        req_pool_idx: int,
+        req: Any | None = None,
+    ) -> tuple[int, Any]:
+        """Prefill native Gemma 4 and retain the final target hidden row."""
+
+        pending = self.prefill_start(
+            req_id=req_id,
+            new_token_ids=new_token_ids,
+            full_token_ids=full_token_ids,
+            prefix_slot_ids=prefix_slot_ids,
+            new_slot_ids=new_slot_ids,
+            req_pool_idx=req_pool_idx,
+            req=req,
+            collect_hidden_states=True,
+        )
+        target_output = pending.target_output
+        assert target_output is not None and target_output.hidden_states is not None
+        mx.eval(
+            pending.lazy_token,
+            target_output.hidden_states,
+            *self._cache_state_arrays([pending.cache]),
+        )
+        token = self.prefill_finalize(pending)
+        self._assert_native_history_lag(req_id)
+        from sglang.srt.hardware_backend.mlx.model_adapter import (
+            MlxTargetForwardOutput,
+        )
+
+        return token, MlxTargetForwardOutput(
+            logits=target_output.logits[:, -1:, :],
+            hidden_states=target_output.hidden_states[:, -1:, :],
+        )
+
     def extend(
         self,
         req_id: str,
@@ -807,6 +876,123 @@ def decode_batch(
         mx.eval(pending.lazy_tokens, *cache_arrays)
         return self.decode_batch_finalize(pending)
 
+    def _forward_native_queries_sequential(
+        self,
+        cache: list[Any],
+        query_token_ids: tuple[int, ...],
+        *,
+        collect_hidden_states: bool,
+    ) -> Any:
+        """Match normal decode's one-token MLX numerical path exactly."""
+
+        if not query_token_ids:
+            raise ValueError("target verification requires at least one query")
+        adapter = self._get_target_adapter()
+        outputs = [
+            adapter.forward(
+                mx.array([[token]], dtype=mx.int32),
+                cache=cache,
+                collect_hidden_states=collect_hidden_states,
+            )
+            for token in query_token_ids
+        ]
+
+        from sglang.srt.hardware_backend.mlx.model_adapter import (
+            MlxTargetForwardOutput,
+        )
+
+        hidden = None
+        if collect_hidden_states:
+            rows = [output.hidden_states for output in outputs]
+            if any(row is None for row in rows):
+                raise RuntimeError("target adapter omitted requested hidden states")
+            hidden = mx.concatenate(rows, axis=1)
+        return MlxTargetForwardOutput(
+            logits=mx.concatenate([output.logits for output in outputs], axis=1),
+            hidden_states=hidden,
+        )
+
+    def verify_start(
+        self, req_id: str, query_token_ids: tuple[int, int]
+    ) -> MlxPendingVerify:
+        """Run two target queries on an isolated native-cache clone."""
+
+        if req_id not in self._req_caches:
+            raise KeyError(f"unknown native-cache request {req_id!r}")
+        queries = tuple(int(token) for token in query_token_ids)
+        if len(queries) != 2:
+            raise ValueError("one-draft verification requires exactly two queries")
+
+        from sglang.srt.hardware_backend.mlx.kv_cache.native_transaction import (
+            MlxNativeCacheTransaction,
+        )
+
+        transaction = MlxNativeCacheTransaction(
+            self._req_caches[req_id],
+            queries,
+            lambda cache, accepted: self._forward_native_queries_sequential(
+                cache, accepted, collect_hidden_states=False
+            ).logits,
+        )
+        speculative_cache = transaction.begin()
+        try:
+            output = self._forward_native_queries_sequential(
+                speculative_cache, queries, collect_hidden_states=True
+            )
+        except BaseException:
+            transaction.abort()
+            raise
+        return MlxPendingVerify(
+            req_id=req_id,
+            query_token_ids=queries,
+            transaction=transaction,
+            speculative_cache=speculative_cache,
+            target_output=output,
+        )
+
+    @staticmethod
+    def verify_materialize(pending: MlxPendingVerify) -> tuple[int, int]:
+        arrays = [pending.target_output.logits]
+        if pending.target_output.hidden_states is not None:
+            arrays.append(pending.target_output.hidden_states)
+        arrays.extend(MlxModelRunner._cache_state_arrays([pending.speculative_cache]))
+        mx.eval(*arrays)
+        target_ids = mx.argmax(pending.target_output.logits, axis=-1)
+        return tuple(int(token) for token in target_ids.reshape(-1).tolist())
+
+    def verify_prepare(self, pending: MlxPendingVerify, decision: Any) -> list[Any]:
+        """Replay accepted queries into a candidate cache without publishing it."""
+
+        if pending.req_id != decision.request_id:
+            pending.transaction.abort()
+            raise ValueError("verification decision belongs to another request")
+        pending.transaction.prepare(decision.committed_query_count)
+        return pending.transaction.candidate_cache
+
+    def verify_commit(self, pending: MlxPendingVerify, decision: Any) -> None:
+        """Publish the prepared cache only after the next proposal succeeds."""
+
+        pending.transaction.commit()
+        self._req_token_ids[pending.req_id].extend(decision.emitted_token_ids)
+        self._assert_native_history_lag(pending.req_id)
+
+    @staticmethod
+    def verify_abort(pending: MlxPendingVerify) -> None:
+        if pending.transaction.active:
+            pending.transaction.abort()
+
+    def _assert_native_history_lag(self, req_id: str) -> None:
+        offsets = {int(entry.offset) for entry in self._req_caches[req_id]}
+        if len(offsets) != 1:
+            raise RuntimeError(f"native cache offsets disagree for {req_id!r}")
+        cache_length = offsets.pop()
+        history_length = len(self._req_token_ids[req_id])
+        if history_length != cache_length + 1:
+            raise RuntimeError(
+                "native Gemma 4 history must remain one token ahead of KV: "
+                f"history={history_length}, cache={cache_length}"
+            )
+
     def prefill_start(
         self,
         req_id: str,
@@ -816,6 +1002,8 @@ def prefill_start(
         new_slot_ids: list[int],
         req_pool_idx: int,
         req: Any | None = None,
+        *,
+        collect_hidden_states: bool = False,
     ) -> MlxPendingPrefill:
         """Queue a prefill forward pass without evaluating.
 
@@ -831,8 +1019,14 @@ def prefill_start(
         if self.disable_radix_cache:
             cache = self._acquire_cache()
             input_ids = mx.array([new_token_ids], dtype=mx.int32)
-            model_output = self.model(input_ids, cache=cache)
-            logits = self._extract_logits(model_output)
+            target_output = None
+            if collect_hidden_states:
+                target_output = self._get_target_adapter().forward(
+                    input_ids, cache=cache, collect_hidden_states=True
+                )
+                logits = target_output.logits
+            else:
+                logits = self._extract_logits(self.model(input_ids, cache=cache))
             lazy_token = mx.argmax(logits[:, -1, :], axis=-1)
             return MlxPendingPrefill(
                 lazy_token=lazy_token,
@@ -841,6 +1035,7 @@ def prefill_start(
                 full_token_ids=list(full_token_ids),
                 req_pool_idx=req_pool_idx,
                 synced_offset=0,
+                target_output=target_output,
             )
 
         assert self._attention_kv_pool is not None

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

@@ -53,6 +53,7 @@ def _init_model_runner(self):
         logger.info("Initializing MlxModelRunner for end-to-end MLX inference")
         init_kwargs = dict(
             model_path=get_model().model_path,
+            revision=getattr(self.server_args, "revision", None),
             trust_remote_code=get_model().trust_remote_code,
             disable_radix_cache=get_memory().disable_radix_cache,
             mem_fraction_static=get_schedule().mem_fraction_static,