dp-attention: add prefix_match load balance for in-instance cache-aware routing

python/sglang/srt/managers/data_parallel_controller.py

@@ -14,6 +14,7 @@
 """A controller that dispatches requests to multiple data parallel workers."""
 
 import faulthandler
+import hashlib
 import logging
 import multiprocessing as mp
 import signal
@@ -76,6 +77,7 @@ class LoadBalanceMethod(Enum):
     FOLLOW_BOOTSTRAP_ROOM = auto()
     TOTAL_REQUESTS = auto()
     TOTAL_TOKENS = auto()
+    PREFIX_MATCH = auto()
 
     @classmethod
     def from_str(cls, method: str):
@@ -149,6 +151,7 @@ def __init__(
             LoadBalanceMethod.FOLLOW_BOOTSTRAP_ROOM: self.follow_bootstrap_room_scheduler,
             LoadBalanceMethod.TOTAL_REQUESTS: self.total_requests_scheduler,
             LoadBalanceMethod.TOTAL_TOKENS: self.total_tokens_scheduler,
+            LoadBalanceMethod.PREFIX_MATCH: self.prefix_match_scheduler,
         }
         self.dispatching = dispatch_lookup[self.load_balance_method]
 
@@ -607,6 +610,33 @@ def total_tokens_scheduler(self, req: Req):
         )
         self.workers[target_worker].send_pyobj(req)
 
+    def prefix_match_scheduler(self, req: Req):
+        """Route by stable hash of the first N input tokens.
+
+        Same prefix (e.g. shared system+tools) -> same DP rank -> radix cache
+        hit. Different prefixes spread evenly across ranks via stable hash.
+        No cross-rank state required; the decision is a pure function of the
+        request's leading tokens, so concurrent requests with the same prefix
+        all land on the same rank without coordination.
+        """
+        if self.maybe_external_dp_rank_routing(req):
+            return
+
+        dp_size = len(self.workers)
+        N = 4096
+        ids = getattr(req, "input_ids", None) or []
+        if len(ids) == 0:
+            target_rank = self.round_robin_counter % dp_size
+            self.round_robin_counter = (self.round_robin_counter + 1) % dp_size
+        else:
+            head = ids[:N]
+            h = hashlib.blake2b(
+                bytes(b & 0xFF for b in head[:256]) + len(head).to_bytes(4, "little"),
+                digest_size=8,
+            ).digest()
+            target_rank = int.from_bytes(h, "little") % dp_size
+        self.workers[target_rank].send_pyobj(req)
+
     def event_loop(self):
         while True:
             while True:

python/sglang/srt/server_args.py

@@ -5291,6 +5291,7 @@ def add_cli_args(parser: argparse.ArgumentParser):
                 "follow_bootstrap_room",
                 "total_requests",
                 "total_tokens",
