From c3da0a2582d32b7ac9c90ab1d193d8b2749165dd Mon Sep 17 00:00:00 2001 From: NOOB <43700604+noob-se7en@users.noreply.github.com> Date: Tue, 7 Jul 2026 09:01:33 +0530 Subject: [PATCH] [MLX] Fix single-token chunked-prefill continuation misrouted as decode (#30181) --- .../srt/hardware_backend/mlx/tp_worker.py | 75 ++++-- .../mlx/test_tp_worker_routing.py | 244 ++++++++++++++++++ 2 files changed, 293 insertions(+), 26 deletions(-) create mode 100644 test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py diff --git a/python/sglang/srt/hardware_backend/mlx/tp_worker.py b/python/sglang/srt/hardware_backend/mlx/tp_worker.py index 0f81881f0..ea51daf7f 100644 --- a/python/sglang/srt/hardware_backend/mlx/tp_worker.py +++ b/python/sglang/srt/hardware_backend/mlx/tp_worker.py @@ -132,6 +132,29 @@ class MlxTpModelWorker(TpModelWorker): # insert. Any older tracked slot is released during component cleanup. req.mamba_last_track_seqlen = None + def _route_extend_request(self, rid: str, decoding_rids: set[str]) -> str: + """Classify a request within an extend / mixed batch. + + Shared by the sync (:meth:`_forward_batch_generation_mlx`) and async + (:meth:`_async_extend_batch`) paths so both route identically. + + Returns one of: + + * ``"prefill"`` -- not seen before; start a fresh prefill. + * ``"decode"`` -- a genuine single-token decode step mixed into + this batch (present in ``batch.decoding_reqs``). + * ``"continuation"`` -- a chunked-prefill continuation. Routing keys on + request state, **not** ``seq_len``: a final continuation chunk can be + exactly one token, which must still extend. Routing it as a decode + would drop the real token and feed the model its own previous-chunk + prediction, silently corrupting the output. + """ + if not self._mlx_runner.has_request(rid): + return "prefill" + if rid in decoding_rids: + return "decode" + return "continuation" + def _forward_batch_generation_mlx( self, batch: ScheduleBatch ) -> GenerationBatchResult: @@ -164,6 +187,9 @@ class MlxTpModelWorker(TpModelWorker): prefill_rids: list[tuple[str, int]] = [] extend_rids: list[tuple[str, int]] = [] decode_rids: list[str] = [] + # Genuine decode steps mixed into this extend batch; see + # _route_extend_request. + decoding_rids = {r.rid for r in (batch.decoding_reqs or [])} for i, req in enumerate(reqs): seq_len = extend_seq_lens[i] @@ -172,18 +198,15 @@ class MlxTpModelWorker(TpModelWorker): offset += seq_len slot_offset += seq_len - if self._mlx_runner.has_request(req.rid): - if seq_len > 1: - # Chunked prefill continuation - next_token = self._mlx_runner.extend( - req.rid, req_token_ids, req_new_slots - ) - extend_rids.append((req.rid, next_token)) - else: - # MIXED mode: single-token decode - decode_rids.append(req.rid) - else: - # New prefill + route = self._route_extend_request(req.rid, decoding_rids) + if route == "continuation": + next_token = self._mlx_runner.extend( + req.rid, req_token_ids, req_new_slots + ) + extend_rids.append((req.rid, next_token)) + elif route == "decode": + decode_rids.append(req.rid) + else: # "prefill" prefix_slot_ids = req.prefix_indices.tolist() full_token_ids = list(req.get_fill_ids()) next_token = self._mlx_runner.prefill( @@ -306,6 +329,9 @@ class MlxTpModelWorker(TpModelWorker): pending_prefills: list[MlxPendingPrefill] = [] pending_extends: list[MlxPendingExtend] = [] mixed_decode_rids: list[str] = [] + # Genuine decode steps mixed into this extend batch; see + # _route_extend_request. + decoding_rids = {r.rid for r in (batch.decoding_reqs or [])} for i, req in enumerate(reqs): seq_len = extend_seq_lens[i] @@ -314,21 +340,18 @@ class MlxTpModelWorker(TpModelWorker): offset += seq_len slot_offset += seq_len - if self._mlx_runner.has_request(req.rid): - if seq_len > 1: - # Chunked prefill continuation - pending_extends.append( - self._mlx_runner.extend_start( - req_id=req.rid, - new_token_ids=req_token_ids, - new_slot_ids=req_new_slots, - ) + route = self._route_extend_request(req.rid, decoding_rids) + if route == "continuation": + pending_extends.append( + self._mlx_runner.extend_start( + req_id=req.rid, + new_token_ids=req_token_ids, + new_slot_ids=req_new_slots, ) - else: - # MIXED mode: single-token decode - mixed_decode_rids.append(req.rid) - else: - # New prefill + ) + elif route == "decode": + mixed_decode_rids.append(req.rid) + else: # "prefill" prefix_slot_ids = req.prefix_indices.tolist() full_token_ids = list(req.get_fill_ids()) pending_prefills.append( diff --git a/test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py b/test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py new file mode 100644 index 000000000..4cb4b5855 --- /dev/null +++ b/test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py @@ -0,0 +1,244 @@ +"""Unit tests for MLX extend-batch request routing in ``MlxTpModelWorker``. + +Regression guard for the single-token chunked-prefill continuation bug: a +continuation whose final chunk is exactly one token (prompt length == +k * chunked_prefill_size + 1) must be routed to the **extend** path, not the +**decode** path. + +The old routing keyed on ``seq_len > 1`` as a proxy for "is this a +continuation"; a 1-token continuation is indistinguishable, by length, from a +genuine single-token decode step mixed into the batch, so it was misrouted to +decode. The decode path ignores the batch's real token and feeds the model its +own stored prediction from the previous chunk -> the true last prompt token is +silently dropped and generation is conditioned on a corrupted prompt. The +correct discriminator is ``batch.decoding_reqs``, not the chunk length. + +The routing decision was duplicated across the sync and async paths (the bug +therefore existed in both). It now lives in the shared +``MlxTpModelWorker._route_extend_request`` helper. These tests cover: + + * the helper decision directly (both paths delegate to it); + * the sync wiring, by driving ``_forward_batch_generation_mlx``; + * the async wiring, by driving ``_async_extend_batch``. + +They mock the MLX runner and load no model. Apple-Silicon-only because +``tp_worker`` imports ``mlx.core`` at module load. +""" + +from __future__ import annotations + +import importlib.util +import platform +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci + +# AST-parsed "this test exists" marker; actual execution is gated by the +# @skipUnless guard below (mirrors test_quantization.py in this directory). +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64" +_HAS_MLX = importlib.util.find_spec("mlx") is not None +_SKIP_REASON = "Apple-Silicon-only (tp_worker imports mlx.core at module load)" + + +class _FakeRunner: + """Records which routing path each request took (sync + async surfaces).""" + + def __init__(self, known_rids): + self._known = set(known_rids) + self.calls: list[tuple[str, str]] = [] # (op, rid) + self._req_caches: dict[str, list] = {} + self._counter = 0 + + # --- shared --- + def has_request(self, rid): + return rid in self._known + + def flush_all_decode_kv(self): + pass + + def ops_for(self, rid): + return [op for op, r in self.calls if r == rid] + + @staticmethod + def _fake_cache_layer(): + import mlx.core as mx + + return SimpleNamespace(state=[mx.array([0.0], dtype=mx.float32)]) + + # --- sync surface --- + def extend(self, rid, new_token_ids, new_slot_ids): + self.calls.append(("extend", rid)) + self._counter += 1 + return 1000 + self._counter + + def decode_batch(self, rids): + for rid in rids: + self.calls.append(("decode", rid)) + return [2000 + i for i in range(len(rids))] + + def prefill( + self, + req_id, + new_token_ids, + full_token_ids, + prefix_slot_ids, + new_slot_ids, + req_pool_idx, + req=None, + ): + self.calls.append(("prefill", req_id)) + return 3000 + + # --- async surface --- + def extend_start(self, req_id, new_token_ids, new_slot_ids): + import mlx.core as mx + + self.calls.append(("extend_start", req_id)) + self._req_caches[req_id] = [self._fake_cache_layer()] + return SimpleNamespace(lazy_token=mx.array([0], dtype=mx.int32), req_id=req_id) + + def prefill_start( + self, + req_id, + new_token_ids, + full_token_ids, + prefix_slot_ids, + new_slot_ids, + req_pool_idx, + req=None, + ): + import mlx.core as mx + + self.calls.append(("prefill_start", req_id)) + return SimpleNamespace( + lazy_token=mx.array([0], dtype=mx.int32), + cache=[self._fake_cache_layer()], + req_id=req_id, + ) + + def decode_batch_start(self, rids): + import mlx.core as mx + + for rid in rids: + self.calls.append(("decode_start", rid)) + return SimpleNamespace( + lazy_tokens=mx.array([0] * len(rids), dtype=mx.int32), + caches=[[self._fake_cache_layer()] for _ in rids], + req_ids=list(rids), + ) + + +class _FakeReq: + def __init__(self, rid, req_pool_idx=0): + self.rid = rid + self.prefix_indices = torch.empty(0, dtype=torch.long) + self.fill_ids = [0] + self.req_pool_idx = req_pool_idx + + def get_fill_ids(self): + return self.fill_ids + + +class _FakeBatch: + def __init__(self, forward_mode, reqs, extend_lens, decoding_reqs=None): + total = sum(extend_lens) + self.forward_mode = forward_mode + self.reqs = reqs + self.extend_lens = list(extend_lens) + self.decoding_reqs = decoding_reqs + # Arbitrary but correctly-sized token / slot arrays. + self.input_ids = torch.arange(total, dtype=torch.long) + self.out_cache_loc = torch.arange(total, dtype=torch.long) + + +@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON) +class TestMlxExtendRouting(unittest.TestCase): + """Routing contract for MlxTpModelWorker: shared helper + sync + async.""" + + @staticmethod + def _worker(known_rids): + from sglang.srt.hardware_backend.mlx.tp_worker import MlxTpModelWorker + + worker = MlxTpModelWorker.__new__(MlxTpModelWorker) + worker._mlx_runner = _FakeRunner(known_rids) + worker._mlx_active_rids = set() + return worker + + # ---------- the shared decision helper ---------- + # The helper takes no seq_len: length cannot distinguish a 1-token + # continuation from a genuine decode -- request state does. + + def test_route_unseen_request_is_prefill(self): + worker = self._worker(known_rids=set()) + self.assertEqual(worker._route_extend_request("r1", set()), "prefill") + + def test_route_seen_non_decode_is_continuation(self): + worker = self._worker(known_rids={"r1"}) + self.assertEqual(worker._route_extend_request("r1", set()), "continuation") + + def test_route_seen_and_in_decoding_reqs_is_decode(self): + worker = self._worker(known_rids={"r1"}) + self.assertEqual(worker._route_extend_request("r1", {"r1"}), "decode") + + # ---------- sync path: _forward_batch_generation_mlx ---------- + + def _run_sync(self, reqs, extend_lens, known_rids, decoding_reqs, forward_mode): + worker = self._worker(known_rids) + batch = _FakeBatch(forward_mode, reqs, extend_lens, decoding_reqs) + result = worker._forward_batch_generation_mlx(batch) + assert result.next_token_ids.numel() == len(reqs) + return worker._mlx_runner + + def test_sync_one_token_continuation_routes_to_extend(self): + """THE REGRESSION (sync): a 1-token continuation must extend, not decode.""" + runner = self._run_sync([_FakeReq("r1")], [1], {"r1"}, None, ForwardMode.EXTEND) + self.assertEqual(runner.ops_for("r1"), ["extend"]) + + def test_sync_multi_token_continuation_routes_to_extend(self): + runner = self._run_sync([_FakeReq("r1")], [4], {"r1"}, None, ForwardMode.EXTEND) + self.assertEqual(runner.ops_for("r1"), ["extend"]) + + def test_sync_genuine_mixed_decode_routes_to_decode(self): + p, d = _FakeReq("p1"), _FakeReq("d1") + runner = self._run_sync([p, d], [4, 1], {"d1"}, [d], ForwardMode.MIXED) + self.assertEqual(runner.ops_for("p1"), ["prefill"]) + self.assertEqual(runner.ops_for("d1"), ["decode"]) + + # ---------- async path: _async_extend_batch ---------- + + def _run_async(self, reqs, extend_lens, known_rids, decoding_reqs, forward_mode): + from sglang.srt.hardware_backend.mlx.tp_worker import MlxTpModelWorker + + worker = MlxTpModelWorker.__new__(MlxTpModelWorker) + worker._mlx_runner = _FakeRunner(known_rids) + batch = _FakeBatch(forward_mode, reqs, extend_lens, decoding_reqs) + # returns (lazy_stacked, pending_prefills, pending_extends, + # pending_mixed_decode, mode) + result = worker._async_extend_batch(batch) + return worker._mlx_runner, result + + def test_async_one_token_continuation_routes_to_extend(self): + """THE REGRESSION (async): a 1-token continuation must extend, not decode.""" + runner, result = self._run_async( + [_FakeReq("r1")], [1], {"r1"}, None, ForwardMode.EXTEND + ) + self.assertEqual(runner.ops_for("r1"), ["extend_start"]) + self.assertEqual(len(result[2]), 1) # one pending extend + self.assertIsNone(result[3]) # no mixed decode + + def test_async_genuine_mixed_decode_routes_to_decode(self): + p, d = _FakeReq("p1"), _FakeReq("d1") + runner, result = self._run_async([p, d], [4, 1], {"d1"}, [d], ForwardMode.MIXED) + self.assertEqual(runner.ops_for("p1"), ["prefill_start"]) + self.assertEqual(runner.ops_for("d1"), ["decode_start"]) + self.assertIsNotNone(result[3]) # pending mixed decode present + + +if __name__ == "__main__": + unittest.main()