diff --git a/.github/workflows/pd-gpu.yaml b/.github/workflows/pd-gpu.yaml index a84fa49569..de8223d7e0 100644 --- a/.github/workflows/pd-gpu.yaml +++ b/.github/workflows/pd-gpu.yaml @@ -41,11 +41,11 @@ jobs: run: | python -m venv .venv-pd .venv-pd/bin/python -m pip install -U pip 'setuptools<82' - .venv-pd/bin/python -m pip install -e '.[dev]' "vllm==$VLLM_VERSION" "nixl==$NIXL_VERSION" openai + .venv-pd/bin/python -m pip install -e '.[dev]' "vllm==$VLLM_VERSION" "nixl==$NIXL_VERSION" "xoscar[nixl]>=0.11.1" openai - name: Verify PD with actual KV transfer run: | source .venv-pd/bin/activate - python -m pytest --timeout=1800 -vv xinference/model/llm/vllm/xavier/test/test_pd_gpu.py xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py + python -m pytest --timeout=1800 -vv xinference/model/llm/vllm/xavier/test/test_pd_gpu.py xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py xinference/model/llm/vllm/xavier/test/test_direct_history.py - name: Upload runtime logs if: always() uses: actions/upload-artifact@v4 diff --git a/.github/workflows/python.yaml b/.github/workflows/python.yaml index a06f0c7d9c..c4df0aeea6 100644 --- a/.github/workflows/python.yaml +++ b/.github/workflows/python.yaml @@ -380,6 +380,8 @@ jobs: xinference/model/llm/vllm/xavier/test/test_v1_engine.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_direct_handoff.py \ + xinference/model/llm/vllm/xavier/test/test_direct_history.py \ xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py \ xinference/model/llm/vllm/xavier/test/test_request_transfer.py \ xinference/model/llm/vllm/xavier/test/test_transfer_blocks.py \ 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 de58fe1b56..05c7755774 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 @@ -299,3 +299,84 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "Übertragungsfehler und inkompatible KV-Cache-Layouts lösen Fehler aus, statt die Anfrage stillschweigend neu zu berechnen." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D verwendet standardmäßig direkte GPU-zu-GPU-Übertragung. " +"Erforderlich sind Linux, NVIDIA-GPUs, vLLM >= 0.21.0 und " +"``xoscar[nixl]>=0.11.1`` in der Worker- und Modellumgebung. Verwenden Sie " +"eine erreichbare Hostadresse (nicht ``0.0.0.0``). Ohne NIXL schlägt der " +"Start fehl; es gibt keinen CPU-Fallback." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"``xavier_gpu_cache_bytes`` beträgt für P/D standardmäßig ``268435456`` (256 " +"MiB pro Replikat). Der Verlauf nutzt zuerst die GPU und wird bei erschöpfter" +" GPU-Kapazität auf die CPU ausgelagert. Der CPU-Verlauf ist durch die KV-" +"Blockkapazität der Engine begrenzt und kann im Hauptspeicher deren Cache-" +"Größe erreichen. CPU-Treffer stellen KV auf der GPU wieder her; Misses " +"werden lokal berechnet. ``0`` deaktiviert den Verlauf, behält aber die " +"direkte Übertragung bei. Übertragungsfehler lösen Fehler aus." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "Direkte Übertragung und gestufter Verlauf" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "Die direkte Übertragung erfordert derzeit ``n=1``." + + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"Der Verlauf wird vor der Übergabe an D auf P wiederhergestellt. Bei langen " +"Prompts wird möglicherweise nur ein Präfix behalten; die Aufbewahrung aller " +"Cache-Daten ist nicht garantiert." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"Nicht angenommene Übergaben verfallen nach 120 Sekunden. Verfällt eine " +"Übergabe, bevor D sie annimmt, berechnet D lokal neu. Nach der Annahme " +"bleibt der Quellcache reserviert, bis das Laden abgeschlossen ist oder die " +"Anfrage vor dem Laden abgebrochen wird." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "" +"Direkte Übertragungen und das Laden des Verlaufs können parallel zur " +"Antwortgenerierung anderer Anfragen erfolgen." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"Nicht angenommene Übergaben verfallen nach 120 Sekunden. Angenommene " +"Übergaben verfallen nach 10 Minuten ohne erneuten Planungsversuch oder " +"Übertragungsfortschritt. Laufende Kopien werden vor der Freigabe " +"abgeschlossen. Verfallene Übergaben werden lokal auf D neu berechnet." 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 1640945a04..45c93832d5 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 @@ -297,3 +297,83 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "Los fallos de transferencia y los formatos de caché KV incompatibles generan errores en lugar de volver a calcular la solicitud silenciosamente." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D utiliza por defecto transferencia directa entre GPU. Requiere " +"Linux, GPU NVIDIA, vLLM >= 0.21.0 y ``xoscar[nixl]>=0.11.1`` en los entornos" +" del worker y del modelo. Use una dirección de host accesible (no " +"``0.0.0.0``). Sin NIXL, el inicio falla; no hay alternativa por CPU." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"Para P/D, ``xavier_gpu_cache_bytes`` vale por defecto ``268435456`` (256 MiB" +" por réplica). El historial usa primero GPU y pasa a CPU cuando se agota su " +"capacidad. El historial CPU está limitado por la capacidad de bloques KV del" +" motor y puede ocupar un tamaño similar en memoria del host. Los aciertos " +"CPU restauran KV en GPU; los fallos se calculan localmente. ``0`` desactiva " +"el historial y mantiene la transferencia directa. Los fallos de " +"transferencia generan errores." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "Transferencia directa e historial por niveles" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "La transferencia directa requiere actualmente ``n=1``." + + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"El historial se restaura en P antes de transferirlo a D. Los prompts largos " +"pueden conservar solo un prefijo; no se garantiza la retención de toda la " +"caché." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"Las transferencias no aceptadas caducan tras 120 segundos. Si caducan antes " +"de que D las acepte, D vuelve a calcular localmente. Tras la aceptación, la " +"caché de origen permanece reservada hasta que termine la carga o se cancele " +"la solicitud antes de cargar." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "" +"Las transferencias directas y las cargas del historial pueden solaparse con " +"la generación de respuestas de otras solicitudes." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"Las transferencias no aceptadas caducan tras 120 segundos. Las aceptadas " +"caducan tras 10 minutos sin reintentos de planificación ni progreso de " +"transferencia. Las copias en curso terminan antes de liberar los recursos. D" +" vuelve a calcular localmente las transferencias caducadas." 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 8cfac69b64..369a021f39 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 @@ -300,3 +300,84 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "Les échecs de transfert et les dispositions incompatibles du cache KV provoquent des erreurs au lieu de recalculer silencieusement la requête." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D utilise par défaut le transfert direct entre GPU. Il nécessite " +"Linux, des GPU NVIDIA, vLLM >= 0.21.0 et ``xoscar[nixl]>=0.11.1`` dans les " +"environnements du worker et du modèle. Utilisez une adresse hôte accessible " +"(pas ``0.0.0.0``). Sans NIXL, le démarrage échoue ; aucun repli CPU n’est " +"prévu." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"Pour P/D, ``xavier_gpu_cache_bytes`` vaut par défaut ``268435456`` (256 MiB " +"par réplique). L’historique utilise d’abord le GPU puis déborde vers le CPU " +"lorsque sa capacité est épuisée. L’historique CPU est limité par la capacité" +" en blocs KV du moteur et peut approcher cette taille de cache en mémoire " +"hôte. Les accès CPU restaurent les KV sur le GPU ; les défauts sont calculés" +" localement. ``0`` désactive l’historique tout en conservant le transfert " +"direct. Les échecs de transfert produisent des erreurs." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "Transfert direct et historique hiérarchisé" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "Le transfert direct nécessite actuellement ``n=1``." + + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"L’historique est restauré sur P avant le transfert vers D. Pour les longs " +"prompts, seul un préfixe peut être conservé ; la conservation de tout le " +"cache n’est pas garantie." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"Les transferts non acceptés expirent après 120 secondes. Si un transfert " +"expire avant que D ne l’accepte, D recalcule localement. Après acceptation, " +"le cache source reste réservé jusqu’à la fin du chargement ou à l’annulation" +" de la requête avant le chargement." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "" +"Les transferts directs et le chargement de l’historique peuvent se dérouler " +"en parallèle de la génération des réponses à d’autres requêtes." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"Les transferts non acceptés expirent après 120 secondes. Les transferts " +"acceptés expirent après 10 minutes sans nouvelle tentative de planification " +"ni progression du transfert. Les copies en cours se terminent avant la " +"libération. Les transferts expirés sont recalculés localement sur D." 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 03dba70583..567aeab68c 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 @@ -297,3 +297,84 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "Gli errori di trasferimento e i layout incompatibili della cache KV generano errori anziché ricalcolare silenziosamente la richiesta." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D usa per impostazione predefinita il trasferimento diretto tra " +"GPU. Richiede Linux, GPU NVIDIA, vLLM >= 0.21.0 e ``xoscar[nixl]>=0.11.1`` " +"negli ambienti del worker e del modello. Usare un indirizzo host " +"raggiungibile (non ``0.0.0.0``). Senza NIXL, l’avvio fallisce; non è " +"previsto il ripiego sulla CPU." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"Per P/D, ``xavier_gpu_cache_bytes`` vale per impostazione predefinita " +"``268435456`` (256 MiB per replica). La cronologia usa prima la GPU e passa " +"alla CPU quando la capacità GPU è esaurita. La cronologia CPU è limitata " +"dalla capacità dei blocchi KV del motore e può avvicinarsi a tale dimensione" +" della cache nella memoria host. Gli hit CPU ripristinano KV sulla GPU; i " +"miss vengono calcolati localmente. ``0`` disabilita la cronologia mantenendo" +" il trasferimento diretto. I guasti di trasferimento generano errori." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "Trasferimento diretto e cronologia a livelli" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "Il trasferimento diretto richiede attualmente ``n=1``." + + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"La cronologia viene ripristinata su P prima del trasferimento a D. Per i " +"prompt lunghi può essere conservato solo un prefisso; la conservazione di " +"tutta la cache non è garantita." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"I trasferimenti non accettati scadono dopo 120 secondi. Se un trasferimento " +"scade prima che D lo accetti, D ricalcola localmente. Dopo l’accettazione, " +"la cache di origine rimane riservata fino al completamento del caricamento o" +" all’annullamento della richiesta prima del caricamento." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "" +"I trasferimenti diretti e il caricamento della cronologia possono avvenire " +"in parallelo alla generazione delle risposte ad altre richieste." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"I trasferimenti non accettati scadono dopo 120 secondi. Quelli accettati " +"scadono dopo 10 minuti senza nuovi tentativi di pianificazione o avanzamento" +" del trasferimento. Le copie in corso terminano prima del rilascio. I " +"trasferimenti scaduti vengono ricalcolati localmente su D." 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 5e6b4c633f..0bf02ed39b 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 @@ -272,3 +272,72 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "転送の失敗や互換性のない KV キャッシュレイアウトは、リクエストを暗黙的に再計算するのではなく、エラーになります。" + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D はデフォルトで GPU 間の直接転送を使用します。Linux、NVIDIA GPU、vLLM >= " +"0.21.0、およびワーカーとモデルの両環境に ``xoscar[nixl]>=0.11.1`` " +"が必要です。到達可能なホストアドレスを使用してください（``0.0.0.0`` は不可）。NIXL がなければ起動に失敗し、CPU " +"へフォールバックしません。" + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"P/D の ``xavier_gpu_cache_bytes`` はデフォルトで ``268435456`` （レプリカ当たり 256 " +"MiB）です。履歴は GPU を優先し、容量不足時に CPU に退避します。CPU 履歴のブロック数はエンジン KV " +"ブロック容量が上限で、ホストメモリ使用量はそのキャッシュサイズに近づく場合があります。CPU ヒットでは KV を GPU " +"に復元し、ミスではローカル計算します。 ``0`` は履歴を無効にし、直接転送は維持します。転送失敗はエラーになります。" + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "直接転送と階層型履歴" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "直接転送は現在 ``n=1`` を必要とします。" + + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"履歴は D への引き渡し前に P で復元されます。長いプロンプトではプレフィックスのみが保持される場合があり、すべてのキャッシュの保持は保証されません。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"未受理の引き渡しは 120 秒で期限切れになります。D が受理する前に期限切れになると、D " +"がローカルで再計算します。受理後は、読み込みが完了するか、読み込み前にリクエストがキャンセルされるまで、転送元のキャッシュが予約されます。" + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "直接転送と履歴の読み込みは、他のリクエストの応答生成と並行して実行できます。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"未受理の引き渡しは 120 秒で期限切れになります。受理済みの引き渡しも、スケジュールの再試行や転送の進行が 10 " +"分間なければ期限切れになります。実行中のコピーは解放前に完了します。期限切れの引き渡しは D がローカルで再計算します。" 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 0c62940bb8..ccad95fec5 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 @@ -272,3 +272,71 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "전송 실패와 호환되지 않는 KV 캐시 레이아웃은 요청을 조용히 다시 계산하는 대신 오류를 발생시킵니다." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D는 기본적으로 GPU 간 직접 전송을 사용합니다. Linux, NVIDIA GPU, vLLM >= 0.21.0이 " +"필요하며 워커와 모델 환경 모두에 ``xoscar[nixl]>=0.11.1`` 을 설치해야 합니다. 연결 가능한 호스트 주소를 " +"사용하세요(``0.0.0.0`` 제외). NIXL이 없으면 시작에 실패하며 CPU로 대체하지 않습니다." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"P/D의 ``xavier_gpu_cache_bytes`` 기본값은 ``268435456`` (복제본당 256 MiB)입니다. 이력은 " +"GPU를 우선 사용하고 GPU 용량이 부족하면 CPU로 이동합니다. CPU 이력 블록 수는 엔진 KV 블록 용량으로 제한되며 호스트 " +"메모리 사용량이 해당 캐시 크기에 근접할 수 있습니다. CPU 적중 시 KV를 GPU로 복원하고 미스 시 로컬에서 계산합니다. ``0``" +" 은 이력을 비활성화하지만 직접 전송은 유지합니다. 전송 실패는 오류를 발생시킵니다." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "직접 전송 및 계층형 이력" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "직접 전송은 현재 ``n=1`` 을 요구합니다." + + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"기록은 D로 전달하기 전에 P에서 복원됩니다. 긴 프롬프트는 접두사만 보관될 수 있으며, 모든 캐시의 보관이 보장되지는 않습니다." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"수락되지 않은 전달은 120초 후 만료됩니다. D가 수락하기 전에 만료되면 D가 로컬에서 다시 계산합니다. D가 수락한 후에는 로드가 " +"완료되거나 로드 전에 요청이 취소될 때까지 원본 캐시가 예약됩니다." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "직접 전송과 기록 로드는 다른 요청의 응답 생성과 동시에 실행될 수 있습니다." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"수락되지 않은 전달은 120초 후 만료됩니다. 수락된 전달도 10분 동안 스케줄링 재시도나 전송 진행이 없으면 만료됩니다. 진행 중인 " +"복사는 해제 전에 완료됩니다. 만료된 전달은 D가 로컬에서 다시 계산합니다." 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 4f97da875f..c04850c3bf 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 @@ -296,3 +296,82 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "Falhas de transferência e layouts incompatíveis do cache KV geram erros em vez de recalcular a solicitação silenciosamente." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D usa transferência direta entre GPUs por padrão. Requer Linux, " +"GPUs NVIDIA, vLLM >= 0.21.0 e ``xoscar[nixl]>=0.11.1`` nos ambientes do " +"worker e do modelo. Use um endereço de host acessível (não ``0.0.0.0``). Sem" +" NIXL, a inicialização falha; não há alternativa por CPU." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"Para P/D, ``xavier_gpu_cache_bytes`` assume ``268435456`` por padrão (256 " +"MiB por réplica). O histórico usa primeiro a GPU e passa para a CPU quando a" +" capacidade GPU se esgota. O histórico CPU é limitado pela capacidade de " +"blocos KV do motor e pode ocupar um tamanho próximo ao desse cache na " +"memória do host. Acertos CPU restauram KV na GPU; ausências são calculadas " +"localmente. ``0`` desativa o histórico e mantém a transferência direta. " +"Falhas de transferência geram erros." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "Transferência direta e histórico em camadas" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "A transferência direta atualmente exige ``n=1``." + + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"O histórico é restaurado em P antes da transferência para D. Prompts longos " +"podem reter apenas um prefixo; a retenção de todo o cache não é garantida." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"Transferências não aceitas expiram após 120 segundos. Se expirarem antes de " +"D aceitá-las, D recalcula localmente. Após a aceitação, o cache de origem " +"permanece reservado até o carregamento terminar ou a solicitação ser " +"cancelada antes do carregamento." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "" +"As transferências diretas e o carregamento do histórico podem ocorrer em " +"paralelo à geração de respostas de outras solicitações." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"Transferências não aceitas expiram após 120 segundos. As aceitas expiram " +"após 10 minutos sem nova tentativa de agendamento ou progresso na " +"transferência. Cópias em andamento terminam antes da liberação. " +"Transferências expiradas são recalculadas localmente em D." 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 4f096f26b7..beb4a2832a 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 @@ -273,3 +273,68 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "传输失败或 KV 缓存布局不兼容时会报错，而不会静默地重新计算请求。" + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D 默认使用 GPU 到 GPU 直传。要求 Linux、NVIDIA GPU、vLLM >= 0.21.0，并在 worker " +"和模型环境中都安装 ``xoscar[nixl]>=0.11.1``。请使用可达的主机地址（不能是 ``0.0.0.0``）。缺少 NIXL " +"时启动失败，不回退 CPU。" + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"P/D 的 ``xavier_gpu_cache_bytes`` 默认为 ``268435456`` （每副本 256 MiB）。历史缓存优先使用 " +"GPU，GPU 容量不足时下沉 CPU。CPU 历史块数受引擎 KV 块容量限制，主机内存占用可能接近该缓存大小。CPU 命中时将 KV 恢复到 " +"GPU，未命中时本地计算。 ``0`` 禁用历史缓存，仍保留直传。传输失败时报错。" + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "直传与分层历史缓存" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "直传目前要求 ``n=1``。" + + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "历史缓存在 P 上恢复，然后交接给 D。长提示词可能只保留前缀；不保证保留所有缓存。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"尚未被接收的交接在 120 秒后过期。如果交接在 D 接收前过期，D 会在本地重新计算。D 接收后，源缓存会保留到加载完成，或请求在加载前被取消。" + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "直传和历史缓存加载可以与其他请求的响应生成重叠执行。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"尚未被接收的交接在 120 秒后过期。已被接收的交接如果 10 分钟内没有调度重试或传输进展，也会过期。正在进行的拷贝会在释放前完成。过期的交接由 D " +"在本地重新计算。" 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 75e8ee75b9..3c36be2e2d 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 @@ -268,3 +268,68 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "傳輸失敗或 KV 快取佈局不相容時會報錯，而不會默默地重新計算請求。" + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D 預設使用 GPU 到 GPU 直傳。需要 Linux、NVIDIA GPU、vLLM >= 0.21.0，並在 worker " +"和模型環境中都安裝 ``xoscar[nixl]>=0.11.1``。請使用可連線的主機位址（不能是 ``0.0.0.0``）。缺少 NIXL " +"時啟動失敗，不退回 CPU。" + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"P/D 的 ``xavier_gpu_cache_bytes`` 預設為 ``268435456`` （每個副本 256 MiB）。歷史快取優先使用 " +"GPU，GPU 容量不足時移至 CPU。CPU 歷史區塊數受引擎 KV 區塊容量限制，主機記憶體用量可能接近該快取大小。CPU 命中時將 KV 還原至 " +"GPU，未命中時在本機計算。 ``0`` 停用歷史快取，仍保留直傳。傳輸失敗時回報錯誤。" + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "直傳與分層歷史快取" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "直傳目前要求 ``n=1``。" + + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "歷史快取在 P 上還原，然後交接給 D。長提示詞可能只保留前綴；不保證保留所有快取。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"尚未被接收的交接在 120 秒後過期。如果交接在 D 接收前過期，D 會在本機重新計算。D 接收後，來源快取會保留到載入完成，或請求在載入前被取消。" + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "直接傳輸和歷史快取載入可以與其他請求的回應生成重疊執行。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"尚未被接收的交接在 120 秒後過期。已被接收的交接如果 10 分鐘內沒有排程重試或傳輸進展，也會過期。正在進行的複製會在釋放前完成。過期的交接由 D " +"在本機重新計算。" diff --git a/doc/source/user_guide/pd_separation.rst b/doc/source/user_guide/pd_separation.rst index 091b7c07fc..f335fffd6d 100644 --- a/doc/source/user_guide/pd_separation.rst +++ b/doc/source/user_guide/pd_separation.rst @@ -11,7 +11,7 @@ No enterprise package or License is required. Launch ------ -Xavier requires NVIDIA GPUs, a reachable host address (not ``0.0.0.0``), and vLLM 0.21.0 or newer. Older deployments can use the legacy Xavier integration with vLLM below 0.11.0; versions 0.11 through 0.20 are not supported. +Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL fails the launch; there is no CPU fallback. The example starts one prefill replica on GPU 0 and one decode replica on GPU 1. Replace the worker address with the full ``ip:port`` reported by @@ -129,18 +129,20 @@ Xavier V1 requires one GPU per replica (TP=1, PP=1) and text-only models without Xavier's CPU cache can consume as much memory as the engine's GPU KV cache. Allow additional host memory for transfers when sizing your deployment. -GPU-first Xavier cache (experimental) -------------------------------------- +Direct handoff and tiered history +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) for P/D. History uses GPU first and spills to CPU when GPU capacity is exhausted. CPU history is bounded by the engine KV block capacity and can approach that cache size in host memory. CPU hits restore KV to GPU; misses compute locally. ``0`` disables history, retaining direct transfer. Transfer failures raise errors. -GPU-first caching requires the Xavier backend and more than one replica. +Direct handoff currently requires ``n=1``. -For the Xavier backend, set ``xavier_gpu_cache_bytes`` to a non-negative integer, such as ``268435456`` for 256 MiB per replica. Omit it to use the existing CPU cache, or set it to ``0`` to use GPU-first caching with no GPU storage budget. This option does not select the native vLLM NIXL backend. +History is restored on P before the handoff to D. Long prompts may retain only a prefix; cache retention is best effort. -Install ``xoscar[nixl]>=0.11.1`` in both the worker and model environments. GPU-first caching requires Linux, NVIDIA CUDA, vLLM 0.21.0 or newer, a full-attention model and one GPU per replica (TP=1, PP=1). Use the same model, cache dtype and engine configuration on all replicas. +Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after 10 minutes without a scheduling retry or transfer progress. Active copies finish before release. Expired handoffs are recomputed locally on D. The budget applies only to Xavier's GPU cache. Leave additional GPU memory for the model, vLLM's own KV cache and transfers. When this cache fills, Xavier stores reusable data in CPU memory; if no cache space is available, new data may not be retained for reuse. -With GPU-first caching, cache transfers can overlap with response generation for other requests. +Direct transfers and history loads can overlap with response generation for other requests. Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request. diff --git a/xinference/core/pd_model.py b/xinference/core/pd_model.py index ae521d512d..00e3f46805 100644 --- a/xinference/core/pd_model.py +++ b/xinference/core/pd_model.py @@ -88,9 +88,11 @@ def __init__( super().__init__() # Prefill request map, used to skip the timeout task for specific request id. self._request_set: Set[str] = set() + self._direct_transfers: dict[str, dict] = {} self._model_uid = model_uid self._transport_backend = transport_backend + self._direct_handoff = transport_backend == "xavier" # 使用字典存储副本：{replica_uid: actor_ref} self._prefill_replicas: Dict[str, xo.ActorRefType["ModelActor"]] = {} @@ -244,21 +246,22 @@ async def free_prefill_model_cache(self, request_id: str): logger.debug( f"[PDModelActor] Free prefill model cache for request {request_id}" ) + handoff = self._direct_transfers.pop(request_id, None) + if handoff and handoff.get("ticket"): + try: + from ..model.llm.vllm.xavier.transfer import TransferActor + + ref = await xo.actor_ref( + address=handoff["address"], + uid=f"{TransferActor.default_uid()}-{handoff['rank']}", + ) + # D owns claimed tickets until its writes drain. Only reclaim + # a handoff that was never accepted by D (router failure/abort). + await ref.abandon_direct_gpu_v1(handoff["ticket"]) + except Exception: + logger.warning("Failed to abandon Xavier handoff", exc_info=True) if request_id in self._request_set: self._request_set.remove(request_id) - if self._prefill_replicas and self._transport_backend == "xavier": - await asyncio.gather( - *[ - actor_call( - model, - "free_model_cache", - request_id, - _rpc_operation_request_id=request_id, - ) - for model in self._prefill_replicas.values() - ], - return_exceptions=True, - ) else: logger.warning( f"[request {request_id}] Prefill model cache has been freed already" @@ -286,15 +289,9 @@ async def _infer(self, method, inputs, *args, **kwargs): args = (kwargs.pop("generate_config"),) if args and args[0] is not None and not isinstance(args[0], dict): raise TypeError("Generation config must be a dict or None") - if ( - self._transport_backend == "nixl" - and args - and args[0] - and args[0].get("n", 1) != 1 - ): - # NIXL releases the producer lease after one decoder (per TP rank) - # acknowledges the read, not after all parallel sampling children. - raise ValueError("Native NIXL PD currently requires n=1") + if args and args[0] and args[0].get("n", 1) != 1: + # Handoff leases cover one decoder, not parallel sampling children. + raise ValueError("PD KV handoff currently requires n=1") prefill = self._prefill_policy.schedule() decode = self._decode_policy.schedule() prefill_args = list(copy.deepcopy(args)) @@ -302,12 +299,11 @@ async def _infer(self, method, inputs, *args, **kwargs): prefill_args = [{}] + prefill_args[1:] prefill_args[0]["max_tokens"] = 1 prefill_args[0]["stream"] = False - if self._transport_backend == "nixl": - prefill_args[0]["n"] = 1 - prefill_args[0]["_pd_kv_transfer_params"] = { - "do_remote_decode": True, - "do_remote_prefill": False, - } + prefill_args[0]["n"] = 1 + prefill_args[0]["_pd_kv_transfer_params"] = { + "do_remote_decode": True, + "do_remote_prefill": False, + } prefill_kwargs = copy.deepcopy(kwargs) if isinstance(prefill_kwargs.get("raw_params"), dict): prefill_kwargs["raw_params"].update(max_tokens=1, stream=False) @@ -320,43 +316,31 @@ async def _infer(self, method, inputs, *args, **kwargs): if hasattr(result, "__aiter__"): async for _ in result: pass - if request_id not in self._request_set: - raise asyncio.CancelledError(f"PD request {request_id} was aborted") logger.debug( "PD prefill complete: request=%s backend=%s elapsed_s=%.6f", request_id, self._transport_backend, time.perf_counter() - prefill_start, ) - if self._transport_backend == "nixl": - payload = ( - json.loads(result) if isinstance(result, (bytes, str)) else result - ) - transfer = ( - payload.get("_pd_kv_transfer_params") - if isinstance(payload, dict) - else None - ) - if not isinstance(transfer, dict) or not transfer.get( - "do_remote_prefill" - ): - raise RuntimeError( - "NIXL prefill did not return KV transfer metadata" - ) - decode_args = list(copy.deepcopy(args)) - if not decode_args or decode_args[0] is None: - decode_args = [{}] + decode_args[1:] - decode_args[0]["_pd_kv_transfer_params"] = transfer - args = tuple(decode_args) - else: - await actor_call( - decode, - "set_unpin_handler", - self._model_uid, - request_id, - self.address, - _rpc_operation_request_id=request_id, - ) + payload = json.loads(result) if isinstance(result, (bytes, str)) else result + transfer = ( + payload.get("_pd_kv_transfer_params") + if isinstance(payload, dict) + else None + ) + if self._direct_handoff and isinstance(transfer, dict): + handoff = transfer.get("xavier_direct") + if handoff: + self._direct_transfers[request_id] = handoff + if not isinstance(transfer, dict) or not transfer.get("do_remote_prefill"): + raise RuntimeError("PD prefill did not return KV transfer metadata") + decode_args = list(copy.deepcopy(args)) + if not decode_args or decode_args[0] is None: + decode_args = [{}] + decode_args[1:] + decode_args[0]["_pd_kv_transfer_params"] = transfer + args = tuple(decode_args) + if request_id not in self._request_set: + raise asyncio.CancelledError(f"PD request {request_id} was aborted") result = await actor_call( decode, method, @@ -369,6 +353,8 @@ async def _infer(self, method, inputs, *args, **kwargs): await self.free_prefill_model_cache(request_id) raise if not hasattr(result, "__aiter__"): + # A successful D response has already completed its handoff. + self._direct_transfers.pop(request_id, None) await self.free_prefill_model_cache(request_id) return result @@ -376,6 +362,7 @@ async def stream(): try: async for chunk in result: yield chunk + self._direct_transfers.pop(request_id, None) finally: try: if hasattr(result, "aclose"): diff --git a/xinference/core/supervisor.py b/xinference/core/supervisor.py index c58025380d..1687d695f9 100644 --- a/xinference/core/supervisor.py +++ b/xinference/core/supervisor.py @@ -3104,6 +3104,8 @@ async def launch_builtin_model( gpu_cache_bytes = validate_gpu_cache_budget( kwargs.pop("xavier_gpu_cache_bytes", None), enable_xavier, replica ) + if pd_enabled and enable_xavier and gpu_cache_bytes is None: + gpu_cache_bytes = 256 * 1024 * 1024 store_address = None store_port = None world_size = None diff --git a/xinference/core/tests/test_pd_launch.py b/xinference/core/tests/test_pd_launch.py index 0fc30e7d98..c913732b9b 100644 --- a/xinference/core/tests/test_pd_launch.py +++ b/xinference/core/tests/test_pd_launch.py @@ -331,3 +331,18 @@ async def test_gpu_pool_options_reach_actual_subpool(monkeypatch): ) assert env["UCX_MEMTYPE_CACHE"] == "n" worker._ensure_subpool_monitor.assert_awaited_once_with() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("budget", [None, 0, 268435456]) +async def test_gpu_pd_defaults_to_direct_handoff(launch_runtime, budget): + supervisor, workers, actors, _ = launch_runtime + await supervisor.launch_builtin_model( + **launch_kwargs(), xavier_gpu_cache_bytes=budget + ) + assert actors["PDModelActor"].constructor_kwargs["transport_backend"] == "xavier" + for worker in workers: + config = worker.launch_builtin_model.call_args.kwargs + assert config["xavier_config"]["gpu_cache_bytes"] == ( + 268435456 if budget is None else budget + ) diff --git a/xinference/core/tests/test_pd_model.py b/xinference/core/tests/test_pd_model.py index c97b2dc0d7..27c97e60d5 100644 --- a/xinference/core/tests/test_pd_model.py +++ b/xinference/core/tests/test_pd_model.py @@ -8,7 +8,7 @@ from ..pd_model import PDModelActor, RoundRobinSchedulingPolicy from ..replica_config import ReplicaConfig, validate_pd_replica_configs -from ..rpc_context import RPC_METADATA_KEY, get_current_rpc_metadata, rpc_context +from ..rpc_context import rpc_context def test_round_robin_updates(): @@ -56,16 +56,12 @@ async def router(): model.abort_request = AsyncMock(return_value="DONE") await actor.add_prefill_actor("p", prefill) await actor.add_decode_actor("d", decode) + prefill.chat.return_value = prefill.generate.return_value = { + "_pd_kv_transfer_params": {"do_remote_prefill": True} + } return actor, prefill, decode -def _assert_free_model_cache_called(model, request_id): - model.free_model_cache.assert_awaited_once() - call = model.free_model_cache.await_args - assert call.args == (request_id,) - assert call.kwargs[RPC_METADATA_KEY]["operation_request_id"] == request_id - - @pytest.mark.asyncio @pytest.mark.parametrize("method", ["chat", "generate"]) async def test_infer_preserves_decode_config(router, method): @@ -83,7 +79,7 @@ async def test_infer_preserves_decode_config(router, method): assert d_call.args[1]["max_tokens"] == 32 assert p_call.kwargs["request_id"] == d_call.kwargs["request_id"] == "r" assert not actor._request_set - _assert_free_model_cache_called(prefill, "r") + prefill.free_model_cache.assert_not_awaited() @pytest.mark.asyncio @@ -105,54 +101,20 @@ async def chunks(): assert await anext(stream) == b"first" await stream.aclose() assert closed and not actor._request_set - _assert_free_model_cache_called(prefill, "r") + prefill.free_model_cache.assert_not_awaited() @pytest.mark.asyncio -@pytest.mark.parametrize("failure", ["prefill", "decode", "handler", "cancel"]) +@pytest.mark.parametrize("failure", ["prefill", "decode", "cancel"]) async def test_failure_releases_cache(router, failure): actor, prefill, decode = router error = asyncio.CancelledError() if failure == "cancel" else RuntimeError("failed") target = prefill.chat if failure == "prefill" else decode.chat - if failure == "handler": - target = decode.set_unpin_handler target.side_effect = error with pytest.raises(type(error)): await actor._infer("chat", [], {"max_tokens": 32}, request_id="r") assert not actor._request_set - _assert_free_model_cache_called(prefill, "r") - - -@pytest.mark.asyncio -async def test_free_prefill_model_cache_propagates_operation_metadata(): - calls = [] - - class Prefill: - @rpc_context - async def free_model_cache(self, request_id, **kwargs): - calls.append((request_id, get_current_rpc_metadata(), kwargs)) - - actor = PDModelActor("pd") - actor._request_set.add("operation-id") - actor._prefill_replicas = {"p": Prefill()} - - await actor.free_prefill_model_cache( - "operation-id", - **{ - RPC_METADATA_KEY: { - "version": 1, - "correlation_id": "http-id", - "actor_call_id": "router-call", - } - }, - ) - - request_id, metadata, kwargs = calls[0] - assert request_id == "operation-id" - assert kwargs == {} - assert metadata.correlation_id == "http-id" - assert metadata.operation_request_id == "operation-id" - assert metadata.parent_call_id == "router-call" + prefill.free_model_cache.assert_not_awaited() @pytest.mark.asyncio @@ -181,6 +143,8 @@ async def free_model_cache(self, request_id): @xo.generator @rpc_context async def chat(self, messages, config, **kwargs): + if config.get("_pd_kv_transfer_params", {}).get("do_remote_decode"): + return {"_pd_kv_transfer_params": {"do_remote_prefill": True}} if not config.get("stream"): return b"non-stream" @@ -215,7 +179,7 @@ async def test_aborted_prefill_does_not_start_decode(router): async def wait(*args, **kwargs): started.set() await finish.wait() - return b"ok" + return {"_pd_kv_transfer_params": {"do_remote_prefill": True}} prefill.chat.side_effect = wait task = asyncio.create_task(actor._infer("chat", [], {}, request_id="r")) @@ -274,16 +238,20 @@ async def test_abort_failure_still_cleans_request(router): with pytest.raises(RuntimeError, match="replica offline"): await actor.abort_request("r") assert not actor._request_set - _assert_free_model_cache_called(prefill, "r") + prefill.free_model_cache.assert_not_awaited() @pytest.mark.asyncio @pytest.mark.parametrize("method", ["chat", "generate"]) -async def test_nixl_handoff_preserves_metadata_and_decode_settings(router, method): +@pytest.mark.parametrize("backend", ["nixl", "xavier"]) +async def test_nixl_handoff_preserves_metadata_and_decode_settings( + router, backend, method +): import json actor, prefill, decode = router - actor._transport_backend = "nixl" + actor._transport_backend = backend + actor._direct_handoff = backend == "xavier" transfer = { "do_remote_prefill": True, "remote_engine_id": "engine-p", @@ -309,9 +277,11 @@ async def test_nixl_handoff_preserves_metadata_and_decode_settings(router, metho @pytest.mark.asyncio -async def test_nixl_parallel_sampling_rejected_before_prefill(router): +@pytest.mark.parametrize("backend", ["nixl", "xavier"]) +async def test_nixl_parallel_sampling_rejected_before_prefill(router, backend): actor, prefill, decode = router - actor._transport_backend = "nixl" + actor._transport_backend = backend + actor._direct_handoff = backend == "xavier" with pytest.raises(ValueError, match="n=1"): await actor._infer("chat", [], {"n": 2}) prefill.chat.assert_not_awaited() @@ -321,9 +291,13 @@ async def test_nixl_parallel_sampling_rejected_before_prefill(router): @pytest.mark.asyncio @pytest.mark.parametrize("payload", [{}, None, b"null", b"[]", b"{}", b"42"]) -async def test_nixl_missing_metadata_does_not_silently_recompute(router, payload): +@pytest.mark.parametrize("backend", ["nixl", "xavier"]) +async def test_nixl_missing_metadata_does_not_silently_recompute( + router, backend, payload +): actor, prefill, decode = router - actor._transport_backend = "nixl" + actor._transport_backend = backend + actor._direct_handoff = backend == "xavier" prefill.chat.return_value = payload with pytest.raises(RuntimeError, match="KV transfer metadata"): await actor._infer("chat", [], {}) @@ -365,9 +339,11 @@ async def test_nixl_multiple_prefillers_and_decoders(): @pytest.mark.asyncio -async def test_nixl_stream_close_releases_decode_slot(router): +@pytest.mark.parametrize("backend", ["nixl", "xavier"]) +async def test_nixl_stream_close_releases_decode_slot(router, backend): actor, prefill, decode = router - actor._transport_backend = "nixl" + actor._transport_backend = backend + actor._direct_handoff = backend == "xavier" prefill.chat.return_value = {"_pd_kv_transfer_params": {"do_remote_prefill": True}} closed = [] @@ -398,3 +374,75 @@ async def test_duplicate_generation_config_rejected_before_dispatch(router): prefill.chat.assert_not_awaited() decode.chat.assert_not_awaited() assert not actor._request_set + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["decode", "cancel", "abort_before_decode"]) +async def test_direct_router_releases_unclaimed_handoff(router, monkeypatch, failure): + actor, prefill, decode = router + actor._direct_handoff = True + transfer = { + "do_remote_prefill": True, + "xavier_direct": {"ticket": "t", "rank": 0, "address": "127.0.0.1:1234"}, + } + + async def finish_prefill(*args, **kwargs): + if failure == "abort_before_decode": + actor._request_set.discard("r") + return {"_pd_kv_transfer_params": transfer} + + prefill.chat.side_effect = finish_prefill + decode.chat.side_effect = ( + RuntimeError("decode failed") + if failure == "decode" + else asyncio.CancelledError() + ) + peer = MagicMock(abandon_direct_gpu_v1=AsyncMock()) + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer)) + with pytest.raises(RuntimeError if failure == "decode" else asyncio.CancelledError): + await actor._infer("chat", [], {}, request_id="r") + peer.abandon_direct_gpu_v1.assert_awaited_once_with("t") + assert not actor._direct_transfers and not actor._request_set + if failure == "abort_before_decode": + decode.chat.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_completed_direct_request_does_not_send_abandon_rpc(router, monkeypatch): + actor, prefill, decode = router + prefill.chat.return_value = { + "_pd_kv_transfer_params": { + "do_remote_prefill": True, + "xavier_direct": {"ticket": "t", "rank": 0, "address": "127.0.0.1:1234"}, + } + } + lookup = AsyncMock() + monkeypatch.setattr(xo, "actor_ref", lookup) + await actor._infer("chat", [], {}, request_id="r") + lookup.assert_not_awaited() + assert not actor._direct_transfers + + +@pytest.mark.asyncio +async def test_completed_direct_stream_clears_handoff_without_abandon( + router, monkeypatch +): + actor, prefill, decode = router + prefill.chat.return_value = { + "_pd_kv_transfer_params": { + "do_remote_prefill": True, + "xavier_direct": {"ticket": "t", "rank": 0, "address": "127.0.0.1:1234"}, + } + } + + async def chunks(): + yield b"one" + yield b"two" + + decode.chat.return_value = chunks() + lookup = AsyncMock() + monkeypatch.setattr(xo, "actor_ref", lookup) + stream = await actor._infer("chat", [], {"stream": True}, request_id="r") + assert [chunk async for chunk in stream] == [b"one", b"two"] + assert not actor._direct_transfers and not actor._request_set + lookup.assert_not_awaited() diff --git a/xinference/model/llm/vllm/core.py b/xinference/model/llm/vllm/core.py index a11b0d44bc..2af1782380 100644 --- a/xinference/model/llm/vllm/core.py +++ b/xinference/model/llm/vllm/core.py @@ -1995,8 +1995,12 @@ async def async_generate( sampling_params = SamplingParams(**sanitized_generate_config) if generate_config and "_pd_kv_transfer_params" in generate_config: - if self._nixl_config is None: - raise ValueError("KV handoff requires the NIXL PD backend") + from .xavier.transport import uses_direct_handoff + + if self._nixl_config is None and not uses_direct_handoff( + self._xavier_config + ): + raise ValueError("KV handoff requires NIXL or Xavier direct PD") sampling_params.extra_args = { **(sampling_params.extra_args or {}), "kv_transfer_params": generate_config["_pd_kv_transfer_params"], @@ -2178,7 +2182,14 @@ async def stream_results() -> AsyncGenerator[CompletionChunk, None]: completion = self._convert_request_output_to_completion( request_id, model=self.model_uid, request_output=final_output ) - if self._nixl_config is not None and self._nixl_config["role"] == "prefill": + from .xavier.transport import uses_direct_handoff + + if ( + self._nixl_config is not None and self._nixl_config["role"] == "prefill" + ) or ( + uses_direct_handoff(self._xavier_config) + and self._xavier_config.get("role") == "prefill" + ): completion["_pd_kv_transfer_params"] = getattr( final_output, "kv_transfer_params", None ) diff --git a/xinference/model/llm/vllm/tests/test_pd.py b/xinference/model/llm/vllm/tests/test_pd.py index a98f87d52d..3a9b6aff7e 100644 --- a/xinference/model/llm/vllm/tests/test_pd.py +++ b/xinference/model/llm/vllm/tests/test_pd.py @@ -62,10 +62,12 @@ def test_backend_selection(): @pytest.mark.parametrize("role", ["prefill", "decode"]) +@pytest.mark.parametrize("backend", ["nixl", "xavier"]) @pytest.mark.asyncio async def test_generate_passes_native_handoff_and_returns_producer_metadata( monkeypatch, role, + backend, ): from .. import core @@ -91,7 +93,10 @@ async def generate(self, prompt, params, request_id, **kwargs): monkeypatch.setattr(core, "VLLM_VERSION", Version("0.21.0")) monkeypatch.setattr(core, "VLLM_INSTALLED", False) model = object.__new__(core.VLLMModel) - model._nixl_config = {"role": role} + model._nixl_config = {"role": role} if backend == "nixl" else None + model._xavier_config = ( + {"role": role, "gpu_cache_bytes": 0} if backend == "xavier" else None + ) model._engine = Engine() model._active_request_ids = set() model.lora_requests = [] diff --git a/xinference/model/llm/vllm/xavier/README.md b/xinference/model/llm/vllm/xavier/README.md index d74065f013..26343a59d9 100644 --- a/xinference/model/llm/vllm/xavier/README.md +++ b/xinference/model/llm/vllm/xavier/README.md @@ -6,13 +6,19 @@ This document covers implementation, diagnostics and contributor validation. For Xavier and vLLM's native `NixlConnector` are separate backends. Selecting `vllm_transfer_backend_type="nixl"` uses the native connector; configuring `xavier_gpu_cache_bytes` does not select it. -The Xavier V1 connector supports vLLM 0.21.0 or newer. The legacy V0 adapter remains for vLLM below 0.11.0; versions 0.11–0.20 are unsupported. GPU integration CI pins vLLM 0.21.0 and NIXL 1.1.0. Match vLLM, Torch and NIXL versions: importing vLLM alone does not validate engine startup. +The Xavier V1 connector supports vLLM 0.21.0 or newer. P/D requires xoscar[nixl]>=0.11.1 in both worker and model environments. The legacy V0 adapter remains for non-P/D deployments with vLLM below 0.11.0; versions 0.11–0.20 are unsupported. GPU integration CI pins vLLM 0.21.0 and NIXL 1.1.0. Match vLLM, Torch and NIXL versions: importing vLLM alone does not validate engine startup. The native backend passes producer handoff metadata to the selected decoder and lets NIXL manage transfer completion and cache release. Failed handoffs fail the request rather than silently recomputing. `VLLM_NIXL_ABORT_REQUEST_TIMEOUT` bounds producer retention if a decoder never consumes a completed prefill. This path does not create Xavier collective actors or force eager execution. Parallel sampling is rejected because child requests cannot safely share one producer transfer lease. ## Cache ownership and transfer -Xavier CPU snapshots are keyed by prompt content, not reusable engine block IDs. Readers reserve complete snapshots during transfer; unavailable snapshots or cache pressure can cause a cache miss and local computation. This differs from a transfer failure, which propagates as an error. Only complete layers are published for remote reuse. CPU capacity is bounded by the engine's KV block capacity. +P/D uses request-scoped GPU-to-GPU handoff by default, with no opt-in flag. The producer keeps its engine blocks until decoder reads and optional history copies finish. Independent history uses 256 MiB of GPU storage per replica by default and spills to CPU, bounded by engine KV block capacity. `xavier_gpu_cache_bytes=0` disables history while keeping direct handoff. CPU-history hits restore to the local GPU in per-layer batches; P-to-D traffic remains xoscar NIXL. Missing NIXL fails launch. + +Unclaimed producer tickets expire after 120 seconds. D atomically claims a live ticket before reporting a remote hit; expiry is a scheduler miss, not a worker transfer error. Claimed tickets have a 600-second idle lease refreshed by scheduling retries and slab transfers. Active reads cannot expire. If a claim expires before or between slabs, the worker reports every destination block of that request through vLLM's load-error callback together with receive completion; the default Xavier recompute policy retries locally without exposing partial KV or crashing EngineCore. Other transfer/layout failures still propagate. Allocation marks the handoff consumed, so preempted decode requests recompute locally rather than rereading the ticket. Router cleanup abandons only unclaimed tickets, while allocated loads drain before release. + +History uses one writer, a 64 MiB candidate-prefix cap (also limited by total history capacity), and a 10 ms soft deadline. Long prompts retain the leading admissible prefix instead of being skipped entirely; all engine blocks remain owned until the active chunk drains. `history_capacity_limited_blocks` counts candidates beyond the cap, `history_deadline_dropped_blocks` and `history_closing_dropped_blocks` count remaining candidates abandoned at the deadline or shutdown, and `history_skipped_requests` counts a busy writer. When both tiers are full, new content must be observed again before replacing retained blocks. Admission stops at the first rejected or capacity-limited prefix block. Cache and probation LRU order retain prefix heads longer than tails. History hit counters count unique blocks actually written by successful loads, not reservations or scheduling retries. Active leases prevent eviction. History is restored on P; routing still visits P before D. Direct handoff supports `n=1`. + +The following snapshot-path details apply to hybrid deployments. Xavier CPU snapshots are keyed by prompt content, not reusable engine block IDs. Readers reserve complete snapshots during transfer; unavailable snapshots or cache pressure can cause a cache miss and local computation. This differs from a transfer failure, which propagates as an error. Only complete layers are published for remote reuse. CPU capacity is bounded by the engine's KV block capacity. In the GPU-first path, engine caches are shared with the local TransferActor through CUDA IPC. Retained GPU snapshots use xoscar NIXL; CPU snapshots use batched Gloo transfers. The worker sets `UCX_MEMTYPE_CACHE=n` and defaults `UCX_TLS` to `tcp,cuda_copy,cuda_ipc`, preserving an explicit `UCX_TLS` setting. @@ -20,7 +26,7 @@ The snapshot budget excludes the model, engine KV cache, persistent transfer buf GPU overflow demotes unleased snapshots to CPU. If both tiers have no evictable space, new snapshots are skipped. Layout mismatches and transfer failures propagate instead of silently converting data or promising recomputation. If staging fails, outstanding CUDA work is fenced before unpublished snapshots are discarded; synchronization errors still propagate. -When `xavier_gpu_cache_bytes` is set (including zero), KV loads execute asynchronously in the TransferActor so ready requests can continue decoding. Omitting the setting retains synchronous CPU loading. Cancellation holds destination blocks and source leases until writes complete. Cleanup attempts every lease release and preserves the original load exception if cleanup also fails. Shutdown drains outstanding operations before releasing buffers and snapshot storage, and logs cache placement and transfer counters. Repeated close calls share one close task; the closed runtime remains a guard against remapping. +P/D always uses asynchronous loads. In the hybrid snapshot path, when `xavier_gpu_cache_bytes` is set (including zero), KV loads execute asynchronously in the TransferActor so ready requests can continue decoding. Omitting the setting for hybrid replicas retains synchronous CPU loading. Cancellation holds destination blocks and source leases until writes complete. Cleanup attempts every lease release and preserves the original load exception if cleanup also fails. Shutdown drains outstanding operations before releasing buffers and snapshot storage, and logs cache placement and transfer counters. Repeated close calls share one close task; the closed runtime remains a guard against remapping. 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. @@ -65,6 +71,6 @@ XINFERENCE_TEST_PD_GPU=1 python -m pytest -v \ xinference/model/llm/vllm/xavier/test/test_pd_gpu.py ``` -The test runs synchronous CPU Xavier, GPU-first asynchronous Xavier 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. CPU Xavier requires producer staging and decoder KV-load log evidence; GPU-first Xavier requires successful async completion logs from `get_finished` after writes and lease cleanup; native NIXL requires `calling _read_blocks` and completed receive logs, including the concurrent requests. +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. diff --git a/xinference/model/llm/vllm/xavier/direct_handoff.py b/xinference/model/llm/vllm/xavier/direct_handoff.py new file mode 100644 index 0000000000..a365bac9c4 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/direct_handoff.py @@ -0,0 +1,267 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Request-scoped GPU handoff with independent tiered history.""" + +import asyncio +import logging +import time +from dataclasses import dataclass, field + +import torch +import xoscar as xo + +from .direct_history import DirectHistoryMixin +from .gpu_transfer import GPUTransfer +from .request_transfer import LayerRead, batch_reads, unpack_reads + +logger = logging.getLogger(__name__) + + +@dataclass +class DirectRequest: + request_id: str + blocks: set[int] + deadline: float + claimed: bool = False + reading: bool = False + released: bool = False + indices: dict[tuple[int, ...], torch.Tensor] = field(default_factory=dict) + hashes: dict[int, int] = field(default_factory=dict) + retaining: bool = False + retention_attempted: bool = False + + +class DirectGPUTransfer(DirectHistoryMixin, GPUTransfer): + def __init__(self, actor, caches, budget): + super().__init__(actor, caches, 0) + self.direct_requests: dict[str, DirectRequest] = {} + self.finished_sending: set[str] = set() + self.metrics.update( + direct_registered=0, direct_finished=0, direct_expired=0, index_uploads=0 + ) + + self._init_history(budget) + + def _index_tensor( + self, indices: dict[tuple[int, ...], torch.Tensor], values: list[int] + ) -> torch.Tensor: + # Reuse across layers and slabs within one request. Constructing a CUDA + # tensor from a Python list otherwise synchronizes an H2D copy each time. + key = tuple(values) + if key not in indices: + indices[key] = torch.tensor(values, dtype=torch.long, device=self.device) + self.metrics["index_uploads"] += 1 + return indices[key] + + async def _close(self): + self._history_closing = True + if self._history_task is not None: + await asyncio.gather(self._history_task, return_exceptions=True) + if self.history is not None: + self.metrics["history"] = self.history.stats() + await super()._close() + if self.history is not None: + for key in list(self.history.blocks): + self.history._drop(key) + self.history.leases.clear() + self.history.evicted.clear() + self._history_leases.clear() + self._history_probation.clear() + for state in self.direct_requests.values(): + state.indices.clear() + self.direct_requests.clear() + self.finished_sending.clear() + + 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 + ): + 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), + time.monotonic() + 120, + hashes=dict(zip(blocks, hashes or [])), + ) + self.metrics["direct_registered"] += 1 + + def _live_direct(self, ticket): + state = self.direct_requests.get(ticket) + if state is None or state.released: + return None + if not state.reading and time.monotonic() >= state.deadline: + self.metrics["direct_expired"] += 1 + state.retention_attempted = True + self.release_direct(ticket) + return None + return state + + def claim_direct(self, ticket): + state = self._live_direct(ticket) + if state is None: + return False + # Reclaim abandoned claims even when D dies before submitting a load. + # A scheduling retry or a completed slab refreshes this idle lease. + state.claimed = True + state.deadline = time.monotonic() + 600 + return True + + def abandon_direct(self, ticket): + state = self.direct_requests.get(ticket) + if state is not None and not state.claimed: + self.release_direct(ticket) + + def release_direct(self, ticket): + state = self.direct_requests.get(ticket) + if state is None: + return + state.released = True + if not state.reading: + if self._schedule_history(ticket, state): + return + state.indices.clear() + self.finished_sending.add(state.request_id) + del self.direct_requests[ticket] + self.metrics["direct_finished"] += 1 + + def poll_direct(self): + for ticket in list(self.direct_requests): + self._live_direct(ticket) + result = self.finished_sending + self.finished_sending = set() + return result + + async def send_direct(self, ticket, reads, remote_ref, slab_bytes): + async with self.send_lock: + state = self._live_direct(ticket) + if state is None: + return False + if slab_bytes not in self.send_buffers: + raise ValueError("Direct peer transfer slab sizes differ") + if sum(read.nbytes for read in reads) > slab_bytes: + raise ValueError("Direct transfer exceeds slab capacity") + for read in reads: + cache = self.caches[read.layer] + if ( + not set(read.keys) <= state.blocks + or cache.dtype != read.dtype + or tuple(cache.shape[1:]) != read.block_shape + ): + raise ValueError("Direct KV source layout or block IDs differ") + state.reading = True + state.deadline = time.monotonic() + 600 + try: + offset = 0 + for read in reads: + end = offset + read.nbytes + target = ( + self.send_buffer[offset:end] + .view(read.dtype) + .reshape(len(read.keys), *read.block_shape) + ) + cache = self.caches[read.layer] + torch.index_select( + cache, + 0, + self._index_tensor(state.indices, read.keys), + out=target, + ) + offset = end + await asyncio.to_thread(torch.cuda.synchronize, self.device) + await xo.copy_to([self.send_buffers[slab_bytes]], [remote_ref]) + self.metrics["wire_bytes"] += slab_bytes + self.metrics["useful_bytes"] += offset + finally: + # Even a failed gather may have queued reads of engine slots. + await asyncio.to_thread(torch.cuda.synchronize, self.device) + state.reading = False + state.deadline = time.monotonic() + 600 + if state.released: + self.release_direct(ticket) + return True + + async def load_direct(self, requests, tickets): + failed = False + try: + return await self._load_direct(requests, tickets) + except BaseException: + failed = True + raise + finally: + + results = await asyncio.gather( + *( + self.actor.release_remote_direct_gpu_v1(rank, ticket) + for ranks, ticket in zip(requests, tickets) + for rank in ranks + ), + return_exceptions=True, + ) + for result in results: + if isinstance(result, BaseException): + if not failed: + raise result + logger.warning("Direct ticket cleanup failed", exc_info=result) + + async def _load_direct(self, requests, tickets): + from .transfer import TransferActor + + self.metrics["load_calls"] += 1 + self.metrics["load_requests"] += len(requests) + invalid_blocks = set() + async with self.recv_lock: + for ranks, ticket in zip(requests, tickets): + if len(ranks) != 1: + raise ValueError("Direct handoff requires one producer") + rank, layers = next(iter(ranks.items())) + sender = await xo.actor_ref( + address=self.actor._world_addresses[rank], + uid=f"{TransferActor.default_uid()}-{rank}", + ) + indices: dict[tuple[int, ...], torch.Tensor] = {} + try: + reads = [ + LayerRead( + name, + list(mapping), + list(mapping.values()), + tuple(self.caches[name].shape[1:]), + self.caches[name].dtype, + ) + for name, mapping in layers.items() + if mapping + ] + for batch in batch_reads(reads, max_bytes=self.slab_bytes): + size = sum(read.nbytes for read in batch) + slab_bytes = min(n for n in self.recv_refs if n >= size) + available = await sender.send_direct_gpu_v1( + ticket, batch, self.recv_refs[slab_bytes], slab_bytes + ) + if not available: + # Never expose a partial request as valid KV. vLLM's + # load-error callback recomputes all its destinations. + invalid_blocks.update( + dest + for mapping in layers.values() + for dest in mapping.values() + ) + break + for read, blocks in unpack_reads( + self.recv_buffer[:size], batch + ): + cache = self.caches[read.layer] + cache[self._index_tensor(indices, read.destinations)] = ( + blocks + ) + await asyncio.to_thread(torch.cuda.synchronize, self.device) + self.metrics["gpu_batches"] += 1 + finally: + await asyncio.to_thread(torch.cuda.synchronize, self.device) + indices.clear() + return invalid_blocks diff --git a/xinference/model/llm/vllm/xavier/direct_history.py b/xinference/model/llm/vllm/xavier/direct_history.py new file mode 100644 index 0000000000..cf58f83545 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/direct_history.py @@ -0,0 +1,361 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Bounded GPU/CPU history retained after request-scoped handoff.""" + +import asyncio +import logging +import time +from collections import OrderedDict + +import torch + +from .request_transfer import LayerRead, batch_reads +from .tiered_snapshot import TieredKVSnapshotStore + +logger = logging.getLogger(__name__) + + +class HistoryStore(TieredKVSnapshotStore): + """Independent GPU-first history with bounded CPU overflow.""" + + def __init__(self, budget, block_bytes, device, cpu_capacity=0): + super().__init__(max(1, cpu_capacity), budget, block_bytes, device) + self.cpu_capacity = cpu_capacity + self.capacity = self.gpu_capacity + cpu_capacity + + def reserve(self, lease, keys): + if not set(keys).issubset(self.ready): + return False + self.leases.setdefault(lease, set()).update(keys) + # Prefix heads outlive their tails. Reservation is not a cache hit: + # scheduling can retry or abort without ever restoring any blocks. + for key in reversed(keys): + self.touch(key) + return True + + def read(self, layer, keys, device=None): + # Batch same-tier blocks before H2D instead of synchronizing once per + # block. load_history partitions reads by tier and holds their leases. + values = [self.blocks[key][layer] for key in keys] + target = torch.device("cpu") if device is None else device + if all(value.device == values[0].device for value in values): + return torch.stack(values).to(target).contiguous() + return super().read(layer, keys, device) + + +class DirectHistoryMixin: + def _init_history(self, budget): + block_bytes = sum(c[0].numel() * c.element_size() for c in self.caches.values()) + cpu_capacity = self.store.cpu_capacity + self.history = ( + HistoryStore(budget, block_bytes, self.device, cpu_capacity) + if budget + else None + ) + self._history_task = None + self._history_closing = False + self._history_leases = {} + self._history_reading = set() + # Metadata only. A second completed request demonstrates reuse before + # new content can displace existing GPU history. Empty slots need no + # probation, preserving first-request retention while there is room. + self._history_probation: OrderedDict[int, None] = OrderedDict() + self._history_probation_limit = min( + 8192, 2 * self.history.capacity if self.history else 0 + ) + self.history_chunk_bytes = 2 * 1024 * 1024 + self.history_fill_chunk_bytes = 8 * 1024 * 1024 + self.history_pending_bytes = min( + self.history.capacity * block_bytes if self.history else 0, 64 * 1024 * 1024 + ) + # A soft admission deadline, not a cancellation timeout: an in-flight + # CUDA chunk always drains before the engine can reuse its blocks. + self.history_retention_seconds = 0.01 + self.metrics.update( + history_saved_blocks=0, + history_skipped_requests=0, + history_failures=0, + history_hit_blocks=0, + history_cpu_hit_blocks=0, + history_gpu_hit_blocks=0, + history_loaded_requests=0, + history_admission_rejected_blocks=0, + history_admission_reused_blocks=0, + history_save_chunks=0, + history_capacity_limited_blocks=0, + history_deadline_dropped_blocks=0, + history_closing_dropped_blocks=0, + ) + + def _expire_history_leases(self): + now = time.monotonic() + for lease, deadline in list(self._history_leases.items()): + if deadline <= now and lease not in self._history_reading: + self.release_history(lease) + + def reserve_history(self, lease, keys): + if self.closing or self.history is None: + return [] + self._expire_history_leases() + if lease in self._history_leases: + raise ValueError("Duplicate history lease") + matched = [] + for key in keys: + if key not in self.history.ready: + break + matched.append(key) + if matched: + if not self.history.reserve(lease, matched): + return [] + self._history_leases[lease] = time.monotonic() + 120 + return matched + + def release_history(self, lease): + if lease not in self._history_reading: + self._history_leases.pop(lease, None) + if self.history is not None: + self.history.release(lease) + + def _schedule_history(self, ticket, state): + if state.retaining: + return True + if state.retention_attempted: + return False + state.retention_attempted = True + if self._history_closing or self.history is None or not state.hashes: + return False + limit = min( + self.history.capacity, + self.history_pending_bytes // self.history.block_bytes, + ) + candidates = [] + free = max(0, self.history.capacity - len(self.history.blocks)) + prefix = [] + blocked = None + for block, key in state.hashes.items(): + if blocked is not None: + if key not in self.history.ready: + self.metrics[blocked] += 1 + continue + if len(prefix) >= self.history.capacity or ( + key not in self.history.ready and len(candidates) >= limit + ): + blocked = "history_capacity_limited_blocks" + if key not in self.history.ready: + self.metrics[blocked] += 1 + elif key in self.history.ready: + prefix.append(key) + elif free or key in self._history_probation: + candidates.append((key, block)) + prefix.append(key) + free = max(0, free - 1) + if key in self._history_probation: + self.metrics["history_admission_reused_blocks"] += 1 + else: + blocked = "history_admission_rejected_blocks" + self.metrics[blocked] += 1 + # Observe the full request, but keep heads newest in both probation and + # cache LRU. A rejected head must not admit unreachable suffix blocks. + for key in reversed(list(state.hashes.values())): + if key in self.history.ready: + self.history.touch(key) + self._history_probation.pop(key, None) + else: + self._history_probation[key] = None + self._history_probation.move_to_end(key) + while len(self._history_probation) > self._history_probation_limit: + self._history_probation.popitem(last=False) + if not candidates: + return False + # Observe completed requests even when the single writer is busy, but + # never queue another writer or extend engine block ownership for it. + if self._history_task is not None and not self._history_task.done(): + self.metrics["history_skipped_requests"] += 1 + return False + state.retaining = True + deadline = time.monotonic() + self.history_retention_seconds + task = self._history_task = asyncio.create_task( + self._retain_history(ticket, state, candidates, deadline) + ) + self.tasks.add(task) + + def done(future): + self.tasks.discard(future) + if not future.cancelled() and future.exception() is not None: + exc = future.exception() + logger.error( + "GPU history retention failed", + exc_info=exc, + ) + + task.add_done_callback(done) + return True + + async def _retain_history(self, ticket, state, candidates, deadline): + pending = [] + prefix_keys = list(state.hashes.values())[: self.history.capacity] + try: + # At most one writer and one bounded chunk. Yield between chunks; + # never queue the writer ahead of an already active handoff or load. + start = 0 + while start < len(candidates): + # Filling unused capacity can amortize per-layer gathers; + # replacements use smaller chunks to limit lock hold time. + free = self.history.gpu_capacity - self.history.counts["gpu"] + chunk_bytes = ( + self.history_fill_chunk_bytes + if free > 0 + else self.history_chunk_bytes + ) + count = max( + 1, + min( + 32, + min(self.slab_bytes, chunk_bytes) // self.history.block_bytes, + ), + ) + if free > 0: + count = min(count, free) + chunk = candidates[start : start + count] + while True: + if self._history_closing or time.monotonic() >= deadline: + reason = "closing" if self._history_closing else "deadline" + self.metrics[f"history_{reason}_dropped_blocks"] += ( + len(candidates) - start + ) + return + if not self.send_lock.locked() and not self.recv_lock.locked(): + break + await asyncio.sleep(0.001) + start += len(chunk) + async with self.send_lock: + self._expire_history_leases() + pairs = [(k, b) for k, b in chunk if k not in self.history.ready] + if not pairs: + continue + pending, ids = map(list, zip(*pairs)) + index = self._index_tensor(state.indices, ids) + layers = { + name: cache.index_select(0, index) + for name, cache in self.caches.items() + } + self.history.stage_blocks(pending, layers) + await asyncio.to_thread(torch.cuda.synchronize, self.device) + available = self.history.publish(pending, set(self.caches)) + self.metrics["history_saved_blocks"] += len(available) + self.metrics["history_save_chunks"] += 1 + for key in available: + self._history_probation.pop(key, None) + for key in reversed(prefix_keys): + if key in self.history.ready: + self.history.touch(key) + complete = len(available) == len(pending) + if not complete: + self.metrics["history_capacity_limited_blocks"] += ( + len(pending) - len(available) + len(candidates) - start + ) + pending = [] + del layers, index + if not complete: + # Capacity may be pinned by another reader. Do not + # continue past a prefix we could not publish. + return + await asyncio.sleep(0) + except Exception: + self.metrics["history_failures"] += 1 + logger.warning("Skipping failed GPU history retention", exc_info=True) + finally: + # Original engine slots stay owned even after D finishes, until all + # optional snapshot reads have drained. Never publish partial data. + await asyncio.to_thread(torch.cuda.synchronize, self.device) + for key in pending: + if key in self.history.blocks and key not in self.history.ready: + self.history._drop(key) + self.history.evicted.clear() + state.retaining = False + self.release_direct(ticket) + + async def load_history(self, requests, leases): + failed = False + try: + await self._load_history(requests, leases) + except BaseException: + failed = True + raise + finally: + errors = [] + for lease in leases: + try: + self.release_history(lease) + except BaseException as error: + errors.append(error) + if errors and not failed: + raise errors[0] + for cleanup_error in errors: + logger.warning("History lease cleanup failed", exc_info=cleanup_error) + + async def _load_history(self, requests, leases): + async with self.recv_lock: + for ranks, lease in zip(requests, leases): + indices = {} + try: + if self.history is None or lease not in self._history_leases: + raise RuntimeError("GPU history lease expired") + self._history_reading.add(lease) + reads = [] + written = set() + for layers in ranks.values(): + for name, mapping in layers.items(): + cache = self.caches[name] + if not set(mapping).issubset(self.history.leases[lease]): + raise ValueError("History read outside lease") + for tier in ("gpu", "cpu"): + pairs = [ + (key, dest) + for key, dest in mapping.items() + if self.history.tiers[key] == tier + ] + if pairs: + keys, destinations = map(list, zip(*pairs)) + reads.append( + LayerRead( + name, + keys, + destinations, + tuple(cache.shape[1:]), + cache.dtype, + ) + ) + for batch in batch_reads(reads, max_bytes=self.slab_bytes): + for read in batch: + blocks = self.history.read( + read.layer, read.keys, self.device + ) + if ( + blocks.dtype != read.dtype + or tuple(blocks.shape[1:]) != read.block_shape + ): + raise ValueError("GPU history layout mismatch") + self.caches[read.layer][ + self._index_tensor(indices, read.destinations) + ] = blocks + written.update(read.keys) + await asyncio.to_thread(torch.cuda.synchronize, self.device) + hits = {"gpu": 0, "cpu": 0} + for key in written: + tier = self.history.tiers[key] + hits[tier] += 1 + self.metrics[f"history_{tier}_hit_blocks"] += 1 + self.history.metrics[f"{tier}_hits"] += 1 + self.metrics["history_hit_blocks"] += len(written) + self.metrics["history_loaded_requests"] += 1 + logger.debug( + "Restored Xavier history: blocks=%s gpu=%s cpu=%s", + len(written), + hits["gpu"], + hits["cpu"], + ) + finally: + await asyncio.to_thread(torch.cuda.synchronize, self.device) + indices.clear() + self._history_reading.discard(lease) diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index 15fb169d83..a4712d8f7d 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/gpu_transfer.py @@ -452,7 +452,7 @@ async def load_requests(self, requests): class GPUTransferMixin: - def map_gpu_caches_v1(self, descriptors, budget): + def map_gpu_caches_v1(self, descriptors, budget, direct_handoff=False): from torch.multiprocessing.reductions import rebuild_cuda_tensor if getattr(self, "_gpu_transfer", None) is not None: @@ -460,7 +460,12 @@ def map_gpu_caches_v1(self, descriptors, budget): caches = { name: rebuild_cuda_tensor(*desc) for name, desc in descriptors.items() } - runtime = GPUTransfer(self, caches, budget) + runtime_type = GPUTransfer + if direct_handoff: + from .direct_handoff import DirectGPUTransfer + + runtime_type = DirectGPUTransfer + runtime = runtime_type(self, caches, budget) self._gpu_transfer = runtime self._snapshot_store = runtime.store @@ -487,3 +492,53 @@ async def close_gpu_caches_v1(self): # Keep the closed runtime as a tombstone: close() reuses its task, # and map_gpu_caches_v1 must not re-register after shutdown/failure. await runtime.close() + + def register_direct_gpu_v1(self, ticket, request_id, blocks, hashes=None): + self._gpu_transfer.register_direct(ticket, request_id, blocks, hashes) + + def claim_direct_gpu_v1(self, ticket): + return self._gpu_transfer.claim_direct(ticket) + + def abandon_direct_gpu_v1(self, ticket): + self._gpu_transfer.abandon_direct(ticket) + + async def claim_remote_direct_gpu_v1(self, rank, ticket): + ref = await xo.actor_ref( + address=self._world_addresses[rank], uid=f"{self.default_uid()}-{rank}" + ) + return await ref.claim_direct_gpu_v1(ticket) + + def release_direct_gpu_v1(self, ticket): + self._gpu_transfer.release_direct(ticket) + + def poll_direct_gpu_v1(self): + return self._gpu_transfer.poll_direct() + + async def send_direct_gpu_v1(self, ticket, reads, remote_ref, slab_bytes): + runtime = self._gpu_transfer + return await runtime.run( + runtime.send_direct, ticket, reads, remote_ref, slab_bytes + ) + + async def load_direct_gpu_v1(self, requests, tickets): + runtime = self._gpu_transfer + return await runtime.run(runtime.load_direct, requests, tickets) + + async def release_remote_direct_gpu_v1(self, rank, ticket): + ref = await xo.actor_ref( + address=self._world_addresses[rank], uid=f"{self.default_uid()}-{rank}" + ) + await ref.release_direct_gpu_v1(ticket) + + def reserve_direct_history_v1(self, lease, keys): + runtime = getattr(self, "_gpu_transfer", None) + return runtime.reserve_history(lease, keys) if runtime is not None else [] + + def release_direct_history_v1(self, lease): + runtime = getattr(self, "_gpu_transfer", None) + if runtime is not None: + runtime.release_history(lease) + + async def load_direct_history_v1(self, requests, leases): + runtime = self._gpu_transfer + return await runtime.run(runtime.load_history, requests, leases) diff --git a/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py new file mode 100644 index 0000000000..c676d16505 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py @@ -0,0 +1,575 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +import torch +import xoscar as xo + +from ..direct_handoff import DirectGPUTransfer +from ..gpu_transfer import GPUTransferMixin +from ..request_transfer import LayerRead +from .test_gpu_transfer import runtime + + +def direct_runtime(monkeypatch): + r = runtime(monkeypatch, gpu_slots=0) + r.__class__ = DirectGPUTransfer + r.direct_requests, r.finished_sending = {}, set() + r.metrics.update( + direct_registered=0, direct_finished=0, direct_expired=0, index_uploads=0 + ) + r.actor.default_uid = lambda: "transfer" + r.actor.release_remote_direct_gpu_v1 = ( + lambda rank, ticket: GPUTransferMixin.release_remote_direct_gpu_v1( + r.actor, rank, ticket + ) + ) + r._init_history(0) + return r + + +def peer_for(source): + async def send(*args): + return await source.run(source.send_direct, *args) + + return SimpleNamespace( + send_direct_gpu_v1=AsyncMock(side_effect=send), + release_direct_gpu_v1=AsyncMock(side_effect=source.release_direct), + ) + + +@pytest.mark.asyncio +async def test_direct_handoff_preserves_bits_and_uses_no_snapshots(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + source.caches["K"].view(torch.int16).copy_( + torch.tensor( + [ + [0, -32768], + [32767, -1], + [32640, -128], + [1, -32767], + [8, 9], + [10, 11], + [12, 13], + [14, 15], + ], + dtype=torch.int16, + ) + ) + source.register_direct("ticket", "producer", [1, 2]) + peer = peer_for(source) + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer)) + + async def copy(buffers, refs): + refs[0].copy_(buffers[0]) + + monkeypatch.setattr(xo, "copy_to", copy) + await dest.run(dest.load_direct, [{0: {"K": {1: 5, 2: 3}}}], ["ticket"]) + assert torch.equal( + dest.caches["K"].view(torch.int16)[[5, 3]], + source.caches["K"].view(torch.int16)[[1, 2]], + ) + assert not source.store.blocks and not dest.store.blocks + assert source.poll_direct() == {"producer"} + assert not source.poll_direct() and not source.direct_requests + + +@pytest.mark.asyncio +async def test_release_during_read_waits_for_copy_and_fence(monkeypatch): + source = direct_runtime(monkeypatch) + source.register_direct("ticket", "producer", [1]) + entered, release = asyncio.Event(), asyncio.Event() + calls = [] + + async def copy(buffers, refs): + entered.set() + await release.wait() + refs[0].copy_(buffers[0]) + calls.append("copy") + + monkeypatch.setattr(xo, "copy_to", copy) + monkeypatch.setattr(torch.cuda, "synchronize", lambda device: calls.append("fence")) + read = LayerRead("K", [1], [0], (2,), torch.bfloat16) + task = asyncio.create_task( + source.run( + source.send_direct, "ticket", [read], torch.zeros(16, dtype=torch.uint8), 16 + ) + ) + await asyncio.wait_for(entered.wait(), 2) + state = source.direct_requests["ticket"] + source.release_direct("ticket") + assert state.indices + assert not source.poll_direct() + release.set() + await task + assert calls == ["fence", "copy", "fence"] + assert not state.indices + assert source.poll_direct() == {"producer"} + + +@pytest.mark.asyncio +async def test_expired_request_cannot_read_reused_engine_blocks(monkeypatch): + source = direct_runtime(monkeypatch) + source.register_direct("ticket", "producer", [1]) + source.direct_requests["ticket"].deadline = 0 + assert source.poll_direct() == {"producer"} + assert not await source.send_direct( + "ticket", [LayerRead("K", [1], [0], (2,), torch.bfloat16)], None, 16 + ) + assert not source.metrics["wire_bytes"] + + +@pytest.mark.asyncio +async def test_cancelled_decoder_drains_before_releasing_source(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + source.caches["K"].fill_(7) + source.register_direct("ticket", "producer", [1]) + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer_for(source))) + entered, release = asyncio.Event(), asyncio.Event() + + async def copy(buffers, refs): + entered.set() + await release.wait() + refs[0].copy_(buffers[0]) + + monkeypatch.setattr(xo, "copy_to", copy) + task = asyncio.create_task( + dest.run(dest.load_direct, [{0: {"K": {1: 5}}}], ["ticket"]) + ) + await asyncio.wait_for(entered.wait(), 2) + task.cancel() + await asyncio.sleep(0) + assert not task.done() and not source.poll_direct() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert dest.caches["K"][5].tolist() == [7, 7] + assert source.poll_direct() == {"producer"} + + +def test_direct_metadata_preserves_local_prefix_offset(connector): + connector._direct_handoff = True + request = SimpleNamespace( + request_id="decoder", + kv_transfer_params={ + "xavier_direct": { + "rank": 0, + "ticket": "ticket", + "tokens": 32, + "blocks": [2, 3], + } + }, + ) + connector._get_transfer_ref = AsyncMock( + return_value=SimpleNamespace( + claim_remote_direct_gpu_v1=AsyncMock(return_value=True) + ) + ) + assert connector.get_num_new_matched_tokens(request, 16) == (16, True) + blocks = SimpleNamespace(get_block_ids=lambda: [[4, 5]]) + connector.update_state_after_alloc(request, blocks, 16) + load = connector._requests_need_load["decoder"] + assert load.local_transfers_by_group == {0: {0: {3: 5}}} + + +@pytest.mark.parametrize("history", [False, True]) +@pytest.mark.parametrize("prompt_tokens", [33, 30]) +def test_direct_producer_retains_original_engine_blocks( + connector, monkeypatch, prompt_tokens, history +): + import sys + + monkeypatch.setitem( + sys.modules, + "vllm.v1.request", + SimpleNamespace(RequestStatus=SimpleNamespace(FINISHED_ABORTED="aborted")), + ) + connector._direct_handoff = True + connector._history_enabled = history + connector._gpu_budget = 32 if history else 0 + actor = SimpleNamespace(register_direct_gpu_v1=AsyncMock()) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = SimpleNamespace( + request_id="producer", + prompt_token_ids=list(range(prompt_tokens)), + status="finished", + kv_transfer_params={"do_remote_decode": True}, + ) + retained, metadata = connector.request_finished(request, [4, 5, 6]) + assert retained + handoff = metadata["xavier_direct"] + assert handoff["blocks"] == [4, 5] and handoff["tokens"] == prompt_tokens - 1 + history_args = () + if history: + hashes = connector._build_xavier_hashes(request.prompt_token_ids[:-1]) + history_args = ([key for key, _ in hashes[:2]],) + actor.register_direct_gpu_v1.assert_awaited_once_with( + handoff["ticket"], "producer", [4, 5], *history_args + ) + metadata = connector.build_connector_meta(SimpleNamespace(num_scheduled_tokens={})) + assert metadata.direct_sends == {"producer"} + assert not metadata.direct_store and not connector._direct_sends + + +@pytest.mark.asyncio +async def test_indices_reused_across_layers_and_slabs_but_not_requests(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + for layer in ("K", "V", "Q"): + source.caches[layer] = torch.arange(16, dtype=torch.bfloat16).reshape(8, 2) + dest.caches[layer] = torch.zeros_like(source.caches[layer]) + 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) + for ticket, destinations in (("first", [5, 3]), ("second", [3, 5])): + source.register_direct(ticket, ticket, [1, 2]) + state = source.direct_requests[ticket] + mapping = {name: dict(zip([1, 2], destinations)) for name in source.caches} + await dest.run(dest.load_direct, [{0: mapping}], [ticket]) + for layer in source.caches: + assert torch.equal( + dest.caches[layer][destinations], source.caches[layer][[1, 2]] + ) + assert not state.indices + assert source.metrics["index_uploads"] == dest.metrics["index_uploads"] == 2 + assert dest.metrics["gpu_batches"] == 4 + assert source.poll_direct() == {"first", "second"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("finish", ["release", "expire", "close"]) +async def test_request_indices_reclaimed_on_lifecycle_end(monkeypatch, finish): + source = direct_runtime(monkeypatch) + source.register_direct("ticket", "request", [1, 2]) + state = source.direct_requests["ticket"] + index = source._index_tensor(state.indices, [1, 2]) + assert source._index_tensor(state.indices, [1, 2]) is index + assert source._index_tensor(state.indices, [2, 1]).tolist() == [2, 1] + assert index.dtype == torch.long + if finish == "release": + source.release_direct("ticket") + elif finish == "expire": + state.deadline = 0 + source.poll_direct() + else: + await source.close() + assert not source.direct_requests + assert not state.indices + + +@pytest.mark.asyncio +async def test_failed_later_slab_releases_both_index_sets(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + for layer in ("V", "Q"): + source.caches[layer] = source.caches["K"].clone() + dest.caches[layer] = dest.caches["K"].clone() + source.register_direct("ticket", "producer", [1, 2]) + state = source.direct_requests["ticket"] + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer_for(source))) + copies = 0 + receiver_indices = [] + original_index = dest._index_tensor + + def index(indices, values): + receiver_indices.append(indices) + return original_index(indices, values) + + async def copy(buffers, refs): + nonlocal copies + copies += 1 + if copies == 2: + raise RuntimeError("transfer failed") + refs[0].copy_(buffers[0]) + + monkeypatch.setattr(dest, "_index_tensor", index) + monkeypatch.setattr(xo, "copy_to", copy) + mapping = {name: {1: 5, 2: 3} for name in source.caches} + with pytest.raises(RuntimeError, match="transfer failed"): + await dest.run(dest.load_direct, [{0: mapping}], ["ticket"]) + assert receiver_indices and all(not indices for indices in receiver_indices) + assert not state.indices + assert source.poll_direct() == {"producer"} + + +@pytest.mark.parametrize("role", ["prefill", "decode", "hybrid"]) +@pytest.mark.parametrize("budget", [None, 0, 268435456]) +def test_direct_handoff_selection_is_per_model(role, budget): + from ..transport import uses_direct_handoff + + assert uses_direct_handoff({"role": role, "gpu_cache_bytes": budget}) is ( + role != "hybrid" and budget is not None + ) + assert not uses_direct_handoff(None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lease", ["ticket", "history:ticket"]) +async def test_direct_submission_failure_releases_correct_lease( + connector, connector_module, lease +): + connector._direct_handoff = True + connector._ensure_gpu_cache_mapping = AsyncMock( + side_effect=RuntimeError("map failed") + ) + transfer = SimpleNamespace( + release_direct_history_v1=AsyncMock(), + release_remote_direct_gpu_v1=AsyncMock(), + release_remote_blocks_v1=AsyncMock(), + ) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + request = connector_module.XavierLoadRequest("r", {2: {1: 0}}, lease=lease) + with pytest.raises(RuntimeError, match="map failed"): + await connector._submit_gpu_requests([request]) + if lease.startswith("history:"): + transfer.release_direct_history_v1.assert_awaited_once_with(lease) + transfer.release_remote_direct_gpu_v1.assert_not_awaited() + else: + transfer.release_remote_direct_gpu_v1.assert_awaited_once_with(2, lease) + transfer.release_direct_history_v1.assert_not_awaited() + transfer.release_remote_blocks_v1.assert_not_awaited() + assert not connector._gpu_load_jobs + + +def decoder_request(): + return SimpleNamespace( + request_id="decoder", + kv_transfer_params={ + "do_remote_prefill": True, + "xavier_direct": { + "rank": 0, + "ticket": "ticket", + "tokens": 32, + "blocks": [2, 3], + }, + }, + ) + + +def test_preempted_decoder_does_not_reload_consumed_ticket(connector): + connector._direct_handoff = True + connector._is_producer = False + actor = SimpleNamespace(claim_remote_direct_gpu_v1=AsyncMock(return_value=True)) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = decoder_request() + assert connector.get_num_new_matched_tokens(request, 0) == (32, True) + connector.update_state_after_alloc( + request, SimpleNamespace(get_block_ids=lambda: [[4, 5]]), 32 + ) + connector._requests_need_load.clear() # worker consumed the scheduled metadata + assert connector.get_num_new_matched_tokens(request, 0) == (0, False) + actor.claim_remote_direct_gpu_v1.assert_awaited_once_with(0, "ticket") + assert not connector._requests_need_load + + +def test_expired_handoff_is_scheduler_miss(connector, monkeypatch): + source = direct_runtime(monkeypatch) + source.register_direct("ticket", "producer", [2, 3]) + source.direct_requests["ticket"].deadline = 0 + connector._direct_handoff = True + connector._is_producer = False + actor = SimpleNamespace( + claim_remote_direct_gpu_v1=AsyncMock( + side_effect=lambda rank, ticket: source.claim_direct(ticket) + ) + ) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = decoder_request() + assert connector.get_num_new_matched_tokens(request, 0) == (0, False) + assert not connector._requests_need_load + assert request.kv_transfer_params["do_remote_prefill"] is False + assert source.poll_direct() == {"producer"} + + +@pytest.mark.asyncio +async def test_claim_pins_queued_handoff_until_decoder_release(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + source.caches["K"].fill_(7) + source.register_direct("ticket", "producer", [2]) + assert source.claim_direct("ticket") + source.abandon_direct("ticket") # router cancellation cannot release D's read + assert not source.poll_direct() + 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) + await dest.load_direct([{0: {"K": {2: 5}}}], ["ticket"]) + assert dest.caches["K"][5].tolist() == [7, 7] + assert source.poll_direct() == {"producer"} + + +@pytest.mark.parametrize("stage", ["unqueried", "claimed", "allocated"]) +def test_abort_before_load_releases_only_unallocated_ticket(connector, stage): + connector._direct_handoff = True + connector._is_producer = False + actor = SimpleNamespace( + claim_remote_direct_gpu_v1=AsyncMock(return_value=True), + release_remote_direct_gpu_v1=AsyncMock(), + ) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = decoder_request() + if stage != "unqueried": + assert connector.get_num_new_matched_tokens(request, 0) == (32, True) + if stage == "allocated": + connector.update_state_after_alloc( + request, SimpleNamespace(get_block_ids=lambda: [[4, 5]]), 32 + ) + assert connector.request_finished(request, []) == (False, None) + if stage == "allocated": + actor.release_remote_direct_gpu_v1.assert_not_awaited() + assert "decoder" in connector._requests_need_load + else: + actor.release_remote_direct_gpu_v1.assert_awaited_once_with(0, "ticket") + + +@pytest.mark.asyncio +async def test_direct_batch_failure_releases_later_tickets_and_preserves_error( + monkeypatch, +): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + for ticket in ("a", "b"): + source.register_direct(ticket, ticket, [1]) + peer = peer_for(source) + + async def release(ticket): + source.release_direct(ticket) + if ticket == "a": + raise RuntimeError("cleanup error") + + peer.release_direct_gpu_v1.side_effect = release + peer.send_direct_gpu_v1.side_effect = ValueError("original read error") + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer)) + with pytest.raises(ValueError, match="original read error"): + await dest.load_direct([{0: {"K": {1: 5}}}] * 2, ["a", "b"]) + assert source.poll_direct() == {"a", "b"} + assert peer.release_direct_gpu_v1.await_count == 2 + + +def test_producer_metadata_gates_idle_rpc_and_fence( + connector, connector_module, monkeypatch +): + connector._direct_handoff = True + connector._is_consumer = False + metadata = connector_module.XavierConnectorMetadata() + monkeypatch.setattr( + connector, "_get_connector_metadata", lambda: metadata, raising=False + ) + sync = Mock() + monkeypatch.setattr(torch.cuda, "synchronize", sync) + connector._ensure_gpu_cache_mapping = AsyncMock() + connector._transfer_ref = SimpleNamespace( + poll_direct_gpu_v1=AsyncMock(return_value={"p"}) + ) + connector.start_load_kv(None) + connector.wait_for_save() + assert connector.get_finished(set()) == (set(), set()) + sync.assert_not_called() + connector._transfer_ref.poll_direct_gpu_v1.assert_not_awaited() + metadata.direct_sends = {"p"} + metadata.direct_store = True + connector.start_load_kv(None) + connector.wait_for_save() + assert connector.get_finished(set()) == ({"p"}, set()) + assert connector.get_finished(set()) == (set(), set()) + sync.assert_called_once() + connector._transfer_ref.poll_direct_gpu_v1.assert_awaited_once() + + +def test_claimed_but_never_read_ticket_expires(monkeypatch): + source = direct_runtime(monkeypatch) + source.register_direct("t", "p", [1]) + assert source.claim_direct("t") + source.direct_requests["t"].deadline = 0 + assert source.poll_direct() == {"p"} + assert not source.direct_requests + assert source.metrics["direct_expired"] == 1 + + +@pytest.mark.asyncio +async def test_active_read_cannot_expire_and_refreshes_idle_lease(monkeypatch): + source = direct_runtime(monkeypatch) + source.register_direct("t", "p", [1]) + assert source.claim_direct("t") + entered, finish = asyncio.Event(), asyncio.Event() + + async def copy(buffers, refs): + entered.set() + await finish.wait() + + monkeypatch.setattr(xo, "copy_to", copy) + task = asyncio.create_task( + source.send_direct( + "t", [LayerRead("K", [1], [0], (2,), torch.bfloat16)], None, 16 + ) + ) + await entered.wait() + source.direct_requests["t"].deadline = 0 + assert not source.poll_direct() + finish.set() + assert await task + assert source.direct_requests["t"].deadline > 0 + assert not source.poll_direct() + + +@pytest.mark.asyncio +async def test_expiry_between_slabs_reports_all_destinations_for_recompute(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + for name in ("V", "Q"): + source.caches[name] = source.caches["K"].clone() + dest.caches[name] = dest.caches["K"].clone() + source.register_direct("t", "p", [1, 2]) + assert source.claim_direct("t") + peer = peer_for(source) + original_send = peer.send_direct_gpu_v1.side_effect + + async def send(*args): + available = await original_send(*args) + if "t" in source.direct_requests: + source.direct_requests["t"].deadline = 0 + return available + + peer.send_direct_gpu_v1.side_effect = send + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer)) + + async def copy(buffers, refs): + refs[0].copy_(buffers[0]) + + monkeypatch.setattr(xo, "copy_to", copy) + invalid = await dest.load_direct( + [{0: {name: {1: 5, 2: 3} for name in source.caches}}], ["t"] + ) + assert invalid == {3, 5} + assert source.metrics["direct_expired"] == 1 + assert source.poll_direct() == {"p"} + + +def test_worker_reports_expired_destinations_with_finished_request(connector): + async def expired(): + return {3, 5} + + async def submit(): + task = asyncio.create_task(expired()) + connector._gpu_load_jobs[task] = ["decoder"] + await asyncio.sleep(0) + + connector._call(submit()) + assert connector.get_finished(set()) == (set(), {"decoder"}) + assert connector.get_block_ids_with_load_errors() == {3, 5} + assert connector.get_block_ids_with_load_errors() == set() + + +def test_full_local_hit_releases_unused_direct_ticket(connector): + connector._direct_handoff = True + actor = SimpleNamespace(release_remote_direct_gpu_v1=AsyncMock()) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = decoder_request() + assert connector.get_num_new_matched_tokens(request, 32) == (0, False) + actor.release_remote_direct_gpu_v1.assert_awaited_once_with(0, "ticket") + assert request.kv_transfer_params["do_remote_prefill"] is False + assert not connector._requests_need_load diff --git a/xinference/model/llm/vllm/xavier/test/test_direct_history.py b/xinference/model/llm/vllm/xavier/test/test_direct_history.py new file mode 100644 index 0000000000..3a2d7db578 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_direct_history.py @@ -0,0 +1,632 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +import asyncio +import threading + +import pytest +import torch + +from ..direct_history import HistoryStore +from .test_direct_handoff import direct_runtime + + +def history_runtime(monkeypatch): + r = direct_runtime(monkeypatch) + r.store.cpu_capacity = 0 + r._init_history(32) + r.history_retention_seconds = 10 + r.caches["K"].copy_(torch.arange(16).reshape(8, 2)) + return r + + +@pytest.mark.asyncio +async def test_snapshot_waits_for_fence_and_survives_engine_reuse(monkeypatch): + r = history_runtime(monkeypatch) + entered, release = threading.Event(), threading.Event() + + def fence(*args): + entered.set() + assert release.wait(5) + + monkeypatch.setattr(torch.cuda, "synchronize", fence) + original = r.caches["K"][[1, 2]].clone() + r.register_direct("t", "p", [1, 2], [101, 102]) + r.release_direct("t") + try: + assert await asyncio.to_thread(entered.wait, 2) + assert not r.history.ready + assert not r.poll_direct() + assert r.direct_requests["t"].retaining + finally: + release.set() + await r._history_task + assert r.poll_direct() == {"p"} + r.caches["K"].zero_() + assert r.reserve_history("h", [101, 102, 999]) == [101, 102] + await r.load_history([{0: {"K": {101: 5, 102: 3}}}], ["h"]) + assert torch.equal(r.caches["K"][[5, 3]], original) + assert not r.history.leases and not r._history_leases + assert r.metrics["history_loaded_requests"] == 1 + assert r.history.stats()["cpu_blocks"] == 0 + + +@pytest.mark.asyncio +async def test_busy_writer_skips_next_request_without_backlog(monkeypatch): + r = history_runtime(monkeypatch) + await r.send_lock.acquire() + r.register_direct("a", "a", [1], [101]) + r.release_direct("a") + writer = r._history_task + r.register_direct("b", "b", [2], [102]) + r.release_direct("b") + assert r.poll_direct() == {"b"} + assert r._history_task is writer + assert r.metrics["history_skipped_requests"] == 1 + r.send_lock.release() + await writer + assert r.poll_direct() == {"a"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reason", ["deadline", "expired"]) +async def test_optional_retention_releases_engine_when_skipped(monkeypatch, reason): + r = history_runtime(monkeypatch) + r.register_direct("t", "p", [1], [101]) + if reason == "deadline": + r.history_retention_seconds = 0 + if reason == "expired": + r.direct_requests["t"].deadline = 0 + assert r.poll_direct() == {"p"} + else: + r.release_direct("t") + if r._history_task: + await r._history_task + assert r.poll_direct() == {"p"} + assert not r.history.blocks and not r.direct_requests + + +def test_gpu_eviction_respects_lease_without_cpu_demotion(): + store = HistoryStore(8, 4, torch.device("cpu")) + + def save(keys): + store.stage_blocks(keys, {"K": torch.ones(len(keys), 2, dtype=torch.bfloat16)}) + store.publish(keys, {"K"}) + + save([1, 2]) + assert store.reserve("h", [1]) + save([3]) + assert store.ready == {1, 3} + assert store.reserve("j", [3]) + save([4]) + assert store.ready == {1, 3} + assert store.stats()["demotions"] == store.stats()["cpu_blocks"] == 0 + assert store.stats()["gpu_reserved_bytes"] == 8 + + +@pytest.mark.asyncio +async def test_failed_staging_never_publishes_partial_snapshot(monkeypatch): + r = history_runtime(monkeypatch) + original = r.history.stage_blocks + + def fail(keys, layers): + original(keys, layers) + raise RuntimeError("staging failed") + + monkeypatch.setattr(r.history, "stage_blocks", fail) + r.register_direct("t", "p", [1, 2], [101, 102]) + r.release_direct("t") + await r._history_task + assert not r.history.ready and not r.history.blocks + assert r.poll_direct() == {"p"} + assert r.metrics["history_failures"] == 1 + + +@pytest.mark.asyncio +async def test_invalid_history_read_releases_lease(monkeypatch): + r = history_runtime(monkeypatch) + r.register_direct("t", "p", [1], [101]) + r.release_direct("t") + await r._history_task + r.reserve_history("h", [101]) + with pytest.raises(ValueError, match="outside lease"): + await r.load_history([{0: {"K": {102: 5}}}], ["h"]) + assert not r._history_leases and not r.history.leases + + +def test_producer_history_maps_prefix_after_local_hit(connector): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + connector._direct_handoff = connector._history_enabled = True + actor = SimpleNamespace( + reserve_direct_history_v1=AsyncMock(), release_direct_history_v1=AsyncMock() + ) + request = SimpleNamespace(request_id="p", prompt_token_ids=list(range(65))) + keys = [ + key for key, _ in connector._build_xavier_hashes(request.prompt_token_ids[:64]) + ] + actor.reserve_direct_history_v1.return_value = keys[1:3] + connector._get_transfer_ref = AsyncMock(return_value=actor) + assert connector.get_num_new_matched_tokens(request, 16) == (32, True) + load = connector._requests_need_load["p"] + actor.reserve_direct_history_v1.assert_awaited_once_with(load.lease, keys[1:]) + connector.update_state_after_alloc( + request, SimpleNamespace(get_block_ids=lambda: [[4, 5, 6, 7]]), 32 + ) + assert connector._requests_need_load["p"].local_transfers_by_group == { + 0: {connector._rank: {keys[1]: 5, keys[2]: 6}} + } + + +def test_unallocated_history_abort_releases_lease(connector, monkeypatch): + import sys + from types import SimpleNamespace + from unittest.mock import AsyncMock + + monkeypatch.setitem( + sys.modules, + "vllm.v1.request", + SimpleNamespace(RequestStatus=SimpleNamespace(FINISHED_ABORTED="aborted")), + ) + connector._direct_handoff = connector._history_enabled = True + actor = SimpleNamespace( + reserve_direct_history_v1=AsyncMock(return_value=[101]), + release_direct_history_v1=AsyncMock(), + ) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = SimpleNamespace( + request_id="p", + prompt_token_ids=list(range(33)), + status="aborted", + kv_transfer_params={"do_remote_decode": True}, + ) + assert connector.get_num_new_matched_tokens(request, 0) == (16, True) + lease = connector._requests_need_load["p"].lease + assert connector.request_finished(request, []) == (False, None) + actor.release_direct_history_v1.assert_awaited_once_with(lease) + assert not connector._requests_need_load + + +@pytest.mark.asyncio +async def test_close_stops_waiting_writer_and_reclaims_snapshots(monkeypatch): + r = history_runtime(monkeypatch) + await r.recv_lock.acquire() + r.register_direct("t", "p", [1], [101]) + r.release_direct("t") + await r.close() + assert not r.direct_requests and not r.history.blocks + assert not r._history_leases + r.recv_lock.release() + + +@pytest.mark.asyncio +async def test_cancelled_restore_holds_lease_until_cuda_drains(monkeypatch): + r = history_runtime(monkeypatch) + r.register_direct("t", "p", [1], [101]) + r.release_direct("t") + await r._history_task + r.reserve_history("h", [101]) + entered, release = threading.Event(), threading.Event() + + def fence(*args): + entered.set() + assert release.wait(5) + + monkeypatch.setattr(torch.cuda, "synchronize", fence) + task = asyncio.create_task(r.run(r.load_history, [{0: {"K": {101: 5}}}], ["h"])) + try: + assert await asyncio.to_thread(entered.wait, 2) + task.cancel() + await asyncio.sleep(0) + r._history_leases["h"] = 0 + r._expire_history_leases() + r.release_history("h") + assert "h" in r.history.leases and not task.done() + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert not r.history.leases and not r._history_leases + assert r.caches["K"][5].tolist() == [2, 3] + + +def test_local_prefix_partial_tail_does_not_query_history(connector): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + connector._direct_handoff = connector._history_enabled = True + connector._get_transfer_ref = AsyncMock() + request = SimpleNamespace(request_id="p", prompt_token_ids=list(range(41))) + assert connector.get_num_new_matched_tokens(request, 32) == (0, False) + connector._get_transfer_ref.assert_not_called() + assert not connector._requests_need_load + + +@pytest.mark.asyncio +async def test_full_history_preserves_hot_content_until_new_content_repeats( + monkeypatch, +): + r = history_runtime(monkeypatch) + r.store.cpu_capacity = 0 + r._init_history(8) + r.history_retention_seconds = 10 + + async def finish(ticket, blocks, hashes): + r.register_direct(ticket, ticket, blocks, hashes) + r.release_direct(ticket) + if r._history_task: + await r._history_task + assert r.poll_direct() == {ticket} + + await finish("warm", [1, 2], [101, 102]) + assert r.history.ready == {101, 102} + await finish("cold", [3, 4], [103, 104]) + assert r.history.ready == {101, 102} + assert r.metrics["history_admission_rejected_blocks"] == 2 + await finish("repeat", [3, 4], [103, 104]) + assert r.history.ready == {103, 104} + assert r.reserve_history("hit", [103, 104]) == [103, 104] + assert not r._history_probation + assert r.metrics["history_saved_blocks"] == 4 + assert r.history.stats()["cpu_blocks"] == 0 + + +@pytest.mark.asyncio +async def test_probation_is_bounded_and_forgets_old_cold_content(monkeypatch): + r = history_runtime(monkeypatch) + r.store.cpu_capacity = 0 + r._init_history(4) + r.history_retention_seconds = 10 + for key in (100, 101, 102, 103, 101): + r.register_direct(str(key), str(key), [1], [key]) + r.release_direct(str(key)) + if r._history_task: + await r._history_task + r.poll_direct() + assert r.history.ready == {100} + assert list(r._history_probation) == [103, 101] + await r.close() + assert not r._history_probation + + +@pytest.mark.asyncio +async def test_snapshot_chunks_respect_byte_budget(monkeypatch): + r = history_runtime(monkeypatch) + r.history_chunk_bytes = r.history_fill_chunk_bytes = 8 + r.register_direct("t", "p", [1, 2, 3, 4, 5], [101, 102, 103, 104, 105]) + original = r.history.stage_blocks + sizes = [] + + def stage(keys, layers): + sizes.append(sum(t.numel() * t.element_size() for t in layers.values())) + return original(keys, layers) + + monkeypatch.setattr(r.history, "stage_blocks", stage) + r.release_direct("t") + await r._history_task + assert sizes == [8, 8, 4] + assert r.metrics["history_save_chunks"] == 3 + assert r.history.ready == {101, 102, 103, 104, 105} + assert r.poll_direct() == {"p"} + + +@pytest.mark.asyncio +async def test_fill_batches_grow_but_replacement_batches_stay_small(monkeypatch): + r = history_runtime(monkeypatch) + r.store.cpu_capacity = 0 + r._init_history(20) + r.history_retention_seconds = 10 + r.history_fill_chunk_bytes = 16 + r.history_chunk_bytes = 8 + sizes = [] + original = r.history.stage_blocks + + def stage(keys, layers): + sizes.append(len(keys)) + return original(keys, layers) + + monkeypatch.setattr(r.history, "stage_blocks", stage) + for ticket, hashes in ( + ("fill", list(range(101, 106))), + ("observe", list(range(201, 206))), + ("replace", list(range(201, 206))), + ): + r.register_direct(ticket, ticket, [1, 2, 3, 4, 5], hashes) + r.release_direct(ticket) + if r._history_task: + await r._history_task + assert r.poll_direct() == {ticket} + assert sizes == [4, 1, 2, 2, 1] + assert r.history.ready == set(range(201, 206)) + + +@pytest.mark.parametrize("matched_blocks", [1, 161, 277, 278]) +def test_fixed_history_prefix_length_and_destination_mapping(connector, matched_blocks): + """Keep the 4441-token divergence reproduction's exact prefix boundaries.""" + from types import SimpleNamespace + from unittest.mock import AsyncMock + + connector._direct_handoff = connector._history_enabled = True + request = SimpleNamespace(request_id="p", prompt_token_ids=list(range(4441))) + keys = [ + key for key, _ in connector._build_xavier_hashes(request.prompt_token_ids[:-1]) + ] + actor = SimpleNamespace( + reserve_direct_history_v1=AsyncMock(return_value=keys[:matched_blocks]) + ) + connector._get_transfer_ref = AsyncMock(return_value=actor) + expected = min(matched_blocks * 16, 4440) + assert connector.get_num_new_matched_tokens(request, 0) == (expected, True) + # Non-contiguous physical allocation must preserve logical token order. + destinations = list(reversed(range(278))) + connector.update_state_after_alloc( + request, SimpleNamespace(get_block_ids=lambda: [destinations]), expected + ) + load = connector._requests_need_load["p"] + assert load.local_transfers_by_group == { + 0: {connector._rank: dict(zip(keys[:matched_blocks], destinations))} + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("prefix_blocks", [1, 161, 277, 278]) +@pytest.mark.parametrize("gpu_blocks", [5, 278]) +async def test_fixed_prefix_restore_preserves_raw_bits( + monkeypatch, device, prefix_blocks, gpu_blocks +): + """Real CUDA fences when available; CPU CI exercises the same history path.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("requires a CUDA GPU") + synchronize = torch.cuda.synchronize + r = direct_runtime(monkeypatch) + if device == "cuda": + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + r.device = torch.device(device) + # 24 layer tensors, each including both K and V, with arbitrary BF16 bits. + generator = torch.Generator().manual_seed(42) + r.caches = { + str(layer): torch.randint( + -32768, 32768, (280, 2, 16, 2), dtype=torch.int16, generator=generator + ) + .view(torch.bfloat16) + .to(device) + for layer in range(24) + } + originals = { + name: cache[:prefix_blocks].clone() for name, cache in r.caches.items() + } + block_bytes = sum( + cache[0].numel() * cache.element_size() for cache in r.caches.values() + ) + r.store.cpu_capacity = 278 + r._init_history(gpu_blocks * block_bytes) + r.history_retention_seconds = 30 + r.slab_bytes = 8 * 1024 * 1024 + keys = list(range(1000, 1000 + prefix_blocks)) + r.register_direct("save", "request", list(range(prefix_blocks)), keys) + r.release_direct("save") + await r._history_task + assert r.poll_direct() == {"request"} + assert r.metrics["history_failures"] == 0 + # Engine reuses every original slot before restoring to a reversed mapping. + for cache in r.caches.values(): + cache.zero_() + assert r.reserve_history("restore", keys + [-1]) == keys + destinations = list(reversed(range(2, prefix_blocks + 2))) + await r.load_history( + [{0: {name: dict(zip(keys, destinations)) for name in r.caches}}], ["restore"] + ) + for name, cache in r.caches.items(): + assert torch.equal( + cache[destinations].view(torch.int16), originals[name].view(torch.int16) + ) + assert not cache[:2].count_nonzero() + assert not r.history.leases and not r._history_leases + assert r.history.stats()["cpu_blocks"] == max(0, prefix_blocks - gpu_blocks) + assert r.metrics["history_cpu_hit_blocks"] == max(0, prefix_blocks - gpu_blocks) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +async def test_history_spills_to_cpu_and_restores_mixed_prefix(monkeypatch, device): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("requires a CUDA GPU") + synchronize = torch.cuda.synchronize + r = direct_runtime(monkeypatch) + if device == "cuda": + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + r.device = torch.device(device) + r.caches = { + name: (torch.arange(64, dtype=torch.int16).reshape(8, 2, 4) + shift) + .view(torch.bfloat16) + .to(device) + for name, shift in [("K", 32620), ("V", -32768)] + } + block_bytes = 32 + r.store.cpu_capacity = 2 + r._init_history(2 * block_bytes) + r.history_retention_seconds = 10 + original = {name: cache[:4].clone() for name, cache in r.caches.items()} + r.register_direct("save", "p", [0, 1, 2, 3], [100, 101, 102, 103]) + r.release_direct("save") + await r._history_task + assert r.poll_direct() == {"p"} + assert r.history.tiers[100] == "gpu" + assert r.history.counts == {"gpu": 2, "cpu": 2} + assert r.history.stats()["demotions"] == 2 + assert r.reserve_history("read", [100, 101, 102, 103, 999]) == [100, 101, 102, 103] + for cache in r.caches.values(): + cache.zero_() + await r.load_history( + [{0: {name: {100: 7, 101: 4, 102: 6, 103: 5} for name in r.caches}}], ["read"] + ) + for name, cache in r.caches.items(): + assert torch.equal( + cache[[7, 4, 6, 5]].view(torch.int16), original[name].view(torch.int16) + ) + assert not cache[:4].count_nonzero() + assert r.metrics["history_cpu_hit_blocks"] == 2 + assert r.metrics["history_gpu_hit_blocks"] == 2 + assert not r.history.leases and not r._history_leases + + +def test_history_eviction_only_after_both_tiers_full_and_preserves_leases(): + store = HistoryStore(8, 4, torch.device("cpu"), cpu_capacity=2) + + def save(key): + store.stage_blocks( + [key], {"K": torch.tensor([[key, -key]], dtype=torch.bfloat16)} + ) + store.publish([key], {"K"}) + + for key in [1, 2, 3, 4]: + save(key) + assert store.ready == {1, 2, 3, 4} + assert store.reserve("reader", [1, 3]) + save(5) + assert store.ready == {1, 3, 4, 5} + assert store.tiers[1] == "cpu" and store.tiers[3] == "gpu" + assert store.counts == {"gpu": 2, "cpu": 2} + assert store.reserve("all", list(store.ready)) + save(6) + assert store.ready == {1, 3, 4, 5} + + +@pytest.mark.asyncio +async def test_long_request_retains_bounded_leading_prefix(monkeypatch): + r = history_runtime(monkeypatch) + # Simulate a request larger than the retention cap, without allocating 8B KV. + r.history_pending_bytes = 2 * r.history.block_bytes + r.register_direct("t", "p", [0, 1, 2, 3], [100, 101, 102, 103]) + r.release_direct("t") + await r._history_task + assert r.history.ready == {100, 101} + assert r.metrics["history_capacity_limited_blocks"] == 2 + assert not r.metrics["history_skipped_requests"] + assert r.poll_direct() == {"p"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("closing", [False, True]) +async def test_abandoned_retention_counts_remaining_candidates(monkeypatch, closing): + r = history_runtime(monkeypatch) + r.register_direct("t", "p", [1, 2], [101, 102]) + r.history_retention_seconds = 0 + r.release_direct("t") + r._history_closing = closing + await r._history_task + reason = "closing" if closing else "deadline" + assert r.metrics[f"history_{reason}_dropped_blocks"] == 2 + assert r.poll_direct() == {"p"} + + +@pytest.mark.asyncio +async def test_history_batch_failure_releases_all_leases_preserving_error(monkeypatch): + r = history_runtime(monkeypatch) + r.register_direct("t", "p", [1, 2], [101, 102]) + r.release_direct("t") + await r._history_task + assert r.reserve_history("a", [101]) == [101] + assert r.reserve_history("b", [102]) == [102] + original = r.release_history + + def release(lease): + original(lease) + if lease == "a": + raise RuntimeError("cleanup error") + + monkeypatch.setattr(r, "release_history", release) + with pytest.raises(ValueError, match="outside lease"): + await r.load_history([{0: {"K": {999: 4}}}, {0: {"K": {102: 5}}}], ["a", "b"]) + assert not r.history.leases and not r._history_leases + + +@pytest.mark.asyncio +async def test_prefix_heads_survive_new_admission_and_reservation(monkeypatch): + r = history_runtime(monkeypatch) + r.store.cpu_capacity = 0 + r._init_history(12) # three blocks + r.history_retention_seconds = 10 + r.register_direct("a", "a", [0, 1, 2], [100, 101, 102]) + r.release_direct("a") + await r._history_task + assert list(r.history.blocks) == [102, 101, 100] + assert r.reserve_history("lease", [100, 101, 102]) == [100, 101, 102] + r.release_history("lease") + r._history_probation[200] = None + r.register_direct("b", "b", [3], [200]) + r.release_direct("b") + await r._history_task + assert r.reserve_history("read", [100, 101, 102]) == [100, 101] + + +@pytest.mark.asyncio +async def test_rejected_head_never_admits_repeated_tail(monkeypatch): + r = history_runtime(monkeypatch) + r._init_history(8) + r.history_retention_seconds = 10 + r.register_direct("a", "a", [0, 1], [100, 101]) + r.release_direct("a") + await r._history_task + r._history_probation[202] = None + r.register_direct("b", "b", [2, 3, 4], [200, 201, 202]) + r.release_direct("b") + assert r.history.ready == {100, 101} + assert r.metrics["history_admission_rejected_blocks"] == 3 + # A request longer than probation capacity retains its head, not its tail. + r._history_probation_limit = 2 + r.register_direct("c", "c", [2, 3, 4], [300, 301, 302]) + r.release_direct("c") + assert list(r._history_probation) == [301, 300] + + +@pytest.mark.asyncio +async def test_history_hit_counts_only_written_keys_once(monkeypatch): + r = history_runtime(monkeypatch) + r.register_direct("a", "a", [0, 1, 2], [100, 101, 102]) + r.release_direct("a") + await r._history_task + for lease in ("retry", "load"): + assert r.reserve_history(lease, [100, 101, 102]) == [100, 101, 102] + if lease == "retry": + r.release_history(lease) + assert r.history.metrics["gpu_hits"] == 0 + await r.load_history([{0: {"K": {101: 5}}}], ["load"]) + assert r.metrics["history_hit_blocks"] == 1 + assert r.metrics["history_gpu_hit_blocks"] == 1 + assert r.history.metrics["gpu_hits"] == 1 + + +def test_prefill_retry_releases_previous_history_reservation(connector): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + connector._direct_handoff = connector._history_enabled = True + actor = SimpleNamespace( + reserve_direct_history_v1=AsyncMock(return_value=[123]), + release_direct_history_v1=AsyncMock(), + ) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = SimpleNamespace( + request_id="p", prompt_token_ids=list(range(33)), kv_transfer_params={} + ) + assert connector.get_num_new_matched_tokens(request, 0) == (16, True) + previous = connector._requests_need_load["p"].lease + assert connector.get_num_new_matched_tokens(request, 0) == (16, True) + actor.release_direct_history_v1.assert_awaited_once_with(previous) + assert connector._requests_need_load["p"].lease != previous + + +@pytest.mark.asyncio +async def test_rejected_prefix_does_not_count_existing_tail_as_dropped(monkeypatch): + r = history_runtime(monkeypatch) + r._init_history(8) + r.history_retention_seconds = 10 + r.register_direct("a", "a", [0, 1], [100, 101]) + r.release_direct("a") + await r._history_task + r.register_direct("b", "b", [2, 3], [200, 101]) + r.release_direct("b") + assert r.metrics["history_admission_rejected_blocks"] == 1 + assert r.history.ready == {100, 101} diff --git a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py index ab9c04d283..d427560e2a 100644 --- a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py @@ -1490,3 +1490,20 @@ def fail(keys, layers): assert set(r.store.blocks) == r.store.ready == {7} assert r.store._block_sizes == {7: 4} assert_gpu_lru_consistent(r.store) + + +@pytest.mark.parametrize("version,nixl", [("0.11.0", True), ("0.11.1", False)]) +def test_gpu_transfer_missing_dependency_fails_without_cpu_fallback( + monkeypatch, version, nixl +): + import importlib.metadata + import importlib.util + + from ..transport import gpu_pool_options + + monkeypatch.setattr(importlib.metadata, "version", lambda name: version) + monkeypatch.setattr( + importlib.util, "find_spec", lambda name: object() if nixl else None + ) + with pytest.raises(RuntimeError, match=r"xoscar\[nixl\]>=0.11.1"): + gpu_pool_options("10.0.0.1:1234", {}) 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 f9bd64f5a4..8b743f03ed 100644 --- a/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py +++ b/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py @@ -71,8 +71,8 @@ def pd_cluster(monkeypatch, tmp_path, backend): @pytest.mark.parametrize( "backend,gpu_cache_bytes", [ - pytest.param("xavier", None, id="xavier-cpu"), - pytest.param("xavier", 256 * 1024**2, id="xavier-gpu-async"), + pytest.param("xavier", None, id="xavier-default-history"), + pytest.param("xavier", 0, id="xavier-no-history"), pytest.param("nixl", None, id="nixl"), ], ) @@ -115,6 +115,7 @@ def test_pd_gpu(pd_cluster, backend, gpu_cache_bytes): max_model_len=2048, gpu_memory_utilization=0.5, dtype="float16", + enable_prefix_caching=False, replica_config=[ { "role": role, @@ -164,18 +165,23 @@ def log_since(offset): responses[stream] = chat(str(stream), stream) evidence = log_since(offset) # Require actual KV transfer, not merely a successful recomputation. - if gpu_cache_bytes is not None: - assert "Finished Xavier async KV load: request=" in evidence - elif backend == "xavier": - assert "Stage Xavier V1 blocks" in evidence - assert "Load Xavier V1 blocks" in evidence + if backend == "xavier": + assert "Register Xavier direct handoff" in evidence + assert "Finished Xavier async KV load" in evidence else: assert "calling _read_blocks" in evidence assert re.search(r"and [1-9]\d* requests done recving", evidence) - # Repeated prompts may hit decode's local prefix cache, but must still - # yield the same deterministic answer without stale or corrupted KV. + # Compare exact text only with fixed zero-prefix computation. Historical + # prefix reuse can change floating-point shapes and greedy output. for stream in (False, True): - assert chat(str(stream), stream) == responses[stream] + offset = Path(log_path).stat().st_size + repeated = chat(str(stream), stream) + if backend == "nixl" or gpu_cache_bytes == 0: + assert repeated == responses[stream] + else: + assert re.search( + r"Restored Xavier history: blocks=[1-9]\d*", log_since(offset) + ) offset = Path(log_path).stat().st_size with ThreadPoolExecutor(max_workers=4) as executor: @@ -187,14 +193,19 @@ def log_since(offset): future.result(timeout=180) pattern = ( r"Finished Xavier async KV load: request=(\S+)" - if gpu_cache_bytes is not None - else ( - r"Load Xavier V1 blocks: request=(\S+)" - if backend == "xavier" - else r"with remote block size \d+ for req (\S+)" - ) + if backend == "xavier" + else r"with remote block size \d+ for req (\S+)" ) - loaded_requests = set(re.findall(pattern, log_since(offset))) + evidence = log_since(offset) + loaded_requests = set(re.findall(pattern, evidence)) + if backend == "xavier": + # P-side history restoration also logs async completion. Exclude + # producer IDs so all four completions must come from D handoffs. + producer_requests = set( + re.findall(r"Register Xavier direct handoff: request=(\S+)", evidence) + ) + assert len(producer_requests) == 4 + loaded_requests -= producer_requests assert len(loaded_requests) == 4 print("PD concurrent requests: 4 completed with remote KV loads", flush=True) client.terminate_model(uid) diff --git a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py b/xinference/model/llm/vllm/xavier/test/test_review_regressions.py index 6b79b84bac..e4e928876d 100644 --- a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py +++ b/xinference/model/llm/vllm/xavier/test/test_review_regressions.py @@ -385,7 +385,9 @@ async def test_gpu_connector_mapping_staging_and_load_fences( connector._gpu_cache_mapped = False connector._registered_kv_caches = {"layer": cache} transfer = SimpleNamespace( - map_gpu_caches_v1=AsyncMock(side_effect=lambda *args: calls.append("map")), + map_gpu_caches_v1=AsyncMock( + side_effect=lambda *args, **kwargs: calls.append("map") + ), stage_gpu_requests_v1=AsyncMock( side_effect=lambda *args: calls.append("stage") ), @@ -437,7 +439,8 @@ async def test_gpu_mapping_serializes_callers_and_retries( release = asyncio.Event() calls = [] - async def register(descriptors, budget): + async def register(descriptors, budget, *, direct_handoff): + assert direct_handoff == connector._direct_handoff calls.append((descriptors, budget)) if len(calls) == 1: entered.set() diff --git a/xinference/model/llm/vllm/xavier/transport.py b/xinference/model/llm/vllm/xavier/transport.py index a4c5eb76e1..ac743f22cf 100644 --- a/xinference/model/llm/vllm/xavier/transport.py +++ b/xinference/model/llm/vllm/xavier/transport.py @@ -88,8 +88,19 @@ def gpu_pool_options(address: str, env: Dict[str, str]) -> Dict[str, str]: if not host: raise ValueError("Cannot determine Xavier NIXL worker host") host = resolve_nixl_host(host, "a reachable worker address") + if not host: + raise ValueError("Cannot determine Xavier NIXL worker host") if ":" in host: host = f"[{host}]" env["UCX_MEMTYPE_CACHE"] = "n" env.setdefault("UCX_TLS", os.environ.get("UCX_TLS", "tcp,cuda_copy,cuda_ipc")) return {"external_address": f"nixl://{host}:0"} + + +def uses_direct_handoff(config: Optional[Dict[str, Any]]) -> bool: + """GPU P/D uses request-scoped handoff; CPU and hybrid caches stay unchanged.""" + return bool( + config + and config.get("gpu_cache_bytes") is not None + and config.get("role") in {"prefill", "decode"} + ) diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 0069c94f68..7078ef38a7 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -34,6 +34,7 @@ from .profiling import profile_stage from .snapshot import block_major_view from .transfer import XAVIER_BF16_TRANSPORT_DTYPE, TransferActor +from .transport import uses_direct_handoff from .utils import hash_block_tokens if TYPE_CHECKING: @@ -79,6 +80,8 @@ class XavierLoadRequest: @dataclass class XavierConnectorMetadata(KVConnectorMetadata): + direct_sends: set[str] = field(default_factory=set) + direct_store: bool = False store_requests: List[XavierStoreRequest] = field(default_factory=list) load_requests: List[XavierLoadRequest] = field(default_factory=list) @@ -123,17 +126,28 @@ def __init__( raise ValueError( "Xavier V1 currently supports text-only models 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: + raise ValueError( + "Direct handoff requires P/D roles, GPU transport and one KV group" + ) 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) + if self._history_enabled and self._is_producer: + # P may restore its own independent history before computing a suffix. + self._is_consumer = True self._requests_need_load: Dict[str, XavierLoadRequest] = {} self._pending_store_requests: Dict[str, XavierStoreRequest] = {} self._leased_requests: Dict[str, XavierLoadRequest] = {} self._gpu_load_jobs: Dict[asyncio.Task, List[str]] = {} + self._direct_sends: set[str] = set() + self._invalid_block_ids: set[int] = set() self._num_cache_blocks = kv_cache_config.num_blocks self._request_staged_layers: Dict[str, set[str]] = {} self._registered_kv_caches: Dict[str, torch.Tensor | Sequence[torch.Tensor]] = ( @@ -175,12 +189,10 @@ def shutdown(self): self._loop = None def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: - if not self._is_consumer: - return - metadata = self._get_connector_metadata() assert isinstance(metadata, XavierConnectorMetadata) - if not metadata.load_requests: + self._direct_sends.update(metadata.direct_sends) + if not self._is_consumer or not metadata.load_requests: return if getattr(self, "_gpu_budget", None) is not None: @@ -217,15 +229,20 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non raise release_error def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: + sent = set() + if self._direct_handoff and self._is_producer and self._direct_sends: + assert self._transfer_ref is not None + sent = self._call(self._transfer_ref.poll_direct_gpu_v1()) + self._direct_sends.difference_update(sent) if not self._gpu_load_jobs: - return set(), set() + return sent, set() # Pump responses on the shared actor loop without waiting for any RPC. # CUDA writes and lease cleanup progress in the independent actor process. self._call(asyncio.sleep(0)) received = set() for task in list(self._gpu_load_jobs): if task.done(): - task.result() # A failed/partial transfer must never become ready. + self._invalid_block_ids.update(task.result() or ()) request_ids = self._gpu_load_jobs.pop(task) received.update(request_ids) for request_id in request_ids: @@ -237,7 +254,12 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: # If the last request aborts while a load is pending, an idle EngineCore # may not poll again until new work arrives. Destination reclamation then # waits for that step; actor-side writes and lease release still progress. - return set(), received + return sent, received + + def get_block_ids_with_load_errors(self) -> set[int]: + invalid = self._invalid_block_ids + self._invalid_block_ids = set() + return invalid def wait_for_layer_load(self, layer_name: str) -> None: return @@ -298,6 +320,12 @@ def save_kv_layer( self._pending_store_requests[request.request_id] = request def wait_for_save(self): + if self._direct_handoff and self._is_producer: + metadata = self._get_connector_metadata() + if metadata.direct_store: + self._call(self._ensure_gpu_cache_mapping()) + torch.cuda.synchronize() + return if not self._is_producer or not self._pending_store_requests: return @@ -324,6 +352,83 @@ def get_num_new_matched_tokens( raise ValueError( "Xavier V1 does not support LoRA, multimodal, embedding or salted prompts" ) + if self._direct_handoff: + params = getattr(request, "kv_transfer_params", None) or {} + handoff = params.get("xavier_direct") + if self._is_producer and self._history_enabled: + previous = self._requests_need_load.get(request.request_id) + if previous is not None: + if previous.local_transfers_by_group: + raise RuntimeError("History load already allocated") + self._requests_need_load.pop(request.request_id) + self._call(self._release_history_request(previous.lease)) + tokens = max(len(request.prompt_token_ids) - 1, 0) + # A local prefix hit can leave only a partial final block. + # Restoring that tiny tail costs more actor/layer work than + # computing it. Keep the local-cache fast path free of RPCs. + if tokens - num_computed_tokens < self._block_size: + return 0, False + start = num_computed_tokens // self._block_size + hashes = self._build_xavier_hashes(request.prompt_token_ids[:tokens])[ + start: + ] + lease = "history:" + uuid.uuid4().hex + + async def reserve(): + ref = await self._get_transfer_ref() + return await ref.reserve_direct_history_v1( + lease, [key for key, _ in hashes] + ) + + matched = self._call(reserve()) + if not matched: + return 0, False + self._requests_need_load[request.request_id] = XavierLoadRequest( + request.request_id, + {self._rank: {key: start + i for i, key in enumerate(matched)}}, + lease=lease, + ) + return ( + min(len(matched) * self._block_size, tokens - num_computed_tokens), + True, + ) + if not self._is_consumer: + return 0, False + if handoff is None: + raise ValueError("Missing direct Xavier handoff metadata") + if params.get("do_remote_prefill") is False: + # A preempted decode recomputes its prompt AND generated tokens. + # Its single-use producer ticket was consumed at allocation. + return 0, False + tokens = handoff["tokens"] + rank, ticket = handoff["rank"], handoff["ticket"] + if tokens <= num_computed_tokens: + if ticket: + + async def release(): + ref = await self._get_transfer_ref() + await ref.release_remote_direct_gpu_v1(rank, ticket) + + self._call(release()) + params["do_remote_prefill"] = False + return 0, False + + async def claim(): + ref = await self._get_transfer_ref() + return await ref.claim_remote_direct_gpu_v1(rank, ticket) + + if not self._call(claim()): + params["do_remote_prefill"] = False + self._requests_need_load.pop(request.request_id, None) + return 0, False + start = num_computed_tokens // self._block_size + blocks = handoff["blocks"] + self._requests_need_load[request.request_id] = XavierLoadRequest( + request.request_id, + {rank: {block: i for i, block in enumerate(blocks) if i >= start}}, + lease=ticket, + ) + return tokens - num_computed_tokens, True self._requests_need_load.pop(request.request_id, None) previous = self._leased_requests.pop(request.request_id, None) if previous is not None: @@ -413,6 +518,8 @@ def update_state_after_alloc( lease=load_request.lease, local_transfers_by_group=local_transfers_by_group, ) + if self._direct_handoff and not load_request.lease.startswith("history:"): + request.kv_transfer_params["do_remote_prefill"] = False logger.debug( "Xavier V1 allocated local blocks: request=%s, transfers=%s", request.request_id, @@ -424,7 +531,11 @@ def build_connector_meta( scheduler_output: SchedulerOutput, ) -> KVConnectorMetadata: meta = XavierConnectorMetadata() - if self._is_producer: + if self._direct_handoff and self._is_producer: + meta.direct_sends = self._direct_sends + self._direct_sends = set() + meta.direct_store = bool(scheduler_output.num_scheduled_tokens) + if self._is_producer and not self._direct_handoff: self._build_store_meta(scheduler_output, meta) if self._is_consumer: self._build_load_meta(scheduler_output, meta) @@ -435,6 +546,76 @@ def request_finished( request: "Request", block_ids: List[int], ) -> tuple[bool, dict[str, Any] | None]: + if self._direct_handoff: + pending = self._requests_need_load.get(request.request_id) + if pending is not None and pending.lease.startswith("history:"): + if not pending.local_transfers_by_group: + self._requests_need_load.pop(request.request_id) + self._call(self._release_history_request(pending.lease)) + else: + # Allocated async loads must still drain after an abort. + return False, None + params = getattr(request, "kv_transfer_params", None) or {} + if self._is_producer and params.get("do_remote_decode"): + from vllm.v1.request import RequestStatus + + if request.status == RequestStatus.FINISHED_ABORTED: + return False, None + token_count = max(len(request.prompt_token_ids) - 1, 0) + count = min( + len(block_ids), + (token_count + self._block_size - 1) // self._block_size, + ) + blocks = block_ids[:count] + ticket = uuid.uuid4().hex if blocks else "" + if blocks: + + async def register(): + ref = await self._get_transfer_ref() + history_args = () + if self._history_enabled: + hashes = self._build_xavier_hashes( + request.prompt_token_ids[:token_count] + ) + history_args = ([key for key, _ in hashes[:count]],) + await ref.register_direct_gpu_v1( + ticket, request.request_id, blocks, *history_args + ) + + self._call(register()) + self._direct_sends.add(request.request_id) + logger.debug( + "Register Xavier direct handoff: request=%s blocks=%s", + request.request_id, + len(blocks), + ) + return bool(blocks), { + "do_remote_prefill": True, + "xavier_direct": { + "rank": self._rank, + "address": self._xavier_config.get("rank_address"), + "ticket": ticket, + "blocks": blocks, + "tokens": min(token_count, count * self._block_size), + }, + } + pending = self._requests_need_load.get(request.request_id) + if pending is not None and pending.local_transfers_by_group: + # The worker must drain all writes before freeing destinations. + return False, None + self._requests_need_load.pop(request.request_id, None) + handoff = params.get("xavier_direct") + if handoff and params.get("do_remote_prefill") is not False: + + async def release(): + ref = await self._get_transfer_ref() + await ref.release_remote_direct_gpu_v1( + handoff["rank"], handoff["ticket"] + ) + + self._call(release()) + params["do_remote_prefill"] = False + return False, None self._chunked_prefill.pop(request.request_id, None) pending = self._requests_need_load.get(request.request_id) if ( @@ -496,13 +677,25 @@ async def _get_transfer_ref(self) -> xo.ActorRefType["TransferActor"]: self._transfer_ref = ref return self._transfer_ref + async def _release_history_request(self, lease): + transfer = await self._get_transfer_ref() + await transfer.release_direct_history_v1(lease) + async def _reserve_load_request(self, request): transfer = await self._get_transfer_ref() return await transfer.reserve_remote_blocks_v1(request.lease, request.transfers) async def _release_load_request(self, request): transfer = await self._get_transfer_ref() - await transfer.release_remote_blocks_v1(request.lease, request.transfers) + if self._direct_handoff: + if request.lease.startswith("history:"): + await transfer.release_direct_history_v1(request.lease) + else: + await transfer.release_remote_direct_gpu_v1( + next(iter(request.transfers)), request.lease + ) + else: + await transfer.release_remote_blocks_v1(request.lease, request.transfers) async def _query_remote_blocks( self, @@ -807,7 +1000,9 @@ async def _ensure_gpu_cache_mapping(self): "Xavier GPU transfer requires registered KV caches" ) torch.cuda.synchronize() - await transfer.map_gpu_caches_v1(descriptors, self._gpu_budget) + await transfer.map_gpu_caches_v1( + descriptors, self._gpu_budget, direct_handoff=self._direct_handoff + ) self._gpu_cache_mapped = True return transfer @@ -843,11 +1038,35 @@ async def _submit_gpu_requests(self, requests): ranks[rank] = layers entries.append(ranks) torch.cuda.synchronize() - task = asyncio.create_task( - transfer.load_gpu_requests_v1( - entries, [(r.lease, r.transfers) for r in requests] - ) - ) + + async def load(): + if self._direct_handoff: + historical = [ + i + for i, r in enumerate(requests) + if r.lease.startswith("history:") + ] + direct = [ + i + for i, r in enumerate(requests) + if not r.lease.startswith("history:") + ] + if historical: + await transfer.load_direct_history_v1( + [entries[i] for i in historical], + [requests[i].lease for i in historical], + ) + if direct: + return await transfer.load_direct_gpu_v1( + [entries[i] for i in direct], + [requests[i].lease for i in direct], + ) + else: + await transfer.load_gpu_requests_v1( + entries, [(r.lease, r.transfers) for r in requests] + ) + + task = asyncio.create_task(load()) except BaseException: # Worker metadata already owns these leases. Until task creation # succeeds, the TransferActor has no operation that can release them.