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

[jianzhichun]

Closes #26611.

Motivation

The single-instance DP controller currently ships only round_robin / follow_bootstrap_room / total_requests / total_tokens. None of them route by request content, so when a single client sends many sequential turns over a shared system + tools prefix (agent serving), the radix cache on each DP rank gets a hit at best once every dp_size requests under round_robin.

Existing cache-aware routing work targets either the sgl-router Rust gateway (#2114, #6869) or P/D-disaggregated decode (#26046, #26561). Neither covers plain --enable-dp-attention --dp-size N serving with no gateway in front: the tokenizer manager dispatches directly to workers, and the only knob is --load-balance-method.

total_tokens accidentally approximates prefix locality only because the just-finished rank tends to be the lowest-loaded, and only for strictly sequential traffic. It degrades under concurrency.

Modifications

• python/sglang/srt/managers/data_parallel_controller.py
  • Add LoadBalanceMethod.PREFIX_MATCH to the enum.
  • Register it in dispatch_lookup.
  • Add DataParallelController.prefix_match_scheduler(req): hashes the first 4k input tokens via blake2b over a 256-token window and routes by int.from_bytes(h) % dp_size. Empty input_ids falls back to round-robin; the existing routed_dp_rank override is preserved (via maybe_external_dp_rank_routing).
  • Add import hashlib.
• python/sglang/srt/server_args.py
  • Add "prefix_match" to the --load-balance-method argparse choices.
• test/registered/dp_attn/test_dp_attention.py
  • New TestDPAttentionPrefixMatchLoadBalance: launches the standard MLA test model with --enable-dp-attention --dp 2 --load-balance-method prefix_match and asserts GSM8K accuracy stays above the same 0.6 threshold as the sibling DP-attention tests.
• test/registered/unit/managers/test_prefix_match_scheduler.py
  • New CPU-only unit test: same prefix -> same rank; distinct prefixes spread to >= 12/16 ranks across 256 prefixes; empty input falls back to round-robin; external routed_dp_rank takes precedence.

The decision is a pure function of req.input_ids, so concurrent requests with the same prefix all land on the same rank with no cross-rank coordination.

Accuracy Tests

No inference output change; load balance method is a routing-only decision over which DP rank receives a request. The new registered DP-attention test (TestDPAttentionPrefixMatchLoadBalance) exercises GSM8K through the new code path with --dp 2 --load-balance-method prefix_match and asserts the same gsm8k_accuracy_thres = 0.6 threshold as the existing TestDPAttentionDP2TP2 / TestDPAttentionMixedChunk tests.

Speed Tests

Measured on a Qwen3-class MoE checkpoint, tp_size=16 dp_size=16 ep_size=16, moe_a2a_backend=deepep, attention_backend=fa3, mem_fraction_static=0.78, schedule_policy=lpm. Workload: one client sending a 22,651-token shared prefix + a short user turn, temperature=0, chat_template_kwargs.enable_thinking=false.

load_balance_method	First request TTFT (cold)	Subsequent same-prefix TTFT
round_robin	7400 ms	7400 ms every request -- each request lands on a different DP rank with a cold radix cache
total_tokens	7400 ms	2740 ms -- only holds for strictly sequential traffic
prefix_match (this PR)	7400 ms	1290 ms steady-state (p50 1292, p99 6012 from cold, 30 requests over 180 s, 0 errors)

That's 5.7x speedup vs round_robin on the exact pattern that matters for agent serving. A 4-agent concurrent same-prefix bench (with enable_mixed_chunk=true) goes from c4 RPS 0.19 (round_robin) to 0.41 (prefix_match).

Checklist

• Format with pre-commit run --files
• New tests added (registered DP-attention smoke + CPU unit)
• No inference output change; routing-only

---

CI States

Latest PR Test (Base): ❌ Run #26603772659
Latest PR Test (Extra): ❌ Run #26603772607

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[jianzhichun]

Hi maintainers 👋

The CI failures here are pr-gate / pr-gate: Require run-ci label — the actual test suite hasn't run yet. Could a maintainer add the run-ci label to trigger CI?

Summary:

• Adds prefix_match load-balance method to data_parallel_controller.py for single-instance DP attention routing
• When --enable-dp-attention, requests sharing a common prefix (e.g. same system prompt + codebase context across a multi-agent session) are routed deterministically to the same DP rank, turning random radix-cache hits into near-deterministic ones
• Unit tests + dp_attention integration test included

Thanks for reviewing!

[github-actions]

Thanks @jianzhichun. Closing this because it has had no updates in 102 days.

Reopen it if the work is still relevant.

Some directories moved recently, so an older branch may need retargeting:
sgl-kernel/ -> python/sglang/kernels/aot/, python/sglang/jit_kernel/
-> python/sglang/kernels/jit/, docs/ -> docs/docs/ (.mdx),
bench_serving.py -> benchmark/serving.py, test/srt/ -> test/registered/.