PERF: reduce Xavier V1 cache overhead and allow CUDA Graph

.github/workflows/python.yaml

@@ -376,6 +376,7 @@ jobs:
               --ignore xinference/model/llm/mlx \
               xinference
             run_pytest \
+              -vv \
               benchmark/tests/test_benchmark_pd.py \
               benchmark/tests/test_benchmark_sglang_pd.py \
               benchmark/tests/test_analyze_pd_profile.py \
@@ -394,6 +395,12 @@ jobs:
               xinference/model/llm/vllm/xavier/test/test_block_tracker_reuse.py \
               xinference/model/llm/vllm/xavier/test/test_remote_kvcache_manager.py \
               xinference/model/llm/vllm/xavier/test/test_v1_engine.py \
+              xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py \
+              xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py \
+              xinference/model/llm/vllm/xavier/test/test_snapshot_gc_lifecycle.py \
+              xinference/model/llm/vllm/xavier/test/test_packed_gather.py \
+              xinference/model/llm/vllm/xavier/test/test_local_directory.py \
+              xinference/model/llm/vllm/xavier/test/test_local_read.py \
               xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py \
               xinference/model/llm/vllm/xavier/test/test_kv_schema.py \
               xinference/model/llm/vllm/xavier/test/test_cross_engine.py \

doc/source/locale/de/LC_MESSAGES/user_guide/vllm_enhancement.po

@@ -130,3 +130,7 @@ msgstr ""
 "wiederverwendete Token unter ``prompt_tokens_details.cached_tokens``. "
 "Explizite Prefill-/Decode-Rollen aktivieren MLX Xavier P/D automatisch. "
 "Engineübergreifendes Cache-Sharing wird noch nicht unterstützt."
+
+#: source/user_guide/vllm_enhancement.rst:18
+msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default."
+msgstr "Bei Modellen mit ausschließlich Attention unter vLLM V1 behält Xavier die Einstellung ``enforce_eager`` des Modells bei und ermöglicht damit die Ausführung mit CUDA Graph. Setzen Sie beim Modellstart ``enforce_eager=True``, um die Eager-Ausführung zu verwenden. Rekurrente Modelle verwenden weiterhin standardmäßig die Eager-Ausführung."

doc/source/locale/es/LC_MESSAGES/user_guide/vllm_enhancement.po

@@ -131,3 +131,7 @@ msgstr ""
 "``prompt_tokens_details.cached_tokens``. Los roles prefill/decode explícitos"
 " activan MLX Xavier P/D automáticamente. Aún no se admite compartir caché "
 "entre motores."
+
+#: source/user_guide/vllm_enhancement.rst:18
+msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default."
+msgstr "Para los modelos que solo usan atención en vLLM V1, Xavier conserva la configuración ``enforce_eager`` del modelo, lo que permite la ejecución con CUDA Graph. Establezca ``enforce_eager=True`` al iniciar el modelo para usar la ejecución eager. Los modelos recurrentes siguen usando la ejecución eager de forma predeterminada."

doc/source/locale/fr/LC_MESSAGES/user_guide/vllm_enhancement.po

@@ -134,3 +134,7 @@ msgstr ""
 "``prompt_tokens_details.cached_tokens``. Les rôles prefill/decode explicites"
 " activent automatiquement MLX Xavier P/D. Le partage de cache entre moteurs "
 "n’est pas encore pris en charge."
+
+#: source/user_guide/vllm_enhancement.rst:18
+msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default."
+msgstr "Pour les modèles utilisant uniquement des couches d’attention sur vLLM V1, Xavier conserve le paramètre ``enforce_eager`` du modèle, permettant l’exécution avec CUDA Graph. Définissez ``enforce_eager=True`` au lancement du modèle pour utiliser l’exécution eager. Les modèles récurrents continuent d’utiliser l’exécution eager par défaut."

doc/source/locale/it/LC_MESSAGES/user_guide/vllm_enhancement.po

@@ -131,3 +131,7 @@ msgstr ""
 "``prompt_tokens_details.cached_tokens``. I ruoli prefill/decode espliciti "
 "attivano automaticamente MLX Xavier P/D. La condivisione della cache tra "
 "motori non è ancora supportata."
