diff --git a/python/sglang/benchmark/one_batch.py b/python/sglang/benchmark/one_batch.py index 8556becec68b..0a0cf06317e2 100644 --- a/python/sglang/benchmark/one_batch.py +++ b/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, diff --git a/python/sglang/srt/hardware_backend/mlx/model_adapter.py b/python/sglang/srt/hardware_backend/mlx/model_adapter.py new file mode 100644 index 000000000000..56a0696309b8 --- /dev/null +++ b/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), + ) diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner.py b/python/sglang/srt/hardware_backend/mlx/model_runner.py index 1173a8e02820..287c8db963d1 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner.py +++ b/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,6 +162,7 @@ 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, @@ -157,6 +170,7 @@ def __init__( 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 diff --git a/python/sglang/srt/hardware_backend/mlx/tp_worker.py b/python/sglang/srt/hardware_backend/mlx/tp_worker.py index 0b8c56fd60a1..3258a578bff0 100644 --- a/python/sglang/srt/hardware_backend/mlx/tp_worker.py +++ b/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, diff --git a/test/registered/unit/hardware_backend/mlx/test_gemma4_target_adapter.py b/test/registered/unit/hardware_backend/mlx/test_gemma4_target_adapter.py new file mode 100644 index 000000000000..1ea1609d5636 --- /dev/null +++ b/test/registered/unit/hardware_backend/mlx/test_gemma4_target_adapter.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +import importlib.util +import unittest +from types import SimpleNamespace +from unittest import mock + +import numpy as np + +from sglang.test.ci.ci_register import register_mlx_ci + +register_mlx_ci(est_time=1, suite="stage-a-unit-test-mlx") + +_HAS_MLX = ( + importlib.util.find_spec("mlx") is not None + and importlib.util.find_spec("mlx_lm.models.gemma4_text") is not None +) + +if _HAS_MLX: + import mlx.core as mx + from registered.unit.hardware_backend.mlx.utils import ( + assert_native_cache_equal, + build_runner, + tiny_gemma4, + ) + + from sglang.srt.hardware_backend.mlx.model_adapter import ( + Gemma4TargetAdapter, + MlxTargetForwardOutput, + ) + from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner + + +@unittest.skipUnless(_HAS_MLX, "requires MLX and mlx-lm Gemma 4") +class TestGemma4TargetAdapter(unittest.TestCase): + def test_wrapped_and_unwrapped_capture_preserve_exact_target_math(self): + for wrapped in (False, True): + with self.subTest(wrapped=wrapped): + model = tiny_gemma4(wrapped=wrapped) + adapter = Gemma4TargetAdapter(model) + ids = mx.array([[1, 2, 3, 4]], dtype=mx.int32) + + normal_cache = model.make_cache() + normal = model(ids, cache=normal_cache) + captured_cache = model.make_cache() + captured = adapter.forward( + ids, cache=captured_cache, collect_hidden_states=True + ) + mx.eval(normal, captured.logits, captured.hidden_states) + + np.testing.assert_array_equal(normal, captured.logits) + np.testing.assert_array_equal( + mx.argmax(normal, axis=-1), mx.argmax(captured.logits, axis=-1) + ) + self.assertEqual(captured.hidden_states.shape, (1, 4, 16)) + assert_native_cache_equal(self, captured_cache, normal_cache) + + def test_prefill_retains_only_the_scaled_final_seed_row(self): + model = tiny_gemma4() + runner = build_runner(model) + token, output = runner.prefill_for_mtp( + "request", [1, 2, 3], [1, 2, 3], [], [], 0 + ) + adapter = runner._get_target_adapter() + seed = adapter.make_seed(output, hidden_row_index=0, emitted_token_id=token) + expected_embedding = ( + model.model.embed_tokens(mx.array([[token]], dtype=mx.int32)) + * model.model.embed_scale + ) + mx.eval(seed.hidden_state, seed.token_embedding, expected_embedding) + + self.assertEqual(output.hidden_states.shape, (1, 1, 16)) + self.assertEqual(seed.hidden_state.shape, (1, 1, 16)) + np.testing.assert_array_equal(seed.token_embedding, expected_embedding) + + def test_verify_and_replay_use_one_token_forward_shapes(self): + calls = [] + + class RecordingAdapter: + def forward(self, input_ids, *, cache, collect_hidden_states): + del cache + calls.append((tuple(input_ids.shape), collect_hidden_states)) + width = input_ids.shape[1] + return MlxTargetForwardOutput( + logits=mx.zeros((1, width, 7)), + hidden_states=( + mx.zeros((1, width, 5)) if collect_hidden_states else None + ), + ) + + runner = object.__new__(MlxModelRunner) + runner._native_cache_fallback = True + runner._target_adapter = RecordingAdapter() + output = runner._forward_native_queries_sequential( + [], (11, 12), collect_hidden_states=True + ) + self.assertEqual(calls, [((1, 1), True), ((1, 1), True)]) + self.assertEqual(output.hidden_states.shape, (1, 2, 5)) + + calls.clear() + runner._forward_native_queries_sequential( + [], (13, 14), collect_hidden_states=False + ) + self.assertEqual(calls, [((1, 1), False), ((1, 1), False)]) + + def test_target_loader_receives_pinned_revision(self): + runner = MlxModelRunner.__new__(MlxModelRunner) + runner.model_path = "target" + runner.revision = "pinned-revision" + runner.trust_remote_code = False + runner._quantization = None + fake_model = SimpleNamespace(parameters=lambda: ()) + with ( + mock.patch( + "sglang.srt.hardware_backend.mlx.model_runner.mlx_lm_load", + return_value=(fake_model, None, {}), + ) as load, + mock.patch("sglang.srt.hardware_backend.mlx.model_runner.mx.eval"), + ): + runner._load_model() + self.assertEqual(load.call_args.kwargs["revision"], "pinned-revision") + + +if __name__ == "__main__": + unittest.main()