Skip to content

fix(router): make the per-class decode wait limit correct for disaggregated serving - #108

Draft
ishandhanani wants to merge 411 commits into
idhanani/base-4cb0a16e4cfrom
idhanani/dyn-16543-disagg-decode-gate
Draft

ishandhanani wants to merge 411 commits into
idhanani/base-4cb0a16e4cfrom
idhanani/dyn-16543-disagg-decode-gate

Conversation

@ishandhanani

Copy link
Copy Markdown
Owner

Summary

PR ai-dynamo#16618 added two separate things, and this PR is about the second:

This PR implements (b) so it is correct in both topologies:

  • Aggregated: unchanged. The router's one queue checks both class fractions on the same workers.
  • Disaggregated with KV routing: the prefill router's queue is the only admission gate. It checks the class prefill fraction on the prefill workers and the class decode fraction on the decode workers. For the decode check, it reads the decode router's booked load directly, with no message to the decode actor. When the decode router frees KV, it fires a separate broadcast signal (RoomChanged) that wakes the prefill queue. That signal does not share the decode actor's notify_one CapacityUpdates. The decode router books each decode leg that follows a remote prefill (RequestPhase::Decode with staged KV) straight away: no class wait and no deadline expiry, because the leg's prefill is already running. Router-wide reject thresholds still apply when the decode worker is selected.

Why this changes behavior: disaggregated KV routing has two policy-class queues in series, prefill then decode, and both get the same policy profile. As merged, the class decode limit did nothing at the prefill router, which never tracks held blocks. At the decode router it could hold a decode leg after its prefill was already dispatched. That held prefill KV in place, and broke SGLang bootstrap, where prefill and decode must run at the same time.

Design:

  • lib/kv-router: a new PairedPoolRoom hook that the host installs on a scheduler (set_paired_pool, like set_selection_trace). In each gate pass the queue checks it together with the existing prefill check, in both the drain path and the never-queue direct path. Only classes with a decode_busy_threshold_frac use it, and without a hook nothing changes. For requests without a LoRA, one answer is cached per class per gate pass. LocalScheduler::any_rank_has_room reads booked load from the slot registry with no actor message, and subscribe_room_changed exposes the broadcast hint. The hint fires on local releases, the actor's cleanup drain, peer-router state changes, and worker config changes. ScheduleRequest / SelectionOperation carry an in-process admitted_upstream flag, so the queue books those requests directly. Wire DTOs (PreprocessedRequest, RouterConfigOverride) are unchanged.
  • lib/llm: the frontend now installs the decode RoutingHost on the PrefillRouter whenever routing is KV, not only under conditional disagg. Every prefill binding's KV router gets a lazily resolved DecodePoolRoom (it holds a Weak<KvRouter>). Until the decode host exists, the pool reports room, so nothing changes. RoutingHost maps a decode leg with staged_kv_cleanup (frontend-only, #[serde(skip)]) to AdmittedUpstream. A conditional-disagg bypass also has RequestPhase::Decode but has nothing staged, so the decode queue stays its gate and its behavior is unchanged.

Decode-room eligibility (what the prefill hop applies):

  • Applied: decode worker reachability and the request's LoRA. Both come from reusing the decode router's existing availability filter. That filter narrows LoRA candidates by the request's pin and allowed workers, which at the prefill hop name prefill workers. So a LoRA request that also carries a prefill pin or allowed set is checked on reachability alone.
  • Not applied: a decode worker pin (decode_worker_id), allowed-worker hints, routing constraints, and the KV-transfer topology domain of the chosen prefill worker.
  • Decode selection still enforces exact decode eligibility (RoutingEligibility, lib/kv-router/src/scheduling/filter.rs). So when the coarse gate admits a request, it is routed at decode exactly as before feat(router): add a decode busy threshold to policy-class queues ai-dynamo/dynamo#16618.

