Merge branch 'main' into wentao-epd-support-for-MRv2

Signed-off-by: yewentao256 <zhyanwentao@126.com>

examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh

@@ -29,7 +29,7 @@ if [[ -z "${DEVICE_AFFINITY_ENV:-}" ]]; then
 fi
 
 EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_cache}"
-TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-12000}"   # wait_for_server timeout
+TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-300}"   # wait_for_server timeout
 
 NUM_PROMPTS="${NUM_PROMPTS:-100}"    # number of prompts to send in benchmark
 

vllm/config/vllm.py

@@ -2207,10 +2207,6 @@ class VllmConfig:
             # Will be added by https://github.com/vllm-project/vllm/pull/35045
             unsupported.append("KV sharing fast prefill")
 
-        if self.ec_transfer_config is not None:
-            # Will be added by https://github.com/vllm-project/vllm/pull/38390
-            unsupported.append("EC transfer")
-
         return unsupported
 
     def _validate_v2_model_runner(self) -> None:

vllm/v1/worker/ec_connector_model_runner_mixin.py

@@ -38,14 +38,21 @@ class ECConnectorModelRunnerMixin:
     def maybe_get_ec_connector_output(
         scheduler_output: "SchedulerOutput",
         encoder_cache: dict[str, torch.Tensor],
+        enabled: bool = True,
+        save_new_caches: bool = False,
         **kwargs,
     ) -> AbstractContextManager[ECConnectorOutput | None]:
