[Router] Extract worker selection into policies::selection (no behavior change) (#39322)

Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

experimental/sgl-router/src/policies/mod.rs

@@ -16,6 +16,7 @@ pub mod random;
 pub mod registry;
 pub mod round_robin;
 pub mod scoring;
+pub mod selection;
 pub mod session_aware;
 pub mod sticky;
 

experimental/sgl-router/src/policies/selection.rs

@@ -0,0 +1,845 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+//! Worker selection: the two ladders that turn a policy proposal into a
+//! committed worker — [`select_prefill_worker`] and [`select_decode_peer`].
+//!
+//! WHY this is a module rather than a block inside the chat handler: each
+//! ladder has several rungs — cache-candidate resolution, the global
+//! session-affinity probe, per-domain admission, the capacity fallback — and
+//! each rung fails into the next. Rungs get added over time, and the
+//! assertions worth writing are almost always about a ladder as a whole ("a
+//! saturated fleet still routes", "a sampled choice never lands on a rejected
+//! worker"), not about one rung in isolation. Written as closures inside an
+//! HTTP handler those assertions can only be expressed as end-to-end HTTP
+//! tests; written here they are unit tests.
+//!
+//! Both ladders live here rather than one per module because `CandidateDomain`
+//! already carries `stage: RoutingStage`: prefill and decode are two
+//! configurations of one idea, so a rung added to one and not the other has to
+//! be visible on one screen.
+//!
+//! The module owns the decision and reports why; it does not own the HTTP
+//! response. Mapping a failed selection onto a status code stays in the route.
+
+use std::sync::Arc;
+
+use crate::config::{DecodePolicyKind, PolicyKind, SessionAffinityMode};
+use crate::discovery::ModelId;
+use crate::policies::admission::{
+    resolve_cache_candidates, resolve_decode, resolve_prefill, resolve_prefill_admitted,
+    CandidateDomain, CandidateRange, DecisionReason,
+};
+use crate::policies::buckets::{BucketRequest, BucketSelector};
+use crate::policies::decode::{
+    build_decode_policy, resolve_decode_with_capacity_fallback, DecodeSelectionContext,
+};
+use crate::policies::engine_load::EngineLoadSnapshot;
+use crate::policies::{
+    ExternalPrefixSignal, Policy, PrefillProposal, ProposalKind, SelectionContext,
+};
+use crate::server::metrics::{MetricsRegistry, PolicySelectionFailureReason};
+use crate::workers::Worker;
+
+/// Everything one prefill selection reads. Collaborators first, then the
+/// per-request facts.
+pub(crate) struct PrefillSelectionInputs<'a> {
+    pub policy: &'a dyn Policy,
+    pub policy_kind: PolicyKind,
+    pub bucket_selector: &'a BucketSelector,
+    pub metrics: &'a MetricsRegistry,
+    /// Names the model in the selection context and in the log lines.
+    pub model_id: &'a ModelId,
+    pub body: Option<&'a [u8]>,
+    pub routing_key: Option<&'a str>,
+    pub session_id: Option<&'a str>,
+    pub request_input_tokens: u64,
+    pub request_tokens: Option<&'a [u32]>,
+    pub external_prefix: Option<&'a ExternalPrefixSignal>,
+    /// Required whenever `policy.uses_shared_prefill_admission()`; the
+    /// per-domain rung panics without it. `Policy::needs_load_snapshot`
+    /// defaults to `uses_shared_prefill_admission`, which is what keeps the
+    /// two in step for the ingress caller.
+    pub load_snapshot: Option<&'a EngineLoadSnapshot>,
+    pub workers: &'a [Arc<Worker>],
+    pub ttft_slo_ms: Option<u64>,
+    pub tps_slo: Option<f64>,
+    /// The configured mode. Without Bucket partitioning all modes reduce to
+    /// the single global domain, and the ladder applies that reduction itself.
+    pub session_affinity_mode: SessionAffinityMode,
+}
+
+/// Runs the prefill selection ladder. `Err` carries the reason the last rung
+/// to record one gave up with; the route maps it onto a status code.
+pub(crate) fn select_prefill_worker(
+    inputs: &PrefillSelectionInputs<'_>,
+) -> Result<Arc<Worker>, PolicySelectionFailureReason> {
+    let mut selector = Selector {
+        inputs,
+        // Prefill reserves no peak sequence room: the decode peer, not the
+        // prefill worker, holds the KV for the tokens still to be generated.
+        bucket_request: BucketRequest {
+            input_tokens: inputs.request_input_tokens,
+            expected_peak_sequence_tokens: None,
+            ttft_slo_ms: inputs.ttft_slo_ms,
+            tps_slo: inputs.tps_slo,
+        },
+        failure_reason: PolicySelectionFailureReason::ProposalEmpty,
+    };
+    let selected = selector.run();
+    selected.ok_or(selector.failure_reason)
+}
+
+/// Carries the per-request bucket request and the failure reason across rungs.
+/// Each rung that gives up overwrites the reason, so the reported one is what
+/// the last rung to record any gave — a rung that returns `None` without
+/// recording leaves the previous reason standing.
+struct Selector<'a> {
+    inputs: &'a PrefillSelectionInputs<'a>,
+    bucket_request: BucketRequest,
+    failure_reason: PolicySelectionFailureReason,
+}
+
+impl<'a> Selector<'a> {
+    fn run(&mut self) -> Option<Arc<Worker>> {
+        let inputs = self.inputs;
+        let bucket_request = self.bucket_request;
+        // Without Bucket partitioning all modes reduce to the single global
+        // domain, so reduce once here rather than trusting every caller to.
+        let session_affinity_mode = if inputs.bucket_selector.is_enabled() {
+            inputs.session_affinity_mode
+        } else {
+            SessionAffinityMode::Bucket
+        };
+        let use_global_affinity_probe = inputs.bucket_selector.is_enabled()
+            && inputs.policy.is_bucket_affinity_policy()
+            && session_affinity_mode != SessionAffinityMode::Bucket;
+
+        // Cache-Aware resolves one bounded global candidate set and returns a final winner.
+        let cache_winner = self.cache_winner();
+
+        let global_affinity_probe = use_global_affinity_probe
+            .then(|| {
+                let snapshot = inputs.load_snapshot?;
+                let global_range = CandidateRange::global(inputs.workers);
+                let probe_ctx = self
+                    .base_context(global_range.id)
+                    .with_load_snapshot(snapshot)
+                    .without_affinity_assignment();
+                inputs.policy.propose(global_range.workers, &probe_ctx)
+            })
+            .flatten();
+        // A new or stale session may create its first assignment in the target Bucket.
+        let global_affinity_missed = global_affinity_probe
+            .as_ref()
+            .is_some_and(|proposal| !matches!(proposal.kind, ProposalKind::SessionAffinity));
+        let global_affinity_worker = global_affinity_probe
+            .and_then(|proposal| {
+                matches!(proposal.kind, ProposalKind::SessionAffinity).then_some(proposal.primary)
+            })
+            .and_then(|primary| {
+                inputs.bucket_selector.prefill_affinity_domain(
+                    inputs.workers,
+                    &primary,
+                    bucket_request,
+                )
+            })
+            // Rebuild the backup inside the primary's own Bucket.
+            .and_then(|domain| self.select_in_domain(&domain, true, false, false));
+
+        cache_winner.or_else(|| {
+            // Materializing the normal domains clones the member list of every
+            // Bucket, so build them only on the rung that actually reads them.
+            let prefill_domains = || {
+                inputs
+                    .bucket_selector
+                    .prefill_domains(inputs.workers, bucket_request)
+            };
+            if inputs.policy_kind == PolicyKind::CacheAware {
+                // Cache miss or failure retries ordered domains with ordinary P2.
+                return self.select_domains(&prefill_domains(), false, false);
+            }
+            if let Some(worker) = global_affinity_worker {
+                return Some(worker);
+            }
+            match session_affinity_mode {
+                SessionAffinityMode::GlobalPreserve if global_affinity_missed => {
+                    self.select_domains(&prefill_domains(), true, true)
+                }
+                SessionAffinityMode::GlobalPreserve => {
+                    self.select_domains(&prefill_domains(), false, false)
+                }
+                SessionAffinityMode::Bucket | SessionAffinityMode::GlobalRebind => {
+                    self.select_domains(&prefill_domains(), true, true)
+                }
+            }
+        })
+    }
+
+    /// The per-request `SelectionContext` every rung starts from.
+    fn base_context(&self, candidate_range_id: &'a str) -> SelectionContext<'a> {
+        let inputs = self.inputs;
+        SelectionContext::with_routing_key(inputs.model_id, inputs.body, inputs.routing_key)
+            .with_session_id(inputs.session_id)
+            .with_candidate_range_id(candidate_range_id)
+            .with_input_tokens(inputs.request_input_tokens)
+            .with_request_tokens(inputs.request_tokens)
+            .with_external_prefix(inputs.external_prefix)
+    }
+
+    fn cache_winner(&mut self) -> Option<Arc<Worker>> {
+        let inputs = self.inputs;
+        let bucket_request = self.bucket_request;
+        if inputs.policy_kind != PolicyKind::CacheAware {
+            return None;
+        }
+        let snapshot = inputs.load_snapshot?;
+        let global_range = CandidateRange::global(inputs.workers);
+        let cache_ctx = self
+            .base_context(global_range.id)
+            .with_load_snapshot(snapshot)
+            .with_prefill_cache_bucket(inputs.bucket_selector, bucket_request);
+        let PrefillProposal::CacheCandidates(proposal) = inputs
+            .policy
+            .propose_prefill(global_range.workers, &cache_ctx)?
+        else {
+            return None;
+        };
+        let bounded_candidate_count = proposal.candidates.len();
+        let cache_decision =
+            resolve_cache_candidates(&proposal, inputs.request_input_tokens, snapshot);
+        inputs
+            .metrics
+            .record_cache_admission_evaluations(cache_decision.admission_evaluated_candidates);
+        inputs
+            .metrics
+            .record_cache_admission_rejections(cache_decision.admission_rejected_candidates);
+        inputs.metrics.record_cache_pressure_guard(
+            cache_decision.pressure_guard_compared_pairs,
+            cache_decision.pressure_guard_overrides,
+        );
+        inputs
+            .metrics
+            .record_cache_monitor_decision(cache_decision.prefill_pressure_source);
+        let Some(decision) = cache_decision.decision else {
+            self.failure_reason = PolicySelectionFailureReason::CacheCandidatesExhausted;
+            return None;
+        };
+        let selected_candidate = proposal
+            .candidates
+            .iter()
+            .find(|candidate| candidate.worker.id == decision.selected.id)?;
+        tracing::debug!(
+            model = %inputs.model_id,
+            policy = ?ProposalKind::CacheAffinity,
+            range = %decision.candidate_range_id,
+            selected = %decision.selected.url,
+            cache_candidates = bounded_candidate_count,
+            input_tokens = inputs.request_input_tokens,
+            matched_prefix_tokens = selected_candidate.matched_prefix_tokens,
+            uncached_tokens = selected_candidate.uncached_tokens,
+            reason = ?decision.reason,
+            load_snapshot_version = decision.load_snapshot_version,
+            prefill_pressure_source = cache_decision.prefill_pressure_source,
+            "cache candidate winner",
+        );
+        inputs
+            .metrics
+            .record_policy_decision("cache_aware", "cache_candidate");
+        Some(decision.selected)
+    }
+
+    /// Ordered domains, first without the capacity fallback and then with it —
+    /// a later domain that admits normally beats an earlier one that only the
+    /// capacity fallback could serve.
+    fn select_domains(
+        &mut self,
+        domains: &[CandidateDomain],
+        affinity_lookup_enabled: bool,
+        affinity_assignment_enabled: bool,
+    ) -> Option<Arc<Worker>> {
+        domains
+            .iter()
+            .find_map(|domain| {
+                self.select_in_domain(
+                    domain,
+                    affinity_lookup_enabled,
+                    affinity_assignment_enabled,
+                    false,
+                )
+            })
+            .or_else(|| {
+                domains.iter().find_map(|domain| {
+                    self.select_in_domain(
+                        domain,
+                        affinity_lookup_enabled,
+                        affinity_assignment_enabled,
+                        true,
+                    )
+                })
+            })
+    }
+
+    fn select_in_domain(
+        &mut self,
+        domain: &CandidateDomain,
+        affinity_lookup_enabled: bool,
+        affinity_assignment_enabled: bool,
+        allow_capacity_fallback: bool,
+    ) -> Option<Arc<Worker>> {
+        let inputs = self.inputs;
+        let candidate_range = domain.prefill_range()?;
+        let mut selection_ctx = self.base_context(candidate_range.id);
+        if let Some(snapshot) = inputs.load_snapshot {
+            selection_ctx = selection_ctx.with_load_snapshot(snapshot);
+        }
+        let selection_ctx = if !affinity_lookup_enabled {
+            selection_ctx.without_affinity_lookup()
+        } else if !affinity_assignment_enabled {
+            selection_ctx.without_affinity_assignment()
+        } else {
+            selection_ctx
+        };
+        let Some(PrefillProposal::Pair(proposal)) = inputs
+            .policy
+            .propose_prefill(candidate_range.workers, &selection_ctx)
+        else {
+            // Domain retries are ordinary pair proposals.
+            return None;
+        };
+        if inputs.policy.uses_shared_prefill_admission() {
+            let snapshot = inputs
+                .load_snapshot
+                .expect("shared prefill admission requires a load snapshot");
+            let decision = if allow_capacity_fallback {
+                resolve_prefill(
+                    &candidate_range,
+                    &proposal,
+                    inputs.request_input_tokens,
+                    snapshot,
+                )
+            } else {
+                resolve_prefill_admitted(
+                    &candidate_range,
+                    &proposal,
+                    inputs.request_input_tokens,
+                    snapshot,
+                )
+            };
+            let Some(decision) = decision else {
+                self.failure_reason = PolicySelectionFailureReason::PrefillAdmissionExhausted;
+                return None;
+            };
+            let reason = prefill_policy_reason(
+                inputs.policy_kind,
+                proposal.kind,
+                decision.reason,
+                inputs.session_id.is_some_and(|value| !value.is_empty()),
+                affinity_lookup_enabled,
+            );
+            inputs.policy.commit_prefill_selection(
+                &selection_ctx,
+                proposal.kind,
+                &decision.selected,
+            );
+            inputs
+                .metrics
+                .record_policy_decision(&inputs.policy_kind.to_string(), reason);
+            tracing::debug!(
+                model = %inputs.model_id,
+                policy = ?proposal.kind,
+                range = %decision.candidate_range_id,
+                primary = %decision.primary.url,
+                backup = ?decision.backup.as_ref().map(|worker| worker.url.as_str()),
+                selected = %decision.selected.url,
+                reason = ?decision.reason,
+                load_snapshot_version = decision.load_snapshot_version,
+                "prefill policy decision",
+            );
+            Some(decision.selected)
+        } else {
+            tracing::debug!(
+                model = %inputs.model_id,
+                policy = ?proposal.kind,
+                range = %candidate_range.id,
+                selected = %proposal.primary.url,
+                "prefill policy decision without shared admission",
+            );
+            Some(proposal.primary)
+        }
+    }
+}
+
+/// Return the low-cardinality reason for the final Prefill decision.
+fn prefill_policy_reason(
+    policy: PolicyKind,
+    proposal: ProposalKind,
+    decision: DecisionReason,
+    has_session_id: bool,
+    affinity_lookup_enabled: bool,
+) -> &'static str {
+    match policy {
+        PolicyKind::SessionAware => match (proposal, decision) {
+            (ProposalKind::SessionAffinity, DecisionReason::Primary) => "session_primary",
+            (ProposalKind::SessionAffinity, DecisionReason::BackupPrimaryAdmission) => {
+                "session_admission_backup"
+            }
+            (ProposalKind::SessionAffinity, DecisionReason::BackupPressureGuard) => {
+                "session_pressure_backup"
+            }
+            (ProposalKind::SessionAffinity, DecisionReason::RangeFallback) => {
+                "session_range_fallback"
+            }
+            (_, DecisionReason::CapacityFallbackPowerOfTwo) => "capacity_fallback_power_of_two",
+            (_, DecisionReason::RangeFallback) => "range_fallback",
+            (_, _) if !affinity_lookup_enabled => "range_fallback",
+            (_, _) if !has_session_id => "no_session",
+            (ProposalKind::PowerOfTwo, _) => "assigned",
+            _ => "primary",
+        },
+        PolicyKind::CacheAware => match (proposal, decision) {
+            (_, DecisionReason::CacheCandidate)
+            | (ProposalKind::CacheAffinity, DecisionReason::Primary) => "cache_candidate",
+            (_, DecisionReason::Primary) => "no_cache_candidate",
+            (_, DecisionReason::BackupPrimaryAdmission) => "no_cache_candidate_admission_backup",
+            (_, DecisionReason::BackupPressureGuard) => "no_cache_candidate_pressure_backup",
+            (_, DecisionReason::RangeFallback) => "no_cache_candidate_range_fallback",
+            (_, DecisionReason::CapacityFallbackPowerOfTwo) => {
+                "no_cache_candidate_capacity_fallback_power_of_two"
+            }
+        },
+        _ => match decision {
+            DecisionReason::Primary => "primary",
+            DecisionReason::CacheCandidate => "cache_candidate",
+            DecisionReason::BackupPrimaryAdmission => "admission_backup",
+            DecisionReason::BackupPressureGuard => "pressure_backup",
+            DecisionReason::RangeFallback => "range_fallback",
+            DecisionReason::CapacityFallbackPowerOfTwo => "capacity_fallback_power_of_two",
+        },
+    }
+}
+
+/// Everything one decode-peer selection reads. Collaborators first, then the
+/// per-request facts.
+pub(crate) struct DecodeSelectionInputs<'a> {
+    pub decode_policy_kind: DecodePolicyKind,
+    pub bucket_selector: &'a BucketSelector,
+    /// Names the model in the log lines.
+    pub model_id: &'a ModelId,
+    /// URL of the committed prefill worker; `legacy_host_affinity` pairs the
+    /// decode peer against it.
+    pub prefill_url: &'a str,
+    pub decode_workers: &'a [Arc<Worker>],
+    pub request_input_tokens: u64,
+    pub requested_max_output_tokens: Option<u64>,
+    pub ttft_slo_ms: Option<u64>,
+    pub tps_slo: Option<f64>,
+    /// Required: every rung resolves its proposal against the snapshot, so
+    /// without one the ladder reports no peer at all rather than picking one
+    /// blind.
+    pub load_snapshot: Option<&'a EngineLoadSnapshot>,
+}
+
+/// Runs the decode selection ladder.
+///
+/// Two rungs, and the order is the point: every Bucket domain is tried on the
+/// strict rung before any domain is retried with the capacity fallback, so a
+/// domain that fits the request is never passed over in favour of an earlier
+/// domain that only fits it by relaxing capacity.
+///
+/// `None` means no domain yielded a peer; the route maps that onto its own
+/// status code, as it does for prefill.
+pub(crate) fn select_decode_peer(inputs: &DecodeSelectionInputs<'_>) -> Option<Arc<Worker>> {
+    // Every rung resolves against the snapshot, so without one no domain can
+    // yield a peer and there is nothing worth building.
+    let snapshot = inputs.load_snapshot?;
+    let request_kv_tokens = projected_decode_kv_tokens(
+        inputs.request_input_tokens,
+        inputs.requested_max_output_tokens,
+    );
+    // Only an explicit output budget justifies reserving peak sequence room;
+    // without one the projection degenerates to the input length and would
+    // bucket every request as if it decoded nothing.
+    let expected_peak_sequence_tokens = inputs
+        .requested_max_output_tokens
+        .map(|_| request_kv_tokens);
+    let decode_domains = inputs.bucket_selector.decode_domains(
+        inputs.decode_workers,
+        BucketRequest {
+            input_tokens: inputs.request_input_tokens,
+            expected_peak_sequence_tokens,
+            ttft_slo_ms: inputs.ttft_slo_ms,
+            tps_slo: inputs.tps_slo,
+        },
+    );
+    let decode_policy = build_decode_policy(inputs.decode_policy_kind);
+    let decode_ctx = DecodeSelectionContext::new()
+        .with_load_snapshot(snapshot)
+        .with_prefill_url(inputs.prefill_url);
+    let select_in_domain = |decode_domain: &CandidateDomain, allow_capacity_fallback: bool| {
+        let decode_proposal = decode_policy.propose(decode_domain, &decode_ctx)?;
+        let decode_decision = if allow_capacity_fallback {
+            resolve_decode_with_capacity_fallback(
+                decode_domain,
+                &decode_proposal,
+                request_kv_tokens,
+                snapshot,
+            )
+        } else {
+            resolve_decode(decode_domain, &decode_proposal, request_kv_tokens, snapshot)
+        }?;
+        tracing::debug!(
+            model = %inputs.model_id,
+            policy = ?inputs.decode_policy_kind,
+            range = %decode_decision.candidate_range_id,
+            primary = %decode_decision.primary.url,
+            backup = ?decode_decision.backup.as_ref().map(|worker| worker.url.as_str()),
+            selected = %decode_decision.selected.url,
+            reason = ?decode_decision.reason,
+            load_snapshot_version = decode_decision.load_snapshot_version,
+            "decode policy decision",
+        );
+        Some(decode_decision.selected)
+    };
+    decode_domains
+        .iter()
+        .find_map(|domain| select_in_domain(domain, false))
+        .or_else(|| {
+            decode_domains
+                .iter()
+                .find_map(|domain| select_in_domain(domain, true))
+        })
+}
+
+/// Project the peak sequence length without integer wraparound.
+fn projected_decode_kv_tokens(input_tokens: u64, max_output_tokens: Option<u64>) -> u64 {
+    max_output_tokens.map_or(input_tokens, |output_tokens| {
+        input_tokens.saturating_add(output_tokens)
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{
+        prefill_policy_reason, projected_decode_kv_tokens, select_decode_peer,
+        select_prefill_worker, DecodeSelectionInputs, PrefillSelectionInputs,
+    };
+    use crate::config::{DecodePolicyKind, PolicyKind, SessionAffinityMode};
+    use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
+    use crate::policies::admission::{resolve_prefill_admitted, CandidateRange, DecisionReason};
+    use crate::policies::buckets::BucketSelector;
+    use crate::policies::engine_load::{EngineLoadSnapshot, NativeCacheWorkerLoad};
+    use crate::policies::power_of_two::PowerOfTwoChoicesPolicy;
+    use crate::policies::{Policy, ProposalKind, SelectionProposal};
+    use crate::server::metrics::{MetricsRegistry, PolicySelectionFailureReason};
+    use crate::workers::Worker;
+    use std::sync::Arc;
+    use std::time::Instant;
+
+    fn worker(id: &str) -> Arc<Worker> {
+        Arc::new(Worker::new(WorkerSpec {
+            id: WorkerId(id.into()),
+            url: format!("http://{id}:30000"),
+            mode: WorkerMode::Plain,
+            model_ids: vec![ModelId("model".into())],
+            bootstrap_port: None,
+        }))
+    }
+
+    /// `(worker, tokens already held, published KV capacity)`.
+    fn snapshot(entries: &[(&Arc<Worker>, u64, u64)]) -> EngineLoadSnapshot {
+        EngineLoadSnapshot::from_native_cache_workers(
+            7,
+            entries
+                .iter()
+                .map(|(worker, used, capacity)| {
+                    (
+                        worker.url.clone(),
+                        NativeCacheWorkerLoad {
+                            num_running_reqs: 0,
+                            num_waiting_reqs: 0,
+                            num_waiting_uncached_tokens: 0,
+                            num_used_tokens: *used,
+                            num_total_tokens: *used,
+                            max_total_num_tokens: *capacity,
+                            max_running_requests: 64,
+                            prefill_throughput_tokens_per_s: None,
+                            estimated_prefill_queue_ms: None,
+                            captured_at: Instant::now(),
+                        },
+                    )
+                })
+                .collect(),
+        )
+    }
+
+    /// Power-of-two prefill, Bucket partitioning off, no session affinity —
+    /// the single global domain, which is what isolates the ladder's rungs.
+    #[allow(clippy::too_many_arguments)]
+    fn prefill_inputs<'a>(
+        policy: &'a dyn Policy,
+        bucket_selector: &'a BucketSelector,
+        metrics: &'a MetricsRegistry,
+        model_id: &'a ModelId,
+        workers: &'a [Arc<Worker>],
+        load_snapshot: Option<&'a EngineLoadSnapshot>,
+        request_input_tokens: u64,
+    ) -> PrefillSelectionInputs<'a> {
+        PrefillSelectionInputs {
+            policy,
+            policy_kind: PolicyKind::PowerOfTwo,
+            bucket_selector,
+            metrics,
+            model_id,
+            body: None,
+            routing_key: None,
+            session_id: None,
+            request_input_tokens,
+            request_tokens: None,
+            external_prefix: None,
+            load_snapshot,
+            workers,
+            ttft_slo_ms: None,
+            tps_slo: None,
+            session_affinity_mode: SessionAffinityMode::Bucket,
+        }
+    }
+
+    fn decode_inputs<'a>(
+        bucket_selector: &'a BucketSelector,
+        model_id: &'a ModelId,
+        decode_workers: &'a [Arc<Worker>],
+        load_snapshot: Option<&'a EngineLoadSnapshot>,
+        request_input_tokens: u64,
+    ) -> DecodeSelectionInputs<'a> {
+        DecodeSelectionInputs {
+            decode_policy_kind: DecodePolicyKind::PowerOfTwo,
+            bucket_selector,
+            model_id,
+            prefill_url: "http://prefill:30000",
+            decode_workers,
+            request_input_tokens,
+            requested_max_output_tokens: None,
+            ttft_slo_ms: None,
+            tps_slo: None,
+            load_snapshot,
+        }
+    }
+
+    #[test]
+    fn an_empty_fleet_reports_the_proposal_empty_failure() {
+        let policy = PowerOfTwoChoicesPolicy::new();
+        let buckets = BucketSelector::new(None);
+        let metrics = MetricsRegistry::new();
+        let model = ModelId("model".into());
+        let workers: Vec<Arc<Worker>> = Vec::new();
+        let loads = snapshot(&[]);
+
+        let outcome = select_prefill_worker(&prefill_inputs(
+            &policy,
+            &buckets,
+            &metrics,
+            &model,
+            &workers,
+            Some(&loads),
+            64,
+        ));
+
+        assert!(matches!(
+            outcome,
+            Err(PolicySelectionFailureReason::ProposalEmpty)
+        ));
+    }
+
+    #[test]
+    fn a_saturated_fleet_still_routes_through_the_capacity_fallback() {
+        let full = worker("full");
+        let also_full = worker("also-full");
+        let workers = vec![Arc::clone(&full), Arc::clone(&also_full)];
+        let loads = snapshot(&[(&full, 100, 100), (&also_full, 100, 100)]);
+
+        // Without this the test would pass even if the strict rung had served
+        // the request, and would prove nothing about the fallback rung.
+        let range = CandidateRange::global(&workers);
+        assert!(
+            resolve_prefill_admitted(
+                &range,
+                &SelectionProposal::with_backup(Arc::clone(&full), Arc::clone(&also_full)),
+                64,
+                &loads,
+            )
+            .is_none(),
+            "fixture must saturate every worker so the strict rung admits none",
+        );
+
+        let policy = PowerOfTwoChoicesPolicy::new();
+        let buckets = BucketSelector::new(None);
+        let metrics = MetricsRegistry::new();
+        let model = ModelId("model".into());
+
+        let selected = select_prefill_worker(&prefill_inputs(
+            &policy,
+            &buckets,
+            &metrics,
+            &model,
+            &workers,
+            Some(&loads),
+            64,
+        ))
+        .expect("the capacity fallback must still place the request");
+        assert!(selected.id == full.id || selected.id == also_full.id);
+    }
+
+    #[test]
+    fn a_sampled_choice_never_lands_on_a_rejected_worker() {
+        let full = worker("full");
+        let roomy = worker("roomy");
+        let workers = vec![Arc::clone(&full), Arc::clone(&roomy)];
+        let loads = snapshot(&[(&full, 100, 100), (&roomy, 0, 100_000)]);
+        let policy = PowerOfTwoChoicesPolicy::new();
+        let buckets = BucketSelector::new(None);
+        let metrics = MetricsRegistry::new();
+        let model = ModelId("model".into());
+
+        // Power-of-two samples its pair at random, so one pass proves nothing
+        // about which rung answered.
+        for _ in 0..32 {
+            let selected = select_prefill_worker(&prefill_inputs(
+                &policy,
+                &buckets,
+                &metrics,
+                &model,
+                &workers,
+                Some(&loads),
+                64,
+            ))
+            .expect("one worker has room");
+            assert_eq!(selected.id, roomy.id);
+        }
+    }
+
+    #[test]
+    fn the_decode_ladder_reports_no_peer_without_a_load_snapshot() {
+        let peer = worker("decode");
+        let workers = vec![Arc::clone(&peer)];
+        let buckets = BucketSelector::new(None);
+        let model = ModelId("model".into());
+
+        assert!(select_decode_peer(&decode_inputs(&buckets, &model, &workers, None, 64)).is_none());
+        assert!(
+            select_decode_peer(&decode_inputs(
+                &buckets,
+                &model,
+                &workers,
+                Some(&snapshot(&[(&peer, 0, 100_000)])),
+                64,
+            ))
+            .is_some(),
+            "the same fleet routes as soon as a snapshot is available",
+        );
+    }
+
+    #[test]
+    fn a_saturated_decode_fleet_still_routes_through_the_capacity_fallback() {
+        let full = worker("decode-full");
+        let also_full = worker("decode-also-full");
+        let workers = vec![Arc::clone(&full), Arc::clone(&also_full)];
+        let loads = snapshot(&[(&full, 100, 100), (&also_full, 100, 100)]);
+        let buckets = BucketSelector::new(None);
+        let model = ModelId("model".into());
+
+        let selected =
+            select_decode_peer(&decode_inputs(&buckets, &model, &workers, Some(&loads), 64))
+                .expect("the capacity fallback must still place the decode peer");
+        assert!(selected.id == full.id || selected.id == also_full.id);
+    }
+
+    #[test]
+    fn decode_kv_projection_includes_the_explicit_output_budget() {
+        assert_eq!(projected_decode_kv_tokens(1_024, Some(512)), 1_536);
+        assert_eq!(projected_decode_kv_tokens(1_024, None), 1_024);
+        assert_eq!(projected_decode_kv_tokens(u64::MAX - 1, Some(8)), u64::MAX);
+    }
+
+    #[test]
+    fn session_reason_distinguishes_hit_assignment_and_keyless_fallback() {
+        assert_eq!(
+            prefill_policy_reason(
+                PolicyKind::SessionAware,
+                ProposalKind::SessionAffinity,
+                DecisionReason::Primary,
+                true,
+                true,
+            ),
+            "session_primary"
+        );
+        assert_eq!(
+            prefill_policy_reason(
+                PolicyKind::SessionAware,
+                ProposalKind::PowerOfTwo,
+                DecisionReason::Primary,
+                true,
+                true,
+            ),
+            "assigned"
+        );
+        assert_eq!(
+            prefill_policy_reason(
+                PolicyKind::SessionAware,
+                ProposalKind::PowerOfTwo,
+                DecisionReason::Primary,
+                false,
+                true,
+            ),
+            "no_session"
+        );
+    }
+
+    #[test]
+    fn session_reason_preserves_admission_and_pressure_escapes() {
+        assert_eq!(
+            prefill_policy_reason(
+                PolicyKind::SessionAware,
+                ProposalKind::SessionAffinity,
+                DecisionReason::BackupPrimaryAdmission,
+                true,
+                true,
+            ),
+            "session_admission_backup"
+        );
+        assert_eq!(
+            prefill_policy_reason(
+                PolicyKind::SessionAware,
+                ProposalKind::SessionAffinity,
+                DecisionReason::BackupPressureGuard,
+                true,
+                true,
+            ),
+            "session_pressure_backup"
+        );
+    }
+
+    #[test]
+    fn cache_no_winner_p2_is_distinct_from_cache_candidate() {
+        assert_eq!(
+            prefill_policy_reason(
+                PolicyKind::CacheAware,
+                ProposalKind::PowerOfTwo,
+                DecisionReason::Primary,
+                false,
+                false,
+            ),
+            "no_cache_candidate"
+        );
+        assert_eq!(
+            prefill_policy_reason(
+                PolicyKind::CacheAware,
+                ProposalKind::CacheAffinity,
+                DecisionReason::Primary,
+                false,
+                true,
+            ),
+            "cache_candidate"
+        );
+    }
+}

