fix: dp_rank always 0 in non-KV router mode

[ishandhanani]

Problem

In non-KV-router mode (e.g. NATS with round-robin), all prefill requests are routed with dp_rank=0, pinning 100% of prefill work to DP rank 0 while ranks 1-3 sit idle. This wastes a shit ton of prefill capacity.

Root cause

PR #6736 attempted to fix this with a u32::MAX sentinel value in prefill_router.rs SimpleRouter path, converted to None on the Python side. However, the sentinel never reaches the Python handler — the value arrives as 0 through the NATS transport path, bypassing the sentinel logic. The more interesting piece is that this regression only came about during benchmarking when we used the next PR #7617. Will investigate why that is in follow ups

Evidence (from GB200 cluster, 4P8D DeepSeek-R1 FP4, C=4096)

HEAD (dp_rank=0 bug):

Prefill batch count by DP rank:
  DP0: 6111 (99.2%)
  DP1: 16
  DP2: 16
  DP3: 16
Peak queue depth on DP0: 771 requests

This PR (dp_rank=None):

Prefill batch count by DP rank:
  DP0: 1936 (25.1%)
  DP1: 1931 (25.0%)
  DP2: 1923 (24.9%)
  DP3: 1929 (25.0%)
Peak queue depth: 83 requests

Decode workers are unaffected — they distribute evenly in both cases.

Fix

Change dp_rank from u32 to Option<u32> in the Rust prefill router types. This properly propagates None when no specific DP rank is selected, instead of relying on a sentinel value that gets lost in transit.

Changes

• types.rs: dp_rank: u32 → dp_rank: Option<u32> in PrefillResolveDecision
• execution.rs: Remove .unwrap_or(0) fallback, pass Option through
• mod.rs: Assign dp_rank directly as Option

Test plan

• Verified on GB200 cluster: dp_rank prints None (not 0) in prefill handler logs
• Verified prefill batches distribute evenly across all 4 DP ranks
• Verified decode workers unaffected
• Benchmark comparison (in progress)

Context

• Related: feat(sglang): disagg DP rank routing + backwards-compatible network imports #6736 (original sentinel fix attempt)
• Related: [Performance Regression] ~2x throughput drop in disaggregated PD mode with Wide-EP (DeepSeek-R1 FP4, GB200) between SGLang v0.5.8 and latest nightly sgl-project/sglang#22095 (disagg throughput regression investigation)

[coderabbitai]

Walkthrough

The changes refactor dp_rank handling in the prefill router subsystem from a mandatory u32 to Option<u32>. This involves updating function return types, enum variant fields, and propagation logic across three files within the same module. The dp_rank value is now optionally wrapped when provided or explicitly set to None when unavailable.

Changes

Cohort / File(s)	Summary
Type Signature Updates lib/llm/src/kv_router/prefill_router/execution.rs, lib/llm/src/kv_router/prefill_router/types.rs	Changed dp_rank from mandatory u32 to Option<u32> in query_prefill_worker return type and PrefillResolveDecision::Resolved enum field. Updated debug logging to use debug formatting for optional dp_rank.
Implementation and Value Handling lib/llm/src/kv_router/prefill_router/execution.rs, lib/llm/src/kv_router/prefill_router/mod.rs	Updated dp_rank propagation: wrapped in Some() for KV routing, replaced hardcoded 0 with None for simple routing, and removed fallback logic. Changed direct assignment in mod.rs from wrapped to direct value assignment.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Title check	✅ Passed	The title accurately describes the main change: fixing dp_rank (which was always 0) in non-KV router mode by making it optional instead of defaulting to 0.
Description check	✅ Passed	The PR description comprehensively covers the problem, root cause, evidence, fix, changes, and test plan with clear context and references.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

lib/llm/src/kv_router/prefill_router/execution.rs (1)

193-207: The unwrap_or(0) on dp_rank is never used in practice.

The function returns Option<(u64, u32)> containing the extracted dp_rank, but both call sites intentionally discard this value:

• mod.rs:171: Uses _worker_info (explicit discard pattern)
• execution.rs:235-240: Uses Ok(_) (discards entire result)

