diff --git a/.sync/vllm-sha b/.sync/vllm-sha index a78f98fe82..5b6f859dbd 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -83762b77b07e97c77f986e4bf5a9474952e47bb3 +4c81772e8bdf9e37ab2b8adfd80e4308b618fa92 diff --git a/aphrodite/compilation/passes/fusion/allreduce_rms_fusion.py b/aphrodite/compilation/passes/fusion/allreduce_rms_fusion.py index 6beb83f84a..e727bf481b 100644 --- a/aphrodite/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/aphrodite/compilation/passes/fusion/allreduce_rms_fusion.py @@ -567,7 +567,14 @@ def replacement( ) return allreduce[1], allreduce[2] - pm.register_replacement(pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass) + pm.register_replacement( + pattern, + replacement, + self.get_inputs(), + pm.fwd_only, + pm_pass, + extra_check=_norm_input_weight_dtype_match, + ) first_return_only = lambda fn: lambda a, b, c: fn(a, b, c)[0] @@ -722,7 +729,14 @@ def replacement( # quant_out, rms_norm_residual return allreduce[4], allreduce[2] - pm.register_replacement(pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass) + pm.register_replacement( + pattern, + replacement, + self.get_inputs(), + pm.fwd_only, + pm_pass, + extra_check=_norm_input_weight_dtype_match, + ) class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern): @@ -899,7 +913,14 @@ def replacement( # quant_out, rms_norm_residual, output_scale return allreduce[4], allreduce[2], allreduce[5] - pm.register_replacement(pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass) + pm.register_replacement( + pattern, + replacement, + self.get_inputs(), + pm.fwd_only, + pm_pass, + extra_check=_norm_input_weight_dtype_match, + ) class AllReduceFusionPass(AphroditePatternMatcherPass): diff --git a/aphrodite/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/aphrodite/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 559707d4f8..f0c20aa214 100644 --- a/aphrodite/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/aphrodite/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1184,6 +1184,7 @@ def reset_cache(self) -> None: for status in self._req_status.values(): for group_state in status.group_states: group_state.next_stored_block_idx = 0 + status.transfer_jobs.clear() # Discard jobs and save job_counter to be able to discard worker responses self._stale_job_threshold = self._job_counter diff --git a/aphrodite/entrypoints/anthropic/protocol.py b/aphrodite/entrypoints/anthropic/protocol.py index 595feba807..40c90c8fc4 100644 --- a/aphrodite/entrypoints/anthropic/protocol.py +++ b/aphrodite/entrypoints/anthropic/protocol.py @@ -137,6 +137,10 @@ class AnthropicMessagesRequest(BaseModel): default=None, description="KVTransfer parameters used for disaggregated serving.", ) + ec_transfer_params: dict[str, Any] | None = Field( + default=None, + description="ECTransfer parameters used for encoder-cache disaggregated serving.", + ) chat_template_kwargs: dict[str, Any] | None = Field( default=None, description=( @@ -208,6 +212,7 @@ class AnthropicMessagesResponse(BaseModel): # Aphrodite-specific fields that are not in Anthropic spec kv_transfer_params: dict[str, Any] | None = Field(default=None, description="KVTransfer parameters.") + ec_transfer_params: dict[str, Any] | None = Field(default=None, description="ECTransfer parameters.") def model_post_init(self, __context): if not self.id: diff --git a/aphrodite/entrypoints/anthropic/serving.py b/aphrodite/entrypoints/anthropic/serving.py index 2c61762c1e..de3c743742 100644 --- a/aphrodite/entrypoints/anthropic/serving.py +++ b/aphrodite/entrypoints/anthropic/serving.py @@ -484,6 +484,7 @@ def _build_base_request( top_p=anthropic_request.top_p, top_k=anthropic_request.top_k, kv_transfer_params=anthropic_request.kv_transfer_params, + ec_transfer_params=anthropic_request.ec_transfer_params, chat_template_kwargs=anthropic_request.chat_template_kwargs, ) @@ -621,6 +622,7 @@ def messages_full_converter( generator.usage, ), kv_transfer_params=generator.kv_transfer_params, + ec_transfer_params=generator.ec_transfer_params, ) choice = generator.choices[0] if choice.finish_reason == "stop": diff --git a/aphrodite/entrypoints/openai/chat_completion/protocol.py b/aphrodite/entrypoints/openai/chat_completion/protocol.py index c7d0766df4..84eb606623 100644 --- a/aphrodite/entrypoints/openai/chat_completion/protocol.py +++ b/aphrodite/entrypoints/openai/chat_completion/protocol.py @@ -160,6 +160,7 @@ class ChatCompletionResponse(OpenAIBaseModel): # ``return_prompt_text=True`` on the request). prompt_text: str | None = None kv_transfer_params: dict[str, Any] | None = Field(default=None, description="KVTransfer parameters.") + ec_transfer_params: dict[str, Any] | None = Field(default=None, description="ECTransfer parameters.") metrics: PerRequestTimingMetrics | None = None @@ -484,6 +485,11 @@ class ChatCompletionRequest(OpenAIBaseModel): description="KVTransfer parameters used for disaggregated serving.", ) + ec_transfer_params: dict[str, Any] | None = Field( + default=None, + description="ECTransfer parameters used for encoder-cache disaggregated serving.", + ) + aphrodite_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field( default=None, validation_alias=AliasChoices("aphrodite_xargs", "aphrodite_xargs"), @@ -701,6 +707,9 @@ def to_sampling_params( if self.kv_transfer_params: # Pass in kv_transfer_params via extra_args extra_args["kv_transfer_params"] = self.kv_transfer_params + if self.ec_transfer_params: + # Pass in ec_transfer_params via extra_args + extra_args["ec_transfer_params"] = self.ec_transfer_params return SamplingParams.from_optional( n=self.n, presence_penalty=self.presence_penalty, diff --git a/aphrodite/entrypoints/openai/chat_completion/serving.py b/aphrodite/entrypoints/openai/chat_completion/serving.py index bcc9c95f71..e2fb56ee8d 100644 --- a/aphrodite/entrypoints/openai/chat_completion/serving.py +++ b/aphrodite/entrypoints/openai/chat_completion/serving.py @@ -973,6 +973,7 @@ async def chat_completion_full_generator( prompt_token_ids=(final_res.prompt_token_ids if request.return_token_ids else None), prompt_text=prompt_text, kv_transfer_params=final_res.kv_transfer_params, + ec_transfer_params=final_res.ec_transfer_params, metrics=per_request_metrics, ) diff --git a/aphrodite/entrypoints/openai/completion/protocol.py b/aphrodite/entrypoints/openai/completion/protocol.py index e6bdfa3e20..f7f0b85db3 100644 --- a/aphrodite/entrypoints/openai/completion/protocol.py +++ b/aphrodite/entrypoints/openai/completion/protocol.py @@ -230,6 +230,11 @@ class CompletionRequest(OpenAIBaseModel): description="KVTransfer parameters used for disaggregated serving.", ) + ec_transfer_params: dict[str, Any] | None = Field( + default=None, + description="ECTransfer parameters used for encoder-cache disaggregated serving.", + ) + aphrodite_xargs: dict[str, str | int | float] | None = Field( default=None, validation_alias=AliasChoices("aphrodite_xargs", "aphrodite_xargs"), @@ -375,6 +380,9 @@ def to_sampling_params( if self.kv_transfer_params: # Pass in kv_transfer_params via extra_args extra_args["kv_transfer_params"] = self.kv_transfer_params + if self.ec_transfer_params: + # Pass in ec_transfer_params via extra_args + extra_args["ec_transfer_params"] = self.ec_transfer_params return SamplingParams.from_optional( n=self.n, presence_penalty=self.presence_penalty, @@ -636,6 +644,7 @@ class CompletionResponse(OpenAIBaseModel): # Aphrodite-specific fields that are not in OpenAI spec kv_transfer_params: dict[str, Any] | None = Field(default=None, description="KVTransfer parameters.") + ec_transfer_params: dict[str, Any] | None = Field(default=None, description="ECTransfer parameters.") metrics: PerRequestTimingMetrics | None = None diff --git a/aphrodite/entrypoints/openai/completion/serving.py b/aphrodite/entrypoints/openai/completion/serving.py index 9eeb650810..842f71e862 100644 --- a/aphrodite/entrypoints/openai/completion/serving.py +++ b/aphrodite/entrypoints/openai/completion/serving.py @@ -479,6 +479,7 @@ def request_output_to_completion_response( num_prompt_tokens = 0 num_generated_tokens = 0 kv_transfer_params = None + ec_transfer_params = None last_final_res = None for final_res in final_res_batch: last_final_res = final_res @@ -584,6 +585,7 @@ def request_output_to_completion_response( if final_res_batch: kv_transfer_params = final_res_batch[0].kv_transfer_params + ec_transfer_params = final_res_batch[0].ec_transfer_params return CompletionResponse( id=request_id, created=created_time, @@ -592,6 +594,7 @@ def request_output_to_completion_response( usage=usage, system_fingerprint=self.system_fingerprint, kv_transfer_params=kv_transfer_params, + ec_transfer_params=ec_transfer_params, metrics=per_request_metrics, ) diff --git a/aphrodite/entrypoints/openai/responses/context.py b/aphrodite/entrypoints/openai/responses/context.py index 460bafbd24..cac4b34c64 100644 --- a/aphrodite/entrypoints/openai/responses/context.py +++ b/aphrodite/entrypoints/openai/responses/context.py @@ -195,6 +195,7 @@ def __init__( self.input_messages: list[ResponseRawMessageAndToken] = [] self.kv_transfer_params: dict[str, Any] | None = None + self.ec_transfer_params: dict[str, Any] | None = None def append_output(self, output) -> None: self.last_output = output @@ -205,6 +206,8 @@ def append_output(self, output) -> None: self.num_output_tokens += len(output.outputs[0].token_ids or []) if output.kv_transfer_params is not None: self.kv_transfer_params = output.kv_transfer_params + if output.ec_transfer_params is not None: + self.ec_transfer_params = output.ec_transfer_params # Accumulate text, token_ids, and logprobs for streaming mode delta_output = output.outputs[0] @@ -323,6 +326,7 @@ def __init__( self.output_messages: list[ResponseRawMessageAndToken] = [] self._accumulated_token_ids: list[int] = [] self.kv_transfer_params: dict[str, Any] | None = None + self.ec_transfer_params: dict[str, Any] | None = None def append_output(self, output: RequestOutput) -> None: self.num_prompt_tokens = len(output.prompt_token_ids or []) @@ -331,6 +335,9 @@ def append_output(self, output: RequestOutput) -> None: if output.kv_transfer_params is not None: self.kv_transfer_params = output.kv_transfer_params + if output.ec_transfer_params is not None: + self.ec_transfer_params = output.ec_transfer_params + completion = output.outputs[0] self.finish_reason = completion.finish_reason @@ -610,6 +617,7 @@ def __init__( self.is_first_turn = True self.first_tok_of_message = True self.kv_transfer_params: dict[str, Any] | None = None + self.ec_transfer_params: dict[str, Any] | None = None def append_output(self, output: RequestOutput) -> None: if self.first_tok_of_message: @@ -625,6 +633,8 @@ def append_output(self, output: RequestOutput) -> None: self._update_decode_token_usage(output) if output.kv_transfer_params is not None: self.kv_transfer_params = output.kv_transfer_params + if output.ec_transfer_params is not None: + self.ec_transfer_params = output.ec_transfer_params if output.finished: self.finish_reason = output.outputs[0].finish_reason diff --git a/aphrodite/entrypoints/openai/responses/protocol.py b/aphrodite/entrypoints/openai/responses/protocol.py index cb12b369d1..cdfba6082b 100644 --- a/aphrodite/entrypoints/openai/responses/protocol.py +++ b/aphrodite/entrypoints/openai/responses/protocol.py @@ -280,6 +280,10 @@ class ResponsesRequest(OpenAIBaseModel): default=None, description="KVTransfer parameters used for disaggregated serving.", ) + ec_transfer_params: dict[str, Any] | None = Field( + default=None, + description="ECTransfer parameters used for encoder-cache disaggregated serving.", + ) chat_template_kwargs: dict[str, Any] | None = Field( default=None, description=( @@ -394,6 +398,8 @@ def to_sampling_params( extra_args: dict[str, Any] = self.aphrodite_xargs if self.aphrodite_xargs else {} if self.kv_transfer_params: extra_args["kv_transfer_params"] = self.kv_transfer_params + if self.ec_transfer_params: + extra_args["ec_transfer_params"] = self.ec_transfer_params return SamplingParams.from_optional( temperature=temperature, @@ -636,6 +642,7 @@ class ResponsesResponse(OpenAIBaseModel): # Aphrodite-specific fields that are not in OpenAI spec kv_transfer_params: dict[str, Any] | None = Field(default=None, description="KVTransfer parameters.") + ec_transfer_params: dict[str, Any] | None = Field(default=None, description="ECTransfer parameters.") # --8<-- [start:responses-response-extra-params] # These are populated when enable_response_messages is set to True @@ -677,6 +684,7 @@ def from_request( input_messages: ResponseInputOutputMessage | None = None, output_messages: ResponseInputOutputMessage | None = None, kv_transfer_params: dict[str, Any] | None = None, + ec_transfer_params: dict[str, Any] | None = None, ) -> "ResponsesResponse": incomplete_details: IncompleteDetails | None = None if status == "incomplete": @@ -715,6 +723,7 @@ def from_request( user=request.user, usage=usage, kv_transfer_params=kv_transfer_params, + ec_transfer_params=ec_transfer_params, ) diff --git a/aphrodite/entrypoints/openai/responses/serving.py b/aphrodite/entrypoints/openai/responses/serving.py index d2b294eeeb..3e8e0cf65f 100644 --- a/aphrodite/entrypoints/openai/responses/serving.py +++ b/aphrodite/entrypoints/openai/responses/serving.py @@ -867,6 +867,7 @@ async def responses_full_generator( status=status, usage=usage, kv_transfer_params=context.kv_transfer_params, + ec_transfer_params=context.ec_transfer_params, ) if request.store: diff --git a/aphrodite/entrypoints/scale_out/token_in_token_out/protocol.py b/aphrodite/entrypoints/scale_out/token_in_token_out/protocol.py index d4683448b8..97d0e022d7 100644 --- a/aphrodite/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/aphrodite/entrypoints/scale_out/token_in_token_out/protocol.py @@ -132,6 +132,10 @@ def validate_token_ids(cls, v: list[int]) -> list[int]: default=None, description="KVTransfer parameters used for disaggregated serving.", ) + ec_transfer_params: dict[str, Any] | None = Field( + default=None, + description="ECTransfer parameters used for encoder-cache disaggregated serving.", + ) # Tracks which keys the caller explicitly set inside ``sampling_params`` # when the request was parsed from a JSON body. Lets the server tell @@ -238,6 +242,10 @@ class GenerateResponse(BaseModel): default=None, description="KVTransfer parameters used for disaggregated serving.", ) + ec_transfer_params: dict[str, Any] | None = Field( + default=None, + description="ECTransfer parameters used for encoder-cache disaggregated serving.", + ) ####### Derender (postprocessing) ####### diff --git a/aphrodite/entrypoints/scale_out/token_in_token_out/serving.py b/aphrodite/entrypoints/scale_out/token_in_token_out/serving.py index fbe7c33e5b..b1c5c9f9fb 100644 --- a/aphrodite/entrypoints/scale_out/token_in_token_out/serving.py +++ b/aphrodite/entrypoints/scale_out/token_in_token_out/serving.py @@ -311,6 +311,7 @@ async def serve_tokens_full_generator( usage=usage, prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs), kv_transfer_params=final_res.kv_transfer_params, + ec_transfer_params=final_res.ec_transfer_params, ) # Log complete response if output logging is enabled diff --git a/aphrodite/envs.py b/aphrodite/envs.py index 78a16d60ce..7988c401b9 100755 --- a/aphrodite/envs.py +++ b/aphrodite/envs.py @@ -207,6 +207,8 @@ APHRODITE_DISABLE_REQUEST_ID_RANDOMIZATION: bool = False APHRODITE_NIXL_SIDE_CHANNEL_HOST: str = "localhost" APHRODITE_NIXL_SIDE_CHANNEL_PORT: int = 5600 + APHRODITE_EC_SIDE_CHANNEL_HOST: str = "localhost" + APHRODITE_EC_SIDE_CHANNEL_PORT: int = 5601 APHRODITE_MOONCAKE_BOOTSTRAP_PORT: int = 8998 APHRODITE_MOONCAKE_STORE_TIER_LOG: bool = False APHRODITE_MOONCAKE_LOAD_RECV_THREADS: int = 1 @@ -1383,6 +1385,12 @@ def _resolve_rust_frontend_path() -> str | None: "APHRODITE_NIXL_SIDE_CHANNEL_HOST": lambda: os.getenv("APHRODITE_NIXL_SIDE_CHANNEL_HOST", "localhost"), # Port used for NIXL handshake between remote agents. "APHRODITE_NIXL_SIDE_CHANNEL_PORT": lambda: int(os.getenv("APHRODITE_NIXL_SIDE_CHANNEL_PORT", "5600")), + # IP address used for the EC connector's ZMQ side channel + # (producer ROUTER bind, consumer DEALER dial). + "APHRODITE_EC_SIDE_CHANNEL_HOST": lambda: os.getenv("APHRODITE_EC_SIDE_CHANNEL_HOST", "localhost"), + # Port for the EC connector's ZMQ side channel; advertised to peers + # via `ec_transfer_params.peer_port` on the producer's response. + "APHRODITE_EC_SIDE_CHANNEL_PORT": lambda: int(os.getenv("APHRODITE_EC_SIDE_CHANNEL_PORT", "5601")), # Port used for Mooncake handshake between remote agents. "APHRODITE_MOONCAKE_BOOTSTRAP_PORT": lambda: int(os.getenv("APHRODITE_MOONCAKE_BOOTSTRAP_PORT", "8998")), # Log per-batch memory/disk tier breakdown on external GETs. diff --git a/aphrodite/outputs.py b/aphrodite/outputs.py index 8c4ed94377..fed5a0553e 100644 --- a/aphrodite/outputs.py +++ b/aphrodite/outputs.py @@ -102,6 +102,7 @@ class RequestOutput: None if decoder-only. num_cached_tokens: The number of tokens with prefix cache hit. kv_transfer_params: The params for remote K/V transfer. + ec_transfer_params: The params for remote encoder-cache transfer. """ def __init__( @@ -119,6 +120,7 @@ def __init__( num_cached_tokens: int | None = None, *, kv_transfer_params: dict[str, Any] | None = None, + ec_transfer_params: dict[str, Any] | None = None, # Forward compatibility, code that uses args added in new release can # still run with older versions of Aphrodite without breaking. **kwargs: Any, @@ -137,12 +139,14 @@ def __init__( self.encoder_prompt_token_ids = encoder_prompt_token_ids self.num_cached_tokens = num_cached_tokens self.kv_transfer_params = kv_transfer_params + self.ec_transfer_params = ec_transfer_params def add(self, next_output: "RequestOutput", aggregate: bool) -> None: """Merge subsequent RequestOutput into this one""" self.finished |= next_output.finished self.kv_transfer_params = next_output.kv_transfer_params + self.ec_transfer_params = next_output.ec_transfer_params for next_completion in next_output.outputs: for i, completion in enumerate(self.outputs): diff --git a/aphrodite/v1/core/sched/interface.py b/aphrodite/v1/core/sched/interface.py index b17814b284..c8b60f6055 100644 --- a/aphrodite/v1/core/sched/interface.py +++ b/aphrodite/v1/core/sched/interface.py @@ -9,6 +9,7 @@ if TYPE_CHECKING: from aphrodite.config import AphroditeConfig + from aphrodite.distributed.ec_transfer.ec_connector.base import ECConnectorBase from aphrodite.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1 from aphrodite.v1.core.sched.output import GrammarOutput, SchedulerOutput from aphrodite.v1.engine import EngineCoreOutputs @@ -244,3 +245,6 @@ def shutdown(self) -> None: def get_kv_connector(self) -> "KVConnectorBase_V1 | None": return None + + def get_ec_connector(self) -> "ECConnectorBase | None": + return None diff --git a/aphrodite/v1/core/sched/scheduler.py b/aphrodite/v1/core/sched/scheduler.py index d0e0ec54b4..f3efb4efca 100644 --- a/aphrodite/v1/core/sched/scheduler.py +++ b/aphrodite/v1/core/sched/scheduler.py @@ -10,6 +10,7 @@ from aphrodite.compilation.cuda_graph import CUDAGraphStat from aphrodite.config import AphroditeConfig from aphrodite.distributed.ec_transfer.ec_connector.base import ( + ECConnectorBase, ECConnectorMetadata, ECConnectorRole, ) @@ -1517,6 +1518,7 @@ def update_from_output( new_token_ids = generated_token_ids pooler_output = pooler_outputs[req_index] if pooler_outputs else None kv_transfer_params = None + ec_transfer_params = None status_before_stop = request.status num_output_tokens_before = len(request._output_token_ids) @@ -1585,7 +1587,7 @@ def update_from_output( finish_reason = request.get_finished_reason() finished = self._handle_stopped_request(request) if finished: - kv_transfer_params = self._free_request(request) + kv_transfer_params, ec_transfer_params = self._free_request(request) if status_before_stop == RequestStatus.RUNNING: stopped_running_reqs.add(request) @@ -1601,7 +1603,7 @@ def update_from_output( # Get prompt logprobs for this request. prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) - if new_token_ids or pooler_output is not None or kv_transfer_params or stopped: + if new_token_ids or pooler_output is not None or kv_transfer_params or ec_transfer_params or stopped: # Add EngineCoreOutput for this Request. outputs[request.client_index].append( EngineCoreOutput( @@ -1615,6 +1617,7 @@ def update_from_output( events=request.take_events(), prefill_stats=request.take_prefill_stats(), kv_transfer_params=kv_transfer_params, + ec_transfer_params=ec_transfer_params, trace_headers=request.trace_headers, routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, @@ -1945,11 +1948,22 @@ def finish_requests( return [(r.request_id, r.client_index) for r in valid_requests] - def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None: + def _free_request( + self, request: Request, delay_free_blocks: bool = False + ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: assert request.is_finished() self._inflight_prefills.discard(request) connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) + + # EC Connector: mirror the KV hook. The contract requires firing + # before the encoder cache is freed so the connector can inspect + # per-request state and emit ec_transfer_params for the response body. + ec_xfer_params: dict[str, Any] | None = None + if self.ec_connector is not None: + ec_delay_free, ec_xfer_params = self.ec_connector.request_finished(request) + connector_delay_free_blocks |= ec_delay_free + self.encoder_cache_manager.free(request) request_id = request.request_id self.finished_req_ids.add(request_id) @@ -1960,7 +1974,7 @@ def _free_request(self, request: Request, delay_free_blocks: bool = False) -> di if not delay_free_blocks: self._free_blocks(request) - return kv_xfer_params + return kv_xfer_params, ec_xfer_params def _free_blocks(self, request: Request): assert request.is_finished() @@ -2184,6 +2198,9 @@ def shutdown(self) -> None: def get_kv_connector(self) -> KVConnectorBase_V1 | None: return self.connector + def get_ec_connector(self) -> ECConnectorBase | None: + return self.ec_connector + def _connector_finished(self, request: Request) -> tuple[bool, dict[str, Any] | None]: """ Invoke the KV connector request_finished() method if applicable. diff --git a/aphrodite/v1/engine/__init__.py b/aphrodite/v1/engine/__init__.py index dcc67b7281..dabcac5eab 100644 --- a/aphrodite/v1/engine/__init__.py +++ b/aphrodite/v1/engine/__init__.py @@ -188,6 +188,7 @@ class EngineCoreOutput( stop_reason: int | str | None = None events: list[EngineCoreEvent] | None = None kv_transfer_params: dict[str, Any] | None = None + ec_transfer_params: dict[str, Any] | None = None trace_headers: Mapping[str, str] | None = None diff --git a/aphrodite/v1/engine/core.py b/aphrodite/v1/engine/core.py index 7c279cf6b3..1235d13d62 100644 --- a/aphrodite/v1/engine/core.py +++ b/aphrodite/v1/engine/core.py @@ -356,6 +356,9 @@ def add_request(self, request: Request, request_wave: int = 0): if request.kv_transfer_params is not None and (not self.scheduler.get_kv_connector()): logger.warning("Got kv_transfer_params, but no KVConnector found. Disabling KVTransfer for this request.") + if request.ec_transfer_params is not None and self.scheduler.get_ec_connector() is None: + logger.warning("Got ec_transfer_params, but no ECConnector found. Disabling ECTransfer for this request.") + self.scheduler.add_request(request) if request.abort_immediately: # Immediately abort so the connector's request_finished hook runs diff --git a/aphrodite/v1/engine/output_processor.py b/aphrodite/v1/engine/output_processor.py index 244423c0a3..c60d00e4f3 100644 --- a/aphrodite/v1/engine/output_processor.py +++ b/aphrodite/v1/engine/output_processor.py @@ -266,6 +266,7 @@ def make_request_output( finish_reason: FinishReason | None, stop_reason: int | str | None, kv_transfer_params: dict[str, Any] | None = None, + ec_transfer_params: dict[str, Any] | None = None, ) -> RequestOutput | PoolingRequestOutput | None: finished = finish_reason is not None final_only = self.output_kind == RequestOutputKind.FINAL_ONLY @@ -313,7 +314,7 @@ def make_request_output( return None external_req_id = self.parent_req.external_req_id - return self._new_request_output(external_req_id, outputs, finished, kv_transfer_params) + return self._new_request_output(external_req_id, outputs, finished, kv_transfer_params, ec_transfer_params) def _new_request_output( self, @@ -321,6 +322,7 @@ def _new_request_output( outputs: list[CompletionOutput] | list[PoolingOutput], finished: bool, kv_transfer_params: dict[str, Any] | None = None, + ec_transfer_params: dict[str, Any] | None = None, ) -> RequestOutput | PoolingRequestOutput: # If prompt embeds were used, put placeholder prompt token ids prompt_token_ids = self.prompt_token_ids @@ -354,6 +356,7 @@ def _new_request_output( outputs=cast(list[CompletionOutput], outputs), finished=finished, kv_transfer_params=kv_transfer_params, + ec_transfer_params=ec_transfer_params, num_cached_tokens=self.num_cached_tokens, metrics=self.stats, ) @@ -480,6 +483,7 @@ def abort_requests(self, request_ids: Iterable[str], internal: bool) -> list[str finish_reason=FinishReason.ABORT, stop_reason=None, kv_transfer_params=None, + ec_transfer_params=None, ) ): req_state.queue.put(request_output) @@ -601,6 +605,7 @@ def process_outputs( finish_reason = engine_core_output.finish_reason stop_reason = engine_core_output.stop_reason kv_transfer_params = engine_core_output.kv_transfer_params + ec_transfer_params = engine_core_output.ec_transfer_params if engine_core_output.routed_experts is not None: req_state.routed_experts_chunks.append(engine_core_output.routed_experts) @@ -629,6 +634,7 @@ def process_outputs( finish_reason, stop_reason, kv_transfer_params, + ec_transfer_params, ): if req_state.streaming_input: request_output.finished = False diff --git a/aphrodite/v1/kv_offload/cpu/shared_offload_region.py b/aphrodite/v1/kv_offload/cpu/shared_offload_region.py index c8fcf03873..51f1660840 100644 --- a/aphrodite/v1/kv_offload/cpu/shared_offload_region.py +++ b/aphrodite/v1/kv_offload/cpu/shared_offload_region.py @@ -31,14 +31,14 @@ class SharedOffloadRegion: the rest open the existing file and wait until it reaches the expected size. Each worker then mmap()s the full file. - File path: /dev/shm/aphrodite_offload_{instance_id}.mmap + File path: /dev/shm/aphrodite_offload_{engine_id}.mmap """ BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE def __init__( self, - instance_id: str, + engine_id: str, num_blocks: int, rank: int | None, kv_bytes_per_block: int, @@ -51,7 +51,7 @@ def __init__( self._row_stride = kv_bytes_per_block self.total_size_bytes = self.num_blocks * self._row_stride - self.mmap_path = f"/dev/shm/aphrodite_offload_{instance_id}.mmap" + self.mmap_path = f"/dev/shm/aphrodite_offload_{engine_id}.mmap" self._creator = False # set True only if this worker creates the file self.rank = rank if rank is not None: diff --git a/aphrodite/v1/kv_offload/tiering/spec.py b/aphrodite/v1/kv_offload/tiering/spec.py index c15c56b6fe..a3243e6f0e 100644 --- a/aphrodite/v1/kv_offload/tiering/spec.py +++ b/aphrodite/v1/kv_offload/tiering/spec.py @@ -106,6 +106,13 @@ def __init__(self, aphrodite_config: AphroditeConfig, kv_cache_config: KVCacheCo # Scheduler-side mmap (rank=None); kept for cleanup self._scheduler_mmap: SharedOffloadRegion | None = None + # engine_id is unique per DP replica (suffixed with _dp{rank} in both + # the Ray and multiprocessing paths), so it names a per-replica offload + # region. Non-None is guaranteed by OffloadingSpec.__init__. + assert aphrodite_config.kv_transfer_config is not None + assert aphrodite_config.kv_transfer_config.engine_id is not None + self._engine_id: str = aphrodite_config.kv_transfer_config.engine_id + @override def get_manager(self) -> OffloadingManager: """ @@ -122,7 +129,7 @@ def get_manager(self) -> OffloadingManager: # Create scheduler-side SharedOffloadRegion (rank=None) so the # primary tier can eagerly create a memoryview over _base. scheduler_mmap = SharedOffloadRegion( - instance_id=self.aphrodite_config.instance_id, + engine_id=self._engine_id, num_blocks=self.num_blocks, rank=None, kv_bytes_per_block=self.kv_bytes_per_offloaded_block, @@ -180,9 +187,12 @@ def get_manager(self) -> OffloadingManager: @override def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: - rank = torch.accelerator.current_device_index() + # Fold the global physical device index into the replica-local + # [0, world_size) slot range. + world_size = self.aphrodite_config.parallel_config.world_size + rank = torch.accelerator.current_device_index() % world_size worker_mmap = SharedOffloadRegion( - instance_id=self.aphrodite_config.instance_id, + engine_id=self._engine_id, num_blocks=self.num_blocks, rank=rank, kv_bytes_per_block=self.kv_bytes_per_offloaded_block, diff --git a/aphrodite/v1/request.py b/aphrodite/v1/request.py index 6754a054f3..1caa3c3434 100644 --- a/aphrodite/v1/request.py +++ b/aphrodite/v1/request.py @@ -96,6 +96,8 @@ def __init__( # P/D: Connector-specific KV transfer parameters. self.kv_transfer_params: dict[str, Any] | None = None + # E/P/D: Connector-specific encoder-cache transfer parameters. + self.ec_transfer_params: dict[str, Any] | None = None self.kv_cache_report_mode = "incremental" if pooling_params is not None: @@ -110,6 +112,7 @@ def __init__( if sampling_params.extra_args is not None: self.kv_transfer_params = sampling_params.extra_args.get("kv_transfer_params") + self.ec_transfer_params = sampling_params.extra_args.get("ec_transfer_params") self.kv_cache_report_mode = sampling_params.extra_args.get("kv_cache_report_mode", "incremental") else: raise ValueError("sampling_params and pooling_params can't both be unset") diff --git a/rust/proto/aphrodite_grpc.proto b/rust/proto/aphrodite_grpc.proto index 36cf628be7..e733e77599 100644 --- a/rust/proto/aphrodite_grpc.proto +++ b/rust/proto/aphrodite_grpc.proto @@ -107,6 +107,9 @@ message KVCacheParameters { // KV Connector transfer parameters google.protobuf.Struct kv_transfer_params = 3; + + // Encoder cache connector transfer parameters + google.protobuf.Struct ec_transfer_params = 4; } // Controls which extra candidate tokens at each position should be returned @@ -173,6 +176,7 @@ message FinishInfo { google.protobuf.Struct kv_transfer_params = 6; //uint64 seed = 7; + google.protobuf.Struct ec_transfer_params = 8; } // Info for candidate tokens other than the input/sampled @@ -193,4 +197,3 @@ message CandidateTokenInfo { message TokenIds { repeated uint32 ids = 1; } - diff --git a/rust/src/chat/src/event.rs b/rust/src/chat/src/event.rs index ee45648f8e..e4e35e3da4 100644 --- a/rust/src/chat/src/event.rs +++ b/rust/src/chat/src/event.rs @@ -202,5 +202,8 @@ pub enum ChatEvent { finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. kv_transfer_params: Option, + /// Connector-specific encoder cache transfer parameters for + /// disaggregated serving. + ec_transfer_params: Option, }, } diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs index d5e800522c..e525cd5ef7 100644 --- a/rust/src/chat/src/output/default/unified.rs +++ b/rust/src/chat/src/output/default/unified.rs @@ -257,6 +257,7 @@ pub(crate) async fn unified_event_stream( usage: finished.usage, finish_reason: finished.finish_reason, kv_transfer_params: finished.kv_transfer_params, + ec_transfer_params: finished.ec_transfer_params, }) .await; } @@ -389,6 +390,7 @@ mod tests { usage: aphrodite_llm::TokenUsage::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, + ec_transfer_params: None, }), } } @@ -630,6 +632,7 @@ mod tests { usage: aphrodite_llm::TokenUsage::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, + ec_transfer_params: None, }, ] ); @@ -673,6 +676,7 @@ mod tests { usage: aphrodite_llm::TokenUsage::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, + ec_transfer_params: None, }, ] ); diff --git a/rust/src/chat/src/output/harmony/mod.rs b/rust/src/chat/src/output/harmony/mod.rs index 42218d5c70..e62472bb17 100644 --- a/rust/src/chat/src/output/harmony/mod.rs +++ b/rust/src/chat/src/output/harmony/mod.rs @@ -370,6 +370,7 @@ async fn harmony_assistant_event_stream( usage: finished.usage, finish_reason: finished.finish_reason, kv_transfer_params: finished.kv_transfer_params, + ec_transfer_params: finished.ec_transfer_params, }) .await; } diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs index eff8675ffe..d89f3163c9 100644 --- a/rust/src/chat/src/output/harmony/tests.rs +++ b/rust/src/chat/src/output/harmony/tests.rs @@ -52,6 +52,7 @@ fn finished() -> Finished { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, } } @@ -115,6 +116,7 @@ fn interrupted_final_message_is_preserved() { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }) ); } @@ -175,6 +177,7 @@ fn interrupted_analysis_message_is_preserved() { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }) ); } diff --git a/rust/src/chat/src/output/mod.rs b/rust/src/chat/src/output/mod.rs index 67829b4fe4..3dac72269d 100644 --- a/rust/src/chat/src/output/mod.rs +++ b/rust/src/chat/src/output/mod.rs @@ -48,6 +48,9 @@ pub(crate) enum AssistantEvent { finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. kv_transfer_params: Option, + /// Connector-specific encoder cache transfer parameters for + /// disaggregated serving. + ec_transfer_params: Option, }, } diff --git a/rust/src/chat/src/output/structured.rs b/rust/src/chat/src/output/structured.rs index 66612eb8f1..ced2119a48 100644 --- a/rust/src/chat/src/output/structured.rs +++ b/rust/src/chat/src/output/structured.rs @@ -146,6 +146,7 @@ impl StructuredEventState { usage: aphrodite_llm::TokenUsage, finish_reason: FinishReason, kv_transfer_params: Option, + ec_transfer_params: Option, ) -> Result> { let mut events = Vec::new(); self.close_open_text_block(&mut events); @@ -155,6 +156,7 @@ impl StructuredEventState { usage, finish_reason, kv_transfer_params, + ec_transfer_params, }); Ok(events) } @@ -296,8 +298,9 @@ pub(crate) async fn structured_chat_event_stream( usage, finish_reason, kv_transfer_params, + ec_transfer_params, } => { - for next in state.finish(usage, finish_reason, kv_transfer_params)? { + for next in state.finish(usage, finish_reason, kv_transfer_params, ec_transfer_params)? { y.yield_ok(next).await; } } @@ -334,6 +337,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); @@ -388,6 +392,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); @@ -439,6 +444,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); @@ -490,6 +496,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); @@ -557,6 +564,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); diff --git a/rust/src/chat/src/renderer/hf/template.rs b/rust/src/chat/src/renderer/hf/template.rs index 38248808c4..11f0be5174 100644 --- a/rust/src/chat/src/renderer/hf/template.rs +++ b/rust/src/chat/src/renderer/hf/template.rs @@ -10,7 +10,7 @@ use std::path::Path; use minijinja::Environment; use serde::{Deserialize, Serialize}; -use serde_json::self; +use serde_json::{self}; use aphrodite_text::backend::hf::HfSpecialTokens; use super::error::TemplateError; diff --git a/rust/src/chat/src/stream.rs b/rust/src/chat/src/stream.rs index b28825f23a..d7a2b50227 100644 --- a/rust/src/chat/src/stream.rs +++ b/rust/src/chat/src/stream.rs @@ -22,6 +22,9 @@ pub struct CollectedAssistantMessage { pub finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. pub kv_transfer_params: Option, + /// Connector-specific encoder cache transfer parameters for disaggregated + /// serving. + pub ec_transfer_params: Option, } /// Per-request stream of chat events. @@ -77,6 +80,7 @@ impl ChatEventStream { usage, finish_reason, kv_transfer_params, + ec_transfer_params, } => { return Ok(CollectedAssistantMessage { message: done, @@ -89,6 +93,7 @@ impl ChatEventStream { usage, finish_reason, kv_transfer_params, + ec_transfer_params, }); } ChatEvent::ToolCallEnd { call, .. } => { @@ -194,6 +199,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]), ); @@ -234,6 +240,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, } ); } diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs index ed858df231..b4e9290e74 100644 --- a/rust/src/chat/tests/chat.rs +++ b/rust/src/chat/tests/chat.rs @@ -48,6 +48,7 @@ fn request_output( stop_reason, events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, @@ -73,6 +74,7 @@ fn request_output_with_logprobs( stop_reason, events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index be42da2349..9194e354d0 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -676,6 +676,7 @@ fn decoded_completion_stream( usage: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), } }); @@ -686,6 +687,7 @@ fn decoded_completion_stream( usage: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }); events.push(DecodedTextEvent::TextDelta { delta: chunk.delta, diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index a35190f905..b32227a527 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -97,6 +97,8 @@ pub struct EngineCoreOutput { #[serde(default)] pub kv_transfer_params: Option, #[serde(default)] + pub ec_transfer_params: Option, + #[serde(default)] pub trace_headers: Option, /// Breakdown of the scheduled prefill computation, set on the first output /// of a newly scheduled prefill and elided for subsequent decode outputs. @@ -362,6 +364,7 @@ mod tests { stop_reason: Some(StopReason::Text("stop".to_string())), events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, @@ -414,6 +417,7 @@ mod tests { stop_reason: None, events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index a8b68c5a37..0123031b1c 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -226,6 +226,7 @@ fn request_output( stop_reason: None, events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, @@ -2522,6 +2523,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { stop_reason: None, events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index 3b628655d0..cac077f94f 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -91,6 +91,7 @@ class EngineCoreOutput( stop_reason: int | str | None = None events: object | None = None kv_transfer_params: object | None = None + ec_transfer_params: object | None = None trace_headers: object | None = None prefill_stats: object | None = None routed_experts: object | None = None diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 999ad5c5b4..e8bbd01679 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -38,6 +38,9 @@ pub struct CollectedGenerateOutput { pub usage: TokenUsage, /// Connector-specific KV transfer parameters for disaggregated serving. pub kv_transfer_params: Option, + /// Connector-specific encoder cache transfer parameters for disaggregated + /// serving. + pub ec_transfer_params: Option, } /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for @@ -146,6 +149,9 @@ pub struct GenerateOutput { pub cached_token_count: usize, /// Connector-specific KV transfer parameters for disaggregated serving. pub kv_transfer_params: Option, + /// Connector-specific encoder cache transfer parameters for disaggregated + /// serving. + pub ec_transfer_params: Option, } impl GenerateOutput { @@ -192,6 +198,7 @@ impl GenerateOutput { finish_reason, cached_token_count: 0, kv_transfer_params: None, + ec_transfer_params: None, } } } @@ -281,6 +288,7 @@ impl Stream for GenerateOutputStream { finish_reason, cached_token_count, kv_transfer_params: raw.kv_transfer_params, + ec_transfer_params: raw.ec_transfer_params, }; Poll::Ready(Some(Ok(output))) @@ -363,6 +371,7 @@ impl> + Send> T { cached_token_count, }, kv_transfer_params: None, + ec_transfer_params: None, }); } @@ -375,6 +384,7 @@ impl> + Send> T { cached_token_count, }; collected.kv_transfer_params = output.kv_transfer_params; + collected.ec_transfer_params = output.ec_transfer_params; return Ok(collected); } } diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 4483402558..aa9586244d 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -48,6 +48,7 @@ fn request_output_with_events( stop_reason: None, events, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, @@ -72,6 +73,7 @@ fn request_output_with_logprobs( stop_reason: None, events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, @@ -86,6 +88,7 @@ fn request_output_with_logprobs_and_kv( new_logprobs: Option, prompt_logprobs: Option, kv_transfer_params: Option, + ec_transfer_params: Option, ) -> EngineCoreOutput { EngineCoreOutput { request_id: request_id.to_string(), @@ -97,6 +100,7 @@ fn request_output_with_logprobs_and_kv( stop_reason: None, events: None, kv_transfer_params, + ec_transfer_params, trace_headers: None, prefill_stats: None, routed_experts: None, @@ -330,7 +334,6 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { send_outputs( push, RequestBatchOutputs { - engine_index: 0, outputs: vec![ EngineCoreOutput { prefill_stats: Some(PrefillStats { @@ -354,9 +357,10 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)), None, Some(serde_json::json!({"connector": "x"})), + None, ), ], - finished_requests: None, + ..Default::default() } .into(), ) diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 7e18f6d134..291a7577a6 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -69,6 +69,11 @@ pub fn to_text_request( let map = sampling_params.aphrodite_xargs.get_or_insert_with(Default::default); map.insert("kv_transfer_params".to_string(), kv_json); } + if let Some(ec_struct) = kv.ec_transfer_params.as_ref() { + let ec_json = proto_struct_to_json(ec_struct); + let map = sampling_params.aphrodite_xargs.get_or_insert_with(Default::default); + map.insert("ec_transfer_params".to_string(), ec_json); + } if kv.bypass_prefix_cache { sampling_params.skip_reading_prefix_cache = Some(true); } @@ -341,6 +346,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo { finish_reason, stop_reason, kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct), + ec_transfer_params: finished.ec_transfer_params.as_ref().and_then(json_to_proto_struct), } } @@ -584,6 +590,7 @@ mod tests { }, finish_reason: reason, kv_transfer_params: None, + ec_transfer_params: None, } } diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index 2f002fd3b8..58ecb5d99b 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -71,6 +71,7 @@ impl pb::generate_server::Generate for GenerateServiceImpl { usage: collected.usage, finish_reason: collected.finish_reason, kv_transfer_params: collected.kv_transfer_params, + ec_transfer_params: collected.ec_transfer_params, }; let outputs = convert::to_sequence_output( diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index bacaf0d619..73ee08acbe 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -102,6 +102,7 @@ fn request_output( stop_reason: None, events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, diff --git a/rust/src/server/src/routes/http_client_tests.rs b/rust/src/server/src/routes/http_client_tests.rs index 788b2b40a0..89bba9f089 100644 --- a/rust/src/server/src/routes/http_client_tests.rs +++ b/rust/src/server/src/routes/http_client_tests.rs @@ -98,6 +98,7 @@ fn request_output( stop_reason: None, events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index c2fd11ec79..5e8f4da71e 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -266,6 +266,7 @@ fn collect_generate( }], prompt_logprobs, kv_transfer_params: collected.kv_transfer_params, + ec_transfer_params: collected.ec_transfer_params, }) } @@ -404,6 +405,7 @@ mod tests { finish_reason: None, cached_token_count: 0, kv_transfer_params: None, + ec_transfer_params: None, }), Ok(GenerateOutput { request_id: String::new(), @@ -416,6 +418,7 @@ mod tests { finish_reason: Some(FinishReason::stop_eos()), cached_token_count: 2, kv_transfer_params: None, + ec_transfer_params: None, }), ]); diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index d7dfbcf027..13c819bf2a 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -4,7 +4,7 @@ use super::types::GenerateRequest; use super::validate; use crate::error::ApiError; use crate::lora::LoraModelResolution; -use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params}; +use crate::utils::{ResolvedRequestContext, merge_ec_transfer_params, merge_kv_transfer_params}; /// Lowered generate request plus the response request ID. #[derive(Debug, Clone, PartialEq)] @@ -56,6 +56,10 @@ pub(super) fn prepare_generate_request( sampling_params.aphrodite_xargs, request.kv_transfer_params.as_ref(), ); + sampling_params.aphrodite_xargs = merge_ec_transfer_params( + sampling_params.aphrodite_xargs, + request.ec_transfer_params.as_ref(), + ); let text_request = TextRequest { request_id: ctx.request_id.clone(), diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index 972b06403e..87ce31d57c 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -22,6 +22,7 @@ pub struct GenerateRequest { #[serde(default)] pub priority: i32, pub kv_transfer_params: Option>, + pub ec_transfer_params: Option>, #[serde(flatten)] pub other: Map, } @@ -66,6 +67,7 @@ pub(super) struct GenerateResponse { pub choices: Vec, pub prompt_logprobs: Option>>>, pub kv_transfer_params: Option, + pub ec_transfer_params: Option, } /// Mirrors the Python Aphrodite `Logprob` class used in prompt-logprobs payloads. diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index af41301365..929b6b6294 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -144,6 +144,7 @@ async fn collect_chat_completion( usage, finish_reason, kv_transfer_params, + ec_transfer_params, } = collected; let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json); let saw_tool_calls = message.tool_calls().next().is_some(); @@ -224,6 +225,7 @@ async fn collect_chat_completion( prompt_logprobs, prompt_token_ids: return_token_ids.then(|| prompt_token_ids.to_vec()), kv_transfer_params, + ec_transfer_params, }) } @@ -951,6 +953,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); @@ -1031,6 +1034,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); @@ -1086,6 +1090,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); @@ -1167,6 +1172,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); @@ -1300,6 +1306,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); @@ -1381,6 +1388,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 4643bca401..a8b15823c6 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -13,7 +13,9 @@ use crate::routes::openai::utils::structured_outputs::convert_from_response_form use crate::routes::openai::utils::types::{ ChatMessage, ContentPart, MessageContent, Tool, ToolChoice, ToolChoiceValue, }; -use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer_params}; +use crate::utils::{ + ResolvedRequestContext, convert_logit_bias, merge_ec_transfer_params, merge_kv_transfer_params, +}; /// Lowered chat request plus the public response metadata carried by every SSE /// chunk. @@ -132,7 +134,7 @@ pub(super) fn prepare_chat_request( structured_outputs, skip_reading_prefix_cache: None, aphrodite_xargs: merge_kv_transfer_params( - request.aphrodite_xargs, + merge_ec_transfer_params(request.aphrodite_xargs, request.ec_transfer_params.as_ref()), request.kv_transfer_params.as_ref(), ), }, diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index 0a0803bcab..0599085254 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -232,6 +232,9 @@ pub struct ChatCompletionRequest { /// KV transfer parameters for disaggregated serving pub kv_transfer_params: Option>, + /// Encoder cache transfer parameters for disaggregated serving + pub ec_transfer_params: Option>, + /// Additional request parameters with string or numeric values for custom /// extensions pub aphrodite_xargs: Option>, @@ -299,6 +302,7 @@ impl Default for ChatCompletionRequest { return_token_ids: None, cache_salt: None, kv_transfer_params: None, + ec_transfer_params: None, aphrodite_xargs: None, repetition_detection: None, } @@ -346,6 +350,7 @@ pub(super) struct ChatCompletionResponse { pub prompt_logprobs: Option>>>, pub prompt_token_ids: Option>, pub kv_transfer_params: Option, + pub ec_transfer_params: Option, } /// Mirrors the Python Aphrodite `ChatCompletionResponseChoice` class. diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index f295c2f720..09cf80b3f5 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -212,6 +212,7 @@ async fn collect_completion( usage: Some(usage), system_fingerprint: None, kv_transfer_params: collected.kv_transfer_params, + ec_transfer_params: collected.ec_transfer_params, }) } @@ -682,6 +683,7 @@ mod tests { "repetition_detected".to_string(), ))), kv_transfer_params: None, + ec_transfer_params: None, }), }), ]); @@ -786,6 +788,7 @@ mod tests { }, finish_reason: FinishReason::Length, kv_transfer_params: None, + ec_transfer_params: None, }), }), ]); @@ -837,6 +840,7 @@ mod tests { }, finish_reason: FinishReason::Length, kv_transfer_params: None, + ec_transfer_params: None, }), }), ]); @@ -891,6 +895,7 @@ mod tests { }, finish_reason: FinishReason::Length, kv_transfer_params: None, + ec_transfer_params: None, }), }), ]); @@ -962,6 +967,7 @@ mod tests { }, finish_reason: FinishReason::Length, kv_transfer_params: None, + ec_transfer_params: None, }), }), ]); @@ -1035,6 +1041,7 @@ mod tests { }, finish_reason: FinishReason::Length, kv_transfer_params: None, + ec_transfer_params: None, }), }), ]); diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 9c247a7690..0c6bcc5fac 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -7,7 +7,9 @@ use crate::error::ApiError; use crate::lora::LoraModelResolution; use crate::routes::openai::completions::validate; use crate::routes::openai::utils::structured_outputs::convert_from_response_format_value; -use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer_params}; +use crate::utils::{ + ResolvedRequestContext, convert_logit_bias, merge_ec_transfer_params, merge_kv_transfer_params, +}; /// Lowered completion request plus the public response metadata carried by /// every SSE chunk. @@ -127,7 +129,7 @@ pub(super) fn prepare_completion_request( structured_outputs, skip_reading_prefix_cache: None, aphrodite_xargs: merge_kv_transfer_params( - request.aphrodite_xargs, + merge_ec_transfer_params(request.aphrodite_xargs, request.ec_transfer_params.as_ref()), request.kv_transfer_params.as_ref(), ), }, diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index d7c1a06141..05d59274ad 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -174,6 +174,9 @@ pub struct CompletionRequest { /// KV transfer parameters for disaggregated serving pub kv_transfer_params: Option>, + /// Encoder cache transfer parameters for disaggregated serving + pub ec_transfer_params: Option>, + /// Additional request parameters with string or numeric values for custom /// extensions pub aphrodite_xargs: Option>, @@ -209,6 +212,7 @@ pub(super) struct CompletionResponse { pub usage: Option, pub system_fingerprint: Option, pub kv_transfer_params: Option, + pub ec_transfer_params: Option, } /// Mirrors the Python Aphrodite `CompletionResponseChoice` class. diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 612d7d060a..5918c645c4 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -73,6 +73,7 @@ fn request_output_with_stop_reason( stop_reason, events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, @@ -98,6 +99,7 @@ fn request_output_with_logprobs( stop_reason, events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, @@ -113,6 +115,7 @@ fn request_output_with_logprobs_and_kv( new_logprobs: Option, new_prompt_logprobs_tensors: Option, kv_transfer_params: Option, + ec_transfer_params: Option, ) -> EngineCoreOutput { EngineCoreOutput { request_id: request_id.to_string(), @@ -124,6 +127,7 @@ fn request_output_with_logprobs_and_kv( stop_reason, events: None, kv_transfer_params, + ec_transfer_params, trace_headers: None, prefill_stats: None, routed_experts: None, @@ -2202,6 +2206,7 @@ async fn non_stream_chat_returns_json_response() { (response_object, "system_fingerprint"), (response_object, "prompt_token_ids"), (response_object, "kv_transfer_params"), + (response_object, "ec_transfer_params"), (choice, "logprobs"), (choice, "stop_reason"), (choice, "token_ids"), @@ -2989,6 +2994,7 @@ async fn non_stream_completions_return_json_response() { for (object, key) in [ (response_object, "system_fingerprint"), (response_object, "kv_transfer_params"), + (response_object, "ec_transfer_params"), (choice, "logprobs"), (choice, "stop_reason"), (choice, "prompt_logprobs"), @@ -3637,6 +3643,7 @@ async fn non_stream_raw_generate_returns_token_output_envelope() { Some(sample_logprobs_for_token(44, 45)), None, Some(json!({"connector": "x"})), + None, ), ], finished_requests: None, diff --git a/rust/src/server/src/utils.rs b/rust/src/server/src/utils.rs index c0ab3a2cba..03af421803 100644 --- a/rust/src/server/src/utils.rs +++ b/rust/src/server/src/utils.rs @@ -45,6 +45,24 @@ pub fn merge_kv_transfer_params( xargs } +/// Merge `ec_transfer_params` into the `aphrodite_xargs` map, mirroring the +/// Python Aphrodite behavior where `ec_transfer_params` is injected into +/// `extra_args` for engine-core consumption. +pub fn merge_ec_transfer_params( + mut xargs: Option>, + ec_transfer_params: Option<&HashMap>, +) -> Option> { + if let Some(ec_params) = ec_transfer_params { + let map = xargs.get_or_insert_with(HashMap::new); + map.insert( + "ec_transfer_params".to_string(), + // This is safe because we know that `ec_params` is already valid JSON. + serde_json::to_value(ec_params).unwrap(), + ); + } + xargs +} + /// Convert OpenAI-style `logit_bias` with string token-ID keys into the /// internal `HashMap` representation, validating that every key /// parses as a `u32`. diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index 9963d253a9..2fee31ebee 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/rust/src/text/src/output/decoded.rs @@ -44,6 +44,9 @@ pub struct Finished { pub finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. pub kv_transfer_params: Option, + /// Connector-specific encoder cache transfer parameters for disaggregated + /// serving. + pub ec_transfer_params: Option, } /// Internal decoded-text event emitted before higher-level assistant @@ -151,6 +154,7 @@ pub async fn decoded_text_event_stream( let decoder = decoder.as_mut().unwrap(); let kv_transfer_params = output.kv_transfer_params; + let ec_transfer_params = output.ec_transfer_params; let mut finish_reason = output.finish_reason; let mut stop_str_matched = false; let suppress_terminal_stop_token = finish_reason.as_ref().is_some_and(|r| r.is_stop()) @@ -275,6 +279,7 @@ pub async fn decoded_text_event_stream( }, finish_reason: reason, kv_transfer_params, + ec_transfer_params, }), }) .await; diff --git a/rust/src/text/src/output/mod.rs b/rust/src/text/src/output/mod.rs index c5c4bc0405..bc79ba70e1 100644 --- a/rust/src/text/src/output/mod.rs +++ b/rust/src/text/src/output/mod.rs @@ -26,6 +26,9 @@ pub struct CollectedTextOutput { pub usage: aphrodite_llm::TokenUsage, /// Connector-specific KV transfer parameters for disaggregated serving. pub kv_transfer_params: Option, + /// Connector-specific encoder cache transfer parameters for disaggregated + /// serving. + pub ec_transfer_params: Option, } #[allow(clippy::manual_async_fn, reason = "specify `Send` bound")] @@ -77,6 +80,7 @@ impl T { finish_reason: FinishReason::Error, usage: aphrodite_llm::TokenUsage::default(), kv_transfer_params: None, + ec_transfer_params: None, }) }; @@ -85,6 +89,7 @@ impl T { collected.finish_reason = finished.finish_reason; collected.usage = finished.usage; collected.kv_transfer_params = finished.kv_transfer_params; + collected.ec_transfer_params = finished.ec_transfer_params; return Ok(collected); } } @@ -156,6 +161,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), }), ]); @@ -273,6 +279,7 @@ mod tests { }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), }), ]); diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index 92652abbdc..8be50e67d3 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -226,6 +226,23 @@ def ops_in_model_before(self): ] +class TestAllReduceGemmaRMSNormStaticQuantFP8Model(TestAllReduceRMSNormStaticQuantFP8Model): + def __init__( + self, + hidden_size=16, + token_num=16, + eps=1e-6, + dtype: torch.dtype = torch.float16, + ): + super().__init__(hidden_size, token_num, eps, dtype) + self.norm = [GemmaRMSNorm(hidden_size, eps) for _ in range(4)] + for norm in self.norm: + norm.weight.requires_grad_(False) + + def ops_in_model_before(self): + return [torch.ops.aphrodite.all_reduce.default] + + class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module): """Exercises the new ROCm AITER AR+RMS+per-group-FP8-quant patterns. @@ -402,6 +419,15 @@ def ops_in_model_before(self): reason="Not supported on ROCm platform", ), ), + pytest.param( + TestAllReduceGemmaRMSNormStaticQuantFP8Model, + True, + False, + marks=pytest.mark.skipif( + current_platform.is_rocm(), + reason="Not supported on ROCm platform", + ), + ), pytest.param( TestAllReduceRMSNormStaticQuantFP8Model, False, @@ -571,7 +597,10 @@ def all_reduce_fusion_pass_on_test_model( assert all_reduce_fusion_pass.matched_count == 4, f"{all_reduce_fusion_pass.matched_count=}" backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False) backend.check_after_ops(model.ops_in_model_after()) - if test_model_cls is TestAllReduceGemmaRMSNormModel: + if test_model_cls in ( + TestAllReduceGemmaRMSNormModel, + TestAllReduceGemmaRMSNormStaticQuantFP8Model, + ): fused_op = torch.ops.aphrodite.flashinfer_trtllm_fused_allreduce_norm.default fused_nodes = list(find_op_nodes(fused_op, backend.graph_post_pass)) assert fused_nodes diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 8040862303..b2c81625d0 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -284,7 +284,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): def free_request(req, delay_free_blocks=False): scheduler.finished_req_ids.add(req.request_id) scheduler.requests.pop(req.request_id, None) - return None + return None, None scheduler._free_request = Mock(side_effect=free_request) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index c458e5d445..7d7322f0c9 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -2867,7 +2867,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): def free_request(req: Request, delay_free_blocks: bool = False): scheduler.finished_req_ids.add(req.request_id) scheduler.requests.pop(req.request_id, None) - return None + return None, None scheduler._free_request = Mock(side_effect=free_request) diff --git a/tests/v1/ec_connector/unit/test_ec_transfer_params.py b/tests/v1/ec_connector/unit/test_ec_transfer_params.py new file mode 100644 index 0000000000..6ac151a569 --- /dev/null +++ b/tests/v1/ec_connector/unit/test_ec_transfer_params.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sanity tests for ec_transfer_params protocol plumbing. + +No running engine required. +""" + +from unittest.mock import MagicMock + +from aphrodite.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from aphrodite.outputs import CompletionOutput, RequestOutput +from aphrodite.sampling_params import SamplingParams +from aphrodite.v1.request import Request, RequestStatus +from tests.v1.core.utils import create_scheduler + +EC_PARAMS: dict = {"mm_hash_abc": {"peer_host": "10.0.0.1", "peer_port": 5501}} + + +def test_ec_transfer_params_routed_to_sampling_params_extra_args(): + """ec_transfer_params on the request must land in SamplingParams.extra_args.""" + req = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + max_completion_tokens=5, + ec_transfer_params=EC_PARAMS, + ) + sp = req.to_sampling_params(max_tokens=5, default_sampling_params={}) + assert sp.extra_args is not None + assert sp.extra_args.get("ec_transfer_params") == EC_PARAMS + + +def test_request_output_add_propagates_ec_transfer_params(): + """RequestOutput.add() must carry ec_transfer_params forward to the caller.""" + + def _out(ec_params): + return RequestOutput( + request_id="r1", + prompt="p", + prompt_token_ids=[1], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=0.0, + logprobs=None, + finish_reason=None, + ) + ], + finished=False, + ec_transfer_params=ec_params, + ) + + accumulated = _out(None) + accumulated.add(_out(EC_PARAMS), aggregate=True) + assert accumulated.ec_transfer_params == EC_PARAMS + + +def test_request_reads_ec_transfer_params_from_extra_args(): + """v1 Request must pull ec_transfer_params out of SamplingParams.extra_args.""" + sp = SamplingParams(extra_args={"ec_transfer_params": EC_PARAMS}) + req = Request( + request_id="r1", + prompt_token_ids=[1, 2, 3], + sampling_params=sp, + pooling_params=None, + ) + assert req.ec_transfer_params == EC_PARAMS + + +def test_free_request_calls_ec_connector_and_surfaces_params(): + """_free_request must call ec_connector.request_finished() and return its params.""" + sp = SamplingParams(max_tokens=1) + sp.update_from_generation_config({}, 50256) + request = Request( + request_id="test-req", + prompt_token_ids=[1, 2, 3], + sampling_params=sp, + pooling_params=None, + client_index=0, + ) + scheduler = create_scheduler(use_ec_connector=True, ec_role="ec_producer") + scheduler.add_request(request) + request.status = RequestStatus.FINISHED_STOPPED + + mock_ec = MagicMock() + mock_ec.request_finished.return_value = (False, EC_PARAMS) + scheduler.ec_connector = mock_ec + + kv_params, ec_params = scheduler._free_request(request) + + mock_ec.request_finished.assert_called_once_with(request) + assert ec_params == EC_PARAMS + assert kv_params is None + + +def test_free_request_without_ec_connector_returns_none(): + """When no EC connector is configured, ec_transfer_params must be None.""" + sp = SamplingParams(max_tokens=1) + sp.update_from_generation_config({}, 50256) + request = Request( + request_id="test-req", + prompt_token_ids=[1, 2, 3], + sampling_params=sp, + pooling_params=None, + client_index=0, + ) + scheduler = create_scheduler(use_ec_connector=True, ec_role="ec_producer") + scheduler.add_request(request) + request.status = RequestStatus.FINISHED_STOPPED + + kv_params, ec_params = scheduler._free_request(request) + + assert ec_params is None + assert kv_params is None diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index 54e07aa45d..c4a7042e03 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -96,7 +96,7 @@ def test_transfer( SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT, ) mmap_region = SharedOffloadRegion( - instance_id=str(uuid.uuid4()), + engine_id=str(uuid.uuid4()), num_blocks=num_cpu_blocks, rank=0, kv_bytes_per_block=cpu_page_size, diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index d830b7d36c..3c5d77692e 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -34,7 +34,7 @@ def _set_spawn_method(monkeypatch): def _make_region( - instance_id: str, + engine_id: str, num_blocks: int = 4, cpu_page_size: int = PAGE_SIZE, num_workers: int = 1, @@ -42,7 +42,7 @@ def _make_region( ) -> SharedOffloadRegion: assert cpu_page_size % PAGE_SIZE == 0 return SharedOffloadRegion( - instance_id=instance_id, + engine_id=engine_id, num_blocks=num_blocks, rank=rank, kv_bytes_per_block=num_workers * cpu_page_size, @@ -57,9 +57,9 @@ def _cleanup_file(path: str) -> None: @contextlib.contextmanager -def _region(instance_id: str, **kwargs): +def _region(engine_id: str, **kwargs): """Context manager: create one region, clean up on exit.""" - r = _make_region(instance_id, **kwargs) + r = _make_region(engine_id, **kwargs) try: yield r finally: @@ -69,7 +69,7 @@ def _region(instance_id: str, **kwargs): @contextlib.contextmanager def _multi_region( - instance_id: str, + engine_id: str, num_workers: int, num_blocks: int = 4, cpu_page_size: int = PAGE_SIZE, @@ -77,7 +77,7 @@ def _multi_region( """Context manager: create one SharedOffloadRegion per rank, clean up on exit.""" regions = [ SharedOffloadRegion( - instance_id=instance_id, + engine_id=engine_id, num_blocks=num_blocks, rank=rank, kv_bytes_per_block=num_workers * cpu_page_size, @@ -94,7 +94,7 @@ def _multi_region( def _race_construct( - instance_id: str, + engine_id: str, num_workers: int, num_blocks: int = 4, cpu_page_size: int = PAGE_SIZE, @@ -108,7 +108,7 @@ def worker(rank: int) -> None: barrier.wait() # all threads start at the same instant try: regions[rank] = SharedOffloadRegion( - instance_id=instance_id, + engine_id=engine_id, num_blocks=num_blocks, rank=rank, kv_bytes_per_block=num_workers * cpu_page_size, @@ -127,7 +127,7 @@ def worker(rank: int) -> None: def _mp_race_construct_and_write( - instance_id: str, + engine_id: str, num_blocks: int, rank: int, num_workers: int, @@ -141,7 +141,7 @@ def _mp_race_construct_and_write( parent a window to read the raw mmap before the creator removes the file.""" try: region = SharedOffloadRegion( - instance_id=instance_id, + engine_id=engine_id, num_blocks=num_blocks, rank=rank, kv_bytes_per_block=num_workers * cpu_page_size, @@ -306,7 +306,7 @@ def test_create_next_view_multiprocess_slots(iid): # Parent is rank 0 (creator); child is rank 1 (joiner). region = SharedOffloadRegion( - instance_id=iid, + engine_id=iid, num_blocks=num_blocks, rank=0, kv_bytes_per_block=num_workers * PAGE_SIZE,