Known limits and follow-ups:

  • Between the decode-room check at the prefill hop and the decode booking (after prefill dispatch), a burst of prefill admissions can overshoot the decode fraction. Each admission sees decode room before the earlier legs have booked. A pending-decode reservation count would close this gap.
  • If either hop is not KV-routed, nothing checks the class decode limit.
  • Request classifiers run only on the decode/aggregated router, so a class a classifier assigns there is not known at the prefill hop. This is unchanged by this PR.

Part of ai-dynamo#16543

Validation

All commands use a per-worktree CARGO_TARGET_DIR.

  • cargo fmt --all --check: clean.
  • cargo clippy -p dynamo-kv-router -p dynamo-llm --all-targets -- -D warnings: pass, no warnings.
  • cargo test -p dynamo-kv-router --lib: 1134 passed, 0 failed, 3 ignored.
  • cargo test -p dynamo-kv-router --lib --features standalone-selection, which also runs the selection-core tests: 1393 passed, 0 failed, 3 ignored.
  • cargo test -p dynamo-llm --lib kv_router::: 467 passed, 0 failed. The new end-to-end test also passed 5 of 5 repeated runs.
  • cargo bench -p dynamo-kv-router --no-run: builds (4 bench targets). With --features standalone-selection,bench: builds (6 targets).
  • cd lib/bindings/python && cargo check: pass.
  • pre-commit run --files <changed files>, run through uvx: 11 hooks passed, the rest skipped (no matching files).
  • python3 docs/fern/scripts/docs_lint.py <changed page>: 0 errors. The 5 NAV warnings are on other, pre-existing pages.

New tests:

  • scheduling::queue::tests::paired_pool_holds_until_its_room_signal_releases_the_request: the request stays held while the paired pool has no room. Room alone does not release it; the pool's room signal does, with no other capacity change.
  • scheduling::queue::tests::paired_pool_rejects_a_never_queue_class_like_its_own_workers: the never-queue direct path and the queue agree.
  • scheduling::queue::tests::paired_pool_is_ignored_by_classes_without_a_decode_limit: a class without a decode limit never consults the hook, and its prefill limit behaves as before.
  • scheduling::queue::tests::admitted_upstream_skips_class_limits_but_not_reject_thresholds: the request is booked past a full class limit and ahead of the backlog, an already-passed deadline does not expire it, and reject thresholds still reject it.
  • scheduling::local::tests::pool_room_reads_booked_decode_blocks_and_signals_releases: the direct read honors the worker filter, and the room signal fires on release even with queueing disabled.
  • kv_router::routing_host::tests::disagg_decode_wait_limit_gates_the_prefill_hop_and_never_the_decode_hop: two real KV hosts share one policy file. Staged decode legs book past the class decode limit. A conditional-bypass-shaped decode request still waits at decode. A new prefill request waits at the prefill hop while decode is full, and freed decode KV releases it.
  • kv_router::prefill_router::tests::kv_decode_host_resolves_the_decode_pool_without_conditional_disagg.

Queue-gate benchmark for the aggregated path (selection_core no_hints_queue_gate, base 4cb0a16 vs this branch, shared host with load average 4 to 7):

workers base this branch
64 39.6 us 36.1 us
512 90.5 us 76.7 us
4096 614.6 us 608.0 us

No regression within noise. The aggregated gate only adds one per-class array reset and one early-returning check per predicate.