The return value is built but never consumed. If the PR theme is standardizing dp_rank handling, consider either removing this extraction entirely or documenting why it's kept for future use.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/kv_router/prefill_router/execution.rs` around lines 193 - 207,
The extracted dp_rank inside the disaggregated_params parsing (variable
prefill_worker_info built from disaggregated_params.get("worker_id") and fields
prefill_worker_id / prefill_dp_rank) is never used; either remove dp_rank
extraction and simplify the tuple to only return the needed worker_id or drop
the entire prefill_worker_info extraction if callers ignore the result. Update
the code in execution.rs where prefill_worker_info is constructed: remove the
prefill_dp_rank.get("prefill_dp_rank") mapping and the map/unwrap_or(0) use, and
adjust the Some((worker_id, dp_rank)) to only return Some(worker_id) (and change
the Option type accordingly) or delete the block and any references to
prefill_worker_info so no unused value is produced; if you keep the field for
future use, add a clear TODO comment near disaggregated_params /
prefill_worker_info explaining why dp_rank is retained.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@lib/llm/src/kv_router/prefill_router/execution.rs`:
- Around line 193-207: The extracted dp_rank inside the disaggregated_params
parsing (variable prefill_worker_info built from
disaggregated_params.get("worker_id") and fields prefill_worker_id /
prefill_dp_rank) is never used; either remove dp_rank extraction and simplify
the tuple to only return the needed worker_id or drop the entire
prefill_worker_info extraction if callers ignore the result. Update the code in
execution.rs where prefill_worker_info is constructed: remove the
prefill_dp_rank.get("prefill_dp_rank") mapping and the map/unwrap_or(0) use, and
adjust the Some((worker_id, dp_rank)) to only return Some(worker_id) (and change
the Option type accordingly) or delete the block and any references to
prefill_worker_info so no unused value is produced; if you keep the field for
future use, add a clear TODO comment near disaggregated_params /
prefill_worker_info explaining why dp_rank is retained.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 47d219dd-e2ae-47a4-b056-384042fd605b

📥 Commits

Reviewing files that changed from the base of the PR and between 61f1671 and b4518e7.

📒 Files selected for processing (3)

• lib/llm/src/kv_router/prefill_router/execution.rs
• lib/llm/src/kv_router/prefill_router/mod.rs
• lib/llm/src/kv_router/prefill_router/types.rs

[YAMY1234]

Verified this fix with a 1P2D disaggregated benchmark
(DeepSeek-R1 FP4, sa-bench ISL=1024 OSL=1024 C=2048):

Config	Output tok/s	Mean TTFT (ms)
main HEAD (without fix)	17,825	70,469
main HEAD + this PR	39,858	16,589

This gives ~2.2x throughput recovery and ~4.2x TTFT improvement.
Performance is fully restored to the pre-regression baseline.

Close the perf gap reported here: #7984 (comment)

[vladnosiv]

wow, what a catch, congrats

[PeaBrane]

according to codex

• PR 6736, merged on March 25, 2026, fixed the original bug by making non-KV prefill routing use “unset DP rank” instead of concrete rank 0. In its merge commit 7edb07b, the old SimpleRouter path returned u32::MAX, and the Python prefill handler already converted that sentinel back to None.
• PR 7617, merged later on March 25, 2026, refactored prefill_router.rs into modules for offline disagg replay. In that refactor, the non-KV SimpleRouter branch changed back from u32::MAX to 0. That is the real regression.
• PR 7617 also introduced the replay/NATS routing path that benchmarking exercises. At merge time, its C routing result did not even carry prefill_dp_rank, so that path could not preserve “unset rank” semantics.
• PR 7741 later added prefill_dp_rank to the C routing result, but by then the refactored router was already producing 0, so the replay path started propagating the wrong value all the way into Python.
• PR 7984 fixes the model properly in Rust: dp_rank becomes Option, non-KV SimpleRouter returns None, and only the C boundary converts None back to u32::MAX so existing Python code can translate it back to None.

regression test here
#7991

[PeaBrane]

We should plan to follow up to cleanup the dp rank plumbing in the code base

• Router bookkeeping uses None -> 0, so it loses the distinction early: push_router.rs (line 193).
• Tracker storage has a separate sentinel for “unset”, so it pretends that distinction exists internally: timing.rs (line 24).
• SGLang prefill has a special transport-level sentinel conversion back to None: prefill_handler.py (line 18), prefill_handler.py (line 141).
• Prefill result parsing sometimes defaults absent rank back to 0 anyway, so the preserved distinction is not stable end-to-end: execution.rs (line 199).