experimental/sgl-router/src/server/routes/chat.rs

@@ -1,22 +1,14 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::config::{PolicyKind, SessionAffinityMode};
+use crate::config::SessionAffinityMode;
 use crate::discovery::{ModelId, WorkerMode};
-use crate::policies::admission::{
-    resolve_cache_candidates, resolve_decode, resolve_prefill, resolve_prefill_admitted,
-    CandidateDomain, CandidateRange, DecisionReason,
-};
-use crate::policies::buckets::BucketRequest;
-use crate::policies::decode::{
-    build_decode_policy, resolve_decode_with_capacity_fallback, DecodeSelectionContext,
-};
 use crate::policies::kv_events::{compute_block_hashes, compute_block_hashes_bigram};
 use crate::policies::registry::{PdPoolResolver, PdResolveError};
-use crate::policies::{
+use crate::policies::selection::{
-    request_tokens_for, ExternalPrefixSignal, PrefillProposal, ProposalKind, RequestTokens,
+    select_decode_peer, select_prefill_worker, DecodeSelectionInputs, PrefillSelectionInputs,
-    SelectionContext,
 };
+use crate::policies::{request_tokens_for, ExternalPrefixSignal, RequestTokens};
 use crate::proxy::sse::StreamEnd;
 use crate::server::app_context::AppContext;
 use crate::server::error::ApiError;
@@ -31,7 +23,6 @@ use axum::http::{HeaderMap, HeaderName, HeaderValue, Response};
 use bytes::Bytes;
 use serde::de::IgnoredAny;
 use serde::Deserialize;
-use std::cell::Cell;
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -59,55 +50,6 @@ const X_SGL_TPS_SLO: HeaderName = HeaderName::from_static("x-sgl-tps-slo");
 /// purpose.
 const CHARS_PER_TOKEN_ESTIMATE: usize = 4;
 
-/// Return the low-cardinality reason for the final Prefill decision.
-fn prefill_policy_reason(
-    policy: PolicyKind,
-    proposal: ProposalKind,
-    decision: DecisionReason,
-    has_session_id: bool,
-    affinity_lookup_enabled: bool,
-) -> &'static str {
-    match policy {
-        PolicyKind::SessionAware => match (proposal, decision) {
-            (ProposalKind::SessionAffinity, DecisionReason::Primary) => "session_primary",
-            (ProposalKind::SessionAffinity, DecisionReason::BackupPrimaryAdmission) => {
-                "session_admission_backup"
-            }
-            (ProposalKind::SessionAffinity, DecisionReason::BackupPressureGuard) => {
-                "session_pressure_backup"
-            }
-            (ProposalKind::SessionAffinity, DecisionReason::RangeFallback) => {
-                "session_range_fallback"
-            }
-            (_, DecisionReason::CapacityFallbackPowerOfTwo) => "capacity_fallback_power_of_two",
-            (_, DecisionReason::RangeFallback) => "range_fallback",
-            (_, _) if !affinity_lookup_enabled => "range_fallback",
-            (_, _) if !has_session_id => "no_session",
-            (ProposalKind::PowerOfTwo, _) => "assigned",
-            _ => "primary",
-        },
-        PolicyKind::CacheAware => match (proposal, decision) {
-            (_, DecisionReason::CacheCandidate)
-            | (ProposalKind::CacheAffinity, DecisionReason::Primary) => "cache_candidate",
-            (_, DecisionReason::Primary) => "no_cache_candidate",
-            (_, DecisionReason::BackupPrimaryAdmission) => "no_cache_candidate_admission_backup",
-            (_, DecisionReason::BackupPressureGuard) => "no_cache_candidate_pressure_backup",
-            (_, DecisionReason::RangeFallback) => "no_cache_candidate_range_fallback",
-            (_, DecisionReason::CapacityFallbackPowerOfTwo) => {
-                "no_cache_candidate_capacity_fallback_power_of_two"
-            }
-        },
-        _ => match decision {
-            DecisionReason::Primary => "primary",
-            DecisionReason::CacheCandidate => "cache_candidate",
-            DecisionReason::BackupPrimaryAdmission => "admission_backup",
-            DecisionReason::BackupPressureGuard => "pressure_backup",
-            DecisionReason::RangeFallback => "range_fallback",
-            DecisionReason::CapacityFallbackPowerOfTwo => "capacity_fallback_power_of_two",
-        },
-    }
-}
-
 /// Maximum buffered chat-completions body (32MiB). Sized for base64 multimodal inputs;
 /// enforced by the `DefaultBodyLimit`, and returns 413 PAYLOAD_TOO_LARGE.
 pub const MAX_CHAT_BODY_BYTES: usize = 32 << 20;
