[llm][kv][11/N] Enable atomic selection & reservation broadcast to avoid herding (ray-project#65010)

Signed-off-by: Jeffrey Wang <jeffreywang@anyscale.com>

Author: jeffreywang88

python/ray/llm/_internal/serve/constants.py

@@ -65,6 +65,9 @@
 RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING = (
     os.environ.get("RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING", "0") == "1"
 )
+RAY_SERVE_LLM_ENABLE_DECODE_BLOCK_PROGRESS = (
+    os.environ.get("RAY_SERVE_LLM_ENABLE_DECODE_BLOCK_PROGRESS", "0") == "1"
+)
 
 MAX_NUM_STOPPING_SEQUENCES = int(os.getenv("RAYLLM_MAX_NUM_STOPPING_SEQUENCES", "8"))
 ENV_VARS_TO_PROPAGATE = {

python/ray/llm/_internal/serve/core/ingress/router.py

@@ -126,11 +126,13 @@ async def __init__(
             # tracked LLMServer deployment.
             from ray.llm._internal.serve.routing_policies.kv_aware.kv_token_tracker import (  # noqa: E501
                 build_kv_token_tracker,
+                get_llm_router_handle,
             )
 
             self._kv_token_tracker = build_kv_token_tracker(
                 llm_config, server.deployment_id
             )
+            self._kv_token_tracker.start_reservation_broadcast(get_llm_router_handle())
             # Lazy import: this module pulls in vLLM's renderer;
             # keep it off the non-KV ingress import path.
             from ray.llm._internal.serve.routing_policies.kv_aware.tokenizer import (
@@ -206,6 +208,10 @@ async def on_lifecycle_events(self, batch):
         """
         return await self._kv_token_tracker.on_lifecycle_events(batch)
 
+    async def on_reservations_created(self, batch):
+        """Ingress-facing intake for already-selected reservation bookings."""
+        return await self._kv_token_tracker.on_reservations_created(batch)
+
     async def _pick_replica(
         self,
         handle: DeploymentHandle,

python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py

@@ -26,6 +26,9 @@
 import ray
 from ray.llm._internal.common.callbacks.base import CallbackCtx
 from ray.llm._internal.common.utils.import_utils import try_import
+from ray.llm._internal.serve.constants import (
+    RAY_SERVE_LLM_ENABLE_DECODE_BLOCK_PROGRESS,
+)
 from ray.llm._internal.serve.core.configs.llm_config import (
     DiskMultiplexConfig,
     LLMConfig,
@@ -571,7 +574,10 @@ def _start_async_llm_engine(
         # resolving it per request would block the engine's event loop.
         engine_cls = AsyncLLM
         if is_kv_aware(self.llm_config):
-            engine_cls = enable_token_tracking(AsyncLLM)
+            engine_cls = enable_token_tracking(
+                AsyncLLM,
+                report_decode_progress=RAY_SERVE_LLM_ENABLE_DECODE_BLOCK_PROGRESS,
+            )
         engine_client = engine_cls(
             vllm_config=vllm_engine_config,
             executor_class=executor_class,

python/ray/llm/_internal/serve/routing_policies/kv_aware/kv_aware_router.py

@@ -19,6 +19,18 @@
 logger = logging.getLogger(SERVE_LOGGER_NAME)
 
 
+def _get_expected_output_tokens(pending_request: PendingRequest) -> Optional[int]:
+    """The request's output cap from the routing payload, if present."""
+    if not pending_request.args:
+        return None
+    payload = pending_request.args[0]
+    for field in ("max_completion_tokens", "max_tokens"):
+        value = getattr(payload, field, None)
+        if isinstance(value, int) and value > 0:
+            return value
+    return None
+
+
 class KVAwareRouter(RequestRouter):
     """Routes each request to the candidate that best balances expected KV-cache
     overlap against the worker's current prefill/decode load.
@@ -83,6 +95,7 @@ async def choose_replicas(
             pending_request.metadata.request_id,
             token_ids,
             list(worker_id_to_replica),
+            _get_expected_output_tokens(pending_request),
         )
         return [[worker_id_to_replica[selection["worker_id"]]]]