diff --git a/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po index 05c7755774..3ab398b385 100644 --- a/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po @@ -149,14 +149,8 @@ msgid "Hybrid/recurrent attention limitation" msgstr "" #: source/user_guide/pd_separation.rst:136 -msgid "" -"Xavier PD does not support hybrid/recurrent attention models such as " -"Qwen3.5. Deploy these models without PD, or choose a full-attention model " -"such as Qwen3 for PD." -msgstr "" -"Xavier PD unterstützt keine hybriden/rekurrenten Aufmerksamkeitsmodelle wie " -"Qwen3.5. Stellen Sie diese Modelle ohne PD bereit oder wählen Sie für PD ein" -" Modell mit vollständiger Aufmerksamkeit wie Qwen3." +msgid "Qwen3.5 text requests are supported with Xavier and native NIXL. Set ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` and ``async_scheduling=False`` on both replicas. Use TP=1 and PP=1. Xavier history caching is not available for these hybrid/recurrent attention models; their configured history budget is not allocated. Image/video inputs and speculative decoding are not supported in this configuration. Native NIXL requires vLLM >= 0.22.0 for this mode." +msgstr "Xavier und natives NIXL unterstützen Qwen3.5-Textanfragen. Setzen Sie auf beiden Replikaten ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` und ``async_scheduling=False``. Verwenden Sie TP=1 und PP=1. Der Xavier-Verlaufscache ist für diese Modelle mit hybrider/rekurrenter Attention nicht verfügbar; sein konfiguriertes Budget wird nicht reserviert. Bild-/Videoeingaben und spekulatives Decoding werden in dieser Konfiguration nicht unterstützt. Natives NIXL erfordert für diesen Modus vLLM >= 0.22.0." #: source/user_guide/pd_separation.rst:143 msgid "Xavier requirements and memory" diff --git a/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po index 45c93832d5..9247367f16 100644 --- a/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po @@ -149,14 +149,8 @@ msgid "Hybrid/recurrent attention limitation" msgstr "" #: source/user_guide/pd_separation.rst:136 -msgid "" -"Xavier PD does not support hybrid/recurrent attention models such as " -"Qwen3.5. Deploy these models without PD, or choose a full-attention model " -"such as Qwen3 for PD." -msgstr "" -"Xavier PD no admite modelos de atención híbrida/recurrente como Qwen3.5. " -"Despliegue estos modelos sin PD o elija un modelo de atención completa como " -"Qwen3 para PD." +msgid "Qwen3.5 text requests are supported with Xavier and native NIXL. Set ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` and ``async_scheduling=False`` on both replicas. Use TP=1 and PP=1. Xavier history caching is not available for these hybrid/recurrent attention models; their configured history budget is not allocated. Image/video inputs and speculative decoding are not supported in this configuration. Native NIXL requires vLLM >= 0.22.0 for this mode." +msgstr "Xavier y NIXL nativo admiten solicitudes de texto de Qwen3.5. Configure ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` y ``async_scheduling=False`` en ambas réplicas. Use TP=1 y PP=1. La caché histórica de Xavier no está disponible para estos modelos de atención híbrida/recurrente; no se asigna su presupuesto configurado. Esta configuración no admite imágenes, vídeos ni decodificación especulativa. NIXL nativo requiere vLLM >= 0.22.0 para este modo." #: source/user_guide/pd_separation.rst:143 msgid "Xavier requirements and memory" diff --git a/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po index 369a021f39..cd1bf921ca 100644 --- a/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po @@ -149,14 +149,8 @@ msgid "Hybrid/recurrent attention limitation" msgstr "" #: source/user_guide/pd_separation.rst:136 -msgid "" -"Xavier PD does not support hybrid/recurrent attention models such as " -"Qwen3.5. Deploy these models without PD, or choose a full-attention model " -"such as Qwen3 for PD." -msgstr "" -"Xavier PD ne prend pas en charge les modèles à attention hybride/récurrente " -"comme Qwen3.5. Déployez ces modèles sans PD ou choisissez un modèle à " -"attention complète comme Qwen3 pour le PD." +msgid "Qwen3.5 text requests are supported with Xavier and native NIXL. Set ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` and ``async_scheduling=False`` on both replicas. Use TP=1 and PP=1. Xavier history caching is not available for these hybrid/recurrent attention models; their configured history budget is not allocated. Image/video inputs and speculative decoding are not supported in this configuration. Native NIXL requires vLLM >= 0.22.0 for this mode." +msgstr "Xavier et NIXL natif prennent en charge les requêtes textuelles Qwen3.5. Définissez ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` et ``async_scheduling=False`` sur les deux réplicas. Utilisez TP=1 et PP=1. Le cache historique Xavier est indisponible pour ces modèles à attention hybride/récurrente ; son budget configuré n’est pas alloué. Les entrées image/vidéo et le décodage spéculatif ne sont pas pris en charge dans cette configuration. Dans ce mode, NIXL natif nécessite vLLM >= 0.22.0." #: source/user_guide/pd_separation.rst:143 msgid "Xavier requirements and memory" diff --git a/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po index 567aeab68c..ee72c84e6d 100644 --- a/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po @@ -149,14 +149,8 @@ msgid "Hybrid/recurrent attention limitation" msgstr "" #: source/user_guide/pd_separation.rst:136 -msgid "" -"Xavier PD does not support hybrid/recurrent attention models such as " -"Qwen3.5. Deploy these models without PD, or choose a full-attention model " -"such as Qwen3 for PD." -msgstr "" -"Xavier PD non supporta modelli con attenzione ibrida/ricorrente come " -"Qwen3.5. Distribuisci questi modelli senza PD oppure scegli un modello con " -"attenzione completa come Qwen3 per il PD." +msgid "Qwen3.5 text requests are supported with Xavier and native NIXL. Set ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` and ``async_scheduling=False`` on both replicas. Use TP=1 and PP=1. Xavier history caching is not available for these hybrid/recurrent attention models; their configured history budget is not allocated. Image/video inputs and speculative decoding are not supported in this configuration. Native NIXL requires vLLM >= 0.22.0 for this mode." +msgstr "Xavier e NIXL nativo supportano le richieste di testo Qwen3.5. Imposta ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` e ``async_scheduling=False`` su entrambe le repliche. Usa TP=1 e PP=1. La cache storica Xavier non è disponibile per questi modelli con attenzione ibrida/ricorrente; il budget configurato non viene allocato. Questa configurazione non supporta immagini, video o decodifica speculativa. In questa modalità NIXL nativo richiede vLLM >= 0.22.0." #: source/user_guide/pd_separation.rst:143 msgid "Xavier requirements and memory" diff --git a/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po index 0bf02ed39b..86a78ef920 100644 --- a/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po @@ -143,13 +143,8 @@ msgid "Hybrid/recurrent attention limitation" msgstr "" #: source/user_guide/pd_separation.rst:136 -msgid "" -"Xavier PD does not support hybrid/recurrent attention models such as " -"Qwen3.5. Deploy these models without PD, or choose a full-attention model " -"such as Qwen3 for PD." -msgstr "" -"Xavier PD は Qwen3.5 などのハイブリッド／再帰型アテンションモデルをサポートしません。これらは PD なしでデプロイするか、PD には" -" Qwen3 などのフルアテンションモデルを選択してください。" +msgid "Qwen3.5 text requests are supported with Xavier and native NIXL. Set ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` and ``async_scheduling=False`` on both replicas. Use TP=1 and PP=1. Xavier history caching is not available for these hybrid/recurrent attention models; their configured history budget is not allocated. Image/video inputs and speculative decoding are not supported in this configuration. Native NIXL requires vLLM >= 0.22.0 for this mode." +msgstr "Xavier とネイティブ NIXL は Qwen3.5 のテキストリクエストに対応しています。両方のレプリカで ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` および ``async_scheduling=False`` を設定し、TP=1、PP=1 を使用してください。これらのハイブリッド／再帰型アテンションモデルでは Xavier の履歴キャッシュは利用できず、設定された履歴キャッシュ容量は確保されません。この構成は画像・動画入力と投機的デコードに対応していません。 このモードでネイティブ NIXL を使うには vLLM >= 0.22.0 が必要です。" #: source/user_guide/pd_separation.rst:143 msgid "Xavier requirements and memory" diff --git a/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po index ccad95fec5..2e4a245fe7 100644 --- a/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po @@ -143,13 +143,8 @@ msgid "Hybrid/recurrent attention limitation" msgstr "" #: source/user_guide/pd_separation.rst:136 -msgid "" -"Xavier PD does not support hybrid/recurrent attention models such as " -"Qwen3.5. Deploy these models without PD, or choose a full-attention model " -"such as Qwen3 for PD." -msgstr "" -"Xavier PD는 Qwen3.5 같은 하이브리드/순환 어텐션 모델을 지원하지 않습니다. 이러한 모델은 PD 없이 배포하거나, PD에는 " -"Qwen3 같은 전체 어텐션 모델을 선택하세요." +msgid "Qwen3.5 text requests are supported with Xavier and native NIXL. Set ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` and ``async_scheduling=False`` on both replicas. Use TP=1 and PP=1. Xavier history caching is not available for these hybrid/recurrent attention models; their configured history budget is not allocated. Image/video inputs and speculative decoding are not supported in this configuration. Native NIXL requires vLLM >= 0.22.0 for this mode." +msgstr "Xavier와 네이티브 NIXL은 Qwen3.5 텍스트 요청을 지원합니다. 두 복제본 모두에서 ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` 및 ``async_scheduling=False`` 를 설정하고 TP=1, PP=1을 사용하세요. 이러한 하이브리드/순환 어텐션 모델에서는 Xavier 기록 캐시를 사용할 수 없으므로 설정된 기록 캐시 용량은 할당되지 않습니다. 이 구성은 이미지·동영상 입력과 추측 디코딩을 지원하지 않습니다. 이 모드에서 네이티브 NIXL을 사용하려면 vLLM >= 0.22.0이 필요합니다." #: source/user_guide/pd_separation.rst:143 msgid "Xavier requirements and memory" diff --git a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po index c04850c3bf..7b43a4d962 100644 --- a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po @@ -149,14 +149,8 @@ msgid "Hybrid/recurrent attention limitation" msgstr "" #: source/user_guide/pd_separation.rst:136 -msgid "" -"Xavier PD does not support hybrid/recurrent attention models such as " -"Qwen3.5. Deploy these models without PD, or choose a full-attention model " -"such as Qwen3 for PD." -msgstr "" -"O Xavier PD não suporta modelos de atenção híbrida/recorrente como Qwen3.5. " -"Implante esses modelos sem PD ou escolha um modelo de atenção completa como " -"Qwen3 para PD." +msgid "Qwen3.5 text requests are supported with Xavier and native NIXL. Set ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` and ``async_scheduling=False`` on both replicas. Use TP=1 and PP=1. Xavier history caching is not available for these hybrid/recurrent attention models; their configured history budget is not allocated. Image/video inputs and speculative decoding are not supported in this configuration. Native NIXL requires vLLM >= 0.22.0 for this mode." +msgstr "Xavier e NIXL nativo suportam solicitações de texto do Qwen3.5. Defina ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` e ``async_scheduling=False`` nas duas réplicas. Use TP=1 e PP=1. O cache histórico do Xavier não está disponível para esses modelos de atenção híbrida/recorrente; seu orçamento configurado não é alocado. Esta configuração não suporta imagens, vídeos ou decodificação especulativa. Nesse modo, o NIXL nativo requer vLLM >= 0.22.0." #: source/user_guide/pd_separation.rst:143 msgid "Xavier requirements and memory" diff --git a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po index beb4a2832a..8c7f26582d 100644 --- a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po @@ -148,13 +148,8 @@ msgid "Hybrid/recurrent attention limitation" msgstr "混合注意力与循环注意力的限制" #: source/user_guide/pd_separation.rst:136 -msgid "" -"Xavier PD does not support hybrid/recurrent attention models such as " -"Qwen3.5. Deploy these models without PD, or choose a full-attention model " -"such as Qwen3 for PD." -msgstr "" -"Xavier PD 不支持 Qwen3.5 等混合或循环注意力模型。请在不启用 PD 的情况下部署这些模型，或选择 Qwen3 等全注意力模型使用 " -"PD。" +msgid "Qwen3.5 text requests are supported with Xavier and native NIXL. Set ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` and ``async_scheduling=False`` on both replicas. Use TP=1 and PP=1. Xavier history caching is not available for these hybrid/recurrent attention models; their configured history budget is not allocated. Image/video inputs and speculative decoding are not supported in this configuration. Native NIXL requires vLLM >= 0.22.0 for this mode." +msgstr "Xavier 和原生 NIXL 支持 Qwen3.5 纯文本请求。两个副本均需设置 ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` 和 ``async_scheduling=False``，并使用 TP=1、PP=1。此类混合或循环注意力模型暂不支持 Xavier 历史缓存，因此不会分配配置的历史缓存预算。该配置不支持图片、视频输入或推测解码。 此模式下原生 NIXL 需要 vLLM >= 0.22.0。" #: source/user_guide/pd_separation.rst:143 msgid "Xavier requirements and memory" diff --git a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po index 3c36be2e2d..ae7c900518 100644 --- a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po @@ -143,13 +143,8 @@ msgid "Hybrid/recurrent attention limitation" msgstr "" #: source/user_guide/pd_separation.rst:136 -msgid "" -"Xavier PD does not support hybrid/recurrent attention models such as " -"Qwen3.5. Deploy these models without PD, or choose a full-attention model " -"such as Qwen3 for PD." -msgstr "" -"Xavier PD 不支援 Qwen3.5 等混合或循環注意力模型。請在不啟用 PD 的情況下部署這些模型，或選擇 Qwen3 等全注意力模型使用 " -"PD。" +msgid "Qwen3.5 text requests are supported with Xavier and native NIXL. Set ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` and ``async_scheduling=False`` on both replicas. Use TP=1 and PP=1. Xavier history caching is not available for these hybrid/recurrent attention models; their configured history budget is not allocated. Image/video inputs and speculative decoding are not supported in this configuration. Native NIXL requires vLLM >= 0.22.0 for this mode." +msgstr "Xavier 和原生 NIXL 支援 Qwen3.5 純文字請求。兩個副本均需設定 ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode=\"none\"``, ``disable_hybrid_kv_cache_manager=False`` 和 ``async_scheduling=False``，並使用 TP=1、PP=1。此類混合或循環注意力模型暫不支援 Xavier 歷史快取，因此不會配置設定的歷史快取預算。此設定不支援圖片、影片輸入或推測解碼。 此模式下原生 NIXL 需要 vLLM >= 0.22.0。" #: source/user_guide/pd_separation.rst:143 msgid "Xavier requirements and memory" diff --git a/doc/source/user_guide/pd_separation.rst b/doc/source/user_guide/pd_separation.rst index f335fffd6d..316548d907 100644 --- a/doc/source/user_guide/pd_separation.rst +++ b/doc/source/user_guide/pd_separation.rst @@ -120,7 +120,7 @@ After a worker restart, relaunch the native NIXL PD deployment. Hybrid/recurrent attention limitation ------------------------------------- -Xavier PD does not support hybrid/recurrent attention models such as Qwen3.5. Deploy these models without PD, or choose a full-attention model such as Qwen3 for PD. +Qwen3.5 text requests are supported with Xavier and native NIXL. Set ``language_model_only=True``, ``enable_prefix_caching=False``, ``mamba_cache_mode="none"``, ``disable_hybrid_kv_cache_manager=False`` and ``async_scheduling=False`` on both replicas. Use TP=1 and PP=1. Xavier history caching is not available for these hybrid/recurrent attention models; their configured history budget is not allocated. Image/video inputs and speculative decoding are not supported in this configuration. Native NIXL requires vLLM >= 0.22.0 for this mode. Xavier requirements and memory ------------------------------ diff --git a/xinference/model/llm/vllm/core.py b/xinference/model/llm/vllm/core.py index 2af1782380..2939543cd8 100644 --- a/xinference/model/llm/vllm/core.py +++ b/xinference/model/llm/vllm/core.py @@ -796,15 +796,23 @@ def load(self): if self._nixl_config is not None: from .pd import configure_nixl_engine - if isinstance(self, VLLMMultiModel): - raise ValueError("Native NIXL PD currently supports text-only models") + if isinstance(self, VLLMMultiModel) and not self._model_config.get( + "language_model_only", False + ): + raise ValueError( + "Native NIXL PD requires text-only models or language_model_only=True" + ) + if isinstance(self, VLLMMultiModel) and VLLM_VERSION < version.parse( + "0.22.0" + ): + raise ValueError( + "Native NIXL PD for language_model_only models requires vLLM >= 0.22.0" + ) configure_nixl_engine(self._model_config, VLLM_VERSION, enable_lora) use_native_mp, native_mp_reason = self._native_mp_route() if self._xavier_config is not None: from .xavier.engine import XavierEngine - # Enabling Xavier means that `enable_prefix_caching` is enabled by default. - self._model_config.setdefault("enable_prefix_caching", True) xavier_transfer_block_num = self._model_config.pop( "xavier_transfer_block_num", 512 ) @@ -2951,7 +2959,13 @@ async def async_chat( ) assert not isinstance(c, AsyncGenerator) if tools: - return self._post_process_completion( + result = self._post_process_completion( self.model_family, self.model_uid, c ) - return self._to_chat_completion(c, self.reasoning_parser) + else: + result = self._to_chat_completion(c, self.reasoning_parser) + if "_pd_kv_transfer_params" in c: + result["_pd_kv_transfer_params"] = cast(Dict[str, Any], c)[ + "_pd_kv_transfer_params" + ] + return result diff --git a/xinference/model/llm/vllm/pd.py b/xinference/model/llm/vllm/pd.py index 3c15da1335..e4b904b0c1 100644 --- a/xinference/model/llm/vllm/pd.py +++ b/xinference/model/llm/vllm/pd.py @@ -41,6 +41,12 @@ def resolve_nixl_host(host: str, host_key: str = "VLLM_NIXL_SIDE_CHANNEL_HOST") def configure_nixl_environment(env: Dict[str, str], worker_address: str) -> None: from xoscar.utils import get_next_port + # Native NIXL's recurrent-state transfer requires dimension-first conv + # states. Select the layout before vLLM initializes model/cache classes. + env.setdefault( + "VLLM_SSM_CONV_STATE_LAYOUT", + os.environ.get("VLLM_SSM_CONV_STATE_LAYOUT", "DS"), + ) # Each model subprocess needs its own listener, including same-host replicas. # Reallocate on recovery as the old listener may not have exited yet. host_key = "VLLM_NIXL_SIDE_CHANNEL_HOST" diff --git a/xinference/model/llm/vllm/tests/test_pd.py b/xinference/model/llm/vllm/tests/test_pd.py index 3a9b6aff7e..60f5339ec2 100644 --- a/xinference/model/llm/vllm/tests/test_pd.py +++ b/xinference/model/llm/vllm/tests/test_pd.py @@ -323,9 +323,12 @@ async def generate(self, *args, **kwargs): model._engine.abort.assert_awaited_once_with("failed") -@pytest.mark.parametrize("multimodal", [False, True]) +@pytest.mark.parametrize( + "multimodal,language_only", [(False, False), (True, False), (True, True)] +) +@pytest.mark.parametrize("engine_version", ["0.21.0", "0.22.0"]) def test_nixl_load_configures_engine_args_or_rejects_multimodal( - monkeypatch, multimodal + monkeypatch, multimodal, language_only, engine_version ): from .. import core @@ -338,7 +341,7 @@ def engine_args(**kwargs): args_factory = Mock(side_effect=engine_args) for name, module in { - "vllm": SimpleNamespace(__version__="0.21.0"), + "vllm": SimpleNamespace(__version__=engine_version), "vllm.engine.arg_utils": SimpleNamespace(AsyncEngineArgs=args_factory), "vllm.engine.async_llm_engine": SimpleNamespace(AsyncLLMEngine=object), "vllm.lora.request": SimpleNamespace(LoRARequest=object), @@ -360,15 +363,61 @@ def engine_args(**kwargs): model._xavier_config = None model.lora_modules = None model._get_cuda_count = lambda: 1 - model._sanitize_model_config = lambda config: {"reasoning_content": False} + model._sanitize_model_config = lambda config: { + "reasoning_content": False, + "language_model_only": language_only, + } model.prepare_parse_reasoning_content = Mock() model.prepare_parse_tool_calls = Mock() model._native_mp_route = lambda: (False, "test") - if multimodal: + if multimodal and not language_only: with pytest.raises(ValueError, match="text-only"): model.load() args_factory.assert_not_called() + elif multimodal and Version(engine_version) < Version("0.22.0"): + with pytest.raises(ValueError, match="requires vLLM >= 0.22.0"): + model.load() + args_factory.assert_not_called() else: with pytest.raises(ReachedEngineArgs): model.load() args_factory.assert_called_once() + + +@pytest.mark.parametrize("tools", [[], [{"type": "function"}]]) +@pytest.mark.asyncio +async def test_multimodal_text_chat_preserves_pd_handoff(monkeypatch, tools): + from .. import core + + model = object.__new__(core.VLLMMultiModel) + model.model_family = SimpleNamespace( + model_family="internvl", model_name="internvl", model_ability=["vision"] + ) + model.model_uid = "p" + model.reasoning_parser = None + model.get_specific_prompt = Mock(return_value=("prompt", None)) + model._sanitize_chat_config = lambda config: config + model._to_chat_completion = Mock(return_value={"choices": []}) + model._post_process_completion = Mock(return_value={"choices": []}) + transfer = {"do_remote_prefill": True, "xavier_direct": {"ticket": "t"}} + model.async_generate = AsyncMock(return_value={"_pd_kv_transfer_params": transfer}) + monkeypatch.setattr(core, "validate_messages_media", Mock()) + result = await core.VLLMMultiModel.async_chat.__wrapped__( + model, [{"role": "user", "content": "hello"}], {"tools": tools} + ) + assert result["_pd_kv_transfer_params"] == transfer + + +@pytest.mark.parametrize( + "worker_layout,launch_layout,expected", + [(None, None, "DS"), ("SD", None, "SD"), ("SD", "DS", "DS")], +) +def test_nixl_conv_layout_default_and_overrides( + monkeypatch, worker_layout, launch_layout, expected +): + monkeypatch.delenv("VLLM_SSM_CONV_STATE_LAYOUT", raising=False) + if worker_layout: + monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", worker_layout) + env = {"VLLM_SSM_CONV_STATE_LAYOUT": launch_layout} if launch_layout else {} + configure_nixl_environment(env, "10.0.0.1:9997") + assert env["VLLM_SSM_CONV_STATE_LAYOUT"] == expected diff --git a/xinference/model/llm/vllm/xavier/README.md b/xinference/model/llm/vllm/xavier/README.md index 26343a59d9..39d1821088 100644 --- a/xinference/model/llm/vllm/xavier/README.md +++ b/xinference/model/llm/vllm/xavier/README.md @@ -30,7 +30,11 @@ P/D always uses asynchronous loads. In the hybrid snapshot path, when `xavier_gp If the last request is aborted while an async load is pending, an idle vLLM EngineCore may not poll completion again until the next request arrives. Actor-side writes and source lease release continue independently, but destination block reclamation and connector-side completion cleanup wait for that next engine step or shutdown. -The connector accepts supported block-first and K/V-first layouts and rejects an ambiguous block axis. Shared null-block positions in allocated cache groups are omitted from destination mappings; conflicting writes to real destinations remain errors. This is not a claim of support for recurrent state: hybrid/recurrent attention models such as Qwen3.5 are rejected. Successful model launch alone does not demonstrate correct recurrent-state transfer. +The connector accepts supported block-first and K/V-first layouts and rejects an ambiguous block axis. Shared null-block positions in allocated cache groups are omitted from destination mappings; conflicting writes to real destinations remain errors. Qwen3.5 text-only P/D additionally transfers each Gated DeltaNet cache group independently. Launch with `language_model_only=True`, `enable_prefix_caching=False`, `mamba_cache_mode="none"`, and `async_scheduling=False`. For recurrent models, Xavier defaults prefix caching and async scheduling off and rejects enabling either. On vLLM 0.21 it enables the hybrid KV cache manager for recurrent models only; explicit user settings are preserved. Keep that manager enabled for recurrent handoff. TP/PP remain 1 and speculative decoding is rejected. + +For recurrent models, P computes the first N-1 prompt tokens; D loads that exact recurrent state and computes the last token. P truncation is idempotent across scheduling retries. Attention tensors may contain several physical kernel blocks per logical block; zero-copy views keep those physical blocks within a single logical transfer unit. Tickets authorize source blocks separately for each layer, including both convolution and SSM state tensors. The engine retains all cache groups until transfer completion. One-token prompts recompute on D without transferring state. + +Recurrent history reuse is currently disabled, even when a history budget is configured. Ordinary attention-block hashes cannot identify a reusable recurrent state at an arbitrary prefix boundary. Supporting historical reuse requires atomic snapshots of attention KV plus the recurrent state at the same token position. Full-attention history caching remains unchanged. In the legacy V0 path, prefill replicas release only the requested completed sequence; ordinary hybrid and decode replicas retain automatic cleanup. @@ -74,3 +78,18 @@ XINFERENCE_TEST_PD_GPU=1 python -m pytest -v \ The test runs default Xavier direct handoff with tiered history, Xavier with history disabled, and native NIXL in real subprocesses, with one GPU per role. It checks streaming and non-streaming responses, repeated prompts and four concurrent requests, then terminates the deployment. Local prefix caching is disabled. Exact repeated-text checks are restricted to native NIXL and Xavier with history disabled; fixed-prefix and raw-bit history tests cover history reuse separately. Xavier requires producer registration and successful async completion logs from `get_finished`, and the default-history case requires a positive P-side history restore for each repeated prompt; native NIXL requires `calling _read_blocks` and completed receive logs, including the concurrent requests. For locally cached weights, set `XINFERENCE_TEST_PD_MODEL_PATH`, `XINFERENCE_TEST_PD_MODEL_NAME` and `XINFERENCE_TEST_PD_MODEL_SIZE` to the path and registered model name/size. The manually triggered **PD GPU integration** GitHub Actions workflow runs the same test on a selected two-GPU runner. + +### Qwen3.5 recurrent-state regression + +Set `XINFERENCE_TEST_PD_RECURRENT_GPU=1` and +`XINFERENCE_TEST_PD_RECURRENT_MODEL_PATH` to a local Qwen3.5-0.8B checkpoint, +then run `test/test_pd_recurrent_gpu.py`. This opt-in test compares Xavier and +native NIXL P/D against standalone greedy outputs for short and long prompts, +including repeated and concurrent requests. It requires actual transfer logs. +The test uses BF16, TP=PP=1, text-only mode and no prefix caching; it does not +establish recurrent history-cache support or a throughput advantage. + +Native NIXL's GDN support landed in vLLM PR #41869 and is included in +vLLM 0.22.0. Stock vLLM 0.21.0 rejects GDN in its convolution-state transfer +setup. The native leg of the regression requires 0.22.0 or newer; Xavier's +custom handoff continues to support 0.21.0. diff --git a/xinference/model/llm/vllm/xavier/engine.py b/xinference/model/llm/vllm/xavier/engine.py index 99e3c6df4c..851c8b0f3d 100644 --- a/xinference/model/llm/vllm/xavier/engine.py +++ b/xinference/model/llm/vllm/xavier/engine.py @@ -104,6 +104,28 @@ def _patch_v1_engine_args( VLLM_VERSION, ) + # vLLM's is_hybrid means attention plus recurrent state, not a mixture + # 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") + and getattr(engine_args, "disable_hybrid_kv_cache_manager", None) is None + ): + engine_args.disable_hybrid_kv_cache_manager = False + if getattr(engine_args, "enable_prefix_caching", None) is None: + engine_args.enable_prefix_caching = not recurrent + if recurrent: + if getattr(engine_args, "async_scheduling", None) is None: + engine_args.async_scheduling = False + if engine_args.enable_prefix_caching or engine_args.async_scheduling: + raise ValueError( + "Xavier recurrent handoff requires enable_prefix_caching=False " + "and async_scheduling=False" + ) + engine_args.kv_transfer_config = KVTransferConfig( kv_connector=XAVIER_CONNECTOR, kv_connector_module_path=XAVIER_CONNECTOR_MODULE, diff --git a/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py b/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py index 8b743f03ed..0b5e5f963c 100644 --- a/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py +++ b/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py @@ -40,6 +40,12 @@ def pd_cluster(monkeypatch, tmp_path, backend): "level": "DEBUG", "propagate": False, } + # Per-op tensor debug messages synchronize CUDA and overwhelm these logs. + logging_config["loggers"]["vllm.ir.op"] = { + "handlers": ["stream_handler", "file_handler"], + "level": "WARNING", + "propagate": False, + } config_path = tmp_path / "vllm-logging.json" config_path.write_text(json.dumps(logging_config)) monkeypatch.setenv("VLLM_LOGGING_CONFIG_PATH", str(config_path)) diff --git a/xinference/model/llm/vllm/xavier/test/test_pd_recurrent_gpu.py b/xinference/model/llm/vllm/xavier/test/test_pd_recurrent_gpu.py new file mode 100644 index 0000000000..50dc6e289e --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_pd_recurrent_gpu.py @@ -0,0 +1,109 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Qwen3.5 text P/D must match a standalone engine, with actual transfers.""" +import os +from pathlib import Path + +import pytest + +from . import test_pd_gpu + +pd_cluster = test_pd_gpu.pd_cluster + +pytestmark = pytest.mark.skipif( + os.environ.get("XINFERENCE_TEST_PD_RECURRENT_GPU") != "1", + reason="requires two CUDA GPUs, vLLM 0.21+ and Qwen3.5-0.8B", +) + + +@pytest.mark.parametrize("backend", ["xavier", "nixl"]) +def test_recurrent_pd_matches_standalone(pd_cluster, backend): + from concurrent.futures import ThreadPoolExecutor + + import openai + + from xinference.client import Client + + if backend == "nixl": + from packaging.version import Version + from vllm import __version__ + + if Version(__version__) < Version("0.22.0"): + pytest.skip("Native NIXL GDN transfer requires vLLM >= 0.22.0") + + endpoint, log_path = pd_cluster + client = Client(endpoint) + worker = client.get_workers_info()[0]["work-ip"] + uid = "recurrent-pd" + config = dict( + model_uid=uid, + model_name="qwen3.5", + model_size_in_billions="0_8", + model_format="pytorch", + quantization="none", + model_engine="vLLM", + model_path=os.environ["XINFERENCE_TEST_PD_RECURRENT_MODEL_PATH"], + enable_virtual_env=False, + max_model_len=2048, + gpu_memory_utilization=0.4, + dtype="bfloat16", + enforce_eager=True, + enable_prefix_caching=False, + language_model_only=True, + mamba_cache_mode="none", + async_scheduling=False, + disable_hybrid_kv_cache_manager=False, + ) + if backend == "nixl": + config["envs"] = {"VLLM_SSM_CONV_STATE_LAYOUT": "DS"} + api = openai.OpenAI(base_url=endpoint + "/v1", api_key="unused", timeout=180) + prompts = [ + "Write exactly this sentence, with no extra text: The answer is four.", + "Water evaporates and condenses in clouds. " * 100 + "Explain rain.", + ] + + def generate(prompt): + response = api.chat.completions.create( + model=uid, + messages=[{"role": "user", "content": prompt}], + max_tokens=32, + temperature=0, + extra_body={"chat_template_kwargs": {"enable_thinking": False}}, + ) + assert response.usage.completion_tokens > 1 + return response.choices[0].message.content + + try: + client.launch_model(**config, n_gpu=1, gpu_idx=[0]) + expected = [generate(prompt) for prompt in prompts] + client.terminate_model(uid) + client.launch_model( + **config, + replica=2, + vllm_transfer_backend_type=backend, + replica_config=[ + { + "role": role, + "devices": [{"worker_ip": worker, "n_gpu": 1, "gpu_idx": [i]}], + } + for i, role in enumerate(["prefill", "decode"]) + ], + ) + offset = Path(log_path).stat().st_size + for _ in range(2): + assert [generate(prompt) for prompt in prompts] == expected + with ThreadPoolExecutor(max_workers=2) as executor: + assert list(executor.map(generate, prompts)) == expected + with open(log_path, "rb") as log: + log.seek(offset) + evidence = log.read().decode(errors="replace") + if backend == "xavier": + assert "Register Xavier direct handoff" in evidence + assert "Finished Xavier async KV load" in evidence + assert "Restored Xavier history" not in evidence + else: + assert "calling _read_blocks" in evidence + print(f"{backend}: six recurrent PD outputs match standalone", flush=True) + finally: + if uid in client.list_models(): + client.terminate_model(uid) diff --git a/xinference/model/llm/vllm/xavier/test/test_recurrent_handoff.py b/xinference/model/llm/vllm/xavier/test/test_recurrent_handoff.py new file mode 100644 index 0000000000..c2a1abb8c5 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_recurrent_handoff.py @@ -0,0 +1,291 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Hybrid attention needs position-correct state and independent cache groups.""" +import sys +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import torch + +from .test_direct_handoff import direct_runtime + + +@pytest.fixture +def recurrent(connector_module, connector_config): + connector_config.kv_transfer_config.get_from_extra_config = lambda *a: { + "rank": 0, + "role": "prefill", + "gpu_cache_bytes": 1024, + } + caches = SimpleNamespace( + num_blocks=8, + kv_cache_groups=[ + SimpleNamespace( + layer_names=["linear"], + kv_cache_spec=SimpleNamespace(mamba_cache_mode="none"), + ), + SimpleNamespace(layer_names=["attention"], kv_cache_spec=SimpleNamespace()), + ], + ) + instance = connector_module.XavierConnector(connector_config, None, caches) + yield instance + instance.shutdown() + + +def test_prefill_state_matches_decoder_position(recurrent, monkeypatch): + monkeypatch.setitem( + sys.modules, + "vllm.v1.request", + SimpleNamespace(RequestStatus=SimpleNamespace(FINISHED_ABORTED="aborted")), + ) + actor = SimpleNamespace(register_direct_gpu_v1=AsyncMock()) + recurrent._get_transfer_ref = AsyncMock(return_value=actor) + request = SimpleNamespace( + request_id="p", + prompt_token_ids=list(range(34)), + _all_token_ids=list(range(34)), + num_prompt_tokens=34, + max_tokens=10, + status="finished", + kv_transfer_params={"do_remote_decode": True}, + ) + assert recurrent.get_num_new_matched_tokens(request, 0) == (0, False) + assert request.num_prompt_tokens == 33 and request.max_tokens == 1 + assert request.prompt_token_ids == request._all_token_ids == list(range(33)) + recurrent.get_num_new_matched_tokens(request, 16) + assert request.num_prompt_tokens == 33 # Retry must not truncate twice. + retained, meta = recurrent.request_finished_all_groups(request, ([7], [2, 3, 4])) + assert retained and not recurrent._history_enabled + handoff = meta["xavier_direct"] + assert handoff["tokens"] == 33 + assert handoff["blocks_by_group"] == [[7], [2, 3, 4]] + actor.register_direct_gpu_v1.assert_awaited_once_with( + handoff["ticket"], "p", {"linear": [7], "attention": [2, 3, 4]} + ) + + +def test_decoder_maps_recurrent_state_separately(recurrent): + recurrent._is_producer = False + recurrent._get_transfer_ref = AsyncMock( + return_value=SimpleNamespace( + claim_remote_direct_gpu_v1=AsyncMock(return_value=True) + ) + ) + request = SimpleNamespace( + request_id="d", + num_computed_tokens=0, + kv_transfer_params={ + "xavier_direct": { + "rank": 0, + "ticket": "t", + "tokens": 33, + "blocks": [2, 3, 4], + "blocks_by_group": [[7], [2, 3, 4]], + } + }, + ) + assert recurrent.get_num_new_matched_tokens(request, 16) == (17, True) + recurrent.update_state_after_alloc( + request, SimpleNamespace(get_block_ids=lambda: ([6], [1, 4, 5])), 17 + ) + load = recurrent._requests_need_load["d"] + assert load.local_transfers_by_group == {0: {0: {7: 6}}, 1: {0: {3: 4, 4: 5}}} + assert request.kv_transfer_params["do_remote_prefill"] is False + + +@pytest.mark.parametrize("kv_first", [False, True]) +def test_logical_blocks_are_zero_copy_views(recurrent, kv_first): + cache = torch.arange(2 * 16 * 4 * 2 * 3).reshape(2, 16, 4, 2, 3) + if not kv_first: + cache = cache.movedim(1, 0).contiguous() + view = recurrent._cache_block_view("attention", cache) + assert view.shape[0] == 8 and view.shape[1] == 2 + view[3].fill_(-1) + physical = cache.movedim(1, 0) if kv_first else cache + assert (physical[6:8] == -1).all() and (physical[5] != -1).all() + state = torch.zeros(8, 2, 4) + assert recurrent._cache_block_view("linear", state) is state + + +def test_ticket_owns_blocks_per_layer(monkeypatch): + runtime = direct_runtime(monkeypatch) + runtime.caches = { + "attention": torch.zeros(8, 2), + "linear#0": torch.zeros(8, 3), + "linear#1": torch.zeros(8, 4), + } + runtime.register_direct("t", "p", {"attention": [1, 2], "linear": [7]}) + assert runtime.direct_requests["t"].layer_blocks == { + "attention": {1, 2}, + "linear#0": {7}, + "linear#1": {7}, + } + + +@pytest.mark.asyncio +async def test_grouped_transfer_preserves_state_and_rejects_other_group_blocks( + monkeypatch, +): + import xoscar as xo + + from ..request_transfer import LayerRead + from .test_direct_handoff import peer_for + + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + source.caches = { + "attention": torch.arange(16, dtype=torch.bfloat16).reshape(8, 2), + "linear#0": torch.arange(16, dtype=torch.float32).reshape(8, 2), + "linear#1": -torch.arange(16, dtype=torch.float32).reshape(8, 2), + } + dest.caches = { + name: torch.zeros_like(cache) for name, cache in source.caches.items() + } + source.register_direct("t", "p", {"attention": [1, 2], "linear": [7]}) + # A block held for attention is not permission to read another request's SSM. + with pytest.raises(ValueError, match="source layout or block IDs"): + await source.send_direct( + "t", + [LayerRead("linear#1", [1], [4], (2,), torch.float32)], + dest.recv_ref, + 16, + ) + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer_for(source))) + + async def copy(buffers, refs): + refs[0].copy_(buffers[0]) + + monkeypatch.setattr(xo, "copy_to", copy) + maps = {"attention": {1: 3, 2: 4}, "linear#0": {7: 6}, "linear#1": {7: 6}} + assert await dest.run(dest.load_direct, [{0: maps}], ["t"]) == set() + for name, mapping in maps.items(): + assert torch.equal( + dest.caches[name][list(mapping.values())], + source.caches[name][list(mapping)], + ) + assert source.poll_direct() == {"p"} + + +def test_single_token_prompt_does_not_transfer_advanced_state(recurrent, monkeypatch): + monkeypatch.setitem( + sys.modules, + "vllm.v1.request", + SimpleNamespace(RequestStatus=SimpleNamespace(FINISHED_ABORTED="aborted")), + ) + request = SimpleNamespace( + request_id="p", + prompt_token_ids=[10], + _all_token_ids=[10], + num_prompt_tokens=1, + status="finished", + kv_transfer_params={"do_remote_decode": True}, + ) + assert recurrent.get_num_new_matched_tokens(request, 0) == (0, False) + retained, metadata = recurrent.request_finished_all_groups(request, ([7], [2])) + assert not retained and metadata["xavier_direct"]["tokens"] == 0 + assert request.prompt_token_ids == [10] + + +def test_recurrent_abort_before_allocation(recurrent, monkeypatch): + monkeypatch.setitem( + sys.modules, + "vllm.v1.request", + SimpleNamespace(RequestStatus=SimpleNamespace(FINISHED_ABORTED="aborted")), + ) + request = SimpleNamespace( + request_id="p", status="aborted", kv_transfer_params={"do_remote_decode": True} + ) + assert recurrent.request_finished_all_groups(request, ()) == (False, None) + + +@pytest.mark.parametrize( + "groups", + [ + None, + [], + [[7]], + [[7], [2, 3, 4], [8]], + [[], [2, 3, 4]], + [[7, 8], [2, 3, 4]], + [[7], [2]], + [[7], [2, 3, "4"]], + ], +) +def test_mismatched_recurrent_payload_recomputes_before_claim(recurrent, groups): + recurrent._is_producer = False + claim = AsyncMock(return_value=True) + recurrent._get_transfer_ref = AsyncMock( + return_value=SimpleNamespace(claim_remote_direct_gpu_v1=claim) + ) + request = SimpleNamespace( + request_id="bad", + num_computed_tokens=0, + kv_transfer_params={ + "do_remote_prefill": True, + "xavier_direct": { + "rank": 0, + "ticket": "t", + "tokens": 33, + "blocks": [2, 3, 4], + "blocks_by_group": groups, + }, + }, + ) + assert recurrent.get_num_new_matched_tokens(request, 0) == (0, False) + assert request.kv_transfer_params["do_remote_prefill"] is False + assert "bad" not in recurrent._requests_need_load + claim.assert_not_awaited() + assert recurrent.get_num_new_matched_tokens(request, 0) == (0, False) + + +@pytest.mark.parametrize( + "case,match", + [ + ("all_recurrent", "full-attention"), + ("speculative", "speculative decoding"), + ("no_direct", "prefill/decode GPU handoff"), + ("prefix", "enable_prefix_caching=False"), + ("async", "async_scheduling=False"), + ], +) +def test_recurrent_configuration_rejections( + connector_module, connector_config, case, match +): + connector_config.kv_transfer_config.get_from_extra_config = lambda *a: { + "rank": 0, + "role": "hybrid" if case == "no_direct" else "prefill", + "gpu_cache_bytes": 1024, + } + groups = [ + SimpleNamespace( + layer_names=["linear"], + kv_cache_spec=SimpleNamespace(mamba_cache_mode="none"), + ) + ] + if case != "all_recurrent": + groups.append( + SimpleNamespace(layer_names=["attention"], kv_cache_spec=SimpleNamespace()) + ) + if case == "speculative": + connector_config.speculative_config = object() + if case == "prefix": + connector_config.cache_config.enable_prefix_caching = True + if case == "async": + connector_config.scheduler_config = SimpleNamespace(async_scheduling=True) + with pytest.raises(ValueError, match=match): + connector_module.XavierConnector( + connector_config, + None, + SimpleNamespace(num_blocks=8, kv_cache_groups=groups), + ) + + +@pytest.mark.parametrize("physical_blocks", [2, 4]) +def test_recurrent_attention_rejects_ambiguous_two_block_axis( + recurrent, physical_blocks +): + recurrent._num_cache_blocks = 2 + with pytest.raises(ValueError, match="Ambiguous"): + recurrent._cache_block_view( + "attention", torch.zeros(2, physical_blocks, 16, 2, 4) + ) diff --git a/xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py b/xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py index 1fddf31304..c92aa329bf 100644 --- a/xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py +++ b/xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py @@ -61,7 +61,7 @@ async def test_read_large_cache_blocks_incrementally(monkeypatch): ) -def test_hybrid_cache_is_rejected_before_serving(monkeypatch): +def test_recurrent_prefix_cache_is_rejected_before_serving(monkeypatch): pytest.importorskip("vllm") from ..v1_connector import KVConnectorBase_V1, XavierConnector @@ -76,8 +76,9 @@ def init_base(self, **kwargs): ) caches = SimpleNamespace( kv_cache_groups=[ - SimpleNamespace(kv_cache_spec=SimpleNamespace(mamba_cache_mode="all")) + SimpleNamespace(kv_cache_spec=SimpleNamespace(mamba_cache_mode="all")), + SimpleNamespace(kv_cache_spec=SimpleNamespace()), ] ) - with pytest.raises(ValueError, match="does not yet support hybrid/recurrent"): + with pytest.raises(ValueError, match="requires mamba_cache_mode=none"): XavierConnector(config, None, caches) diff --git a/xinference/model/llm/vllm/xavier/test/test_v1_engine.py b/xinference/model/llm/vllm/xavier/test/test_v1_engine.py index bc8944ee9e..c82a1f8a01 100644 --- a/xinference/model/llm/vllm/xavier/test/test_v1_engine.py +++ b/xinference/model/llm/vllm/xavier/test/test_v1_engine.py @@ -50,7 +50,11 @@ def engine(monkeypatch): ) def test_v1_connector_configuration(engine, role, kv_role): module, factory = engine - args = SimpleNamespace(additional_config={"custom": 1}, enforce_eager=False) + args = SimpleNamespace( + additional_config={"custom": 1}, + enforce_eager=False, + create_model_config=lambda: SimpleNamespace(is_hybrid=False), + ) config = {"role": role, "rank": 2, "block_tracker_uid": b"tracker"} module.XavierEngine.from_engine_args(args, xavier_config=config) assert args.kv_transfer_config.kv_role == kv_role @@ -58,6 +62,8 @@ def test_v1_connector_configuration(engine, role, kv_role): assert args.additional_config["custom"] == 1 json.dumps(args.additional_config) assert args.enforce_eager + assert getattr(args, "disable_hybrid_kv_cache_manager", None) is None + assert args.enable_prefix_caching is True assert config == {"role": role, "rank": 2, "block_tracker_uid": b"tracker"} factory.assert_called_once() @@ -96,7 +102,9 @@ def create(*args): assert os.environ["VLLM_WORKER_MULTIPROC_METHOD"] == (configured or "spawn") factory.side_effect = create - module.XavierEngine.from_engine_args(SimpleNamespace()) + module.XavierEngine.from_engine_args( + SimpleNamespace(create_model_config=lambda: SimpleNamespace(is_hybrid=False)) + ) factory.assert_called_once() @@ -120,3 +128,41 @@ def test_gpu_path_rejects_legacy_before_engine_start(engine): object(), xavier_config={"gpu_cache_bytes": 0} ) factory.assert_not_called() + + +@pytest.mark.parametrize("recurrent", [False, True]) +@pytest.mark.parametrize("vllm_version", ["0.21.0", "0.22.0"]) +@pytest.mark.parametrize("disable_hma", [None, True, False]) +def test_hma_defaults_only_for_recurrent_models( + engine, recurrent, vllm_version, disable_hma +): + module, _ = engine + module.VLLM_VERSION = vllm_version + args = SimpleNamespace( + create_model_config=lambda: SimpleNamespace(is_hybrid=recurrent), + disable_hybrid_kv_cache_manager=disable_hma, + ) + module.XavierEngine.from_engine_args(args) + expected = ( + False + if recurrent and vllm_version == "0.21.0" and disable_hma is None + else disable_hma + ) + assert args.disable_hybrid_kv_cache_manager is expected + assert args.enable_prefix_caching is (not recurrent) + if recurrent: + assert args.async_scheduling is False + + +@pytest.mark.parametrize("setting", ["enable_prefix_caching", "async_scheduling"]) +def test_recurrent_rejects_unsupported_settings_before_start(engine, setting): + module, factory = engine + args = SimpleNamespace( + create_model_config=lambda: SimpleNamespace(is_hybrid=True), + **{setting: True}, + ) + with pytest.raises( + ValueError, match="enable_prefix_caching=False and async_scheduling=False" + ): + module.XavierEngine.from_engine_args(args) + factory.assert_not_called() diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 1ba9ae1324..8f421160ce 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -76,6 +76,8 @@ class XavierLoadRequest: local_transfers_by_group: Dict[int, Dict[int, Dict[int, int]]] = field( default_factory=dict ) + remote_blocks_by_group: List[List[int]] = field(default_factory=list) + local_hit_tokens: int = 0 @dataclass @@ -106,39 +108,84 @@ def __init__( ) self._block_size = vllm_config.cache_config.block_size self._cache_dtype = vllm_config.cache_config.cache_dtype - has_recurrent_cache = any( - hasattr(group.kv_cache_spec, "mamba_cache_mode") - for group in kv_cache_config.kv_cache_groups - ) - if has_recurrent_cache: - raise ValueError( - "Xavier PD does not yet support hybrid/recurrent attention caches " - "(for example Qwen3.5). Launch this model without PD/Xavier, " - "or use a full-attention model such as Qwen3 for PD." - ) + self._recurrent_groups = { + i + for i, group in enumerate(kv_cache_config.kv_cache_groups) + if hasattr(group.kv_cache_spec, "mamba_cache_mode") + } + self._has_recurrent_cache = bool(self._recurrent_groups) + if self._has_recurrent_cache: + if getattr( + vllm_config.cache_config, "enable_prefix_caching", False + ) or getattr( + getattr(vllm_config, "scheduler_config", None), + "async_scheduling", + False, + ): + raise ValueError( + "Xavier recurrent handoff requires enable_prefix_caching=False " + "and async_scheduling=False" + ) + if len(self._recurrent_groups) == len(kv_cache_config.kv_cache_groups): + raise ValueError( + "Xavier recurrent handoff requires a full-attention group" + ) + if getattr(vllm_config, "speculative_config", None) is not None: + raise ValueError( + "Xavier recurrent handoff does not support speculative decoding" + ) + for i in self._recurrent_groups: + if ( + kv_cache_config.kv_cache_groups[i].kv_cache_spec.mamba_cache_mode + != "none" + ): + raise ValueError( + "Xavier recurrent handoff requires mamba_cache_mode=none" + ) parallel = vllm_config.parallel_config if parallel.tensor_parallel_size != 1 or parallel.pipeline_parallel_size != 1: raise ValueError("Xavier V1 currently requires TP=1 and PP=1") - if ( - vllm_config.lora_config is not None - or vllm_config.model_config.is_multimodal_model + if vllm_config.lora_config is not None or ( + vllm_config.model_config.is_multimodal_model + and not getattr( + getattr(vllm_config.model_config, "multimodal_config", None), + "language_model_only", + False, + ) ): raise ValueError( - "Xavier V1 currently supports text-only models without LoRA" + "Xavier V1 requires text-only models or language_model_only=True, without LoRA" ) self._direct_handoff = uses_direct_handoff(self._xavier_config) self._gpu_budget = self._xavier_config.get("gpu_cache_bytes") - if self._direct_handoff and len(kv_cache_config.kv_cache_groups) != 1: + if ( + self._direct_handoff + and len(kv_cache_config.kv_cache_groups) != 1 + and not self._has_recurrent_cache + ): raise ValueError( "Direct handoff requires P/D roles, GPU transport and one KV group" ) + if self._has_recurrent_cache and not self._direct_handoff: + raise ValueError( + "Xavier recurrent caches require prefill/decode GPU handoff" + ) self._gpu_cache_mapped = False self._gpu_mapping_lock = asyncio.Lock() self._kv_schema: Optional[XavierKVSchema] = None self._rank = int(self._xavier_config.get("rank", 0)) self._is_producer = self._kv_transfer_config.is_kv_producer self._is_consumer = self._kv_transfer_config.is_kv_consumer - self._history_enabled = self._direct_handoff and bool(self._gpu_budget) + self._history_enabled = ( + self._direct_handoff + and bool(self._gpu_budget) + and not self._has_recurrent_cache + ) + if self._has_recurrent_cache and self._gpu_budget: + logger.info( + "Xavier recurrent handoff uses engine GPU states; historical prefix " + "reuse is unavailable, so no history cache budget is allocated." + ) if self._history_enabled and self._is_producer: # P may restore its own independent history before computing a suffix. self._is_consumer = True @@ -270,7 +317,7 @@ def register_kv_caches( self._registered_kv_caches = dict(kv_caches) schema = { name: ( - tuple(block_major_view(tensor, self._num_cache_blocks).shape[1:]), + tuple(self._cache_block_view(layer, tensor).shape[1:]), tensor.dtype, ) for layer, cache in self._registered_kv_caches.items() @@ -355,6 +402,21 @@ def get_num_new_matched_tokens( if self._direct_handoff: params = getattr(request, "kv_transfer_params", None) or {} handoff = params.get("xavier_direct") + if ( + self._has_recurrent_cache + and self._is_producer + and params.get("do_remote_decode") + ): + if ( + not params.get("_xavier_truncated") + and request.num_prompt_tokens > 1 + ): + request.prompt_token_ids.pop() + request._all_token_ids.pop() + request.num_prompt_tokens -= 1 + request.max_tokens = 1 + params["_xavier_truncated"] = True + return 0, False if self._is_producer and self._history_enabled: previous = self._requests_need_load.get(request.request_id) if previous is not None: @@ -413,6 +475,29 @@ async def release(): params["do_remote_prefill"] = False return 0, False + remote_groups = handoff.get("blocks_by_group", []) + if self._has_recurrent_cache or remote_groups: + expected_blocks = (tokens + self._block_size - 1) // self._block_size + valid_groups = ( + isinstance(remote_groups, (list, tuple)) + and len(remote_groups) == len(self._kv_cache_config.kv_cache_groups) + and all( + isinstance(group, (list, tuple)) + and len(group) + == (1 if i in self._recurrent_groups else expected_blocks) + and all(type(block) is int and block >= 0 for block in group) + for i, group in enumerate(remote_groups) + ) + ) + if not valid_groups: + logger.warning( + "Incompatible Xavier cache groups; recomputing request %s", + request.request_id, + ) + params["do_remote_prefill"] = False + self._requests_need_load.pop(request.request_id, None) + return 0, False + async def claim(): ref = await self._get_transfer_ref() return await ref.claim_remote_direct_gpu_v1(rank, ticket) @@ -427,6 +512,8 @@ async def claim(): request.request_id, {rank: {block: i for i, block in enumerate(blocks) if i >= start}}, lease=ticket, + remote_blocks_by_group=remote_groups, + local_hit_tokens=num_computed_tokens, ) return tokens - num_computed_tokens, True self._requests_need_load.pop(request.request_id, None) @@ -501,6 +588,22 @@ def update_state_after_alloc( for group_id, group_block_ids in enumerate(block_ids_by_group): updated: Dict[int, Dict[int, int]] = {} for from_rank, remote_to_placeholder in load_request.transfers.items(): + if load_request.remote_blocks_by_group: + remote = load_request.remote_blocks_by_group[group_id] + if group_id in self._recurrent_groups: + if len(remote) != 1 or len(group_block_ids) != 1: + raise ValueError( + "Expected one recurrent state block per request" + ) + updated[from_rank] = {remote[0]: group_block_ids[0]} + else: + start = load_request.local_hit_tokens // self._block_size + updated[from_rank] = { + block: group_block_ids[i] + for i, block in enumerate(remote) + if i >= start and (group_id, i) not in null_positions + } + continue # Placeholder indices are absolute positions in the prompt. # Allocation includes locally cached prefix blocks: remote # suffix blocks must not overwrite that prefix. Indexing by @@ -545,6 +648,7 @@ def request_finished( self, request: "Request", block_ids: List[int], + block_ids_by_group: Optional[Tuple[List[int], ...]] = None, ) -> tuple[bool, dict[str, Any] | None]: if self._direct_handoff: pending = self._requests_need_load.get(request.request_id) @@ -562,11 +666,34 @@ def request_finished( if request.status == RequestStatus.FINISHED_ABORTED: return False, None token_count = max(len(request.prompt_token_ids) - 1, 0) + if self._has_recurrent_cache: + if not params.get("_xavier_truncated"): + return False, { + "do_remote_prefill": True, + "xavier_direct": { + "rank": self._rank, + "ticket": "", + "blocks": [], + "tokens": 0, + }, + } + token_count = len(request.prompt_token_ids) count = min( len(block_ids), (token_count + self._block_size - 1) // self._block_size, ) blocks = block_ids[:count] + grouped = None + registered: List[int] | Dict[str, List[int]] = blocks + if block_ids_by_group is not None: + grouped = [ + list(ids) if i in self._recurrent_groups else list(ids[:count]) + for i, ids in enumerate(block_ids_by_group) + ] + registered = { + name: grouped[group_id] + for name, group_id in self._layer_group_ids.items() + } ticket = uuid.uuid4().hex if blocks else "" if blocks: @@ -579,7 +706,7 @@ async def register(): ) history_args = ([key for key, _ in hashes[:count]],) await ref.register_direct_gpu_v1( - ticket, request.request_id, blocks, *history_args + ticket, request.request_id, registered, *history_args ) self._call(register()) @@ -596,6 +723,7 @@ async def register(): "address": self._xavier_config.get("rank_address"), "ticket": ticket, "blocks": blocks, + **({"blocks_by_group": grouped} if grouped is not None else {}), "tokens": min(token_count, count * self._block_size), }, } @@ -638,6 +766,11 @@ def request_finished_all_groups( request: "Request", block_ids: Tuple[List[int], ...], ) -> tuple[bool, dict[str, Any] | None]: + if self._has_recurrent_cache and block_ids: + full_group = next( + i for i in range(len(block_ids)) if i not in self._recurrent_groups + ) + return self.request_finished(request, block_ids[full_group], block_ids) return self.request_finished(request, block_ids[0] if block_ids else []) def take_events(self) -> Iterable: @@ -993,7 +1126,7 @@ async def _ensure_gpu_cache_mapping(self): descriptors = {} for name, cache in self._registered_kv_caches.items(): for layer, tensor in self._iter_kv_tensors(name, cache): - tensor = block_major_view(tensor, self._num_cache_blocks) + tensor = self._cache_block_view(name, tensor) if not tensor.is_cuda: raise ValueError( "Xavier GPU transfer requires CUDA KV caches" @@ -1005,7 +1138,9 @@ async def _ensure_gpu_cache_mapping(self): ) torch.cuda.synchronize() await transfer.map_gpu_caches_v1( - descriptors, self._gpu_budget, direct_handoff=self._direct_handoff + descriptors, + 0 if self._has_recurrent_cache else self._gpu_budget, + direct_handoff=self._direct_handoff, ) self._gpu_cache_mapped = True return transfer @@ -1090,6 +1225,19 @@ async def load(): # transfer completion or introducing another polling RPC per engine step. await asyncio.sleep(0) + def _cache_block_view(self, layer_name: str, tensor: torch.Tensor) -> torch.Tensor: + if not self._has_recurrent_cache: + return block_major_view(tensor, self._num_cache_blocks) + if self._get_layer_group_id(layer_name) in self._recurrent_groups: + return block_major_view(tensor, self._num_cache_blocks) + # HMA attention kernels may split a logical block into multiple physical + # blocks. Keep that extra dimension inside the transfer unit, as a view. + view = block_major_view(tensor, self._num_cache_blocks, allow_multiple=True) + physical = view.shape[0] + return view.unflatten( + 0, (self._num_cache_blocks, physical // self._num_cache_blocks) + ) + @staticmethod def _iter_kv_tensors( layer_name: str, kv_layer: torch.Tensor | Sequence[torch.Tensor] diff --git a/xinference/model/llm/xavier/backends/torch/direct_handoff.py b/xinference/model/llm/xavier/backends/torch/direct_handoff.py index 8fe2b9669d..9774630d28 100644 --- a/xinference/model/llm/xavier/backends/torch/direct_handoff.py +++ b/xinference/model/llm/xavier/backends/torch/direct_handoff.py @@ -20,7 +20,7 @@ @dataclass class DirectRequest: request_id: str - blocks: set[int] + layer_blocks: dict[str, set[int]] deadline: float claimed: bool = False reading: bool = False @@ -75,17 +75,28 @@ async def _close(self): def register_direct(self, ticket, request_id, blocks, hashes=None): if self.closing or ticket in self.direct_requests: raise ValueError("Invalid direct handoff registration") - if not blocks or any( - block < 0 or block >= len(cache) - for cache in self.caches.values() - for block in blocks + if isinstance(blocks, dict): + if hashes is not None: + raise ValueError( + "Grouped recurrent handoff cannot retain block history" + ) + layer_blocks = { + name: set(blocks[name.split("#", 1)[0]]) for name in self.caches + } + else: + allowed = set(blocks) + layer_blocks = {name: allowed for name in self.caches} + if not layer_blocks or any( + not ids + or any(block < 0 or block >= len(self.caches[name]) for block in ids) + for name, ids in layer_blocks.items() ): raise ValueError("Invalid direct KV block IDs") if hashes is not None and len(hashes) != len(blocks): raise ValueError("History hashes and engine blocks differ") self.direct_requests[ticket] = DirectRequest( request_id, - set(blocks), + layer_blocks, time.monotonic() + 120, hashes=dict(zip(blocks, hashes or [])), ) @@ -149,7 +160,7 @@ async def send_direct(self, ticket, reads, remote_ref, slab_bytes): for read in reads: cache = self.caches[read.layer] if ( - not set(read.keys) <= state.blocks + not set(read.keys) <= state.layer_blocks[read.layer] or cache.dtype != read.dtype or tuple(cache.shape[1:]) != read.block_shape ): diff --git a/xinference/model/llm/xavier/backends/torch/snapshot.py b/xinference/model/llm/xavier/backends/torch/snapshot.py index 7d516a2a9f..d04f09a6a2 100644 --- a/xinference/model/llm/xavier/backends/torch/snapshot.py +++ b/xinference/model/llm/xavier/backends/torch/snapshot.py @@ -75,13 +75,21 @@ def read(self, layer: str, keys: List[int]) -> torch.Tensor: return torch.stack([self.blocks[key][layer] for key in keys]).contiguous() -def block_major_view(tensor: torch.Tensor, num_blocks: int) -> torch.Tensor: +def block_major_view( + tensor: torch.Tensor, num_blocks: int, *, allow_multiple: bool = False +) -> torch.Tensor: """Normalize supported attention layouts without guessing ambiguous axes.""" - if tensor.ndim == 5 and tensor.shape[0] == 2: - if tensor.shape[1] == num_blocks: - if num_blocks == 2: - raise ValueError("Ambiguous Xavier KV block axis with only two blocks") - return tensor.movedim(1, 0) - if tensor.shape[0] != num_blocks: + + def matches(size: int) -> bool: + return size == num_blocks or ( + allow_multiple and num_blocks > 0 and size > 0 and size % num_blocks == 0 + ) + + first_is_blocks = matches(tensor.shape[0]) + if tensor.ndim == 5 and tensor.shape[0] == 2 and matches(tensor.shape[1]): + if first_is_blocks: + raise ValueError("Ambiguous Xavier KV block axis with only two blocks") + return tensor.movedim(1, 0) + if not first_is_blocks: raise ValueError(f"Unsupported Xavier KV layout {tuple(tensor.shape)}") return tensor