@@ -143,13 +85,6 @@ impl RequestProbe {
     }
 }
 
-/// Project the peak sequence length without integer wraparound.
-fn projected_decode_kv_tokens(input_tokens: u64, max_output_tokens: Option<u64>) -> u64 {
-    max_output_tokens.map_or(input_tokens, |output_tokens| {
-        input_tokens.saturating_add(output_tokens)
-    })
-}
-
 /// RAII guard that records `sgl_router_request_duration_seconds` when
 /// dropped. For streaming requests the handler returns at response-headers
 /// time, so recording end-to-end latency at the dispatch site would capture
@@ -345,266 +280,35 @@ pub async fn chat_completions(
         .and_then(|config| headers.get(config.session_id_header.as_str()))
         .and_then(|value| value.to_str().ok())
         .filter(|value| !value.is_empty());
-    // Each Bucket retry rebuilds the proposal and reruns Admission/Guard.
+    // `select_prefill_worker` reduces this to `Bucket` when Bucket
-    let prefill_bucket_request = BucketRequest {
+    // partitioning is off.
-        input_tokens: request_input_tokens,
+    let session_affinity_mode = ctx
-        expected_peak_sequence_tokens: None,
-        ttft_slo_ms,
-        tps_slo,
-    };
-    let configured_session_affinity_mode = ctx
         .config
         .model
         .affinity
         .as_ref()
         .map(|config| config.session_affinity_mode)
         .unwrap_or(SessionAffinityMode::Bucket);
-    // Without Bucket partitioning all modes reduce to the single global domain.
+    // Each Bucket retry rebuilds the proposal and reruns Admission/Guard.
-    let session_affinity_mode = if ctx.bucket_selector.is_enabled() {
+    let worker = select_prefill_worker(&PrefillSelectionInputs {
-        configured_session_affinity_mode
+        policy: policy.as_ref(),
-    } else {
+        policy_kind: ctx.config.model.policy,
-        SessionAffinityMode::Bucket
+        bucket_selector: ctx.bucket_selector.as_ref(),
-    };
+        metrics: ctx.metrics.as_ref(),
-    let use_global_affinity_probe = ctx.bucket_selector.is_enabled()
+        model_id: &model_id,
-        && policy.is_bucket_affinity_policy()
+        body: Some(&body),
-        && session_affinity_mode != SessionAffinityMode::Bucket;
+        routing_key,
-    let worker = {
+        session_id,
-        let selection_failure_reason = Cell::new(PolicySelectionFailureReason::ProposalEmpty);
-        let select_prefill_in_domain = |domain: &CandidateDomain,
-                                        affinity_lookup_enabled: bool,
-                                        affinity_assignment_enabled: bool,
-                                        allow_capacity_fallback: bool|
-         -> Option<Arc<Worker>> {
-            let candidate_range = domain.prefill_range()?;
-            let mut selection_ctx =
-                SelectionContext::with_routing_key(&model_id, Some(&body), routing_key)
-                    .with_session_id(session_id)
-                    .with_candidate_range_id(candidate_range.id)
-                    .with_input_tokens(request_input_tokens)
-                    .with_request_tokens(
-                        request_tokens.as_ref().map(|tokens| tokens.ids.as_slice()),
-                    )
-                    .with_external_prefix(external_prefix.as_ref());
-            if let Some(snapshot) = load_snapshot.as_ref() {
-                selection_ctx = selection_ctx.with_load_snapshot(snapshot);
-            }
-            let selection_ctx = if !affinity_lookup_enabled {
-                selection_ctx.without_affinity_lookup()
-            } else if !affinity_assignment_enabled {
-                selection_ctx.without_affinity_assignment()
-            } else {
-                selection_ctx
-            };
-            let Some(PrefillProposal::Pair(proposal)) =
-                policy.propose_prefill(candidate_range.workers, &selection_ctx)
-            else {
-                // Domain retries are ordinary pair proposals.
-                return None;
-            };
-            if policy.uses_shared_prefill_admission() {
-                let snapshot = load_snapshot
-                    .as_ref()
-                    .expect("shared prefill admission requires a load snapshot");
-                let decision = if allow_capacity_fallback {
-                    resolve_prefill(&candidate_range, &proposal, request_input_tokens, snapshot)
-                } else {
-                    resolve_prefill_admitted(
-                        &candidate_range,
-                        &proposal,
         request_input_tokens,
-                        snapshot,
+        request_tokens: request_tokens.as_ref().map(|tokens| tokens.ids.as_slice()),
-                    )
+        external_prefix: external_prefix.as_ref(),
-                };
+        load_snapshot: load_snapshot.as_ref(),
-                let Some(decision) = decision else {
+        workers: &workers,
-                    selection_failure_reason
+        ttft_slo_ms,
-                        .set(PolicySelectionFailureReason::PrefillAdmissionExhausted);
+        tps_slo,
-                    return None;
+        session_affinity_mode,
-                };
-                let reason = prefill_policy_reason(
-                    ctx.config.model.policy,
-                    proposal.kind,
-                    decision.reason,
-                    session_id.is_some_and(|value| !value.is_empty()),
-                    affinity_lookup_enabled,
-                );
-                policy.commit_prefill_selection(&selection_ctx, proposal.kind, &decision.selected);
-                ctx.metrics
-                    .record_policy_decision(&ctx.config.model.policy.to_string(), reason);
-                tracing::debug!(
-                    model = %model_str,
-                    policy = ?proposal.kind,
-                    range = %decision.candidate_range_id,
-                    primary = %decision.primary.url,
-                    backup = ?decision.backup.as_ref().map(|worker| worker.url.as_str()),
-                    selected = %decision.selected.url,
-                    reason = ?decision.reason,
-                    load_snapshot_version = decision.load_snapshot_version,
-                    "prefill policy decision",
-                );
-                Some(decision.selected)
-            } else {
-                tracing::debug!(
-                    model = %model_str,
-                    policy = ?proposal.kind,
-                    range = %candidate_range.id,
-                    selected = %proposal.primary.url,
-                    "prefill policy decision without shared admission",
-                );
-                Some(proposal.primary)
-            }
-        };
-        let select_prefill_domains =
-            |domains: &[CandidateDomain],
-             affinity_lookup_enabled: bool,
-             affinity_assignment_enabled: bool| {
-                domains
-                    .iter()
-                    .find_map(|domain| {
-                        select_prefill_in_domain(
-                            domain,
-                            affinity_lookup_enabled,
-                            affinity_assignment_enabled,
-                            false,
-                        )
     })
-                    .or_else(|| {
+    .map_err(|reason| policy_selection_failed(&ctx, &model_str, reason))?;
-                        domains.iter().find_map(|domain| {
-                            select_prefill_in_domain(
-                                domain,
-                                affinity_lookup_enabled,
-                                affinity_assignment_enabled,
-                                true,
-                            )
-                        })
-                    })
-            };
-
-        // Cache-Aware resolves one bounded global candidate set and returns a final winner.
-        let cache_winner = (ctx.config.model.policy == PolicyKind::CacheAware)
-            .then(|| {
-                let snapshot = load_snapshot.as_ref()?;
-                let global_range = CandidateRange::global(&workers);
-                let cache_ctx =
-                    SelectionContext::with_routing_key(&model_id, Some(&body), routing_key)
-                        .with_session_id(session_id)
-                        .with_candidate_range_id(global_range.id)
-                        .with_input_tokens(request_input_tokens)
-                        .with_request_tokens(
-                            request_tokens.as_ref().map(|tokens| tokens.ids.as_slice()),
-                        )
-                        .with_external_prefix(external_prefix.as_ref())
-                        .with_load_snapshot(snapshot)
-                        .with_prefill_cache_bucket(&ctx.bucket_selector, prefill_bucket_request);
-                let PrefillProposal::CacheCandidates(proposal) =
-                    policy.propose_prefill(global_range.workers, &cache_ctx)?
-                else {
-                    return None;
-                };
-                let bounded_candidate_count = proposal.candidates.len();
-                let cache_decision =
-                    resolve_cache_candidates(&proposal, request_input_tokens, snapshot);
-                ctx.metrics.record_cache_admission_evaluations(
-                    cache_decision.admission_evaluated_candidates,
-                );
-                ctx.metrics.record_cache_admission_rejections(
-                    cache_decision.admission_rejected_candidates,
-                );
-                ctx.metrics.record_cache_pressure_guard(
-                    cache_decision.pressure_guard_compared_pairs,
-                    cache_decision.pressure_guard_overrides,
-                );
-                ctx.metrics
-                    .record_cache_monitor_decision(cache_decision.prefill_pressure_source);
-                let Some(decision) = cache_decision.decision else {
-                    selection_failure_reason
-                        .set(PolicySelectionFailureReason::CacheCandidatesExhausted);
-                    return None;
-                };
-                let selected_candidate = proposal
-                    .candidates
-                    .iter()
-                    .find(|candidate| candidate.worker.id == decision.selected.id)?;
-                tracing::debug!(
-                    model = %model_str,
-                    policy = ?ProposalKind::CacheAffinity,
-                    range = %decision.candidate_range_id,
-                    selected = %decision.selected.url,
-                    cache_candidates = bounded_candidate_count,
-                    input_tokens = request_input_tokens,
-                    matched_prefix_tokens = selected_candidate.matched_prefix_tokens,
-                    uncached_tokens = selected_candidate.uncached_tokens,
-                    reason = ?decision.reason,
-                    load_snapshot_version = decision.load_snapshot_version,
-                    prefill_pressure_source = cache_decision.prefill_pressure_source,
-                    "cache candidate winner",
-                );
-                ctx.metrics
-                    .record_policy_decision("cache_aware", "cache_candidate");
-                Some(decision.selected)
-            })
-            .flatten();
-
-        let global_affinity_probe = use_global_affinity_probe
-            .then(|| {
-                let snapshot = load_snapshot.as_ref()?;
-                let global_range = CandidateRange::global(&workers);
-                let probe_ctx =
-                    SelectionContext::with_routing_key(&model_id, Some(&body), routing_key)
-                        .with_session_id(session_id)
-                        .with_candidate_range_id(global_range.id)
-                        .with_input_tokens(request_input_tokens)
-                        .with_request_tokens(
-                            request_tokens.as_ref().map(|tokens| tokens.ids.as_slice()),
-                        )
-                        .with_external_prefix(external_prefix.as_ref())
-                        .with_load_snapshot(snapshot)
-                        .without_affinity_assignment();
-                policy.propose(global_range.workers, &probe_ctx)
-            })
-            .flatten();
-        // A new or stale session may create its first assignment in the target Bucket.
-        let global_affinity_missed = global_affinity_probe
-            .as_ref()
-            .is_some_and(|proposal| !matches!(proposal.kind, ProposalKind::SessionAffinity));
-        let global_affinity_worker = global_affinity_probe
-            .and_then(|proposal| {
-                matches!(proposal.kind, ProposalKind::SessionAffinity).then_some(proposal.primary)
-            })
-            .and_then(|primary| {
-                ctx.bucket_selector.prefill_affinity_domain(
-                    &workers,
-                    &primary,
-                    prefill_bucket_request,
-                )
-            })
-            // Rebuild the backup inside the primary's own Bucket.
-            .and_then(|domain| select_prefill_in_domain(&domain, true, false, false));
-        cache_winner
-            .or_else(|| {
-                // Materialize normal domains only when Cache-Aware has no winner.
-                let prefill_domains = ctx
-                    .bucket_selector
-                    .prefill_domains(&workers, prefill_bucket_request);
-                if ctx.config.model.policy == PolicyKind::CacheAware {
-                    // Cache miss or failure retries ordered domains with ordinary P2.
-                    return select_prefill_domains(&prefill_domains, false, false);
-                }
-                global_affinity_worker.or_else(|| match session_affinity_mode {
-                    SessionAffinityMode::GlobalPreserve if global_affinity_missed => {
-                        select_prefill_domains(&prefill_domains, true, true)
-                    }
-                    SessionAffinityMode::GlobalPreserve => {
-                        select_prefill_domains(&prefill_domains, false, false)
-                    }
-                    SessionAffinityMode::Bucket | SessionAffinityMode::GlobalRebind => {
-                        select_prefill_domains(&prefill_domains, true, true)
-                    }
-                })
-            })
-            .ok_or_else(|| {
-                policy_selection_failed(&ctx, &model_str, selection_failure_reason.get())
-            })?
-    };
 
     // Decode selection starts after Final P.
     //
@@ -624,61 +328,23 @@ pub async fn chat_completions(
                 model: model_str.clone(),
             },
         })?;
-        let request_kv_tokens =
+        Some(
-            projected_decode_kv_tokens(request_input_tokens, requested_max_output_tokens);
+            select_decode_peer(&DecodeSelectionInputs {
-        let expected_peak_sequence_tokens = requested_max_output_tokens.map(|_| request_kv_tokens);
+                decode_policy_kind: ctx.config.model.decode_policy,
-        let decode_domains = ctx.bucket_selector.decode_domains(
+                bucket_selector: ctx.bucket_selector.as_ref(),
-            &decode_workers,
+                model_id: &model_id,
-            BucketRequest {
+                prefill_url: &worker.url,
-                input_tokens: request_input_tokens,
+                decode_workers: &decode_workers,
-                expected_peak_sequence_tokens,
+                request_input_tokens,
+                requested_max_output_tokens,
                 ttft_slo_ms,
                 tps_slo,
-            },
+                load_snapshot: load_snapshot.as_ref(),
-        );
-        let decode_policy = build_decode_policy(ctx.config.model.decode_policy);
-        let select_decode_in_domain =
-            |decode_domain: &CandidateDomain, allow_capacity_fallback: bool| {
-                let snapshot = load_snapshot.as_ref()?;
-                let decode_ctx = DecodeSelectionContext::new()
-                    .with_load_snapshot(snapshot)
-                    .with_prefill_url(&worker.url);
-                let decode_proposal = decode_policy.propose(decode_domain, &decode_ctx)?;
-                let decode_decision = if allow_capacity_fallback {
-                    resolve_decode_with_capacity_fallback(
-                        decode_domain,
-                        &decode_proposal,
-                        request_kv_tokens,
-                        snapshot,
-                    )
-                } else {
-                    resolve_decode(decode_domain, &decode_proposal, request_kv_tokens, snapshot)
-                }?;
-                tracing::debug!(
-                    model = %model_str,
-                    policy = ?ctx.config.model.decode_policy,
-                    range = %decode_decision.candidate_range_id,
-                    primary = %decode_decision.primary.url,
-                    backup = ?decode_decision.backup.as_ref().map(|worker| worker.url.as_str()),
-                    selected = %decode_decision.selected.url,
-                    reason = ?decode_decision.reason,
-                    load_snapshot_version = decode_decision.load_snapshot_version,
-                    "decode policy decision",
-                );
-                Some(decode_decision.selected)
-            };
-        decode_domains
-            .iter()
-            .find_map(|domain| select_decode_in_domain(domain, false))
-            .or_else(|| {
-                decode_domains
-                    .iter()
-                    .find_map(|domain| select_decode_in_domain(domain, true))
             })
             .ok_or_else(|| ApiError::NoDecodeWorkersAvailable {
                 model: model_str.clone(),
-            })
+            })?,
-            .map(Some)?
+        )
     } else {
         None
     };