-        return (
-            ECConnectorModelRunnerMixin._get_ec_connector_output(
-                scheduler_output, encoder_cache, **kwargs
-            )
-            if has_ec_transfer()
-            else nullcontext()
+        if (
+            not enabled
+            or scheduler_output.ec_connector_metadata is None
+            or not has_ec_transfer()
+        ):
+            return nullcontext()
+        return ECConnectorModelRunnerMixin._get_ec_connector_output(
+            scheduler_output,
+            encoder_cache,
+            save_new_caches=save_new_caches,
+            **kwargs,
         )
 
     # This context manager must be used within an active forward context.
@@ -55,6 +62,7 @@ class ECConnectorModelRunnerMixin:
     def _get_ec_connector_output(
         scheduler_output: "SchedulerOutput",
         encoder_cache: dict[str, torch.Tensor],
+        save_new_caches: bool = False,
         **kwargs,
     ) -> Generator[ECConnectorOutput, None, None]:
         output = ECConnectorOutput()
@@ -68,8 +76,14 @@ class ECConnectorModelRunnerMixin:
         if ec_connector.is_consumer:
             ec_connector.start_load_caches(encoder_cache, **kwargs)
 
+        cached_hashes = set(encoder_cache) if save_new_caches else None
         try:
             yield output
+            if cached_hashes is not None:
+                for mm_hash in encoder_cache.keys() - cached_hashes:
+                    ec_connector.save_caches(
+                        encoder_cache=encoder_cache, mm_hash=mm_hash
+                    )
         finally:
             output.finished_sending, output.finished_recving = (
                 ec_connector.get_finished(scheduler_output.finished_req_ids)

vllm/v1/worker/gpu/block_table.py

@@ -120,6 +120,8 @@ class BlockTables:
             self.num_blocks.np[i, req_index] = start + len(block_ids)
 
     def apply_staged_writes(self) -> None:
+        if self.num_kv_cache_groups == 0:
+            return
         if self.num_kv_cache_groups == 1:
             # Single group: write directly, skipping the per-write group lookup.
             self.block_tables[0].apply_write()
@@ -138,6 +140,8 @@ class BlockTables:
         out: tuple[torch.Tensor, ...] | None = None,
         out_ptrs: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, ...]:
+        if self.num_kv_cache_groups == 0:
+            return ()
         if out is None:
             out = tuple(self.input_block_tables)
             out_ptrs = self.input_block_table_ptrs
@@ -173,6 +177,8 @@ class BlockTables:
         num_tokens_padded: int,
         out: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        if self.num_kv_cache_groups == 0:
+            return (self.slot_mappings if out is None else out)[:, :num_tokens_padded]
         num_reqs = idx_mapping.shape[0]
         num_groups = self.num_kv_cache_groups
         slot_mappings = self.slot_mappings if out is None else out

vllm/v1/worker/gpu/mm/encoder_runner.py

@@ -51,6 +51,8 @@ class EncoderRunner:
                 mm_feature = mm_features[mm_input_id]
                 if mm_feature.data is None:
                     continue
+                if mm_feature.identifier in self.encoder_cache.encoder_outputs:
+                    continue
                 mm_hashes.append(mm_feature.identifier)
                 mm_kwargs.append((mm_feature.modality, mm_feature.data))
 

vllm/v1/worker/gpu/model_runner.py

@@ -54,8 +54,15 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
-from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
+from vllm.v1.outputs import (
+    DraftTokenIds,
+    ModelRunnerOutput,
+    make_empty_encoder_model_runner_output,
+)
 from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
+from vllm.v1.worker.ec_connector_model_runner_mixin import (
+    ECConnectorModelRunnerMixin,
+)
 from vllm.v1.worker.gpu import pcp_manager as pcp
 from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput
 from vllm.v1.worker.gpu.attn_utils import (
@@ -122,7 +129,7 @@ from vllm.v1.worker.utils import KVBlockZeroer, copy_kv_cache_blocks_inplace
 logger = init_logger(__name__)
 
 
-class GPUModelRunner(LoRAModelRunnerMixin):
+class GPUModelRunner(LoRAModelRunnerMixin, ECConnectorModelRunnerMixin):
     def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.vllm_config = vllm_config
         self.model_config = vllm_config.model_config
@@ -137,6 +144,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         self.device = device
         self.dtype = self.model_config.dtype
+        ec_config = vllm_config.ec_transfer_config
+        self.is_ec_producer_only = (
+            ec_config is not None
+            and ec_config.is_ec_producer
+            and not ec_config.is_ec_consumer
+        )
+        mm_config = self.model_config.multimodal_config
+        self.is_encoder_only = self.is_ec_producer_only or bool(
+            mm_config and mm_config.mm_encoder_only
+        )
         self.kv_cache_dtype = self.dtype
         if self.cache_config.cache_dtype != "auto":
             # Quantized KV cache.
@@ -410,6 +427,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         return torch.cuda.current_stream(self.device)
 
     def get_kv_cache_spec(self):
+        if self.is_encoder_only:
+            return {}
         return get_kv_cache_spec(self.vllm_config)
 
     def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
@@ -539,6 +558,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         is_profile: bool = False,
         **kwargs,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        if self.is_encoder_only:
+            empty = torch.empty(0, device=self.device)
+            return empty, empty
         if skip_attn and not is_profile:
             raise ValueError(
                 "skip_attn must only be True for initial memory profiling."
@@ -689,6 +711,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     dummy_mm_inputs, mm_budget
                 )
 
+        if self.is_encoder_only:
+            torch.accelerator.synchronize()
+            self.reset_encoder_cache()
+            gc.collect()
+            return
+
         hidden_states, sample_hidden_states = self._dummy_run(
             self.max_num_tokens, skip_attn=True, is_profile=True
         )
@@ -727,6 +755,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
     @torch.inference_mode()
     def capture_model(self) -> int:
+        if self.is_encoder_only:
+            return 0
+
         assert self.cudagraph_manager is not None
         if not self.cudagraph_manager.needs_capture():
             logger.warning(
@@ -1284,7 +1315,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         input_ids = input_batch.input_ids
         inputs_embeds = None
+        ec_connector_output = None
         if self.supports_mm_inputs and self.is_first_pp_rank:
+            assert self.encoder_cache is not None
             # Run MM encoder (if needed) and get multimodal embeddings.
             # Only first PP rank prepares multimodal embeddings.
             if dummy_run:
@@ -1303,12 +1336,25 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                         lora_state=self.lora_state,
                         scheduled_encoder_inputs=scheduled_encoder_inputs,
                     )
-                inputs_embeds = self.model_state.get_mm_embeddings(
-                    scheduled_encoder_inputs, input_batch, self.req_states
-                )
+
+                with self.maybe_get_ec_connector_output(
+                    scheduler_output,
+                    encoder_cache=self.encoder_cache.encoder_outputs,
+                    enabled=not self.is_encoder_decoder,
+                    save_new_caches=self.is_ec_producer_only,
+                ) as ec_connector_output:
+                    inputs_embeds = self.model_state.get_mm_embeddings(
+                        scheduled_encoder_inputs, input_batch, self.req_states
+                    )
+
             if inputs_embeds is not None and not self.model.requires_raw_input_tokens:
                 input_ids = None
 
+        if self.is_encoder_only:
+            output = make_empty_encoder_model_runner_output(scheduler_output)
+            output.ec_connector_output = ec_connector_output
+            return output
+
         model_inputs = {
             "input_ids": input_ids,
             "positions": input_batch.positions,

vllm/v1/worker/gpu/warmup.py

@@ -165,6 +165,9 @@ def warmup_kernels(
     decode_query_len + 1 prompt tokens each. The second iteration simulates
     a decode step with all requests generating decode_query_len tokens.
     """
+    if model_runner.is_encoder_only:
+        return
+
     num_spec_steps = model_runner.num_speculative_steps
     decode_query_len = model_runner.decode_query_len
     # Use decode_query_len + 1 tokens so the prefill batch's per-request query