+
+#: source/user_guide/vllm_enhancement.rst:18
+msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default."
+msgstr "Per i modelli che usano solo attention su vLLM V1, Xavier mantiene l’impostazione ``enforce_eager`` del modello, consentendo l’esecuzione con CUDA Graph. Impostare ``enforce_eager=True`` all’avvio del modello per usare l’esecuzione eager. I modelli ricorrenti continuano a usare l’esecuzione eager per impostazione predefinita."

doc/source/locale/ja/LC_MESSAGES/user_guide/vllm_enhancement.po

@@ -118,3 +118,7 @@ msgstr ""
 "キャッシュも保持します。通常レプリカはミスやストレージ利用不可時にローカル計算し、互換性のないデプロイ間ではページを共有できません。使用量の "
 "``prompt_tokens_details.cached_tokens`` に再利用トークン数を返します。明示的な prefill/decode "
 "ロールは MLX Xavier P/D を自動で有効にします。異なるエンジン間のキャッシュ共有は未対応です。"
+
+#: source/user_guide/vllm_enhancement.rst:18
+msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default."
+msgstr "vLLM V1 の attention のみを使用するモデルでは、Xavier はモデルの ``enforce_eager`` 設定を維持し、CUDA Graph による実行を可能にします。eager 実行を使用するには、モデル起動時に ``enforce_eager=True`` を設定してください。再帰型モデルでは、引き続きデフォルトで eager 実行を使用します。"

doc/source/locale/ko/LC_MESSAGES/user_guide/vllm_enhancement.po

@@ -122,3 +122,7 @@ msgstr ""
 "페이지를 공유할 수 없습니다. 사용량의 ``prompt_tokens_details.cached_tokens`` 는 재사용된 토큰 수를 "
 "보고합니다. 명시적인 prefill/decode 역할은 MLX Xavier P/D를 자동 활성화합니다. 엔진 간 캐시 공유는 아직 "
 "지원하지 않습니다."
+
+#: source/user_guide/vllm_enhancement.rst:18
+msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default."
+msgstr "vLLM V1에서 attention만 사용하는 모델의 경우 Xavier는 모델의 ``enforce_eager`` 설정을 유지하여 CUDA Graph 실행을 허용합니다. eager 실행을 사용하려면 모델을 시작할 때 ``enforce_eager=True`` 설정을 사용하세요. 순환 모델은 계속해서 기본적으로 eager 실행을 사용합니다."

doc/source/locale/pt_BR/LC_MESSAGES/user_guide/vllm_enhancement.po

@@ -129,3 +129,7 @@ msgstr ""
 "tokens reutilizados em ``prompt_tokens_details.cached_tokens``. Papéis "
 "prefill/decode explícitos ativam MLX Xavier P/D automaticamente. "
 "Compartilhamento de cache entre engines ainda não é suportado."
+
+#: source/user_guide/vllm_enhancement.rst:18
+msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default."
+msgstr "Para modelos que usam apenas atenção no vLLM V1, o Xavier preserva a configuração ``enforce_eager`` do modelo, permitindo a execução com CUDA Graph. Defina ``enforce_eager=True`` ao iniciar o modelo para usar a execução eager. Os modelos recorrentes continuam usando a execução eager por padrão."

doc/source/locale/zh_CN/LC_MESSAGES/user_guide/vllm_enhancement.po

@@ -119,3 +119,7 @@ msgstr ""
 "缓存。普通副本在未命中或存储不可用时本地计算；不兼容的部署不能共享缓存页。响应通过 "
 "``prompt_tokens_details.cached_tokens`` 报告复用 token 数。显式设置 prefill/decode "
 "角色会自动启用 MLX Xavier P/D。目前不支持跨引擎缓存共享。"
+
+#: source/user_guide/vllm_enhancement.rst:18
+msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default."
+msgstr "对于 vLLM V1 上仅使用注意力层的模型，Xavier 保留模型的 ``enforce_eager`` 设置，允许使用 CUDA Graph 执行。若要使用 eager 执行，请在启动模型时设置 ``enforce_eager=True``。循环模型仍默认使用 eager 执行。"

