From 044a2eb0480600002a89ec5fff09d8e5389d1173 Mon Sep 17 00:00:00 2001 From: qinxuye Date: Wed, 7 Oct 2026 03:46:30 +0800 Subject: [PATCH 1/6] PERF: reduce Xavier V1 prefix cache overhead --- .github/workflows/python.yaml | 1 + .../user_guide/vllm_enhancement.po | 4 + .../user_guide/vllm_enhancement.po | 4 + .../user_guide/vllm_enhancement.po | 4 + .../user_guide/vllm_enhancement.po | 4 + .../user_guide/vllm_enhancement.po | 4 + .../user_guide/vllm_enhancement.po | 4 + .../user_guide/vllm_enhancement.po | 4 + .../user_guide/vllm_enhancement.po | 4 + .../user_guide/vllm_enhancement.po | 4 + doc/source/user_guide/vllm_enhancement.rst | 2 + xinference/model/llm/vllm/xavier/engine.py | 15 +- .../xavier/test/test_prefix_cache_hot_path.py | 223 ++++++++++++++++++ .../xavier/test/test_review_regressions.py | 2 +- .../vllm/xavier/test/test_v1_cache_layout.py | 1 + .../llm/vllm/xavier/test/test_v1_engine.py | 19 +- .../model/llm/vllm/xavier/test/test_xavier.py | 1 - xinference/model/llm/vllm/xavier/transfer.py | 11 + .../model/llm/vllm/xavier/v1_connector.py | 70 ++++-- 19 files changed, 347 insertions(+), 34 deletions(-) create mode 100644 xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py diff --git a/.github/workflows/python.yaml b/.github/workflows/python.yaml index f6d57b86e3..ae156ca3a5 100644 --- a/.github/workflows/python.yaml +++ b/.github/workflows/python.yaml @@ -394,6 +394,7 @@ jobs: xinference/model/llm/vllm/xavier/test/test_block_tracker_reuse.py \ xinference/model/llm/vllm/xavier/test/test_remote_kvcache_manager.py \ xinference/model/llm/vllm/xavier/test/test_v1_engine.py \ + xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py \ xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py \ xinference/model/llm/vllm/xavier/test/test_kv_schema.py \ xinference/model/llm/vllm/xavier/test/test_direct_handoff.py \ diff --git a/doc/source/locale/de/LC_MESSAGES/user_guide/vllm_enhancement.po b/doc/source/locale/de/LC_MESSAGES/user_guide/vllm_enhancement.po index 54ab14e090..0c2d1fcd42 100644 --- a/doc/source/locale/de/LC_MESSAGES/user_guide/vllm_enhancement.po +++ b/doc/source/locale/de/LC_MESSAGES/user_guide/vllm_enhancement.po @@ -130,3 +130,7 @@ msgstr "" "wiederverwendete Token unter ``prompt_tokens_details.cached_tokens``. " "Explizite Prefill-/Decode-Rollen aktivieren MLX Xavier P/D automatisch. " "Engineübergreifendes Cache-Sharing wird noch nicht unterstützt." + +#: source/user_guide/vllm_enhancement.rst:18 +msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default." +msgstr "Bei Modellen mit ausschließlich Attention unter vLLM V1 behält Xavier die Einstellung ``enforce_eager`` des Modells bei und ermöglicht damit die Ausführung mit CUDA Graph. Setzen Sie beim Modellstart ``enforce_eager=True``, um die Eager-Ausführung zu verwenden. Rekurrente Modelle verwenden weiterhin standardmäßig die Eager-Ausführung." diff --git a/doc/source/locale/es/LC_MESSAGES/user_guide/vllm_enhancement.po b/doc/source/locale/es/LC_MESSAGES/user_guide/vllm_enhancement.po index 0ff9050dd5..97815de6df 100644 --- a/doc/source/locale/es/LC_MESSAGES/user_guide/vllm_enhancement.po +++ b/doc/source/locale/es/LC_MESSAGES/user_guide/vllm_enhancement.po @@ -131,3 +131,7 @@ msgstr "" "``prompt_tokens_details.cached_tokens``. Los roles prefill/decode explícitos" " activan MLX Xavier P/D automáticamente. Aún no se admite compartir caché " "entre motores." + +#: source/user_guide/vllm_enhancement.rst:18 +msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default." +msgstr "Para los modelos que solo usan atención en vLLM V1, Xavier conserva la configuración ``enforce_eager`` del modelo, lo que permite la ejecución con CUDA Graph. Establezca ``enforce_eager=True`` al iniciar el modelo para usar la ejecución eager. Los modelos recurrentes siguen usando la ejecución eager de forma predeterminada." diff --git a/doc/source/locale/fr/LC_MESSAGES/user_guide/vllm_enhancement.po b/doc/source/locale/fr/LC_MESSAGES/user_guide/vllm_enhancement.po index 92a0e2ad09..db736ac524 100644 --- a/doc/source/locale/fr/LC_MESSAGES/user_guide/vllm_enhancement.po +++ b/doc/source/locale/fr/LC_MESSAGES/user_guide/vllm_enhancement.po @@ -134,3 +134,7 @@ msgstr "" "``prompt_tokens_details.cached_tokens``. Les rôles prefill/decode explicites" " activent automatiquement MLX Xavier P/D. Le partage de cache entre moteurs " "n’est pas encore pris en charge." + +#: source/user_guide/vllm_enhancement.rst:18 +msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default." +msgstr "Pour les modèles utilisant uniquement des couches d’attention sur vLLM V1, Xavier conserve le paramètre ``enforce_eager`` du modèle, permettant l’exécution avec CUDA Graph. Définissez ``enforce_eager=True`` au lancement du modèle pour utiliser l’exécution eager. Les modèles récurrents continuent d’utiliser l’exécution eager par défaut." diff --git a/doc/source/locale/it/LC_MESSAGES/user_guide/vllm_enhancement.po b/doc/source/locale/it/LC_MESSAGES/user_guide/vllm_enhancement.po index a1d87ec162..d87a515824 100644 --- a/doc/source/locale/it/LC_MESSAGES/user_guide/vllm_enhancement.po +++ b/doc/source/locale/it/LC_MESSAGES/user_guide/vllm_enhancement.po @@ -131,3 +131,7 @@ msgstr "" "``prompt_tokens_details.cached_tokens``. I ruoli prefill/decode espliciti " "attivano automaticamente MLX Xavier P/D. La condivisione della cache tra " "motori non è ancora supportata." + +#: source/user_guide/vllm_enhancement.rst:18 +msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default." +msgstr "Per i modelli che usano solo attention su vLLM V1, Xavier mantiene l’impostazione ``enforce_eager`` del modello, consentendo l’esecuzione con CUDA Graph. Impostare ``enforce_eager=True`` all’avvio del modello per usare l’esecuzione eager. I modelli ricorrenti continuano a usare l’esecuzione eager per impostazione predefinita." diff --git a/doc/source/locale/ja/LC_MESSAGES/user_guide/vllm_enhancement.po b/doc/source/locale/ja/LC_MESSAGES/user_guide/vllm_enhancement.po index 89028d5b22..c71d0d71bd 100644 --- a/doc/source/locale/ja/LC_MESSAGES/user_guide/vllm_enhancement.po +++ b/doc/source/locale/ja/LC_MESSAGES/user_guide/vllm_enhancement.po @@ -118,3 +118,7 @@ msgstr "" "キャッシュも保持します。通常レプリカはミスやストレージ利用不可時にローカル計算し、互換性のないデプロイ間ではページを共有できません。使用量の " "``prompt_tokens_details.cached_tokens`` に再利用トークン数を返します。明示的な prefill/decode " "ロールは MLX Xavier P/D を自動で有効にします。異なるエンジン間のキャッシュ共有は未対応です。" + +#: source/user_guide/vllm_enhancement.rst:18 +msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default." +msgstr "vLLM V1 の attention のみを使用するモデルでは、Xavier はモデルの ``enforce_eager`` 設定を維持し、CUDA Graph による実行を可能にします。eager 実行を使用するには、モデル起動時に ``enforce_eager=True`` を設定してください。再帰型モデルでは、引き続きデフォルトで eager 実行を使用します。" diff --git a/doc/source/locale/ko/LC_MESSAGES/user_guide/vllm_enhancement.po b/doc/source/locale/ko/LC_MESSAGES/user_guide/vllm_enhancement.po index 24f37f39a5..a731290f8d 100644 --- a/doc/source/locale/ko/LC_MESSAGES/user_guide/vllm_enhancement.po +++ b/doc/source/locale/ko/LC_MESSAGES/user_guide/vllm_enhancement.po @@ -122,3 +122,7 @@ msgstr "" "페이지를 공유할 수 없습니다. 사용량의 ``prompt_tokens_details.cached_tokens`` 는 재사용된 토큰 수를 " "보고합니다. 명시적인 prefill/decode 역할은 MLX Xavier P/D를 자동 활성화합니다. 엔진 간 캐시 공유는 아직 " "지원하지 않습니다." + +#: source/user_guide/vllm_enhancement.rst:18 +msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default." +msgstr "vLLM V1에서 attention만 사용하는 모델의 경우 Xavier는 모델의 ``enforce_eager`` 설정을 유지하여 CUDA Graph 실행을 허용합니다. eager 실행을 사용하려면 모델을 시작할 때 ``enforce_eager=True`` 설정을 사용하세요. 순환 모델은 계속해서 기본적으로 eager 실행을 사용합니다." diff --git a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/vllm_enhancement.po b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/vllm_enhancement.po index 8b943c529f..392067c5a1 100644 --- a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/vllm_enhancement.po +++ b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/vllm_enhancement.po @@ -129,3 +129,7 @@ msgstr "" "tokens reutilizados em ``prompt_tokens_details.cached_tokens``. Papéis " "prefill/decode explícitos ativam MLX Xavier P/D automaticamente. " "Compartilhamento de cache entre engines ainda não é suportado." + +#: source/user_guide/vllm_enhancement.rst:18 +msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default." +msgstr "Para modelos que usam apenas atenção no vLLM V1, o Xavier preserva a configuração ``enforce_eager`` do modelo, permitindo a execução com CUDA Graph. Defina ``enforce_eager=True`` ao iniciar o modelo para usar a execução eager. Os modelos recorrentes continuam usando a execução eager por padrão." diff --git a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/vllm_enhancement.po b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/vllm_enhancement.po index c8cc717221..9b306d5fd5 100644 --- a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/vllm_enhancement.po +++ b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/vllm_enhancement.po @@ -119,3 +119,7 @@ msgstr "" "缓存。普通副本在未命中或存储不可用时本地计算；不兼容的部署不能共享缓存页。响应通过 " "``prompt_tokens_details.cached_tokens`` 报告复用 token 数。显式设置 prefill/decode " "角色会自动启用 MLX Xavier P/D。目前不支持跨引擎缓存共享。" + +#: source/user_guide/vllm_enhancement.rst:18 +msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default." +msgstr "对于 vLLM V1 上仅使用注意力层的模型，Xavier 保留模型的 ``enforce_eager`` 设置，允许使用 CUDA Graph 执行。若要使用 eager 执行，请在启动模型时设置 ``enforce_eager=True``。循环模型仍默认使用 eager 执行。" diff --git a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/vllm_enhancement.po b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/vllm_enhancement.po index 9d07718f0f..01bd2d5a08 100644 --- a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/vllm_enhancement.po +++ b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/vllm_enhancement.po @@ -118,3 +118,7 @@ msgstr "" "快取。普通副本在未命中或儲存不可用時於本機計算；不相容的部署不能共用快取頁面。回應透過 " "``prompt_tokens_details.cached_tokens`` 報告重用 token 數。明確設定 prefill/decode " "角色會自動啟用 MLX Xavier P/D。目前不支援跨引擎快取共用。" + +#: source/user_guide/vllm_enhancement.rst:18 +msgid "For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default." +msgstr "對於 vLLM V1 上僅使用注意力層的模型，Xavier 保留模型的 ``enforce_eager`` 設定，允許使用 CUDA Graph 執行。若要使用 eager 執行，請在啟動模型時設定 ``enforce_eager=True``。循環模型仍預設使用 eager 執行。" diff --git a/doc/source/user_guide/vllm_enhancement.rst b/doc/source/user_guide/vllm_enhancement.rst index 3b26a529ae..b3c0f8019a 100644 --- a/doc/source/user_guide/vllm_enhancement.rst +++ b/doc/source/user_guide/vllm_enhancement.rst @@ -15,6 +15,8 @@ Usage ***** Simply add the parameter ``enable_xavier=True`` when starting the vllm model. +For attention-only models on vLLM V1, Xavier preserves the model's ``enforce_eager`` setting, allowing CUDA Graph execution. Set ``enforce_eager=True`` at model launch to use eager execution. Recurrent models continue to use eager execution by default. + *********** Limitations *********** diff --git a/xinference/model/llm/vllm/xavier/engine.py b/xinference/model/llm/vllm/xavier/engine.py index 851c8b0f3d..cfab819755 100644 --- a/xinference/model/llm/vllm/xavier/engine.py +++ b/xinference/model/llm/vllm/xavier/engine.py @@ -30,8 +30,6 @@ logger = logging.getLogger(__name__) -XAVIER_EAGER_VLLM_VERSION = version.parse("0.21.0") - class XavierEngine: _xavier_config: Optional[Dict] = None @@ -90,14 +88,12 @@ def _patch_v1_engine_args( extra_config = dict(xavier_config.get("kv_connector_extra_config") or {}) extra_config["xavier_config"] = xavier_config - if ( - version.parse(VLLM_VERSION) >= XAVIER_EAGER_VLLM_VERSION - and xavier_config.get("enforce_eager", True) - and not getattr(engine_args, "enforce_eager", False) + recurrent = engine_args.create_model_config().is_hybrid + if xavier_config.get("enforce_eager", recurrent) and not getattr( + engine_args, "enforce_eager", False ): - # vLLM V1 may create XavierConnector twice during CUDA graph - # setup. In CUDA-heavy processes this can trip glibc static TLS - # allocation, so keep Xavier on the eager execution path. + # Attention snapshots are exported after forward and support CUDA + # graphs. Keep the existing eager default for recurrent handoff. engine_args.enforce_eager = True logger.info( "Set enforce_eager=True for Xavier V1 on vLLM %s.", @@ -108,7 +104,6 @@ def _patch_v1_engine_args( # of sliding-window and full attention. Only recurrent models need this # override: 0.21 disables HMA for every KV connector by default; 0.22+ # checks SupportsHMA itself. Preserve an explicit user setting. - recurrent = engine_args.create_model_config().is_hybrid if ( recurrent and version.parse(VLLM_VERSION) < version.parse("0.22.0") diff --git a/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py b/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py new file mode 100644 index 0000000000..cf696b210c --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py @@ -0,0 +1,223 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +import torch + +from ..snapshot import KVSnapshotStore +from ..tiered_snapshot import TieredKVSnapshotStore +from ..transfer import TransferActor + + +@pytest.fixture +def snapshot_export(connector): + store = KVSnapshotStore(8) + actor = SimpleNamespace(_snapshot_store=store, _rank=1) + transfer = SimpleNamespace( + ready_blocks_for_export_v1=AsyncMock( + side_effect=lambda keys: TransferActor.ready_blocks_for_export_v1( + actor, keys + ) + ), + stage_layer_blocks_v1=AsyncMock( + side_effect=lambda *args: TransferActor.stage_layer_blocks_v1(actor, *args) + ), + publish_blocks_v1=AsyncMock( + side_effect=lambda *args: TransferActor.publish_blocks_v1(actor, *args) + ), + ) + tracker = SimpleNamespace( + register_blocks=AsyncMock(), unregister_blocks=AsyncMock() + ) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + connector._get_tracker_ref = AsyncMock(return_value=tracker) + connector._registered_kv_caches = { + "K": torch.arange(16).view(8, 2).float(), + "V": torch.arange(16, 32).view(8, 2).float(), + } + connector._layer_group_ids = {"K": 0, "V": 1} + return store, transfer, tracker + + +def bind_store_requests(connector, module, *requests): + metadata = module.XavierConnectorMetadata(store_requests=list(requests)) + connector._get_connector_metadata = lambda: metadata + return metadata + + +@pytest.mark.parametrize("remaining", [0, 1, 15, 16]) +def test_local_prefix_tail_skips_lookup(connector, remaining): + connector._query_remote_blocks = AsyncMock(return_value={}) + connector._build_xavier_hashes = Mock(return_value=[(1, 0), (2, 1)]) + request = SimpleNamespace( + request_id="r", prompt_token_ids=list(range(16 + remaining + 1)) + ) + assert connector.get_num_new_matched_tokens(request, 16) == (0, False) + assert connector._query_remote_blocks.await_count == (remaining == 16) + assert connector._build_xavier_hashes.call_count == (remaining == 16) + + +def test_skipped_tail_releases_previous_lease(connector, connector_module): + load = connector_module.XavierLoadRequest("r", {2: {111: 0}}, lease="1:old") + connector._leased_requests["r"] = load + connector._requests_need_load["r"] = load + connector._release_load_request = AsyncMock() + connector._query_remote_blocks = AsyncMock() + request = SimpleNamespace(request_id="r", prompt_token_ids=list(range(18))) + assert connector.get_num_new_matched_tokens(request, 16) == (0, False) + connector._release_load_request.assert_awaited_once_with(load) + connector._query_remote_blocks.assert_not_awaited() + assert not connector._requests_need_load and not connector._leased_requests + + +def test_export_uses_post_forward_data_and_current_metadata( + connector, connector_module, snapshot_export +): + store, transfer, tracker = snapshot_export + capture = connector_module.XavierStoreRequest("capture", [0], [111], [[0], [1]]) + metadata = bind_store_requests(connector, connector_module, capture) + for name, cache in connector._registered_kv_caches.items(): + connector.save_kv_layer(name, cache, None) + transfer.stage_layer_blocks_v1.assert_not_awaited() + transfer.ready_blocks_for_export_v1.assert_not_awaited() + current = connector_module.XavierStoreRequest("replay", [2], [222], [[2], [3]]) + metadata.store_requests = [current] + connector._registered_kv_caches["K"][2].fill_(42) + connector._registered_kv_caches["V"][3].fill_(84) + # Full CUDA graph replay need not invoke any per-layer Python callback. + connector.wait_for_save() + assert store.ready == {222} + assert store.read("K", [222]).tolist() == [[42, 42]] + assert store.read("V", [222]).tolist() == [[84, 84]] + tracker.register_blocks.assert_awaited_once_with(0, [(222, 222)], 1) + assert not connector._request_staged_layers + + +def test_repeated_export_skips_payload_but_refreshes_discovery( + connector, connector_module, snapshot_export +): + store, transfer, tracker = snapshot_export + request = connector_module.XavierStoreRequest("r", [2], [111], [[2], [3]]) + bind_store_requests(connector, connector_module, request) + connector.wait_for_save() + original = store.read("K", [111]).clone() + connector._registered_kv_caches["K"][2].fill_(99) + connector.wait_for_save() + assert transfer.stage_layer_blocks_v1.await_count == 2 + assert tracker.register_blocks.await_count == 2 + assert torch.equal(store.read("K", [111]), original) + + +def test_batch_dedup_preserves_positions_in_every_cache_group( + connector, connector_module, snapshot_export +): + store, transfer, _ = snapshot_export + store.stage("K", [111], torch.ones(1, 2)) + store.stage("V", [111], torch.ones(1, 2)) + store.publish([111], {"K", "V"}) + first = connector_module.XavierStoreRequest( + "a", [0, 2, 4], [111, 222, 333], [[0, 2, 4], [1, 3, 5]] + ) + second = connector_module.XavierStoreRequest("b", [6], [222], [[6], [7]]) + bind_store_requests(connector, connector_module, first, second) + connector.wait_for_save() + assert store.ready == {111, 222, 333} + assert store.read("K", [222, 333]).tolist() == [[4, 5], [8, 9]] + assert store.read("V", [222, 333]).tolist() == [[22, 23], [26, 27]] + assert transfer.stage_layer_blocks_v1.await_count == 2 + assert transfer.stage_layer_blocks_v1.await_args_list[0].args[2] == [222, 333] + assert first.block_ids == [0, 2, 4] + assert first.block_ids_by_group == [[0, 2, 4], [1, 3, 5]] + + +def test_partial_export_failure_retries_unpublished_content( + connector, connector_module, snapshot_export +): + store, transfer, _ = snapshot_export + request = connector_module.XavierStoreRequest("r", [2], [111], [[2], [3]]) + bind_store_requests(connector, connector_module, request) + stage = transfer.stage_layer_blocks_v1.side_effect + + def fail_v(*args): + if args[1] == "V": + raise RuntimeError("copy failed") + return stage(*args) + + transfer.stage_layer_blocks_v1.side_effect = fail_v + with pytest.raises(RuntimeError, match="copy failed"): + connector.wait_for_save() + assert not store.ready + assert set(store.blocks[111]) == {"K"} + assert not connector._request_staged_layers + transfer.stage_layer_blocks_v1.side_effect = stage + connector.wait_for_save() + assert store.ready == {111} + assert store.read("V", [111]).tolist() == [[22, 23]] + + +def test_tracker_failure_retries_without_exporting_payload( + connector, connector_module, snapshot_export +): + store, transfer, tracker = snapshot_export + request = connector_module.XavierStoreRequest("r", [2], [111], [[2], [3]]) + bind_store_requests(connector, connector_module, request) + tracker.register_blocks.side_effect = RuntimeError("tracker unavailable") + with pytest.raises(RuntimeError, match="tracker unavailable"): + connector.wait_for_save() + assert store.ready == {111} + tracker.register_blocks.side_effect = None + connector.wait_for_save() + assert tracker.register_blocks.await_count == 2 + assert transfer.stage_layer_blocks_v1.await_count == 2 + + +def test_evicted_content_is_exported_again( + connector, connector_module, snapshot_export +): + store, transfer, tracker = snapshot_export + store.capacity = 1 + first = connector_module.XavierStoreRequest("a", [2], [111], [[2], [3]]) + second = connector_module.XavierStoreRequest("b", [4], [222], [[4], [5]]) + for request in [first, second, first]: + bind_store_requests(connector, connector_module, request) + connector.wait_for_save() + assert store.ready == set(request.block_hashes) + assert transfer.stage_layer_blocks_v1.await_count == 6 + assert tracker.unregister_blocks.await_count == 2 + assert store.read("K", [111]).tolist() == [[4, 5]] + + +@pytest.mark.parametrize("tiered", [False, True]) +def test_ready_query_excludes_partial_blocks_and_refreshes_lru(tiered): + store = ( + TieredKVSnapshotStore(2, 0, 16, torch.device("cpu")) + if tiered + else KVSnapshotStore(2) + ) + actor = SimpleNamespace(_snapshot_store=store) + for key in [111, 222]: + store.stage("K", [key], torch.ones(1, 2)) + store.publish([key], {"K"}) + assert TransferActor.ready_blocks_for_export_v1(actor, [111, 111, 333]) == [111] + store.stage("K", [333], torch.ones(1, 2)) + assert store.ready == {111} + assert TransferActor.ready_blocks_for_export_v1(actor, [111, 222, 333]) == [111] + store.publish([333], {"K"}) + assert TransferActor.ready_blocks_for_export_v1(actor, [333]) == [333] + + +def test_unconfigured_store_has_no_ready_blocks(): + actor = SimpleNamespace(_snapshot_store=None) + assert TransferActor.ready_blocks_for_export_v1(actor, [111]) == [] + + +@pytest.mark.asyncio +async def test_empty_gpu_export_avoids_mapping_and_fence(connector, monkeypatch): + connector._ensure_gpu_cache_mapping = AsyncMock() + fence = Mock() + monkeypatch.setattr(torch.cuda, "synchronize", fence) + await connector._stage_gpu_requests([]) + connector._ensure_gpu_cache_mapping.assert_not_awaited() + fence.assert_not_called() 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 e4e928876d..352b3a764c 100644 --- a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py +++ b/xinference/model/llm/vllm/xavier/test/test_review_regressions.py @@ -425,7 +425,7 @@ async def test_gpu_connector_mapping_staging_and_load_fences( side_effect=AssertionError("GPU mode stages once in wait_for_save") ) connector.save_kv_layer("layer", cache, None) - assert connector._pending_store_requests["r"] is request + connector._stage_kv_layer_for_request.assert_not_called() @pytest.mark.asyncio diff --git a/xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py b/xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py index c92aa329bf..095e6b22fd 100644 --- a/xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py +++ b/xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py @@ -16,6 +16,7 @@ def test_load_uses_registered_packed_hybrid_cache(): load = Mock() connector = SimpleNamespace( _is_consumer=True, + _direct_sends=set(), _registered_kv_caches={"linear_attn": packed}, _get_connector_metadata=lambda: XavierConnectorMetadata( load_requests=[request] diff --git a/xinference/model/llm/vllm/xavier/test/test_v1_engine.py b/xinference/model/llm/vllm/xavier/test/test_v1_engine.py index c82a1f8a01..288e4f18d3 100644 --- a/xinference/model/llm/vllm/xavier/test/test_v1_engine.py +++ b/xinference/model/llm/vllm/xavier/test/test_v1_engine.py @@ -61,13 +61,30 @@ def test_v1_connector_configuration(engine, role, kv_role): assert args.kv_transfer_config.kv_connector == "XavierConnector" assert args.additional_config["custom"] == 1 json.dumps(args.additional_config) - assert args.enforce_eager + assert not args.enforce_eager assert getattr(args, "disable_hybrid_kv_cache_manager", None) is None assert args.enable_prefix_caching is True assert config == {"role": role, "rank": 2, "block_tracker_uid": b"tracker"} factory.assert_called_once() +@pytest.mark.parametrize("recurrent", [False, True]) +@pytest.mark.parametrize("launch_eager", [False, True]) +@pytest.mark.parametrize("xavier_eager", [None, False, True]) +def test_eager_defaults_and_explicit_settings( + engine, recurrent, launch_eager, xavier_eager +): + module, _ = engine + args = SimpleNamespace( + enforce_eager=launch_eager, + create_model_config=lambda: SimpleNamespace(is_hybrid=recurrent), + ) + config = {} if xavier_eager is None else {"enforce_eager": xavier_eager} + module.XavierEngine.from_engine_args(args, xavier_config=config) + expected = launch_eager or (recurrent if xavier_eager is None else xavier_eager) + assert args.enforce_eager is expected + + def test_v0_uses_legacy_adapter(engine, monkeypatch): module, factory = engine module.VLLM_VERSION = "0.7.3" diff --git a/xinference/model/llm/vllm/xavier/test/test_xavier.py b/xinference/model/llm/vllm/xavier/test/test_xavier.py index 1dc1c69308..cea57c19c5 100644 --- a/xinference/model/llm/vllm/xavier/test/test_xavier.py +++ b/xinference/model/llm/vllm/xavier/test/test_xavier.py @@ -105,7 +105,6 @@ async def test_block_tracker(actor_pool_context): 123: { (rank, 0), }, - 456: set(), 789: {(rank, 2)}, } diff --git a/xinference/model/llm/vllm/xavier/transfer.py b/xinference/model/llm/vllm/xavier/transfer.py index 46999d21aa..f04fa98bc1 100644 --- a/xinference/model/llm/vllm/xavier/transfer.py +++ b/xinference/model/llm/vllm/xavier/transfer.py @@ -215,6 +215,17 @@ def stage_layer_blocks_v1(self, request_id, layer_name, block_ids, blocks): block_ids, ) + def ready_blocks_for_export_v1(self, keys: List[int]) -> List[int]: + if self._snapshot_store is None: + return [] + ready = [] + for key in dict.fromkeys(keys): + if key in self._snapshot_store.ready: + ready.append(key) + # A reused snapshot is still hot even without another write. + self._snapshot_store.blocks.move_to_end(key) + return ready + def publish_blocks_v1(self, keys, layers): available = self._snapshot_store.publish(keys, set(layers)) evicted = list(self._snapshot_store.evicted) diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 8f421160ce..bd95b5ce2f 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -190,7 +190,6 @@ def __init__( # 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() @@ -351,20 +350,10 @@ def save_kv_layer( attn_metadata: "AttentionMetadata", **kwargs: Any, ) -> None: - if not self._is_producer: - return - - metadata = self._get_connector_metadata() - assert isinstance(metadata, XavierConnectorMetadata) - if not metadata.store_requests: - return - - for request in metadata.store_requests: - if not request.block_ids: - continue - if getattr(self, "_gpu_budget", None) is None: - self._stage_kv_layer_for_request(request, layer_name, kv_layer) - self._pending_store_requests[request.request_id] = request + # This callback runs inside attention and CUDA graph capture/replay. + # Export registered caches in wait_for_save, outside the model forward, + # using the current metadata rather than capture-time request IDs. + return def wait_for_save(self): if self._direct_handoff and self._is_producer: @@ -373,13 +362,20 @@ def wait_for_save(self): self._call(self._ensure_gpu_cache_mapping()) torch.cuda.synchronize() return - if not self._is_producer or not self._pending_store_requests: + if not self._is_producer: + return + metadata = self._get_connector_metadata() + assert isinstance(metadata, XavierConnectorMetadata) + pending = metadata.store_requests + if not pending: return - - pending = list(self._pending_store_requests.values()) - self._pending_store_requests.clear() try: - self._stage_missing_registered_layers(pending) + missing = self._call(self._filter_store_requests(pending)) + if missing: + self._stage_missing_registered_layers(missing) + # Re-publish all requested keys, including reused snapshots. This + # refreshes discovery and retries a failed tracker registration + # without copying already published KV payloads again. self._call(self._register_blocks(pending)) finally: for request in pending: @@ -525,7 +521,7 @@ async def claim(): token_ids = list(request.prompt_token_ids or []) external_token_count = max(len(token_ids) - 1, 0) - if external_token_count <= num_computed_tokens: + if external_token_count - num_computed_tokens < self._block_size: return 0, False hashes = self._build_xavier_hashes(token_ids[:external_token_count]) @@ -859,6 +855,36 @@ async def _stage_layer_blocks( blocks, ) + async def _filter_store_requests( + self, requests: List[XavierStoreRequest] + ) -> List[XavierStoreRequest]: + transfer_ref = await self._get_transfer_ref() + ready = set( + await transfer_ref.ready_blocks_for_export_v1( + [key for request in requests for key in request.block_hashes] + ) + ) + missing = [] + for request in requests: + indices = [] + for i, key in enumerate(request.block_hashes): + if key not in ready: + indices.append(i) + ready.add(key) + if indices: + missing.append( + XavierStoreRequest( + request.request_id, + [request.block_ids[i] for i in indices], + [request.block_hashes[i] for i in indices], + [ + [group[i] for i in indices] + for group in request.block_ids_by_group + ], + ) + ) + return missing + async def _register_blocks(self, requests: List[XavierStoreRequest]): tracker_ref = await self._get_tracker_ref() transfer_ref = await self._get_transfer_ref() @@ -1146,6 +1172,8 @@ async def _ensure_gpu_cache_mapping(self): return transfer async def _stage_gpu_requests(self, requests): + if not requests: + return transfer = await self._ensure_gpu_cache_mapping() entries = [] for request in requests: From 80d5c11e99bffa53d48a4e77f6c365eb27b1b0c6 Mon Sep 17 00:00:00 2001 From: qinxuye Date: Wed, 7 Oct 2026 10:30:35 +0800 Subject: [PATCH 2/6] PERF: batch Xavier V1 snapshot export and registration --- .../xavier/test/test_prefix_cache_hot_path.py | 272 +++++++++++++++++- xinference/model/llm/vllm/xavier/transfer.py | 39 +++ .../model/llm/vllm/xavier/v1_connector.py | 202 +++++++++++-- 3 files changed, 480 insertions(+), 33 deletions(-) diff --git a/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py b/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py index cf696b210c..0c96df4f23 100644 --- a/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py +++ b/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py @@ -1,5 +1,8 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. +import asyncio +import socket +import weakref from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -24,6 +27,9 @@ def snapshot_export(connector): stage_layer_blocks_v1=AsyncMock( side_effect=lambda *args: TransferActor.stage_layer_blocks_v1(actor, *args) ), + stage_layer_batches_v1=AsyncMock( + side_effect=lambda *args: TransferActor.stage_layer_batches_v1(actor, *args) + ), publish_blocks_v1=AsyncMock( side_effect=lambda *args: TransferActor.publish_blocks_v1(actor, *args) ), @@ -81,6 +87,7 @@ def test_export_uses_post_forward_data_and_current_metadata( for name, cache in connector._registered_kv_caches.items(): connector.save_kv_layer(name, cache, None) transfer.stage_layer_blocks_v1.assert_not_awaited() + transfer.stage_layer_batches_v1.assert_not_awaited() transfer.ready_blocks_for_export_v1.assert_not_awaited() current = connector_module.XavierStoreRequest("replay", [2], [222], [[2], [3]]) metadata.store_requests = [current] @@ -105,7 +112,7 @@ def test_repeated_export_skips_payload_but_refreshes_discovery( original = store.read("K", [111]).clone() connector._registered_kv_caches["K"][2].fill_(99) connector.wait_for_save() - assert transfer.stage_layer_blocks_v1.await_count == 2 + assert transfer.stage_layer_batches_v1.await_count == 1 assert tracker.register_blocks.await_count == 2 assert torch.equal(store.read("K", [111]), original) @@ -113,7 +120,7 @@ def test_repeated_export_skips_payload_but_refreshes_discovery( def test_batch_dedup_preserves_positions_in_every_cache_group( connector, connector_module, snapshot_export ): - store, transfer, _ = snapshot_export + store, transfer, tracker = snapshot_export store.stage("K", [111], torch.ones(1, 2)) store.stage("V", [111], torch.ones(1, 2)) store.publish([111], {"K", "V"}) @@ -126,8 +133,11 @@ def test_batch_dedup_preserves_positions_in_every_cache_group( assert store.ready == {111, 222, 333} assert store.read("K", [222, 333]).tolist() == [[4, 5], [8, 9]] assert store.read("V", [222, 333]).tolist() == [[22, 23], [26, 27]] - assert transfer.stage_layer_blocks_v1.await_count == 2 - assert transfer.stage_layer_blocks_v1.await_args_list[0].args[2] == [222, 333] + transfer.stage_layer_batches_v1.assert_awaited_once() + assert transfer.stage_layer_batches_v1.await_args.args[0][0][1] == [222, 333] + tracker.register_blocks.assert_awaited_once_with( + 0, [(111, 111), (222, 222), (333, 333)], 1 + ) assert first.block_ids == [0, 2, 4] assert first.block_ids_by_group == [[0, 2, 4], [1, 3, 5]] @@ -138,20 +148,20 @@ def test_partial_export_failure_retries_unpublished_content( store, transfer, _ = snapshot_export request = connector_module.XavierStoreRequest("r", [2], [111], [[2], [3]]) bind_store_requests(connector, connector_module, request) - stage = transfer.stage_layer_blocks_v1.side_effect + stage = store.stage - def fail_v(*args): - if args[1] == "V": + def fail_v(*args, **kwargs): + if args[0] == "V": raise RuntimeError("copy failed") - return stage(*args) + return stage(*args, **kwargs) - transfer.stage_layer_blocks_v1.side_effect = fail_v + store.stage = fail_v with pytest.raises(RuntimeError, match="copy failed"): connector.wait_for_save() assert not store.ready assert set(store.blocks[111]) == {"K"} assert not connector._request_staged_layers - transfer.stage_layer_blocks_v1.side_effect = stage + store.stage = stage connector.wait_for_save() assert store.ready == {111} assert store.read("V", [111]).tolist() == [[22, 23]] @@ -170,7 +180,7 @@ def test_tracker_failure_retries_without_exporting_payload( tracker.register_blocks.side_effect = None connector.wait_for_save() assert tracker.register_blocks.await_count == 2 - assert transfer.stage_layer_blocks_v1.await_count == 2 + assert transfer.stage_layer_batches_v1.await_count == 1 def test_evicted_content_is_exported_again( @@ -184,7 +194,7 @@ def test_evicted_content_is_exported_again( bind_store_requests(connector, connector_module, request) connector.wait_for_save() assert store.ready == set(request.block_hashes) - assert transfer.stage_layer_blocks_v1.await_count == 6 + assert transfer.stage_layer_batches_v1.await_count == 3 assert tracker.unregister_blocks.await_count == 2 assert store.read("K", [111]).tolist() == [[4, 5]] @@ -221,3 +231,241 @@ async def test_empty_gpu_export_avoids_mapping_and_fence(connector, monkeypatch) await connector._stage_gpu_requests([]) connector._ensure_gpu_cache_mapping.assert_not_awaited() fence.assert_not_called() + + +def test_export_batches_bound_bytes_across_layers_and_requests( + connector, connector_module, snapshot_export, monkeypatch +): + store, transfer, tracker = snapshot_export + # Each block has two float32 values in each of K and V: 16 bytes total. + monkeypatch.setattr(connector_module, "_MAX_EXPORT_BYTES", 32) + first = connector_module.XavierStoreRequest( + "a", [0, 2, 4], [111, 222, 333], [[0, 2, 4], [1, 3, 5]] + ) + second = connector_module.XavierStoreRequest("b", [6], [444], [[6], [7]]) + bind_store_requests(connector, connector_module, first, second) + connector.wait_for_save() + batches = [call.args[0] for call in transfer.stage_layer_batches_v1.await_args_list] + assert len(batches) == 2 + assert [batch[0][1] for batch in batches] == [[111, 222], [333, 444]] + assert all( + sum(t.numel() * t.element_size() for _, _, t in batch) <= 32 + for batch in batches + ) + assert store.read("K", [333, 444]).tolist() == [[8, 9], [12, 13]] + assert store.read("V", [333, 444]).tolist() == [[26, 27], [30, 31]] + transfer.publish_blocks_v1.assert_awaited_once_with( + [111, 222, 333, 444], {"K", "V"} + ) + tracker.register_blocks.assert_awaited_once_with( + 0, [(key, key) for key in [111, 222, 333, 444]], 1 + ) + + +def test_oversized_export_block_keeps_all_layers_together( + connector, connector_module, snapshot_export, monkeypatch +): + store, transfer, _ = snapshot_export + monkeypatch.setattr(connector_module, "_MAX_EXPORT_BYTES", 1) + request = connector_module.XavierStoreRequest("r", [2], [111], [[2], [3]]) + bind_store_requests(connector, connector_module, request) + connector.wait_for_save() + transfer.stage_layer_batches_v1.assert_awaited_once() + assert store.ready == {111} + + +def test_batch_larger_than_available_capacity_keeps_complete_blocks( + connector, connector_module, snapshot_export +): + store, _, tracker = snapshot_export + store.capacity = 2 + store.stage("K", [999], torch.zeros(1, 2)) + store.stage("V", [999], torch.zeros(1, 2)) + store.publish([999], {"K", "V"}) + assert store.reserve("lease", [999]) + request = connector_module.XavierStoreRequest( + "r", [2, 4], [111, 222], [[2, 4], [3, 5]] + ) + bind_store_requests(connector, connector_module, request) + connector.wait_for_save() + assert store.ready == {999, 222} + assert store.read("V", [222]).tolist() == [[26, 27]] + tracker.unregister_blocks.assert_awaited_once_with(0, 1, [111]) + tracker.register_blocks.assert_awaited_once_with(0, [(222, 222)], 1) + + +def test_batched_export_preserves_noncontiguous_bf16_and_snapshot_ownership( + connector, connector_module, snapshot_export +): + store, _, _ = snapshot_export + values = torch.arange(32, dtype=torch.int16).view(8, 4) + tensor = values.view(torch.bfloat16)[:, ::2] + connector._registered_kv_caches = {"K": tensor} + request = connector_module.XavierStoreRequest("r", [2], [111], [[2]]) + bind_store_requests(connector, connector_module, request) + connector.wait_for_save() + assert store.logical_dtypes[111]["K"] == torch.bfloat16 + original = tensor[2].view(torch.int16).clone() + tensor[2].zero_() + assert torch.equal(store.read("K", [111]).view(torch.int16)[0], original) + + +@pytest.mark.asyncio +async def test_batched_registration_keeps_distinct_layer_requirements( + connector, connector_module +): + connector._registered_kv_caches = {} + connector._request_staged_layers = {"a": {"K"}, "b": {"K", "V"}} + transfer = SimpleNamespace( + publish_blocks_v1=AsyncMock(side_effect=[([111], [222]), ([222], [])]) + ) + tracker = SimpleNamespace( + register_blocks=AsyncMock(), unregister_blocks=AsyncMock() + ) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + connector._get_tracker_ref = AsyncMock(return_value=tracker) + request_type = connector_module.XavierStoreRequest + await connector._register_blocks( + [ + request_type("a", [0], [111], [[0]]), + request_type("b", [1], [222], [[1]]), + ] + ) + assert [call.args for call in transfer.publish_blocks_v1.await_args_list] == [ + ([111], {"K"}), + ([222], {"K", "V"}), + ] + tracker.unregister_blocks.assert_awaited_once_with(0, 1, [222]) + tracker.register_blocks.assert_awaited_once_with(0, [(111, 111), (222, 222)], 1) + + +@pytest.mark.asyncio +async def test_export_failure_drains_next_copy_without_publishing( + connector, snapshot_export, monkeypatch +): + _, transfer, tracker = snapshot_export + events = [Mock(), Mock()] + batches = [([], [torch.ones(1)], [event]) for event in events] + monkeypatch.setattr(connector, "_cpu_export_batches", lambda _: iter(batches)) + transfer.stage_layer_batches_v1.side_effect = RuntimeError("actor failed") + with pytest.raises(RuntimeError, match="actor failed"): + await connector._stage_cpu_requests([Mock()]) + for event in events: + event.synchronize.assert_called_once() + transfer.publish_blocks_v1.assert_not_awaited() + tracker.register_blocks.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_export_releases_previous_buffers_before_preparing_next_batch( + connector, snapshot_export, monkeypatch +): + owners = [] + + def batches(_): + for i in range(4): + # Only the preceding batch may still own GPU storage at this point. + assert sum(ref() is not None for ref in owners) <= 1 + owner = torch.ones(1) + owners.append(weakref.ref(owner)) + yield [("K", [i], torch.ones(1))], [owner], [] + del owner + + monkeypatch.setattr(connector, "_cpu_export_batches", batches) + await connector._stage_cpu_requests([Mock()]) + assert all(ref() is None for ref in owners) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_cuda_export_fences_current_stream_and_owns_snapshot( + connector, connector_module, snapshot_export, monkeypatch, dtype +): + store, transfer, _ = snapshot_export + connector._registered_kv_caches = { + name: tensor.to(device="cuda", dtype=dtype) + for name, tensor in connector._registered_kv_caches.items() + } + # Force two batches so the second copy is in flight during the first RPC. + monkeypatch.setattr(connector_module, "_MAX_EXPORT_BYTES", 8) + request = connector_module.XavierStoreRequest( + "r", [2, 4], [111, 222], [[2, 4], [3, 5]] + ) + bind_store_requests(connector, connector_module, request) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + connector._registered_kv_caches["K"][2].fill_(42) + connector._registered_kv_caches["V"][5].fill_(84) + connector.wait_for_save() + for tensor in connector._registered_kv_caches.values(): + tensor.zero_() + stream.synchronize() + assert store.ready == {111, 222} + assert store.read("K", [111, 222]).view(dtype).tolist() == [[42, 42], [8, 9]] + assert store.read("V", [111, 222]).view(dtype).tolist() == [[22, 23], [84, 84]] + assert transfer.stage_layer_batches_v1.await_count == 2 + assert all( + t.is_pinned() + for call in transfer.stage_layer_batches_v1.await_args_list + for _, _, t in call.args[0] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +async def test_batch_export_socket_backpressure_preserves_dtype_and_bytes( + connector, connector_module, snapshot_export, dtype +): + from xoscar.backends.communication.socket import SocketChannel + from xoscar.serialization.aio import AioSerializer + + store, transfer, _ = snapshot_export + source = torch.arange(8 * 32768).view(8, 32768).to(dtype) + connector._registered_kv_caches = {f"layer{i}": source for i in range(4)} + actor = SimpleNamespace(_snapshot_store=store, _rank=1) + received = asyncio.get_running_loop().create_future() + + async def receive(reader, writer): + try: + # Force the sender to queue tensor buffers, rather than writing + # the entire message immediately into the socket's send buffer. + await asyncio.sleep(0.01) + channel = SocketChannel(reader, writer) + entries, metadata = await channel.recv() + TransferActor.stage_layer_batches_v1(actor, entries, metadata) + await channel.send(True) + received.set_result(None) + except Exception as exc: + received.set_exception(exc) + finally: + writer.close() + await writer.wait_closed() + + server = await asyncio.start_server(receive, "127.0.0.1", 0) + reader, writer = await asyncio.open_connection(*server.sockets[0].getsockname()) + writer.get_extra_info("socket").setsockopt( + socket.SOL_SOCKET, socket.SO_SNDBUF, 4096 + ) + channel = SocketChannel(reader, writer) + + async def send(entries, metadata): + buffers = await AioSerializer((entries, metadata)).run() + assert all(len(buf) == getattr(buf, "nbytes", len(buf)) for buf in buffers) + await channel.send((entries, metadata)) + assert await channel.recv() is True + + transfer.stage_layer_batches_v1.side_effect = send + request = connector_module.XavierStoreRequest("r", [2, 4], [111, 222], [[2, 4]]) + try: + await asyncio.wait_for(connector._stage_cpu_requests([request]), timeout=20) + await asyncio.wait_for(received, timeout=20) + for name in connector._registered_kv_caches: + actual = store.read(name, [111, 222]).view(dtype) + assert torch.equal(actual, source[[2, 4]]) + assert store.logical_dtypes[111][name] == dtype + finally: + writer.close() + await writer.wait_closed() + server.close() + await server.wait_closed() diff --git a/xinference/model/llm/vllm/xavier/transfer.py b/xinference/model/llm/vllm/xavier/transfer.py index f04fa98bc1..ac048a9e21 100644 --- a/xinference/model/llm/vllm/xavier/transfer.py +++ b/xinference/model/llm/vllm/xavier/transfer.py @@ -226,6 +226,45 @@ def ready_blocks_for_export_v1(self, keys: List[int]) -> List[int]: self._snapshot_store.blocks.move_to_end(key) return ready + def stage_layer_batches_v1( + self, + entries: List[Tuple[str, List[int], torch.Tensor]], + metadata: List[Tuple[Tuple[int, ...], torch.dtype]], + ) -> None: + if len(entries) != len(metadata): + raise ValueError("Xavier export metadata does not match payload") + entries = [ + (name, keys, blocks.view(dtype).view(shape)) + for (name, keys, blocks), (shape, dtype) in zip(entries, metadata) + ] + store = self._snapshot_store + assert store is not None + batch_keys = set(key for _, ids, _ in entries for key in ids) + pinned = set().union(*store.leases.values()) + if len(batch_keys | pinned) <= store.capacity: + # All batch keys can stay resident across layers, even if staging + # first needs to evict older unleased content. + for name, ids, blocks in entries: + TransferActor.stage_layer_blocks_v1(self, "batch", name, ids, blocks) + return + # Stage a complete block before moving to the next one. Layer-major + # staging could evict every partial block when the batch exceeds the + # available capacity (including capacity held by transfer leases). + positions = [ + (name, dict(zip(keys, range(len(keys)))), blocks) + for name, keys, blocks in entries + ] + ordered_keys = dict.fromkeys( + key for _, indices, _ in positions for key in indices + ) + for key in ordered_keys: + for name, indices, blocks in positions: + if key in indices: + i = indices[key] + TransferActor.stage_layer_blocks_v1( + self, "batch", name, [key], blocks[i : i + 1] + ) + def publish_blocks_v1(self, keys, layers): available = self._snapshot_store.publish(keys, set(layers)) evicted = list(self._snapshot_store.evicted) diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index bd95b5ce2f..d359ac6e5c 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -16,7 +16,17 @@ import math import uuid from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Sequence, Tuple +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterable, + Iterator, + List, + Optional, + Sequence, + Tuple, +) import torch import xoscar as xo @@ -49,6 +59,12 @@ _MAX_READ_BYTES = 1024 * 1024 _MAX_READ_BLOCKS = 64 +_MAX_EXPORT_BYTES = 32 * 1024 * 1024 +_CPUExportBatch = Tuple[ + List[Tuple[str, List[int], torch.Tensor]], + List[torch.Tensor], + List[torch.cuda.Event], +] @dataclass @@ -886,6 +902,8 @@ async def _filter_store_requests( return missing async def _register_blocks(self, requests: List[XavierStoreRequest]): + if not requests: + return tracker_ref = await self._get_tracker_ref() transfer_ref = await self._get_transfer_ref() expected_layers = { @@ -893,25 +911,48 @@ async def _register_blocks(self, requests: List[XavierStoreRequest]): for layer, cache in self._registered_kv_caches.items() for name, _ in self._iter_kv_tensors(layer, cache) } - for request in requests: - layers = expected_layers or self._request_staged_layers.get( - request.request_id, set() + if expected_layers: + # Registered caches share one schema. Avoid constructing layer + # groups and copying the actor's already unique result on hot hits. + keys = list( + dict.fromkeys( + key for request in requests for key in request.block_hashes + ) ) available, evicted = await transfer_ref.publish_blocks_v1( - request.block_hashes, layers - ) - if evicted: - await tracker_ref.unregister_blocks(0, self._rank, evicted) - # Transport addresses identify immutable content, not recyclable GPU slots. - await tracker_ref.register_blocks( - 0, [(key, key) for key in available], self._rank - ) - logger.debug( - "Xavier V1 registered blocks: request=%s, rank=%s, blocks=%s", - request.request_id, - self._rank, - available, + keys, expected_layers ) + else: + groups: Dict[frozenset[str], Dict[int, None]] = {} + for request in requests: + layers = frozenset( + self._request_staged_layers.get(request.request_id, set()) + ) + groups.setdefault(layers, {}).update( + dict.fromkeys(request.block_hashes) + ) + available_keys: Dict[int, None] = {} + evicted_keys: Dict[int, None] = {} + for layers, group_keys in groups.items(): + ready, removed = await transfer_ref.publish_blocks_v1( + list(group_keys), set(layers) + ) + available_keys.update(dict.fromkeys(ready)) + evicted_keys.update(dict.fromkeys(removed)) + available, evicted = list(available_keys), list(evicted_keys) + if evicted: + await tracker_ref.unregister_blocks(0, self._rank, evicted) + # Retry discovery even when ready snapshots required no payload export. + # Transport addresses identify immutable content, not recyclable GPU slots. + await tracker_ref.register_blocks( + 0, [(key, key) for key in available], self._rank + ) + logger.debug( + "Xavier V1 registered blocks: requests=%s, rank=%s, blocks=%s", + len(requests), + self._rank, + available, + ) async def _read_layer_blocks( self, @@ -1137,11 +1178,130 @@ def _stage_missing_registered_layers( if not self._registered_kv_caches: return + self._call(self._stage_cpu_requests(requests)) + + def _cpu_export_batches( + self, requests: List[XavierStoreRequest] + ) -> Iterator[_CPUExportBatch]: + caches = { + name: block_major_view(tensor, self._num_cache_blocks) + for layer, cache in self._registered_kv_caches.items() + for name, tensor in self._iter_kv_tensors(layer, cache) + } + block_bytes = sum( + math.prod(tensor.shape[1:]) * tensor.element_size() + for tensor in caches.values() + ) + if not block_bytes: + return + # Keep every layer of a block in the same RPC. One oversized block is + # indivisible, just as on the read path. + batch_size = max(1, _MAX_EXPORT_BYTES // block_bytes) + keys: List[int] = [] + sources: Dict[str, List[int]] = {name: [] for name in caches} for request in requests: - if not request.block_ids: - continue - for layer_name, kv_layer in self._registered_kv_caches.items(): - self._stage_kv_layer_for_request(request, layer_name, kv_layer) + layer_ids = { + name: self._get_source_block_ids(request, name) for name in caches + } + count = min( + len(request.block_hashes), + len(request.block_ids), + *(len(ids) for ids in layer_ids.values()), + ) + for i in range(count): + keys.append(request.block_hashes[i]) + for name, ids in layer_ids.items(): + sources[name].append(ids[i]) + if len(keys) == batch_size: + yield self._prepare_cpu_export(caches, keys, sources) + keys = [] + sources = {name: [] for name in caches} + if keys: + yield self._prepare_cpu_export(caches, keys, sources) + + def _prepare_cpu_export( + self, + caches: Dict[str, torch.Tensor], + keys: List[int], + sources: Dict[str, List[int]], + ) -> _CPUExportBatch: + entries = [] + owners = [] + events = {} + try: + for name, tensor in caches.items(): + with profile_stage( + "store_d2h", + device=tensor.device, + layer=name, + blocks=len(keys), + rank=self._rank, + ): + indices = torch.tensor( + sources[name], device=tensor.device, dtype=torch.long + ) + gathered = tensor.index_select(0, indices).detach() + if tensor.is_cuda: + if tensor.device not in events: + events[tensor.device] = torch.cuda.Event() + event = events[tensor.device] + blocks = torch.empty_like( + gathered, device="cpu", pin_memory=True + ) + # Retain both sides until the event has completed, including + # when allocation, staging or an actor RPC fails. + owners.append(gathered) + entries.append((name, keys, blocks)) + blocks.copy_(gathered, non_blocking=True) + event.record(torch.cuda.current_stream(tensor.device)) + else: + entries.append((name, keys, gathered.contiguous())) + except BaseException: + for event in events.values(): + event.synchronize() + raise + return entries, owners, list(events.values()) + + async def _stage_cpu_requests(self, requests: List[XavierStoreRequest]) -> None: + if not requests or not self._registered_kv_caches: + return + transfer = await self._get_transfer_ref() + pending: Optional[_CPUExportBatch] = None + + async def stage(batch: _CPUExportBatch) -> None: + entries, _, events = batch + for event in events: + event.synchronize() + with profile_stage( + "store_rpc", + rank=self._rank, + nbytes=sum(t.numel() * t.element_size() for _, _, t in entries), + ): + # xoscar's Torch serializer exposes a NumPy memoryview. Socket + # backpressure requires len(buffer) to count bytes, so send + # flat uint8 views rather than multidimensional typed buffers. + await transfer.stage_layer_batches_v1( + [ + (name, keys, t.view(torch.uint8).reshape(-1)) + for name, keys, t in entries + ], + [(tuple(t.shape), t.dtype) for _, _, t in entries], + ) + + try: + for current in self._cpu_export_batches(requests): + previous, pending = pending, current + if previous is not None: + # The next D2H copy overlaps the previous batch's actor RPC. + # At most two bounded batches own GPU and pinned CPU buffers. + await stage(previous) + previous = None + if pending is not None: + await stage(pending) + finally: + if pending is not None: + for event in pending[2]: + event.synchronize() async def _ensure_gpu_cache_mapping(self): from torch.multiprocessing.reductions import reduce_tensor From 1dffa46908d07bbfdd50a897257f169432828481 Mon Sep 17 00:00:00 2001 From: qinxuye Date: Wed, 7 Oct 2026 14:26:15 +0800 Subject: [PATCH 3/6] PERF: streamline Xavier V1 CPU snapshot pipeline --- .github/workflows/python.yaml | 5 + xinference/core/model.py | 4 + xinference/model/llm/vllm/xavier/engine.py | 22 +- .../model/llm/vllm/xavier/gc_lifecycle.py | 25 + .../xavier/test/test_block_tracker_reuse.py | 71 ++ .../vllm/xavier/test/test_local_directory.py | 181 +++++ .../llm/vllm/xavier/test/test_local_read.py | 225 ++++++ .../vllm/xavier/test/test_packed_gather.py | 148 ++++ .../xavier/test/test_prefix_cache_hot_path.py | 254 ++++++- .../vllm/xavier/test/test_request_transfer.py | 46 +- .../xavier/test/test_review_regressions.py | 6 +- .../xavier/test/test_snapshot_export_ipc.py | 508 +++++++++++++ .../xavier/test/test_snapshot_gc_lifecycle.py | 108 +++ .../llm/vllm/xavier/test/test_v1_engine.py | 30 + xinference/model/llm/vllm/xavier/transfer.py | 215 +++++- .../model/llm/vllm/xavier/v1_connector.py | 697 +++++++++++++++++- .../llm/xavier/backends/torch/gpu_export.py | 56 ++ .../llm/xavier/backends/torch/local_read.py | 151 ++++ .../xavier/backends/torch/packed_gather.py | 156 ++++ .../llm/xavier/backends/torch/snapshot.py | 205 +++++- .../xavier/backends/torch/tiered_snapshot.py | 4 +- xinference/model/llm/xavier/block_tracker.py | 144 +++- .../model/llm/xavier/local_directory.py | 139 ++++ .../model/llm/xavier/tests/test_token_hash.py | 20 + xinference/model/llm/xavier/utils.py | 27 +- 25 files changed, 3332 insertions(+), 115 deletions(-) create mode 100644 xinference/model/llm/vllm/xavier/gc_lifecycle.py create mode 100644 xinference/model/llm/vllm/xavier/test/test_local_directory.py create mode 100644 xinference/model/llm/vllm/xavier/test/test_local_read.py create mode 100644 xinference/model/llm/vllm/xavier/test/test_packed_gather.py create mode 100644 xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py create mode 100644 xinference/model/llm/vllm/xavier/test/test_snapshot_gc_lifecycle.py create mode 100644 xinference/model/llm/xavier/backends/torch/gpu_export.py create mode 100644 xinference/model/llm/xavier/backends/torch/local_read.py create mode 100644 xinference/model/llm/xavier/backends/torch/packed_gather.py create mode 100644 xinference/model/llm/xavier/local_directory.py create mode 100644 xinference/model/llm/xavier/tests/test_token_hash.py diff --git a/.github/workflows/python.yaml b/.github/workflows/python.yaml index ae156ca3a5..5a33297cbd 100644 --- a/.github/workflows/python.yaml +++ b/.github/workflows/python.yaml @@ -395,6 +395,11 @@ jobs: xinference/model/llm/vllm/xavier/test/test_remote_kvcache_manager.py \ xinference/model/llm/vllm/xavier/test/test_v1_engine.py \ xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py \ + xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py \ + xinference/model/llm/vllm/xavier/test/test_snapshot_gc_lifecycle.py \ + xinference/model/llm/vllm/xavier/test/test_packed_gather.py \ + xinference/model/llm/vllm/xavier/test/test_local_directory.py \ + xinference/model/llm/vllm/xavier/test/test_local_read.py \ xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py \ xinference/model/llm/vllm/xavier/test/test_kv_schema.py \ xinference/model/llm/vllm/xavier/test/test_direct_handoff.py \ diff --git a/xinference/core/model.py b/xinference/core/model.py index 4ad3e3c770..625b92cdce 100644 --- a/xinference/core/model.py +++ b/xinference/core/model.py @@ -399,6 +399,10 @@ async def start_transfer_for_vllm(self, rank_addresses: List[str]): store_address=self._xavier_config.get("store_address"), # type: ignore store_port=self._xavier_config.get("store_port"), # type: ignore world_addresses=rank_addresses, + freeze_initialization_gc=( + self._model._is_vllm_v1() + and self._xavier_config.get("gpu_cache_bytes") is None + ), ) await self._model.init_xavier() logger.debug( diff --git a/xinference/model/llm/vllm/xavier/engine.py b/xinference/model/llm/vllm/xavier/engine.py index cfab819755..d81ca64b03 100644 --- a/xinference/model/llm/vllm/xavier/engine.py +++ b/xinference/model/llm/vllm/xavier/engine.py @@ -34,18 +34,6 @@ class XavierEngine: _xavier_config: Optional[Dict] = None - @staticmethod - def _json_safe(value: Any) -> Any: - if isinstance(value, bytes): - return value.decode(errors="replace") - if isinstance(value, dict): - return {str(k): XavierEngine._json_safe(v) for k, v in value.items()} - if isinstance(value, (list, tuple)): - return [XavierEngine._json_safe(v) for v in value] - if isinstance(value, (str, int, float, bool)) or value is None: - return value - return str(value) - @classmethod def _patch_v1_engine_args( cls, engine_args: AsyncEngineArgs, xavier_config: Optional[Dict] @@ -67,13 +55,6 @@ def _patch_v1_engine_args( xavier_config = dict(xavier_config) set_xavier_transport_backend(xavier_config, XAVIER_TRANSPORT_XAVIER) - additional_config = dict(getattr(engine_args, "additional_config", {}) or {}) - # vLLM V1 includes additional_config in its compilation cache hash and - # serializes it as JSON. Keep the connector's full config below, but - # only expose a JSON-safe copy here. - additional_config["xavier_config"] = cls._json_safe(xavier_config) - engine_args.additional_config = additional_config - role = xavier_config.get("role") if role == "prefill": kv_role = "kv_producer" @@ -86,6 +67,9 @@ def _patch_v1_engine_args( kv_rank = xavier_config.get("rank", 0) extra_config = dict(xavier_config.get("kv_connector_extra_config") or {}) + # Actor addresses, IDs and ranks affect transport, not compiled model + # operations. Keep them in the connector config: additional_config is + # included in vLLM's graph hash and would recompile on every launch. extra_config["xavier_config"] = xavier_config recurrent = engine_args.create_model_config().is_hybrid diff --git a/xinference/model/llm/vllm/xavier/gc_lifecycle.py b/xinference/model/llm/vllm/xavier/gc_lifecycle.py new file mode 100644 index 0000000000..db482dadcb --- /dev/null +++ b/xinference/model/llm/vllm/xavier/gc_lifecycle.py @@ -0,0 +1,25 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Exclude initialized libraries from snapshot allocation cycle scans.""" + +import gc + + +class InitializationGCFreeze: + def __init__(self) -> None: + self._owned = False + self._started = False + + def start(self) -> None: + if self._started: + return + self._owned = gc.get_freeze_count() == 0 + gc.collect() + gc.freeze() + self._started = True + + def close(self) -> None: + if self._owned: + gc.unfreeze() + self._owned = False + self._started = False diff --git a/xinference/model/llm/vllm/xavier/test/test_block_tracker_reuse.py b/xinference/model/llm/vllm/xavier/test/test_block_tracker_reuse.py index c6514277b8..2706d43a7c 100644 --- a/xinference/model/llm/vllm/xavier/test/test_block_tracker_reuse.py +++ b/xinference/model/llm/vllm/xavier/test/test_block_tracker_reuse.py @@ -78,3 +78,74 @@ def test_unregister_block_still_works(self): tracker.unregister_block(0, rank=1, block_id=5) assert _query_rank_block(tracker, 1001) == set() assert tracker._rank_to_hash_and_block_id[0][1] == set() + + def test_multiple_hashes_in_one_batch_are_removed_on_reuse(self): + tracker = VLLMBlockTracker() + tracker.register_blocks(0, [(1001, 5), (1002, 5), (2001, 6)], rank=1) + tracker.register_blocks(0, [(3001, 5)], rank=1) + assert _query_rank_block(tracker, 1001) == set() + assert _query_rank_block(tracker, 1002) == set() + assert _query_rank_block(tracker, 2001) == {(1, 6)} + assert _query_rank_block(tracker, 3001) == {(1, 5)} + tracker.unregister_blocks(0, 1, [5, 5, 99]) + assert _query_rank_block(tracker, 3001) == set() + assert tracker._rank_to_block_hashes[0][1] == {6: {2001}} + + def test_recovery_clears_slot_index_for_every_engine(self): + tracker = VLLMBlockTracker() + for engine in [0, 1]: + tracker.register_blocks(engine, [(1001, 5)], rank=1) + tracker.register_blocks(engine, [(2001, 5)], rank=2) + tracker.unregister_rank(1) + tracker.register_rank(1) + for engine in [0, 1]: + tracker.register_blocks(engine, [(3001, 5)], rank=1) + assert tracker._rank_to_block_hashes[engine] == { + 1: {5: {3001}}, + 2: {5: {2001}}, + } + assert tracker.query_blocks(engine, [(1001, 0)]) == {} + + +class TestSnapshotDirectory: + def test_duplicate_publication_and_rank_exclusion(self): + tracker = VLLMBlockTracker() + tracker.register_snapshot_blocks(0, [101, 202, 101], 1) + tracker.register_snapshot_blocks(0, [101], 2) + assert tracker.query_blocks(0, [(101, 7), (202, 8)], exclude_rank=1) == { + 2: {(101, 101, 7)} + } + tracker.unregister_blocks(0, 1, [101, 101, 999]) + assert _query_rank_block(tracker, 101) == {(2, 101)} + assert _query_rank_block(tracker, 202) == {(1, 202)} + tracker.unregister_blocks(0, 2, [101]) + assert _query_rank_block(tracker, 101) == set() + + def test_recovery_clears_all_engines_without_removing_other_ranks(self): + tracker = VLLMBlockTracker() + for engine in (0, 1): + tracker.register_snapshot_blocks(engine, [101], 1) + tracker.register_snapshot_blocks(engine, [101, 202], 2) + tracker.unregister_rank(1) + for engine in (0, 1): + assert tracker.query_blocks(engine, [(101, 0)], exclude_rank=2) == {} + tracker.register_rank(1) + for engine in (0, 1): + assert tracker.query_blocks(engine, [(101, 0)], exclude_rank=1) == { + 2: {(101, 101, 0)} + } + tracker.register_snapshot_blocks(engine, [303], 1) + assert tracker.query_blocks(engine, [(101, 0)], exclude_rank=2) == {} + assert tracker.query_blocks(engine, [(303, 1)], exclude_rank=2) == { + 1: {(303, 303, 1)} + } + + def test_switching_address_kinds_replaces_stale_locations(self): + tracker = VLLMBlockTracker() + tracker.register_blocks(0, [(999, 101)], 1) + tracker.register_snapshot_blocks(0, [101], 1) + assert _query_rank_block(tracker, 999) == set() + assert _query_rank_block(tracker, 101) == {(1, 101)} + tracker.register_blocks(0, [(303, 101)], 1) + assert _query_rank_block(tracker, 101) == set() + assert _query_rank_block(tracker, 303) == {(1, 101)} diff --git a/xinference/model/llm/vllm/xavier/test/test_local_directory.py b/xinference/model/llm/vllm/xavier/test/test_local_directory.py new file mode 100644 index 0000000000..b6aa696b1e --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_local_directory.py @@ -0,0 +1,181 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +import asyncio +import multiprocessing +import os +import sys +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from ..block_tracker import VLLMBlockTracker + +pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="Linux memfd required") + + +def _read_in_child(metadata, key, queue): + from ....xavier.local_directory import LocalBlockDirectory + + directory = LocalBlockDirectory.attach(metadata) + queue.put(None if directory is None else directory.contains(key)) + if directory is not None: + directory.close() + + +def test_directory_publication_collisions_eviction_and_recovery(): + from ....xavier.local_directory import LocalBlockDirectory + + tracker = VLLMBlockTracker() + collision = 1 << 51 # Both filter positions coincide with key zero. + tracker.register_snapshot_blocks(0, [0, collision], 1) + directory = LocalBlockDirectory.attach(tracker.get_snapshot_directory(0)) + assert directory is not None + try: + assert directory.contains(0) is True and directory.contains(999) is False + tracker.register_snapshot_blocks(0, [0], 2) + tracker.unregister_blocks(0, 1, [0, collision, collision]) + assert directory.contains(0) is True + tracker.unregister_rank(2) + tracker.register_rank(2) + assert directory.contains(0) is False + # Publication is visible immediately, without a refresh interval. + tracker.register_snapshot_blocks(0, [999], 2) + assert directory.contains(999) is True + tracker.register_blocks(0, [(123, 999)], 2) + assert directory.contains(123) is None and not directory.valid + assert tracker.get_snapshot_directory(0) is None + finally: + directory.close() + asyncio.run(tracker.__pre_destroy__()) + + +def test_busy_writer_falls_back_and_owner_close_invalidates_reader(): + from ....xavier.local_directory import LocalBlockDirectory, SharedBlockDirectory + + owner = SharedBlockDirectory() + directory = LocalBlockDirectory.attach(owner.metadata()) + assert directory is not None + try: + with owner.update(): + assert directory.contains(111) is None and directory.valid + owner.add(111) + assert directory.contains(111) is True + owner.close() + assert directory.contains(111) is None and not directory.valid + finally: + owner.close() + directory.close() + + +def test_spawned_reader_and_foreign_or_reused_owner_rejection(): + from ....xavier.local_directory import LocalBlockDirectory, SharedBlockDirectory + + owner = SharedBlockDirectory() + try: + with owner.update(): + owner.add(111) + metadata = owner.metadata() + ctx = multiprocessing.get_context("spawn") + queue = ctx.Queue() + proc = ctx.Process(target=_read_in_child, args=(metadata, 111, queue)) + proc.start() + try: + assert queue.get(timeout=30) is True + proc.join(timeout=30) + assert proc.exitcode == 0 + finally: + if proc.is_alive(): + proc.terminate() + proc.join(timeout=5) + queue.close() + assert ( + LocalBlockDirectory.attach({**metadata, "boot_id": "foreign-host"}) is None + ) + assert LocalBlockDirectory.attach({**metadata, "token": b"wrong-owner"}) is None + assert ( + LocalBlockDirectory.attach( + {**metadata, "path": f"/proc/{os.getpid()}/fd/999999"} + ) + is None + ) + finally: + owner.close() + + +def test_negative_shortcut_and_invalid_owner_reconnect(connector): + directory = SimpleNamespace( + contains=Mock(return_value=False), valid=True, close=Mock() + ) + connector._local_directory = directory + connector._query_remote_blocks = AsyncMock(return_value={}) + request = SimpleNamespace(request_id="r", prompt_token_ids=list(range(33))) + assert connector.get_num_new_matched_tokens(request, 0) == (0, False) + connector._query_remote_blocks.assert_not_awaited() + directory.contains.return_value = None + directory.valid = False + assert connector.get_num_new_matched_tokens(request, 0) == (0, False) + directory.close.assert_called_once() + assert connector._local_directory is None + connector._query_remote_blocks.assert_awaited_once() + + +def _owner_in_child(queue, stop): + from ....xavier.local_directory import SharedBlockDirectory + + owner = SharedBlockDirectory() + with owner.update(): + owner.add(111) + queue.put(owner.metadata()) + queue.close() + queue.join_thread() + stop.wait(30) + os._exit(0) # Simulate owner death without publishing its closed flag. + + +def test_owner_process_death_invalidates_retained_mapping(): + from ....xavier.local_directory import LocalBlockDirectory + + ctx = multiprocessing.get_context("spawn") + queue = ctx.Queue() + stop = ctx.Event() + proc = ctx.Process(target=_owner_in_child, args=(queue, stop)) + directory = None + proc.start() + try: + directory = LocalBlockDirectory.attach(queue.get(timeout=30)) + assert directory is not None and directory.contains(111) is True + stop.set() + proc.join(timeout=30) + assert proc.exitcode == 0 + assert directory.contains(111) is None and not directory.valid + finally: + if proc.is_alive(): + proc.terminate() + proc.join(timeout=5) + if directory is not None: + directory.close() + queue.close() + + +@pytest.mark.asyncio +async def test_shared_memory_unavailable_keeps_authoritative_query( + connector, monkeypatch +): + from ....xavier import local_directory + + def unavailable(): + raise OSError("shared memory unavailable") + + monkeypatch.setattr(local_directory, "SharedBlockDirectory", unavailable) + tracker = VLLMBlockTracker() + tracker.register_snapshot_blocks(0, [111], 2) + ref = SimpleNamespace( + get_snapshot_directory=AsyncMock(side_effect=tracker.get_snapshot_directory), + query_blocks=AsyncMock(side_effect=tracker.query_blocks), + ) + connector._get_tracker_ref = AsyncMock(return_value=ref) + assert await connector._query_remote_blocks("r", [(111, 0)]) == {2: {(111, 111, 0)}} + assert await connector._query_remote_blocks("s", [(222, 0)]) == {} + assert ref.get_snapshot_directory.await_count == 1 + assert ref.query_blocks.await_count == 2 diff --git a/xinference/model/llm/vllm/xavier/test/test_local_read.py b/xinference/model/llm/vllm/xavier/test/test_local_read.py new file mode 100644 index 0000000000..b9d781f245 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_local_read.py @@ -0,0 +1,225 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +import asyncio +import multiprocessing +import sys +from types import MethodType, SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import torch + +from ..request_transfer import LayerRead +from ..transfer import TransferActor + +pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="Linux tmpfs required") + + +def _copy_in_child(metadata, lease, queue): + from ....xavier.backends.torch.local_read import LocalReadBuffer + + view = LocalReadBuffer.attach(metadata) + queue.put(view.copy(lease, 8, pin_memory=False).tolist()) + view.close() + + +def test_shared_receive_is_leased_and_copied_before_slot_reuse(): + from ....xavier.backends.torch.local_read import LocalReadBuffer, SharedReadBuffer + + owner = SharedReadBuffer(64) + view = LocalReadBuffer.attach(owner.metadata()) + try: + lease, payload = owner.acquire(8) + payload.copy_(torch.arange(8, dtype=torch.uint8)) + expected = view.copy(lease, 8, pin_memory=False) + assert owner.acquire(8) is None + owner.release(b"wrong") + assert owner.acquire(8) is None + ctx = multiprocessing.get_context("spawn") + queue = ctx.Queue() + child = ctx.Process( + target=_copy_in_child, args=(owner.metadata(), lease, queue) + ) + child.start() + assert queue.get(timeout=30) == list(range(8)) + child.join(30) + assert child.exitcode == 0 + owner.release(lease) + next_lease, replacement = owner.acquire(8) + replacement.fill_(99) + assert expected.tolist() == list(range(8)) + assert view.copy(next_lease, 8, pin_memory=False).tolist() == [99] * 8 + with pytest.raises(RuntimeError, match="lease"): + view.copy(lease, 8, pin_memory=False) + owner.release(next_lease) + assert owner.acquire(65) is None + assert owner.acquire(0) is None + del payload, replacement + owner.close() + assert not view.valid() + finally: + view.close() + owner.close() + view.close() + + +@pytest.mark.parametrize("changed", ["boot_id", "token", "path", "capacity"]) +def test_foreign_stale_or_invalid_receive_metadata_falls_back(changed): + from ....xavier.backends.torch.local_read import LocalReadBuffer, SharedReadBuffer + + owner = SharedReadBuffer(64) + try: + metadata = owner.metadata() + metadata[changed] = { + "boot_id": "foreign", + "token": bytes(16), + "path": "/proc/0/fd/0", + "capacity": 65, + }[changed] + assert LocalReadBuffer.attach(metadata) is None + finally: + owner.close() + + +@pytest.mark.parametrize("cancel", [False, True]) +@pytest.mark.asyncio +async def test_receive_failure_or_cancel_fences_native_write_before_release(cancel): + from ....xavier.backends.torch.local_read import SharedReadBuffer + + owner = SharedReadBuffer(64) + entered, finish = asyncio.Event(), asyncio.Event() + + async def receive(*_, _buffer): + entered.set() + await finish.wait() + _buffer.fill_(99) + if not cancel: + raise RuntimeError("receive failed") + + actor = SimpleNamespace( + _local_read_buffer_v1=owner, + _local_read_tasks_v1=set(), + read_request_blocks_v1=receive, + ) + reads = [LayerRead("K", [1], [2], (8,), torch.uint8)] + task = asyncio.create_task( + TransferActor.read_request_blocks_local_v1(actor, 0, reads, owner.token) + ) + try: + await entered.wait() + if cancel: + task.cancel() + await asyncio.sleep(0) + task.cancel() + assert owner.acquire(8) is None and not task.done() + finish.set() + with pytest.raises(asyncio.CancelledError if cancel else RuntimeError): + await task + assert not actor._local_read_tasks_v1 + acquired = owner.acquire(8) + assert acquired is not None and acquired[1].tolist() == [99] * 8 + owner.release(acquired[0]) + del acquired + finally: + finish.set() + await asyncio.gather(task, return_exceptions=True) + owner.close() + + +@pytest.mark.parametrize("copy_failure", [False, True]) +def test_connector_owns_receive_bytes_and_releases_even_on_copy_failure( + connector, monkeypatch, copy_failure +): + from ....xavier.backends.torch.local_read import LocalReadBuffer, SharedReadBuffer + + owner = SharedReadBuffer(64) + actor = SimpleNamespace(_local_read_buffer_v1=owner) + + async def read_local(*_): + acquired = owner.acquire(8) + acquired[1].fill_(7) + return acquired[0] + + transfer = SimpleNamespace( + get_local_read_metadata_v1=AsyncMock(return_value=owner.metadata()), + read_request_blocks_local_v1=AsyncMock(side_effect=read_local), + release_local_read_v1=AsyncMock( + side_effect=MethodType(TransferActor.release_local_read_v1, actor) + ), + ) + reads = [LayerRead("K", [1], [2], (8,), torch.uint8)] + try: + if copy_failure: + monkeypatch.setattr( + LocalReadBuffer, + "copy", + lambda *_, **__: (_ for _ in ()).throw(RuntimeError("copy failed")), + ) + with pytest.raises(RuntimeError, match="copy failed"): + connector._call( + connector._read_request_payload( + transfer, 0, reads, pin_memory=False + ) + ) + else: + copied = connector._call( + connector._read_request_payload(transfer, 0, reads, pin_memory=False) + ) + assert copied.tolist() == [7] * 8 + transfer.release_local_read_v1.assert_awaited_once() + lease, payload = owner.acquire(8) + payload.fill_(99) + if not copy_failure: + assert copied.tolist() == [7] * 8 + owner.release(lease) + del payload + finally: + connector.shutdown() + owner.close() + + +def test_shared_receive_unavailable_uses_normal_tensor_rpc(connector, monkeypatch): + from ....xavier.backends.torch import local_read + + monkeypatch.setattr( + local_read, + "SharedReadBuffer", + lambda: (_ for _ in ()).throw(OSError("no shared memory")), + ) + assert TransferActor.get_local_read_metadata_v1(SimpleNamespace()) is None + expected = torch.arange(8, dtype=torch.uint8) + transfer = SimpleNamespace( + get_local_read_metadata_v1=AsyncMock(return_value=None), + read_request_blocks_v1=AsyncMock(return_value=expected), + ) + result = connector._call( + connector._read_request_payload(transfer, 0, [], pin_memory=False) + ) + assert result is expected + transfer.read_request_blocks_v1.assert_awaited_once_with(0, []) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_pinned_receive_retains_exact_bits_after_shared_slot_is_reused(dtype): + from ....xavier.backends.torch.local_read import LocalReadBuffer, SharedReadBuffer + + owner = SharedReadBuffer(64) + view = LocalReadBuffer.attach(owner.metadata()) + values = torch.tensor([1e30, -1e30, 3, -4]).to(dtype) + bits = values.view(torch.uint8) + try: + lease, payload = owner.acquire(bits.numel()) + payload.copy_(bits) + copied = view.copy(lease, bits.numel(), pin_memory=True) + assert copied.is_pinned() + owner.release(lease) + new_lease, replacement = owner.acquire(bits.numel()) + replacement.zero_() + result = copied.to("cuda", non_blocking=True).cpu() + assert torch.equal(result, bits) + owner.release(new_lease) + del payload, replacement + finally: + view.close() + owner.close() diff --git a/xinference/model/llm/vllm/xavier/test/test_packed_gather.py b/xinference/model/llm/vllm/xavier/test/test_packed_gather.py new file mode 100644 index 0000000000..f745b8dd52 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_packed_gather.py @@ -0,0 +1,148 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +import torch + + +def test_fused_export_queues_owned_payload_without_repacking( + connector, connector_module, monkeypatch +): + from torch.multiprocessing import reductions + + packed = torch.arange(32, dtype=torch.uint8).view(2, 16) + layers = [("K", (4,), torch.float32, 0, 16)] + + class Gather: + device = torch.device("cpu") + + def __call__(self, sources): + assert sources == {"K": [2, 4]} + return packed, [] + + gather = Gather() + gather.layers = layers + connector._packed_gather = gather + transfer = SimpleNamespace(enqueue_snapshot_export_v1=AsyncMock()) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + monkeypatch.setattr(reductions, "reduce_tensor", lambda value: (None, (value,))) + monkeypatch.setattr(torch, "cat", lambda *a, **kw: pytest.fail("payload repacked")) + batch = connector._prepare_cpu_export({}, [111, 222], {"K": [2, 4]}, gpu_only=True) + connector._call(connector._queue_ipc_export([batch], 32)) + ticket, descriptors, _, _ = transfer.enqueue_snapshot_export_v1.await_args.args + keys, (payload,), schema = descriptors[0] + assert keys == [111, 222] and payload is packed and schema == layers + assert connector._ipc_export_jobs[ticket][2] == [packed] + connector._ipc_export_jobs.clear() + + +@pytest.mark.parametrize("outcome", ["recovery", "lost-ack", "failed-copy"]) +def test_arena_slots_survive_recovery_and_release_only_after_completion( + connector, connector_module, outcome +): + packed = torch.zeros(2, 16, dtype=torch.uint8) + arena = SimpleNamespace( + source="source", metadata=Mock(return_value=["fresh"]), release=Mock() + ) + connector._gpu_export_arena = arena + connector._gpu_export_registered = outcome == "recovery" + connector._exported_keys = {999} + connector._export_refresh_at = 100 + transfer = SimpleNamespace( + register_snapshot_export_source_v1=AsyncMock(), + enqueue_snapshot_export_v1=AsyncMock(), + poll_snapshot_exports_v1=AsyncMock(return_value={}), + ) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + entries = connector_module._PackedExportEntries( + [111, 222], packed, [("K", (4,), torch.float32, 0, 16)], 0 + ) + if outcome == "recovery": + transfer.enqueue_snapshot_export_v1.side_effect = [False, None] + elif outcome == "lost-ack": + transfer.enqueue_snapshot_export_v1.side_effect = RuntimeError("lost ack") + try: + if outcome == "lost-ack": + with pytest.raises(RuntimeError, match="lost ack"): + connector._call(connector._queue_ipc_export([(entries, [], [])], 32)) + assert not connector._gpu_export_registered + else: + connector._call(connector._queue_ipc_export([(entries, [], [])], 32)) + ticket = next(iter(connector._ipc_export_jobs)) + assert connector._gpu_export_tickets[ticket] == [0] + connector._poll_ipc_exports(wait=True) + arena.release.assert_not_called() + assert ticket in connector._ipc_export_jobs + transfer.register_snapshot_export_source_v1.assert_awaited_once_with( + "source", ["fresh"] + ) + if outcome == "recovery": + assert not connector._exported_keys and connector._export_refresh_at == 0 + if outcome == "failed-copy": + transfer.poll_snapshot_exports_v1.return_value = { + ticket: (None, "copy failed") + } + with pytest.raises(RuntimeError, match="copy failed"): + connector._poll_ipc_exports(wait=True) + else: + transfer.poll_snapshot_exports_v1.return_value = {ticket: ([111, 222], [])} + connector._poll_ipc_exports(wait=True) + arena.release.assert_called_once_with([0]) + assert not connector._ipc_export_jobs and not connector._gpu_export_tickets + finally: + connector._gpu_export_arena = None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.uint8]) +def test_fused_gather_preserves_mixed_group_bits_and_stream_order(dtype): + from ....xavier.backends.torch.packed_gather import PackedGather + + base = torch.arange(2 * 8 * 16 * 2 * 64, device="cuda", dtype=torch.float32) + k = base.remainder(197).to(dtype).view(2, 8, 16, 2, 64).permute(1, 0, 2, 3, 4) + v = torch.randn(8, 1, 8, 4, 64, device="cuda", dtype=torch.float32) + caches = {"K": k, "V": v} + sources = {"K": [5, 0, 6, 2], "V": [1, 7, 3, 0]} + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + gather = PackedGather.try_create(caches, {"K": 0, "V": 1}) + assert gather is not None + gather.warmup() + packed, owners = gather(sources) + expected = torch.cat( + [ + cache.index_select(0, torch.tensor(sources[name], device="cuda")) + .view(torch.uint8) + .reshape(4, -1) + for name, cache in caches.items() + ], + dim=1, + ) + # A later model step can overwrite every original source slot. + k.zero_() + v.zero_() + stream.synchronize() + assert torch.equal(packed, expected) + assert owners and packed.dtype == torch.uint8 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_unsupported_geometry_uses_per_layer_fallback(): + from ....xavier.backends.torch.packed_gather import PackedGather + + assert ( + PackedGather.try_create( + {"K": torch.ones(8, 3, device="cuda", dtype=torch.uint8)}, {"K": 0} + ) + is None + ) + assert ( + PackedGather.try_create( + {"K": torch.ones(8, 16, device="cuda")[:, ::2]}, {"K": 0} + ) + is None + ) + assert PackedGather.try_create({"K": torch.ones(8, 16)}, {"K": 0}) is None diff --git a/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py b/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py index 0c96df4f23..21a033da2b 100644 --- a/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py +++ b/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py @@ -35,7 +35,7 @@ def snapshot_export(connector): ), ) tracker = SimpleNamespace( - register_blocks=AsyncMock(), unregister_blocks=AsyncMock() + register_snapshot_blocks=AsyncMock(), unregister_blocks=AsyncMock() ) connector._get_transfer_ref = AsyncMock(return_value=transfer) connector._get_tracker_ref = AsyncMock(return_value=tracker) @@ -78,6 +78,57 @@ def test_skipped_tail_releases_previous_lease(connector, connector_module): assert not connector._requests_need_load and not connector._leased_requests +def test_cold_long_prompt_probes_only_first_uncached_block(connector): + request = SimpleNamespace(request_id="cold", prompt_token_ids=list(range(4097))) + connector._query_remote_blocks = AsyncMock(return_value={}) + connector._build_xavier_hashes = Mock(wraps=connector._build_xavier_hashes) + assert connector.get_num_new_matched_tokens(request, 0) == (0, False) + connector._build_xavier_hashes.assert_called_once_with(list(range(16))) + queried = connector._query_remote_blocks.await_args.args[1] + assert len(queried) == 1 and queried[0][1] == 0 + + +def test_positive_probe_loads_contiguous_suffix_from_multiple_ranks(connector): + request = SimpleNamespace(request_id="hit", prompt_token_ids=list(range(65))) + hashes = connector._build_xavier_hashes(request.prompt_token_ids[:-1]) + connector._query_remote_blocks = AsyncMock( + side_effect=[ + {2: {(hashes[1][0], 101, 1)}}, + {3: {(hashes[2][0], 202, 2)}, 2: {(hashes[3][0], 303, 3)}}, + ] + ) + connector._reserve_load_request = AsyncMock(return_value=True) + assert connector.get_num_new_matched_tokens(request, 16) == (48, False) + assert [ + call.args[1] for call in connector._query_remote_blocks.await_args_list + ] == [ + hashes[1:2], + hashes[2:], + ] + assert connector._requests_need_load["hit"].transfers == { + 2: {101: 1, 303: 3}, + 3: {202: 2}, + } + + +def test_positive_probe_stops_at_first_gap_and_reservation_failure_recomputes( + connector, +): + request = SimpleNamespace(request_id="gap", prompt_token_ids=list(range(65))) + hashes = connector._build_xavier_hashes(request.prompt_token_ids[:-1]) + connector._query_remote_blocks = AsyncMock( + side_effect=[ + {2: {(hashes[0][0], 101, 0)}}, + {3: {(hashes[2][0], 303, 2)}}, + ] + ) + connector._reserve_load_request = AsyncMock(return_value=False) + assert connector.get_num_new_matched_tokens(request, 0) == (0, False) + reserved = connector._reserve_load_request.await_args.args[0] + assert reserved.transfers == {2: {101: 0}} + assert not connector._requests_need_load and not connector._leased_requests + + def test_export_uses_post_forward_data_and_current_metadata( connector, connector_module, snapshot_export ): @@ -98,11 +149,11 @@ def test_export_uses_post_forward_data_and_current_metadata( assert store.ready == {222} assert store.read("K", [222]).tolist() == [[42, 42]] assert store.read("V", [222]).tolist() == [[84, 84]] - tracker.register_blocks.assert_awaited_once_with(0, [(222, 222)], 1) + tracker.register_snapshot_blocks.assert_awaited_once_with(0, [222], 1) assert not connector._request_staged_layers -def test_repeated_export_skips_payload_but_refreshes_discovery( +def test_repeated_export_skips_control_rpcs_until_discovery_refresh( connector, connector_module, snapshot_export ): store, transfer, tracker = snapshot_export @@ -113,8 +164,41 @@ def test_repeated_export_skips_payload_but_refreshes_discovery( connector._registered_kv_caches["K"][2].fill_(99) connector.wait_for_save() assert transfer.stage_layer_batches_v1.await_count == 1 - assert tracker.register_blocks.await_count == 2 + assert tracker.register_snapshot_blocks.await_count == 1 + assert transfer.ready_blocks_for_export_v1.await_count == 1 assert torch.equal(store.read("K", [111]), original) + # Periodic authoritative refresh restores lost discovery entries without + # copying a ready snapshot again. + connector._export_refresh_at = 0 + connector.wait_for_save() + assert tracker.register_snapshot_blocks.await_count == 2 + assert transfer.ready_blocks_for_export_v1.await_count == 2 + assert transfer.stage_layer_batches_v1.await_count == 1 + + +def test_cold_export_skips_readiness_rpc_and_still_refreshes_recovered_hits( + connector, connector_module, snapshot_export +): + store, transfer, _ = snapshot_export + connector._exported_keys = {111} + connector._export_refresh_at = 1.0 # Refresh is due for existing hits. + request = connector_module.XavierStoreRequest("new", [4], [222], [[4], [5]]) + bind_store_requests(connector, connector_module, request) + connector.wait_for_save() + transfer.ready_blocks_for_export_v1.assert_not_awaited() + assert 222 in store.ready + store.blocks.clear() + store.logical_dtypes.clear() + store.ready.clear() + connector._export_refresh_at = 1.0 + connector.wait_for_save() + transfer.ready_blocks_for_export_v1.assert_not_awaited() + # Publication supplies the authoritative ready set, removes stale mirrored + # keys after actor recovery, then republishes the newly copied snapshot. + assert transfer.publish_blocks_v1.await_count == 3 + assert transfer.stage_layer_batches_v1.await_count == 2 + assert 222 in store.ready + assert transfer.stage_layer_batches_v1.await_count == 2 def test_batch_dedup_preserves_positions_in_every_cache_group( @@ -135,9 +219,7 @@ def test_batch_dedup_preserves_positions_in_every_cache_group( assert store.read("V", [222, 333]).tolist() == [[22, 23], [26, 27]] transfer.stage_layer_batches_v1.assert_awaited_once() assert transfer.stage_layer_batches_v1.await_args.args[0][0][1] == [222, 333] - tracker.register_blocks.assert_awaited_once_with( - 0, [(111, 111), (222, 222), (333, 333)], 1 - ) + tracker.register_snapshot_blocks.assert_awaited_once_with(0, [111, 222, 333], 1) assert first.block_ids == [0, 2, 4] assert first.block_ids_by_group == [[0, 2, 4], [1, 3, 5]] @@ -173,13 +255,13 @@ def test_tracker_failure_retries_without_exporting_payload( store, transfer, tracker = snapshot_export request = connector_module.XavierStoreRequest("r", [2], [111], [[2], [3]]) bind_store_requests(connector, connector_module, request) - tracker.register_blocks.side_effect = RuntimeError("tracker unavailable") + tracker.register_snapshot_blocks.side_effect = RuntimeError("tracker unavailable") with pytest.raises(RuntimeError, match="tracker unavailable"): connector.wait_for_save() assert store.ready == {111} - tracker.register_blocks.side_effect = None + tracker.register_snapshot_blocks.side_effect = None connector.wait_for_save() - assert tracker.register_blocks.await_count == 2 + assert tracker.register_snapshot_blocks.await_count == 2 assert transfer.stage_layer_batches_v1.await_count == 1 @@ -194,6 +276,7 @@ def test_evicted_content_is_exported_again( bind_store_requests(connector, connector_module, request) connector.wait_for_save() assert store.ready == set(request.block_hashes) + assert connector._exported_keys == store.ready assert transfer.stage_layer_batches_v1.await_count == 3 assert tracker.unregister_blocks.await_count == 2 assert store.read("K", [111]).tolist() == [[4, 5]] @@ -257,8 +340,8 @@ def test_export_batches_bound_bytes_across_layers_and_requests( transfer.publish_blocks_v1.assert_awaited_once_with( [111, 222, 333, 444], {"K", "V"} ) - tracker.register_blocks.assert_awaited_once_with( - 0, [(key, key) for key in [111, 222, 333, 444]], 1 + tracker.register_snapshot_blocks.assert_awaited_once_with( + 0, [111, 222, 333, 444], 1 ) @@ -291,7 +374,7 @@ def test_batch_larger_than_available_capacity_keeps_complete_blocks( assert store.ready == {999, 222} assert store.read("V", [222]).tolist() == [[26, 27]] tracker.unregister_blocks.assert_awaited_once_with(0, 1, [111]) - tracker.register_blocks.assert_awaited_once_with(0, [(222, 222)], 1) + tracker.register_snapshot_blocks.assert_awaited_once_with(0, [222], 1) def test_batched_export_preserves_noncontiguous_bf16_and_snapshot_ownership( @@ -320,7 +403,7 @@ async def test_batched_registration_keeps_distinct_layer_requirements( publish_blocks_v1=AsyncMock(side_effect=[([111], [222]), ([222], [])]) ) tracker = SimpleNamespace( - register_blocks=AsyncMock(), unregister_blocks=AsyncMock() + register_snapshot_blocks=AsyncMock(), unregister_blocks=AsyncMock() ) connector._get_transfer_ref = AsyncMock(return_value=transfer) connector._get_tracker_ref = AsyncMock(return_value=tracker) @@ -336,7 +419,7 @@ async def test_batched_registration_keeps_distinct_layer_requirements( ([222], {"K", "V"}), ] tracker.unregister_blocks.assert_awaited_once_with(0, 1, [222]) - tracker.register_blocks.assert_awaited_once_with(0, [(111, 111), (222, 222)], 1) + tracker.register_snapshot_blocks.assert_awaited_once_with(0, [111, 222], 1) @pytest.mark.asyncio @@ -353,7 +436,7 @@ async def test_export_failure_drains_next_copy_without_publishing( for event in events: event.synchronize.assert_called_once() transfer.publish_blocks_v1.assert_not_awaited() - tracker.register_blocks.assert_not_awaited() + tracker.register_snapshot_blocks.assert_not_awaited() @pytest.mark.asyncio @@ -382,6 +465,7 @@ def test_cuda_export_fences_current_stream_and_owns_snapshot( connector, connector_module, snapshot_export, monkeypatch, dtype ): store, transfer, _ = snapshot_export + monkeypatch.setattr(connector, "_can_queue_ipc_export", lambda _: False) connector._registered_kv_caches = { name: tensor.to(device="cuda", dtype=dtype) for name, tensor in connector._registered_kv_caches.items() @@ -469,3 +553,141 @@ async def send(entries, metadata): await writer.wait_closed() server.close() await server.wait_closed() + + +def test_prompt_hash_cache_preserves_input_and_result_ownership( + connector, connector_module, monkeypatch +): + hashing = Mock(wraps=connector_module.hash_block_tokens) + monkeypatch.setattr(connector_module, "hash_block_tokens", hashing) + tokens = list(range(65)) + expected = connector._build_xavier_hashes(tokens) + assert hashing.call_count == 5 + again = connector._build_xavier_hashes(tokens.copy()) + assert again == expected and again is not expected + assert hashing.call_count == 5 + again.clear() + assert connector._build_xavier_hashes(tokens) == expected + tokens[0] = 999 + assert connector._build_xavier_hashes(tokens) != expected + assert hashing.call_count == 10 + connector._block_size = 32 + assert len(connector._build_xavier_hashes(tokens)) == 3 + assert hashing.call_count == 13 + + +def test_prompt_hash_cache_bounds_tokens_and_entries( + connector, connector_module, monkeypatch +): + monkeypatch.setattr(connector_module, "_MAX_HASH_CACHE_TOKENS", 64) + monkeypatch.setattr(connector_module, "_MAX_HASH_CACHE_ENTRIES", 2) + first = list(range(32)) + expected = connector._build_xavier_hashes(first) + for i in range(4): + connector._build_xavier_hashes([i + 100] * 24) + assert connector._hash_cache_tokens <= 64 and len(connector._hash_cache) <= 2 + before = list(connector._hash_cache) + connector._build_xavier_hashes(list(range(65))) + assert list(connector._hash_cache) == before + assert connector._build_xavier_hashes(first) == expected + + +def test_background_export_poll_does_not_block_and_queries_only_issued_tickets( + connector, +): + gate = asyncio.Event() + + async def poll(tickets, *, wait): + if not wait: + await gate.wait() + return {ticket: ([111 if ticket == "a" else 222], []) for ticket in tickets} + + transfer = SimpleNamespace(poll_snapshot_exports_v1=AsyncMock(side_effect=poll)) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + connector._call(asyncio.sleep(0)) + connector._ipc_export_jobs["a"] = ({111}, 8, []) + connector._poll_ipc_exports() + pending = connector._ipc_export_poll + connector._poll_ipc_exports() + assert pending is connector._ipc_export_poll and not pending.done() + # A new export must not enter the earlier poll before its acknowledgement. + connector._ipc_export_jobs["b"] = ({222}, 8, []) + transfer.poll_snapshot_exports_v1.assert_awaited_once_with(["a"], wait=False) + gate.set() + connector._poll_ipc_exports() + assert set(connector._ipc_export_jobs) == {"b"} + assert connector._exported_keys == {111} + assert connector._ipc_export_poll is None + # Backpressure drains an outstanding poll before waiting for the remainder. + gate.clear() + connector._ipc_export_poll_at = 0 + connector._poll_ipc_exports() + assert not connector._ipc_export_poll.done() + gate.set() + connector._poll_ipc_exports(wait=True) + assert not connector._ipc_export_jobs and connector._exported_keys == {111, 222} + + +def test_hot_and_cold_export_do_not_copy_entire_published_directory( + connector, connector_module, snapshot_export +): + class PublishedKeys(set): + def __or__(self, other): + pytest.fail("Export must not copy the entire published-key mirror") + + _, transfer, _ = snapshot_export + connector._exported_keys = PublishedKeys(range(100000)) + connector._export_refresh_at = float("inf") + hot = connector_module.XavierStoreRequest("hot", [2], [111], [[2], [3]]) + bind_store_requests(connector, connector_module, hot) + connector.wait_for_save() + transfer.ready_blocks_for_export_v1.assert_not_awaited() + transfer.stage_layer_batches_v1.assert_not_awaited() + cold = connector_module.XavierStoreRequest("cold", [4], [100001], [[4], [5]]) + bind_store_requests(connector, connector_module, cold) + connector.wait_for_save() + transfer.ready_blocks_for_export_v1.assert_not_awaited() + transfer.stage_layer_batches_v1.assert_awaited_once() + + +@pytest.mark.parametrize("actor_recovered", [False, True]) +def test_mixed_cold_prefix_repairs_expired_publication_once_then_uses_mirror( + connector, connector_module, snapshot_export, monkeypatch, actor_recovered +): + store, transfer, tracker = snapshot_export + for name in ["K", "V"]: + store.stage(name, [111], torch.ones(1, 2)) + store.publish([111], {"K", "V"}) + connector._exported_keys = {111} + connector._export_refresh_at = 10.0 + clock = [10.5] + monkeypatch.setattr( + connector_module, "time", SimpleNamespace(monotonic=lambda: clock[0]) + ) + monkeypatch.setattr(connector, "_can_queue_ipc_export", lambda _: True) + queued = Mock() + monkeypatch.setattr(connector, "_enqueue_ipc_export", queued) + if actor_recovered: + store.blocks.clear() + store.logical_dtypes.clear() + store.ready.clear() + first = connector_module.XavierStoreRequest( + "first", [2, 4], [111, 222], [[2, 4], [3, 5]] + ) + bind_store_requests(connector, connector_module, first) + connector.wait_for_save() + first_keys = queued.call_args.args[0][0][0][0][1] + assert first_keys == ([111, 222] if actor_recovered else [222]) + assert connector._exported_keys == (set() if actor_recovered else {111}) + assert connector._export_refresh_at == 11.5 + clock[0] = 10.6 + second = connector_module.XavierStoreRequest( + "second", [2, 6], [111, 333], [[2, 6], [3, 7]] + ) + bind_store_requests(connector, connector_module, second) + connector.wait_for_save() + assert queued.call_count == 2 + transfer.ready_blocks_for_export_v1.assert_not_awaited() + transfer.publish_blocks_v1.assert_awaited_once() + tracker.register_snapshot_blocks.assert_awaited_once() + assert connector._export_refresh_at == 11.5 diff --git a/xinference/model/llm/vllm/xavier/test/test_request_transfer.py b/xinference/model/llm/vllm/xavier/test/test_request_transfer.py index 22f44ad633..7682e90875 100644 --- a/xinference/model/llm/vllm/xavier/test/test_request_transfer.py +++ b/xinference/model/llm/vllm/xavier/test/test_request_transfer.py @@ -141,9 +141,10 @@ def test_connector_writes_correct_destinations( ) source.publish([77], set(caches)) transfer = SimpleNamespace( + get_local_read_metadata_v1=AsyncMock(return_value=None), read_request_blocks_v1=AsyncMock( side_effect=lambda rank, reads: pack_reads(source, reads) - ) + ), ) connector._get_transfer_ref = AsyncMock(return_value=transfer) request = connector_module.XavierLoadRequest( @@ -174,6 +175,46 @@ def test_connector_writes_correct_destinations( assert cache[:3].count_nonzero() == 0 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_connector_uploads_packed_payload_on_current_cuda_stream( + connector, connector_module, dtype +): + caches = { + name: torch.zeros(8, 4, dtype=dtype, device="cuda") for name in ["K", "V"] + } + connector._registered_kv_caches = caches + connector._layer_group_ids = {"K": 0, "V": 1} + source = KVSnapshotStore(2) + producer = SimpleNamespace(_snapshot_store=source, _rank=2) + values = torch.tensor([[1.0, -2.0, 32.0, 0.125], [9.0, 8.0, 7.0, 6.0]], dtype=dtype) + for name in caches: + TransferActor.stage_layer_blocks_v1(producer, "r", name, [77, 88], values) + source.publish([77, 88], set(caches)) + transfer = SimpleNamespace( + get_local_read_metadata_v1=AsyncMock(return_value=None), + read_request_blocks_v1=AsyncMock( + side_effect=lambda rank, reads: pack_reads(source, reads) + ), + ) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + request = connector_module.XavierLoadRequest( + "r", + {2: {77: 0, 88: 1}}, + local_transfers_by_group={0: {2: {77: 2, 88: 4}}, 1: {2: {77: 3, 88: 5}}}, + ) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + connector._load_request_blocks(request) + stream.synchronize() + transfer.read_request_blocks_v1.assert_awaited_once() + for name, destinations in (("K", [2, 4]), ("V", [3, 5])): + actual = caches[name][destinations].cpu() + assert torch.equal(actual.view(torch.uint8), values.view(torch.uint8)) + assert caches[name][0].count_nonzero() == 0 + + def test_payload_preserves_all_bf16_bit_patterns(): # Raw packing preserves the bits independently of the staging carrier. bits = torch.arange(65536, dtype=torch.int32).to(torch.int16).reshape(8, 8192) @@ -273,9 +314,10 @@ def test_connector_cross_layer_preserves_all_bf16_bits(connector, connector_modu TransferActor.stage_layer_blocks_v1(producer, "r", "K", [77], values) source.publish([77], {"K"}) transfer = SimpleNamespace( + get_local_read_metadata_v1=AsyncMock(return_value=None), read_request_blocks_v1=AsyncMock( side_effect=lambda rank, reads: pack_reads(source, reads) - ) + ), ) connector._get_transfer_ref = AsyncMock(return_value=transfer) request = connector_module.XavierLoadRequest( 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 352b3a764c..5c9fc8c50d 100644 --- a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py +++ b/xinference/model/llm/vllm/xavier/test/test_review_regressions.py @@ -126,7 +126,7 @@ def test_unsupported_configs_rejected(connector_module, connector_config, field) @pytest.mark.asyncio async def test_reused_request_id_publishes_new_content(connector, connector_module): tracker = SimpleNamespace( - register_blocks=AsyncMock(), unregister_blocks=AsyncMock() + register_snapshot_blocks=AsyncMock(), unregister_blocks=AsyncMock() ) transfer = SimpleNamespace( publish_blocks_v1=AsyncMock(side_effect=[([111], []), ([222], [111])]) @@ -138,8 +138,8 @@ async def test_reused_request_id_publishes_new_content(connector, connector_modu await connector._register_blocks( [connector_module.XavierStoreRequest("r", [1], [key], [[1]])] ) - assert tracker.register_blocks.await_count == 2 - tracker.register_blocks.assert_awaited_with(0, [(222, 222)], 1) + assert tracker.register_snapshot_blocks.await_count == 2 + tracker.register_snapshot_blocks.assert_awaited_with(0, [222], 1) tracker.unregister_blocks.assert_awaited_once_with(0, 1, [111]) assert not hasattr(connector, "_stored_requests") diff --git a/xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py b/xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py new file mode 100644 index 0000000000..2b2a08d8bb --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py @@ -0,0 +1,508 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +import asyncio +import sys +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +import torch +import xoscar as xo +from xoscar.backends.allocate_strategy import ProcessIndex + +from ..block_tracker import VLLMBlockTracker +from ..snapshot import KVSnapshotStore +from ..transfer import TransferActor + + +class SnapshotExportActor(xo.StatelessActor): + enqueue_snapshot_export_v1 = TransferActor.enqueue_snapshot_export_v1 + poll_snapshot_exports_v1 = TransferActor.poll_snapshot_exports_v1 + publish_blocks_v1 = TransferActor.publish_blocks_v1 + register_snapshot_export_source_v1 = ( + TransferActor.register_snapshot_export_source_v1 + ) + close_snapshot_export_source_v1 = TransferActor.close_snapshot_export_source_v1 + + def __init__(self): + self._rank = 1 + self._snapshot_store = KVSnapshotStore(8) + self._snapshot_export_jobs_v1 = {} + self._snapshot_export_lock_v1 = asyncio.Lock() + + def snapshot_bits(self, layer, keys): + store = self._snapshot_store + return store.read(layer, keys).view(torch.uint8).reshape(-1).tolist() + + async def __pre_destroy__(self): + await asyncio.gather( + *self._snapshot_export_jobs_v1.values(), return_exceptions=True + ) + getattr(self, "_snapshot_export_sources_v1", {}).clear() + + +@pytest.mark.skipif( + sys.platform != "linux" or not torch.cuda.is_available(), + reason="Linux CUDA required", +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.asyncio +async def test_gpu_arena_reuse_and_fresh_registration_after_actor_recovery(dtype): + from ....xavier.backends.torch.gpu_export import GPUExportArena + + pool = await xo.create_actor_pool( + "127.0.0.1", n_process=1, subprocess_start_method="spawn" + ) + async with pool: + tracker = await xo.create_actor( + VLLMBlockTracker, address=pool.external_address, uid="arena-tracker" + ) + actor = await xo.create_actor( + SnapshotExportActor, + address=pool.external_address, + uid="arena-export", + allocate_strategy=ProcessIndex(1), + ) + arena = GPUExportArena(4096, 1, torch.device("cuda", 0)) + expected = torch.ones(2, 2, dtype=dtype).view(torch.uint8).reshape(-1).tolist() + try: + for ticket, keys, fill in (("a", [111, 222], 1), ("b", [333, 444], 9)): + slot, value = arena.allocate((2, 4)) + value.view(dtype).fill_(float(fill)) + descriptor = arena.record(slot, tuple(value.shape)) + batch = [(keys, descriptor, [("K", (2,), dtype, 0, 4)])] + # Missing mappings consume no descriptors and start no writes. + if ticket == "a": + assert ( + await actor.enqueue_snapshot_export_v1( + ticket, batch, tracker.address, tracker.uid + ) + is False + ) + await actor.register_snapshot_export_source_v1( + arena.source, arena.metadata() + ) + await actor.enqueue_snapshot_export_v1( + ticket, batch, tracker.address, tracker.uid + ) + result = await actor.poll_snapshot_exports_v1([ticket], wait=True) + assert set(result[ticket][0]) == set(keys) + arena.release([slot]) + assert await actor.snapshot_bits("K", [111, 222]) == expected + await xo.destroy_actor(actor) + actor = await xo.create_actor( + SnapshotExportActor, + address=pool.external_address, + uid="arena-export", + allocate_strategy=ProcessIndex(1), + ) + slot, value = arena.allocate((2, 4)) + value.view(dtype).fill_(1) + descriptor = arena.record(slot, tuple(value.shape)) + batch = [([555, 666], descriptor, [("K", (2,), dtype, 0, 4)])] + assert ( + await actor.enqueue_snapshot_export_v1( + "c", batch, tracker.address, tracker.uid + ) + is False + ) + await actor.register_snapshot_export_source_v1( + arena.source, arena.metadata() + ) + await actor.enqueue_snapshot_export_v1( + "c", batch, tracker.address, tracker.uid + ) + result = await actor.poll_snapshot_exports_v1(["c"], wait=True) + assert set(result["c"][0]) == {555, 666} + arena.release([slot]) + await actor.close_snapshot_export_source_v1(arena.source) + finally: + arena.close() + assert await actor.snapshot_bits("K", [555, 666]) == expected + + +def test_packed_export_preserves_mixed_dtypes_and_block_ownership( + connector, connector_module, monkeypatch +): + from torch.multiprocessing import reductions + + monkeypatch.setattr(reductions, "reduce_tensor", lambda value: (None, (value,))) + transfer = SimpleNamespace( + enqueue_snapshot_export_v1=AsyncMock(), poll_snapshot_exports_v1=AsyncMock() + ) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + connector._registered_kv_caches = { + "K": torch.arange(8, dtype=torch.uint8).view(8, 1), + "V": torch.tensor([1e30, -1e30] * 8).to(torch.bfloat16).view(8, 2), + } + connector._layer_group_ids = {"K": 0, "V": 0} + request = connector_module.XavierStoreRequest("r", [2, 4], [111, 222], [[2, 4]]) + expected = connector._registered_kv_caches["V"][[2, 4]].view(torch.uint8).clone() + connector._call( + connector._queue_ipc_export(list(connector._cpu_export_batches([request])), 32) + ) + ticket, batches, _, _ = transfer.enqueue_snapshot_export_v1.await_args.args + keys, (packed,), layers = batches[0] + assert packed.shape == (2, 16) # dtype-aligned rows and layer offsets. + store = KVSnapshotStore(2) + store.stage_packed_blocks(keys, packed, layers) + store.publish(keys, {"K", "V"}) + packed.zero_() + assert store.read("K", keys).tolist() == [[2], [4]] + assert torch.equal(store.read("V", keys).view(torch.uint8), expected) + assert store.logical_dtypes[111]["V"] == torch.bfloat16 + assert store.blocks[111]["K"].untyped_storage().nbytes() == 16 + assert ( + store.blocks[111]["K"].untyped_storage().data_ptr() + != store.blocks[222]["K"].untyped_storage().data_ptr() + ) + transfer.poll_snapshot_exports_v1.return_value = {ticket: (keys, [])} + connector._poll_ipc_exports(wait=True) + assert not connector._ipc_export_jobs + assert connector._exported_keys == set(keys) + + +def test_ipc_export_failure_releases_completed_owners_and_invalidates_mirror(connector): + transfer = SimpleNamespace( + poll_snapshot_exports_v1=AsyncMock(return_value={"bad": (None, "copy failed")}) + ) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + connector._ipc_export_jobs["bad"] = ({111}, 16, [torch.ones(16)]) + connector._exported_keys = {111, 222} + with pytest.raises(RuntimeError, match="copy failed"): + connector._poll_ipc_exports(wait=True) + assert not connector._ipc_export_jobs and not connector._exported_keys + + +def test_new_exports_do_not_postpone_periodic_discovery_refresh(connector): + transfer = SimpleNamespace( + poll_snapshot_exports_v1=AsyncMock(return_value={"new": ([111], [])}) + ) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + connector._ipc_export_jobs["new"] = ({111}, 16, [torch.ones(16)]) + connector._export_refresh_at = ( + 1.0 # already due, even under continuous cold traffic + ) + connector._poll_ipc_exports(wait=True) + assert connector._export_refresh_at == 1.0 + + +def test_lost_enqueue_ack_keeps_ipc_producer_storage( + connector, connector_module, monkeypatch +): + from torch.multiprocessing import reductions + + monkeypatch.setattr(reductions, "reduce_tensor", lambda value: (None, (value,))) + transfer = SimpleNamespace( + enqueue_snapshot_export_v1=AsyncMock( + side_effect=RuntimeError("connection lost") + ), + poll_snapshot_exports_v1=AsyncMock(), + ) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + connector._registered_kv_caches = {"K": torch.ones(8, 2)} + request = connector_module.XavierStoreRequest("r", [2], [111], [[2]]) + with pytest.raises(RuntimeError, match="connection lost"): + connector._call( + connector._queue_ipc_export( + list(connector._cpu_export_batches([request])), 8 + ) + ) + ticket = transfer.enqueue_snapshot_export_v1.await_args.args[0] + assert connector._ipc_export_jobs[ticket][2][0].numel() == 8 + transfer.poll_snapshot_exports_v1.return_value = {ticket: ([111], [])} + connector._poll_ipc_exports(wait=True) + assert not connector._ipc_export_jobs + + +@pytest.mark.parametrize("lost_ack", [False, True]) +def test_async_arena_submission_excludes_unacknowledged_jobs_and_retains_owners( + connector, connector_module, lost_ack +): + gate = asyncio.Event() + actor_jobs = {} + + async def enqueue(ticket, batches, *_): + actor_jobs[ticket] = batches + await gate.wait() + if lost_ack: + raise RuntimeError("lost acknowledgement") + + async def poll(tickets, *, wait): + assert all(ticket in actor_jobs for ticket in tickets) + return {ticket: ([111], []) for ticket in tickets} + + transfer = SimpleNamespace( + enqueue_snapshot_export_v1=AsyncMock(side_effect=enqueue), + poll_snapshot_exports_v1=AsyncMock(side_effect=poll), + close_snapshot_export_source_v1=AsyncMock(), + ) + arena = SimpleNamespace(source="source", release=Mock(), close=Mock()) + connector._gpu_export_arena = arena + connector._gpu_export_registered = True + connector._get_transfer_ref = AsyncMock(return_value=transfer) + packed = torch.ones(1, 8, dtype=torch.uint8) + indices = torch.tensor([2]) + entries = connector_module._PackedExportEntries( + [111], packed, [("K", (8,), torch.uint8, 0, 8)], slot=0 + ) + connector._enqueue_ipc_export([(entries, [indices], [])], 8) + ticket = next(iter(connector._ipc_export_enqueues)) + assert not connector._ipc_export_enqueues[ticket].done() + assert connector._ipc_export_jobs[ticket][2] == [packed, indices] + connector._poll_ipc_exports() + transfer.poll_snapshot_exports_v1.assert_not_awaited() + arena.release.assert_not_called() + gate.set() + if lost_ack: + with pytest.raises(RuntimeError, match="lost acknowledgement"): + connector._poll_ipc_exports(wait=True) + assert connector._ipc_export_jobs[ticket][2] == [packed, indices] + assert connector._gpu_export_tickets[ticket] == [0] + arena.release.assert_not_called() + # The source-close handshake must precede freeing its GPU owners. + connector._transfer_ref = transfer + transfer.close_snapshot_export_source_v1.side_effect = ( + lambda _: arena.close.assert_not_called() + ) + connector.shutdown() + transfer.close_snapshot_export_source_v1.assert_awaited_once_with("source") + arena.close.assert_called_once() + else: + connector._poll_ipc_exports(wait=True) + assert not connector._ipc_export_jobs and not connector._ipc_export_enqueues + assert connector._exported_keys == {111} + arena.release.assert_called_once_with([0]) + + +def test_failed_submission_drain_waits_for_all_late_tasks(connector): + completed = [] + + async def submit(ticket): + await asyncio.sleep(0) + completed.append(ticket) + if ticket == "bad": + raise RuntimeError("lost acknowledgement") + + connector._call(asyncio.sleep(0)) + for ticket in ("bad", "late"): + connector._ipc_export_enqueues[ticket] = connector._loop.create_task( + submit(ticket) + ) + connector._ipc_export_jobs[ticket] = ({111}, 8, [torch.ones(8)]) + with pytest.raises(RuntimeError, match="lost acknowledgement"): + connector._poll_ipc_exports(wait=True) + assert completed == ["bad", "late"] and not connector._ipc_export_enqueues + assert set(connector._ipc_export_jobs) == {"bad", "late"} + connector._ipc_export_jobs.clear() + + +def test_packed_snapshot_respects_leases_and_immutable_partial_content(): + store = KVSnapshotStore(1) + layers = [("K", (1,), torch.float32, 0, 4), ("V", (1,), torch.float32, 4, 4)] + payload = torch.tensor([[1.0, 2.0], [3.0, 4.0]]).view(torch.uint8) + store.stage("K", [111], torch.tensor([[99.0]])) + store.stage_packed_blocks([111], payload[:1], layers) + store.publish([111], {"K", "V"}) + assert store.read("K", [111]).item() == 99.0 + assert store.read("V", [111]).item() == 2.0 + assert store.reserve("reader", [111]) + store.stage_packed_blocks([222], payload[1:], layers) + assert store.ready == {111} + assert set(store.blocks) == {111} + store.release("reader") + store.stage_packed_blocks([222], payload[1:], layers) + assert store.publish([222], {"K", "V"}) == [222] + assert store.ready == {222} and store.evicted == {111} + + +def test_retained_slabs_charge_and_evict_whole_allocations_with_leases(): + store = KVSnapshotStore(4) + layers = [("K", (2,), torch.float32, 0, 8)] + payloads = [torch.full((2, 2), float(i)).view(torch.uint8) for i in range(3)] + for keys, payload in zip(([1, 2], [3, 4]), payloads): + store.stage_packed_blocks(keys, payload, layers, retain_storage=True) + store.publish(keys, {"K"}) + assert store._used_slots() == 4 + assert store.reserve("reader", [1]) + store.stage_packed_blocks([5, 6], payloads[2], layers, retain_storage=True) + store.publish([5, 6], {"K"}) + assert store.ready == {1, 2, 5, 6} + assert store.evicted == {3, 4} + assert store._used_slots() == 4 + assert torch.equal(store.read("K", [1, 2]), payloads[0].view(torch.float32)) + assert store.blocks[5]["K"].untyped_storage().data_ptr() == ( + payloads[2].untyped_storage().data_ptr() + ) + store.release("reader") + store.stage("K", [7], torch.ones(1, 2)) + assert set(store.blocks) == {5, 6, 7} + assert store._used_slots() == 3 + + +def test_retained_payload_charges_backing_storage_and_falls_back_when_oversized(): + store = KVSnapshotStore(4) + layers = [("K", (2,), torch.float32, 0, 8)] + backing = torch.ones(10, 2).view(torch.uint8) + store.stage_packed_blocks([1, 2], backing[:2], layers, retain_storage=True) + store.publish([1, 2], {"K"}) + backing.zero_() + assert store.read("K", [1, 2]).tolist() == [[1.0, 1.0], [1.0, 1.0]] + assert store._used_slots() == 2 and not store._packed_slabs + assert store.blocks[1]["K"].untyped_storage().nbytes() == 8 + + +def test_packed_reads_preserve_order_and_bits_without_materializing_layer_views(): + store = KVSnapshotStore(2) + payload = torch.tensor([[1e30, -1e30], [3e30, -3e30]]).to(torch.bfloat16) + packed = payload.view(torch.uint8) + layers = [("K", (2,), torch.bfloat16, 0, 4)] + store.stage_packed_blocks([1, 2], packed, layers, retain_storage=True) + store.publish([1, 2], {"K"}) + result = store.read("K", [2, 1]) + assert torch.equal(result.view(torch.uint8), packed[[1, 0]]) + assert all(not block._layers for block in store.blocks.values()) + # MutableMapping updates must not keep using the original packed bytes. + store.blocks[1]["K"] = torch.zeros(2, dtype=torch.float16) + assert store.read("K", [1, 2])[0].count_nonzero() == 0 + + +def test_compact_blocks_share_one_slab_tensor_and_preserve_mapping_mutations(): + store = KVSnapshotStore(2) + packed = torch.arange(8).reshape(2, 4).float().view(torch.uint8) + layers = [("K", (2,), torch.float32, 0, 8), ("V", (2,), torch.float32, 8, 8)] + store.stage_packed_blocks([1, 2], packed, layers, retain_storage=True) + first, second = store.blocks.values() + assert first._packed is packed and second._packed is packed + assert [first._row, second._row] == [0, 1] + assert not first._layers and first._overrides is None and first._removed is None + first["K"] = torch.tensor([99.0, 100.0]) + assert store.read("K", [1, 2]).tolist() == [[99.0, 100.0], [4.0, 5.0]] + del first["V"] + assert set(first) == {"K"} and len(first) == 1 + assert store.publish([1, 2], {"K", "V"}) == [2] + first["V"] = torch.tensor([88.0, 89.0]) + assert set(first) == {"K", "V"} and len(first) == 2 + assert store.publish([1], {"K", "V"}) == [1] + + +def test_packed_arena_submission_does_not_create_unused_layer_tensor_views( + connector, connector_module, monkeypatch +): + transfer = SimpleNamespace(enqueue_snapshot_export_v1=AsyncMock()) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + connector._gpu_export_arena = SimpleNamespace(source="source") + connector._gpu_export_registered = True + packed = torch.ones(2, 16, dtype=torch.uint8) + + def unexpected_view(*_): + pytest.fail("Packed IPC metadata must not materialize layer tensor views") + + monkeypatch.setattr(torch.Tensor, "__getitem__", unexpected_view) + try: + entries = connector_module._PackedExportEntries( + [111, 222], packed, [("K", (4,), torch.float32, 0, 16)], slot=0 + ) + connector._call(connector._queue_ipc_export([(entries, [], [])], 32)) + transfer.enqueue_snapshot_export_v1.assert_awaited_once() + finally: + connector._gpu_export_arena = None + connector._gpu_export_tickets.clear() + connector._ipc_export_jobs.clear() + + +def test_ipc_backpressure_waits_before_gathering( + connector, connector_module, monkeypatch +): + from torch.multiprocessing import reductions + + transfer = SimpleNamespace(ready_blocks_for_export_v1=AsyncMock(return_value=[])) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + connector._registered_kv_caches = {"K": torch.ones(8, 2), "V": torch.ones(8, 2)} + connector._layer_group_ids = {"K": 0, "V": 1} + transfer.enqueue_snapshot_export_v1 = AsyncMock() + transfer.poll_snapshot_exports_v1 = AsyncMock(return_value={"old": ([111], [])}) + connector._ipc_export_jobs["old"] = ({111}, 16, [torch.ones(16)]) + connector._ipc_export_poll_at = float("inf") + monkeypatch.setattr(connector_module, "_MAX_PENDING_EXPORT_BYTES", 16) + connector._ipc_export_budget = 16 + monkeypatch.setattr(connector, "_can_queue_ipc_export", lambda _: True) + monkeypatch.setattr(reductions, "reduce_tensor", lambda value: (None, (value,))) + original_prepare = connector._prepare_cpu_export + + def prepare(*args, **kwargs): + assert not connector._ipc_export_jobs + return original_prepare(*args, **kwargs) + + monkeypatch.setattr(connector, "_prepare_cpu_export", prepare) + request = connector_module.XavierStoreRequest("new", [4], [222], [[4], [5]]) + connector._get_connector_metadata = ( + lambda: connector_module.XavierConnectorMetadata(store_requests=[request]) + ) + connector.wait_for_save() + transfer.poll_snapshot_exports_v1.assert_awaited_once_with(["old"], wait=True) + ticket = transfer.enqueue_snapshot_export_v1.await_args.args[0] + assert connector._ipc_export_jobs[ticket][0] == {222} + transfer.poll_snapshot_exports_v1.return_value = {ticket: ([222], [])} + connector._poll_ipc_exports(wait=True) + + +@pytest.mark.skipif( + sys.platform != "linux" or not torch.cuda.is_available(), + reason="Linux CUDA IPC required", +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.asyncio +async def test_actor_publishes_gpu_snapshot_without_producer_polling(dtype): + from torch.multiprocessing.reductions import reduce_tensor + + pool = await xo.create_actor_pool( + "127.0.0.1", n_process=1, subprocess_start_method="spawn" + ) + async with pool: + tracker = await xo.create_actor( + VLLMBlockTracker, address=pool.external_address, uid="snapshot-tracker" + ) + actor = await xo.create_actor( + SnapshotExportActor, + address=pool.external_address, + uid="snapshot-export", + allocate_strategy=ProcessIndex(1), + ) + source = ( + torch.arange(16, dtype=torch.float32, device="cuda").to(dtype).view(8, 2) + ) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + gathered = source.index_select(0, torch.tensor([2, 4], device="cuda")) + packed = gathered.view(torch.uint8).reshape(2, -1) + descriptor = reduce_tensor(packed)[1] + expected = ( + torch.arange(16) + .to(dtype) + .view(8, 2)[[2, 4]] + .view(torch.uint8) + .reshape(-1) + .tolist() + ) + await actor.enqueue_snapshot_export_v1( + "export", + [([111, 222], descriptor, [("K", (2,), dtype, 0, 4)])], + tracker.address, + tracker.uid, + ) + with torch.cuda.stream(stream): + source.zero_() + # Do not drive any producer connector loop. The actor must publish and + # answer discovery while the producer is idle. + for _ in range(100): + remote = await tracker.query_blocks(0, [(111, 0), (222, 1)]) + if remote: + break + await asyncio.sleep(0.01) + assert remote == {1: {(111, 111, 0), (222, 222, 1)}} + assert await actor.snapshot_bits("K", [111, 222]) == expected + result = await actor.poll_snapshot_exports_v1(["export"], wait=True) + assert set(result["export"][0]) == {111, 222} + assert result["export"][1] == [] diff --git a/xinference/model/llm/vllm/xavier/test/test_snapshot_gc_lifecycle.py b/xinference/model/llm/vllm/xavier/test/test_snapshot_gc_lifecycle.py new file mode 100644 index 0000000000..c4b47a6ed4 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_snapshot_gc_lifecycle.py @@ -0,0 +1,108 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + + +@pytest.mark.parametrize("mode", ["enabled", "disabled", "already-frozen"]) +def test_snapshot_gc_preserves_new_request_cycles_and_external_owner(mode): + source = Path(__file__).parents[1] / "gc_lifecycle.py" + script = r""" +import gc, runpy, sys, weakref +guard = runpy.run_path(sys.argv[1])["InitializationGCFreeze"]() +mode = sys.argv[2] +if mode == "disabled": + gc.disable() +if mode == "already-frozen": + gc.collect() + gc.freeze() +was_enabled, was_frozen = gc.isenabled(), bool(gc.get_freeze_count()) +class Library: + pass +library = Library() +library.cycle = library +library_ref = weakref.ref(library) +try: + guard.start() + assert gc.isenabled() == was_enabled + del library + gc.collect() + assert library_ref() is not None + request = Library() + request.cycle = request + request_ref = weakref.ref(request) + guard.start() + del request + gc.collect() + assert request_ref() is None, "new request cycles must remain collectible" +finally: + guard.close() + guard.close() +assert gc.isenabled() == was_enabled +assert bool(gc.get_freeze_count()) == was_frozen +if was_frozen: + gc.unfreeze() +gc.collect() +assert library_ref() is None +""" + result = subprocess.run( + [sys.executable, "-c", script, str(source), mode], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.asyncio +async def test_actor_starts_gc_guard_and_closes_it_after_cleanup_error(): + from ..transfer import TransferActor + + guard = SimpleNamespace(start=Mock(), close=Mock()) + actor = SimpleNamespace( + init_rank=Mock(), + _snapshot_gc_freeze=guard, + close_gpu_caches_v1=AsyncMock(), + _layer_send_tasks_v1=set(), + _context=SimpleNamespace( + closeConnections=Mock(side_effect=RuntimeError("stop")) + ), + ) + await TransferActor.__post_create__(actor) + guard.start.assert_called_once() + with pytest.raises(RuntimeError, match="stop"): + await TransferActor.__pre_destroy__(actor) + guard.close.assert_called_once() + + +@pytest.mark.parametrize("v1,gpu_budget", [(True, None), (True, 64), (False, None)]) +@pytest.mark.asyncio +async def test_model_only_freezes_ordinary_v1_snapshot_runtime( + monkeypatch, v1, gpu_budget +): + import xoscar as xo + + from xinference.core.model import ModelActor + from xinference.model.llm.vllm import core + + class Model: + _is_vllm_v1 = Mock(return_value=v1) + init_xavier = AsyncMock() + + monkeypatch.setattr(core, "VLLMModel", Model) + create = AsyncMock(return_value=SimpleNamespace(address="local")) + monkeypatch.setattr(xo, "create_actor", create) + actor = SimpleNamespace( + _model=Model(), + address="local", + _xavier_config={"rank": 1, "gpu_cache_bytes": gpu_budget}, + ) + await ModelActor.start_transfer_for_vllm(actor, ["local"]) + assert create.await_args.kwargs["freeze_initialization_gc"] == ( + v1 and gpu_budget is None + ) diff --git a/xinference/model/llm/vllm/xavier/test/test_v1_engine.py b/xinference/model/llm/vllm/xavier/test/test_v1_engine.py index 288e4f18d3..510dc8258b 100644 --- a/xinference/model/llm/vllm/xavier/test/test_v1_engine.py +++ b/xinference/model/llm/vllm/xavier/test/test_v1_engine.py @@ -60,6 +60,13 @@ def test_v1_connector_configuration(engine, role, kv_role): assert args.kv_transfer_config.kv_role == kv_role assert args.kv_transfer_config.kv_connector == "XavierConnector" assert args.additional_config["custom"] == 1 + assert args.additional_config == {"custom": 1} + assert ( + args.kv_transfer_config.kv_connector_extra_config["xavier_config"][ + "block_tracker_uid" + ] + == b"tracker" + ) json.dumps(args.additional_config) assert not args.enforce_eager assert getattr(args, "disable_hybrid_kv_cache_manager", None) is None @@ -68,6 +75,29 @@ def test_v1_connector_configuration(engine, role, kv_role): factory.assert_called_once() +def test_transport_changes_preserve_user_compilation_config(engine): + module, _ = engine + additional = {"custom": {"graph_option": 1}} + for rank in [1, 2]: + args = SimpleNamespace( + additional_config=additional, + create_model_config=lambda: SimpleNamespace(is_hybrid=False), + ) + module.XavierEngine._patch_v1_engine_args( + args, + { + "rank": rank, + "rank_address": f"localhost:{1000 + rank}", + "block_tracker_uid": b"runtime-only", + }, + ) + assert args.additional_config is additional + assert ( + args.kv_transfer_config.kv_connector_extra_config["xavier_config"]["rank"] + == rank + ) + + @pytest.mark.parametrize("recurrent", [False, True]) @pytest.mark.parametrize("launch_eager", [False, True]) @pytest.mark.parametrize("xavier_eager", [None, False, True]) diff --git a/xinference/model/llm/vllm/xavier/transfer.py b/xinference/model/llm/vllm/xavier/transfer.py index ac048a9e21..421d9c80f8 100644 --- a/xinference/model/llm/vllm/xavier/transfer.py +++ b/xinference/model/llm/vllm/xavier/transfer.py @@ -13,6 +13,8 @@ # limitations under the License. import asyncio import logging +import math +import sys from functools import lru_cache from queue import Queue from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, no_type_check @@ -25,6 +27,7 @@ from ...xavier.collective import CollectiveRank from ...xavier.constants import DEFAULT_TRANSFER_ACTOR_UID from ...xavier.profiling import profile_stage +from .gc_lifecycle import InitializationGCFreeze try: from vllm.utils import TORCH_DTYPE_TO_NUMPY_DTYPE, Device @@ -136,6 +139,7 @@ def __init__( store_address: str, store_port: int, world_addresses: List[str], + freeze_initialization_gc: bool = False, ): super().__init__() CollectiveRank.__init__( @@ -153,10 +157,22 @@ def __init__( self._kv_schema_v1: Optional[Tuple[int, dict, str]] = None self._schema_mismatch_warnings: Set[Tuple[str, str]] = set() self._layer_send_tasks_v1: Set[asyncio.Task[Any]] = set() + self._snapshot_export_jobs_v1: Dict[str, asyncio.Task] = {} + self._snapshot_export_lock_v1 = asyncio.Lock() + self._snapshot_export_sources_v1: Dict[ + str, List[Tuple[torch.Tensor, torch.cuda.Event]] + ] = {} + self._snapshot_gc_freeze = ( + InitializationGCFreeze() if freeze_initialization_gc else None + ) self._swap_stream = torch.cuda.Stream() async def __post_create__(self): self.init_rank() + if self._snapshot_gc_freeze is not None: + # Actor creation follows model loading and precedes serving. Freeze + # once here; snapshots and request cycles allocated later still GC. + self._snapshot_gc_freeze.start() def setup( self, @@ -176,15 +192,33 @@ def setup( async def __pre_destroy__(self): try: - await self.close_gpu_caches_v1() - except Exception: - logger.warning( - "Failed to close Xavier GPU caches; continuing transfer cleanup", - exc_info=True, - ) - for task in self._layer_send_tasks_v1: - task.cancel() - self._context.closeConnections() + if getattr(self, "_local_read_tasks_v1", None): + await asyncio.gather(*self._local_read_tasks_v1, return_exceptions=True) + read_buffer = getattr(self, "_local_read_buffer_v1", None) + if read_buffer is not None: + read_buffer.close() + self._local_read_buffer_v1 = None + if getattr(self, "_snapshot_export_jobs_v1", None): + await asyncio.gather( + *self._snapshot_export_jobs_v1.values(), return_exceptions=True + ) + self._snapshot_export_jobs_v1.clear() + if getattr(self, "_snapshot_export_sources_v1", None): + self._snapshot_export_sources_v1.clear() + try: + await self.close_gpu_caches_v1() + except Exception: + logger.warning( + "Failed to close Xavier GPU caches; continuing transfer cleanup", + exc_info=True, + ) + for task in self._layer_send_tasks_v1: + task.cancel() + self._context.closeConnections() + finally: + guard = getattr(self, "_snapshot_gc_freeze", None) + if guard is not None: + guard.close() def _get_cache_engine(self, virtual_engine: int) -> CacheEngine: return self._cache_engine[virtual_engine] # type: ignore @@ -278,6 +312,115 @@ def configure_kv_schema_v1(self, block_size: int, layers: dict, cache_dtype: str raise ValueError("Xavier registered KV schema changed") self._kv_schema_v1 = schema + def register_snapshot_export_source_v1(self, source, metadata): + from torch.multiprocessing.reductions import rebuild_cuda_tensor + + mapped = [] + for descriptor, handle in metadata: + value = rebuild_cuda_tensor(*descriptor) + ready = torch.cuda.Event.from_ipc_handle(value.device, handle) + mapped.append((value, ready)) + sources = getattr(self, "_snapshot_export_sources_v1", None) + if sources is None: + sources = self._snapshot_export_sources_v1 = {} + sources[source] = mapped + + async def close_snapshot_export_source_v1(self, source): + if getattr(self, "_snapshot_export_jobs_v1", None): + await asyncio.gather( + *self._snapshot_export_jobs_v1.values(), return_exceptions=True + ) + sources = getattr(self, "_snapshot_export_sources_v1", None) + if sources is not None: + sources.pop(source, None) + + def enqueue_snapshot_export_v1(self, ticket, batches, tracker_address, tracker_uid): + """Own immutable GPU gathers before acknowledging the producer.""" + from torch.multiprocessing.reductions import rebuild_cuda_tensor + + sources = getattr(self, "_snapshot_export_sources_v1", {}) + if any( + isinstance(descriptor, dict) and descriptor["gpu_export"] not in sources + for _, descriptor, _ in batches + ): + return False # Actor recovery requires fresh IPC refcounts. + mapped = [] + for keys, descriptor, layers in batches: + ready = None + if isinstance(descriptor, dict): + buffer, ready = sources[descriptor["gpu_export"]][descriptor["slot"]] + shape = descriptor["shape"] + value = buffer[: math.prod(shape)].view(shape) + else: + value = rebuild_cuda_tensor(*descriptor) + mapped.append((keys, value, layers, ready)) + batches = mapped + + async def export(): + available, evicted = set(), set() + async with self._snapshot_export_lock_v1: + tracker = await xo.actor_ref(address=tracker_address, uid=tracker_uid) + for keys, value, layers, ready_event in batches: + event = torch.cuda.Event() + recorded = False + try: + if ready_event is not None: + torch.cuda.current_stream(value.device).wait_event( + ready_event + ) + host = torch.empty_like(value, device="cpu", pin_memory=True) + host.copy_(value, non_blocking=True) + event.record(torch.cuda.current_stream(value.device)) + recorded = True + # Publication progresses independently of EngineCore. + while not event.query(): + await asyncio.sleep(0.001) + # Own this immutable D2H allocation without copying it + # again. The store charges and evicts whole slabs, so + # retained memory remains bounded by its capacity. + self._snapshot_store.stage_packed_blocks( + keys, host, layers, retain_storage=True + ) + ready, removed = self.publish_blocks_v1( + keys, {layer[0] for layer in layers} + ) + if removed: + await tracker.unregister_blocks(0, self._rank, removed) + await tracker.register_snapshot_blocks(0, ready, self._rank) + available.difference_update(removed) + available.update(ready) + evicted.update(removed) + evicted.difference_update(ready) + finally: + if recorded: + event.synchronize() + else: + torch.cuda.current_stream(value.device).synchronize() + return list(available), list(evicted) + + self._snapshot_export_jobs_v1[ticket] = asyncio.create_task(export()) + + async def poll_snapshot_exports_v1(self, tickets, *, wait=False): + tasks = { + ticket: self._snapshot_export_jobs_v1.get(ticket) for ticket in tickets + } + if wait: + await asyncio.gather( + *(task for task in tasks.values() if task is not None), + return_exceptions=True, + ) + results = {} + for ticket, task in tasks.items(): + if task is None: + results[ticket] = (None, "Unknown snapshot export ticket") + elif task.done(): + self._snapshot_export_jobs_v1.pop(ticket) + try: + results[ticket] = task.result() + except BaseException as error: + results[ticket] = (None, repr(error)) + return results + def reserve_blocks_v1(self, lease, keys, expected_schema): schema = self._kv_schema_v1 if not schema or not schema[1]: @@ -560,7 +703,53 @@ def completed(future): task.add_done_callback(completed) - async def read_request_blocks_v1(self, from_rank, reads): + def get_local_read_metadata_v1(self): + if sys.platform != "linux": + return None + from ...xavier.backends.torch.local_read import SharedReadBuffer + + owner = getattr(self, "_local_read_buffer_v1", None) + try: + if owner is None: + owner = self._local_read_buffer_v1 = SharedReadBuffer() + self._local_read_tasks_v1 = set() + return owner.metadata() + except OSError: + if owner is not None: + owner.close() + self._local_read_buffer_v1 = None + return None + + async def read_request_blocks_local_v1(self, from_rank, reads, token): + from ...xavier.backends.torch.gpu_transfer import finish_before_cancel + + owner = getattr(self, "_local_read_buffer_v1", None) + if owner is None or owner.token != token: + return None + acquired = owner.acquire(sum(read.nbytes for read in reads)) + if acquired is None: + return None # Busy or oversized messages use the ordinary RPC path. + lease, payload = acquired + task = asyncio.create_task( + self.read_request_blocks_v1(from_rank, reads, _buffer=payload) + ) + self._local_read_tasks_v1.add(task) + try: + # Cancellation cannot release/unmap a buffer while Gloo still writes. + await finish_before_cancel(task) + except BaseException: + owner.release(lease) + raise + finally: + self._local_read_tasks_v1.discard(task) + return lease + + def release_local_read_v1(self, lease): + owner = getattr(self, "_local_read_buffer_v1", None) + if owner is not None: + owner.release(lease) + + async def read_request_blocks_v1(self, from_rank, reads, *, _buffer=None): from xoscar.collective import xoscar_pygloo as xp fields = dict( @@ -568,7 +757,11 @@ async def read_request_blocks_v1(self, from_rank, reads): nbytes=sum(read.nbytes for read in reads), blocks=sum(len(read.keys) for read in reads), ) - payload = torch.empty(fields["nbytes"], dtype=torch.uint8) + payload = ( + torch.empty(fields["nbytes"], dtype=torch.uint8) + if _buffer is None + else _buffer + ) with profile_stage("actor_control", **fields): sender = await xo.actor_ref( address=self._world_addresses[from_rank], diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index d359ac6e5c..f02e884155 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -14,7 +14,10 @@ import asyncio import logging import math +import sys +import time import uuid +from collections import OrderedDict from dataclasses import dataclass, field from typing import ( TYPE_CHECKING, @@ -26,6 +29,7 @@ Optional, Sequence, Tuple, + Union, ) import torch @@ -55,13 +59,49 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request + from ...xavier.backends.torch.gpu_export import GPUExportArena + from ...xavier.backends.torch.local_read import LocalReadBuffer + from ...xavier.backends.torch.packed_gather import PackedGather + from ...xavier.local_directory import LocalBlockDirectory + logger = logging.getLogger(__name__) _MAX_READ_BYTES = 1024 * 1024 _MAX_READ_BLOCKS = 64 _MAX_EXPORT_BYTES = 32 * 1024 * 1024 +_MAX_PENDING_EXPORT_BYTES = 8 * _MAX_EXPORT_BYTES +_EXPORT_REFRESH_SECONDS = 1.0 +_MAX_HASH_CACHE_TOKENS = 128 * 1024 +_MAX_HASH_CACHE_ENTRIES = 128 + + +@dataclass +class _PackedExportEntries: + keys: List[int] + packed: torch.Tensor + layers: List[Tuple[str, Tuple[int, ...], torch.dtype, int, int]] + slot: Optional[int] = None + + def __len__(self) -> int: + return len(self.layers) + + def __getitem__(self, index: int) -> Tuple[str, List[int], torch.Tensor]: + name, shape, dtype, offset, nbytes = self.layers[index] + return ( + name, + self.keys, + self.packed[:, offset : offset + nbytes] + .view(dtype) + .view(len(self.keys), *shape), + ) + + def __iter__(self) -> Iterator[Tuple[str, List[int], torch.Tensor]]: + for index in range(len(self)): + yield self[index] + + _CPUExportBatch = Tuple[ - List[Tuple[str, List[int], torch.Tensor]], + Union[List[Tuple[str, List[int], torch.Tensor]], _PackedExportEntries], List[torch.Tensor], List[torch.cuda.Event], ] @@ -220,11 +260,62 @@ def __init__( self._tracker_ref: Optional[xo.ActorRefType["VLLMBlockTracker"]] = None self._transfer_ref: Optional[xo.ActorRefType["TransferActor"]] = None self._loop: Optional[asyncio.AbstractEventLoop] = None + self._exported_keys: set[int] = set() + self._export_refresh_at = 0.0 + self._cpu_export_streams: Dict[torch.device, torch.cuda.Stream] = {} + self._ipc_export_jobs: Dict[str, Tuple[set[int], int, List[torch.Tensor]]] = {} + self._ipc_export_poll_at = 0.0 + self._ipc_export_poll: Optional[asyncio.Task] = None + self._ipc_export_enqueues: Dict[str, asyncio.Task] = {} + self._ipc_export_budget = _MAX_PENDING_EXPORT_BYTES + self._packed_gather: Optional["PackedGather"] = None + self._gpu_export_arena: Optional["GPUExportArena"] = None + self._gpu_export_tickets: Dict[str, List[int]] = {} + self._gpu_export_registered = False + self._hash_cache: OrderedDict[ + Tuple[int, int, Tuple[int, ...]], Tuple[Tuple[int, int], ...] + ] = OrderedDict() + self._hash_cache_tokens = 0 + self._local_directory: Optional["LocalBlockDirectory"] = None + self._directory_retry_at = 0.0 + self._local_read: Optional["LocalReadBuffer"] = None + self._local_read_checked = False def shutdown(self): + local_read = getattr(self, "_local_read", None) + if local_read is not None: + local_read.close() + self._local_read = None + directory = getattr(self, "_local_directory", None) + if directory is not None: + directory.close() + self._local_directory = None if self._loop is None: + arena = getattr(self, "_gpu_export_arena", None) + if arena is not None: + arena.close() + self._gpu_export_arena = None return try: + if getattr(self, "_ipc_export_jobs", None): + try: + self._poll_ipc_exports(wait=True) + except Exception: + logger.warning( + "Xavier IPC export failed during shutdown", exc_info=True + ) + arena = getattr(self, "_gpu_export_arena", None) + if arena is not None: + try: + if self._transfer_ref is not None: + self._call( + self._transfer_ref.close_snapshot_export_source_v1( + arena.source + ) + ) + finally: + arena.close() + self._gpu_export_arena = None if getattr(self, "_gpu_load_jobs", None): results = self._call( asyncio.gather(*self._gpu_load_jobs, return_exceptions=True) @@ -292,6 +383,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: sent = set() + self._poll_ipc_exports() 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()) @@ -330,6 +422,54 @@ def register_kv_caches( self, kv_caches: Dict[str, torch.Tensor | Sequence[torch.Tensor]] ): self._registered_kv_caches = dict(kv_caches) + if ( + sys.platform == "linux" + and self._gpu_budget is None + and not self._direct_handoff + ): + devices = { + tensor.device + for layer, cache in kv_caches.items() + for _, tensor in self._iter_kv_tensors(layer, cache) + if tensor.is_cuda + } + if devices: + # Two GPU copies can briefly coexist during packing. Leave + # headroom for graph capture, activations and allocator segments. + free = min(torch.cuda.mem_get_info(device)[0] for device in devices) + self._ipc_export_budget = min( + _MAX_PENDING_EXPORT_BYTES, max(2 * _MAX_EXPORT_BYTES, free // 8) + ) + try: + from ...xavier.backends.torch.packed_gather import PackedGather + + caches = { + name: block_major_view(tensor, self._num_cache_blocks) + for layer, cache in kv_caches.items() + for name, tensor in self._iter_kv_tensors(layer, cache) + } + self._packed_gather = PackedGather.try_create( + caches, + {name: self._get_layer_group_id(name) for name in caches}, + ) + if self._packed_gather is not None: + self._packed_gather.warmup() + if self._packed_gather.row_bytes <= _MAX_EXPORT_BYTES: + from ...xavier.backends.torch.gpu_export import ( + GPUExportArena, + ) + + self._gpu_export_arena = GPUExportArena( + _MAX_EXPORT_BYTES, + self._ipc_export_budget // _MAX_EXPORT_BYTES, + self._packed_gather.device, + ) + except Exception: + self._packed_gather = None + logger.warning( + "Fused Xavier export unavailable; using per-layer gathers", + exc_info=True, + ) schema = { name: ( tuple(self._cache_block_view(layer, tensor).shape[1:]), @@ -386,7 +526,71 @@ def wait_for_save(self): if not pending: return try: - missing = self._call(self._filter_store_requests(pending)) + self._poll_ipc_exports() + keys = {key for request in pending for key in request.block_hashes} + pending_keys = { + key for keys, _, _ in self._ipc_export_jobs.values() for key in keys + } + if keys.issubset(pending_keys): + return + # Only this worker writes its actor's snapshots. Publication reports + # every eviction; successful keys need no per-request control RPC. + # Periodic authoritative refresh also repairs discovery after recovery. + if time.monotonic() < self._export_refresh_at and keys.issubset( + self._exported_keys + ): + return + if time.monotonic() < self._export_refresh_at and ( + keys - self._exported_keys + ).issubset(pending_keys): + return + # New content cannot require authoritative recovery of a mirrored + # hit. Avoid a readiness RPC on every cold prefill; existing keys + # still use the periodic refresh above and below. + ready = None + if self._export_refresh_at: + if time.monotonic() >= self._export_refresh_at and ( + keys & self._exported_keys + ): + # Cold traffic may reuse a small published prefix. Repair + # its authoritative publication once per refresh interval, + # rather than leaving the deadline expired and querying on + # every subsequent cold request. + self._call(self._register_blocks(pending)) + if time.monotonic() < self._export_refresh_at: + ready = (keys & self._exported_keys) | pending_keys + elif not (keys & self._exported_keys): + ready = pending_keys.copy() + if ready is not None and keys.issubset(ready): + return + missing = ( + self._missing_store_requests(pending, ready) + if ready is not None + else self._call(self._filter_store_requests(pending)) + ) + if missing and self._can_queue_ipc_export(missing): + nbytes = self._cpu_export_size(missing) + count = sum(len(request.block_hashes) for request in missing) + per_batch = max(1, _MAX_EXPORT_BYTES // (nbytes // count)) + slots = (count + per_batch - 1) // per_batch + if ( + nbytes + sum(entry[1] for entry in self._ipc_export_jobs.values()) + > self._ipc_export_budget + ) or ( + self._gpu_export_arena is not None + and slots > len(self._gpu_export_arena.available) + ): + self._poll_ipc_exports(wait=True) + missing = self._call(self._filter_store_requests(pending)) + nbytes = self._cpu_export_size(missing) + # Gather every source before returning to vLLM. Later engine + # steps may preempt requests or reuse their original GPU slots. + batches = list(self._cpu_export_batches(missing, gpu_only=True)) + self._enqueue_ipc_export(batches, nbytes) + return + if self._ipc_export_jobs: + self._poll_ipc_exports(wait=True) + missing = self._call(self._filter_store_requests(pending)) if missing: self._stage_missing_registered_layers(missing) # Re-publish all requested keys, including reused snapshots. This @@ -540,16 +744,40 @@ async def claim(): if external_token_count - num_computed_tokens < self._block_size: return 0, False - hashes = self._build_xavier_hashes(token_ids[:external_token_count]) + local_hit_blocks = num_computed_tokens // self._block_size + probe_tokens = min( + external_token_count, (local_hit_blocks + 1) * self._block_size + ) + hashes = self._build_xavier_hashes(token_ids[:probe_tokens]) if not hashes: return 0, False - local_hit_blocks = num_computed_tokens // self._block_size query_hashes = hashes[local_hit_blocks:] if not query_hashes: return 0, False + if self._local_directory is not None: + present = self._local_directory.contains(query_hashes[0][0]) + if present is False: + return 0, False + if not self._local_directory.valid: + self._local_directory.close() + self._local_directory = None + self._directory_retry_at = 0.0 remote = self._call(self._query_remote_blocks(request.request_id, query_hashes)) + # A remote prefix must contain the first uncached block. Cold requests + # need neither hashes nor a directory RPC payload for the whole prompt. + # Positive probes still use authoritative queries and snapshot leases. + if not self._build_contiguous_transfers(query_hashes, remote)[1]: + return 0, False + if probe_tokens < external_token_count: + hashes = self._build_xavier_hashes(token_ids[:external_token_count]) + query_hashes = hashes[local_hit_blocks:] + remainder = self._call( + self._query_remote_blocks(request.request_id, query_hashes[1:]) + ) + for rank, blocks in remainder.items(): + remote.setdefault(rank, set()).update(blocks) transfers, matched_blocks = self._build_contiguous_transfers( query_hashes, remote ) @@ -848,6 +1076,23 @@ async def _query_remote_blocks( query_hashes: List[Tuple[int, int]], ) -> Dict[int, set[Tuple[int, int, int]]]: tracker_ref = await self._get_tracker_ref() + if ( + sys.platform == "linux" + and self._gpu_budget is None + and self._local_directory is None + and time.monotonic() >= self._directory_retry_at + ): + from ...xavier.local_directory import LocalBlockDirectory + + self._directory_retry_at = time.monotonic() + 1.0 + self._local_directory = LocalBlockDirectory.attach( + await tracker_ref.get_snapshot_directory(0) + ) + if ( + self._local_directory is not None + and self._local_directory.contains(query_hashes[0][0]) is False + ): + return {} remote = await tracker_ref.query_blocks( 0, query_hashes, @@ -872,14 +1117,24 @@ async def _stage_layer_blocks( ) async def _filter_store_requests( - self, requests: List[XavierStoreRequest] + self, requests: List[XavierStoreRequest], ready: Optional[set[int]] = None ) -> List[XavierStoreRequest]: - transfer_ref = await self._get_transfer_ref() - ready = set( - await transfer_ref.ready_blocks_for_export_v1( - [key for request in requests for key in request.block_hashes] + if ready is None: + transfer_ref = await self._get_transfer_ref() + ready = set( + await transfer_ref.ready_blocks_for_export_v1( + [key for request in requests for key in request.block_hashes] + ) ) + ready.update( + key for keys, _, _ in self._ipc_export_jobs.values() for key in keys ) + return self._missing_store_requests(requests, ready) + + @staticmethod + def _missing_store_requests( + requests: List[XavierStoreRequest], ready: set[int] + ) -> List[XavierStoreRequest]: missing = [] for request in requests: indices = [] @@ -940,13 +1195,16 @@ async def _register_blocks(self, requests: List[XavierStoreRequest]): available_keys.update(dict.fromkeys(ready)) evicted_keys.update(dict.fromkeys(removed)) available, evicted = list(available_keys), list(evicted_keys) + requested_keys = {key for request in requests for key in request.block_hashes} + self._exported_keys.difference_update(requested_keys - set(available)) if evicted: + self._exported_keys.difference_update(evicted) await tracker_ref.unregister_blocks(0, self._rank, evicted) # Retry discovery even when ready snapshots required no payload export. # Transport addresses identify immutable content, not recyclable GPU slots. - await tracker_ref.register_blocks( - 0, [(key, key) for key in available], self._rank - ) + await tracker_ref.register_snapshot_blocks(0, available, self._rank) + self._exported_keys.update(available) + self._export_refresh_at = time.monotonic() + _EXPORT_REFRESH_SECONDS logger.debug( "Xavier V1 registered blocks: requests=%s, rank=%s, blocks=%s", len(requests), @@ -982,6 +1240,33 @@ async def _read_layer_blocks( ) return torch.cat(blocks, dim=0) + async def _read_request_payload(self, transfer, rank, reads, *, pin_memory): + local = self._local_read + if local is not None and not local.valid(): + local.close() + self._local_read = None + self._local_read_checked = False + if sys.platform == "linux" and not self._local_read_checked: + from ...xavier.backends.torch.local_read import LocalReadBuffer + + self._local_read = LocalReadBuffer.attach( + await transfer.get_local_read_metadata_v1() + ) + self._local_read_checked = True + local = self._local_read + if local is not None: + lease = await transfer.read_request_blocks_local_v1( + rank, reads, local.token + ) + if lease is not None: + try: + return local.copy( + lease, sum(read.nbytes for read in reads), pin_memory=pin_memory + ) + finally: + await transfer.release_local_read_v1(lease) + return await transfer.read_request_blocks_v1(rank, reads) + def _load_request_blocks(self, request): from ...xavier.backends.torch.request_transfer import ( LayerRead, @@ -997,6 +1282,7 @@ def _load_request_blocks(self, request): async def load(): transfer = await self._get_transfer_ref() + indices_by_destination = {} for rank in request.transfers: reads = [] for layer, tensor in caches.items(): @@ -1018,6 +1304,8 @@ async def load(): ) ) for batch in batch_reads(reads): + devices = {caches[read.layer].device for read in batch} + cuda = len(devices) == 1 and next(iter(devices)).type == "cuda" with profile_stage( "load_rpc", request_id=request.request_id, @@ -1025,7 +1313,21 @@ async def load(): nbytes=sum(read.nbytes for read in batch), blocks=sum(len(read.keys) for read in batch), ): - payload = await transfer.read_request_blocks_v1(rank, batch) + payload = await self._read_request_payload( + transfer, rank, batch, pin_memory=cuda + ) + if cuda: + device = next(iter(devices)) + with profile_stage( + "load_h2d", + device=device, + request_id=request.request_id, + rank=self._rank, + nbytes=payload.numel(), + ): + # Upload one packed payload instead of synchronizing + # a pageable copy and index upload for every layer. + payload = payload.pin_memory().to(device, non_blocking=True) for read, blocks in unpack_reads(payload, batch): cache = caches[read.layer] with profile_stage( @@ -1046,9 +1348,21 @@ async def load(): raise RuntimeError( f"Unexpected Xavier KV dtype {blocks.dtype} for cache {cache.dtype}" ) - cache[ - torch.tensor(read.destinations, device=cache.device) - ] = blocks.to(cache.device, non_blocking=True) + destination = (cache.device, tuple(read.destinations)) + if destination not in indices_by_destination: + if cache.is_cuda: + indices_by_destination[destination] = torch.tensor( + read.destinations, + dtype=torch.long, + pin_memory=True, + ).to(cache.device, non_blocking=True) + else: + indices_by_destination[destination] = torch.tensor( + read.destinations, device=cache.device + ) + cache[indices_by_destination[destination]] = blocks.to( + cache.device, non_blocking=True + ) self._call(load()) @@ -1181,7 +1495,7 @@ def _stage_missing_registered_layers( self._call(self._stage_cpu_requests(requests)) def _cpu_export_batches( - self, requests: List[XavierStoreRequest] + self, requests: List[XavierStoreRequest], *, gpu_only: bool = False ) -> Iterator[_CPUExportBatch]: caches = { name: block_major_view(tensor, self._num_cache_blocks) @@ -1189,7 +1503,11 @@ def _cpu_export_batches( for name, tensor in self._iter_kv_tensors(layer, cache) } block_bytes = sum( - math.prod(tensor.shape[1:]) * tensor.element_size() + ( + ((math.prod(tensor.shape[1:]) * tensor.element_size() + 7) // 8 * 8) + if gpu_only + else math.prod(tensor.shape[1:]) * tensor.element_size() + ) for tensor in caches.values() ) if not block_bytes: @@ -1213,21 +1531,68 @@ def _cpu_export_batches( for name, ids in layer_ids.items(): sources[name].append(ids[i]) if len(keys) == batch_size: - yield self._prepare_cpu_export(caches, keys, sources) + if gpu_only: + yield self._prepare_cpu_export( + caches, keys, sources, gpu_only=True + ) + else: + yield self._prepare_cpu_export(caches, keys, sources) keys = [] sources = {name: [] for name in caches} if keys: - yield self._prepare_cpu_export(caches, keys, sources) + if gpu_only: + yield self._prepare_cpu_export(caches, keys, sources, gpu_only=True) + else: + yield self._prepare_cpu_export(caches, keys, sources) def _prepare_cpu_export( self, caches: Dict[str, torch.Tensor], keys: List[int], sources: Dict[str, List[int]], + *, + gpu_only: bool = False, ) -> _CPUExportBatch: + if gpu_only and self._packed_gather is not None: + slot, output = None, None + arena = self._gpu_export_arena + if ( + arena is not None + and len(keys) * self._packed_gather.row_bytes <= arena.slot_bytes + ): + slot, output = arena.allocate( + (len(keys), self._packed_gather.row_bytes) + ) + with profile_stage( + "store_gather", + device=self._packed_gather.device, + blocks=len(keys), + rank=self._rank, + ): + try: + if output is None: + packed, owners = self._packed_gather(sources) + else: + packed, owners = self._packed_gather(sources, output) + assert arena is not None and slot is not None + arena.record(slot, tuple(packed.shape)) + except BaseException: + if slot is not None: + torch.cuda.current_stream( + self._packed_gather.device + ).synchronize() + assert arena is not None + arena.release([slot]) + raise + return ( + _PackedExportEntries(keys, packed, self._packed_gather.layers, slot), + owners, + [], + ) entries = [] owners = [] events = {} + indices_by_source = {} try: for name, tensor in caches.items(): with profile_stage( @@ -1237,10 +1602,25 @@ def _prepare_cpu_export( blocks=len(keys), rank=self._rank, ): - indices = torch.tensor( - sources[name], device=tensor.device, dtype=torch.long - ) + source = (tensor.device, tuple(sources[name])) + if source not in indices_by_source: + if tensor.is_cuda: + host_indices = torch.tensor( + sources[name], dtype=torch.long, pin_memory=True + ) + indices_by_source[source] = host_indices.to( + tensor.device, non_blocking=True + ) + owners.append(host_indices) + else: + indices_by_source[source] = torch.tensor( + sources[name], device=tensor.device, dtype=torch.long + ) + indices = indices_by_source[source] gathered = tensor.index_select(0, indices).detach() + if gpu_only: + entries.append((name, keys, gathered)) + continue if tensor.is_cuda: if tensor.device not in events: events[tensor.device] = torch.cuda.Event() @@ -1252,8 +1632,18 @@ def _prepare_cpu_export( # when allocation, staging or an actor RPC fails. owners.append(gathered) entries.append((name, keys, blocks)) - blocks.copy_(gathered, non_blocking=True) - event.record(torch.cuda.current_stream(tensor.device)) + stream = self._cpu_export_streams.get(tensor.device) + if stream is None: + stream = torch.cuda.Stream(device=tensor.device) + self._cpu_export_streams[tensor.device] = stream + # Source gathers precede future writes on the model + # stream. Only immutable gathered buffers cross to the + # copy stream, so subsequent forwards need not wait for + # PCIe D2H traffic. + stream.wait_stream(torch.cuda.current_stream(tensor.device)) + with torch.cuda.stream(stream): + blocks.copy_(gathered, non_blocking=True) + event.record(stream) else: entries.append((name, keys, gathered.contiguous())) except BaseException: @@ -1262,6 +1652,247 @@ def _prepare_cpu_export( raise return entries, owners, list(events.values()) + def _cpu_export_size(self, requests: List[XavierStoreRequest]) -> int: + return sum(len(request.block_hashes) for request in requests) * sum( + ( + math.prod(block_major_view(tensor, self._num_cache_blocks).shape[1:]) + * tensor.element_size() + + 7 + ) + // 8 + * 8 + for layer, cache in self._registered_kv_caches.items() + for _, tensor in self._iter_kv_tensors(layer, cache) + ) + + def _can_queue_ipc_export(self, requests: List[XavierStoreRequest]) -> bool: + return ( + sys.platform == "linux" + and self._gpu_budget is None + and bool(self._registered_kv_caches) + and all( + tensor.is_cuda + for layer, cache in self._registered_kv_caches.items() + for _, tensor in self._iter_kv_tensors(layer, cache) + ) + and self._cpu_export_size(requests) <= self._ipc_export_budget + ) + + def _prepare_ipc_export( + self, batches: List[_CPUExportBatch], nbytes: int + ) -> Tuple[str, list, List[int]]: + from torch.multiprocessing.reductions import reduce_tensor + + ticket = uuid.uuid4().hex + descriptors = [] + owners = [] + slots = [] + for entries, indices, _ in batches: + if not entries: + continue + if isinstance(entries, _PackedExportEntries): + packed, layers = entries.packed, entries.layers + else: + count = len(entries[0][1]) + pieces, layers, offset = [], [], 0 + for name, keys, value in entries: + if offset % 8: + padding = 8 - offset % 8 + pieces.append( + torch.zeros( + (count, padding), dtype=torch.uint8, device=value.device + ) + ) + offset += padding + part = value.view(torch.uint8).reshape(count, -1) + layers.append( + ( + name, + tuple(value.shape[1:]), + value.dtype, + offset, + part.shape[1], + ) + ) + pieces.append(part) + offset += part.shape[1] + if offset % 8: + pieces.append( + torch.zeros( + (count, 8 - offset % 8), + dtype=torch.uint8, + device=pieces[0].device, + ) + ) + packed = torch.cat(pieces, dim=1) + if isinstance(entries, _PackedExportEntries) and entries.slot is not None: + assert self._gpu_export_arena is not None + slots.append(entries.slot) + descriptor = { + "gpu_export": self._gpu_export_arena.source, + "slot": entries.slot, + "shape": tuple(packed.shape), + } + else: + descriptor = reduce_tensor(packed)[1] + keys = ( + entries.keys + if isinstance(entries, _PackedExportEntries) + else entries[0][1] + ) + descriptors.append((keys, descriptor, layers)) + owners.append(packed) + owners.extend(indices) + all_keys = {key for keys, _, _ in descriptors for key in keys} + # Keep ownership even if an enqueue acknowledgement is lost; the actor + # may already have started copying these IPC buffers. + self._ipc_export_jobs[ticket] = (all_keys, nbytes, owners) + if slots: + self._gpu_export_tickets[ticket] = slots + + return ticket, descriptors, slots + + async def _send_ipc_export( + self, ticket: str, descriptors: list, slots: List[int] + ) -> None: + transfer = await self._get_transfer_ref() + # Serialize registration and recovery. Each mapping owns one set of CUDA + # IPC references for the lifetime of the actor, shared by all exports. + async with self._gpu_mapping_lock: + + async def register(): + assert self._gpu_export_arena is not None + await transfer.register_snapshot_export_source_v1( + self._gpu_export_arena.source, self._gpu_export_arena.metadata() + ) + self._gpu_export_registered = True + + async def enqueue(): + return await transfer.enqueue_snapshot_export_v1( + ticket, + descriptors, + self._xavier_config.get("block_tracker_address"), + self._xavier_config.get("block_tracker_uid"), + ) + + try: + if slots and not self._gpu_export_registered: + await register() + if await enqueue() is False: + self._exported_keys.clear() + self._export_refresh_at = 0 + await register() + if await enqueue() is False: + raise RuntimeError("Xavier export source registration lost") + except BaseException: + self._gpu_export_registered = False + raise + + async def _queue_ipc_export( + self, batches: List[_CPUExportBatch], nbytes: int + ) -> None: + await self._send_ipc_export(*self._prepare_ipc_export(batches, nbytes)) + + def _enqueue_ipc_export(self, batches: List[_CPUExportBatch], nbytes: int) -> None: + # Only the reusable arena has a source-close handshake that drains late + # submissions before freeing GPU storage. Legacy per-batch CUDA IPC + # exports retain their synchronous acknowledgement and cleanup behavior. + if self._gpu_export_arena is None or not all( + isinstance(entries, _PackedExportEntries) and entries.slot is not None + for entries, _, _ in batches + ): + self._call(self._queue_ipc_export(batches, nbytes)) + return + ticket, descriptors, slots = self._prepare_ipc_export(batches, nbytes) + if self._loop is None: + self._loop = acquire_actor_loop() + # Own slots and index tensors before scheduling the RPC. Advance the + # existing EngineCore loop to submit it without waiting for an ACK. + # Subsequent polling progresses the task; the actor copies and publishes + # independently once it receives the immutable arena descriptor. + self._ipc_export_enqueues[ticket] = self._loop.create_task( + self._send_ipc_export(ticket, descriptors, slots) + ) + self._call(asyncio.sleep(0)) + + def _poll_ipc_exports(self, *, wait: bool = False) -> None: + if self._ipc_export_enqueues or self._ipc_export_poll is not None: + self._call(asyncio.sleep(0)) + submission_error = None + for ticket, submitted in list(self._ipc_export_enqueues.items()): + if wait or submitted.done(): + del self._ipc_export_enqueues[ticket] + # On a lost acknowledgement keep the GPU owners until the + # source-close handshake has drained any accepted copy. + try: + if wait and not submitted.done(): + self._call(asyncio.shield(submitted)) + submitted.result() + except Exception as exc: + submission_error = submission_error or exc + if submission_error is not None: + # Drain every submission before shutdown can invalidate the source. + # A late registration must never reopen a freed arena mapping. + raise submission_error + future = self._ipc_export_poll + if not self._ipc_export_jobs or ( + future is None + and (not wait and time.monotonic() < self._ipc_export_poll_at) + ): + return + + tickets = [ + ticket + for ticket in self._ipc_export_jobs + if ticket not in self._ipc_export_enqueues + ] + if not tickets and future is None: + return + + async def poll(): + transfer = await self._get_transfer_ref() + return await transfer.poll_snapshot_exports_v1(tickets, wait=wait) + + if future is not None: + if not wait and not future.done(): + return + try: + if wait and not future.done(): + self._call(asyncio.shield(future)) + results = future.result() + finally: + self._ipc_export_poll = None + elif not wait and self._loop is not None: + self._ipc_export_poll = self._loop.create_task(poll()) + self._call(asyncio.sleep(0)) + return + else: + results = self._call(poll()) + error = None + for ticket, (available, evicted) in results.items(): + slots = self._gpu_export_tickets.get(ticket) + if slots is not None: + assert self._gpu_export_arena is not None + self._gpu_export_arena.release(slots) + del self._gpu_export_tickets[ticket] + self._ipc_export_jobs.pop(ticket) + if available is None: + self._exported_keys.clear() + self._export_refresh_at = 0 + error = error or RuntimeError( + f"Xavier snapshot export failed: {evicted}" + ) + continue + self._exported_keys.difference_update(evicted) + self._exported_keys.update(available) + if not self._export_refresh_at: + self._export_refresh_at = time.monotonic() + _EXPORT_REFRESH_SECONDS + self._ipc_export_poll_at = time.monotonic() + 0.01 + if error is not None: + raise error + if wait and future is not None and self._ipc_export_jobs: + self._poll_ipc_exports(wait=True) + async def _stage_cpu_requests(self, requests: List[XavierStoreRequest]) -> None: if not requests or not self._registered_kv_caches: return @@ -1630,6 +2261,13 @@ def _get_local_transfer_map( return dict(group_transfers.get(from_rank, {})) def _build_xavier_hashes(self, token_ids: List[int]) -> List[Tuple[int, int]]: + cacheable = len(token_ids) <= _MAX_HASH_CACHE_TOKENS + cache_key = ( + (self._block_size, self._none_hash, tuple(token_ids)) if cacheable else None + ) + if cache_key is not None and cache_key in self._hash_cache: + self._hash_cache.move_to_end(cache_key) + return list(self._hash_cache[cache_key]) hashes: List[Tuple[int, int]] = [] prev_hash: Optional[int] = None for block_idx, start in enumerate(range(0, len(token_ids), self._block_size)): @@ -1645,6 +2283,15 @@ def _build_xavier_hashes(self, token_ids: List[int]) -> List[Tuple[int, int]]: ) hashes.append((content_hash, block_idx)) prev_hash = content_hash + if cache_key is not None: + self._hash_cache[cache_key] = tuple(hashes) + self._hash_cache_tokens += len(token_ids) + while ( + self._hash_cache_tokens > _MAX_HASH_CACHE_TOKENS + or len(self._hash_cache) > _MAX_HASH_CACHE_ENTRIES + ): + key, _ = self._hash_cache.popitem(last=False) + self._hash_cache_tokens -= len(key[2]) return hashes @staticmethod diff --git a/xinference/model/llm/xavier/backends/torch/gpu_export.py b/xinference/model/llm/xavier/backends/torch/gpu_export.py new file mode 100644 index 0000000000..33c7a7aafc --- /dev/null +++ b/xinference/model/llm/xavier/backends/torch/gpu_export.py @@ -0,0 +1,56 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Bounded reusable GPU exports with explicit producer/consumer completion.""" + +import uuid +from typing import List, Tuple + +import torch + + +class GPUExportArena: + def __init__(self, slot_bytes: int, slots: int, device: torch.device): + self.source = uuid.uuid4().hex + self.slot_bytes = slot_bytes + self.device = device + self.buffers = [ + torch.empty(slot_bytes, dtype=torch.uint8, device=device) + for _ in range(slots) + ] + self.available = list(range(slots)) + self.events = [torch.cuda.Event(interprocess=True) for _ in range(slots)] + for event in self.events: + event.record(torch.cuda.current_stream(device)) + self.handles = [event.ipc_handle() for event in self.events] + + def allocate(self, shape: Tuple[int, int]) -> Tuple[int, torch.Tensor]: + if shape[0] * shape[1] > self.slot_bytes: + raise ValueError("GPU export exceeds one arena slot") + slot = self.available.pop() + return slot, self.buffers[slot][: shape[0] * shape[1]].view(shape) + + def record(self, slot: int, shape: Tuple[int, int]) -> dict: + self.events[slot].record(torch.cuda.current_stream(self.device)) + return {"gpu_export": self.source, "slot": slot, "shape": shape} + + def metadata(self) -> list: + from torch.multiprocessing.reductions import reduce_tensor + + # Issue one IPC refcount per actor lifetime. Reusing an old descriptor + # after actor recovery would reuse a released refcount. + return [ + (reduce_tensor(buffer)[1], handle) + for buffer, handle in zip(self.buffers, self.handles) + ] + + def release(self, slots: List[int]) -> None: + for slot in slots: + self.events[slot].synchronize() + if slot not in self.available: + self.available.append(slot) + + def close(self) -> None: + for event in self.events: + event.synchronize() + self.buffers.clear() + self.available.clear() diff --git a/xinference/model/llm/xavier/backends/torch/local_read.py b/xinference/model/llm/xavier/backends/torch/local_read.py new file mode 100644 index 0000000000..4e8e35a586 --- /dev/null +++ b/xinference/model/llm/xavier/backends/torch/local_read.py @@ -0,0 +1,151 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Bounded, leased same-host receive buffers; remote actors use normal RPCs.""" + +import mmap +import os +import tempfile +import uuid +from pathlib import Path +from typing import Optional, Tuple + +import torch + +from .request_transfer import MAX_REQUEST_BYTES + +_HEADER = 64 + + +def _boot_id() -> str: + return Path("/proc/sys/kernel/random/boot_id").read_text().strip() + + +class SharedReadBuffer: + def __init__(self, capacity: int = MAX_REQUEST_BYTES): + if not 0 < capacity <= MAX_REQUEST_BYTES: + raise ValueError("Invalid shared receive capacity") + self.capacity = capacity + self.token = uuid.uuid4().bytes + self._file = tempfile.TemporaryFile( + dir="/dev/shm", prefix="xinference-xavier-read-" + ) + try: + os.ftruncate(self._file.fileno(), _HEADER + capacity) + self._data = mmap.mmap(self._file.fileno(), _HEADER + capacity) + self._data[:16] = self.token + self._data[16] = 1 + self._payload: Optional[torch.Tensor] = torch.frombuffer( + self._data, dtype=torch.uint8, offset=_HEADER + ) + except BaseException: + if hasattr(self, "_data"): + self._data.close() + self._file.close() + raise + self._lease: Optional[bytes] = None + + def metadata(self) -> dict: + return { + "boot_id": _boot_id(), + "path": f"/proc/{os.getpid()}/fd/{self._file.fileno()}", + "token": self.token, + "capacity": self.capacity, + } + + def acquire(self, nbytes: int) -> Optional[Tuple[bytes, torch.Tensor]]: + if self._lease is not None or self._payload is None: + return None + if not 0 < nbytes <= self.capacity: + return None + self._lease = uuid.uuid4().bytes + self._data[24:40] = self._lease + return self._lease, self._payload[:nbytes] + + def release(self, lease: bytes) -> None: + if lease == self._lease and not self._file.closed: + self._lease = None + self._data[24:40] = bytes(16) + + def close(self) -> None: + if not self._file.closed: + self._data[16] = 0 + self._payload = None + self._data.close() + self._file.close() + + +class LocalReadBuffer: + @classmethod + def attach(cls, metadata: Optional[dict]) -> Optional["LocalReadBuffer"]: + if metadata is None: + return None + try: + if metadata["boot_id"] != _boot_id(): + return None + return cls(metadata) + except (OSError, ValueError): + return None + + def __init__(self, metadata: dict): + capacity = metadata["capacity"] + if not 0 < capacity <= MAX_REQUEST_BYTES: + raise ValueError("Invalid shared receive capacity") + self.capacity = capacity + self.token = metadata["token"] + self._path = metadata["path"] + self._fd = os.open(self._path, os.O_RDONLY | os.O_CLOEXEC) + try: + stat = os.fstat(self._fd) + if stat.st_size != _HEADER + capacity: + raise ValueError("Shared receive buffer size changed") + self._identity = (stat.st_dev, stat.st_ino) + # Private writable views let torch read without a read-only-buffer + # warning. Consumer writes cannot modify the actor's shared payload. + self._data = mmap.mmap( + self._fd, _HEADER + capacity, access=mmap.ACCESS_COPY + ) + if self._data[:16] != self.token or not self._data[16]: + raise ValueError("Shared receive owner changed") + self._payload: Optional[torch.Tensor] = torch.frombuffer( + self._data, dtype=torch.uint8, offset=_HEADER + ) + except BaseException: + if hasattr(self, "_data"): + self._data.close() + os.close(self._fd) + raise + + def valid(self) -> bool: + if self._fd < 0: + return False + try: + stat = os.stat(self._path) + except OSError: + return False + return ( + (stat.st_dev, stat.st_ino) == self._identity + and self._data[:16] == self.token + and bool(self._data[16]) + ) + + def copy(self, lease: bytes, nbytes: int, *, pin_memory: bool) -> torch.Tensor: + if ( + not self.valid() + or self._data[24:40] != lease + or not 0 < nbytes <= self.capacity + ): + raise RuntimeError("Shared receive lease or owner changed") + assert self._payload is not None + payload = self._payload[:nbytes] + # The actor can reuse its receive slot immediately after the caller's + # release RPC. H2D must own an independent pinned allocation first. + if pin_memory: + return payload.pin_memory() + return torch.from_numpy(payload.numpy().copy()) + + def close(self) -> None: + if self._fd >= 0: + self._payload = None + self._data.close() + os.close(self._fd) + self._fd = -1 diff --git a/xinference/model/llm/xavier/backends/torch/packed_gather.py b/xinference/model/llm/xavier/backends/torch/packed_gather.py new file mode 100644 index 0000000000..fd13a6d674 --- /dev/null +++ b/xinference/model/llm/xavier/backends/torch/packed_gather.py @@ -0,0 +1,156 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Gather and pack aligned CUDA KV blocks without per-layer intermediates.""" + +import math +from collections import defaultdict +from typing import Dict, List, Optional, Tuple + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _gather_words( + meta, + indices, + output, + count, + ROW_WORDS: tl.constexpr, + SIZE: tl.constexpr, + SHAPE: tl.constexpr, + STRIDES: tl.constexpr, + BLOCK_STRIDE: tl.constexpr, + BLOCK: tl.constexpr, +): + layer, row = tl.program_id(2), tl.program_id(1) + position = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + address = tl.load(meta + layer * 3).to(tl.pointer_type(tl.uint64)) + offset = tl.load(meta + layer * 3 + 1) + group = tl.load(meta + layer * 3 + 2) + source_block = tl.load(indices + group * count + row) + source_position = source_block * BLOCK_STRIDE + remaining = position + for dimension in tl.static_range(len(SHAPE) - 1, -1, -1): + source_position += (remaining % SHAPE[dimension]) * STRIDES[dimension] + remaining = remaining // SHAPE[dimension] + value = tl.load(address + source_position, mask=position < SIZE, other=0) + target = output.to(tl.pointer_type(tl.uint64)) + tl.store(target + row * ROW_WORDS + offset + position, value, mask=position < SIZE) + + +class PackedGather: + @classmethod + def try_create( + cls, caches: Dict[str, torch.Tensor], source_groups: Dict[str, int] + ) -> Optional["PackedGather"]: + if not caches or len({tensor.device for tensor in caches.values()}) != 1: + return None + for tensor in caches.values(): + size = tensor.element_size() + if ( + not tensor.is_cuda + or tensor.ndim < 2 + or tensor.stride(-1) != 1 + or tensor.shape[-1] * size % 8 + or tensor.data_ptr() % 8 + or any(stride * size % 8 for stride in tensor.stride()[:-1]) + ): + return None + return cls(caches, source_groups) + + def __init__(self, caches: Dict[str, torch.Tensor], source_groups: Dict[str, int]): + self.caches = caches + self.layers: List[Tuple[str, Tuple[int, ...], torch.dtype, int, int]] = [] + self.source_names = list(dict.fromkeys(source_groups.values())) + group_indices = {group: i for i, group in enumerate(self.source_names)} + self.representatives = { + group: next(name for name in caches if source_groups[name] == group) + for group in self.source_names + } + self.source_groups = source_groups + self.source_limits = { + group: min( + tensor.shape[0] + for name, tensor in caches.items() + if source_groups[name] == group + ) + for group in self.source_names + } + geometries = defaultdict(list) + offset = 0 + for name, tensor in caches.items(): + size = tensor.element_size() + shape = (*tensor.shape[1:-1], tensor.shape[-1] * size // 8) + strides = tuple(step * size // 8 for step in tensor.stride()[1:-1]) + (1,) + block_stride = tensor.stride(0) * size // 8 + nbytes = math.prod(tensor.shape[1:]) * size + geometries[(shape, strides, block_stride)].append( + [tensor.data_ptr(), offset // 8, group_indices[source_groups[name]]] + ) + self.layers.append( + (name, tuple(tensor.shape[1:]), tensor.dtype, offset, nbytes) + ) + offset += nbytes + self.row_bytes = offset + self.device = next(iter(caches.values())).device + self.metadata = [] + for geometry, rows in geometries.items(): + host = torch.tensor(rows, dtype=torch.int64, pin_memory=True) + self.metadata.append( + (host, host.to(self.device, non_blocking=True), geometry) + ) + + def __call__( + self, sources: Dict[str, List[int]], output: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, List[torch.Tensor]]: + count = len(next(iter(sources.values()))) + for name in self.caches: + representative = self.representatives[self.source_groups[name]] + if sources[name] != sources[representative]: + raise ValueError("Inconsistent source indices in one KV group") + for group, name in self.representatives.items(): + if len(sources[name]) != count: + raise ValueError("Inconsistent KV group batch lengths") + limit = self.source_limits[group] + if any(index < 0 or index >= limit for index in sources[name]): + raise IndexError("Invalid Xavier source block index") + host = torch.tensor( + [sources[self.representatives[group]] for group in self.source_names], + dtype=torch.int64, + pin_memory=True, + ) + indices = host.to(self.device, non_blocking=True) + if output is None: + output = torch.empty( + (count, self.row_bytes), dtype=torch.uint8, device=self.device + ) + elif ( + output.shape != (count, self.row_bytes) + or output.dtype != torch.uint8 + or output.device != self.device + or not output.is_contiguous() + ): + raise ValueError("Invalid packed export destination") + for _, meta, (shape, strides, block_stride) in self.metadata: + size = math.prod(shape) + _gather_words[(triton.cdiv(size, 1024), count, meta.shape[0])]( + meta, + indices, + output, + count, + self.row_bytes // 8, + size, + shape, + strides, + block_stride, + 1024, + num_warps=4, + ) + return output, [host, indices] + + def warmup(self) -> None: + # Compile before CUDA graph capture and request processing. Count is a + # runtime argument, so subsequent batch sizes reuse the same kernel. + self({name: [0] for name in self.caches}) + torch.cuda.current_stream(self.device).synchronize() diff --git a/xinference/model/llm/xavier/backends/torch/snapshot.py b/xinference/model/llm/xavier/backends/torch/snapshot.py index d04f09a6a2..82a2e4f6c7 100644 --- a/xinference/model/llm/xavier/backends/torch/snapshot.py +++ b/xinference/model/llm/xavier/backends/torch/snapshot.py @@ -1,25 +1,136 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. from collections import OrderedDict -from typing import Dict, List, Set +from collections.abc import MutableMapping +from typing import Dict, Iterator, List, Optional, Set, Tuple +import numpy as np import torch +class _PackedKVBlock(MutableMapping[str, torch.Tensor]): + """Index an owned slab; create tensor views only when a reader needs them.""" + + __slots__ = ( + "_packed", + "_schema", + "_row", + "_layers", + "_overrides", + "_removed", + "_numpy", + ) + + def __init__( + self, + packed: torch.Tensor, + schema: Dict[str, Tuple[Tuple[int, ...], torch.dtype, int, int]], + row: Optional[int] = None, + ): + self._packed, self._schema, self._row = packed, schema, row + self._layers: Dict[str, Optional[torch.Tensor]] = {} + self._overrides: Optional[Set[str]] = None + self._removed: Optional[Set[str]] = None + self._numpy: Optional[np.ndarray] = None + + def packed_layer( + self, name: str + ) -> Optional[Tuple[Tuple[int, ...], torch.dtype, np.ndarray]]: + if name not in self or ( + self._overrides is not None and name in self._overrides + ): + return None + if self._numpy is None: + value = self._packed.numpy() + self._numpy = value if self._row is None else value[self._row] + shape, dtype, offset, nbytes = self._schema[name] + return shape, dtype, self._numpy[offset : offset + nbytes] + + def __getitem__(self, name: str) -> torch.Tensor: + if name not in self: + raise KeyError(name) + value = self._layers.get(name) + if value is None: + shape, dtype, offset, nbytes = self._schema[name] + carrier = torch.float16 if dtype == torch.bfloat16 else dtype + row = self._packed if self._row is None else self._packed[self._row] + value = row[offset : offset + nbytes].view(carrier).view(shape) + self._layers[name] = value + return value + + def __setitem__(self, name: str, value: torch.Tensor) -> None: + self._layers[name] = value + if self._overrides is None: + self._overrides = set() + self._overrides.add(name) + if self._removed is not None: + self._removed.discard(name) + + def __delitem__(self, name: str) -> None: + if name not in self: + raise KeyError(name) + self._layers.pop(name, None) + if self._removed is None: + self._removed = set() + self._removed.add(name) + + def __contains__(self, name: object) -> bool: + return (name in self._schema or name in self._layers) and ( + self._removed is None or name not in self._removed + ) + + def __iter__(self) -> Iterator[str]: + for name in self._schema: + if self._removed is None or name not in self._removed: + yield name + for name in self._layers: + if name not in self._schema: + yield name + + def __len__(self) -> int: + return sum(1 for _ in self) + + class KVSnapshotStore: """Bounded, content-addressed CPU blocks with atomic publication and read leases. - All mutation runs on the TransferActor's event loop. Tensor storage is owned - by each block, never by a GPU slot or a view retaining an entire request. + All mutation runs on the TransferActor's event loop. CPU payloads never + alias recyclable GPU slots. Shared packed allocations are charged in full + and evicted together, so sparse views cannot retain unbounded batches. """ def __init__(self, capacity: int): self.capacity = capacity - self.blocks: OrderedDict[int, Dict[str, torch.Tensor]] = OrderedDict() + self.blocks: OrderedDict[int, MutableMapping[str, torch.Tensor]] = OrderedDict() self.logical_dtypes: Dict[int, Dict[str, torch.dtype]] = {} self.ready: Set[int] = set() self.leases: Dict[str, Set[int]] = {} self.evicted: Set[int] = set() + self._packed_slabs: Dict[int, Set[int]] = {} + self._packed_slab_slots: Dict[int, int] = {} + self._block_slabs: Dict[int, int] = {} + self._reserved_slots = 0 + + def _used_slots(self) -> int: + return len(self.blocks) - len(self._block_slabs) + self._reserved_slots + + def _evict_one(self, pinned: Set[int]) -> bool: + for key in self.blocks: + slab = self._block_slabs.get(key) + victims = self._packed_slabs[slab] if slab is not None else {key} + if victims & pinned: + continue + for victim in victims: + del self.blocks[victim] + del self.logical_dtypes[victim] + self.ready.discard(victim) + self.evicted.add(victim) + self._block_slabs.pop(victim, None) + if slab is not None: + del self._packed_slabs[slab] + self._reserved_slots -= self._packed_slab_slots.pop(slab) + return True + return False def stage( self, @@ -31,14 +142,8 @@ def stage( pinned = set().union(*self.leases.values()) if self.leases else set() for key, tensor in zip(keys, tensors): if key not in self.blocks: - if len(self.blocks) >= self.capacity: - victim = next((k for k in self.blocks if k not in pinned), None) - if victim is None: - continue # No room: leave this block to local recomputation. - del self.blocks[victim] - del self.logical_dtypes[victim] - self.ready.discard(victim) - self.evicted.add(victim) + if self._used_slots() >= self.capacity and not self._evict_one(pinned): + continue # No room: leave this block to local recomputation. self.blocks[key] = {} self.logical_dtypes[key] = {} self.blocks.move_to_end(key) @@ -47,6 +152,65 @@ def stage( self.blocks[key][layer] = tensor.clone().contiguous() self.logical_dtypes[key][layer] = logical_dtype or tensors.dtype + def stage_packed_blocks( + self, + keys: List[int], + packed: torch.Tensor, + layers: List[Tuple[str, Tuple[int, ...], torch.dtype, int, int]], + *, + retain_storage: bool = False, + ) -> None: + """Own complete blocks, with bounded whole-slab ownership if requested.""" + if packed.dtype != torch.uint8 or packed.ndim != 2 or len(packed) != len(keys): + raise ValueError("Invalid packed snapshot payload") + pinned = set().union(*self.leases.values()) if self.leases else set() + schema = { + name: (shape, dtype, offset, nbytes) + for name, shape, dtype, offset, nbytes in layers + } + dtypes = {name: dtype for name, _, dtype, _, _ in layers} + if retain_storage and keys and all(key not in self.blocks for key in keys): + storage = packed.untyped_storage() + slab = storage.data_ptr() + row_bytes = packed.shape[1] + slots = (storage.nbytes() + row_bytes - 1) // row_bytes + if slots <= self.capacity and slab not in self._packed_slabs: + while self._used_slots() + slots > self.capacity: + if not self._evict_one(pinned): + break + else: + self._packed_slabs[slab] = set(keys) + self._packed_slab_slots[slab] = slots + self._reserved_slots += slots + for index, key in enumerate(keys): + self.blocks[key] = _PackedKVBlock(packed, schema, index) + self.logical_dtypes[key] = dtypes.copy() + self._block_slabs[key] = slab + return + rows = packed.numpy() + for index, key in enumerate(keys): + fresh = key not in self.blocks + if key not in self.blocks: + if self._used_slots() >= self.capacity and not self._evict_one(pinned): + continue + self.blocks[key] = {} + self.logical_dtypes[key] = {} + self.blocks.move_to_end(key) + # NumPy memcpy avoids a CPU thread-pool launch for every block. + # This allocation never retains any other block in the batch. + block = torch.from_numpy(rows[index].copy()) + if fresh: + self.blocks[key] = _PackedKVBlock(block, schema) + self.logical_dtypes[key] = dtypes.copy() + continue + for name, shape, dtype, offset, nbytes in layers: + if name not in self.blocks[key]: + carrier = torch.float16 if dtype == torch.bfloat16 else dtype + self.blocks[key][name] = ( + block[offset : offset + nbytes].view(carrier).view(shape) + ) + self.logical_dtypes[key][name] = dtype + def publish(self, keys: List[int], layers: Set[str]) -> List[int]: available = [] for key in keys: @@ -72,6 +236,23 @@ def release_consumer(self, rank: int): self.release(lease) def read(self, layer: str, keys: List[int]) -> torch.Tensor: + blocks = [self.blocks[key] for key in keys] + rows = [] + for block in blocks: + if not isinstance(block, _PackedKVBlock): + break + row = block.packed_layer(layer) + if row is None: + break + rows.append(row) + if rows and len(rows) == len(blocks): + shape, dtype, _ = rows[0] + if all(row[:2] == (shape, dtype) for row in rows): + # Copy a read in NumPy without materializing and stacking a + # tensor view for every layer of every block. + carrier = torch.float16 if dtype == torch.bfloat16 else dtype + packed = torch.from_numpy(np.stack([row[2] for row in rows])) + return packed.view(carrier).view(len(keys), *shape) return torch.stack([self.blocks[key][layer] for key in keys]).contiguous() diff --git a/xinference/model/llm/xavier/backends/torch/tiered_snapshot.py b/xinference/model/llm/xavier/backends/torch/tiered_snapshot.py index 457fd2a77f..42ea875cd2 100644 --- a/xinference/model/llm/xavier/backends/torch/tiered_snapshot.py +++ b/xinference/model/llm/xavier/backends/torch/tiered_snapshot.py @@ -3,7 +3,7 @@ """Budgeted hot snapshots with CPU overflow and shared publication/read leases.""" from collections import OrderedDict -from typing import Dict, List, Tuple +from typing import Dict, List, Mapping, Tuple import torch @@ -71,7 +71,7 @@ def _cpu_room(self, pinned: set) -> bool: return True @staticmethod - def _copy_to_cpu(layers: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + def _copy_to_cpu(layers: Mapping[str, torch.Tensor]) -> Dict[str, torch.Tensor]: # A synchronous D2H copy per layer stalls the producer for every layer # of every evicted block. Pack compatible layers and synchronize once # per group. Each host allocation belongs to one block, so evicting a diff --git a/xinference/model/llm/xavier/block_tracker.py b/xinference/model/llm/xavier/block_tracker.py index 582d025daa..5675b66a44 100644 --- a/xinference/model/llm/xavier/block_tracker.py +++ b/xinference/model/llm/xavier/block_tracker.py @@ -12,10 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. import random -from typing import Dict, List, Optional, Set, Tuple +import sys +from contextlib import nullcontext +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple import xoscar as xo +if TYPE_CHECKING: + from .local_directory import SharedBlockDirectory + class BlockTracker(xo.StatelessActor): @classmethod @@ -28,11 +33,90 @@ def __init__(self): self._hash_to_rank_and_block_id: Dict[int, Dict[int, Set[Tuple[int, int]]]] = {} # type: ignore # engine -> rank -> (hash, block_id) self._rank_to_hash_and_block_id: Dict[int, Dict[int, Set[Tuple[int, int]]]] = {} # type: ignore + # engine -> rank -> block id -> content hashes. Updates must scale with + # the incoming batch rather than scanning a rank's entire cache. + self._rank_to_block_hashes: Dict[int, Dict[int, Dict[int, Set[int]]]] = {} self._unavailable_ranks: Set[int] = set() # type: ignore + # Immutable V1 snapshots use their content hash as their address. Keep + # only integer keys/values on that path: per-block sets and location + # tuples otherwise trigger collections of the serving process's heap. + self._snapshot_ranks: Dict[int, Dict[int, Dict[int, None]]] = {} + self._rank_snapshots: Dict[int, Dict[int, Dict[int, None]]] = {} + self._snapshot_directories: Dict[int, "SharedBlockDirectory"] = {} + self._snapshot_directory_disabled: Set[int] = set() + + def get_snapshot_directory(self, virtual_engine: int): + if ( + sys.platform != "linux" + or virtual_engine in self._snapshot_directory_disabled + ): + return None + directory = self._snapshot_directories.get(virtual_engine) + if directory is None: + from .local_directory import SharedBlockDirectory + + try: + directory = SharedBlockDirectory() + with directory.update(): + for key in self._snapshot_ranks.get(virtual_engine, {}): + directory.add(key) + self._snapshot_directories[virtual_engine] = directory + except OSError: + return None + try: + return directory.metadata() + except OSError: + return None + + async def __pre_destroy__(self): + for directory in self._snapshot_directories.values(): + directory.close() + self._snapshot_directories.clear() + + def register_snapshot_blocks(self, virtual_engine: int, keys: List[int], rank: int): + if self._rank_to_hash_and_block_id.get(virtual_engine, {}).get(rank): + self.unregister_blocks(virtual_engine, rank, keys) + locations = self._snapshot_ranks.setdefault(virtual_engine, {}) + owned = self._rank_snapshots.setdefault(virtual_engine, {}).setdefault(rank, {}) + directory = self._snapshot_directories.get(virtual_engine) + with directory.update() if directory is not None else nullcontext(): + for key in keys: + if key not in locations: + if directory is not None: + directory.add(key) + locations[key] = {} + locations[key][rank] = None + owned[key] = None + + def _unregister_snapshots(self, virtual_engine: int, rank: int, keys): + owned = self._rank_snapshots.get(virtual_engine, {}).get(rank) + if not owned: + return + locations = self._snapshot_ranks[virtual_engine] + directory = self._snapshot_directories.get(virtual_engine) + with directory.update() if directory is not None else nullcontext(): + for key in keys: + if key in owned: + del owned[key] + ranks = locations[key] + del ranks[rank] + if not ranks: + del locations[key] + if directory is not None: + directory.remove(key) def register_blocks( self, virtual_engine: int, block_infos: List[Tuple[int, int]], rank: int ): + # Physical-slot registrations are a separate protocol. Disable the + # negative shortcut before admitting keys outside the snapshot index. + directory = self._snapshot_directories.pop(virtual_engine, None) + if directory is not None: + directory.close() + self._snapshot_directory_disabled.add(virtual_engine) + self._unregister_snapshots( + virtual_engine, rank, (block_id for _, block_id in block_infos) + ) if virtual_engine not in self._hash_to_rank_and_block_id: self._hash_to_rank_and_block_id[virtual_engine] = {} hash_to_rank_and_block_id = self._hash_to_rank_and_block_id[virtual_engine] @@ -51,21 +135,21 @@ def register_blocks( # keeps pointing at a block that now holds different KV, and a later # query on that hash would load the wrong KV. This also bounds the # tracker to one entry per (rank, block_id). - new_block_ids = {block_id for _, block_id in block_infos} - new_pairs = set(block_infos) - stale = [ - (stale_hash, stale_block_id) - for stale_hash, stale_block_id in hash_and_block_id - if stale_block_id in new_block_ids - and (stale_hash, stale_block_id) not in new_pairs - ] - for stale_hash, stale_block_id in stale: - hash_and_block_id.discard((stale_hash, stale_block_id)) - stale_entries = hash_to_rank_and_block_id.get(stale_hash) - if stale_entries is not None: - stale_entries.discard((rank, stale_block_id)) - if not stale_entries: - hash_to_rank_and_block_id.pop(stale_hash, None) + by_block = self._rank_to_block_hashes.setdefault(virtual_engine, {}).setdefault( + rank, {} + ) + incoming: Dict[int, Set[int]] = {} + for content_hash, block_id in block_infos: + incoming.setdefault(block_id, set()).add(content_hash) + for block_id, new_hashes in incoming.items(): + for stale_hash in by_block.get(block_id, set()) - new_hashes: + hash_and_block_id.discard((stale_hash, block_id)) + stale_entries = hash_to_rank_and_block_id.get(stale_hash) + if stale_entries is not None: + stale_entries.discard((rank, block_id)) + if not stale_entries: + hash_to_rank_and_block_id.pop(stale_hash, None) + by_block[block_id] = new_hashes for hash_content, block_id in block_infos: hash_to_rank_and_block_id.setdefault(hash_content, set()).add( @@ -79,11 +163,21 @@ def query_blocks( hash_contents: List[Tuple[int, int]], exclude_rank: Optional[int] = None, ) -> Dict[int, Set[Tuple[int, int, int]]]: - if virtual_engine not in self._hash_to_rank_and_block_id: - return {} - hash_to_rank_and_block_id = self._hash_to_rank_and_block_id[virtual_engine] + hash_to_rank_and_block_id = self._hash_to_rank_and_block_id.get( + virtual_engine, {} + ) + snapshots = self._snapshot_ranks.get(virtual_engine, {}) remote: Dict[int, Set[Tuple[int, int, int]]] = {} for hash_content, _id in hash_contents: + snapshot_ranks = [ + rank + for rank in snapshots.get(hash_content, ()) + if rank not in self._unavailable_ranks and rank != exclude_rank + ] + if snapshot_ranks: + rank = random.choice(snapshot_ranks) + remote.setdefault(rank, set()).add((hash_content, hash_content, _id)) + continue if ( hash_content in hash_to_rank_and_block_id ) and hash_to_rank_and_block_id[hash_content]: @@ -109,15 +203,16 @@ def unregister_block(self, virtual_engine: int, rank: int, block_id: int): self.unregister_blocks(virtual_engine, rank, [block_id]) def unregister_blocks(self, virtual_engine: int, rank: int, block_ids: List[int]): + self._unregister_snapshots(virtual_engine, rank, block_ids) entries = self._rank_to_hash_and_block_id.get(virtual_engine, {}).get(rank) hashes = self._hash_to_rank_and_block_id.get(virtual_engine, {}) if not entries: return - removed_ids = set(block_ids) + by_block = self._rank_to_block_hashes.get(virtual_engine, {}).get(rank, {}) removed = { (content_hash, block_id) - for content_hash, block_id in entries - if block_id in removed_ids + for block_id in set(block_ids) + for content_hash in by_block.pop(block_id, ()) } entries.difference_update(removed) for content_hash, block_id in removed: @@ -137,8 +232,13 @@ def register_rank(self, rank: int): """ After recovery is successful, clear all stale data of the rank and mark the rank as available. """ + for engine, owned in self._rank_snapshots.items(): + self._unregister_snapshots(engine, rank, list(owned.get(rank, ()))) + owned.pop(rank, None) for _, rank_to_hash_and_block_id in self._rank_to_hash_and_block_id.items(): rank_to_hash_and_block_id.pop(rank, None) + for rank_to_block_hashes in self._rank_to_block_hashes.values(): + rank_to_block_hashes.pop(rank, None) for _, hash_to_rank_and_block_id in self._hash_to_rank_and_block_id.items(): for _, rank_and_block_id in hash_to_rank_and_block_id.items(): diff --git a/xinference/model/llm/xavier/local_directory.py b/xinference/model/llm/xavier/local_directory.py new file mode 100644 index 0000000000..63e493909d --- /dev/null +++ b/xinference/model/llm/xavier/local_directory.py @@ -0,0 +1,139 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""A bounded, same-host negative directory for immutable snapshot addresses.""" + +import fcntl +import mmap +import os +import tempfile +import uuid +from array import array +from contextlib import contextmanager +from pathlib import Path +from typing import Iterator, Optional, Tuple + +_BYTES = 64 * 1024 +_HEADER = 32 + + +def _indices(key: int) -> Tuple[int, int]: + return key % (_BYTES * 8), ((key >> 32) ^ (key * 0x9E3779B1)) % (_BYTES * 8) + + +def _boot_id() -> str: + return Path("/proc/sys/kernel/random/boot_id").read_text().strip() + + +class SharedBlockDirectory: + def __init__(self): + # Some bundled Linux Python builds omit os.memfd_create. TemporaryFile + # on tmpfs is anonymous (or unlinked immediately) and works there too. + self._file = tempfile.TemporaryFile( + dir="/dev/shm", prefix="xinference-xavier-directory-" + ) + self._fd = self._file.fileno() + try: + os.ftruncate(self._fd, _HEADER + _BYTES) + self._data = mmap.mmap(self._fd, _HEADER + _BYTES) + except BaseException: + self._file.close() + raise + self._token = uuid.uuid4().bytes + self._data[:16] = self._token + self._data[16] = 1 + self._counts = array("I", [0]) * (_BYTES * 8) + + def metadata(self) -> dict: + return { + "boot_id": _boot_id(), + "path": f"/proc/{os.getpid()}/fd/{self._fd}", + "token": self._token, + } + + @contextmanager + def update(self) -> Iterator[None]: + fcntl.flock(self._fd, fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(self._fd, fcntl.LOCK_UN) + + def add(self, key: int) -> None: + for index in _indices(key): + self._counts[index] += 1 + if self._counts[index] == 1: + offset = _HEADER + (index >> 3) + self._data[offset] |= 1 << (index & 7) + + def remove(self, key: int) -> None: + for index in _indices(key): + self._counts[index] -= 1 + if not self._counts[index]: + offset = _HEADER + (index >> 3) + self._data[offset] &= ~(1 << (index & 7)) + + def close(self) -> None: + if self._fd >= 0: + with self.update(): + self._data[16] = 0 + self._data.close() + self._file.close() + self._fd = -1 + + +class LocalBlockDirectory: + @classmethod + def attach(cls, metadata: Optional[dict]) -> Optional["LocalBlockDirectory"]: + if metadata is None: + return None + try: + if metadata["boot_id"] != _boot_id(): + return None + return cls(metadata) + except (OSError, ValueError): + return None + + def __init__(self, metadata: dict): + self._path = metadata["path"] + self.valid = True + self._fd = os.open(self._path, os.O_RDONLY | os.O_CLOEXEC) + try: + self._data = mmap.mmap(self._fd, _HEADER + _BYTES, access=mmap.ACCESS_READ) + if self._data[:16] != metadata["token"]: + raise ValueError("Directory owner changed") + stat = os.fstat(self._fd) + self._identity = (stat.st_dev, stat.st_ino) + except BaseException: + if hasattr(self, "_data"): + self._data.close() + os.close(self._fd) + raise + + def contains(self, key: int) -> Optional[bool]: + # Independent read-only opens have independent flock ownership. Never + # block EngineCore behind publication; a busy writer uses the RPC path. + try: + fcntl.flock(self._fd, fcntl.LOCK_SH | fcntl.LOCK_NB) + except BlockingIOError: + return None + try: + try: + stat = os.stat(self._path) + except OSError: + self.valid = False + return None + if (stat.st_dev, stat.st_ino) != self._identity or not self._data[16]: + self.valid = False + return None + return all( + self._data[_HEADER + (index >> 3)] & (1 << (index & 7)) + for index in _indices(key) + ) + finally: + fcntl.flock(self._fd, fcntl.LOCK_UN) + + def close(self) -> None: + if self._fd >= 0: + self._data.close() + os.close(self._fd) + self._fd = -1 diff --git a/xinference/model/llm/xavier/tests/test_token_hash.py b/xinference/model/llm/xavier/tests/test_token_hash.py new file mode 100644 index 0000000000..a3866f8414 --- /dev/null +++ b/xinference/model/llm/xavier/tests/test_token_hash.py @@ -0,0 +1,20 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +import pytest + +from ..utils import hash_block_tokens + + +@pytest.mark.parametrize( + "first,previous,tokens,expected", + [ + (True, None, [], 16765256144593249530), + (True, -1, [1, 2, 3], 3926270711534698589), + (False, 2**64 - 1, [0, 2147483647, -2147483648], 682864629166071906), + (False, 123, [1] * 16, 9100343863099402539), + ], +) +def test_token_hash_preserves_v1_wire_format(first, previous, tokens, expected): + # Golden hashes produced before vectorizing token encoding. Existing + # snapshots and ranks running older releases must keep matching. + assert hash_block_tokens(first, previous, tokens) == expected diff --git a/xinference/model/llm/xavier/utils.py b/xinference/model/llm/xavier/utils.py index adc963164a..2af8224550 100644 --- a/xinference/model/llm/xavier/utils.py +++ b/xinference/model/llm/xavier/utils.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import struct -from typing import List, Optional +from functools import lru_cache +from typing import List, Optional, Tuple import xxhash @@ -35,6 +36,25 @@ def hash_block_tokens( if extra_hash is None: extra_hash = none_hash + return _hash_block_tokens( + is_first_block, + int(prev_block_hash) & 0xFFFFFFFFFFFFFFFF, + tuple(cur_block_token_ids), + int(extra_hash) & 0xFFFFFFFFFFFFFFFF, + ) + + +@lru_cache(maxsize=16384) +def _hash_block_tokens( + is_first_block: bool, + prev_block_hash: int, + cur_block_token_ids: Tuple[int, ...], + extra_hash: int, +) -> int: + # Cache whole hash inputs, including the predecessor. Identical token + # chunks at different prefix positions must not share a content address. + # The bounded cache retains no request or tensor storage. + buf = bytearray() # 0. hash version @@ -50,8 +70,9 @@ def hash_block_tokens( buf += struct.pack(" Date: Wed, 7 Oct 2026 19:55:18 +0800 Subject: [PATCH 4/6] perf: overlap Xavier prefill snapshots and reduce pinned history --- .../xavier/test/test_snapshot_export_ipc.py | 260 +++++++++++++++++- xinference/model/llm/vllm/xavier/transfer.py | 82 +++--- .../model/llm/vllm/xavier/v1_connector.py | 111 +++++++- 3 files changed, 407 insertions(+), 46 deletions(-) diff --git a/xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py b/xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py index e7ae02f2e0..77c79ffbf6 100644 --- a/xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py +++ b/xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py @@ -2,6 +2,8 @@ # Licensed under the Apache License, Version 2.0. import asyncio import sys +import threading +from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -12,7 +14,7 @@ from ..block_tracker import VLLMBlockTracker from ..snapshot import KVSnapshotStore -from ..transfer import TransferActor +from ..transfer import TransferActor, _copy_snapshot_to_cpu @pytest.mark.parametrize("wait", [False, True]) @@ -26,6 +28,252 @@ def test_exportless_completion_poll_needs_no_cpu_export_state(connector_module, connector._poll_ipc_exports(wait=wait) +@pytest.mark.asyncio +async def test_snapshot_copy_owns_independent_pageable_bytes(): + value = torch.tensor([1e30, -1e30]).to(torch.bfloat16).view(torch.uint8) + expected = value.clone() + snapshot = await _copy_snapshot_to_cpu(value, None, None) + value.zero_() + assert torch.equal(snapshot, expected) + assert snapshot.data_ptr() != value.data_ptr() + assert not snapshot.is_pinned() + + +@pytest.mark.asyncio +async def test_snapshot_copy_retains_source_after_repeated_cancellation(monkeypatch): + started, release = threading.Event(), threading.Event() + tensor_to = torch.Tensor.to + + def blocked_copy(value, *args, **kwargs): + started.set() + assert release.wait(5) + return tensor_to(value, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "to", blocked_copy) + task = asyncio.create_task(_copy_snapshot_to_cpu(torch.arange(8), None, None)) + try: + assert await asyncio.to_thread(started.wait, 5) + for _ in range(2): + task.cancel() + await asyncio.sleep(0) + assert not task.done() + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +async def test_failed_snapshot_copy_drains_its_stream(monkeypatch): + stream, ready = Mock(), Mock() + value = SimpleNamespace(to=Mock(side_effect=RuntimeError("copy failed"))) + monkeypatch.setattr(torch.cuda, "stream", lambda _: nullcontext()) + with pytest.raises(RuntimeError, match="copy failed"): + await _copy_snapshot_to_cpu(value, stream, ready) + stream.wait_event.assert_called_once_with(ready) + stream.synchronize.assert_called_once_with() + + +def test_next_forward_fences_exports_even_without_loads( + connector, connector_module, monkeypatch +): + fence, stream = Mock(), Mock() + connector._gpu_export_fence = fence + connector._packed_gather = SimpleNamespace(device="cuda:0") + connector._get_connector_metadata = connector_module.XavierConnectorMetadata + connector._is_consumer = False + monkeypatch.setattr(torch.cuda, "current_stream", Mock(return_value=stream)) + connector.start_load_kv(SimpleNamespace()) + stream.wait_event.assert_called_once_with(fence) + assert connector._gpu_export_fence is None + connector.start_load_kv(SimpleNamespace()) + stream.wait_event.assert_called_once_with(fence) + + +@pytest.mark.skipif( + sys.platform != "linux" or not torch.cuda.is_available(), + reason="Linux CUDA required", +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_side_stream_gather_finishes_before_source_overwrite( + connector, connector_module, dtype +): + from ....xavier.backends.torch.gpu_export import GPUExportArena + from ....xavier.backends.torch.packed_gather import PackedGather + + cache = torch.zeros(8, 256, dtype=dtype, device="cuda") + gather = PackedGather.try_create({"K": cache}, {"K": 0}) + assert gather is not None + gather.warmup() + model_stream = torch.cuda.current_stream(cache.device) + + class DelayedGather: + row_bytes, device, layers = gather.row_bytes, gather.device, gather.layers + + def __call__(self, sources, output): + assert torch.cuda.current_stream(self.device) != model_stream + # Without the next-forward fence, a model-stream overwrite would + # finish before this delayed read and corrupt the immutable snapshot. + torch.cuda._sleep(20_000_000) + return gather(sources, output) + + connector._packed_gather = DelayedGather() + connector._gpu_export_arena = GPUExportArena(4096, 1, cache.device) + connector._get_connector_metadata = connector_module.XavierConnectorMetadata + connector._is_consumer = False + try: + # The gather stream must also wait for this model-stream write. + torch.cuda._sleep(20_000_000) + cache.fill_(7) + entries, owners, _ = connector._prepare_cpu_export( + {"K": cache}, [111, 222], {"K": [0, 1]}, gpu_only=True + ) + connector.start_load_kv(SimpleNamespace()) + cache.fill_(9) + torch.cuda.synchronize(cache.device) + assert owners + assert torch.equal( + entries.packed.view(dtype).view(2, 256), + torch.full((2, 256), 7, dtype=dtype, device=cache.device), + ) + assert torch.all(cache == 9) + finally: + connector._gpu_export_arena.close() + connector._gpu_export_arena = None + + +@pytest.mark.parametrize("slots,budget", [(0, 64), (2, 16)]) +def test_optional_export_congestion_never_waits_or_gathers( + connector, connector_module, slots, budget +): + connector._gpu_export_arena = SimpleNamespace( + available=list(range(slots)), slot_bytes=32 + ) + connector._ipc_export_budget = budget + connector._ipc_export_jobs["old"] = ({111}, 16, []) + connector._export_refresh_at = float("inf") + connector._registered_kv_caches = {"K": torch.ones(8, 4)} + connector._layer_group_ids = {"K": 0} + connector._poll_ipc_exports = Mock() + connector._prepare_cpu_export = Mock( + side_effect=AssertionError("unexpected gather") + ) + connector._stage_missing_registered_layers = Mock( + side_effect=AssertionError("unexpected synchronous export") + ) + connector._get_connector_metadata = ( + lambda: connector_module.XavierConnectorMetadata( + store_requests=[ + connector_module.XavierStoreRequest("new", [4], [222], [[4]]) + ] + ) + ) + try: + connector.wait_for_save() + connector._poll_ipc_exports.assert_called_once_with() + assert connector._ipc_export_jobs == {"old": ({111}, 16, [])} + assert 222 not in connector._exported_keys + finally: + connector._gpu_export_arena = None + connector._ipc_export_jobs.clear() + + +def test_partial_export_admission_bounds_bytes_and_slots_without_losing_group_ids( + connector, connector_module, monkeypatch +): + monkeypatch.setattr(connector_module, "_MAX_EXPORT_BYTES", 32) + connector._gpu_export_arena = SimpleNamespace(available=[0, 1], slot_bytes=32) + connector._ipc_export_budget = 100 + connector._ipc_export_jobs["old"] = ({999}, 16, []) + connector._registered_kv_caches = {"K": torch.ones(8, 4)} + requests = [ + connector_module.XavierStoreRequest("a", [0, 1], [100, 101], [[0, 1], [2, 3]]), + connector_module.XavierStoreRequest( + "b", [2, 3, 4, 5], [102, 103, 104, 105], [[2, 3, 4, 5], [4, 5, 6, 7]] + ), + ] + try: + admitted = connector._admit_cpu_exports(requests) + assert admitted[0] == requests[0] + assert admitted[1] == connector_module.XavierStoreRequest( + "b", [2, 3], [102, 103], [[2, 3], [4, 5]] + ) + assert connector._cpu_export_size(admitted) == 64 + assert len(requests[1].block_hashes) == 4 + # Congestion did not mark skipped content as stored or pending. A later + # request may export every missing block when ownership becomes available. + connector._ipc_export_jobs.clear() + connector._gpu_export_arena.available.extend([2]) + assert connector._admit_cpu_exports(requests) == requests + finally: + connector._gpu_export_arena = None + connector._ipc_export_jobs.clear() + + +@pytest.mark.parametrize("gpu_budget", [None, 1024]) +def test_chunked_cpu_export_uses_only_completed_blocks_and_resume_sources( + connector, connector_module, gpu_budget +): + connector._gpu_budget = gpu_budget + tokens = list(range(65)) + scheduler = SimpleNamespace( + scheduled_new_reqs=[ + SimpleNamespace( + req_id="r", + prompt_token_ids=tokens, + block_ids=([0, 1], [10, 11]), + num_computed_tokens=0, + ) + ], + num_scheduled_tokens={"r": 17}, + scheduled_cached_reqs=SimpleNamespace(req_ids=[]), + ) + try: + meta = connector_module.XavierConnectorMetadata() + connector._build_store_meta(scheduler, meta) + hashes = [key for key, _ in connector._build_xavier_hashes(tokens[:-1])] + if gpu_budget is None: + assert meta.store_requests == [ + connector_module.XavierStoreRequest("r", [0], hashes[:1], [[0], [10]]) + ] + else: + assert not meta.store_requests + scheduler.scheduled_new_reqs = [] + scheduler.scheduled_cached_reqs = SimpleNamespace( + req_ids=["r"], + num_computed_tokens=[17], + new_block_ids=[([2], [12])], + resumed_req_ids=set(), + ) + scheduler.num_scheduled_tokens["r"] = 16 + meta = connector_module.XavierConnectorMetadata() + connector._build_store_meta(scheduler, meta) + if gpu_budget is None: + assert meta.store_requests[0].block_hashes == hashes[:2] + assert meta.store_requests[0].block_ids_by_group == [[0, 1], [10, 11]] + else: + assert not meta.store_requests + # Preemption replaced every source block. Resume must gather from the + # new allocation while preserving the same content-addressed prefix. + scheduler.scheduled_cached_reqs = SimpleNamespace( + req_ids=["r"], + num_computed_tokens=[0], + new_block_ids=[([4, 5, 6, 7], [14, 15, 16, 17])], + resumed_req_ids={"r"}, + ) + scheduler.num_scheduled_tokens["r"] = 65 + meta = connector_module.XavierConnectorMetadata() + connector._build_store_meta(scheduler, meta) + assert meta.store_requests == [ + connector_module.XavierStoreRequest( + "r", [4, 5, 6, 7], hashes, [[4, 5, 6, 7], [14, 15, 16, 17]] + ) + ] + assert not connector._chunked_prefill + finally: + connector._gpu_budget = None + + class SnapshotExportActor(xo.StatelessActor): enqueue_snapshot_export_v1 = TransferActor.enqueue_snapshot_export_v1 poll_snapshot_exports_v1 = TransferActor.poll_snapshot_exports_v1 @@ -40,11 +288,19 @@ def __init__(self): self._snapshot_store = KVSnapshotStore(8) self._snapshot_export_jobs_v1 = {} self._snapshot_export_lock_v1 = asyncio.Lock() + self._swap_stream = torch.cuda.Stream() def snapshot_bits(self, layer, keys): store = self._snapshot_store return store.read(layer, keys).view(torch.uint8).reshape(-1).tolist() + def snapshots_are_pageable(self): + return all( + not layer.is_pinned() + for block in self._snapshot_store.blocks.values() + for layer in block.values() + ) + async def __pre_destroy__(self): await asyncio.gather( *self._snapshot_export_jobs_v1.values(), return_exceptions=True @@ -98,6 +354,7 @@ async def test_gpu_arena_reuse_and_fresh_registration_after_actor_recovery(dtype ) result = await actor.poll_snapshot_exports_v1([ticket], wait=True) assert set(result[ticket][0]) == set(keys) + assert await actor.snapshots_are_pageable() arena.release([slot]) assert await actor.snapshot_bits("K", [111, 222]) == expected await xo.destroy_actor(actor) @@ -486,6 +743,7 @@ async def test_actor_publishes_gpu_snapshot_without_producer_polling(dtype): stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): + torch.cuda._sleep(200_000_000) gathered = source.index_select(0, torch.tensor([2, 4], device="cuda")) packed = gathered.view(torch.uint8).reshape(2, -1) descriptor = reduce_tensor(packed)[1] diff --git a/xinference/model/llm/vllm/xavier/transfer.py b/xinference/model/llm/vllm/xavier/transfer.py index 421d9c80f8..306f1f6e8c 100644 --- a/xinference/model/llm/vllm/xavier/transfer.py +++ b/xinference/model/llm/vllm/xavier/transfer.py @@ -15,6 +15,7 @@ import logging import math import sys +from contextlib import nullcontext from functools import lru_cache from queue import Queue from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, no_type_check @@ -121,7 +122,31 @@ def get_gloo_dtype(self, input_dtype: torch.dtype): return TypeMappingGloo[TORCH_DTYPE_TO_NUMPY_DTYPE[input_dtype]] -from ...xavier.backends.torch.gpu_transfer import GPUTransferMixin +from ...xavier.backends.torch.gpu_transfer import GPUTransferMixin, finish_before_cancel + + +async def _copy_snapshot_to_cpu( + value: torch.Tensor, + stream: Optional[torch.cuda.Stream], + ready_event: Optional[torch.cuda.Event], +) -> torch.Tensor: + def copy(): + # Allocation and pageable D2H can block the caller. Keep both off the + # actor loop shared by streaming requests; no extra CPU memcpy is needed. + with torch.cuda.stream(stream) if stream is not None else nullcontext(): + if ready_event is not None: + assert stream is not None + stream.wait_event(ready_event) + try: + return value.to(device="cpu", non_blocking=False, copy=True) + except BaseException: + if stream is not None: + stream.synchronize() + raise + + task = asyncio.create_task(asyncio.to_thread(copy)) + # A cancelled export must retain its GPU source until the copy thread exits. + return await finish_before_cancel(task) class TransferActor( @@ -353,6 +378,10 @@ def enqueue_snapshot_export_v1(self, ticket, batches, tracker_address, tracker_u value = buffer[: math.prod(shape)].view(shape) else: value = rebuild_cuda_tensor(*descriptor) + # Legacy CUDA IPC reconstruction waits on the actor's current + # stream. Carry that dependency to the background copy stream. + ready = torch.cuda.Event() + ready.record(torch.cuda.current_stream(value.device)) mapped.append((keys, value, layers, ready)) batches = mapped @@ -361,41 +390,22 @@ async def export(): async with self._snapshot_export_lock_v1: tracker = await xo.actor_ref(address=tracker_address, uid=tracker_uid) for keys, value, layers, ready_event in batches: - event = torch.cuda.Event() - recorded = False - try: - if ready_event is not None: - torch.cuda.current_stream(value.device).wait_event( - ready_event - ) - host = torch.empty_like(value, device="cpu", pin_memory=True) - host.copy_(value, non_blocking=True) - event.record(torch.cuda.current_stream(value.device)) - recorded = True - # Publication progresses independently of EngineCore. - while not event.query(): - await asyncio.sleep(0.001) - # Own this immutable D2H allocation without copying it - # again. The store charges and evicts whole slabs, so - # retained memory remains bounded by its capacity. - self._snapshot_store.stage_packed_blocks( - keys, host, layers, retain_storage=True - ) - ready, removed = self.publish_blocks_v1( - keys, {layer[0] for layer in layers} - ) - if removed: - await tracker.unregister_blocks(0, self._rank, removed) - await tracker.register_snapshot_blocks(0, ready, self._rank) - available.difference_update(removed) - available.update(ready) - evicted.update(removed) - evicted.difference_update(ready) - finally: - if recorded: - event.synchronize() - else: - torch.cuda.current_stream(value.device).synchronize() + host = await _copy_snapshot_to_cpu( + value, self._swap_stream, ready_event + ) + self._snapshot_store.stage_packed_blocks( + keys, host, layers, retain_storage=True + ) + ready, removed = self.publish_blocks_v1( + keys, {layer[0] for layer in layers} + ) + if removed: + await tracker.unregister_blocks(0, self._rank, removed) + await tracker.register_snapshot_blocks(0, ready, self._rank) + available.difference_update(removed) + available.update(ready) + evicted.update(removed) + evicted.difference_update(ready) return list(available), list(evicted) self._snapshot_export_jobs_v1[ticket] = asyncio.create_task(export()) diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index ca6a708214..1beffc301b 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -18,6 +18,7 @@ import time import uuid from collections import OrderedDict +from contextlib import nullcontext from dataclasses import dataclass, field from typing import ( TYPE_CHECKING, @@ -272,6 +273,7 @@ def __init__( self._gpu_export_arena: Optional["GPUExportArena"] = None self._gpu_export_tickets: Dict[str, List[int]] = {} self._gpu_export_registered = False + self._gpu_export_fence: Optional[torch.cuda.Event] = None self._hash_cache: OrderedDict[ Tuple[int, int, Tuple[int, ...]], Tuple[Tuple[int, int], ...] ] = OrderedDict() @@ -342,6 +344,14 @@ def shutdown(self): self._loop = None def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: + # Packing may overlap logits/sampling on the model stream. Before the + # next forward or KV load can overwrite source slots, finish those reads. + # Only the GPU gather is fenced; D2H and publication remain independent. + fence = getattr(self, "_gpu_export_fence", None) + if fence is not None: + assert self._packed_gather is not None + torch.cuda.current_stream(self._packed_gather.device).wait_event(fence) + self._gpu_export_fence = None metadata = self._get_connector_metadata() assert isinstance(metadata, XavierConnectorMetadata) self._direct_sends.update(metadata.direct_sends) @@ -568,6 +578,17 @@ def wait_for_save(self): if ready is not None else self._call(self._filter_store_requests(pending)) ) + if self._gpu_export_arena is not None: + # Ordinary replica snapshots are optional. Admit only what the + # owned arena and byte budget can hold, without waiting for D2H + # or gathering payloads that will immediately need backpressure. + # Skipped blocks remain eligible for a later request to export. + missing = self._admit_cpu_exports(missing) + if missing: + nbytes = self._cpu_export_size(missing) + batches = list(self._cpu_export_batches(missing, gpu_only=True)) + self._enqueue_ipc_export(batches, nbytes) + return if missing and self._can_queue_ipc_export(missing): nbytes = self._cpu_export_size(missing) count = sum(len(request.block_hashes) for request in missing) @@ -1563,6 +1584,7 @@ def _prepare_cpu_export( slot, output = arena.allocate( (len(keys), self._packed_gather.row_bytes) ) + stream = None with profile_stage( "store_gather", device=self._packed_gather.device, @@ -1570,17 +1592,39 @@ def _prepare_cpu_export( rank=self._rank, ): try: - if output is None: - packed, owners = self._packed_gather(sources) - else: - packed, owners = self._packed_gather(sources, output) - assert arena is not None and slot is not None - arena.record(slot, tuple(packed.shape)) + if output is not None: + model_stream = torch.cuda.current_stream( + self._packed_gather.device + ) + stream = self._cpu_export_streams.get( + self._packed_gather.device + ) + if stream is None: + stream = torch.cuda.Stream( + device=self._packed_gather.device + ) + self._cpu_export_streams[self._packed_gather.device] = ( + stream + ) + stream.wait_stream(model_stream) + with ( + torch.cuda.stream(stream) + if stream is not None + else nullcontext() + ): + if output is None: + packed, owners = self._packed_gather(sources) + else: + packed, owners = self._packed_gather(sources, output) + assert arena is not None and slot is not None + arena.record(slot, tuple(packed.shape)) + fence = torch.cuda.Event() + fence.record(stream) + self._gpu_export_fence = fence except BaseException: if slot is not None: - torch.cuda.current_stream( - self._packed_gather.device - ).synchronize() + if stream is not None: + stream.synchronize() assert arena is not None arena.release([slot]) raise @@ -1678,6 +1722,40 @@ def _can_queue_ipc_export(self, requests: List[XavierStoreRequest]) -> bool: and self._cpu_export_size(requests) <= self._ipc_export_budget ) + def _admit_cpu_exports( + self, requests: List[XavierStoreRequest] + ) -> List[XavierStoreRequest]: + arena = self._gpu_export_arena + assert arena is not None + count = sum(len(request.block_hashes) for request in requests) + if not count or not arena.available: + return [] + block_bytes = self._cpu_export_size(requests) // count + pending_bytes = sum(entry[1] for entry in self._ipc_export_jobs.values()) + max_blocks = min( + max(0, self._ipc_export_budget - pending_bytes) // block_bytes, + len(arena.available) + * (min(arena.slot_bytes, _MAX_EXPORT_BYTES) // block_bytes), + ) + if max_blocks >= count: + return requests + admitted = [] + for request in requests: + take = min(len(request.block_hashes), max_blocks) + if take: + admitted.append( + XavierStoreRequest( + request.request_id, + request.block_ids[:take], + request.block_hashes[:take], + [group[:take] for group in request.block_ids_by_group], + ) + ) + max_blocks -= take + if not max_blocks: + break + return admitted + def _prepare_ipc_export( self, batches: List[_CPUExportBatch], nbytes: int ) -> Tuple[str, list, List[int]]: @@ -2102,6 +2180,10 @@ def _build_store_meta( block_ids_by_group, token_ids, ) + if self._gpu_budget is None: + self._add_store_request( + meta, new_req.req_id, token_ids, block_ids_by_group, num_tokens + ) continue self._add_store_request(meta, new_req.req_id, token_ids, block_ids_by_group) @@ -2128,6 +2210,10 @@ def _build_store_meta( block_ids_by_group, prompt_token_ids, ) + if self._gpu_budget is None: + self._add_store_request( + meta, req_id, prompt_token_ids, block_ids_by_group, num_tokens + ) continue self._add_store_request(meta, req_id, prompt_token_ids, block_ids_by_group) self._chunked_prefill.pop(req_id, None) @@ -2161,8 +2247,15 @@ def _add_store_request( request_id: str, token_ids: List[int], block_ids_by_group: List[List[int]], + num_tokens: Optional[int] = None, ) -> None: external_token_count = max(len(token_ids) - 1, 0) + if num_tokens is not None: + # The current chunk may end inside a block that a subsequent + # forward still writes. Publish only complete, computed blocks. + external_token_count = ( + min(external_token_count, num_tokens) // self._block_size + ) * self._block_size hashes = self._build_xavier_hashes(token_ids[:external_token_count]) if not hashes: return From c0e2e1ce6804e9b12c13501439041908751291c2 Mon Sep 17 00:00:00 2001 From: qinxuye Date: Thu, 8 Oct 2026 00:33:59 +0800 Subject: [PATCH 5/6] FIX: serialize Xavier snapshot refresh and recover optional exports --- .../model/llm/vllm/xavier/gc_lifecycle.py | 25 -- .../vllm/xavier/test/test_local_directory.py | 31 ++ .../llm/vllm/xavier/test/test_local_read.py | 19 ++ .../vllm/xavier/test/test_packed_gather.py | 3 +- .../xavier/test/test_prefix_cache_hot_path.py | 59 ++-- .../xavier/test/test_review_regressions.py | 33 +- .../xavier/test/test_snapshot_export_ipc.py | 140 +++++++- .../xavier/test/test_snapshot_gc_lifecycle.py | 2 +- xinference/model/llm/vllm/xavier/transfer.py | 16 +- .../model/llm/vllm/xavier/v1_connector.py | 313 +++++++----------- .../llm/xavier/backends/torch/local_read.py | 20 +- .../llm/xavier/backends/torch/snapshot.py | 2 +- 12 files changed, 368 insertions(+), 295 deletions(-) delete mode 100644 xinference/model/llm/vllm/xavier/gc_lifecycle.py diff --git a/xinference/model/llm/vllm/xavier/gc_lifecycle.py b/xinference/model/llm/vllm/xavier/gc_lifecycle.py deleted file mode 100644 index db482dadcb..0000000000 --- a/xinference/model/llm/vllm/xavier/gc_lifecycle.py +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright 2022-2026 Xinference Holdings Pte. Ltd -# Licensed under the Apache License, Version 2.0. -"""Exclude initialized libraries from snapshot allocation cycle scans.""" - -import gc - - -class InitializationGCFreeze: - def __init__(self) -> None: - self._owned = False - self._started = False - - def start(self) -> None: - if self._started: - return - self._owned = gc.get_freeze_count() == 0 - gc.collect() - gc.freeze() - self._started = True - - def close(self) -> None: - if self._owned: - gc.unfreeze() - self._owned = False - self._started = False diff --git a/xinference/model/llm/vllm/xavier/test/test_local_directory.py b/xinference/model/llm/vllm/xavier/test/test_local_directory.py index b6aa696b1e..3ac8514dbf 100644 --- a/xinference/model/llm/vllm/xavier/test/test_local_directory.py +++ b/xinference/model/llm/vllm/xavier/test/test_local_directory.py @@ -179,3 +179,34 @@ def unavailable(): assert await connector._query_remote_blocks("s", [(222, 0)]) == {} assert ref.get_snapshot_directory.await_count == 1 assert ref.query_blocks.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata", [None, {"boot_id": "foreign-host"}]) +async def test_permanently_unavailable_directory_stops_discovery(connector, metadata): + ref = SimpleNamespace( + get_snapshot_directory=AsyncMock(return_value=metadata), + query_blocks=AsyncMock(return_value={}), + ) + connector._get_tracker_ref = AsyncMock(return_value=ref) + for _ in range(3): + connector._directory_retry_at = 0 + assert await connector._query_remote_blocks("r", [(111, 0)]) == {} + assert connector._directory_disabled + ref.get_snapshot_directory.assert_awaited_once_with(0) + assert ref.query_blocks.await_count == 3 + + +@pytest.mark.asyncio +async def test_transient_directory_error_retries_without_losing_query(connector): + ref = SimpleNamespace( + get_snapshot_directory=AsyncMock(side_effect=[OSError("disconnected"), None]), + query_blocks=AsyncMock(return_value={2: {(111, 111, 0)}}), + ) + connector._get_tracker_ref = AsyncMock(return_value=ref) + assert await connector._query_remote_blocks("r", [(111, 0)]) == {2: {(111, 111, 0)}} + assert not connector._directory_disabled + connector._directory_retry_at = 0 + assert await connector._query_remote_blocks("s", [(111, 0)]) == {2: {(111, 111, 0)}} + assert connector._directory_disabled + assert ref.get_snapshot_directory.await_count == 2 diff --git a/xinference/model/llm/vllm/xavier/test/test_local_read.py b/xinference/model/llm/vllm/xavier/test/test_local_read.py index b9d781f245..b91fb930a4 100644 --- a/xinference/model/llm/vllm/xavier/test/test_local_read.py +++ b/xinference/model/llm/vllm/xavier/test/test_local_read.py @@ -223,3 +223,22 @@ def test_pinned_receive_retains_exact_bits_after_shared_slot_is_reused(dtype): finally: view.close() owner.close() + + +def test_busy_receive_lease_warns_once_until_released(caplog): + from ....xavier.backends.torch.local_read import SharedReadBuffer + + owner = SharedReadBuffer(64) + try: + lease, payload = owner.acquire(8) + for _ in range(3): + assert owner.acquire(8) is None + assert caplog.text.count("Xavier shared receive lease is still held") == 1 + owner.release(lease) + lease, payload = owner.acquire(8) + assert owner.acquire(8) is None + assert caplog.text.count("Xavier shared receive lease is still held") == 2 + owner.release(lease) + del payload + finally: + owner.close() diff --git a/xinference/model/llm/vllm/xavier/test/test_packed_gather.py b/xinference/model/llm/vllm/xavier/test/test_packed_gather.py index f745b8dd52..e5a10f72bd 100644 --- a/xinference/model/llm/vllm/xavier/test/test_packed_gather.py +++ b/xinference/model/llm/vllm/xavier/test/test_packed_gather.py @@ -84,8 +84,7 @@ def test_arena_slots_survive_recovery_and_release_only_after_completion( transfer.poll_snapshot_exports_v1.return_value = { ticket: (None, "copy failed") } - with pytest.raises(RuntimeError, match="copy failed"): - connector._poll_ipc_exports(wait=True) + connector._poll_ipc_exports(wait=True) else: transfer.poll_snapshot_exports_v1.return_value = {ticket: ([111, 222], [])} connector._poll_ipc_exports(wait=True) diff --git a/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py b/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py index 21a033da2b..4ba9d1bc4d 100644 --- a/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py +++ b/xinference/model/llm/vllm/xavier/test/test_prefix_cache_hot_path.py @@ -15,9 +15,11 @@ @pytest.fixture -def snapshot_export(connector): +def snapshot_export(connector, monkeypatch): store = KVSnapshotStore(8) - actor = SimpleNamespace(_snapshot_store=store, _rank=1) + actor = SimpleNamespace( + _snapshot_store=store, _rank=1, _snapshot_export_lock_v1=asyncio.Lock() + ) transfer = SimpleNamespace( ready_blocks_for_export_v1=AsyncMock( side_effect=lambda keys: TransferActor.ready_blocks_for_export_v1( @@ -30,13 +32,23 @@ def snapshot_export(connector): stage_layer_batches_v1=AsyncMock( side_effect=lambda *args: TransferActor.stage_layer_batches_v1(actor, *args) ), - publish_blocks_v1=AsyncMock( + publish_blocks_v1=Mock( side_effect=lambda *args: TransferActor.publish_blocks_v1(actor, *args) ), ) tracker = SimpleNamespace( register_snapshot_blocks=AsyncMock(), unregister_blocks=AsyncMock() ) + monkeypatch.setattr( + "xinference.model.llm.vllm.xavier.transfer.xo.actor_ref", + AsyncMock(return_value=tracker), + ) + actor.publish_blocks_v1 = transfer.publish_blocks_v1 + + async def refresh(*args): + return await TransferActor.refresh_snapshot_blocks_v1(actor, *args) + + transfer.refresh_snapshot_blocks_v1 = AsyncMock(side_effect=refresh) connector._get_transfer_ref = AsyncMock(return_value=transfer) connector._get_tracker_ref = AsyncMock(return_value=tracker) connector._registered_kv_caches = { @@ -150,7 +162,6 @@ def test_export_uses_post_forward_data_and_current_metadata( assert store.read("K", [222]).tolist() == [[42, 42]] assert store.read("V", [222]).tolist() == [[84, 84]] tracker.register_snapshot_blocks.assert_awaited_once_with(0, [222], 1) - assert not connector._request_staged_layers def test_repeated_export_skips_control_rpcs_until_discovery_refresh( @@ -195,7 +206,7 @@ def test_cold_export_skips_readiness_rpc_and_still_refreshes_recovered_hits( transfer.ready_blocks_for_export_v1.assert_not_awaited() # Publication supplies the authoritative ready set, removes stale mirrored # keys after actor recovery, then republishes the newly copied snapshot. - assert transfer.publish_blocks_v1.await_count == 3 + assert transfer.publish_blocks_v1.call_count == 3 assert transfer.stage_layer_batches_v1.await_count == 2 assert 222 in store.ready assert transfer.stage_layer_batches_v1.await_count == 2 @@ -242,7 +253,6 @@ def fail_v(*args, **kwargs): connector.wait_for_save() assert not store.ready assert set(store.blocks[111]) == {"K"} - assert not connector._request_staged_layers store.stage = stage connector.wait_for_save() assert store.ready == {111} @@ -337,9 +347,7 @@ def test_export_batches_bound_bytes_across_layers_and_requests( ) assert store.read("K", [333, 444]).tolist() == [[8, 9], [12, 13]] assert store.read("V", [333, 444]).tolist() == [[26, 27], [30, 31]] - transfer.publish_blocks_v1.assert_awaited_once_with( - [111, 222, 333, 444], {"K", "V"} - ) + transfer.publish_blocks_v1.assert_called_once_with([111, 222, 333, 444], {"K", "V"}) tracker.register_snapshot_blocks.assert_awaited_once_with( 0, [111, 222, 333, 444], 1 ) @@ -393,35 +401,6 @@ def test_batched_export_preserves_noncontiguous_bf16_and_snapshot_ownership( assert torch.equal(store.read("K", [111]).view(torch.int16)[0], original) -@pytest.mark.asyncio -async def test_batched_registration_keeps_distinct_layer_requirements( - connector, connector_module -): - connector._registered_kv_caches = {} - connector._request_staged_layers = {"a": {"K"}, "b": {"K", "V"}} - transfer = SimpleNamespace( - publish_blocks_v1=AsyncMock(side_effect=[([111], [222]), ([222], [])]) - ) - tracker = SimpleNamespace( - register_snapshot_blocks=AsyncMock(), unregister_blocks=AsyncMock() - ) - connector._get_transfer_ref = AsyncMock(return_value=transfer) - connector._get_tracker_ref = AsyncMock(return_value=tracker) - request_type = connector_module.XavierStoreRequest - await connector._register_blocks( - [ - request_type("a", [0], [111], [[0]]), - request_type("b", [1], [222], [[1]]), - ] - ) - assert [call.args for call in transfer.publish_blocks_v1.await_args_list] == [ - ([111], {"K"}), - ([222], {"K", "V"}), - ] - tracker.unregister_blocks.assert_awaited_once_with(0, 1, [222]) - tracker.register_snapshot_blocks.assert_awaited_once_with(0, [111, 222], 1) - - @pytest.mark.asyncio async def test_export_failure_drains_next_copy_without_publishing( connector, snapshot_export, monkeypatch @@ -435,7 +414,7 @@ async def test_export_failure_drains_next_copy_without_publishing( await connector._stage_cpu_requests([Mock()]) for event in events: event.synchronize.assert_called_once() - transfer.publish_blocks_v1.assert_not_awaited() + transfer.publish_blocks_v1.assert_not_called() tracker.register_snapshot_blocks.assert_not_awaited() @@ -688,6 +667,6 @@ def test_mixed_cold_prefix_repairs_expired_publication_once_then_uses_mirror( connector.wait_for_save() assert queued.call_count == 2 transfer.ready_blocks_for_export_v1.assert_not_awaited() - transfer.publish_blocks_v1.assert_awaited_once() + transfer.publish_blocks_v1.assert_called_once() tracker.register_snapshot_blocks.assert_awaited_once() assert connector._export_refresh_at == 11.5 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 5c9fc8c50d..d43ae1334d 100644 --- a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py +++ b/xinference/model/llm/vllm/xavier/test/test_review_regressions.py @@ -55,12 +55,18 @@ def test_cache_layout_round_trip(connector, connector_module, kv_first): source = canonical.movedim(0, 1).contiguous() if kv_first else canonical sent = [] - async def stage(request, layer, keys, blocks): - sent.append((keys, blocks.clone())) + async def stage(entries, layouts): + for (_, keys, blocks), (shape, dtype) in zip(entries, layouts): + sent.append((keys, blocks.view(dtype).view(shape).clone())) - connector._stage_layer_blocks = stage + connector._get_transfer_ref = AsyncMock( + return_value=SimpleNamespace( + stage_layer_batches_v1=AsyncMock(side_effect=stage) + ) + ) + connector.register_kv_caches({"layer": source}) request = connector_module.XavierStoreRequest("r", [2, 5], [111, 222], [[2, 5]]) - connector._stage_kv_layer_for_request(request, "layer", source) + connector._call(connector._stage_cpu_requests([request])) assert sent[0][0] == [111, 222] assert torch.equal(sent[0][1], canonical[[2, 5]]) destination = torch.zeros_like(source) @@ -125,22 +131,18 @@ def test_unsupported_configs_rejected(connector_module, connector_config, field) @pytest.mark.asyncio async def test_reused_request_id_publishes_new_content(connector, connector_module): - tracker = SimpleNamespace( - register_snapshot_blocks=AsyncMock(), unregister_blocks=AsyncMock() - ) transfer = SimpleNamespace( - publish_blocks_v1=AsyncMock(side_effect=[([111], []), ([222], [111])]) + refresh_snapshot_blocks_v1=AsyncMock(side_effect=[([111], []), ([222], [111])]) ) - connector._get_tracker_ref = AsyncMock(return_value=tracker) connector._get_transfer_ref = AsyncMock(return_value=transfer) connector._registered_kv_caches = {"layer": torch.zeros(8, 2, 16, 2, 4)} for key in [111, 222]: await connector._register_blocks( [connector_module.XavierStoreRequest("r", [1], [key], [[1]])] ) - assert tracker.register_snapshot_blocks.await_count == 2 - tracker.register_snapshot_blocks.assert_awaited_with(0, [222], 1) - tracker.unregister_blocks.assert_awaited_once_with(0, 1, [111]) + assert transfer.refresh_snapshot_blocks_v1.await_count == 2 + assert transfer.refresh_snapshot_blocks_v1.await_args.args[:2] == ([222], {"layer"}) + assert connector._exported_keys == {222} assert not hasattr(connector, "_stored_requests") @@ -406,7 +408,6 @@ async def test_gpu_connector_mapping_staging_and_load_fences( transfer.stage_gpu_requests_v1.assert_awaited_once_with( [{"layer": ([111, 222], [2, 5])}] ) - assert not connector._request_staged_layers calls.clear() load = connector_module.XavierLoadRequest( "r", {1: {111: 0}}, local_transfers_by_group={0: {1: {111: 3}}} @@ -421,11 +422,9 @@ async def test_gpu_connector_mapping_staging_and_load_fences( connector._get_connector_metadata = ( lambda: connector_module.XavierConnectorMetadata(store_requests=[request]) ) - connector._stage_kv_layer_for_request = Mock( - side_effect=AssertionError("GPU mode stages once in wait_for_save") - ) + calls.clear() connector.save_kv_layer("layer", cache, None) - connector._stage_kv_layer_for_request.assert_not_called() + assert not calls @pytest.mark.asyncio diff --git a/xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py b/xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py index 77c79ffbf6..3eee13a16a 100644 --- a/xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py +++ b/xinference/model/llm/vllm/xavier/test/test_snapshot_export_ipc.py @@ -430,16 +430,27 @@ def test_packed_export_preserves_mixed_dtypes_and_block_ownership( assert connector._exported_keys == set(keys) -def test_ipc_export_failure_releases_completed_owners_and_invalidates_mirror(connector): +@pytest.mark.parametrize("error", ["copy failed", "Unknown snapshot export ticket"]) +@pytest.mark.parametrize("arena", [False, True]) +def test_ipc_export_failure_releases_completed_owners_and_invalidates_mirror( + connector, caplog, error, arena +): transfer = SimpleNamespace( - poll_snapshot_exports_v1=AsyncMock(return_value={"bad": (None, "copy failed")}) + poll_snapshot_exports_v1=AsyncMock(return_value={"bad": (None, error)}) ) connector._get_transfer_ref = AsyncMock(return_value=transfer) connector._ipc_export_jobs["bad"] = ({111}, 16, [torch.ones(16)]) connector._exported_keys = {111, 222} - with pytest.raises(RuntimeError, match="copy failed"): - connector._poll_ipc_exports(wait=True) + if arena: + connector._gpu_export_arena = SimpleNamespace(release=Mock(), close=Mock()) + connector._gpu_export_tickets["bad"] = [0] + connector._poll_ipc_exports(wait=True) assert not connector._ipc_export_jobs and not connector._exported_keys + assert not connector._gpu_export_tickets + assert connector._export_refresh_at == 0 + assert error in caplog.text + if arena: + connector._gpu_export_arena.release.assert_called_once_with([0]) def test_new_exports_do_not_postpone_periodic_discovery_refresh(connector): @@ -775,3 +786,124 @@ async def test_actor_publishes_gpu_snapshot_without_producer_polling(dtype): result = await actor.poll_snapshot_exports_v1(["export"], wait=True) assert set(result["export"][0]) == {111, 222} assert result["export"][1] == [] + + +@pytest.mark.parametrize("free", [128 << 20, 512 << 20]) +def test_export_arena_requires_capture_headroom(connector, monkeypatch, free): + from ....xavier.backends.torch import gpu_export + + gather = SimpleNamespace(row_bytes=16, warmup=Mock(), device=torch.device("cpu")) + arena = Mock(return_value=SimpleNamespace(close=Mock())) + monkeypatch.setattr(sys, "platform", "linux") + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda _: True)) + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (free, free)) + monkeypatch.setitem( + sys.modules, + "xinference.model.llm.xavier.backends.torch.packed_gather", + SimpleNamespace( + PackedGather=SimpleNamespace(try_create=Mock(return_value=gather)) + ), + ) + monkeypatch.setattr(gpu_export, "GPUExportArena", arena) + connector.register_kv_caches({"layer": torch.zeros(8, 2, 16, 2, 4)}) + assert connector._ipc_export_budget == free // 8 + assert arena.call_count == int(free >= 512 << 20) + + +def test_export_size_reuses_registered_block_bytes(connector, connector_module): + calculate = Mock(wraps=connector._cache_export_block_bytes) + connector._cache_export_block_bytes = calculate + request = connector_module.XavierStoreRequest("r", [0, 1], [111, 222], [[0, 1]]) + connector.register_kv_caches({"layer": torch.zeros(8, 2, 16, 2, 4)}) + assert connector._cpu_export_size([request]) == 2 * 2 * 16 * 2 * 4 * 4 + assert connector._cpu_export_size([request]) == 2 * 2 * 16 * 2 * 4 * 4 + assert calculate.call_count == 1 + connector.register_kv_caches( + {"layer": torch.zeros(8, 2, 16, 2, 4, dtype=torch.float16)} + ) + assert connector._cpu_export_size([request]) == 2 * 2 * 16 * 2 * 4 * 2 + assert calculate.call_count == 2 + + +@pytest.mark.asyncio +async def test_refresh_serializes_with_export_eviction(monkeypatch): + from types import MethodType + + from torch.multiprocessing import reductions + + store = KVSnapshotStore(1) + store.stage("K", [111], torch.ones(1, 2)) + tracker = VLLMBlockTracker() + entered, finish = asyncio.Event(), asyncio.Event() + register = tracker.register_snapshot_blocks + + async def paused_register(virtual_engine, keys, rank): + if keys == [111]: + entered.set() + await finish.wait() + register(virtual_engine, keys, rank) + + ref = SimpleNamespace( + register_snapshot_blocks=AsyncMock(side_effect=paused_register), + unregister_blocks=AsyncMock(side_effect=tracker.unregister_blocks), + ) + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=ref)) + monkeypatch.setattr(reductions, "rebuild_cuda_tensor", lambda value: value) + monkeypatch.setattr(torch.cuda, "Event", Mock()) + monkeypatch.setattr(torch.cuda, "current_stream", Mock()) + + async def copy(value, stream, ready): + return await _copy_snapshot_to_cpu(value, None, None) + + monkeypatch.setattr( + "xinference.model.llm.vllm.xavier.transfer._copy_snapshot_to_cpu", copy + ) + actor = SimpleNamespace( + _rank=1, + _snapshot_store=store, + _snapshot_export_lock_v1=asyncio.Lock(), + _snapshot_export_jobs_v1={}, + _swap_stream=None, + ) + actor.publish_blocks_v1 = MethodType(TransferActor.publish_blocks_v1, actor) + refresh = asyncio.create_task( + TransferActor.refresh_snapshot_blocks_v1(actor, [111], {"K"}, "tracker", "uid") + ) + try: + await asyncio.wait_for(entered.wait(), 5) + TransferActor.enqueue_snapshot_export_v1( + actor, + "next", + [ + ( + [222], + (torch.full((1, 8), 2, dtype=torch.uint8),), + [("K", (2,), torch.float32, 0, 8)], + ) + ], + "tracker", + "uid", + ) + await asyncio.sleep(0) + assert not actor._snapshot_export_jobs_v1["next"].done() + assert 111 in store.blocks + finish.set() + await refresh + result = await TransferActor.poll_snapshot_exports_v1( + actor, ["next"], wait=True + ) + assert result["next"] == ([222], [111]) + assert tracker.query_blocks(0, [(111, 0)]) == {} + assert tracker.query_blocks(0, [(222, 0)]) == {1: {(222, 222, 0)}} + # A later refresh removes stale positives if this rank lost its snapshot. + tracker.register_snapshot_blocks(0, [333], 1) + await TransferActor.refresh_snapshot_blocks_v1( + actor, [333], {"K"}, "tracker", "uid" + ) + assert tracker.query_blocks(0, [(333, 0)]) == {} + finally: + finish.set() + await asyncio.gather( + refresh, *actor._snapshot_export_jobs_v1.values(), return_exceptions=True + ) + await tracker.__pre_destroy__() diff --git a/xinference/model/llm/vllm/xavier/test/test_snapshot_gc_lifecycle.py b/xinference/model/llm/vllm/xavier/test/test_snapshot_gc_lifecycle.py index c4b47a6ed4..eba75a361f 100644 --- a/xinference/model/llm/vllm/xavier/test/test_snapshot_gc_lifecycle.py +++ b/xinference/model/llm/vllm/xavier/test/test_snapshot_gc_lifecycle.py @@ -11,7 +11,7 @@ @pytest.mark.parametrize("mode", ["enabled", "disabled", "already-frozen"]) def test_snapshot_gc_preserves_new_request_cycles_and_external_owner(mode): - source = Path(__file__).parents[1] / "gc_lifecycle.py" + source = Path(__file__).parents[3] / "sglang" / "gc_lifecycle.py" script = r""" import gc, runpy, sys, weakref guard = runpy.run_path(sys.argv[1])["InitializationGCFreeze"]() diff --git a/xinference/model/llm/vllm/xavier/transfer.py b/xinference/model/llm/vllm/xavier/transfer.py index 306f1f6e8c..267c544561 100644 --- a/xinference/model/llm/vllm/xavier/transfer.py +++ b/xinference/model/llm/vllm/xavier/transfer.py @@ -24,11 +24,11 @@ import torch import xoscar as xo +from ...sglang.gc_lifecycle import InitializationGCFreeze from ...xavier.backends.torch.snapshot import KVSnapshotStore from ...xavier.collective import CollectiveRank from ...xavier.constants import DEFAULT_TRANSFER_ACTOR_UID from ...xavier.profiling import profile_stage -from .gc_lifecycle import InitializationGCFreeze try: from vllm.utils import TORCH_DTYPE_TO_NUMPY_DTYPE, Device @@ -330,6 +330,20 @@ def publish_blocks_v1(self, keys, layers): self._snapshot_store.evicted.clear() return available, evicted + async def refresh_snapshot_blocks_v1( + self, keys, layers, tracker_address, tracker_uid + ): + # Refresh and background export must finish tracker publication before + # either path can evict the other's snapshots. + async with self._snapshot_export_lock_v1: + tracker = await xo.actor_ref(address=tracker_address, uid=tracker_uid) + available, evicted = self.publish_blocks_v1(keys, layers) + removed = set(evicted) | (set(keys) - set(available)) + if removed: + await tracker.unregister_blocks(0, self._rank, list(removed)) + await tracker.register_snapshot_blocks(0, available, self._rank) + return available, evicted + def configure_kv_schema_v1(self, block_size: int, layers: dict, cache_dtype: str): schema = (block_size, dict(layers), cache_dtype) previous = self._kv_schema_v1 diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 1beffc301b..3819f37434 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -252,7 +252,6 @@ def __init__( 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]] = ( {} ) @@ -269,6 +268,7 @@ def __init__( self._ipc_export_poll: Optional[asyncio.Task] = None self._ipc_export_enqueues: Dict[str, asyncio.Task] = {} self._ipc_export_budget = _MAX_PENDING_EXPORT_BYTES + self._cpu_export_block_bytes: Optional[int] = None self._packed_gather: Optional["PackedGather"] = None self._gpu_export_arena: Optional["GPUExportArena"] = None self._gpu_export_tickets: Dict[str, List[int]] = {} @@ -280,6 +280,7 @@ def __init__( self._hash_cache_tokens = 0 self._local_directory: Optional["LocalBlockDirectory"] = None self._directory_retry_at = 0.0 + self._directory_disabled = False self._local_read: Optional["LocalReadBuffer"] = None self._local_read_checked = False @@ -447,9 +448,7 @@ def register_kv_caches( # Two GPU copies can briefly coexist during packing. Leave # headroom for graph capture, activations and allocator segments. free = min(torch.cuda.mem_get_info(device)[0] for device in devices) - self._ipc_export_budget = min( - _MAX_PENDING_EXPORT_BYTES, max(2 * _MAX_EXPORT_BYTES, free // 8) - ) + self._ipc_export_budget = min(_MAX_PENDING_EXPORT_BYTES, free // 8) try: from ...xavier.backends.torch.packed_gather import PackedGather @@ -464,7 +463,10 @@ def register_kv_caches( ) if self._packed_gather is not None: self._packed_gather.warmup() - if self._packed_gather.row_bytes <= _MAX_EXPORT_BYTES: + if ( + self._packed_gather.row_bytes <= _MAX_EXPORT_BYTES + and self._ipc_export_budget >= 2 * _MAX_EXPORT_BYTES + ): from ...xavier.backends.torch.gpu_export import ( GPUExportArena, ) @@ -490,6 +492,7 @@ def register_kv_caches( } self._kv_schema = XavierKVSchema(self._block_size, schema, self._cache_dtype) + self._cpu_export_block_bytes = self._cache_export_block_bytes() cache_groups = { layer_name: self._get_layer_group_id(layer_name) for layer_name in self._registered_kv_caches @@ -535,92 +538,82 @@ def wait_for_save(self): pending = metadata.store_requests if not pending: return - try: - self._poll_ipc_exports() - keys = {key for request in pending for key in request.block_hashes} - pending_keys = { - key for keys, _, _ in self._ipc_export_jobs.values() for key in keys - } - if keys.issubset(pending_keys): - return - # Only this worker writes its actor's snapshots. Publication reports - # every eviction; successful keys need no per-request control RPC. - # Periodic authoritative refresh also repairs discovery after recovery. - if time.monotonic() < self._export_refresh_at and keys.issubset( - self._exported_keys + self._poll_ipc_exports() + keys = {key for request in pending for key in request.block_hashes} + pending_keys = { + key for keys, _, _ in self._ipc_export_jobs.values() for key in keys + } + if keys.issubset(pending_keys): + return + # Only this worker writes its actor's snapshots. Publication reports + # every eviction; successful keys need no per-request control RPC. + # Periodic authoritative refresh also repairs discovery after recovery. + if time.monotonic() < self._export_refresh_at and keys.issubset( + self._exported_keys + ): + return + if time.monotonic() < self._export_refresh_at and ( + keys - self._exported_keys + ).issubset(pending_keys): + return + # New content cannot require authoritative recovery of a mirrored + # hit. Avoid a readiness RPC on every cold prefill; existing keys + # still use the periodic refresh above and below. + ready = None + if self._export_refresh_at: + if time.monotonic() >= self._export_refresh_at and ( + keys & self._exported_keys ): - return - if time.monotonic() < self._export_refresh_at and ( - keys - self._exported_keys - ).issubset(pending_keys): - return - # New content cannot require authoritative recovery of a mirrored - # hit. Avoid a readiness RPC on every cold prefill; existing keys - # still use the periodic refresh above and below. - ready = None - if self._export_refresh_at: - if time.monotonic() >= self._export_refresh_at and ( - keys & self._exported_keys - ): - # Cold traffic may reuse a small published prefix. Repair - # its authoritative publication once per refresh interval, - # rather than leaving the deadline expired and querying on - # every subsequent cold request. - self._call(self._register_blocks(pending)) - if time.monotonic() < self._export_refresh_at: - ready = (keys & self._exported_keys) | pending_keys - elif not (keys & self._exported_keys): - ready = pending_keys.copy() - if ready is not None and keys.issubset(ready): - return - missing = ( - self._missing_store_requests(pending, ready) - if ready is not None - else self._call(self._filter_store_requests(pending)) - ) - if self._gpu_export_arena is not None: - # Ordinary replica snapshots are optional. Admit only what the - # owned arena and byte budget can hold, without waiting for D2H - # or gathering payloads that will immediately need backpressure. - # Skipped blocks remain eligible for a later request to export. - missing = self._admit_cpu_exports(missing) - if missing: - nbytes = self._cpu_export_size(missing) - batches = list(self._cpu_export_batches(missing, gpu_only=True)) - self._enqueue_ipc_export(batches, nbytes) - return - if missing and self._can_queue_ipc_export(missing): + # Cold traffic may reuse a small published prefix. Repair + # its authoritative publication once per refresh interval, + # rather than leaving the deadline expired and querying on + # every subsequent cold request. + self._call(self._register_blocks(pending)) + if time.monotonic() < self._export_refresh_at: + ready = (keys & self._exported_keys) | pending_keys + elif not (keys & self._exported_keys): + ready = pending_keys.copy() + if ready is not None and keys.issubset(ready): + return + missing = ( + self._missing_store_requests(pending, ready) + if ready is not None + else self._call(self._filter_store_requests(pending)) + ) + if self._gpu_export_arena is not None: + # Ordinary replica snapshots are optional. Admit only what the + # owned arena and byte budget can hold, without waiting for D2H + # or gathering payloads that will immediately need backpressure. + # Skipped blocks remain eligible for a later request to export. + missing = self._admit_cpu_exports(missing) + if missing: nbytes = self._cpu_export_size(missing) - count = sum(len(request.block_hashes) for request in missing) - per_batch = max(1, _MAX_EXPORT_BYTES // (nbytes // count)) - slots = (count + per_batch - 1) // per_batch - if ( - nbytes + sum(entry[1] for entry in self._ipc_export_jobs.values()) - > self._ipc_export_budget - ) or ( - self._gpu_export_arena is not None - and slots > len(self._gpu_export_arena.available) - ): - self._poll_ipc_exports(wait=True) - missing = self._call(self._filter_store_requests(pending)) - nbytes = self._cpu_export_size(missing) - # Gather every source before returning to vLLM. Later engine - # steps may preempt requests or reuse their original GPU slots. batches = list(self._cpu_export_batches(missing, gpu_only=True)) self._enqueue_ipc_export(batches, nbytes) - return - if self._ipc_export_jobs: + return + if missing and self._can_queue_ipc_export(missing): + nbytes = self._cpu_export_size(missing) + if ( + nbytes + sum(entry[1] for entry in self._ipc_export_jobs.values()) + > self._ipc_export_budget + ): self._poll_ipc_exports(wait=True) missing = self._call(self._filter_store_requests(pending)) - if missing: - self._stage_missing_registered_layers(missing) - # Re-publish all requested keys, including reused snapshots. This - # refreshes discovery and retries a failed tracker registration - # without copying already published KV payloads again. - self._call(self._register_blocks(pending)) - finally: - for request in pending: - self._request_staged_layers.pop(request.request_id, None) + nbytes = self._cpu_export_size(missing) + # Gather every source before returning to vLLM. Later engine + # steps may preempt requests or reuse their original GPU slots. + batches = list(self._cpu_export_batches(missing, gpu_only=True)) + self._enqueue_ipc_export(batches, nbytes) + return + if self._ipc_export_jobs: + self._poll_ipc_exports(wait=True) + missing = self._call(self._filter_store_requests(pending)) + if missing: + self._stage_missing_registered_layers(missing) + # Re-publish all requested keys, including reused snapshots. This + # refreshes discovery and retries a failed tracker registration + # without copying already published KV payloads again. + self._call(self._register_blocks(pending)) def get_num_new_matched_tokens( self, @@ -785,6 +778,7 @@ async def claim(): self._local_directory.close() self._local_directory = None self._directory_retry_at = 0.0 + self._directory_disabled = False remote = self._call(self._query_remote_blocks(request.request_id, query_hashes)) # A remote prefix must contain the first uncached block. Cold requests # need neither hashes nor a directory RPC payload for the whole prompt. @@ -1101,14 +1095,21 @@ async def _query_remote_blocks( sys.platform == "linux" and self._gpu_budget is None and self._local_directory is None + and not self._directory_disabled and time.monotonic() >= self._directory_retry_at ): - from ...xavier.local_directory import LocalBlockDirectory + from ...xavier.local_directory import LocalBlockDirectory, _boot_id self._directory_retry_at = time.monotonic() + 1.0 - self._local_directory = LocalBlockDirectory.attach( - await tracker_ref.get_snapshot_directory(0) - ) + try: + metadata = await tracker_ref.get_snapshot_directory(0) + if metadata is None or metadata["boot_id"] != _boot_id(): + self._directory_disabled = True + else: + self._local_directory = LocalBlockDirectory.attach(metadata) + except OSError: + # A temporarily unavailable tracker or procfs is retryable. + logger.debug("Xavier local directory discovery failed", exc_info=True) if ( self._local_directory is not None and self._local_directory.contains(query_hashes[0][0]) is False @@ -1122,21 +1123,6 @@ async def _query_remote_blocks( logger.debug("Xavier V1 query result for request %s: %s", request_id, remote) return remote - async def _stage_layer_blocks( - self, - request_id: str, - layer_name: str, - block_ids: List[int], - blocks: torch.Tensor, - ): - transfer_ref = await self._get_transfer_ref() - await transfer_ref.stage_layer_blocks_v1( - request_id, - layer_name, - block_ids, - blocks, - ) - async def _filter_store_requests( self, requests: List[XavierStoreRequest], ready: Optional[set[int]] = None ) -> List[XavierStoreRequest]: @@ -1180,50 +1166,27 @@ def _missing_store_requests( async def _register_blocks(self, requests: List[XavierStoreRequest]): if not requests: return - tracker_ref = await self._get_tracker_ref() transfer_ref = await self._get_transfer_ref() expected_layers = { name for layer, cache in self._registered_kv_caches.items() for name, _ in self._iter_kv_tensors(layer, cache) } - if expected_layers: - # Registered caches share one schema. Avoid constructing layer - # groups and copying the actor's already unique result on hot hits. - keys = list( - dict.fromkeys( - key for request in requests for key in request.block_hashes - ) - ) - available, evicted = await transfer_ref.publish_blocks_v1( - keys, expected_layers - ) - else: - groups: Dict[frozenset[str], Dict[int, None]] = {} - for request in requests: - layers = frozenset( - self._request_staged_layers.get(request.request_id, set()) - ) - groups.setdefault(layers, {}).update( - dict.fromkeys(request.block_hashes) - ) - available_keys: Dict[int, None] = {} - evicted_keys: Dict[int, None] = {} - for layers, group_keys in groups.items(): - ready, removed = await transfer_ref.publish_blocks_v1( - list(group_keys), set(layers) - ) - available_keys.update(dict.fromkeys(ready)) - evicted_keys.update(dict.fromkeys(removed)) - available, evicted = list(available_keys), list(evicted_keys) + if not expected_layers: + return + keys = list( + dict.fromkeys(key for request in requests for key in request.block_hashes) + ) + available, evicted = await transfer_ref.refresh_snapshot_blocks_v1( + keys, + expected_layers, + self._xavier_config.get("block_tracker_address"), + self._xavier_config.get("block_tracker_uid"), + ) requested_keys = {key for request in requests for key in request.block_hashes} self._exported_keys.difference_update(requested_keys - set(available)) if evicted: self._exported_keys.difference_update(evicted) - await tracker_ref.unregister_blocks(0, self._rank, evicted) - # Retry discovery even when ready snapshots required no payload export. - # Transport addresses identify immutable content, not recyclable GPU slots. - await tracker_ref.register_snapshot_blocks(0, available, self._rank) self._exported_keys.update(available) self._export_refresh_at = time.monotonic() + _EXPORT_REFRESH_SECONDS logger.debug( @@ -1451,59 +1414,6 @@ def _load_layer_blocks( len(local_block_ids), ) - def _stage_kv_layer_for_request( - self, - request: XavierStoreRequest, - layer_name: str, - kv_layer: torch.Tensor | Sequence[torch.Tensor], - ) -> None: - staged_layers = self._request_staged_layers.setdefault( - request.request_id, set() - ) - for kv_layer_name, kv_tensor in self._iter_kv_tensors(layer_name, kv_layer): - if kv_layer_name in staged_layers: - continue - kv_tensor = block_major_view(kv_tensor, self._num_cache_blocks) - source_block_ids = self._get_source_block_ids(request, kv_layer_name) - num_blocks = min(len(request.block_ids), len(source_block_ids)) - if num_blocks <= 0: - continue - transport_block_ids = request.block_hashes[:num_blocks] - source_block_ids = source_block_ids[:num_blocks] - with profile_stage( - "store_d2h", - device=kv_tensor.device, - request_id=request.request_id, - layer=kv_layer_name, - blocks=num_blocks, - rank=self._rank, - ): - block_ids_tensor = torch.tensor( - source_block_ids, device=kv_tensor.device, dtype=torch.long - ) - blocks = ( - kv_tensor.index_select(0, block_ids_tensor) - .detach() - .cpu() - .contiguous() - ) - with profile_stage( - "store_rpc", - request_id=request.request_id, - layer=kv_layer_name, - nbytes=blocks.numel() * blocks.element_size(), - rank=self._rank, - ): - self._call( - self._stage_layer_blocks( - request.request_id, - kv_layer_name, - transport_block_ids, - blocks, - ) - ) - staged_layers.add(kv_layer_name) - def _stage_missing_registered_layers( self, requests: List[XavierStoreRequest] ) -> None: @@ -1696,8 +1606,8 @@ def _prepare_cpu_export( raise return entries, owners, list(events.values()) - def _cpu_export_size(self, requests: List[XavierStoreRequest]) -> int: - return sum(len(request.block_hashes) for request in requests) * sum( + def _cache_export_block_bytes(self) -> int: + return sum( ( math.prod(block_major_view(tensor, self._num_cache_blocks).shape[1:]) * tensor.element_size() @@ -1709,6 +1619,14 @@ def _cpu_export_size(self, requests: List[XavierStoreRequest]) -> int: for _, tensor in self._iter_kv_tensors(layer, cache) ) + def _cpu_export_size(self, requests: List[XavierStoreRequest]) -> int: + if self._cpu_export_block_bytes is None: + self._cpu_export_block_bytes = self._cache_export_block_bytes() + return ( + sum(len(request.block_hashes) for request in requests) + * self._cpu_export_block_bytes + ) + def _can_queue_ipc_export(self, requests: List[XavierStoreRequest]) -> bool: return ( sys.platform == "linux" @@ -1950,7 +1868,6 @@ async def poll(): return else: results = self._call(poll()) - error = None for ticket, (available, evicted) in results.items(): slots = self._gpu_export_tickets.get(ticket) if slots is not None: @@ -1961,8 +1878,10 @@ async def poll(): if available is None: self._exported_keys.clear() self._export_refresh_at = 0 - error = error or RuntimeError( - f"Xavier snapshot export failed: {evicted}" + logger.warning( + "Optional Xavier snapshot export failed: %s; " + "future requests can export it again", + evicted, ) continue self._exported_keys.difference_update(evicted) @@ -1970,8 +1889,6 @@ async def poll(): if not self._export_refresh_at: self._export_refresh_at = time.monotonic() + _EXPORT_REFRESH_SECONDS self._ipc_export_poll_at = time.monotonic() + 0.01 - if error is not None: - raise error if wait and future is not None and self._ipc_export_jobs: self._poll_ipc_exports(wait=True) diff --git a/xinference/model/llm/xavier/backends/torch/local_read.py b/xinference/model/llm/xavier/backends/torch/local_read.py index 4e8e35a586..534724568f 100644 --- a/xinference/model/llm/xavier/backends/torch/local_read.py +++ b/xinference/model/llm/xavier/backends/torch/local_read.py @@ -2,22 +2,20 @@ # Licensed under the Apache License, Version 2.0. """Bounded, leased same-host receive buffers; remote actors use normal RPCs.""" +import logging import mmap import os import tempfile import uuid -from pathlib import Path from typing import Optional, Tuple import torch +from ...local_directory import _boot_id from .request_transfer import MAX_REQUEST_BYTES _HEADER = 64 - - -def _boot_id() -> str: - return Path("/proc/sys/kernel/random/boot_id").read_text().strip() +logger = logging.getLogger(__name__) class SharedReadBuffer: @@ -43,6 +41,7 @@ def __init__(self, capacity: int = MAX_REQUEST_BYTES): self._file.close() raise self._lease: Optional[bytes] = None + self._busy_warned = False def metadata(self) -> dict: return { @@ -53,7 +52,15 @@ def metadata(self) -> dict: } def acquire(self, nbytes: int) -> Optional[Tuple[bytes, torch.Tensor]]: - if self._lease is not None or self._payload is None: + if self._lease is not None: + if not self._busy_warned: + logger.warning( + "Xavier shared receive lease is still held; using tensor RPC. " + "A lost read or release reply can keep this lease busy." + ) + self._busy_warned = True + return None + if self._payload is None: return None if not 0 < nbytes <= self.capacity: return None @@ -64,6 +71,7 @@ def acquire(self, nbytes: int) -> Optional[Tuple[bytes, torch.Tensor]]: def release(self, lease: bytes) -> None: if lease == self._lease and not self._file.closed: self._lease = None + self._busy_warned = False self._data[24:40] = bytes(16) def close(self) -> None: diff --git a/xinference/model/llm/xavier/backends/torch/snapshot.py b/xinference/model/llm/xavier/backends/torch/snapshot.py index 82a2e4f6c7..00781fbba8 100644 --- a/xinference/model/llm/xavier/backends/torch/snapshot.py +++ b/xinference/model/llm/xavier/backends/torch/snapshot.py @@ -190,7 +190,7 @@ def stage_packed_blocks( rows = packed.numpy() for index, key in enumerate(keys): fresh = key not in self.blocks - if key not in self.blocks: + if fresh: if self._used_slots() >= self.capacity and not self._evict_one(pinned): continue self.blocks[key] = {} From 143460177a4fa0c80abe316f2b3099fa33b9d099 Mon Sep 17 00:00:00 2001 From: qinxuye Date: Thu, 8 Oct 2026 02:52:09 +0800 Subject: [PATCH 6/6] fix: stop Xavier lease renewal after release despite lost cancellation --- .github/workflows/python.yaml | 1 + .../model/llm/sglang/tests/test_xavier_pd.py | 55 +++++++++++++++++++ xinference/model/llm/sglang/xavier/pd.py | 4 +- 3 files changed, 59 insertions(+), 1 deletion(-) diff --git a/.github/workflows/python.yaml b/.github/workflows/python.yaml index 4104582ad3..c1ab06cd78 100644 --- a/.github/workflows/python.yaml +++ b/.github/workflows/python.yaml @@ -376,6 +376,7 @@ jobs: --ignore xinference/model/llm/mlx \ xinference run_pytest \ + -vv \ benchmark/tests/test_benchmark_pd.py \ benchmark/tests/test_benchmark_sglang_pd.py \ benchmark/tests/test_analyze_pd_profile.py \ diff --git a/xinference/model/llm/sglang/tests/test_xavier_pd.py b/xinference/model/llm/sglang/tests/test_xavier_pd.py index 6ce52218a3..4f3a57c00d 100644 --- a/xinference/model/llm/sglang/tests/test_xavier_pd.py +++ b/xinference/model/llm/sglang/tests/test_xavier_pd.py @@ -562,6 +562,61 @@ async def test_decode_renews_completed_lease_and_stops_on_release( assert (await directory.get_stats())["active_handoffs"] == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("failed", [False, True]) +async def test_release_finishes_when_renewal_rpc_swallows_cancellation( + monkeypatch, failed +): + from ..xavier.settings import TRANSFER_TIMEOUT_ENV + + monkeypatch.setenv(TRANSFER_TIMEOUT_ENV, "0.03") + handoff = fingerprint_handoff() + entered = asyncio.Event() + + async def renew(room): + entered.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + # Model Python 3.10 wait_for returning the completed RPC's result + # or exception instead of propagating the caller's cancellation. + if failed: + raise RuntimeError("RPC completed while cancelling") + return True + + handoff._directory_actor.renew_completed = AsyncMock(side_effect=renew) + handoff._directory_actor.release = AsyncMock() + original_call = handoff._call + + async def call(method, *args, **kwargs): + if method == "renew_completed": + return await handoff._directory_actor.renew_completed(*args) + return await original_call(method, *args, **kwargs) + + monkeypatch.setattr(handoff, "_call", call) + accepted = await handoff.accept( + "prompt", config("decode")["_pd_kv_transfer_params"] + ) + task = handoff._lease_tasks[123] + releasing = None + try: + await asyncio.wait_for(entered.wait(), 1) + releasing = asyncio.create_task(handoff.release(accepted)) + done, _ = await asyncio.wait([releasing], timeout=1) + assert releasing in done + await releasing + assert task.done() and not handoff._lease_tasks + handoff._directory_actor.renew_completed.assert_awaited_once_with(123) + handoff._directory_actor.release.assert_awaited_once_with(123, "decode") + finally: + handoff._directory_actor.renew_completed.side_effect = None + handoff._directory_actor.renew_completed.return_value = False + task.cancel() + await asyncio.gather(task, return_exceptions=True) + if releasing is not None: + await asyncio.gather(releasing, return_exceptions=True) + + @pytest.mark.asyncio async def test_unconsumed_stream_reserves_no_request_or_directory_room(deployment): directory, _, decode = deployment diff --git a/xinference/model/llm/sglang/xavier/pd.py b/xinference/model/llm/sglang/xavier/pd.py index a9ec388f03..5ef6c8f406 100644 --- a/xinference/model/llm/sglang/xavier/pd.py +++ b/xinference/model/llm/sglang/xavier/pd.py @@ -79,7 +79,9 @@ async def accept(self, prompt: str, transfer: dict) -> dict: async def _renew_completed(self, room): # Completed records are needed by non-streaming decode after generation. # Keep them while D is alive; lost cleanup/router processes remain bounded. - while True: + # release() removes the owner before cancelling. Python 3.10 wait_for + # can swallow that cancellation when the RPC has just completed. + while self._lease_tasks.get(room) is asyncio.current_task(): await asyncio.sleep(min(60, transfer_timeout() / 3)) try: if not await self._call("renew_completed", room):