[john-liu-dev]

Draft. Not done yet: a replay benchmark on public data, an on-hardware (end-to-end, GPU) A/B, a selection microbenchmark, and a docs entry. Their results will be added here before this is marked ready.

Overview

Under data-parallel attention, the default decode cost D + W·N ignores each rank's KV capacity: a rank can be booked past what it holds while others have room, and no single W balances both request counts and KV blocks. This PR adds dynamo-kv-capacity-aware, an opt-in builtin worker-selection policy for decode pools, selected through --router-policy-config. It places a decode request only on a rank the request fits on, judged by the rank's advertised total_kv_blocks, preferring fewer active requests and then D + W·N; when no rank fits, it picks the rank with the smallest resulting footprint instead of waiting.

Key Changes

About 300 lines of core logic: 296 in the new kv_capacity_aware.rs and 2 in lib.rs, excluding tests, comments and blank lines. They cover the selection rule, the picker, and the parameters, types and registration. No router-core change and no new dependencies.

• kv_capacity_aware.rs: the new picker.
  • A rank fits when decode_cost_blocks ≤ capacity_fraction × capacity.
  • Capacity is WorkerInfo::total_kv_blocks, or capacity_blocks when a rank does not report it.
  • When none fits, the picker takes the smallest decode_cost_blocks.
• lib.rs: registers dynamo-kv-capacity-aware in the custom catalog. Default routing is unchanged.
• Docs entry in docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md: pending.

Scope and compatibility:

• Decode pools only.
  • On prefill pools, decode_cost_blocks is 0 on every rank, so the policy degrades to fewest requests.
  • On aggregated pools, or with conditional disaggregation, it ignores the default's cache-overlap and prefill terms.
• Session affinity: it does not enforce session-affinity targets; the default policy does.
• Output blocks: booked blocks count prompts only, unless the router runs with --router-track-output-blocks.
• Where capacity comes from:
  • vLLM with data parallelism advertises the process total divided by the rank count.
  • A reported 0 counts as not reported, so those ranks are filtered only when capacity_blocks is set.
  • SGLang advertises one representative rank's value.

Parameters and an example

All parameters are optional; an unknown key fails startup.

Parameter	Default	Meaning
capacity_fraction	1.0	Share of a rank's capacity a booking may fill; null disables the filter
count_priority	true	Prefer fewer active requests before the block cost
decode_active_request_weight	0	W in D + W·N. The --router-decode-active-request-weight flag does not set it
capacity_blocks	unset	Capacity for ranks that do not report total_kv_blocks

worker_selection:
  decode: kvca
  instances:
    - name: kvca
      type: dynamo-kv-capacity-aware
      parameters:
        capacity_fraction: 1.0

Validation

Case	Default D + W·N	This policy
This request would push one rank past its capacity; others have room	can pick that rank	picks a rank it fits on
No rank has room	lowest D + W·N	smallest decode_cost_blocks; never errors or waits
capacity_fraction: null, count_priority: false	minimum D + W·N, random among ties	the same tie set
A rank reports no total_kv_blocks, or 0	capacity not used	uses capacity_blocks; with neither, not filtered, one warning per rank
Preferred-taint multipliers	scale the cost	scale the cost among fitting ranks; the fallback ignores them

Key tests and the regression each catches:

• boundary_vectors_match_the_reference_tie_sets (37 cases: at capacity, all over, unknown capacity, taints): any change in the tie set, the fallback or the warned ranks.
• degenerate_parameters_tie_exactly_like_the_default_cost: drift from the default's tie set with the filter and count priority off.
• never_fails_and_prefers_fitting_ranks: a failed pick, or a non-fitting rank beating a fitting one.
• capacity_is_total_kv_blocks_then_the_parameter: capacity_blocks overriding an advertised total_kv_blocks, or a reported 0 used as a capacity.
• resolves_alongside_the_decode_weight_flag: the policy ignored or rejected alongside the deprecated decode-weight flag.

Not run yet (draft):

• Public-data replay benchmark. Replay on main rejects custom policies, so this needs the replay support in feat(router): port published routing policies as worker-selection plugins #15450 or a separate harness.
• On-hardware A/B, end to end on GPUs.
• Selection microbenchmark: kv_capacity_aware cases in lib/router-plugins/builtin/benches/default_selection.rs, against the default decode policy.
• Docs entry.

Issues and Ongoing Work

• Closes [FEATURE]: KV-capacity-aware decode worker selection (place a request only on a rank it fits) #18223
• DEP (light): Unify router admission capacity checks #16543 plans its capacity checks on the queue side. This PR is closest to its placement room filter, with its own capacity_fraction; it never waits or rejects.
• feat(router): add a decode busy threshold to policy-class queues #16618 added active_decode_blocks() and a queueing decode busy threshold. This PR reads the same row, but counts the incoming request and only picks the rank.
• feat(router): enforce busy-reject thresholds in the KV router queue #16619 added the router-wide busy-reject thresholds, including --active-decode-blocks-threshold. Selection skips a worker past them and returns an overload error when none remains. This PR filters per request and never rejects.
• fix(router): make the per-class decode wait limit correct for disaggregated serving #18056 (open draft) makes the per-class decode wait limit work for disaggregated serving. That is a queue-side change; this PR only chooses the rank.
• fix(router): exclude ranks at their request capacity #16605 (open) caps requests per rank and rejects when all ranks are full. This PR caps KV blocks, never rejects, and runs after it.
• KLXHY-RM-3252:feat(kv-router) decode length pooling, memory guard, co… #11590 (open) adds to the router core a per-DP-rank projected-usage guard, a max_kv_cache_tokens override, an active-request decode cost and ISL rank pools. This PR offers a similar guard, a fallback capacity and request ordering as an opt-in plugin that never rejects.
• feat(kv-router): expose prefix hashes, policy class, and worker capacity to selection policies #15449 (open) proposed worker_capacity(). main now exposes that capacity through WorkerInfo, which this PR reads.
• feat(kv-router): expose per-rank resident block counts to selection plugins #16078 (open) exposes per-rank resident block counts, including cached blocks no request uses. This PR uses booked active blocks against advertised capacity.
• feat(router): port published routing policies as worker-selection plugins #15450 (open) ports published routing policies into this crate and adds replay support for catalog policies. This PR adds one more policy to the same catalog.
• TODO(rank-aware-kv-capacity) on main: total_kv_blocks may be an aggregate or representative-rank estimate, so this PR only steers with it.

Pull Request Checklist

• I used Dynamo's /pr-checklist skill.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

👋 Hi john-liu-dev! Thank you for contributing to ai-dynamo/dynamo.

Just a reminder: The NVIDIA Test Github Validation CI runs an essential subset of the testing framework to quickly catch errors.Your PR reviewers may elect to test the changes comprehensively before approving your changes.

🚀