doc/source/locale/zh_TW/LC_MESSAGES/user_guide/vllm_enhancement.po

@@ -118,3 +118,7 @@ msgstr ""
 "快取。普通副本在未命中或儲存不可用時於本機計算；不相容的部署不能共用快取頁面。回應透過 "
 "``prompt_tokens_details.cached_tokens`` 報告重用 token 數。明確設定 prefill/decode "
 "角色會自動啟用 MLX Xavier P/D。目前不支援跨引擎快取共用。"
+
+#: source/user_guide/vllm_enhancement.rst:18
+msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default."
+msgstr "對於 vLLM V1 上僅使用注意力層的模型，Xavier 保留模型的 ``enforce_eager`` 設定，允許使用 CUDA Graph 執行。若要使用 eager 執行，請在啟動模型時設定 ``enforce_eager=True``。循環模型仍預設使用 eager 執行。"

doc/source/user_guide/vllm_enhancement.rst

@@ -15,6 +15,8 @@ Usage
 *****
 Simply add the parameter ``enable_xavier=True`` when starting the vllm model.
 
+For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default.
+
 ***********
 Limitations
 ***********

xinference/core/model.py

@@ -399,6 +399,10 @@ async def start_transfer_for_vllm(self, rank_addresses: List[str]):
             store_address=self._xavier_config.get("store_address"),  # type: ignore
             store_port=self._xavier_config.get("store_port"),  # type: ignore
             world_addresses=rank_addresses,