@@ -1498,87 +1164,6 @@ mod tests {
         assert!(!should_tokenize_request(false, false, false));
     }
 
-    #[test]
-    fn session_reason_distinguishes_hit_assignment_and_keyless_fallback() {
-        assert_eq!(
-            prefill_policy_reason(
-                PolicyKind::SessionAware,
-                ProposalKind::SessionAffinity,
-                DecisionReason::Primary,
-                true,
-                true,
-            ),
-            "session_primary"
-        );
-        assert_eq!(
-            prefill_policy_reason(
-                PolicyKind::SessionAware,
-                ProposalKind::PowerOfTwo,
-                DecisionReason::Primary,
-                true,
-                true,
-            ),
-            "assigned"
-        );
-        assert_eq!(
-            prefill_policy_reason(
-                PolicyKind::SessionAware,
-                ProposalKind::PowerOfTwo,
-                DecisionReason::Primary,
-                false,
-                true,
-            ),
-            "no_session"
-        );
-    }
-
-    #[test]
-    fn session_reason_preserves_admission_and_pressure_escapes() {
-        assert_eq!(
-            prefill_policy_reason(
-                PolicyKind::SessionAware,
-                ProposalKind::SessionAffinity,
-                DecisionReason::BackupPrimaryAdmission,
-                true,
-                true,
-            ),
-            "session_admission_backup"
-        );
-        assert_eq!(
-            prefill_policy_reason(
-                PolicyKind::SessionAware,
-                ProposalKind::SessionAffinity,
-                DecisionReason::BackupPressureGuard,
-                true,
-                true,
-            ),
-            "session_pressure_backup"
-        );
-    }
-
-    #[test]
-    fn cache_no_winner_p2_is_distinct_from_cache_candidate() {
-        assert_eq!(
-            prefill_policy_reason(
-                PolicyKind::CacheAware,
-                ProposalKind::PowerOfTwo,
-                DecisionReason::Primary,
-                false,
-                false,
-            ),
-            "no_cache_candidate"
-        );
-        assert_eq!(
-            prefill_policy_reason(
-                PolicyKind::CacheAware,
-                ProposalKind::CacheAffinity,
-                DecisionReason::Primary,
-                false,
-                true,
-            ),
-            "cache_candidate"
-        );
-    }
     /// `generate_room_id` MUST return values in `[0, i64::MAX]`. The
     /// SGLang prefill stores `bootstrap_room` as `torch.int64`; a u64
     /// with the top bit set would wrap negative on the engine side.
@@ -1867,13 +1452,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn decode_kv_projection_includes_the_explicit_output_budget() {
-        assert_eq!(projected_decode_kv_tokens(1_024, Some(512)), 1_536);
-        assert_eq!(projected_decode_kv_tokens(1_024, None), 1_024);
-        assert_eq!(projected_decode_kv_tokens(u64::MAX - 1, Some(8)), u64::MAX);
-    }
-
     #[test]
     fn parse_probe_rejects_non_object_shapes() {
         // Pin the contract: degenerate JSON (valid JSON but wrong shape)