Auto-enable router dp_aware routing when DP attention is on

miles/ray/rollout/router_manager.py

@@ -1,5 +1,6 @@
 import logging
 import multiprocessing
+import os
 import random
 import uuid
 
@@ -16,6 +17,40 @@
 logger = logging.getLogger(__name__)
 
 
+def _maybe_enable_router_dp_aware(args, router_args) -> None:
+    """Auto-enable sgl-router DP-aware routing when DP attention shards the KV cache.
+
+    With ``sglang_dp_size > 1`` each DP rank owns an independent KV pool and radix
+    prefix tree. Unless the router routes at DP-rank granularity, requests scatter
+    across ranks by load and the prefix cache fragments, sharply lowering the
+    KV-cache hit rate. Enabling ``dp_aware`` keeps prefix-sharing requests on the
+    same rank.
+
+    Honors an explicit ``--router-dp-aware`` (already True -> left untouched) and the
+    ``MILES_DISABLE_AUTO_DP_AWARE=1`` opt-out, warning loudly in the opt-out case
+    since it reintroduces the prefix-cache fragmentation.
+    """
+    dp_size = getattr(args, "sglang_dp_size", 1) or 1
+    if dp_size <= 1 or getattr(router_args, "dp_aware", False):
+        return
+    if os.environ.get("MILES_DISABLE_AUTO_DP_AWARE") == "1":
+        logger.warning(
+            "DP attention is enabled (sglang_dp_size=%d) but router dp_aware routing is "
+            "explicitly disabled via MILES_DISABLE_AUTO_DP_AWARE=1. The router will dispatch "
+            "per-engine and SGLang will scatter requests across DP ranks by load, fragmenting "
+            "the radix prefix cache and likely reducing the KV-cache hit rate. Unset the env "
+            "var to restore prefix-affinity routing.",
+            dp_size,
+        )
+        return
+    router_args.dp_aware = True
+    logger.info(
+        "DP attention enabled (sglang_dp_size=%d); auto-enabling router dp_aware routing to "
+        "preserve prefix-cache locality.",
+        dp_size,
+    )
+
+
 def start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool = False) -> tuple[str, int]:
     """Start sgl router or miles router and return (router_ip, router_port).
 
@@ -58,6 +93,8 @@ def start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool =
         if args.sglang_router_policy:
             router_args.policy = args.sglang_router_policy
 
+        _maybe_enable_router_dp_aware(args, router_args)
+
         if has_pd_disaggregation:
             router_args.pd_disaggregation = True
 

tests/fast/router/test_router_dp_aware.py

@@ -0,0 +1,50 @@
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="stage-a-fast")
+
+from types import SimpleNamespace
+
+from miles.ray.rollout.router_manager import _maybe_enable_router_dp_aware
+
+
+def _make(dp_size: int | None, dp_aware: bool = False):
+    args = SimpleNamespace(sglang_dp_size=dp_size)
+    router_args = SimpleNamespace(dp_aware=dp_aware)
+    return args, router_args
+
+
+def test_auto_enables_when_dp_attention_on(monkeypatch):
+    monkeypatch.delenv("MILES_DISABLE_AUTO_DP_AWARE", raising=False)
+    args, router_args = _make(dp_size=4)
+    _maybe_enable_router_dp_aware(args, router_args)
+    assert router_args.dp_aware is True
+
+
+def test_noop_when_dp_size_one(monkeypatch):
+    monkeypatch.delenv("MILES_DISABLE_AUTO_DP_AWARE", raising=False)
+    args, router_args = _make(dp_size=1)
+    _maybe_enable_router_dp_aware(args, router_args)
+    assert router_args.dp_aware is False
+
+
+def test_noop_when_dp_size_none(monkeypatch):
+    monkeypatch.delenv("MILES_DISABLE_AUTO_DP_AWARE", raising=False)
+    args, router_args = _make(dp_size=None)
+    _maybe_enable_router_dp_aware(args, router_args)
+    assert router_args.dp_aware is False
+
+
+def test_respects_explicit_dp_aware(monkeypatch):
+    monkeypatch.delenv("MILES_DISABLE_AUTO_DP_AWARE", raising=False)
+    args, router_args = _make(dp_size=4, dp_aware=True)
+    _maybe_enable_router_dp_aware(args, router_args)
+    assert router_args.dp_aware is True
+
+
+def test_opt_out_warns_and_keeps_disabled(monkeypatch, caplog):
+    monkeypatch.setenv("MILES_DISABLE_AUTO_DP_AWARE", "1")
+    args, router_args = _make(dp_size=4)
+    with caplog.at_level("WARNING"):
+        _maybe_enable_router_dp_aware(args, router_args)
+    assert router_args.dp_aware is False
+    assert "MILES_DISABLE_AUTO_DP_AWARE" in caplog.text