[sgl-router] Fail fast on prefill errors and isolate PD transfer groups

experimental/sgl-router/README.md

@@ -51,6 +51,14 @@ Omit `--service-discovery-namespace` to watch all namespaces (requires
 cluster-wide RBAC). For prefill/decode disaggregation, replace `--selector`
 with `--prefill-selector` and `--decode-selector`.
 
+PD workers can be partitioned with the EndpointSlice label `sglang.ai/transfer-group`
+(or `--transfer-group-label`). Prefill and decode must have the same label value;
+unlabeled workers pair only with other unlabeled workers. The router skips prefill
+workers without a compatible decode peer and reports unready without a complete
+pair. Prefill failures cancel decode promptly: engine 4xx responses are preserved,
+while transport and server failures return 502. Prefill continues after client
+disconnect so KV transfer can finish.
+
 External KV indexer as the cache-aware signal source:
 
 ```bash

experimental/sgl-router/benches/policy_select.rs

@@ -21,6 +21,7 @@ fn workers(n: usize, model: &str) -> Vec<Arc<Worker>> {
     for i in 0..n {
         registry
             .add(WorkerSpec {
+                transfer_group: None,
                 id: WorkerId(format!("w{i}")),
                 url: format!("http://w{i}:30000"),
                 mode: WorkerMode::Plain,

experimental/sgl-router/src/buckets_reorg.rs

@@ -25,7 +25,7 @@ use std::sync::Arc;
 
 use crate::discovery::{ModelId, WorkerId};
 use crate::policies_reorg::{Pick, PickError, PickRequest, Policy, Stage};
-use crate::workers::WorkerRegistry;
+use crate::workers::{Worker, WorkerRegistry};
 
 #[derive(Debug, Clone, Copy, Default)]
 pub struct TokenLimits {
@@ -60,16 +60,28 @@ impl EngineGroup {
         workers: &WorkerRegistry,
         request: &PickRequest<'_>,
     ) -> Result<Pick, PickError> {
-        let mut engines: Vec<_> = workers
-            .healthy_workers_for(request.model)
+        self.pick_from(self.members(workers, request.model, request.stage), request)
+            .await
+    }
+
+    fn members(&self, workers: &WorkerRegistry, model: &ModelId, stage: Stage) -> Vec<Arc<Worker>> {
+        workers
+            .healthy_workers_for(model)
             .into_iter()
-            .filter(|engine| engine.mode() == request.stage)
+            .filter(|engine| engine.mode() == stage)
             .filter(|engine| {
                 self.worker_ids
                     .as_ref()
                     .is_none_or(|ids| ids.contains(&engine.id))
             })
-            .collect();
+            .collect()
+    }
+
+    async fn pick_from(
+        &self,
+        mut engines: Vec<Arc<Worker>>,
+        request: &PickRequest<'_>,
+    ) -> Result<Pick, PickError> {
         // Stable order so cursor-based policies see a consistent candidate list.
         engines.sort_by(|left, right| left.id.0.cmp(&right.id.0));
         if engines.is_empty() {
@@ -168,16 +180,35 @@ impl Bucket {
     ) -> Result<BucketPick, (Stage, PickError)> {
         let (prefill, decode) = match &self.groups {
             BucketGroups::Plain(group) => (
-                self.pick_from_group(group, Stage::Plain, workers, request)
-                    .await?,
+                self.pick_from_group(
+                    group,
+                    Stage::Plain,
+                    group.members(workers, request.model, Stage::Plain),
+                    request,
+                )
+                .await?,
                 None,
             ),
             BucketGroups::Pd { prefill, decode } => {
+                let mut decoders = decode.members(workers, request.model, Stage::Decode);
+                let mut prefills = prefill.members(workers, request.model, Stage::Prefill);
+                if prefills.is_empty() {
+                    return Err((Stage::Prefill, PickError::NoCandidates));
+                }
+                prefills.retain(|p| {
+                    decoders
+                        .iter()
+                        .any(|d| d.transfer_group == p.transfer_group)
+                });
+                if prefills.is_empty() {
+                    return Err((Stage::Decode, PickError::NoCandidates));
+                }
                 let prefill = self
-                    .pick_from_group(prefill, Stage::Prefill, workers, request)
+                    .pick_from_group(prefill, Stage::Prefill, prefills, request)
                     .await?;
+                decoders.retain(|d| d.transfer_group == prefill.engine.transfer_group);
                 let decode = self
-                    .pick_from_group(decode, Stage::Decode, workers, request)
+                    .pick_from_group(decode, Stage::Decode, decoders, request)
                     .await?;
                 (prefill, Some(decode))
             }
@@ -190,7 +221,7 @@ impl Bucket {
         &self,
         group: &EngineGroup,
         stage: Stage,
-        workers: &WorkerRegistry,
+        engines: Vec<Arc<Worker>>,
         request: &BucketRequest<'_>,
     ) -> Result<Pick, (Stage, PickError)> {
         let request = PickRequest {
@@ -205,7 +236,7 @@ impl Bucket {
             routing_key: request.routing_key,
         };
         group
-            .pick(workers, &request)
+            .pick_from(engines, &request)
             .await
             .map_err(|error| (stage, error))
     }

experimental/sgl-router/src/config/cli.rs

@@ -135,6 +135,10 @@ pub struct ServerArgs {
 
 #[derive(clap::Args, Debug)]
 pub struct DiscoveryArgs {
+    /// EndpointSlice label identifying compatible prefill/decode workers.
+    #[arg(long, default_value = "sglang.ai/transfer-group")]
+    pub transfer_group_label: String,
+
     /// Static worker URLs, space-separated or repeated. Conflicts with --service-discovery.
     #[arg(long, num_args = 1..)]
     pub worker_urls: Vec<String>,
@@ -439,6 +443,7 @@ impl DiscoveryArgs {
                     join_selector(&self.decode_selector).as_deref(),
                 )?;
                 DiscoveryBackend::K8s(K8sDiscoveryConfig {
+                    transfer_group_label: self.transfer_group_label,
                     namespace: self.service_discovery_namespace.unwrap_or_default(),
                     mode,
                 })

experimental/sgl-router/src/config/types.rs

@@ -570,6 +570,7 @@ pub struct StaticUrlsDiscoveryConfig {
 /// and bootstrap ports come from `/server_info` introspection.
 #[derive(Debug, Clone)]
 pub struct K8sDiscoveryConfig {
+    pub transfer_group_label: String,
     pub namespace: String,
     /// Resolved + validated selector mode (plain vs PD).
     pub mode: K8sDiscoveryMode,