@github-actions github-actions Bot added fix documentation Improvements or additions to documentation router labels Oct 29, 2026
JulienDarve and others added 26 commits November 3, 2026 23:55
Signed-off-by: tanmayv25 <tanmay2592@gmail.com>
…ps serving (ai-dynamo#17973)

Signed-off-by: Julien Darve <jdarve@NVIDIA.com>
…i-dynamo#17713)

Signed-off-by: Karen Chung <karenc@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com>
…amo#17475)

Signed-off-by: Keiven Chang <keivenchang@users.noreply.github.com>
Co-authored-by: Keiven Chang <keivenchang@users.noreply.github.com>
…hot restore (ai-dynamo#17246)

Signed-off-by: Sai Krishna Pulipaka <skpulipaka@gmail.com>
Co-authored-by: Biswa Panda <biswa.panda@gmail.com>
…ynamo#17869)

Signed-off-by: Yuewei Na <nv-yna@users.noreply.github.com>
Co-authored-by: Yuewei Na <nv-yna@users.noreply.github.com>
…i-dynamo#17483)

Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: Pranjul Gupta <pranjulgupta0203@gmail.com>
Signed-off-by: Dr. Stefan Schimanski <sschimanski@nvidia.com>
Signed-off-by: devivasudevan <49675305+devivasudevan@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Pranjul Gupta <pranjulgupta0203@gmail.com>
Co-authored-by: Dr. Stefan Schimanski <sschimanski@nvidia.com>
Co-authored-by: devivasudevan <49675305+devivasudevan@users.noreply.github.com>
)

Signed-off-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com>
…i-dynamo#17988)

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
…#17710)

Signed-off-by: Alec Flowers <aflowers@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
…7.0 removal (ai-dynamo#17899)

Signed-off-by: Karen Chung <karenc@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com>
…#16871)

Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com>
Signed-off-by: Matthew Wittwer <mwittwer@nvidia.com>
Signed-off-by: furionw <qiwa@nvidia.com>
Co-authored-by: Qi Wang <qiwa@nvidia.com>
…-dynamo#17932)

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
Signed-off-by: GLAMR <svc-glamr@nvidia.com>
…gg (ai-dynamo#17977)

Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
… to 16Gi (ai-dynamo#17979)

Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Biswa Panda <biswa.panda@gmail.com>
Add a PairedPoolRoom hook a host installs on a scheduler (set_paired_pool).
A policy class with a decode_busy_threshold_frac then also waits until the
paired pool, the decode pool behind a prefill router, has a rank within that
fraction. The queue checks the hook with its own prefill check in the drain
path and in the never-queue direct path, and caches one answer per class per
gate pass for requests without a LoRA. Classes without a decode fraction
never consult it, and a scheduler with no hook is unchanged.

LocalScheduler::any_rank_has_room reads booked load from the slot registry
without a turn of the actor. subscribe_room_changed exposes a broadcast hint
(RoomChanged) fired on releases, cleanup drains, peer-router state and worker
config changes. It is separate from the actor's notify_one CapacityUpdates so
a paired queue never steals the actor's permits.

ScheduleRequest and SelectionOperation carry an in-process admitted_upstream
flag: the queue books such a request directly, with no class wait and no
deadline expiry, while selection still applies reject thresholds.

Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
…egated serving

In disaggregated KV routing the prefill and decode routers both run a
policy-class queue with the same profile. The class decode wait limit did
nothing on the prefill router, which tracks no held blocks, and on the
decode router it could hold a decode leg after its prefill was dispatched,
stranding prefill KV and breaking SGLang bootstrap.

Make the prefill hop the only gate. The frontend now installs the decode
RoutingHost on the PrefillRouter whenever routing is KV, and every prefill
binding's KV router judges class decode limits on a lazily resolved
DecodePoolRoom over the decode router: its booked load, reachable workers and
the request's LoRA through the decode router's own availability filter. Until
the decode host exists the pool reports room. RoutingHost books every decode-phase
request as AdmittedUpstream, so the decode queue never holds or expires one:
a decode leg whose prefill is running and a conditional-disagg bypass alike.
Only the prefill router sets the decode phase, so aggregated requests keep
the full class gate.

Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
State where decode_busy_threshold_frac is judged: on the same workers in
aggregated serving, and at the prefill hop against the decode pool in
disaggregated KV routing, where the decode hop books legs without waiting.
List what the pool-level decode check applies and what decode selection
still enforces.

Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
Automated refresh from the public Dynamo Google Calendar.

Generated by .github/workflows/community-events-refresh.yml.

Signed-off-by: dynamo-ops <170655669+dynamo-ops@users.noreply.github.com>
The decode hop booked every decode-phase request as AdmittedUpstream,
including a conditional-disagg bypass. A bypass never passes the prefill
queue and holds no staged KV, so it lost every class limit, queue limit
and deadline. Book directly only decode legs with staged KV; the bypass
keeps the decode queue's full class gate, as before this change.

Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
jkiczka-nvidia and others added 29 commits November 12, 2026 11:01
…eo first (ai-dynamo#18434)

Signed-off-by: Julia Kiczka <jkiczka@nvidia.com>
…amo#17991)

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
Signed-off-by: GLAMR <svc-glamr@nvidia.com>
…m path (ai-dynamo#18482)

Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
…amo#18130)

Signed-off-by: doxu <doxu@nvidia.com>
Co-authored-by: Krishnan Prashanth <140860868+KrishnanPrash@users.noreply.github.com>
…i-dynamo#18687)

Signed-off-by: Krishnan Prashanth <kprashanth@nvidia.com>
…revision labels (ai-dynamo#18322)

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
Signed-off-by: GLAMR <svc-glamr@nvidia.com>
…er (ai-dynamo#18419)

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
…sh (ai-dynamo#18585)

Signed-off-by: Krishnan Prashanth <kprashanth@nvidia.com>
…s a core plugin (ai-dynamo#17597)

Signed-off-by: nnshah1 <neelays@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
…i-dynamo#18377)

Signed-off-by: Hoang Minh An <89495585+anhalu@users.noreply.github.com>
Signed-off-by: Krishnan Prashanth <kprashanth@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Krishnan Prashanth <kprashanth@nvidia.com>
…tputs (ai-dynamo#18408)

Signed-off-by: Ryan Olson <rolson@nvidia.com>
ai-dynamo#18436)

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
Signed-off-by: GLAMR <svc-glamr@nvidia.com>
…k dynamo (ai-dynamo#18653)

Signed-off-by: Shahar Mor <smor@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Automated refresh from the public Dynamo Google Calendar.

Generated by .github/workflows/community-events-refresh.yml.

Signed-off-by: dynamo-ops <170655669+dynamo-ops@users.noreply.github.com>
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
…8151)

Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
Signed-off-by: jthomson04 <jothomson@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
…mo#18574)

Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
Co-authored-by: Claude Fable 5.2 <noreply@anthropic.com>
…Q1BXR) (ai-dynamo#18016)

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
Signed-off-by: GLAMR <svc-glamr@nvidia.com>
Co-authored-by: Indrajit Bhosale <iamindrajitb@gmail.com>
…sconnect (ai-dynamo#18572)

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
Signed-off-by: GLAMR <svc-glamr@nvidia.com>
…ests (ai-dynamo#18177)

Signed-off-by: xianlubird <xianlubird@gmail.com>
Automated refresh from the public Dynamo Google Calendar.

Generated by .github/workflows/community-events-refresh.yml.

Signed-off-by: dynamo-ops <170655669+dynamo-ops@users.noreply.github.com>
…18611)

Signed-off-by: Jan Bernloehr <jbernloehr@nvidia.com>
Automated refresh from the public Dynamo Google Calendar.

Generated by .github/workflows/community-events-refresh.yml.

Signed-off-by: dynamo-ops <170655669+dynamo-ops@users.noreply.github.com>
…lang:*-xpu base (ai-dynamo#18449)

Signed-off-by: Ellie Jan <ellie.jan@intel.com>
Co-authored-by: ellie.jan <ejan@clx4892.jf.intel.com>
…i-dynamo#18703)

Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
…i-dynamo#17311)

Signed-off-by: JasperChan929 <wc3042@nyu.edu>
Co-authored-by: jthomson04 <jwillthomson19@gmail.com>
…sagg-decode-gate

# Conflicts:
#	lib/router/scheduler/src/queue.rs

This branch had an error being deployed

1 failed deployment
external_collaborator — 45fbe15b Deployed Nov 13, 2026 by ishandhanani via ok-to-test #671
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.