[Feature] Enable cache aware DP routing for standalone ATOMesh

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

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
@@ -935,6 +927,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`.
 
@@ -976,6 +969,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
@@ -1080,6 +1074,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())
 
@@ -1115,6 +1110,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(
@@ -1133,6 +1129,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,
@@ -1157,6 +1154,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,
@@ -1173,6 +1171,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,
@@ -1193,6 +1192,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,
@@ -1209,6 +1209,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,
@@ -1256,6 +1257,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())
 
@@ -1268,6 +1270,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(
@@ -1284,6 +1287,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,
@@ -1667,6 +1671,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."""

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."""