+                "prefix_match",
             ],
         )
         parser.add_argument(

test/registered/dp_attn/test_dp_attention.py

@@ -96,6 +96,46 @@ def tearDownClass(cls):
         kill_process_tree(cls.process.pid)
 
 
+class TestDPAttentionPrefixMatchLoadBalance(
+    CustomTestCase,
+    GSM8KMixin,
+):
+    """Smoke test for --load-balance-method prefix_match.
+
+    The new method routes by a stable hash of the leading input tokens, so
+    requests sharing a prefix land on the same DP rank. This test only
+    asserts that the server boots with the option and that GSM8K accuracy
+    stays above the same threshold the other DP-attention configs use; the
+    routing decision is a pure function and does not affect inference output.
+    """
+
+    gsm8k_accuracy_thres = 0.6
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--tp",
+                "2",
+                "--enable-dp-attention",
+                "--dp",
+                "2",
+                "--load-balance-method",
+                "prefix_match",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+
 class TestDPRetract(
     CustomTestCase,
     JSONConstrainedMixin,

test/registered/unit/managers/test_prefix_match_scheduler.py

@@ -0,0 +1,103 @@
+"""Unit tests for the prefix_match load balance method.
+
+The prefix_match scheduler routes requests by a stable hash of the leading
+input tokens, so requests sharing a prefix land on the same DP rank. These
+tests assert determinism (same prefix -> same rank), even spread across
+distinct prefixes, and the empty-input fallback path -- all without booting
+a real server.
+"""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.managers.data_parallel_controller import (
+    DataParallelController,
+    LoadBalanceMethod,
+)
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _make_controller(dp_size: int = 8) -> DataParallelController:
+    """Construct a controller with just enough state to exercise the scheduler.
+
+    The real ``__init__`` spawns workers and binds ZMQ sockets, which we do
+    not need for a routing-only unit test. We bypass it and stitch in the
+    minimum attributes the scheduler reads.
+    """
+    ctrl = DataParallelController.__new__(DataParallelController)
+    ctrl.server_args = SimpleNamespace(dp_size=dp_size)
+    ctrl.workers = [MagicMock(name=f"worker-{i}") for i in range(dp_size)]
+    ctrl.round_robin_counter = 0
+    return ctrl
+
+
+def _make_req(input_ids, routed_dp_rank=None):
+    """Build a minimal req object recognised by the scheduler."""
+    return SimpleNamespace(input_ids=input_ids, routed_dp_rank=routed_dp_rank)
+
+
+class TestPrefixMatchScheduler(CustomTestCase):
+    def test_enum_registered(self):
+        self.assertIs(
+            LoadBalanceMethod.from_str("prefix_match"),
+            LoadBalanceMethod.PREFIX_MATCH,
+        )
+
+    def test_same_prefix_routes_to_same_rank(self):
+        ctrl = _make_controller(dp_size=16)
+        prefix = list(range(1, 4001))
+        req_a = _make_req(prefix + [9001, 9002])
+        req_b = _make_req(prefix + [7777])
+
+        ctrl.prefix_match_scheduler(req_a)
+        ctrl.prefix_match_scheduler(req_b)
+
+        sent_a = [i for i, w in enumerate(ctrl.workers) if w.send_pyobj.called]
+        called_on_b = [
+            i for i, w in enumerate(ctrl.workers) if w.send_pyobj.call_count >= 2
+        ]
+        self.assertEqual(len(sent_a), 1)
+        self.assertEqual(sent_a, called_on_b)
+
+    def test_distinct_prefixes_spread_across_ranks(self):
+        dp_size = 16
+        ctrl = _make_controller(dp_size=dp_size)
+        # 256 distinct prefixes -> we expect coverage of most ranks.
+        for seed in range(256):
+            req = _make_req([seed] * 4096)
+            ctrl.prefix_match_scheduler(req)
+
+        used = sum(1 for w in ctrl.workers if w.send_pyobj.called)
+        # With 256 prefixes uniformly hashed into 16 buckets the empty-bucket
+        # probability is vanishingly small; require at least 12/16 to allow
+        # generous slack while still catching a degenerate hash.
+        self.assertGreaterEqual(used, 12)
+
+    def test_empty_input_falls_back_to_round_robin(self):
+        dp_size = 4
+        ctrl = _make_controller(dp_size=dp_size)
+        for _ in range(dp_size):
+            ctrl.prefix_match_scheduler(_make_req([]))
+        for w in ctrl.workers:
+            self.assertEqual(w.send_pyobj.call_count, 1)
+
+    def test_external_routed_dp_rank_takes_precedence(self):
+        ctrl = _make_controller(dp_size=8)
+        req = _make_req([1, 2, 3, 4], routed_dp_rank=5)
+        ctrl.prefix_match_scheduler(req)
+        for i, w in enumerate(ctrl.workers):
+            if i == 5:
+                w.send_pyobj.assert_called_once_with(req)
+            else:
+                w.send_pyobj.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()