diff --git a/atom/entrypoints/atomesh/atom_standalone_service.py b/atom/entrypoints/atomesh/atom_standalone_service.py index 411ee347b4..7654076b89 100644 --- a/atom/entrypoints/atomesh/atom_standalone_service.py +++ b/atom/entrypoints/atomesh/atom_standalone_service.py @@ -24,14 +24,14 @@ load_custom_message_encoder, ) from atom.entrypoints.openai.protocol import ( + CHAT_COMPLETION_CHUNK_OBJECT, DEFAULT_MAX_TOKENS, DEFAULT_TEMPERATURE, DEFAULT_TOP_K, DEFAULT_TOP_P, - CHAT_COMPLETION_CHUNK_OBJECT, - CompletionRequest, STREAM_DONE_MESSAGE, TEXT_COMPLETION_OBJECT, + CompletionRequest, ) from atom.entrypoints.openai.reasoning import ReasoningFilter from atom.entrypoints.openai.serving_chat import ( @@ -57,6 +57,7 @@ class EngineRequest: effective_n: int future: concurrent.futures.Future[list[dict[str, Any]]] kv_transfer_params: dict[str, Any] | None = None + data_parallel_rank: int | None = None @dataclasses.dataclass @@ -67,6 +68,7 @@ class EngineStreamRequest: effective_n: int stream_queue: queue.Queue[dict[str, Any]] kv_transfer_params: dict[str, Any] | None = None + data_parallel_rank: int | None = None class AtomEngineService: @@ -95,6 +97,7 @@ def generate( request_id: str, effective_n: int, kv_transfer_params: dict[str, Any] | None = None, + data_parallel_rank: int | None = None, ) -> list[dict[str, Any]]: if self._closed.is_set(): raise RuntimeError("ATOM standalone engine service is closed") @@ -113,6 +116,7 @@ def generate( effective_n=effective_n, future=future, kv_transfer_params=kv_transfer_params, + data_parallel_rank=data_parallel_rank, ) ) try: @@ -184,6 +188,7 @@ def start_stream( request_id: str, effective_n: int, kv_transfer_params: dict[str, Any] | None = None, + data_parallel_rank: int | None = None, ) -> queue.Queue[dict[str, Any]]: if self._closed.is_set(): raise RuntimeError("ATOM standalone engine service is closed") @@ -197,6 +202,7 @@ def start_stream( effective_n=effective_n, stream_queue=stream_queue, kv_transfer_params=kv_transfer_params, + data_parallel_rank=data_parallel_rank, ) ) return stream_queue @@ -224,6 +230,7 @@ def completion_callback(request_output: Any) -> None: request.sampling_params, stream_callback=completion_callback, kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=request.data_parallel_rank, ) def _preprocess_fanout_stream_request( @@ -250,6 +257,7 @@ def completion_callback(request_output: Any) -> None: ], kv_transfer_params=request.kv_transfer_params, parent_request_id=request.request_id, + data_parallel_rank=request.data_parallel_rank, ) def _preprocess_single_request(self, request: EngineRequest) -> Any: @@ -267,6 +275,7 @@ def completion_callback(request_output: Any) -> None: request.sampling_params, stream_callback=completion_callback, kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=request.data_parallel_rank, ) state.set_num_tokens_input(seq.num_prompt_tokens) return seq @@ -293,6 +302,7 @@ def completion_callback(request_output: Any) -> None: ], kv_transfer_params=request.kv_transfer_params, parent_request_id=request.request_id, + data_parallel_rank=request.data_parallel_rank, ) if seqs: state.set_num_tokens_input(seqs[0].num_prompt_tokens) @@ -914,10 +924,15 @@ def chat_completions(self, request_data: dict[str, Any]) -> dict[str, Any]: request_data.get("temperature", DEFAULT_TEMPERATURE), ) sampling_params = self._build_sampling_params(request_data, effective_n) + data_parallel_rank = self._get_data_parallel_rank(request_data) request_id = f"chatcmpl-{uuid.uuid4().hex}" if effective_n > 1: outputs = self.engine_service.generate( - prompt, sampling_params, request_id, effective_n + prompt, + sampling_params, + request_id, + effective_n, + data_parallel_rank=data_parallel_rank, ) if not outputs: raise RuntimeError("No output generated") @@ -926,7 +941,11 @@ def chat_completions(self, request_data: dict[str, Any]) -> dict[str, Any]: ) else: outputs = self.engine_service.generate( - prompt, sampling_params, request_id, effective_n + prompt, + sampling_params, + request_id, + effective_n, + data_parallel_rank=data_parallel_rank, ) if not outputs: raise RuntimeError("No output generated") @@ -967,6 +986,7 @@ def completions(self, request_data: dict[str, Any]) -> dict[str, Any]: request_id, effective_n, kv_transfer_params=request_data.get("kv_transfer_params"), + data_parallel_rank=self._get_data_parallel_rank(request_data), ) if not outputs: raise RuntimeError("No output generated") @@ -1007,6 +1027,7 @@ def start_completions_stream(self, request_data: dict[str, Any]) -> str: request_id, effective_n, kv_transfer_params=request_data.get("kv_transfer_params"), + data_parallel_rank=self._get_data_parallel_rank(request_data), ) stream_state = CompletionStreamState( request_id=request_id, @@ -1089,6 +1110,7 @@ def start_chat_completions_stream(self, request_data: dict[str, Any]) -> str: sampling_params, request_id, effective_n, + data_parallel_rank=self._get_data_parallel_rank(request_data), ) stream_state = ChatCompletionStreamState( request_id=request_id, @@ -1159,6 +1181,21 @@ def _normalize_chat_request(request_data: dict[str, Any]) -> dict[str, Any]: normalized["max_tokens"] = normalized["max_completion_tokens"] return normalized + @staticmethod + def _get_data_parallel_rank(request_data: dict[str, Any]) -> int | None: + """Extract the DP-attention rank, if available. + + Routers can inject a ``data_parallel_rank`` field into the request body to + indicate which DP rank to route to. Falls back to round-robin if not available. + """ + raw = request_data.get("data_parallel_rank") + if raw is None: + return None + try: + return int(raw) + except (TypeError, ValueError): + raise ValueError(f"data_parallel_rank must be an integer, got {raw!r}") + def _validate_model_name(self, request_model: str | None) -> None: if ( request_model is not None diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index c8134d8500..dec914c9f8 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -425,14 +425,10 @@ def do_preprocess(): sampling_params, stream_callback=completion_callback, kv_transfer_params=kv_transfer_params, + data_parallel_rank=data_parallel_rank, ) seq = await loop.run_in_executor(None, do_preprocess) - if data_parallel_rank is not None: - seq.data_parallel_rank = data_parallel_rank - logger.info( - "Request %s pinned to data_parallel_rank=%s", seq.id, data_parallel_rank - ) try: _validate_sequence_context_length(seq) except Exception: @@ -508,6 +504,7 @@ async def generate_async_multimodal( multimodal_data: Dict[str, Any], sampling_params: SamplingParams, request_id: str, + data_parallel_rank: Optional[int] = None, ) -> AsyncGenerator[Dict[str, Any], None]: """Generate text asynchronously for one multimodal request.""" global engine, tokenizer @@ -540,6 +537,7 @@ def do_preprocess(): sampling_params, stream_callback=completion_callback, multimodal_data=multimodal_data, + data_parallel_rank=data_parallel_rank, ) seq = await loop.run_in_executor(None, do_preprocess) @@ -657,18 +655,10 @@ def do_preprocess(): kv_transfer_params=kv_transfer_params, multimodal_data=multimodal_data, parent_request_id=request_id, + data_parallel_rank=data_parallel_rank, ) seqs = await loop.run_in_executor(None, do_preprocess) - if data_parallel_rank is not None: - for seq in seqs: - seq.data_parallel_rank = data_parallel_rank - logger.info( - "Request %s fanout pinned %d sequence(s) to data_parallel_rank=%s", - request_id, - len(seqs), - data_parallel_rank, - ) try: _validate_sequence_context_length(seqs[0]) except Exception: @@ -756,6 +746,7 @@ async def setup_streaming_request( request_id: str, kv_transfer_params: Optional[Dict[str, Any]] = None, multimodal_data: Optional[Dict[str, Any]] = None, + data_parallel_rank: Optional[int] = None, ) -> Tuple[int, asyncio.Queue, int]: """Set up a streaming request with the engine. @@ -784,6 +775,7 @@ def do_preprocess(): stream_callback=stream_callback, kv_transfer_params=kv_transfer_params, multimodal_data=multimodal_data, + data_parallel_rank=data_parallel_rank, ) _seq_id_to_request_id[seq.id] = request_id return seq @@ -942,6 +934,7 @@ async def setup_streaming_request_fanout( request_id: str, kv_transfer_params: Optional[Dict[str, Any]] = None, multimodal_data: Optional[Dict[str, Any]] = None, + data_parallel_rank: Optional[int] = None, ) -> Tuple[List[int], asyncio.Queue, int]: """Fan-out variant of :func:`setup_streaming_request`. @@ -983,6 +976,7 @@ def do_preprocess(): kv_transfer_params=kv_transfer_params, multimodal_data=multimodal_data, parent_request_id=request_id, + data_parallel_rank=data_parallel_rank, ) for seq in seqs: _seq_id_to_request_id[seq.id] = request_id @@ -1087,6 +1081,7 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request) ) request_id = f"chatcmpl-{uuid.uuid4().hex}" + dp_rank = request.data_parallel_rank _log_request_event("request", request_id, request.model_dump()) @@ -1122,6 +1117,7 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request) request_id, multimodal_data=stream_multimodal_data, kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=dp_rank, ) ) gen = stream_chat_response_fanout( @@ -1140,6 +1136,7 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request) request_id, multimodal_data=stream_multimodal_data, kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=dp_rank, ) gen = stream_chat_response( request_id, @@ -1164,6 +1161,7 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request) request_id, multimodal_data=multimodal_data, kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=dp_rank, ), raw_request, request_id, @@ -1180,6 +1178,7 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request) multimodal_data, sampling_params, request_id, + data_parallel_rank=dp_rank, ), raw_request, request_id, @@ -1200,6 +1199,7 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request) sampling_params, request_id, kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=dp_rank, ), raw_request, request_id, @@ -1216,6 +1216,7 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request) sampling_params, request_id, kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=dp_rank, ), raw_request, request_id, @@ -1263,6 +1264,7 @@ async def completions(request: CompletionRequest, raw_request: Request): ) request_id = f"cmpl-{uuid.uuid4().hex}" + dp_rank = request.data_parallel_rank _log_request_event("request", request_id, request.model_dump()) @@ -1275,6 +1277,7 @@ async def completions(request: CompletionRequest, raw_request: Request): sampling_params, request_id, kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=dp_rank, ) ) gen = stream_completion_response_fanout( @@ -1291,6 +1294,7 @@ async def completions(request: CompletionRequest, raw_request: Request): sampling_params, request_id, kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=dp_rank, ) gen = stream_completion_response( request_id, @@ -1679,6 +1683,24 @@ async def kv_transfer_info(): } +@app.get("/server_info") +async def server_info(): + """Server metadata for the Atomesh router. + + The router's dp-aware discovery reads ``dp_size`` here to expand the + per-DP-rank worker set and enable cache-aware routing to the rank that + holds a request's prefix. + """ + global engine, model_name + cfg = engine.config + return { + "model_id": model_name, + "served_model_name": model_name, + "tp_size": cfg.tensor_parallel_size, + "dp_size": cfg.parallel_config.data_parallel_size, + } + + @app.post("/start_profile") async def start_profile(): """Start profiling the engine.""" diff --git a/atom/entrypoints/openai/protocol.py b/atom/entrypoints/openai/protocol.py index 75227cefdf..5219c182ff 100644 --- a/atom/entrypoints/openai/protocol.py +++ b/atom/entrypoints/openai/protocol.py @@ -121,6 +121,7 @@ class ChatCompletionRequest(BaseModel): n: Optional[int] = 1 # Optional KV-transfer metadata for P/D disaggregation. kv_transfer_params: Optional[Dict[str, Any]] = None + data_parallel_rank: Optional[int] = None def get_max_tokens(self) -> int: """Return the effective generation cap for OpenAI chat requests.""" diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index eddebc82ce..691c103180 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -12,6 +12,7 @@ import zmq import zmq.asyncio + from atom.config import Config from atom.model_engine.engine_core import EngineCore, EngineCoreRequestType from atom.model_engine.sequence import Sequence @@ -393,19 +394,22 @@ def add_request(self, seqs: List[Sequence]): copy=False, ) else: - # DP ranks: honor an explicit atomesh DPA routing hint when present; - # otherwise keep the existing round-robin behavior. + # DP ranks. A seq with an explicit target_dp_rank (set by an + # external cache-aware router) is dispatched to that rank so it + # lands on the DP rank holding its prefix cache. Seqs without a + # target fall back to round-robin load balancing. dp_seqs = [[] for _ in range(self.local_engine_count)] for seq in seqs: - requested_dp_rank = getattr(seq, "data_parallel_rank", None) - if requested_dp_rank is not None: - dp_rank = int(requested_dp_rank) - if not 0 <= dp_rank < self.local_engine_count: - raise ValueError( - f"Invalid data_parallel_rank={dp_rank}; " - f"local_engine_count={self.local_engine_count}" - ) + target = seq.target_dp_rank + if target is not None and 0 <= target < self.local_engine_count: + dp_rank = target else: + if target is not None: + logger.warning( + f"{self.label}: seq {seq.id} target_dp_rank {target} " + f"out of range [0, {self.local_engine_count}); " + f"falling back to round-robin" + ) dp_rank = self._rr_counter % self.local_engine_count self._rr_counter += 1 dp_seqs[dp_rank].append(seq) diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index a20b019dc2..948e0d3f34 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -350,6 +350,7 @@ def preprocess( kv_transfer_params=None, multimodal_data=None, request_id: Optional[str] = None, + data_parallel_rank: Optional[int] = None, ): """responsible for: 1) Tokenize @@ -371,6 +372,7 @@ def preprocess( kv_transfer_params=kv_transfer_params, multimodal_data=multimodal_data, parent_request_id=request_id, + data_parallel_rank=data_parallel_rank, ) return seqs[0] @@ -383,6 +385,7 @@ def preprocess_fanout( kv_transfer_params=None, multimodal_data=None, parent_request_id: Optional[str] = None, + data_parallel_rank: Optional[int] = None, ) -> List[Sequence]: """Tokenize once and materialize ``sampling_params.n`` Sequences. @@ -454,6 +457,7 @@ def preprocess_fanout( parent_request_id=parent_request_id, sibling_index=i, request_id=parent_request_id if n == 1 else None, + target_dp_rank=data_parallel_rank, ) seq.arrive_time = time.time() self.requests[seq.id] = seq diff --git a/atom/model_engine/sequence.py b/atom/model_engine/sequence.py index cfa58e5cf0..f4c127d6bb 100644 --- a/atom/model_engine/sequence.py +++ b/atom/model_engine/sequence.py @@ -7,6 +7,7 @@ from typing import Any, Callable, Optional import numpy as np + from atom.sampling_params import SamplingParams @@ -56,6 +57,7 @@ def __init__( multimodal_data: Optional[dict] = None, mrope_positions: Optional[np.ndarray] = None, mrope_position_delta: int = 0, + target_dp_rank: Optional[int] = None, ): self.block_size = block_size self.id = id or next(Sequence.counter) @@ -143,6 +145,8 @@ def __init__( # to safe values for single-sample requests. self.parent_request_id = parent_request_id self.sibling_index = sibling_index + # Explicitly requested DP rank, e.g. for cache aware DP routing + self.target_dp_rank = target_dp_rank def __len__(self): return self._num_tokens