+            freeze_initialization_gc=(
+                self._model._is_vllm_v1()
+                and self._xavier_config.get("gpu_cache_bytes") is None
+            ),
         )
         await self._model.init_xavier()
         logger.debug(

xinference/model/llm/sglang/tests/test_xavier_pd.py

@@ -562,6 +562,61 @@ async def test_decode_renews_completed_lease_and_stops_on_release(
     assert (await directory.get_stats())["active_handoffs"] == 0
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failed", [False, True])
+async def test_release_finishes_when_renewal_rpc_swallows_cancellation(
+    monkeypatch, failed
+):
+    from ..xavier.settings import TRANSFER_TIMEOUT_ENV
+
+    monkeypatch.setenv(TRANSFER_TIMEOUT_ENV, "0.03")
+    handoff = fingerprint_handoff()
+    entered = asyncio.Event()
+
+    async def renew(room):
+        entered.set()
+        try:
+            await asyncio.Future()
+        except asyncio.CancelledError:
+            # Model Python 3.10 wait_for returning the completed RPC's result
+            # or exception instead of propagating the caller's cancellation.
+            if failed:
+                raise RuntimeError("RPC completed while cancelling")
+            return True
+
+    handoff._directory_actor.renew_completed = AsyncMock(side_effect=renew)
+    handoff._directory_actor.release = AsyncMock()
+    original_call = handoff._call
+
+    async def call(method, *args, **kwargs):
+        if method == "renew_completed":
+            return await handoff._directory_actor.renew_completed(*args)
+        return await original_call(method, *args, **kwargs)
+
+    monkeypatch.setattr(handoff, "_call", call)
+    accepted = await handoff.accept(
+        "prompt", config("decode")["_pd_kv_transfer_params"]
+    )
+    task = handoff._lease_tasks[123]
+    releasing = None
+    try:
+        await asyncio.wait_for(entered.wait(), 1)
+        releasing = asyncio.create_task(handoff.release(accepted))
+        done, _ = await asyncio.wait([releasing], timeout=1)
+        assert releasing in done
+        await releasing
+        assert task.done() and not handoff._lease_tasks
+        handoff._directory_actor.renew_completed.assert_awaited_once_with(123)
+        handoff._directory_actor.release.assert_awaited_once_with(123, "decode")
+    finally:
+        handoff._directory_actor.renew_completed.side_effect = None
+        handoff._directory_actor.renew_completed.return_value = False
+        task.cancel()
+        await asyncio.gather(task, return_exceptions=True)
+        if releasing is not None:
+            await asyncio.gather(releasing, return_exceptions=True)
+
+
 @pytest.mark.asyncio
 async def test_unconsumed_stream_reserves_no_request_or_directory_room(deployment):
     directory, _, decode = deployment

xinference/model/llm/sglang/xavier/pd.py

@@ -79,7 +79,9 @@ async def accept(self, prompt: str, transfer: dict) -> dict:
     async def _renew_completed(self, room):
         # Completed records are needed by non-streaming decode after generation.
         # Keep them while D is alive; lost cleanup/router processes remain bounded.
-        while True:
+        # release() removes the owner before cancelling. Python 3.10 wait_for
+        # can swallow that cancellation when the RPC has just completed.
+        while self._lease_tasks.get(room) is asyncio.current_task():
             await asyncio.sleep(min(60, transfer_timeout() / 3))
             try:
                 if not await self._call("renew_completed", room):

xinference/model/llm/vllm/xavier/engine.py

@@ -30,24 +30,10 @@
 
 logger = logging.getLogger(__name__)
 
-XAVIER_EAGER_VLLM_VERSION = version.parse("0.21.0")
-
 
 class XavierEngine:
     _xavier_config: Optional[Dict] = None
 
-    @staticmethod
-    def _json_safe(value: Any) -> Any:
-        if isinstance(value, bytes):
-            return value.decode(errors="replace")
-        if isinstance(value, dict):
-            return {str(k): XavierEngine._json_safe(v) for k, v in value.items()}
-        if isinstance(value, (list, tuple)):
-            return [XavierEngine._json_safe(v) for v in value]
-        if isinstance(value, (str, int, float, bool)) or value is None:
-            return value
-        return str(value)
-
     @classmethod
     def _patch_v1_engine_args(
         cls, engine_args: AsyncEngineArgs, xavier_config: Optional[Dict]
@@ -69,13 +55,6 @@ def _patch_v1_engine_args(
         xavier_config = dict(xavier_config)
         set_xavier_transport_backend(xavier_config, XAVIER_TRANSPORT_XAVIER)
 
-        additional_config = dict(getattr(engine_args, "additional_config", {}) or {})
-        # vLLM V1 includes additional_config in its compilation cache hash and
-        # serializes it as JSON. Keep the connector's full config below, but
-        # only expose a JSON-safe copy here.
-        additional_config["xavier_config"] = cls._json_safe(xavier_config)
-        engine_args.additional_config = additional_config
-
         role = xavier_config.get("role")
         if role == "prefill":
             kv_role = "kv_producer"
@@ -88,16 +67,17 @@ def _patch_v1_engine_args(
             kv_rank = xavier_config.get("rank", 0)
 
         extra_config = dict(xavier_config.get("kv_connector_extra_config") or {})
+        # Actor addresses, IDs and ranks affect transport, not compiled model
+        # operations. Keep them in the connector config: additional_config is
+        # included in vLLM's graph hash and would recompile on every launch.
         extra_config["xavier_config"] = xavier_config
 
-        if (
-            version.parse(VLLM_VERSION) >= XAVIER_EAGER_VLLM_VERSION
-            and xavier_config.get("enforce_eager", True)
-            and not getattr(engine_args, "enforce_eager", False)
+        recurrent = engine_args.create_model_config().is_hybrid
+        if xavier_config.get("enforce_eager", recurrent) and not getattr(
+            engine_args, "enforce_eager", False
         ):
-            # vLLM V1 may create XavierConnector twice during CUDA graph
-            # setup. In CUDA-heavy processes this can trip glibc static TLS
-            # allocation, so keep Xavier on the eager execution path.
+            # Attention snapshots are exported after forward and support CUDA
+            # graphs. Keep the existing eager default for recurrent handoff.
             engine_args.enforce_eager = True
             logger.info(
                 "Set enforce_eager=True for Xavier V1 on vLLM %s.",
@@ -108,7 +88,6 @@ def _patch_v1_engine_args(
         # of sliding-window and full attention. Only recurrent models need this
         # override: 0.21 disables HMA for every KV connector by default; 0.22+
         # checks SupportsHMA itself. Preserve an explicit user setting.
-        recurrent = engine_args.create_model_config().is_hybrid
         if (
             recurrent
             and version.parse(VLLM_VERSION) < version.parse("0.22.0")