fix: repair MLX Xavier recovery and prefix publication

Author: qinxuye

xinference/core/supervisor.py

@@ -3095,6 +3095,14 @@ async def launch_builtin_model(
             raise ValueError("NIXL requires explicit prefill and decode replica roles")
         # Xavier-related
         requested_xavier = bool(kwargs.pop("enable_xavier", False))
+        if (
+            requested_xavier
+            and not pd_enabled
+            and replica <= 1
+            and (model_engine or "").lower() in ("mlx", "sglang")
+        ):
+            logger.warning("Enabling xavier when replica<=1 is meaningless.")
+            requested_xavier = False
         mlx_xavier = (
             (requested_xavier or pd_enabled)
             and transport_backend == "xavier"
@@ -3108,7 +3116,7 @@ async def launch_builtin_model(
             if (
                 model_type not in (None, "LLM")
                 or model_format != "mlx"
-                or quantization not in (None, "none")
+                or quantization not in (None, "none", "fp16", "bf16")
             ):
                 raise ValueError("MLX Xavier requires unquantized MLX text weights")
             kwargs["_xavier_cache_config"] = {
@@ -3121,9 +3129,6 @@ async def launch_builtin_model(
             and model_engine is not None
             and model_engine.lower() == "sglang"
         )
-        if sglang_xavier and not pd_enabled and replica <= 1:
-            logger.warning("Enabling xavier when replica<=1 is meaningless.")
-            sglang_xavier = False
         sglang_nixl = (
             pd_enabled
             and transport_backend == "nixl"

xinference/core/tests/test_pd_launch.py

@@ -69,10 +69,16 @@ def launch_kwargs():
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("pd", [False, True])
-async def test_mlx_xavier_launch_and_cleanup(launch_runtime, pd):
+@pytest.mark.parametrize("quantization", [None, "none", "fp16", "bf16"])
+async def test_mlx_xavier_launch_and_cleanup(launch_runtime, pd, quantization):
     supervisor, workers, actors, destroy = launch_runtime
     kwargs = launch_kwargs()
-    kwargs.update(model_engine="MLX", model_format="mlx", xavier_cache_bytes=123456)
+    kwargs.update(
+        model_engine="MLX",
+        model_format="mlx",
+        quantization=quantization,
+        xavier_cache_bytes=123456,
+    )
     if not pd:
         kwargs["enable_xavier"] = True
         for replica in kwargs["replica_config"]:
@@ -99,6 +105,59 @@ async def test_mlx_xavier_launch_and_cleanup(launch_runtime, pd):
     assert destroy.await_count == (2 if pd else 1)
 
 
+@pytest.mark.asyncio
+async def test_single_mlx_replica_disables_shared_xavier(launch_runtime, caplog):
+    supervisor, workers, actors, _ = launch_runtime
+    supervisor._resolve_replica_config.return_value = ([(workers[0], [0], 1)], {0: "p"})
+    kwargs = launch_kwargs()
+    kwargs.update(model_engine="MLX", model_format="mlx", enable_xavier=True, replica=1)
+    kwargs["replica_config"] = kwargs["replica_config"][:1]
+    kwargs["replica_config"][0].role = None
+    await supervisor.launch_builtin_model(**kwargs)
+    assert "replica<=1" in caplog.text
+    assert not actors and not supervisor._xavier_cache_mapping
+    assert (
+        "_xavier_cache_config" not in workers[0].launch_builtin_model.call_args.kwargs
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("role", ["prefill", "decode"])
+async def test_mlx_pd_recovery_keeps_bytes_cache_and_registers_replacement(role):
+    import xoscar as xo
+
+    from ...model.llm.xavier.backends.bytes.storage import XavierBytesCacheActor
+    from ...model.llm.xavier.tests.test_bytes_storage import contract
+    from ..worker import WorkerActor
+
+    pool = await xo.create_actor_pool("127.0.0.1", n_process=0)
+    async with pool:
+        cache = await xo.create_actor(
+            XavierBytesCacheActor, address=pool.external_address, uid="bytes-cache"
+        )
+        namespace = await cache.configure(contract().to_dict())
+        worker = MagicMock()
+        supervisor = AsyncMock()
+        worker.get_supervisor_ref = AsyncMock(return_value=supervisor)
+        worker.launch_builtin_model = AsyncMock(return_value="replacement:1234")
+        worker.wait_for_load = AsyncMock()
+        replacement = MagicMock()
+        worker._model_uid_to_model = {"pd-rep0": replacement}
+        config = dict(role=role, address=cache.address, uid=cache.uid)
+        await WorkerActor.recover_model(
+            worker, dict(model_uid="pd-rep0", _xavier_cache_config=config)
+        )
+        worker.launch_builtin_model.assert_awaited_once_with(
+            model_uid="pd-rep0", _xavier_cache_config=config
+        )
+        worker.wait_for_load.assert_awaited_once_with("pd-rep0")
+        supervisor.unregister_pd_replica.assert_awaited_once_with("pd", "pd-rep0")
+        supervisor.register_pd_replica.assert_awaited_once_with(
+            "pd", "pd-rep0", replacement
+        )
+        assert (await cache.get_stats())["namespace"] == namespace
+
+
 @pytest.mark.asyncio
 async def test_mlx_failed_launch_cleans_bytes_actor(launch_runtime):
     supervisor, workers, actors, destroy = launch_runtime

xinference/core/worker.py

@@ -6756,7 +6756,10 @@ async def recover_model(self, launch_args: Dict[str, Any]):
         ).get("role") in ("prefill", "decode"):
             await supervisor_ref.unregister_pd_replica(origin_uid, rep_model_uid)
             cache_config = launch_args.get("_xavier_cache_config", {})
-            if cache_config.get("role") in ("prefill", "decode"):
+            if (
+                cache_config.get("role") in ("prefill", "decode")
+                and "rank" in cache_config
+            ):
                 directory = await xo.actor_ref(
                     address=cache_config["address"], uid=cache_config["uid"]
                 )

xinference/model/llm/mlx/core.py

@@ -226,6 +226,7 @@ def _get_or_create_generator(
                 "active": set(),  # active uids
                 "cache_boundaries": {},  # uid -> stable prefix token count
                 "xavier_prompts": {},
+                "xavier_writes": {},
                 "task": None,
             }
 
@@ -345,11 +346,16 @@ def _store_xavier_prompt_caches(self, batch_generator, prompt_results, gen_dict)
         for result in prompt_results:
             if not result.end_of_prompt or result.uid not in prompts:
                 continue
-            extracted = batch_generator.extract_cache([result.uid]).get(result.uid)
-            if extracted is not None:
-                cache, _ = extracted
-                tokens = prompts.pop(result.uid)
-                self._xavier.publish(cache, tokens[:-1])
+            tokens, cached_tokens = prompts.pop(result.uid)
+            try:
+                extracted = batch_generator.extract_cache([result.uid]).get(result.uid)
+                if extracted is not None:
+                    cache, _ = extracted
+                    gen_dict["xavier_writes"][result.uid] = self._xavier.publish(
+                        cache, tokens[:-1], cached_tokens
+                    )
+            except Exception:
+                logger.warning("MLX Xavier prefix publication failed", exc_info=True)
 
     async def _background_worker(self, gen_dict):
         """Background worker that continuously calls next() and distributes results."""
@@ -425,7 +431,6 @@ async def generate_stream(
         request_id: Optional[str] = None,
         skip_special_tokens: bool = True,
         prompt_cache_prefix_len: Optional[int] = None,
-        kv_transfer_params: Optional[dict] = None,
         prepared_cache=None,
         prompt_token_ids: Optional[List[int]] = None,
     ) -> AsyncGenerator[CompletionChunk, None]:
@@ -450,8 +455,6 @@ async def generate_stream(
         queue: asyncio.Queue = asyncio.Queue()
 
         external_cache = prepared_cache
-        if external_cache is None and getattr(self, "_xavier", None) is not None:
-            external_cache = await self._xavier.fetch(prompt_tokens, kv_transfer_params)
         insert_args = (
             batch_generator,
             prompt_tokens,
@@ -463,11 +466,15 @@ async def generate_stream(
             if external_cache is not None
             else self._insert_request(*insert_args)
         )
+        remote_cached_tokens = external_cache[1] if external_cache is not None else 0
         if (
             getattr(self, "_xavier", None) is not None
-            and cached_prompt_tokens < input_echo_len - 1
+            and remote_cached_tokens < input_echo_len - 1
         ):
-            gen_dict["xavier_prompts"][inserted_uid] = prompt_tokens
+            gen_dict["xavier_prompts"][inserted_uid] = (
+                prompt_tokens,
+                remote_cached_tokens,
+            )
         if cache_boundary is not None:
             gen_dict["cache_boundaries"][inserted_uid] = cache_boundary
 
@@ -606,7 +613,8 @@ async def generate_stream(
                     gen_dict["xavier_prompts"].pop(inserted_uid, None)
                     batch_generator.remove([inserted_uid])
             if getattr(self, "_xavier", None) is not None:
-                await self._xavier.flush()
+                write = gen_dict["xavier_writes"].pop(inserted_uid, None)
+                await self._xavier.flush([write] if write is not None else [])
 
     async def generate(
         self,
@@ -618,7 +626,6 @@ async def generate(
         stream: bool = False,
         skip_special_tokens: bool = True,
         prompt_cache_prefix_len: Optional[int] = None,
-        kv_transfer_params: Optional[dict] = None,
         prepared_cache=None,
         prompt_token_ids: Optional[List[int]] = None,
     ) -> Tuple[str, CompletionUsage]:
@@ -640,7 +647,6 @@ async def generate(
             request_id=None,
             skip_special_tokens=skip_special_tokens,
             prompt_cache_prefix_len=prompt_cache_prefix_len,
-            kv_transfer_params=kv_transfer_params,
             prepared_cache=prepared_cache,
             prompt_token_ids=prompt_token_ids,
         ):
@@ -1119,6 +1125,10 @@ def wait_for_load(self):
                 raise ValueError(
                     "MLX Xavier requires the continuous batching text engine"
                 )
+            assert self._loop is not None, "Service not started correctly"
+            asyncio.run_coroutine_threadsafe(
+                self._xavier.initialize(), self._loop
+            ).result()
 
         # Update allow_batch based on distributed inference
         # Only enable continuous batching for non-distributed inference (single worker)