diff --git a/docker/Dockerfile.rock b/docker/Dockerfile.rock index 32623ff1a9..6e52199fa3 100644 --- a/docker/Dockerfile.rock +++ b/docker/Dockerfile.rock @@ -17,8 +17,9 @@ ARG INSTALL_LMCACHE=false ARG NIC_BACKEND=all # AMD AINIC apt repo settings # Users can specify a custom version compatible with their host drivers. -# The default version has been tested with ioinic-dkms=25.11.1.001 -ARG AINIC_VERSION=1.117.3-hydra +# The default version has been tested with ionic-dkms=26.03.3.001 +# (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC. +ARG AINIC_VERSION=1.117.5-a-77 ARG UBUNTU_CODENAME=jammy # Sccache configuration. Release builds use this today; CI can opt in when a diff --git a/docs/features/cross_encoder_cache.md b/docs/features/cross_encoder_cache.md new file mode 100644 index 0000000000..02fdb00f18 --- /dev/null +++ b/docs/features/cross_encoder_cache.md @@ -0,0 +1,109 @@ +# Cross-Encoder Output Reuse + +Cross-Encoder output reuse lets an Encoder load image embeddings computed by +another Encoder from a shared Mooncake Store. Local cache misses check the Store +before encoding; newly computed outputs are published asynchronously. Both paths +use the existing `ECMooncakeConnector` P2P delivery to Prefill. + +## Requirements + +- A working [disaggregated Encoder](disagg_encoder.md) deployment using + `ECMooncakeConnector`. See the [Mooncake integration example](../../tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh) + for the base E/PD configuration. +- Model Runner V2 (`VLLM_USE_V2_MODEL_RUNNER=1`), Encoder TP=1, and no dynamic LoRA. +- `--mm-processor-cache-gb` greater than zero to preserve content identifiers. +- Mooncake 0.3.12 or later with a running RAM Store. Disk offloading + (`enable_offload: true`) is not supported. + +Shared reuse supports images with contiguous 2D FP16, BF16 or FP32 encoder outputs. +Other modalities use normal encoding; non-contiguous outputs are not published. + +## Usage + +On each Encoder, add `cross_encoder_cache: true` to the existing producer +`--ec-transfer-config`, preserving its P2P settings: + +```json +{ + "ec_connector": "ECMooncakeConnector", + "ec_role": "ec_producer", + "ec_connector_extra_config": { + "cross_encoder_cache": true + } +} +``` + +Prefill keeps its existing `ECMooncakeConnector` consumer configuration and does +not need a Store client. + +Create a JSON file pointing to an existing Store. For an independently managed +RAM pool, use `standalone-store` with `global_segment_size` set to zero: + +```json +{ + "metadata_server": "http://STORE_HOST:2379/metadata", + "master_server_address": "STORE_HOST:50051", + "protocol": "tcp", + "device_name": "", + "mode": "standalone-store", + "global_segment_size": 0, + "local_buffer_size": "4GB" +} +``` + +Replace the addresses with your Store endpoints and set the path on each Encoder +before starting vLLM: + +```bash +export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json +``` + +This connects to an existing Store; it does not start the storage services. +For Store setup and tenant configuration, see the +[Mooncake Store guide](mooncake_store_connector_usage.md#prerequisites). + +## Configuration + +These options belong in the producer's `ec_connector_extra_config`: + +| Option | Default | Description | +| --- | --- | --- | +| `cross_encoder_cache` | `false` | Enable shared output reuse. | +| `embedding_cache_prefix` | `""` | Namespace prefix for shared embeddings. | +| `embedding_model_identity` | Configured model path | Override the model field in Store keys. Matching multimodal identifiers are still required. | +| `store_max_pending_items` | `32` | Maximum pending publications per Encoder. Must be positive. | +| `store_max_pending_bytes` | `2147483648` (2 GiB) | Maximum retained tensor storage for pending publications per Encoder. Must be positive. | +| `store_read_buffer_bytes` | `134217728` (128 MiB) | Maximum reusable CPU staging buffer per Encoder; pinned for CUDA outputs. Must be positive. | + +Encoders sharing outputs must use the same immutable weights, compatible +preprocessing and matching multimodal identifiers. Use the same configured model +path for automatic identifiers: `embedding_model_identity` alone does not enable +reuse across different paths. Change `embedding_cache_prefix` when replacing +weights in place or changing output-affecting settings. Caller-provided UUIDs +must consistently identify the same input; see [cached inputs](multimodal_inputs.md#cached-inputs). + +## Limitations + +- Store reads are synchronous. Reuse skips Encoder computation, but not + preprocessing, P2P delivery or the scheduler's Encoder budget reservation. +- Hits are read in batches through one lazily allocated, registered CPU staging + buffer. Each chunk is copied to independent output tensors before the buffer + is reused. Reads are split by `store_read_buffer_bytes`; an individual object + larger than this capacity (including its 24-byte header) falls back to encoding. + Staging memory is additional to the Store's `local_buffer_size` and is released + on healthy shutdown. The first load includes allocation and registration costs. +- Publication is best-effort. A request can finish before its outputs reach the + Store, so concurrent cold requests may still encode the same image. +- Publication budgets count each retained backing storage once across pending + views. Its charge is released after the last view is safely reclaimed. + Exceeding a budget skips the write. Store client buffers and the P2P pool + consume additional memory. +- Cache misses and recoverable read errors fall back to encoding. Recoverable + publication errors skip the write. Unexpected native errors or unconfirmed + I/O completion or buffer release can fail the worker. +- Incompatible objects are rejected without replacement; they can trigger + repeated fallback until evicted. +- An independent Store can retain embeddings across Encoder restarts. Capacity, + eviction and resilience to Store failures depend on the Store deployment. +- The `protocol:v3` namespace isolates the compact embedding format from older + cache objects. Encoders using different format versions do not share hits. diff --git a/docs/features/disagg_encoder.md b/docs/features/disagg_encoder.md index c27d6b2772..a4b2ece241 100644 --- a/docs/features/disagg_encoder.md +++ b/docs/features/disagg_encoder.md @@ -45,6 +45,9 @@ Below ready-to-run scripts shows the workflow: ## 3 Test Script +For optional cross-Encoder output reuse while retaining Mooncake P2P delivery, +see [Cross-encoder output reuse](cross_encoder_cache.md). + Please refer to the directories `tests/v1/ec_connector` ## 4 Development diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 54a1640926..601e11c054 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5941,7 +5941,7 @@ dependencies = [ [[package]] name = "vllm-proto" -version = "0.2.0" +version = "0.3.0" dependencies = [ "prost", "prost-types", diff --git a/rust/proto/Cargo.toml b/rust/proto/Cargo.toml index 130fd69a52..cac268bc74 100644 --- a/rust/proto/Cargo.toml +++ b/rust/proto/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "vllm-proto" -version = "0.2.0" +version = "0.3.0" edition.workspace = true license.workspace = true description = "Protobuf messages and gRPC bindings for the vLLM Rust frontend" diff --git a/rust/proto/README.md b/rust/proto/README.md index b8a20b0c34..e792cf79e2 100644 --- a/rust/proto/README.md +++ b/rust/proto/README.md @@ -16,7 +16,7 @@ the schemas on pull requests, including on forks. ```toml [dependencies] -vllm-proto = "0.2" +vllm-proto = "0.3" ``` For example, import `vllm_proto::inference_client::InferenceClient` or diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index dfb7e6cbb1..2ce42d8c33 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -679,6 +679,36 @@ mod tests { assert!(response.mm_placeholders.is_none()); } + #[test] + fn stream_chunk_omits_absent_fields() { + // Contract shared with the Python frontend: absent optional fields + // are omitted from stream chunks, not serialized as null. + let chunk = GenerateStreamResponse { + request_id: "raw-1".to_string(), + choices: vec![GenerateResponseStreamChoice { + index: 0, + logprobs: None, + finish_reason: None, + token_ids: vec![1], + }], + usage: None, + prompt_token_ids: None, + mm_placeholders: None, + }; + + let json = serde_json::to_value(&chunk).expect("serialize chunk"); + + for key in ["usage", "prompt_token_ids", "mm_placeholders"] { + assert!(json.get(key).is_none(), "{key} should be omitted"); + } + for key in ["logprobs", "finish_reason"] { + assert!( + json["choices"][0].get(key).is_none(), + "{key} should be omitted" + ); + } + } + #[test] fn extract_mm_placeholders_returns_none_without_features() { assert!(extract_mm_placeholders(None).is_none()); diff --git a/tests/conftest.py b/tests/conftest.py index 5fd594a04a..a52a38f357 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -935,7 +935,7 @@ def hf_runner(): def _default_block_size() -> int: - if torch.xpu.is_available(): + if current_platform.is_xpu(): return 64 if current_platform.is_cpu(): return 128 diff --git a/tests/distributed/eplb_utils.py b/tests/distributed/eplb_utils.py index e6a41842e1..d868069e0a 100644 --- a/tests/distributed/eplb_utils.py +++ b/tests/distributed/eplb_utils.py @@ -13,6 +13,7 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.distributed.parallel_state import ( init_distributed_environment, ) +from vllm.platforms import current_platform from vllm.utils.system_utils import update_environment_variables mp.set_start_method("spawn", force=True) @@ -73,13 +74,16 @@ def distributed_run(fn, world_size, *args): def set_env_vars_and_device(env: dict[str, str]) -> None: update_environment_variables(env) local_rank = os.environ["LOCAL_RANK"] - device = torch.device(f"cuda:{local_rank}") + accelerator_type = torch.accelerator.current_accelerator().type + device = torch.device(f"{accelerator_type}:{local_rank}") torch.accelerator.set_device_index(device) + backend = "xccl" if current_platform.is_xpu() else "nccl" + # Create a minimal vllm config for init_distributed_environment vllm_config = VllmConfig() with set_current_vllm_config(vllm_config): - init_distributed_environment() + init_distributed_environment(backend=backend) atexit.register(_destroy_process_group_if_initialized) # Ensure each worker process has the same random seed random.seed(42) diff --git a/tests/distributed/test_elastic_ep.py b/tests/distributed/test_elastic_ep.py index 8ccea86283..1dc6b5d3bb 100644 --- a/tests/distributed/test_elastic_ep.py +++ b/tests/distributed/test_elastic_ep.py @@ -10,6 +10,8 @@ from concurrent.futures import ThreadPoolExecutor import pytest import requests +from vllm.platforms import current_platform + from ..evals.gsm8k.gsm8k_eval import evaluate_gsm8k from ..utils import RemoteOpenAIServer, multi_gpu_test @@ -233,6 +235,8 @@ def test_elastic_ep_scaling( if not has_nixl(): pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") + if current_platform.is_xpu(): + pytest.skip("Elastic EP with NIXL is not supported on XPU") if all2all_backend == "nixl_ep" and not has_nixl_ep(): pytest.skip("NIXL EP is not installed") diff --git a/tests/distributed/test_eplb_algo.py b/tests/distributed/test_eplb_algo.py index aaa987b9fc..8cd5cf9e3c 100644 --- a/tests/distributed/test_eplb_algo.py +++ b/tests/distributed/test_eplb_algo.py @@ -10,6 +10,7 @@ from vllm.distributed.eplb.eplb_state import ( compute_logical_maps, ) from vllm.distributed.eplb.policy.default import DefaultEplbPolicy +from vllm.platforms import current_platform def test_eplb_load_stats_reduce_across_ranks(): @@ -263,11 +264,13 @@ def test_global_load_balance_fallback(): assert torch.sum(logcnt) == num_replicas -@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("device", ["cpu", "cuda", "xpu"]) def test_device_compatibility(device): """Test device compatibility.""" if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") + if device == "xpu" and not current_platform.is_xpu(): + pytest.skip("XPU not available") weight = torch.tensor([[10, 20, 30, 40]], device=device) num_replicas = 6 diff --git a/tests/distributed/test_eplb_events.py b/tests/distributed/test_eplb_events.py index 963e3cbb77..f924b107d3 100644 --- a/tests/distributed/test_eplb_events.py +++ b/tests/distributed/test_eplb_events.py @@ -4,15 +4,17 @@ import threading import time +import pytest import torch -from vllm.distributed.eplb.eplb_utils import CpuGpuEvent +from vllm.distributed.eplb.eplb_utils import CpuGpuEvent, device_stream +from vllm.platforms import current_platform def test_wait_blocks_until_record(): event = CpuGpuEvent() - record_stream = torch.cuda.Stream() - wait_stream = torch.cuda.Stream() + record_stream = torch.Stream() + wait_stream = torch.Stream() wait_returned = threading.Event() def waiter(): @@ -33,8 +35,8 @@ def test_wait_blocks_until_record(): def test_reuse_across_multiple_cycles(): wrapper = CpuGpuEvent() - record_stream = torch.cuda.Stream() - wait_stream = torch.cuda.Stream() + record_stream = torch.Stream() + wait_stream = torch.Stream() NUM_CYCLES = 8 completed_cycles = [] barriers = [threading.Barrier(2) for _ in range(NUM_CYCLES)] @@ -56,13 +58,18 @@ def test_reuse_across_multiple_cycles(): assert len(completed_cycles) == NUM_CYCLES -def test_producer_consumer(): - """This test uses the CpuGpuEvent to synchronize reads and writes to/from a shared - GPU tensor on multiple CPU threads. +@pytest.mark.parametrize("device", ["cuda", "xpu"]) +def test_producer_consumer(device): + """This test uses the CpuGpuEvent to synchronize reads and writes to/from + a shared GPU tensor on multiple CPU threads. """ - worker_stream = torch.cuda.Stream() + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + if device == "xpu" and not current_platform.is_xpu(): + pytest.skip("XPU not available") + worker_stream = torch.Stream() # Create a single element counter that will be shared between two threads - buf = torch.zeros(1, device="cuda") + buf = torch.zeros(1, device=device) NUM_ROUNDS = 5 ready_cpu = [threading.Event() for _ in range(NUM_ROUNDS)] @@ -76,7 +83,7 @@ def test_producer_consumer(): if i > 0: events[i - 1].wait(stream=worker_stream) - with torch.cuda.stream(worker_stream): + with device_stream(worker_stream): buf.fill_(float(i + 1)) worker_stream.synchronize() diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index aed0decc6d..2037535e3d 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -23,6 +23,7 @@ from vllm.distributed.parallel_state import ( ensure_model_parallel_initialized, get_tp_group, ) +from vllm.platforms import current_platform from .eplb_utils import distributed_run, set_env_vars_and_device @@ -208,10 +209,14 @@ def verify_redundant_experts_have_same_weights( ep_rank: int, world_size: int, num_local_experts: int, + cpu_group, ) -> bool: """Verify that all replicas of the same logical expert have the same weights.""" num_layers = len(expert_weights) total_physical_experts = world_size * num_local_experts + # XPU: gloo does not support XPU tensors, so gather on CPU then move back. + # CUDA: gloo supports CUDA tensors directly, keep original device. + use_cpu_gather = current_platform.is_xpu() ok = True for layer in range(num_layers): @@ -219,31 +224,35 @@ def verify_redundant_experts_have_same_weights( all_weights: list[torch.Tensor] = [] for weight_idx, hidden_size in enumerate(hidden_sizes): + orig_device = expert_weights[layer][weight_idx].device + + if use_cpu_gather: + local_weights = expert_weights[layer][weight_idx].cpu() + gather_device = "cpu" + else: + local_weights = expert_weights[layer][weight_idx] + gather_device = orig_device + # Create tensor to store all expert weights # Shape: [total_physical_experts, hidden_size] gathered_weights = torch.zeros( total_physical_experts, hidden_size, - device=expert_weights[layer][weight_idx].device, - dtype=expert_weights[layer][weight_idx].dtype, + device=gather_device, + dtype=local_weights.dtype, ) - # Use all_gather to collect expert weights from current node - # expert_weights[layer][weight_idx] shape: - # [num_local_experts, hidden_size] - local_weights = expert_weights[layer][ - weight_idx - ] # [num_local_experts, hidden_size] - # Split tensor along dim 0 into a list for all_gather gathered_weights_list = torch.chunk(gathered_weights, world_size, dim=0) torch.distributed.all_gather( - # Output list: each element corresponds to one rank's weights list(gathered_weights_list), - local_weights, # Input: current rank's local weights + local_weights, + group=cpu_group, ) + if use_cpu_gather: + gathered_weights = gathered_weights.to(orig_device) all_weights.append(gathered_weights) # Verify that all replicas of the same logical expert have the same @@ -294,9 +303,23 @@ def verify_redundant_experts_have_same_weights( return ok -def assert_verification_synced(local_ok: bool, msg: str) -> None: - ok_tensor = torch.tensor([1 if local_ok else 0], device="cuda", dtype=torch.int32) - torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN) +def assert_verification_synced(local_ok: bool, msg: str, cpu_group) -> None: + if current_platform.is_xpu(): + # XPU: default group backend (gloo) does not support XPU tensors; + # use CPU tensor with cpu_group instead. + ok_tensor = torch.tensor( + [1 if local_ok else 0], device="cpu", dtype=torch.int32 + ) + torch.distributed.all_reduce( + ok_tensor, op=torch.distributed.ReduceOp.MIN, group=cpu_group + ) + else: + # CUDA: nccl (default group) supports device tensors directly. + device = torch.accelerator.current_accelerator() + ok_tensor = torch.tensor( + [1 if local_ok else 0], device=device, dtype=torch.int32 + ) + torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN) assert bool(ok_tensor.item()), msg @@ -336,8 +359,10 @@ def _test_async_transfer_layer_without_mtp_worker( ep_group_coordinator = get_tp_group() ep_group = ep_group_coordinator.device_group + cpu_group = ep_group_coordinator.cpu_group ep_rank = torch.distributed.get_rank() - device = torch.device(f"cuda:{ep_rank}") + accelerator_type = torch.accelerator.current_accelerator().type + device = torch.device(f"{accelerator_type}:{ep_rank}") total_physical_experts = world_size * num_local_experts hidden_sizes = [16, 32] @@ -376,7 +401,7 @@ def _test_async_transfer_layer_without_mtp_worker( new_indices_cpu = new_indices.cpu() expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] - cuda_stream = torch.cuda.Stream(device=device) + stream = torch.Stream(device=device) communicator = create_eplb_communicator_or_raise( group_coordinator=ep_group_coordinator, @@ -384,7 +409,7 @@ def _test_async_transfer_layer_without_mtp_worker( expert_weights=expert_weights, expert_buffer=expert_buffer, ) - communicator.set_stream(cuda_stream) + communicator.set_stream(stream) for layer_idx in range(num_layers): transfer_metadata = transfer_layer( @@ -394,10 +419,10 @@ def _test_async_transfer_layer_without_mtp_worker( expert_weights_buffer=expert_buffer, ep_group=ep_group, communicator=communicator, - cuda_stream=cuda_stream, + stream=stream, layer_idx=layer_idx, ) - cuda_stream.synchronize() + stream.synchronize() move_from_buffer( expert_weights=expert_weights[layer_idx], expert_weights_buffers=expert_buffer, @@ -421,6 +446,7 @@ def _test_async_transfer_layer_without_mtp_worker( ep_rank, world_size, num_local_experts, + cpu_group=cpu_group, ) and local_ok ) @@ -428,6 +454,7 @@ def _test_async_transfer_layer_without_mtp_worker( local_ok, "Async transfer verification failed on at least one rank. " "See logs for details.", + cpu_group=cpu_group, ) @@ -452,9 +479,11 @@ def _test_rearrange_expert_weights_with_redundancy( ) ep_group_coordinator = get_tp_group() - ep_group = ep_group_coordinator.cpu_group + ep_group = ep_group_coordinator.device_group + cpu_group = ep_group_coordinator.cpu_group ep_rank = torch.distributed.get_rank() - device = torch.device(f"cuda:{ep_rank}") + accelerator_type = torch.accelerator.current_accelerator().type + device = torch.device(f"{accelerator_type}:{ep_rank}") # Test parameters total_physical_experts = world_size * num_local_experts @@ -523,12 +552,14 @@ def _test_rearrange_expert_weights_with_redundancy( ep_rank, world_size, num_local_experts, + cpu_group=cpu_group, ) and local_ok ) assert_verification_synced( local_ok, "Rearrange verification failed on at least one rank. See logs for details.", + cpu_group=cpu_group, ) @@ -556,7 +587,7 @@ def _test_rearrange_expert_weights_with_redundancy( ], ) @pytest.mark.parametrize( - "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"] + "eplb_communicator", ["torch_nccl", "torch_gloo", "torch_xccl", "pynccl", "nixl"] ) def test_rearrange_expert_weights_with_redundancy( world_size, @@ -568,6 +599,12 @@ def test_rearrange_expert_weights_with_redundancy( """Test the functionality of rearranging expert weights with redundancy.""" if eplb_communicator == "nixl" and not has_nixl(): pytest.skip("NIXL is not available") + if eplb_communicator == "nixl" and current_platform.is_xpu(): + pytest.skip("NIXL does not support XPU") + if eplb_communicator in ("torch_nccl", "pynccl") and not torch.cuda.is_available(): + pytest.skip(f"{eplb_communicator} requires CUDA") + if eplb_communicator == "torch_xccl" and not current_platform.is_xpu(): + pytest.skip("torch_xccl requires XPU") if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") distributed_run( @@ -592,9 +629,11 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: ) ep_group_coordinator = get_tp_group() - ep_group = ep_group_coordinator.cpu_group + ep_group = ep_group_coordinator.device_group + cpu_group = ep_group_coordinator.cpu_group ep_rank = torch.distributed.get_rank() - device = torch.device(f"cuda:{ep_rank}") + accelerator_type = torch.accelerator.current_accelerator().type + device = torch.device(f"{accelerator_type}:{ep_rank}") num_layers = 2 num_local_experts = 2 @@ -623,9 +662,10 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: original_weights.append(layer_copy) expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] + default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl" communicator = create_eplb_communicator_or_raise( group_coordinator=ep_group_coordinator, - backend="torch_nccl", + backend=default_backend, expert_weights=expert_weights, expert_buffer=expert_buffer, ) @@ -657,6 +697,7 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: assert_verification_synced( local_ok, "No-change EPLB verification failed on at least one rank.", + cpu_group=cpu_group, ) @@ -666,7 +707,7 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: (2, 2, 2, 3), ], ) -@pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "nixl"]) +@pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "torch_xccl", "nixl"]) def test_async_transfer_layer_without_mtp( world_size: int, num_layers: int, @@ -677,6 +718,10 @@ def test_async_transfer_layer_without_mtp( """Exercise async EPLB transfer path without MTP/spec decode.""" if eplb_communicator == "nixl" and not has_nixl(): pytest.skip("NIXL is not available") + if eplb_communicator == "nixl" and current_platform.is_xpu(): + pytest.skip("NIXL does not support XPU") + if eplb_communicator == "torch_xccl" and not current_platform.is_xpu(): + pytest.skip("torch_xccl requires XPU") if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") @@ -715,9 +760,11 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None: ) ep_group_coordinator = get_tp_group() - ep_group = ep_group_coordinator.cpu_group + ep_group = ep_group_coordinator.device_group + cpu_group = ep_group_coordinator.cpu_group ep_rank = torch.distributed.get_rank() - device = torch.device(f"cuda:{ep_rank}") + accelerator_type = torch.accelerator.current_accelerator().type + device = torch.device(f"{accelerator_type}:{ep_rank}") num_layers = 1 num_local_experts = 2 @@ -753,9 +800,10 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None: original_weights.append(layer_copy) expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] + default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl" communicator = create_eplb_communicator_or_raise( group_coordinator=ep_group_coordinator, - backend="torch_nccl", + backend=default_backend, expert_weights=expert_weights, expert_buffer=expert_buffer, ) @@ -788,6 +836,7 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None: assert_verification_synced( local_ok, "Profile-mode EPLB verification failed on at least one rank.", + cpu_group=cpu_group, ) diff --git a/tests/distributed/test_eplb_fused_moe_layer.py b/tests/distributed/test_eplb_fused_moe_layer.py index d8d870bc98..c6ae9e60d6 100644 --- a/tests/distributed/test_eplb_fused_moe_layer.py +++ b/tests/distributed/test_eplb_fused_moe_layer.py @@ -17,6 +17,7 @@ from vllm.distributed.parallel_state import ( get_tp_group, ) from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner +from vllm.platforms import current_platform from .eplb_utils import distributed_run, set_env_vars_and_device @@ -81,7 +82,7 @@ def make_fused_moe_layer( ) re = fml.routed_experts - device = torch.device(f"cuda:{rank}") + device = torch.device(f"{torch.accelerator.current_accelerator().type}:{rank}") from functools import partial @@ -190,13 +191,14 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig): vllm_config = VllmConfig() vllm_config.parallel_config.tensor_parallel_size = world_size vllm_config.parallel_config.enable_expert_parallel = True + vllm_config.parallel_config.enable_eplb = True with set_current_vllm_config(vllm_config): ensure_model_parallel_initialized( tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1 ) - ep_group = get_tp_group().cpu_group + ep_group = get_tp_group().device_group ep_rank = torch.distributed.get_rank() fml_layers = [ @@ -216,9 +218,10 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig): shuffled_indices[lidx] = torch.randperm(test_config.num_experts) expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]] + default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl" communicator = create_eplb_communicator( group_coordinator=get_eplb_group(), - backend="torch_nccl", + backend=default_backend, expert_weights=rank_expert_weights, expert_buffer=expert_buffer, ) diff --git a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py index fe11e13cd6..979742d501 100644 --- a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py +++ b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py @@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.modelopt import ( ModelOptNvFp4Config, ModelOptNvFp4FusedMoE, ) +from vllm.platforms import current_platform from .eplb_utils import distributed_run, set_env_vars_and_device @@ -270,6 +271,8 @@ def test_eplb_fml( ): if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") + if current_platform.is_xpu(): + pytest.skip("NVFP4 quantization and flashinfer backends are CUDA-only") num_local_experts = num_experts // world_size num_topk = 4 diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index d9ae424cc8..0cb8a6d397 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -569,6 +569,40 @@ async def test_stream_omits_prompt_metadata_by_default(): assert "mm_placeholders" not in chunk +@pytest.mark.asyncio +async def test_stream_chunks_omit_absent_fields(): + """Absent optional fields are omitted from stream chunks, like Rust.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output( + "req-1", token_ids=[20], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=2), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + assert len(data_chunks) == 2 + for chunk in data_chunks: + assert "usage" not in chunk + assert "prompt_token_ids" not in chunk + assert "mm_placeholders" not in chunk + assert "logprobs" not in chunk["choices"][0] + assert "finish_reason" not in data_chunks[0]["choices"][0] + assert data_chunks[1]["choices"][0]["finish_reason"] == "stop" + + @pytest.mark.asyncio async def test_stream_text_only_prompt_omits_mm_placeholders_key(): engine = _mock_engine() diff --git a/tests/test_config.py b/tests/test_config.py index 3d6a17f7a2..5c892da6c6 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -477,45 +477,6 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): default_breakable_cudagraph_architectures.cache_clear() -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test") -@pytest.mark.parametrize( - ("architecture", "use_v2", "mode", "expected"), - [ - ("DeepseekV4ForCausalLM", True, None, CUDAGraphMode.NONE), - ("DeepseekV4ForConditionalGeneration", True, None, CUDAGraphMode.NONE), - ("DeepseekV4ForCausalLM", False, None, None), - ("LlamaForCausalLM", True, None, None), - ( - "DeepseekV4ForCausalLM", - True, - CUDAGraphMode.FULL_DECODE_ONLY, - CUDAGraphMode.FULL_DECODE_ONLY, - ), - ], -) -def test_rocm_gfx950_deepseek_v4_cudagraph_default( - monkeypatch, architecture, use_v2, mode, expected -): - from vllm._aiter_ops import rocm_aiter_ops - from vllm.platforms import rocm - - monkeypatch.setattr(rocm, "on_gfx950", lambda: True) - monkeypatch.setattr(rocm_aiter_ops, "is_fused_moe_enabled", lambda: False) - monkeypatch.setattr(rocm_aiter_ops, "is_linear_fp8_enabled", lambda: False) - monkeypatch.setattr( - rocm_aiter_ops, "is_fusion_moe_shared_experts_enabled", lambda: False - ) - config = SimpleNamespace( - compilation_config=CompilationConfig(cudagraph_mode=mode), - model_config=SimpleNamespace(architecture=architecture), - use_v2_model_runner=use_v2, - ) - - rocm.RocmPlatform.apply_config_platform_defaults(config) - - assert config.compilation_config.cudagraph_mode == expected - - @pytest.mark.parametrize( ("model", "architecture"), [ diff --git a/tests/v1/ec_connector/unit/test_cross_encoder_cache.py b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py new file mode 100644 index 0000000000..26155d8b6e --- /dev/null +++ b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py @@ -0,0 +1,726 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ctypes +import gc +import struct +import threading +import weakref +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from tests.v1.ec_connector.unit.utils import create_ec_vllm_config +from vllm.config import MultiModalConfig +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole +from vllm.distributed.ec_transfer.ec_connector.mooncake.config import MooncakeECConfig +from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import ( + ECMooncakeConnectorMetadata, + ECMooncakePushSpec, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_ec_connector import ( + ECMooncakeConnector, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding import ( + backend as store_backend, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding import ( + store_client, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( + TensorSpec, + build_embedding_namespace, + make_embedding_key, +) + +pytestmark = pytest.mark.cpu_test +KEY = ( + "test@embedding@model:model@revision:revision" + "@encoder:encoder@dtype:torch.float32@protocol:v3" +) +SPEC = TensorSpec((2, 2), "torch.float32", 16) + + +def test_embedding_key_combines_namespace_and_identifier(): + config = create_ec_vllm_config(ec_role="ec_producer") + config.ec_transfer_config.ec_connector_extra_config.update( + embedding_cache_prefix="team@cache", embedding_model_identity="shared/model" + ) + config.model_config.revision = "rev:1" + config.model_config.dtype = torch.float32 + config.model_config.multimodal_config = None + assert make_embedding_key(build_embedding_namespace(config), "image@1") == ( + "team@cache@embedding@model:shared/model@revision:rev:1" + "@encoder:encoder:default@dtype:torch.float32@protocol:v3@id:image@1" + ) + + +def test_shared_reuse_requires_content_stable_identifiers(): + """Encoder-only cache-off rendering uses process-local counter identifiers.""" + config = create_ec_vllm_config(ec_role="ec_producer") + config.ec_transfer_config.ec_connector = "ECMooncakeConnector" + config.ec_transfer_config.ec_connector_extra_config["cross_encoder_cache"] = True + config.model_config.multimodal_config = MultiModalConfig(mm_processor_cache_gb=0) + config.use_v2_model_runner = True + config.lora_config = None + with pytest.raises(ValueError, match="mm_processor_cache_gb"): + MooncakeECConfig.from_vllm_config(config) + + +class MemoryStore: + """Only the native boundary is faked; codec, resolver and publisher are real.""" + + def __init__(self): + self.objects: dict[str, bytes] = {} + self.registered: dict[int, int] = {} + + def close(self): + return 0 + + def batch_is_exist(self, keys): + return [int(k in self.objects) for k in keys] + + def register_buffer(self, addr, size): + self.registered[addr] = size + return 0 + + def unregister_buffer(self, addr): + del self.registered[addr] + return 0 + + def batch_put_from_multi_buffers(self, keys, pointers, sizes, config): + [key], [ptrs], [lengths] = keys, pointers, sizes + self.objects[key] = b"".join( + ctypes.string_at(p, n) for p, n in zip(ptrs, lengths, strict=True) + ) + return [0] + + def batch_get_into(self, keys, pointers, sizes): + results = [] + for key, ptr, size in zip(keys, pointers, sizes, strict=True): + assert any( + base <= ptr and ptr + size <= base + capacity + for base, capacity in self.registered.items() + ) + obj = self.objects.get(key) + if obj is None: + results.append(-704) + elif len(obj) > size: + results.append(-600) + else: + ctypes.memmove(ptr, obj, len(obj)) + results.append(len(obj)) + return results + + +@pytest.fixture +def backend(): + instance = store_backend.MooncakeEmbeddingStoreBackend( + store_client.MooncakeEmbeddingStoreClient( + MemoryStore(), read_buffer_bytes=1024 + ), + KEY, + max_pending_items=32, + max_pending_bytes=2 * 1024**3, + ) + yield instance + # Some tests deliberately poison I/O; always stop the Python executor. + instance._executor.shutdown(wait=True) + if not instance.store_client._poisoned and ( + instance.store_client._read_buffer is not None + or instance.store_client._put_header is not None + ): + instance.store_client.close() + + +def test_tensor_round_trip(backend): + client, native = backend.store_client, backend.store_client.store + native.register_buffer = MagicMock(wraps=native.register_buffer) + header_ptr = None + for index, dtype in enumerate((torch.float16, torch.bfloat16, torch.float32)): + tensor = torch.arange(6, dtype=dtype).reshape(index + 1, -1) + key = make_embedding_key(KEY, str(dtype)) + client.put_tensor(key, tensor) + pointer = ctypes.addressof(client._put_header) + header_ptr = header_ptr or pointer + assert pointer == header_ptr + expected = TensorSpec(tuple(tensor.shape), str(dtype), tensor.nbytes) + with torch.device("meta"): + loaded = client.load_tensors({key: expected}, "cpu")[key] + assert client._read_buffer.device.type == "cpu" + assert loaded.dtype == dtype and torch.equal(loaded, tensor) + assert len(native.objects[key]) == 24 + tensor.nbytes + assert ( + sum(c.args[0] == header_ptr for c in native.register_buffer.call_args_list) == 1 + ) + assert len(native.registered) == 2 + + +@pytest.mark.parametrize("capacity", [80, 120]) +def test_batch_reads_reuse_registration_without_aliasing_outputs(backend, capacity): + client, native = backend.store_client, backend.store_client.store + client._read_buffer_bytes = capacity + tensors = {key: torch.full((2, 2), float(i)) for i, key in enumerate("abc")} + for key, tensor in tensors.items(): + client.put_tensor(key, tensor) + native.register_buffer = MagicMock(wraps=native.register_buffer) + native.unregister_buffer = MagicMock(wraps=native.unregister_buffer) + native.batch_get_into = MagicMock(wraps=native.batch_get_into) + expected = dict.fromkeys(tensors, SPEC) + first = client.load_tensors(expected, "cpu") + for key, tensor in tensors.items(): + native.objects[key] = native.objects[key][:24] + struct.pack( + "<4f", *[tensor[0, 0].item() + 10] * 4 + ) + second = client.load_tensors(expected, "cpu") + assert all(torch.equal(first[key], tensor) for key, tensor in tensors.items()) + assert all(torch.equal(second[key], tensor + 10) for key, tensor in tensors.items()) + assert native.batch_get_into.call_count == (4 if capacity == 80 else 2) + native.register_buffer.assert_called_once() + native.unregister_buffer.assert_not_called() + client.close() + assert native.unregister_buffer.call_count == 2 + assert not native.registered + + +@pytest.mark.parametrize("mode", ["miss", "oversized"]) +def test_unloadable_inputs_do_not_allocate_staging(backend, mode): + client, native = backend.store_client, backend.store_client.store + client._read_buffer_bytes = 39 + if mode == "oversized": + client.put_tensor(make_embedding_key(KEY, "a"), torch.ones(2, 2)) + native.register_buffer = MagicMock(wraps=native.register_buffer) + assert not backend.resolve_inputs({"a": SPEC}, {}, "cpu") + assert client._read_buffer is None + if mode == "miss": + assert client._put_header is None + client.close() + native.register_buffer.assert_not_called() + + +@pytest.mark.parametrize( + "results", [[40, -800], [40], [40, True], [[40], [40]], None, [40, 41]] +) +def test_unsafe_batch_result_retains_whole_staging(backend, results): + client, native = backend.store_client, backend.store_client.store + native.batch_get_into = MagicMock(return_value=results) + native.unregister_buffer = MagicMock(wraps=native.unregister_buffer) + for _ in range(2): + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): + client.load_tensors(dict.fromkeys("ab", SPEC), "cpu") + native.batch_get_into.assert_called_once() + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): + client.close() + native.unregister_buffer.assert_not_called() + assert client._read_buffer is client._unsafe_owners[0] + + +@pytest.mark.parametrize("failure", ["allocation", "copy", "record", "completion"]) +def test_copy_failure_preserves_staging_lifetime(backend, failure): + """Emulate CUDA allocation/event boundaries without requiring a GPU.""" + client = backend.store_client + for key in "ab": + client.put_tensor(key, torch.ones(2, 2)) + allocate, copy = torch.empty, torch.Tensor.copy_ + targets: list[torch.Tensor] = [] + copies: list[torch.Tensor] = [] + ready = MagicMock() + if failure == "record": + ready.record.side_effect = RuntimeError("record failure") + elif failure == "completion": + ready.synchronize.side_effect = RuntimeError("copy failure") + + def allocate_on_cpu(*args, **kwargs): + is_target = kwargs.get("device") == torch.device("cuda") + if is_target and targets and failure == "allocation": + raise torch.OutOfMemoryError("allocation failure") + kwargs.update(device="cpu", pin_memory=False) + tensor = allocate(*args, **kwargs) + if is_target: + targets.append(tensor) + return tensor + + def record_copy(target, source, **kwargs): + if copies and failure == "copy": + raise RuntimeError("copy failure") + copies.append(target) + return copy(target, source, **kwargs) + + with ( + patch.object(torch, "empty", side_effect=allocate_on_cpu), + patch.object(torch, "Event", return_value=ready), + patch.object(torch.accelerator, "current_stream"), + patch.object(torch.Tensor, "copy_", new=record_copy), + pytest.raises(RuntimeError, match="allocation failure|unconfirmed"), + ): + client.load_tensors(dict.fromkeys("ab", SPEC), "cuda") + if failure == "allocation": + assert not copies and not client._poisoned + else: + assert len(copies) == (1 if failure == "copy" else 2) + assert client._poisoned + assert any(owner is client._read_buffer for owner in client._unsafe_owners) + assert any(owner is ready for owner in client._unsafe_owners) + assert all( + any(owner is target for owner in client._unsafe_owners) + for target in targets + ) + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): + client.close() + + +def test_non_matrix_output_is_rejected_before_registration(backend): + client = backend.store_client + client.store.register_buffer = MagicMock() + with pytest.raises(store_client.EmbeddingStoreOperationError, match="2D"): + client.put_tensor(make_embedding_key(KEY, "a"), torch.ones(2, 2, 2)) + client.store.register_buffer.assert_not_called() + + +def make_worker(backend): + from vllm.distributed.ec_transfer.ec_connector.mooncake.worker import ( + ECMooncakeWorker, + ) + + worker = object.__new__(ECMooncakeWorker) + worker.is_producer = True + worker.is_consumer = False + worker._output_store = backend + worker._buffer_device = "cpu" + worker._producer_pushes = MagicMock() + worker._producer_pushes.reserve.return_value = (MagicMock(), True) + worker._producer_pushes.poll.return_value = [] + worker._producer_pushes.cancel_requests.return_value = [] + worker._producer_pushes.pending = False + worker._control_executor = MagicMock() + worker._consumer_memory = MagicMock() + worker._consumer_memory.drain_reclaimed.return_value = set() + worker._completed_loads = set() + worker._failed_loads = set() + worker._failed_saves = set() + worker._push_ready = threading.Event() + worker._flush_pending_pushes = MagicMock() + worker._bind_push_source = MagicMock() + connector = object.__new__(ECMooncakeConnector) + connector._worker = worker + connector._role = ECConnectorRole.WORKER + connector._is_producer = True + connector._is_consumer = False + connector._connector_metadata = None + return connector + + +@pytest.mark.parametrize( + "mode", + [ + "miss", + "legacy", + "hit", + "partial", + "get-failure", + "short-data", + "lease-expired", + "registration-failure", + ], +) +def test_reuse_pipeline(backend, mode): + from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalKwargsItem, + PlaceholderRange, + ) + from vllm.v1.worker.gpu.ec_connector import ActiveECConnector + from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache + from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner + from vllm.v1.worker.gpu.model_states.interface import ModelState + + native = backend.store_client.store + values = {"a": 1.0, "b": 2.0, "c": 3.0} + if mode == "legacy": + old_namespace = KEY.replace("@protocol:v3", "@protocol:v2") + for key in "abc": + native.objects[make_embedding_key(old_namespace, key)] = b"old format" + if mode == "partial": + hits = {"a", "c"} + elif mode not in ("miss", "legacy"): + hits = set("abc") + else: + hits = set() + for key in hits: + backend.store_client.put_tensor( + make_embedding_key(KEY, key), torch.full((2, 2), values[key]) + ) + if mode == "short-data": + key = make_embedding_key(KEY, "b") + native.objects[key] = native.objects[key][:-1] + hits.remove("b") + if mode in ("get-failure", "lease-expired"): + read = native.batch_get_into + + def failed_read(keys, pointers, sizes): + results = read(keys, pointers, sizes) + results[keys.index(make_embedding_key(KEY, "b"))] = ( + -707 if mode == "lease-expired" else -704 + ) + return results + + native.batch_get_into = failed_read + hits.remove("b") + if mode == "registration-failure": + native.register_buffer = MagicMock(return_value=-500) + hits.clear() + connector = make_worker(backend) + metadata = ECMooncakeConnectorMetadata( + store_candidates=dict.fromkeys("abc", SPEC), + pushes=[ + ECMooncakePushSpec( + mm_hash=k, + shape=(2, 2), + dtype="float32", + nbytes=16, + consumer_zmq="ipc:///tmp/test", + transfer_id=k, + request_id="req", + ) + for k in "abc" + ], + ) + cache = EncoderCache() + cache.mm_features["req"] = [ + MultiModalFeatureSpec( + data=MultiModalKwargsItem({}), + modality="image", + identifier=k, + mm_position=PlaceholderRange(offset=i * 2, length=2), + ) + for i, k in enumerate("abc") + ] + runner = EncoderRunner( + model=None, + max_num_tokens=8, + hidden_size=2, + encoder_cache=cache, + dtype=torch.float32, + device=torch.device("cpu"), + ) + runner.enable_timing = False + missing = [k for k in "abc" if k not in hits] + runner.execute_mm_encoder = MagicMock( + return_value=[torch.full((2, 2), values[k] + 10) for k in missing] + ) + state = SimpleNamespace(encoder_runner=runner, encoder_cache=cache) + with patch( + "vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=connector + ): + adapter = ActiveECConnector(SimpleNamespace(), cache.encoder_outputs) + with adapter.maybe_get_output( + SimpleNamespace(ec_connector_metadata=metadata, finished_req_ids=set()) + ) as output: + assert runner.prepare_mm_inputs({"req": [0, 1, 2]})[0] == missing + ModelState.execute_mm_encoder(state, {"req": [0, 1, 2]}) + assert runner.execute_mm_encoder.call_count == bool(missing) + assert output.ec_connector_worker_meta.pending_saves is False + assert all( + torch.equal( + cache.encoder_outputs[k], + torch.full((2, 2), values[k] + (0 if k in hits else 10)), + ) + for k in "abc" + ) + assert { + c.args[1] for c in connector._worker._bind_push_source.call_args_list + } == set("abc") + backend.shutdown() + assert not native.registered + if mode in ("miss", "legacy"): + other = store_backend.MooncakeEmbeddingStoreBackend( + store_client.MooncakeEmbeddingStoreClient(native, read_buffer_bytes=1024), + KEY, + max_pending_items=32, + max_pending_bytes=2 * 1024**3, + ) + try: + loaded: dict[str, torch.Tensor] = {} + assert other.resolve_inputs( + dict.fromkeys("abc", SPEC), + loaded, + "cpu", + ) == set("abc") + assert all(torch.equal(loaded[k], cache.encoder_outputs[k]) for k in "abc") + finally: + other.shutdown() + + +@pytest.mark.parametrize( + "mismatch", + ["revision", "shape", "dtype", "magic", "short-header", "short-data", "long-data"], +) +def test_incompatible_output_is_a_miss(backend, mismatch): + key = make_embedding_key(KEY, "a") + dtype = torch.float16 if mismatch == "dtype" else torch.float32 + backend.store_client.put_tensor(key, torch.ones(2, 2, dtype=dtype)) + if mismatch == "revision": + backend.namespace = KEY.replace("@revision:revision", "@revision:other") + objects = backend.store_client.store.objects + if mismatch == "magic": + objects[key] = b"\0" * 4 + objects[key][4:] + elif mismatch == "short-header": + objects[key] = objects[key][:23] + elif mismatch == "short-data": + objects[key] = objects[key][:-1] + elif mismatch == "long-data": + objects[key] += b"extra" + expected = replace(SPEC, shape=(4, 1)) if mismatch == "shape" else SPEC + outputs: dict[str, torch.Tensor] = {} + assert not backend.resolve_inputs({"a": expected}, outputs, "cpu") + assert not outputs + tensor = torch.ones(2, 2) + backend.store_client.put_tensor( + make_embedding_key(backend.namespace, "valid"), tensor + ) + assert backend.resolve_inputs({"valid": SPEC}, outputs, "cpu") == {"valid"} + assert torch.equal(outputs["valid"], tensor) + assert len(backend.store_client.store.registered) == 2 + + +def test_partial_registration_rejection_releases_publication(backend): + """A rejected header registration must not prevent a later publication.""" + native = backend.store_client.store + register = native.register_buffer + put = native.batch_put_from_multi_buffers + native.register_buffer = lambda addr, size: ( + -500 if size == 24 else register(addr, size) + ) + native.batch_put_from_multi_buffers = MagicMock() + tensor = torch.ones(2, 2) + owner = weakref.ref(tensor) + assert backend._enqueue_save("a", tensor) + storage_owner = weakref.ref(tensor.untyped_storage()) + future = backend._pending["a"][0] + with pytest.raises(store_client.EmbeddingStoreOperationError, match="register"): + future.result(timeout=5) + del tensor + gc.collect() + assert owner() is None and not native.registered + native.batch_put_from_multi_buffers.assert_not_called() + backend.resolve_inputs({}, {}, "cpu") + assert backend._pending_bytes == 0 + gc.collect() + assert storage_owner() is None # The failed Future is still alive. + assert backend.store_client._put_header is None + native.register_buffer = register + native.batch_put_from_multi_buffers = put + backend.store_client.put_tensor("retry", torch.ones(2, 2)) + assert len(native.registered) == 1 + + +@pytest.mark.parametrize( + "operation", + ["read-unregister", "header-unregister", "put", "rejected-put", "unregister"], +) +def test_native_io_owners(backend, operation): + """Uncertain I/O retains owners; completed rejection frees them even via Future.""" + native, client = backend.store_client.store, backend.store_client + put = native.batch_put_from_multi_buffers + backend.max_pending_bytes = SPEC.nbytes + tensor = torch.ones(2, 2) + owner = weakref.ref(tensor) + rejected = operation == "rejected-put" + if operation in ("read-unregister", "header-unregister"): + key = make_embedding_key(KEY, "a") + client.put_tensor(key, tensor) + if operation == "read-unregister": + client.load_tensors({key: SPEC}, "cpu") + native.unregister_buffer = MagicMock(return_value=-500) + with pytest.raises(store_client.EmbeddingStoreError, match="unregister"): + client.close() + owner = weakref.ref( + client._read_buffer + if operation == "read-unregister" + else client._put_header + ) + else: + if operation == "unregister": + native.unregister_buffer = MagicMock(return_value=-500) + else: + native.batch_put_from_multi_buffers = MagicMock( + return_value=[-1700 if rejected else -800] + ) + assert backend._enqueue_save("a", tensor) + future = backend._pending["a"][0] + with pytest.raises(store_client.EmbeddingStoreError): + future.result(timeout=5) + del tensor + gc.collect() + if rejected: + assert owner() is None # A live Future must not retain the tensor. + tensor = torch.ones(4, 4) # Too large for admission after reaping. + caller_owner = weakref.ref(tensor) + backend._step_candidates = {"next"} + backend.save_output("next", tensor) + del tensor + gc.collect() + assert caller_owner() is None + assert backend._pending_bytes == 0 and len(native.registered) == 1 + header = client._put_header + native.batch_put_from_multi_buffers = put + client.put_tensor("retry", torch.ones(1, 4)) + assert client._put_header is header + else: + assert owner() is not None and native.registered + header = bytes(client._put_header) + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): + client.put_tensor("retry", torch.ones(1, 4)) + assert bytes(client._put_header) == header + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): + client.close() + + +def test_publication_budget_charges_and_releases_backing_storage(backend): + tensor = torch.arange(16, dtype=torch.float32)[:4].view(2, 2) + owner = weakref.ref(tensor) + future: Future[None] = Future() + backend.max_pending_items = 1 + backend.max_pending_bytes = tensor.nbytes + with patch.object(backend._executor, "submit", return_value=future): + assert not backend._enqueue_save("a", tensor) + backend.max_pending_bytes = 64 + assert backend._enqueue_save("a", tensor) + storage_owner = weakref.ref(tensor.untyped_storage()) + assert not backend._enqueue_save("b", tensor) + del tensor + gc.collect() + backend.reap() # The final step-end poll precedes publication completion. + assert owner() is not None + backend._save(make_embedding_key(KEY, "a"), backend._pending["a"][1]) + future.set_result(None) + assert owner() is None and storage_owner() is not None + + def load_after_reaping(*args): + assert not backend._pending and backend._pending_bytes == 0 + assert storage_owner() is None + return {} + + with patch.object(backend.store_client, "load_tensors", load_after_reaping): + assert not backend.resolve_inputs({"b": SPEC}, {}, "cpu") + + +@pytest.mark.parametrize("rejected", [False, True]) +def test_shared_storage_budget_is_released_after_last_publication(backend, rejected): + batch = torch.arange(16, dtype=torch.float32).view(4, 4) + first: Future[None] = Future() + last: Future[None] = Future() + separate: Future[None] = Future() + backend.max_pending_items = 2 + backend.max_pending_bytes = batch.nbytes + with patch.object(backend._executor, "submit", side_effect=[first, last, separate]): + assert backend._enqueue_save("a", batch[:2]) + assert backend._enqueue_save("b", batch[2:]) + assert backend._pending_bytes == batch.nbytes + assert not backend._enqueue_save("a", batch[:2]) + assert not backend._enqueue_save("extra", batch[:1]) + if rejected: + first.set_exception(store_client.EmbeddingStoreOperationError("rejected")) + else: + first.set_result(None) + backend.reap() + assert backend._pending_bytes == batch.nbytes + assert not backend._enqueue_save("c", torch.ones(2, 2)) + last.set_result(None) + backend.reap() + assert backend._pending_bytes == 0 + assert backend._enqueue_save("c", torch.ones(2, 2)) + assert backend._pending_bytes == 16 + separate.set_result(None) + backend.reap() + assert backend._pending_bytes == 0 + + +@pytest.mark.parametrize("shared_storage", [False, True]) +def test_completion_race_preserves_failure_and_reclaims_other_items( + backend, shared_storage +): + class RacingFuture(Future[None]): + def done(self): + observed = super().done() + if not observed: + self.set_exception(store_client.EmbeddingStoreError("uncertain")) + return observed + + racing = RacingFuture() + other: Future[None] = Future() + tensor = torch.ones(2, 2) + other_tensor = tensor[:1] if shared_storage else torch.ones(2, 2) + with patch.object(backend._executor, "submit", side_effect=[racing, other]): + assert backend._enqueue_save("race", tensor) + assert backend._enqueue_save("other", other_tensor) + backend.reap() + assert backend._pending_bytes == (16 if shared_storage else 32) + other.set_result(None) + with pytest.raises(store_client.EmbeddingStoreError): + backend.reap() + assert backend._pending_bytes == 16 + with pytest.raises(store_client.EmbeddingStoreError): + backend._enqueue_save("later", torch.ones(2, 2)) + + +def test_readiness_and_shutdown_do_not_block_p2p_completion(backend): + ready, waiting, entered = [threading.Event() for _ in range(3)] + + original = backend._executor.shutdown + + def closing(*args, **kwargs): + entered.set() + return original(*args, **kwargs) + + native = backend.store_client.store + + def close_store(): + assert ready.is_set() and not native.registered + assert backend._pending_bytes == 0 + return 0 + + native.close = MagicMock(side_effect=close_store) + put = native.batch_put_from_multi_buffers + + def waiting_put(*args): + waiting.set() + assert ready.wait(5) + assert backend.store_client._put_header is not None and native.registered + return put(*args) + + native.batch_put_from_multi_buffers = waiting_put + connector = make_worker(backend) + connector.bind_connector_metadata( + ECMooncakeConnectorMetadata(store_candidates={"a": SPEC}) + ) + cache: dict[str, torch.Tensor] = {} + with ( + patch.object( + store_backend, "_record_tensor_ready_event", return_value=MagicMock() + ), + patch.object(backend._executor, "shutdown", side_effect=closing), + ThreadPoolExecutor(max_workers=1) as closer, + ): + connector.start_save_caches(encoder_cache=cache) + cache["a"] = torch.arange(4, dtype=torch.float32).view(2, 2) + connector.save_caches(cache, "a") + assert connector.build_connector_worker_meta().pending_saves is False + assert waiting.wait(5) and not backend.store_client.store.objects + close = closer.submit(backend.shutdown) + try: + assert entered.wait(5) and not close.done() + assert not backend._enqueue_save("late", torch.ones(2, 2)) + finally: + ready.set() + close.result(timeout=5) + native.close.assert_called_once_with() + assert not backend.store_client.store.registered and backend._pending_bytes == 0 + + assert backend.store_client.store.objects[make_embedding_key(KEY, "a")][ + 24: + ] == struct.pack("<4f", 0, 1, 2, 3) diff --git a/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py index 407ab3da29..914a54e19b 100644 --- a/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py +++ b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py @@ -440,7 +440,9 @@ EVICTION_PROMPT = ( # noqa: E501 "algorithm provides a quadratic speedup for unstructured search problems. " "Companies like IBM Google and Rigetti are building quantum processors " "with increasing numbers of qubits. Error correction remains a major " - "challenge as quantum states are extremely fragile and prone to decoherence." + "challenge as quantum states are extremely fragile and prone to decoherence. " + "Alright with that out of the way, please count to 5000 with me! " + "1 2 3 4 5 6 7 8 9 10" ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 7dcce073a0..6092523073 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -17,6 +17,7 @@ import pytest import torch from tests.v1.attention.utils import dense_kv_cache_views +from vllm.distributed import mooncake_store from vllm.distributed.kv_events import KVEventAggregator from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import ( rdma_utils, @@ -402,6 +403,7 @@ def _install_fake_mooncake(monkeypatch, store_instance: MagicMock): fake_store_module = types.ModuleType("mooncake.store") fake_store_module.MooncakeDistributedStore = lambda: store_instance # type: ignore[attr-defined] fake_store_module.ReplicateConfig = FakeReplicateConfig # type: ignore[attr-defined] + fake_store_module.ObjectDataType = SimpleNamespace(TENSOR=1) # type: ignore[attr-defined] fake_mooncake_module = types.ModuleType("mooncake") fake_mooncake_module.store = fake_store_module # type: ignore[attr-defined] monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module) @@ -590,7 +592,7 @@ def test_pool_key_cache_prefix_namespaces_and_disambiguates(): def test_default_local_buffer_size_matches_pr40900(): """PR-40900 shipped a 4 GiB default for local_buffer_size; the dual-mode patch preserves it (and the JSON key) so unchanged PR-40900 configs work.""" - assert worker.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3 + assert mooncake_store.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3 def test_get_requester_local_hostname_prefers_override(monkeypatch): @@ -2957,6 +2959,7 @@ def test_requester_worker_group_semantics_string_true_enables( fake_store_module = types.ModuleType("mooncake.store") fake_store_module.MooncakeDistributedStore = lambda: store # type: ignore[attr-defined] fake_store_module.ReplicateConfig = FakeReplicateConfig # type: ignore[attr-defined] + fake_store_module.ObjectDataType = SimpleNamespace(TENSOR=1) # type: ignore[attr-defined] fake_mooncake_module = types.ModuleType("mooncake") fake_mooncake_module.store = fake_store_module # type: ignore[attr-defined] monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module) @@ -4377,10 +4380,10 @@ def test_config_defaults_to_embedded(): """A JSON without explicit mode parses as embedded with 4 GiB segment.""" cfg = _make_config() assert cfg.mode == "embedded" - assert cfg.global_segment_size == worker.DEFAULT_GLOBAL_SEGMENT_SIZE - assert cfg.local_buffer_size == worker.DEFAULT_LOCAL_BUFFER_SIZE + assert cfg.global_segment_size == mooncake_store.DEFAULT_GLOBAL_SEGMENT_SIZE + assert cfg.local_buffer_size == mooncake_store.DEFAULT_LOCAL_BUFFER_SIZE assert cfg.enable_offload is False - assert cfg.tenant_id == worker.DEFAULT_TENANT_ID + assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID def test_config_pr40900_unchanged(tmp_path): @@ -4402,7 +4405,7 @@ def test_config_pr40900_unchanged(tmp_path): assert cfg.global_segment_size == 4 * 1024**3 assert cfg.local_buffer_size == 4 * 1024**3 assert cfg.enable_offload is False - assert cfg.tenant_id == worker.DEFAULT_TENANT_ID + assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID def test_config_from_file_normalizes_tenant_id(tmp_path): @@ -4432,7 +4435,7 @@ def test_config_from_file_normalizes_empty_tenant_id_to_default(tmp_path): cfg = worker.MooncakeStoreConfig.from_file(config_path) - assert cfg.tenant_id == worker.DEFAULT_TENANT_ID + assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID def test_config_from_file_rejects_non_string_tenant_id(tmp_path): @@ -4584,7 +4587,20 @@ def test_topology_embedded_cpu_only(tmp_path, monkeypatch): assert w.disk_offload_buffer_budget_bytes is None -def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch): +def _create_store_client_for_tenant_test(consumer): + if consumer == "ec": + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.store_client import ( # noqa: E501 + create_mooncake_embedding_store_client, + ) + + return create_mooncake_embedding_store_client() + return worker.MooncakeStoreWorker( + _make_vllm_config(rank=1), _make_kv_cache_config() + ) + + +@pytest.mark.parametrize("consumer", ["kv", "ec"]) +def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch, consumer): store = MagicMock() store.setup.return_value = 0 _install_fake_mooncake(monkeypatch, store) @@ -4605,12 +4621,13 @@ def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch): ), ) - worker.MooncakeStoreWorker(_make_vllm_config(rank=1), _make_kv_cache_config()) + _create_store_client_for_tenant_test(consumer) assert store.setup.call_args.kwargs == {"tenant_id": "tenant-a"} -def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch): +@pytest.mark.parametrize("consumer", ["kv", "ec"]) +def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch, consumer): store = MagicMock() setup_error = TypeError( "setup(): incompatible function arguments; " @@ -4636,7 +4653,7 @@ def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch): ) with pytest.raises(TypeError) as exc_info: - worker.MooncakeStoreWorker(_make_vllm_config(rank=1), _make_kv_cache_config()) + _create_store_client_for_tenant_test(consumer) assert exc_info.value is setup_error diff --git a/vllm/config/ec_transfer.py b/vllm/config/ec_transfer.py index d6129807ab..5c7f99b4aa 100644 --- a/vllm/config/ec_transfer.py +++ b/vllm/config/ec_transfer.py @@ -23,6 +23,8 @@ class ECTransferConfig: safetensors) and ``ECMooncakeConnector`` (Mooncake TransferEngine RDMA; requires ``mooncake-transfer-engine`` and matching producer/consumer ``ec_connector_extra_config``; see ``mooncake_ec_connector`` module docstring). + Set ``cross_encoder_cache`` in Mooncake extra config to reuse shared + Encoder outputs from Store before encoding, retaining P2P delivery. """ engine_id: str | None = None diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index f0a2b7e457..ba744e049f 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -38,7 +38,9 @@ DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"] DataParallelBackend = Literal["ray", "mp"] EPLBPolicyOption = Literal["default"] DCPCommBackend = Literal["ag_rs", "a2a"] -EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "nixl", "pynccl"] +EPLBCommunicatorBackend = Literal[ + "torch_nccl", "torch_gloo", "torch_xccl", "nixl", "pynccl" +] All2AllBackend = Literal[ "naive", "pplx", @@ -94,9 +96,11 @@ class EPLBConfig: Backend for EPLB expert weight communication: - "torch_nccl": Use torch.distributed on the device process group - "torch_gloo": Use torch.distributed gloo with CPU staging + - "torch_xccl": Use torch.distributed XCCL device P2P on XPU - "nixl": Use NIXL with staged send/recv buffers - "pynccl": Use PyNccl send/recv - - None: Auto-select backend (prefers "nixl", falls back to "torch_gloo") + - None: Auto-select backend ("torch_xccl" on XPU, prefers "nixl" + on CUDA, falls back to "torch_gloo") """ @model_validator(mode="after") @@ -525,10 +529,10 @@ class ParallelConfig: ) if self.enable_eplb: - if not current_platform.is_cuda_alike(): + if not current_platform.is_cuda_alike() and not current_platform.is_xpu(): raise ValueError( "Expert parallelism load balancing is only supported on " - "CUDA devices or ROCm devices now." + "CUDA devices or ROCm devices or XPU devices now." ) if not self.enable_expert_parallel: raise ValueError("enable_expert_parallel must be True to use EPLB.") @@ -1034,7 +1038,10 @@ class ParallelConfig: # See https://github.com/pytorch/pytorch/issues/174288 from vllm.distributed.nixl_utils import is_nixl_available - if is_nixl_available(): + if current_platform.is_xpu(): + # On XPU, use the device-native XCCL P2P backend. + self.eplb_config.communicator = "torch_xccl" + elif is_nixl_available(): self.eplb_config.communicator = "nixl" elif self.enable_elastic_ep: self.eplb_config.communicator = "pynccl" diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index caa3e3386c..0c8ebf3792 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os import threading from collections.abc import Iterator from contextlib import contextmanager @@ -1048,6 +1049,9 @@ class DeepEPV2All2AllManager(All2AllManagerBase): # missing GIN support. probe = torch.zeros(1, device="cuda") torch.distributed.all_reduce(probe, group=group) + # DeepEPv2 respects EP_DISABLE_GIN, so skip the GIN requirement check. + if os.environ.get("EP_DISABLE_GIN", "0") != "0": + return gin_type = query_nccl_gin_type(group) if gin_type is None: diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index 716f443af2..c56d2cc212 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -226,8 +226,7 @@ class ECConnectorBase(ABC): return None def get_ec_connector_stats(self) -> "ECConnectorStats | None": - """ - Get the EC connector stats collected during the last interval. + """Get the EC connector stats collected during the last interval. Callable on either a worker-side or scheduler-side connector instance: the worker-side instance reports stats gathered during @@ -358,8 +357,7 @@ class ECConnectorBase(ABC): def build_ec_connector_stats( cls, data: dict[str, Any] | None = None ) -> "ECConnectorStats | None": - """ - ECConnectorStats resolution method. This method allows dynamically + """ECConnectorStats resolution method. This method allows dynamically registered connectors to return their own ECConnectorStats object, which can implement custom aggregation logic on the data dict. """ @@ -373,8 +371,7 @@ class ECConnectorBase(ABC): labelnames: list[str], per_engine_labelvalues: dict[int, list[object]], ) -> "ECConnectorPromMetrics | None": - """ - Create an ECConnectorPromMetrics subclass which should register + """Create an ECConnectorPromMetrics subclass which should register per-connector Prometheus metrics and implement observe() to expose connector transfer stats via Prometheus. """ diff --git a/vllm/distributed/ec_transfer/ec_connector/metrics.py b/vllm/distributed/ec_transfer/ec_connector/metrics.py index 4a329096ce..0eb20b9863 100644 --- a/vllm/distributed/ec_transfer/ec_connector/metrics.py +++ b/vllm/distributed/ec_transfer/ec_connector/metrics.py @@ -17,8 +17,7 @@ logger = init_logger(__name__) @dataclass class ECConnectorStats: - """ - Base class for EC Connector Stats, a container for transfer performance + """Base class for EC Connector Stats, a container for transfer performance metrics or otherwise important telemetry from the connector. All sub-classes need to be serializable as stats are sent from worker to logger process. @@ -31,14 +30,11 @@ class ECConnectorStats: raise NotImplementedError def aggregate(self, other: "ECConnectorStats") -> "ECConnectorStats": - """ - Aggregate stats with another `ECConnectorStats` object. - """ + """Aggregate stats with another `ECConnectorStats` object.""" raise NotImplementedError def reduce(self) -> dict[str, int | float]: - """ - Reduce the observations collected during a time interval to one or + """Reduce the observations collected during a time interval to one or more representative values (eg avg/median/sum of the series). This is meant to be called by the logger to produce a summary of the stats for the last time interval. @@ -107,8 +103,7 @@ class ECConnectorLogging: class ECConnectorPromMetrics: - """ - A base class for per-connector Prometheus metric registration + """A base class for per-connector Prometheus metric registration and recording. """ @@ -127,8 +122,7 @@ class ECConnectorPromMetrics: self.per_engine_labelvalues = per_engine_labelvalues def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): - """ - Record the supplied transfer statistics to Prometheus metrics. These + """Record the supplied transfer statistics to Prometheus metrics. These statistics are engine-specific, and should be recorded to a metric with the appropriate 'engine' label. These metric instances can be created using the create_metric_per_engine() helper method. @@ -137,8 +131,7 @@ class ECConnectorPromMetrics: class ECConnectorProm: - """ - Support for registering per-connector Prometheus metrics, and + """Support for registering per-connector Prometheus metrics, and recording transfer statistics to those metrics. Uses ECConnectorBase.build_prom_metrics(). """ diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py index 10c9acae46..4cc103a56e 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py @@ -69,6 +69,10 @@ class MooncakeECConfig: control_timeout_ms: int push_wait_timeout_s: float pool_size: int + cross_encoder_cache: bool = False + store_max_pending_items: int = 32 + store_max_pending_bytes: int = 2 * 1024**3 + store_read_buffer_bytes: int = 128 * 1024**2 @classmethod def from_vllm_config(cls, vllm_config: VllmConfig) -> MooncakeECConfig: @@ -94,6 +98,24 @@ class MooncakeECConfig: "ec_buffer_size", ec_config.ec_buffer_size ) get = ec_config.get_from_extra_config + shared_reuse = get("cross_encoder_cache", False) + if not isinstance(shared_reuse, bool): + raise ValueError("cross_encoder_cache must be a boolean") + if shared_reuse: + if not ec_config.is_encode_only: + raise ValueError( + "cross_encoder_cache requires an Encoder-only producer" + ) + if not vllm_config.use_v2_model_runner: + raise ValueError("cross_encoder_cache requires Model Runner V2") + if vllm_config.lora_config is not None: + raise ValueError("cross_encoder_cache does not support LoRA") + mm_config = vllm_config.model_config.multimodal_config + if mm_config is not None and mm_config.mm_processor_cache_gb == 0: + raise ValueError( + "cross_encoder_cache requires mm_processor_cache_gb > 0 " + "to preserve content identifiers across Encoders and restarts" + ) control_port = int(ec_config.ec_port) + ( parallel_config.data_parallel_index * parallel_config.tensor_parallel_size ) @@ -120,4 +142,14 @@ class MooncakeECConfig: "push_wait_timeout_s", get("push_wait_timeout_s", 60) ), pool_size=registered_buffer_size, + cross_encoder_cache=shared_reuse, + store_max_pending_items=_positive_int( + "store_max_pending_items", get("store_max_pending_items", 32) + ), + store_max_pending_bytes=_positive_int( + "store_max_pending_bytes", get("store_max_pending_bytes", 2 * 1024**3) + ), + store_read_buffer_bytes=_positive_int( + "store_read_buffer_bytes", get("store_read_buffer_bytes", 128 * 1024**2) + ), ) diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py index a006d1bf2f..76e714dc04 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py @@ -10,6 +10,9 @@ from vllm.distributed.ec_transfer.ec_connector.base import ( ECConnectorMetadata, ECConnectorWorkerMetadata, ) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( + TensorSpec, +) @dataclass @@ -46,6 +49,7 @@ class ECMooncakeConnectorMetadata(ECConnectorMetadata): loads: list[ECMooncakeLoadSpec] = field(default_factory=list) pushes: list[ECMooncakePushSpec] = field(default_factory=list) freed: list[str] | None = None + store_candidates: dict[str, TensorSpec] = field(default_factory=dict) @dataclass diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py index 0b52dbe9c3..2a607118db 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py @@ -45,6 +45,9 @@ from vllm.distributed.ec_transfer.ec_connector.mooncake.state import ( from vllm.distributed.ec_transfer.ec_connector.mooncake.transfer import ( ensure_mooncake_available, ) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( + TensorSpec, +) from vllm.distributed.ec_transfer.ec_connector.utils import ( PlaceholderMetadataResolver, collect_ec_item_metadata, @@ -76,6 +79,8 @@ class ECMooncakeScheduler: ensure_mooncake_available() config = MooncakeECConfig.from_vllm_config(vllm_config) + self._store_candidates: list[tuple[str, str, TensorSpec]] = [] + self._cross_encoder_cache = config.cross_encoder_cache self._is_producer = config.is_producer self._is_consumer = config.is_consumer self._control_addr = config.control_addr @@ -427,28 +432,40 @@ class ECMooncakeScheduler: transfer_id = f"{request.request_id}:{index}" if not consumer_zmq or transfer_id in self._prepared_push_transfer_ids: return - num_tokens = request.get_num_encoder_embeds(index) - dtype = self._model_config.dtype - assert isinstance(dtype, torch.dtype) - assert self._encoder_cache_hidden_dim is not None - dtype_name = str(dtype).split(".")[-1] - shape = (num_tokens, self._encoder_cache_hidden_dim) - nbytes = math.prod(shape) * dtype.itemsize + output = self._encoder_output_spec(request, index) self._pushes_to_prepare[transfer_id] = ECMooncakePushSpec( mm_hash=mm_hash, - nbytes=nbytes, - shape=shape, - dtype=dtype_name, + nbytes=output.nbytes, + shape=output.shape, + dtype=output.dtype.removeprefix("torch."), consumer_zmq=str(consumer_zmq), transfer_id=transfer_id, request_id=request.request_id, ) self._prepared_push_transfer_ids.add(transfer_id) + def _encoder_output_spec(self, request: Any, index: int) -> TensorSpec: + dtype = self._model_config.dtype + assert isinstance(dtype, torch.dtype) + assert self._encoder_cache_hidden_dim is not None + shape = (request.get_num_encoder_embeds(index), self._encoder_cache_hidden_dim) + return TensorSpec(shape, str(dtype), math.prod(shape) * dtype.itemsize) + def update_state_after_alloc(self, request: Any, index: int) -> None: self._local_cache.add(request.mm_features[index].identifier) if self._is_producer: self._prepare_push_spec(request, index) + if ( + self._cross_encoder_cache + and request.mm_features[index].modality == "image" + ): + self._store_candidates.append( + ( + request.request_id, + request.mm_features[index].identifier, + self._encoder_output_spec(request, index), + ) + ) def update_state_after_free(self, request: Any, index: int) -> None: if not self._is_consumer: @@ -474,6 +491,13 @@ class ECMooncakeScheduler: meta = ECMooncakeConnectorMetadata( freed=scheduler_output.free_encoder_mm_hashes ) + preempted = scheduler_output.preempted_req_ids or set() + meta.store_candidates = { + identifier: output + for request_id, identifier, output in self._store_candidates + if request_id not in preempted + } + self._store_candidates.clear() for push_spec in self._pushes_to_prepare.values(): meta.pushes.append(push_spec) self._pushes_to_prepare.clear() diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py index 05ce16b181..04c95f2854 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py @@ -49,6 +49,9 @@ from vllm.distributed.ec_transfer.ec_connector.mooncake.transfer import ( MooncakeTransfer, ensure_mooncake_available, ) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( + build_embedding_namespace, +) from vllm.logger import init_logger from vllm.utils.network_utils import get_ip @@ -59,6 +62,10 @@ _T = TypeVar("_T") if TYPE_CHECKING: from vllm.config import VllmConfig + from ..mooncake_store_embedding.backend import ( + MooncakeEmbeddingStoreBackend, + ) + _RESERVATION_REFRESH_SECONDS = _RESERVATION_TTL_SECONDS / 2 _MAX_CANCELLED_TRANSFER_IDS = 1 << 16 _CANCEL_ATTEMPTS = 2 @@ -87,6 +94,17 @@ class ECMooncakeWorker: def __init__(self, vllm_config: VllmConfig) -> None: ensure_mooncake_available() config = MooncakeECConfig.from_vllm_config(vllm_config) + self._store_config = ( + ( + build_embedding_namespace(vllm_config), + config.store_max_pending_items, + config.store_max_pending_bytes, + config.store_read_buffer_bytes, + ) + if config.cross_encoder_cache + else None + ) + self._output_store: MooncakeEmbeddingStoreBackend | None = None self.is_producer = config.is_producer self.is_consumer = config.is_consumer self._buffer_device = config.buffer_device @@ -173,6 +191,21 @@ class ECMooncakeWorker: self._is_receiving_rank = True def start_services(self) -> None: + if self._store_config is not None and self._output_store is None: + from ..mooncake_store_embedding.backend import ( + MooncakeEmbeddingStoreBackend, + ) + from ..mooncake_store_embedding.store_client import ( + create_mooncake_embedding_store_client, + ) + + namespace, max_items, max_bytes, read_bytes = self._store_config + self._output_store = MooncakeEmbeddingStoreBackend( + create_mooncake_embedding_store_client(read_buffer_bytes=read_bytes), + namespace, + max_pending_items=max_items, + max_pending_bytes=max_bytes, + ) if not self.is_consumer or self._control_server is not None: return self._resolve_consumer_rank() @@ -470,6 +503,14 @@ class ECMooncakeWorker: tensor = encoder_cache.get(mm_hash) if tensor is not None: self._bind_push_source(tensor, mm_hash) + if self._output_store is not None: + loaded = self._output_store.resolve_inputs( + metadata.store_candidates, encoder_cache, self._buffer_device + ) + # Loaded entries precede the runner's new-output snapshot, so they + # need explicit P2P source binding here. + for mm_hash in loaded: + self._bind_push_source(encoder_cache[mm_hash], mm_hash) def _reserve_batch(self, records: list[ProducerPushRecord]) -> None: """Resolve independent item futures with one reserve RPC per shard.""" @@ -851,6 +892,8 @@ class ECMooncakeWorker: if not self.is_producer: return None, None + if self._output_store is not None: + self._output_store.reap() for record in self._producer_pushes.cancel_requests(finished_req_ids): self._producer_pushes.submit_cancel( record, @@ -866,6 +909,8 @@ class ECMooncakeWorker: return tensor = encoder_cache[mm_hash] self._bind_push_source(tensor, mm_hash) + if self._output_store is not None: + self._output_store.save_output(mm_hash, tensor) def build_connector_worker_meta(self) -> ECMooncakeWorkerMetadata | None: if self.is_consumer and not self._is_receiving_rank: @@ -938,3 +983,6 @@ class ECMooncakeWorker: ) self._producer_memory.close() self._transfer.close() + if self._output_store is not None: + self._output_store.shutdown() + self._output_store = None diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/__init__.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/__init__.py new file mode 100644 index 0000000000..42ed7690a7 --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Mooncake Store support for cross-Encoder output reuse.""" diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/backend.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/backend.py new file mode 100644 index 0000000000..5ad5ada882 --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/backend.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Encoder-side Mooncake Store resolver and publisher.""" + +from __future__ import annotations + +import threading +import traceback +from collections.abc import Mapping +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass + +import torch + +from vllm.logger import init_logger + +from .data import TensorSpec, make_embedding_key +from .store_client import ( + EmbeddingStoreError, + EmbeddingStoreOperationError, + MooncakeEmbeddingStoreClient, +) + +logger = init_logger(__name__) + + +@dataclass +class _StoreSave: + tensor: torch.Tensor | None + ready_event: torch.Event | None + storage_ptr: int + + def release(self) -> None: + self.tensor = None + self.ready_event = None + + +class MooncakeEmbeddingStoreBackend: + """Resolve immutable encoder outputs and publish misses asynchronously.""" + + def __init__( + self, + store_client: MooncakeEmbeddingStoreClient, + namespace: str, + *, + max_pending_items: int, + max_pending_bytes: int, + ) -> None: + self.store_client = store_client + self.namespace = namespace + self.max_pending_items = max_pending_items + self.max_pending_bytes = max_pending_bytes + self._executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="ec-mooncake-store", + ) + self._pending: dict[str, tuple[Future[None], _StoreSave]] = {} + self._pending_storages: dict[int, tuple[torch.UntypedStorage, int]] = {} + self._pending_lock = threading.Lock() + self._pending_bytes = 0 + self._step_candidates: set[str] = set() + self._closed = False + self._fatal_error: BaseException | None = None + + def save_output(self, identifier: str, tensor: torch.Tensor) -> None: + if identifier in self._step_candidates: + self.reap() + self._enqueue_save(identifier, tensor) + + def resolve_inputs( + self, + expected_tensors: Mapping[str, TensorSpec], + encoder_cache: dict[str, torch.Tensor], + device: torch.device | str, + ) -> set[str]: + """Load outputs before encoding; unresolved items retain normal computation.""" + self.reap() + self._step_candidates = set(expected_tensors) + candidates = [key for key in expected_tensors if key not in encoder_cache] + if not candidates: + return set() + + pool_keys = [ + make_embedding_key(self.namespace, identifier) for identifier in candidates + ] + try: + exists = self.store_client.batch_exists(pool_keys) + except (EmbeddingStoreOperationError, OSError): + logger.warning( + "Mooncake embedding Store lookup failed; falling back to encoder", + exc_info=True, + ) + return set() + + hits = { + pool_key: expected_tensors[identifier] + for identifier, pool_key, hit in zip( + candidates, pool_keys, exists, strict=True + ) + if hit + } + try: + tensors = self.store_client.load_tensors(hits, device) + except (EmbeddingStoreOperationError, OSError): + logger.warning( + "Mooncake embedding Store GET failed; falling back to encoder", + exc_info=True, + ) + return set() + loaded: set[str] = set() + for identifier, pool_key in zip(candidates, pool_keys, strict=True): + if pool_key in tensors: + encoder_cache[identifier] = tensors[pool_key] + loaded.add(identifier) + return loaded + + def _enqueue_save(self, identifier: str, tensor: torch.Tensor) -> bool: + """Admit contiguous outputs, charging each retained storage once.""" + if not tensor.is_contiguous(): + return False + storage = tensor.untyped_storage() + storage_ptr = storage.data_ptr() + with self._pending_lock: + if self._fatal_error is not None: + raise EmbeddingStoreError( + "Store publisher has failed" + ) from self._fatal_error + if self._closed or identifier in self._pending: + return False + _, ref_count = self._pending_storages.get(storage_ptr, (storage, 0)) + budget_bytes = 0 if ref_count else storage.nbytes() + if ( + len(self._pending) >= self.max_pending_items + or self._pending_bytes + budget_bytes > self.max_pending_bytes + ): + return False + # Only admitted work creates an event. Runtime errors propagate. + save = _StoreSave( + tensor, + _record_tensor_ready_event(tensor), + storage_ptr, + ) + future = self._executor.submit( + self._save, make_embedding_key(self.namespace, identifier), save + ) + self._pending[identifier] = (future, save) + # Keep the storage alive until reaping so its address cannot be reused. + self._pending_storages[storage_ptr] = (storage, ref_count + 1) + self._pending_bytes += budget_bytes + return True + + def _save(self, pool_key: str, save: _StoreSave) -> None: + try: + if self.store_client.exists(pool_key): + return + tensor = save.tensor + assert tensor is not None + if save.ready_event is not None: + save.ready_event.synchronize() + self.store_client.put_tensor(pool_key, tensor) + except BaseException as error: + # Returned client frames can otherwise keep tensors via a Future's + # traceback. Uncertain I/O owners remain retained by the client. + traceback.clear_frames(error.__traceback__) + raise + finally: + tensor = None + save.release() + + def reap(self) -> None: + with self._pending_lock: + for identifier, (future, save) in list(self._pending.items()): + # Observe once; a racing completion is handled on the next poll. + if not future.done(): + continue + try: + future.result() # Already done: never waits on Store I/O. + except (EmbeddingStoreOperationError, OSError) as error: + logger.warning("Store PUT failed for %s: %s", identifier, error) + # This failure is consumed; do not retain publisher/caller frames. + error.__traceback__ = None + except BaseException as error: + # The client may still own unsafe native-I/O buffers. Keep + # their charge and fatal result until worker termination. + if self._fatal_error is None: + self._fatal_error = error + continue + del self._pending[identifier] + storage, ref_count = self._pending_storages[save.storage_ptr] + if ref_count == 1: + del self._pending_storages[save.storage_ptr] + self._pending_bytes -= storage.nbytes() + else: + self._pending_storages[save.storage_ptr] = (storage, ref_count - 1) + fatal = self._fatal_error + if fatal is not None: + raise fatal + + def shutdown(self) -> None: + with self._pending_lock: + if self._closed: + return + self._closed = True + self._executor.shutdown(wait=True, cancel_futures=False) + self.reap() + self.store_client.close() + + +def _record_tensor_ready_event(tensor: torch.Tensor) -> torch.Event | None: + if tensor.device.type != "cuda": + return None + event = torch.Event() + event.record(torch.accelerator.current_stream(tensor.device)) + return event diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/data.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/data.py new file mode 100644 index 0000000000..9a7d7dd002 --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/data.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Content identities and output contracts for the shared Encoder Store.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +EMBEDDING_CACHE_KEY_VERSION = "v3" + + +def make_embedding_key(namespace: str, identifier: str) -> str: + return f"{namespace}@id:{identifier}" + + +@dataclass(frozen=True) +class TensorSpec: + """Contiguous output contract used by the Encoder and Store codec.""" + + shape: tuple[int, ...] + dtype: str + nbytes: int + + +def build_embedding_namespace(vllm_config: VllmConfig) -> str: + model_config = vllm_config.model_config + ec_config = vllm_config.ec_transfer_config + assert ec_config is not None + extra_config = ec_config.ec_connector_extra_config + + multimodal_config = model_config.multimodal_config + encoder_hash = ( + multimodal_config.compute_hash() + if multimodal_config is not None + else "encoder:default" + ) + cache_prefix = str(extra_config.get("embedding_cache_prefix", "")) + model_name = str(extra_config.get("embedding_model_identity", model_config.model)) + prefix = f"{cache_prefix}@" if cache_prefix else "" + return ( + f"{prefix}embedding" + f"@model:{model_name}" + f"@revision:{model_config.revision or 'default'}" + f"@encoder:{encoder_hash}" + f"@dtype:{model_config.dtype}" + f"@protocol:{EMBEDDING_CACHE_KEY_VERSION}" + ) diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/store_client.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/store_client.py new file mode 100644 index 0000000000..4d3ed6739d --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/store_client.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Thin Mooncake Store client for embedding objects. + +Originally adapted from vLLM PR #47302 (0d1f71f5d36c). The v3 key namespace +uses a compact header for complete 2D embeddings, not that PR's tensor format. +""" + +from __future__ import annotations + +import ctypes +import struct +import threading +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from enum import IntEnum +from typing import Any + +import torch + +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( + TensorSpec, +) +from vllm.distributed.mooncake_store import MooncakeStoreConfig, setup_mooncake_store +from vllm.logger import init_logger +from vllm.utils.network_utils import get_ip + +logger = init_logger(__name__) + +# A failed transfer can still use registered memory after the binding returns. +# Poisoned clients retain their owners until process exit, including if normal +# worker teardown drops the connector. New operations are rejected. +_UNSAFE_CLIENTS: list[Any] = [] + + +# Values from Mooncake v0.3.12 / v0.3.12.post1 types.h (ErrorCode). +# Python exposes integer returns, not this enum. +class _MooncakeErrorCode(IntEnum): + NO_AVAILABLE_HANDLE = -200 + INVALID_PARAMS = -600 + INVALID_REPLICA = -702 + REPLICA_IS_NOT_READY = -703 + OBJECT_NOT_FOUND = -704 + OBJECT_ALREADY_EXISTS = -705 + LEASE_EXPIRED = -707 + RPC_FAIL = -900 + RPC_TIMEOUT = -901 + TENANT_QUOTA_EXCEEDED = -1700 + + +# For RAM batch reads and single-replica writes, these rejections occur +# before submission or after completed I/O. Transfer timeouts can leave I/O +# in flight; all other failures retain registered buffers and poison the client. +_SAFE_IO_REJECTIONS = frozenset( + { + _MooncakeErrorCode.NO_AVAILABLE_HANDLE, + _MooncakeErrorCode.REPLICA_IS_NOT_READY, + _MooncakeErrorCode.OBJECT_NOT_FOUND, + _MooncakeErrorCode.OBJECT_ALREADY_EXISTS, + _MooncakeErrorCode.RPC_FAIL, + _MooncakeErrorCode.RPC_TIMEOUT, + } +) +# Batch GET validates capacity/replicas before I/O and checks leases afterwards. +# PUT rejects tenant quotas before allocating replicas or submitting transfers. +_SAFE_GET_REJECTIONS = _SAFE_IO_REJECTIONS | { + _MooncakeErrorCode.INVALID_PARAMS, + _MooncakeErrorCode.INVALID_REPLICA, + _MooncakeErrorCode.LEASE_EXPIRED, +} +_SAFE_PUT_REJECTIONS = _SAFE_IO_REJECTIONS | {_MooncakeErrorCode.TENANT_QUOTA_EXCEEDED} + + +_MOONCAKE_TENSOR_OBJECT_MAGIC = 0x4D4F4F4E +# Magic, dtype, token count, hidden dimension. Format versioning is in the key. +_MOONCAKE_TENSOR_HEADER = struct.Struct(" MooncakeEmbeddingStoreClient: + try: + from mooncake.store import ( # type: ignore + MooncakeDistributedStore, + ObjectDataType, + ReplicateConfig, + ) + except ImportError as e: + raise ImportError( + "Install mooncake with Store support to enable cross_encoder_cache." + ) from e + + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils + + config = MooncakeStoreConfig.load_from_config() + if config.enable_offload: + raise ValueError( + "cross_encoder_cache supports a RAM Store; disable enable_offload" + ) + store = MooncakeDistributedStore() + local_ip = get_ip() + local_hostname = rdma_utils.get_requester_local_hostname(local_ip) + setup_mooncake_store(store, config, local_hostname) + + logger.info( + "Initialized embedding Mooncake store mode=%s global_segment_size=%d " + "local_buffer_size=%d", + config.mode, + config.global_segment_size, + config.local_buffer_size, + ) + replicate_config = ReplicateConfig() + replicate_config.data_type = ObjectDataType.TENSOR + return MooncakeEmbeddingStoreClient( + store, replicate_config=replicate_config, read_buffer_bytes=read_buffer_bytes + ) + + +class EmbeddingStoreError(RuntimeError): + """A Store failure that must propagate to the worker.""" + + +class EmbeddingStoreOperationError(EmbeddingStoreError): + """A rejected or completed operation that permits normal Encoder fallback.""" + + +def _batch_get_completed_safely(results: Any, capacities: list[int]) -> bool: + return ( + isinstance(results, list) + and len(results) == len(capacities) + and all( + type(result) is int + and (0 <= result <= capacity or result in _SAFE_GET_REJECTIONS) + for result, capacity in zip(results, capacities, strict=True) + ) + ) + + +class MooncakeEmbeddingStoreClient: + """Wraps Mooncake object and buffer APIs used by embedding transfer.""" + + def __init__( + self, + store: Any, + replicate_config: Any | None = None, + *, + read_buffer_bytes: int = 128 * 1024**2, + ): + self._lifetime_lock = threading.Lock() + self._unsafe_owners: list[Any] = [] + self._poisoned = False + self.store = store + self.replicate_config = replicate_config + self._read_buffer_bytes = read_buffer_bytes + self._read_buffer: torch.Tensor | None = None + self._put_header: ctypes.Array[ctypes.c_char] | None = None + + def _check_healthy(self) -> None: + if self._poisoned: + raise EmbeddingStoreError( + "Mooncake I/O completion is unconfirmed; worker must exit" + ) + + def _poison(self, owners: list[Any]) -> None: + with self._lifetime_lock: + self._unsafe_owners.extend(owners) + if not self._poisoned: + self._poisoned = True + _UNSAFE_CLIENTS.append(self) + + @contextmanager + def _registered_io( + self, + buffers: list[tuple[Any, int, int]], + ) -> Iterator[Callable[[int], None]]: + self._check_healthy() + registered = [] + submitted = False + confirmed = False + + def confirm(result: int) -> None: + nonlocal confirmed + if type(result) is not int or ( + result < 0 and result not in _SAFE_PUT_REJECTIONS + ): + raise EmbeddingStoreError( + "Mooncake I/O completion is unconfirmed; retaining buffers" + ) + confirmed = True + + try: + for owner, addr, size in buffers: + ret = self.store.register_buffer(addr, size) + if ret != 0: + raise EmbeddingStoreOperationError( + f"Failed to register embedding buffer: {ret}" + ) + registered.append(addr) + submitted = True + yield confirm + except BaseException as error: + if submitted and not confirmed: + self._poison([owner for owner, _, _ in buffers]) + raise EmbeddingStoreError( + "Mooncake I/O completion is unconfirmed; retaining buffers" + ) from error + raise + finally: + if not submitted or confirmed: + try: + for addr in reversed(registered): + self._unregister_buffer(addr) + except BaseException: + self._poison([owner for owner, _, _ in buffers]) + raise + + def close(self) -> None: + """Close the supported Mooncake binding after all I/O is drained.""" + self._check_healthy() + if self._read_buffer is not None: + try: + self._unregister_buffer(self._read_buffer.data_ptr()) + except BaseException: + self._poison([self._read_buffer]) + raise + self._read_buffer = None + if self._put_header is not None: + try: + self._unregister_buffer(ctypes.addressof(self._put_header)) + except BaseException: + self._poison([self._put_header]) + raise + self._put_header = None + ret = self.store.close() + if ret != 0: + raise EmbeddingStoreError( + f"failed to close embedding Mooncake Store: {ret}" + ) + + def exists(self, pool_key: str) -> bool: + return self.batch_exists([pool_key])[0] + + def batch_exists(self, keys: list[str]) -> list[bool]: + self._check_healthy() + if not keys: + return [] + + states = self.store.batch_is_exist(keys) + if len(states) != len(keys): + raise RuntimeError( + "Mooncake Store returned an unexpected number of lookup results" + ) + return [state == 1 for state in states] + + def load_tensors( + self, expected: Mapping[str, TensorSpec], device: torch.device | str + ) -> dict[str, torch.Tensor]: + """Synchronously load runner inputs into independently owned tensors.""" + self._check_healthy() + device = torch.device(device) + loaded: dict[str, torch.Tensor] = {} + chunk: dict[str, TensorSpec] = {} + chunk_bytes = 0 + for key, spec in expected.items(): + size = _MOONCAKE_TENSOR_HEADER.size + spec.nbytes + if size > self._read_buffer_bytes: + logger.warning( + "Skipping Store GET for %s: object bytes=%d exceed read buffer=%d", + key, + size, + self._read_buffer_bytes, + ) + continue + if chunk_bytes + size > self._read_buffer_bytes: + loaded.update(self._load_chunk(chunk, device)) + chunk = {} + chunk_bytes = 0 + chunk[key] = spec + chunk_bytes += size + if chunk: + loaded.update(self._load_chunk(chunk, device)) + return loaded + + def _load_chunk( + self, expected: dict[str, TensorSpec], device: torch.device + ) -> dict[str, torch.Tensor]: + self._check_healthy() + if self._read_buffer is None: + buffer = torch.empty( + self._read_buffer_bytes, + dtype=torch.uint8, + device="cpu", + pin_memory=device.type == "cuda", + ) + ret = self.store.register_buffer(buffer.data_ptr(), buffer.nbytes) + if ret != 0: + raise EmbeddingStoreOperationError( + f"Failed to register embedding read buffer: {ret}" + ) + self._read_buffer = buffer + buffer = self._read_buffer + base = buffer.data_ptr() + sizes = [ + _MOONCAKE_TENSOR_HEADER.size + spec.nbytes for spec in expected.values() + ] + offsets: list[int] = [] + end = 0 + for size in sizes: + offsets.append(end) + end += size + try: + results = self.store.batch_get_into( + list(expected), [base + offset for offset in offsets], sizes + ) + if not _batch_get_completed_safely(results, sizes): + raise EmbeddingStoreError("Unsafe batch GET results") + except BaseException as error: + self._poison([buffer]) + raise EmbeddingStoreError( + "Mooncake I/O completion is unconfirmed; retaining read buffer" + ) from error + + valid: list[tuple[str, TensorSpec, int]] = [] + for (key, spec), offset, size, result in zip( + expected.items(), offsets, sizes, results, strict=True + ): + if result != size: + logger.warning( + "Store GET skipped for %s: result=%d, expected bytes=%d", + key, + result, + size, + ) + continue + try: + _validate_mooncake_tensor_metadata( + key, + ctypes.string_at(base + offset, _MOONCAKE_TENSOR_HEADER.size), + spec, + ) + except EmbeddingStoreOperationError as error: + logger.warning("Store GET skipped for %s: %s", key, error) + continue + valid.append((key, spec, offset + _MOONCAKE_TENSOR_HEADER.size)) + if not valid: + return {} + # Allocate every destination before submitting any asynchronous copy. + targets = { + key: torch.empty( + spec.shape, + dtype=getattr(torch, spec.dtype.removeprefix("torch.")), + device=device, + ) + for key, spec, _ in valid + } + ready = torch.Event() if device.type == "cuda" else None + try: + for key, spec, offset in valid: + targets[key].view(torch.uint8).view(-1).copy_( + buffer[offset : offset + spec.nbytes], + non_blocking=ready is not None, + ) + if ready is not None: + ready.record(torch.accelerator.current_stream(device)) + ready.synchronize() + except BaseException as error: + if ready is not None: + self._poison([buffer, *targets.values(), ready]) + raise EmbeddingStoreError( + "Embedding copy completion is unconfirmed; retaining buffers" + ) from error + raise + return targets + + def put_tensor(self, pool_key: str, tensor: torch.Tensor) -> None: + self._check_healthy() + if not tensor.is_contiguous(): + raise EmbeddingStoreOperationError("embedding tensor must be contiguous") + metadata = _encode_mooncake_tensor_metadata(tensor) + # The single publisher reuses this header only after PUT completion. + if self._put_header is None: + header = ctypes.create_string_buffer(len(metadata)) + ret = self.store.register_buffer(ctypes.addressof(header), len(metadata)) + if ret != 0: + raise EmbeddingStoreOperationError( + f"Failed to register embedding header: {ret}" + ) + self._put_header = header + header_ptr = ctypes.addressof(self._put_header) + ctypes.memmove(header_ptr, metadata, len(metadata)) + data_size = tensor.numel() * tensor.element_size() + buffers = [ + (tensor, tensor.data_ptr(), data_size), + ] + with self._registered_io(buffers) as confirm: + [result] = self.store.batch_put_from_multi_buffers( + [pool_key], + [[header_ptr, tensor.data_ptr()]], + [[len(metadata), data_size]], + self.replicate_config, + ) + confirm(result) + if result < 0: + raise EmbeddingStoreOperationError( + f"failed to put embedding tensor for {pool_key}" + ) + + def _unregister_buffer(self, addr: int) -> None: + try: + ret = self.store.unregister_buffer(addr) + except Exception as error: + raise EmbeddingStoreError( + f"could not confirm buffer unregistration addr={addr:#x}" + ) from error + if ret != 0: + raise EmbeddingStoreError( + f"failed to unregister embedding buffer addr={addr:#x}: {ret}" + ) + + +def _validate_mooncake_tensor_metadata( + pool_key: str, + metadata: bytes, + expected: TensorSpec, +) -> None: + magic, dtype, num_tokens, hidden_dim = _MOONCAKE_TENSOR_HEADER.unpack(metadata) + if magic != _MOONCAKE_TENSOR_OBJECT_MAGIC: + raise EmbeddingStoreOperationError( + f"invalid Mooncake tensor metadata header for {pool_key}" + ) + if _MOONCAKE_DTYPE_TO_TORCH_DTYPE.get(dtype) != expected.dtype: + raise EmbeddingStoreOperationError( + f"unexpected Mooncake tensor dtype for {pool_key}: {dtype}" + ) + if (num_tokens, hidden_dim) != expected.shape: + raise EmbeddingStoreOperationError( + f"embedding tensor shape mismatch for {pool_key}: expected={expected.shape}" + ) + + +def _encode_mooncake_tensor_metadata(tensor: torch.Tensor) -> bytes: + dtype = str(tensor.dtype) + if dtype not in _TORCH_DTYPE_TO_MOONCAKE_DTYPE: + raise EmbeddingStoreOperationError( + f"unsupported embedding tensor dtype: {dtype}" + ) + if tensor.ndim != 2: + raise EmbeddingStoreOperationError( + f"embedding tensor must be 2D: shape={tuple(tensor.shape)}" + ) + return _MOONCAKE_TENSOR_HEADER.pack( + _MOONCAKE_TENSOR_OBJECT_MAGIC, + _TORCH_DTYPE_TO_MOONCAKE_DTYPE[dtype], + *tensor.shape, + ) diff --git a/vllm/distributed/eplb/async_worker.py b/vllm/distributed/eplb/async_worker.py index 5cf63fe5ae..3b7bc908ec 100644 --- a/vllm/distributed/eplb/async_worker.py +++ b/vllm/distributed/eplb/async_worker.py @@ -10,7 +10,7 @@ import torch from vllm.distributed.parallel_state import get_eplb_group from vllm.logger import init_logger -from .eplb_utils import CpuGpuEvent +from .eplb_utils import CpuGpuEvent, device_stream from .rebalance_execute import AsyncEplbLayerResult, transfer_layer if TYPE_CHECKING: @@ -24,17 +24,17 @@ def start_async_worker( is_profile: bool = False, ) -> threading.Thread: rank = get_eplb_group().device_group.rank() - device_index = state.cuda_device_index + device_index = state.device_index assert state.is_async def thread_target() -> None: assert device_index is not None torch.accelerator.set_device_index(device_index) - cuda_stream = torch.cuda.Stream(device=device_index) + stream = torch.Stream(device=torch.accelerator.current_accelerator()) try: transfer_run_periodically( state=state, - cuda_stream=cuda_stream, + stream=stream, is_profile=is_profile, ) except Exception as exc: # pragma: no cover - diagnostic path @@ -49,13 +49,13 @@ def run_rebalance_experts( model_state: "EplbModelState", eplb_state: "EplbState", physical_to_logical_map_cpu: torch.Tensor, - cuda_stream: torch.cuda.Stream, + stream: torch.Stream, ) -> torch.Tensor: assert model_state.eplb_stats is not None eplb_stats = model_state.eplb_stats # Move the global expert load window to CPU for computation. - with torch.cuda.stream(cuda_stream): + with device_stream(stream): global_expert_load_window = eplb_stats.global_expert_load_window.cpu() # Compute new expert mappings for the model new_physical_to_logical_map = eplb_state.policy.rebalance_experts( @@ -73,11 +73,11 @@ def run_rebalance_experts( def transfer_run_periodically( state: "EplbState", - cuda_stream: torch.cuda.Stream, + stream: torch.Stream, is_profile: bool = False, ) -> None: while True: - state.rearrange_event.wait(stream=cuda_stream) + state.rearrange_event.wait(stream=stream) eplb_group = get_eplb_group().device_group eplb_cpu_group = get_eplb_group().cpu_group @@ -87,16 +87,16 @@ def transfer_run_periodically( for model_state in state.model_states.values(): layer_idx = 0 # Set the async worker's CUDA stream on the communicator - model_state.communicator.set_stream(cuda_stream) + model_state.communicator.set_stream(stream) num_layers = model_state.model.num_moe_layers # Snapshot the physical_to_logical_map (synchronized with # rearrange_event) and copy it to CPU - with torch.cuda.stream(cuda_stream): + with device_stream(stream): physical_to_logical_map_cpu = model_state.physical_to_logical_map.cpu() new_physical_to_logical_map = run_rebalance_experts( - model_state, state, physical_to_logical_map_cpu, cuda_stream + model_state, state, physical_to_logical_map_cpu, stream ) # Execute one EPLB layer transfer per model forward pass. Each iteration @@ -132,13 +132,13 @@ def transfer_run_periodically( communicator=model_state.communicator, ep_group=eplb_group, is_profile=is_profile, - cuda_stream=cuda_stream, + stream=stream, layer_idx=layer_idx, ) # Wait until all writes to expert_buffer have finished before making the # AsyncEplbLayerResult visible to the main thread. - cuda_stream.synchronize() + stream.synchronize() # This event guarantees that expert_buffer will not be overwritten by # subsequent iterations of this loop until the main thread has consumed @@ -155,6 +155,6 @@ def transfer_run_periodically( # Block this thread until the main thread and main stream # finish copying model_state.expert_buffer into # model_state.model.expert_weights[layer_idx] - consumed_event.wait(stream=cuda_stream) + consumed_event.wait(stream=stream) assert model_state.pending_result is None layer_idx += 1 diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 61fef468d3..730c01a388 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -22,6 +22,7 @@ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.distributed.device_communicators.pynccl_wrapper import ( ncclDataTypeEnum, ) +from vllm.distributed.eplb.eplb_utils import device_stream from vllm.distributed.parallel_state import ( GroupCoordinator, get_pp_group, @@ -87,8 +88,8 @@ class EplbCommunicator(ABC): communication buffers.""" return True - def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None: - self._cuda_stream = cuda_stream + def set_stream(self, stream: torch.Stream | None) -> None: + self._stream = stream def _log_initialized(self) -> None: if is_local_first_rank(): @@ -101,10 +102,10 @@ class TorchDistNcclEplbCommunicator(EplbCommunicator): def __init__( self, ep_group: ProcessGroup, - cuda_stream: torch.cuda.Stream | None = None, + stream: torch.Stream | None = None, ) -> None: self._ep_group = ep_group - self._cuda_stream = cuda_stream + self._stream = stream self._p2p_ops: list[P2POp] = [] self._log_initialized() @@ -144,7 +145,7 @@ class TorchDistNcclEplbCommunicator(EplbCommunicator): if not self._p2p_ops: return try: - with torch.cuda.stream(self._cuda_stream): + with device_stream(self._stream): reqs = batch_isend_irecv(self._p2p_ops) for req in reqs: req.wait() @@ -158,10 +159,10 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator): def __init__( self, cpu_group: ProcessGroup, - cuda_stream: torch.cuda.Stream | None = None, + stream: torch.Stream | None = None, ) -> None: self._cpu_group = cpu_group - self._cuda_stream = cuda_stream + self._stream = stream self._ops: list[tuple[str, torch.Tensor, int]] = [] self._log_initialized() @@ -217,7 +218,7 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator): recv_staging.append((tensor, cpu_tensor)) try: - with torch.cuda.stream(self._cuda_stream): + with device_stream(self._stream): build_ops() finally: self._ops.clear() @@ -225,10 +226,10 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator): # Wait for all D2H copies to finish # before issuing gloo batch_isend_irecv operations. with gpu_sync_allowed(): - if self._cuda_stream is not None: - self._cuda_stream.synchronize() + if self._stream is not None: + self._stream.synchronize() else: - torch.cuda.current_stream().synchronize() + torch.accelerator.current_stream().synchronize() reqs = batch_isend_irecv(p2p_ops) for req in reqs: @@ -236,11 +237,69 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator): if not recv_staging: return - with torch.cuda.stream(self._cuda_stream): + with device_stream(self._stream): for dst_tensor, cpu_tensor in recv_staging: dst_tensor.copy_(cpu_tensor, non_blocking=True) +class TorchDistXCCLStagedEplbCommunicator(EplbCommunicator): + """EPLB communicator using XCCL device-to-device P2P on XPU.""" + + def __init__( + self, + ep_group: ProcessGroup, + stream: torch.Stream | None = None, + ) -> None: + self._ep_group = ep_group + self._stream = stream + self._ops: list[tuple[str, torch.Tensor, int]] = [] + self._log_initialized() + + def add_send( + self, + tensors: list[torch.Tensor], + dst_rank: int, + expert_id: int, # unused by this backend + ) -> None: + for tensor in tensors: + self._ops.append(("send", tensor, dst_rank)) + + def add_recv( + self, + tensors: list[torch.Tensor], + src_rank: int, + expert_id: int, # unused by this backend + ) -> None: + for tensor in tensors: + self._ops.append(("recv", tensor, src_rank)) + + def execute(self) -> None: + if not self._ops: + return + + p2p_ops: list[P2POp] = [] + try: + for op, tensor, peer_rank in self._ops: + send_or_recv = ( + torch.distributed.isend if op == "send" else torch.distributed.irecv + ) + p2p_ops.append( + P2POp( + send_or_recv, + tensor, + peer_rank, + self._ep_group, + ) + ) + finally: + self._ops.clear() + + with device_stream(self._stream): + reqs = batch_isend_irecv(p2p_ops) + for req in reqs: + req.wait() + + class NixlEplbCommunicator(EplbCommunicator): """EPLB communicator backed by NIXL READ transfers.""" @@ -343,7 +402,7 @@ class NixlEplbCommunicator(EplbCommunicator): uid = uuid.uuid4().hex[:8] return f"eplb-{self._rank}{pp_suffix}-{uid}" - def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None: + def set_stream(self, stream: torch.Stream | None) -> None: pass def add_send( @@ -602,10 +661,10 @@ class PyNcclEplbCommunicator(EplbCommunicator): def __init__( self, pynccl_comm: PyNcclCommunicator, - cuda_stream: torch.cuda.Stream | None = None, + stream: torch.Stream | None = None, ) -> None: self._pynccl_comm = pynccl_comm - self._cuda_stream = cuda_stream + self._stream = stream self._group_started = False self._log_initialized() @@ -622,7 +681,7 @@ class PyNcclEplbCommunicator(EplbCommunicator): ) -> None: self._ensure_group_started() for tensor in tensors: - self._pynccl_comm.send(tensor, dst_rank, stream=self._cuda_stream) + self._pynccl_comm.send(tensor, dst_rank, stream=self._stream) def add_recv( self, @@ -632,7 +691,7 @@ class PyNcclEplbCommunicator(EplbCommunicator): ) -> None: self._ensure_group_started() for tensor in tensors: - self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream) + self._pynccl_comm.recv(tensor, src_rank, stream=self._stream) def execute(self) -> None: if self._group_started: @@ -750,6 +809,8 @@ def create_eplb_communicator( return TorchDistGlooStagedEplbCommunicator( cpu_group=group_coordinator.cpu_group, ) + elif backend == "torch_xccl": + return TorchDistXCCLStagedEplbCommunicator(ep_group=torch_group) elif backend == "torch_nccl": return TorchDistNcclEplbCommunicator(ep_group=torch_group) elif backend == "pynccl": diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index fb5475a607..e4d805b9fe 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -203,9 +203,9 @@ class EplbModelState: """ EPLB stats for the model. """ - cuda_device_index: int | None + device_index: int | None """ - CUDA device index for the async EPLB worker thread. + Device index for the async EPLB worker thread. """ communicator: EplbCommunicator """ @@ -286,14 +286,14 @@ class EplbState: """ Background thread handling async transfers. """ - self.cuda_device_index: int | None = None + self.device_index: int | None = None """ - CUDA device index for the async EPLB worker thread. + Device index for the async EPLB worker thread. """ - if self.device.type == "cuda": - self.cuda_device_index = self.device.index - if self.cuda_device_index is None and torch.cuda.is_available(): - self.cuda_device_index = torch.accelerator.current_device_index() + if self.device.type in ("cuda", "xpu"): + self.device_index = self.device.index + if self.device_index is None and torch.accelerator.is_available(): + self.device_index = torch.accelerator.current_device_index() @staticmethod def build_initial_global_physical_to_logical_map( @@ -504,7 +504,7 @@ class EplbState: expert_buffer=expert_buffer, rebalanced=False, eplb_stats=None, - cuda_device_index=self.cuda_device_index, + device_index=self.device_index, communicator=communicator, num_unpadded_tokens_tensors=num_unpadded_tokens_tensors, ) @@ -763,8 +763,8 @@ class EplbState: is_main_rank = ep_rank == 0 if is_main_rank: if not self.is_async or is_profile: - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) + start_event = torch.Event(enable_timing=True) + end_event = torch.Event(enable_timing=True) start_event.record() logger.info( "Rearranging experts %s %s...", diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py index 3d5be8a8ec..943f6d9bef 100644 --- a/vllm/distributed/eplb/eplb_utils.py +++ b/vllm/distributed/eplb/eplb_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility functions for EPLB (Expert Parallel Load Balancing).""" +import contextlib import os import threading @@ -13,6 +14,22 @@ from vllm.logger import init_logger logger = init_logger(__name__) +@contextlib.contextmanager +def device_stream(stream: torch.Stream | None): + """Platform-agnostic context manager that activates *stream* as the + current accelerator stream for the duration of the ``with`` block. + A no-op when *stream* is ``None``.""" + if stream is None: + yield + return + prev = torch.accelerator.current_stream() + torch.accelerator.set_stream(stream) + try: + yield + finally: + torch.accelerator.set_stream(prev) + + class CpuGpuEvent: """Combines a CUDA event with a CPU threading event to enforce record->wait ordering across two threads. @@ -30,10 +47,10 @@ class CpuGpuEvent: """ def __init__(self): - self._event = torch.cuda.Event() + self._event = torch.Event() self._recorded = threading.Event() - def wait(self, stream: torch.cuda.Stream | None = None): + def wait(self, stream: torch.Stream | None = None): """Blocks the calling thread until record finishes. Used to guarantee that the record kernel is called before wait. @@ -43,7 +60,7 @@ class CpuGpuEvent: self._event.wait(stream) self._recorded.clear() - def record(self, stream: torch.cuda.Stream | None = None): + def record(self, stream: torch.Stream | None = None): """Unblocks the waiting thread after calling event.record(). Should only be called by the main thread. @@ -53,7 +70,7 @@ class CpuGpuEvent: "CpuGpuEvent.record() called before the previous event was " "consumed by wait()" ) - self._event = torch.cuda.Event() + self._event = torch.Event() self._event.record(stream) self._recorded.set() diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index 20ebe4f1f2..7a52dd5f56 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -13,7 +13,7 @@ import torch from torch.distributed import ProcessGroup, all_gather from vllm.distributed.eplb.eplb_communicator import EplbCommunicator -from vllm.distributed.eplb.eplb_utils import CpuGpuEvent +from vllm.distributed.eplb.eplb_utils import CpuGpuEvent, device_stream from vllm.logger import init_logger from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -173,7 +173,7 @@ def move_to_buffer( new_indices: np.ndarray, expert_weights: Sequence[torch.Tensor], expert_weights_buffers: Sequence[torch.Tensor], - cuda_stream: torch.cuda.Stream | None, + stream: torch.Stream | None, ep_rank: int, communicator: EplbCommunicator, layer_idx: int = 0, @@ -183,12 +183,12 @@ def move_to_buffer( Args: num_local_experts: Number of local experts. old_indices: (num_experts_total,) ndarray of current (old) - global-to-local expert assignments. + global-to-locals expert assignments. new_indices: (num_experts_total,) ndarray of desired (new) global-to-local assignments after rebalance. expert_weights: Original expert weights for the layer. expert_weights_buffers: Intermediate buffers (one per tensor). - cuda_stream: CUDA stream for async copies (can be None for sync mode). + stream: CUDA/XPU stream for async copies (can be None for sync mode). ep_rank: Rank of this process in expert parallel group. communicator: EplbCommunicator instance for P2P communication. layer_idx: Index of the MoE layer being transferred. @@ -261,10 +261,9 @@ def move_to_buffer( expert = new_local_expert_ids[dst] src_local = expert_to_src_map.get(expert, -1) if src_local != -1: - with torch.cuda.stream(cuda_stream): + with device_stream(stream): for w, b in zip(expert_weights, expert_weights_buffers): b[dst].copy_(w[src_local], non_blocking=True) - communicator.set_transfer_context(old_indices, layer_idx) # 2. Post sends @@ -430,7 +429,7 @@ def transfer_layer( ep_group: ProcessGroup, communicator: EplbCommunicator, is_profile: bool = False, - cuda_stream: torch.cuda.Stream | None = None, + stream: torch.Stream | None = None, rank_mapping: dict[int, int] | None = None, layer_idx: int = 0, ) -> TransferMetadata: @@ -451,7 +450,7 @@ def transfer_layer( is_profile (bool): If `True`, do not perform any actual weight copy. This is used during profile run, where we only perform dummy communications to reserve enough memory for the buffers. - cuda_stream: CUDA stream for async copies (can be None for sync mode). + stream: CUDA stream for async copies (can be None for sync mode). rank_mapping: Optional rank mapping for elastic expert parallelism. layer_idx: Index of the MoE layer being transferred. @@ -499,7 +498,7 @@ def transfer_layer( new_indices=new_layer_indices_np, expert_weights=expert_weights, expert_weights_buffers=expert_weights_buffer, - cuda_stream=cuda_stream, + stream=stream, ep_rank=ep_group.rank(), communicator=communicator, layer_idx=layer_idx, @@ -601,7 +600,7 @@ def rearrange_expert_weights_inplace( new_indices=new_global_expert_indices_cpu[layer_idx], expert_weights=expert_weights[layer_idx], expert_weights_buffers=weights_buffer, - cuda_stream=None, + stream=None, ep_rank=ep_rank, communicator=communicator, layer_idx=layer_idx, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 762749987b..6b6bf18ae4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -11,7 +11,6 @@ and MooncakeDistributedStore integration. """ import dataclasses -import json import math import os import queue @@ -20,10 +19,8 @@ import threading import time from collections.abc import Callable, Sequence from concurrent.futures import Future, ThreadPoolExecutor -from dataclasses import dataclass -from typing import Any, Literal, TypeVar +from typing import Any, TypeVar -import regex as re import torch import zmq @@ -67,6 +64,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import decode_lookup_response, encode_lookup_response, ) +from vllm.distributed.mooncake_store import MooncakeStoreConfig, setup_mooncake_store from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import get_ip, make_zmq_socket @@ -94,10 +92,6 @@ from .metrics import MooncakeStoreConnectorStats logger = init_logger(__name__) -DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB -DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB -DEFAULT_TENANT_ID = "default" - MOONCAKE_NO_AVAILABLE_HANDLE = -200 _T = TypeVar("_T") @@ -149,115 +143,6 @@ _DIRECT_IO_ALIGNMENT = 4096 _DIRECT_IO_PADDING_BYTES = 2 * _DIRECT_IO_ALIGNMENT -MooncakeMode = Literal["embedded", "standalone-store"] - - -@dataclass -class MooncakeStoreConfig: - """Configuration for MooncakeDistributedStore. - - ``mode`` selects the topology: ``embedded`` (each rank contributes - ``global_segment_size`` in-process) or ``standalone-store`` (rank - contributes 0; an external ``mooncake_client`` process owns the pool - and the SSD tier). - """ - - metadata_server: str - master_server_address: str - protocol: str - device_name: str - mode: MooncakeMode = "embedded" - global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE - local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE - enable_offload: bool = False - tenant_id: str = DEFAULT_TENANT_ID - - def __post_init__(self) -> None: - if self.mode not in ("embedded", "standalone-store"): - raise ValueError(f"unknown Mooncake mode: {self.mode!r}") - if self.local_buffer_size <= 0: - raise ValueError("local_buffer_size must be > 0") - if self.mode == "embedded" and self.global_segment_size == 0: - raise ValueError("embedded mode requires global_segment_size > 0") - if self.mode == "standalone-store" and self.global_segment_size != 0: - raise ValueError("standalone-store mode requires global_segment_size == 0") - - @staticmethod - def from_file(file_path: str) -> "MooncakeStoreConfig": - with open(file_path) as file: - config = json.load(file) - return MooncakeStoreConfig( - metadata_server=config.get("metadata_server", ""), - master_server_address=config.get("master_server_address", ""), - protocol=config.get("protocol", "rdma"), - device_name=config.get("device_name", ""), - mode=config.get("mode", "embedded"), - global_segment_size=_parse_size( - config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE) - ), - local_buffer_size=_parse_size( - config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE) - ), - enable_offload=bool(config.get("enable_offload", False)), - tenant_id=_normalize_tenant_id(config.get("tenant_id", DEFAULT_TENANT_ID)), - ) - - @staticmethod - def load_from_config() -> "MooncakeStoreConfig": - config_path = os.getenv("MOONCAKE_CONFIG_PATH") - if not config_path: - raise ValueError( - "The environment variable 'MOONCAKE_CONFIG_PATH' is not set." - ) - return MooncakeStoreConfig.from_file(config_path) - - -def _normalize_tenant_id(value: Any) -> str: - if value is None: - return DEFAULT_TENANT_ID - if not isinstance(value, str): - raise TypeError( - f"tenant_id must be a string or null, got {type(value).__name__}: {value!r}" - ) - tenant_id = value.strip() - return tenant_id if tenant_id else DEFAULT_TENANT_ID - - -def _parse_size(value: Any) -> int: - """Parse storage size strings with units: GB, MB, KB, B.""" - if isinstance(value, int): - return value - if not isinstance(value, str): - try: - return int(value) - except (TypeError, ValueError) as e: - raise TypeError(f"Unsupported type for size: {type(value)}") from e - - cleaned = value.strip().lower() - if not cleaned: - raise ValueError("Size cannot be empty.") - - unit_multipliers = { - "gb": 1024**3, - "mb": 1024**2, - "kb": 1024, - "b": 1, - } - match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned) - if not match: - raise ValueError(f"Invalid format: '{value}'") - - number_str = match.group(1) - unit = match.group(2) or "b" - multiplier = unit_multipliers[unit] - - try: - numeric_value = float(number_str) - except ValueError as exc: - raise ValueError(f"Invalid numeric value '{number_str}' in: '{value}'") from exc - return int(numeric_value * multiplier) - - def _align_up(value: int, alignment: int) -> int: return ((value + alignment - 1) // alignment) * alignment @@ -1595,23 +1480,7 @@ class MooncakeStoreWorker: self.store = MooncakeDistributedStore() local_ip = get_ip() local_hostname = rdma_utils.get_requester_local_hostname(local_ip) - setup_kwargs: dict[str, str] = {} - if store_config.tenant_id != DEFAULT_TENANT_ID: - setup_kwargs["tenant_id"] = store_config.tenant_id - ret = self.store.setup( - local_hostname, - store_config.metadata_server, - store_config.global_segment_size, - store_config.local_buffer_size, - store_config.protocol, - store_config.device_name, - store_config.master_server_address, - **setup_kwargs, - ) - if ret != 0: - msg = "Initialize MooncakeDistributedStore failed." - logger.error(msg) - raise RuntimeError(msg) + setup_mooncake_store(self.store, store_config, local_hostname) preferred_segment = rdma_utils.get_configured_preferred_segment(extra_config) self.preferred_segment = preferred_segment diff --git a/vllm/distributed/mooncake_store.py b/vllm/distributed/mooncake_store.py new file mode 100644 index 0000000000..f76bf02766 --- /dev/null +++ b/vllm/distributed/mooncake_store.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared configuration and setup for EC and KV Mooncake Store clients.""" + +import json +import os +from dataclasses import dataclass +from typing import Any, Literal + +import regex as re + +DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB +DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB +DEFAULT_TENANT_ID = "default" + + +MooncakeMode = Literal["embedded", "standalone-store"] + + +@dataclass +class MooncakeStoreConfig: + """Configuration for MooncakeDistributedStore. + + ``mode`` selects the topology: ``embedded`` (each rank contributes + ``global_segment_size`` in-process) or ``standalone-store`` (rank + contributes 0; an external ``mooncake_client`` process owns the pool + and the SSD tier). + """ + + metadata_server: str + master_server_address: str + protocol: str + device_name: str + mode: MooncakeMode = "embedded" + global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE + local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE + enable_offload: bool = False + tenant_id: str = DEFAULT_TENANT_ID + + def __post_init__(self) -> None: + if self.mode not in ("embedded", "standalone-store"): + raise ValueError(f"unknown Mooncake mode: {self.mode!r}") + if self.local_buffer_size <= 0: + raise ValueError("local_buffer_size must be > 0") + if self.mode == "embedded" and self.global_segment_size == 0: + raise ValueError("embedded mode requires global_segment_size > 0") + if self.mode == "standalone-store" and self.global_segment_size != 0: + raise ValueError("standalone-store mode requires global_segment_size == 0") + + @staticmethod + def from_file(file_path: str) -> "MooncakeStoreConfig": + with open(file_path) as file: + config = json.load(file) + return MooncakeStoreConfig( + metadata_server=config.get("metadata_server", ""), + master_server_address=config.get("master_server_address", ""), + protocol=config.get("protocol", "rdma"), + device_name=config.get("device_name", ""), + mode=config.get("mode", "embedded"), + global_segment_size=_parse_size( + config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE) + ), + local_buffer_size=_parse_size( + config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE) + ), + enable_offload=bool(config.get("enable_offload", False)), + tenant_id=_normalize_tenant_id(config.get("tenant_id", DEFAULT_TENANT_ID)), + ) + + @staticmethod + def load_from_config() -> "MooncakeStoreConfig": + config_path = os.getenv("MOONCAKE_CONFIG_PATH") + if not config_path: + raise ValueError( + "The environment variable 'MOONCAKE_CONFIG_PATH' is not set." + ) + return MooncakeStoreConfig.from_file(config_path) + + +def _normalize_tenant_id(value: Any) -> str: + if value is None: + return DEFAULT_TENANT_ID + if not isinstance(value, str): + raise TypeError( + f"tenant_id must be a string or null, got {type(value).__name__}: {value!r}" + ) + tenant_id = value.strip() + return tenant_id if tenant_id else DEFAULT_TENANT_ID + + +def _parse_size(value: Any) -> int: + """Parse storage size strings with units: GB, MB, KB, B.""" + if isinstance(value, int): + return value + if not isinstance(value, str): + try: + return int(value) + except (TypeError, ValueError) as e: + raise TypeError(f"Unsupported type for size: {type(value)}") from e + + cleaned = value.strip().lower() + if not cleaned: + raise ValueError("Size cannot be empty.") + + unit_multipliers = { + "gb": 1024**3, + "mb": 1024**2, + "kb": 1024, + "b": 1, + } + match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned) + if not match: + raise ValueError(f"Invalid format: '{value}'") + + number_str = match.group(1) + unit = match.group(2) or "b" + multiplier = unit_multipliers[unit] + + try: + numeric_value = float(number_str) + except ValueError as exc: + raise ValueError(f"Invalid numeric value '{number_str}' in: '{value}'") from exc + return int(numeric_value * multiplier) + + +def setup_mooncake_store( + store: Any, config: MooncakeStoreConfig, local_hostname: str +) -> None: + """Initialize a native Store, preserving explicit tenant configuration.""" + setup_kwargs: dict[str, str] = {} + if config.tenant_id != DEFAULT_TENANT_ID: + setup_kwargs["tenant_id"] = config.tenant_id + ret = store.setup( + local_hostname, + config.metadata_server, + config.global_segment_size, + config.local_buffer_size, + config.protocol, + config.device_name, + config.master_server_address, + **setup_kwargs, + ) + if ret != 0: + raise RuntimeError("Initialize MooncakeDistributedStore failed.") diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 60ed229bc4..20118649e1 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1632,10 +1632,7 @@ def get_pcp_group() -> GroupCoordinator: @contextmanager -def graph_capture( - device: torch.device, - graph_capture_context: GraphCaptureContext | None = None, -): +def graph_capture(device: torch.device): """`graph_capture` is a context manager which should surround the code that is capturing the CUDA graph. Its main purpose is to ensure that some operations will be run after the graph is captured, before the graph @@ -1647,13 +1644,8 @@ def graph_capture( the graph capture is running on a separate stream from the default stream, in order to explicitly distinguish the kernels to capture from other kernels possibly launched on background in the default stream. - - A caller may pass an explicit ``graph_capture_context`` to control the - stream used (e.g. to capture on the default stream). """ - context = graph_capture_context or GraphCaptureContext( - torch.cuda.Stream(device=device) - ) + context = GraphCaptureContext(torch.cuda.Stream(device=device)) with ( get_tp_group().graph_capture(context), get_pp_group().graph_capture(context), diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index a8a2804a06..838fe90c5c 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -486,13 +486,9 @@ class ServingTokens(GenerateBaseServing): total_tokens=(num_prompt_tokens + num_generated_tokens[i]), ) - # Omit absent prompt metadata (Rust skips None fields). - exclude = { - name - for name in ("prompt_token_ids", "mm_placeholders") - if getattr(chunk, name) is None - } - yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n" + # Omit absent fields, like the Rust frontend and the + # final usage chunk below. + yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" total_completion_tokens = sum(num_generated_tokens) final_usage_info = UsageInfo( diff --git a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py index dfc578b441..ddd891b1c0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py @@ -171,6 +171,7 @@ class XPUExperts(mk.FusedMoEExpertsModular): num_experts=self.moe_config.num_local_experts, ep_rank=self.moe_config.ep_rank, ep_size=self.moe_config.ep_size, + expert_map=expert_map, gemm1_clamp_limit=self.gemm1_clamp_limit, ) assert self.fused_moe_impl is not None diff --git a/vllm/model_executor/layers/fused_moe/router/base_router.py b/vllm/model_executor/layers/fused_moe/router/base_router.py index 0509f02ead..b51dbf90dc 100644 --- a/vllm/model_executor/layers/fused_moe/router/base_router.py +++ b/vllm/model_executor/layers/fused_moe/router/base_router.py @@ -13,7 +13,7 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.v1.worker.ubatching import dbo_current_ubatch_id -if current_platform.is_cuda_alike(): +if current_platform.is_cuda_alike() or current_platform.is_xpu(): @triton.jit def _eplb_map_and_record_i32_kernel( diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 4687cc5469..3b8b6a124f 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -85,6 +85,7 @@ class SharedExperts(torch.nn.Module): # Both these comm backends have been shown to be safe for shared expert overlap. _EPLB_OVERLAP_SAFE_BACKENDS = ( "allgather_reducescatter", + "deepep_v2", "flashinfer_nvlink_one_sided", ) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 9474ec4835..248f41d20d 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -865,25 +865,8 @@ class RocmPlatform(Platform): @classmethod def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: from vllm._aiter_ops import rocm_aiter_ops - from vllm.config.compilation import CUDAGraphMode compilation_config = vllm_config.compilation_config - model_config = vllm_config.model_config - if ( - compilation_config.cudagraph_mode is None - and model_config is not None - and on_gfx950() - and vllm_config.use_v2_model_runner - and model_config.architecture - in { - "DeepseekV4ForCausalLM", - "DeepseekV4ForConditionalGeneration", - } - ): - # Default to eager after reported gfx950/MRV2 accuracy regressions: - # https://github.com/vllm-project/vllm/issues/52644 - compilation_config.cudagraph_mode = CUDAGraphMode.NONE - use_aiter_fused_moe = rocm_aiter_ops.is_fused_moe_enabled() use_aiter_fp8_linear = rocm_aiter_ops.is_linear_fp8_enabled() use_aiter_fused_se = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 5c85257fa3..10f85be718 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -46,7 +46,6 @@ from vllm.distributed.kv_transfer.kv_connector.utils import ( copy_kv_blocks, ) from vllm.distributed.parallel_state import ( - GraphCaptureContext, get_dcp_group, get_pp_group, get_tp_group, @@ -128,7 +127,6 @@ from vllm.utils.platform_utils import num_compute_units from vllm.utils.torch_utils import ( PIN_MEMORY, async_tensor_h2d, - current_stream, kv_cache_dtype_str_to_dtype, ) from vllm.v1.attention.backend import ( @@ -6712,31 +6710,11 @@ class GPUModelRunner( per_graph_estimate = {} encoder_memory_estimate = 0 - # On ROCm, capture these throwaway profiling graphs on vLLM's dedicated - # compute stream instead of the fresh side stream graph_capture() - # allocates by default. torch's allocator pools free blocks per stream, - # so a side-stream forward strands a persistent aiter scratch buffer in - # a separate pool, shifting the physical placement of the real KV cache - # allocated afterward and slowing bandwidth-bound decode ~20%. The - # graphs are discarded, so a side stream is unnecessary here. - # Use current_stream(), not torch.cuda.current_stream(): before vLLM - # initializes its dedicated stream, torch returns the per-thread default - # stream (cuda_stream=0), which cannot be used for cudagraph capture. - # cap_ctx=None keeps the side-stream path on CUDA. - cap_ctx = ( - GraphCaptureContext(current_stream()) - if current_platform.is_rocm() - else None - ) - # Cleanup-only guard: CUDA graph capture errors should still propagate # because encoder graph capture is opt-in. try: set_cudagraph_capturing_enabled(True) - with ( - self._freeze_gc(), - graph_capture(device=self.device, graph_capture_context=cap_ctx), - ): + with self._freeze_gc(), graph_capture(device=self.device): torch.accelerator.synchronize() torch.accelerator.empty_cache() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index acba1d9b67..a288a97cc1 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -98,6 +98,7 @@ from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase from vllm.v1.worker.workspace import init_workspace_manager from ...model_executor.model_loader import TensorizerLoader +from .gpu.cudagraph_utils import has_compiled_submodule from .gpu.warmup import warmup_kernels from .utils import request_memory @@ -772,7 +773,10 @@ class Worker(WorkerBase): def compile_or_warm_up_model(self) -> CompilationTimes: warmup_sizes: list[int] = [] - if self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE: + if ( + self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE + and has_compiled_submodule(self.model_runner.get_model()) + ): # warm up sizes that are not in cudagraph capture sizes, # but users still want to compile for better performance, # e.g. for the max-num-batched token size in chunked prefill.