feat(router): add per-worker PD Prefill admission

bindings/golang/src/policy.rs

@@ -172,6 +172,14 @@ impl Worker for GrpcWorker {
         self.load.fetch_add(1, Ordering::Relaxed);
     }
 
+    fn try_increment_load(&self, max: usize) -> bool {
+        self.load
+            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
+                current.checked_add(1).filter(|next| *next <= max)
+            })
+            .is_ok()
+    }
+
     fn decrement_load(&self) {
         self.load.fetch_sub(1, Ordering::Relaxed);
     }

bindings/python/src/lib.rs

@@ -541,6 +541,9 @@ struct Router {
     rdma_listen_ip: Option<String>,
     rdma_slot_ttl_s: Option<u64>,
     log_mm_timing: bool,
+    prefill_max_inflight_requests_per_worker: i32,
+    prefill_queue_size: Option<usize>,
+    prefill_queue_timeout_secs: Option<u64>,
 }
 
 impl Router {
@@ -878,6 +881,9 @@ impl Router {
             .max_concurrent_requests(self.max_concurrent_requests)
             .queue_size(self.queue_size)
             .queue_timeout_secs(self.queue_timeout_secs)
+            .prefill_max_inflight_requests_per_worker(self.prefill_max_inflight_requests_per_worker)
+            .prefill_queue_size(self.prefill_queue_size)
+            .prefill_queue_timeout_secs(self.prefill_queue_timeout_secs)
             .cors_allowed_origins(self.cors_allowed_origins.clone())
             .retry_config(config::RetryConfig {
                 max_retries: self.retry_max_retries,
@@ -1151,6 +1157,9 @@ impl Router {
         rdma_listen_ip = None,
         rdma_slot_ttl_s = None,
         log_mm_timing = false,
+        prefill_max_inflight_requests_per_worker = -1,
+        prefill_queue_size = None,
+        prefill_queue_timeout_secs = None,
     ))]
     #[expect(clippy::too_many_arguments)]
     #[expect(
@@ -1317,6 +1326,9 @@ impl Router {
         rdma_listen_ip: Option<String>,
         rdma_slot_ttl_s: Option<u64>,
         log_mm_timing: bool,
+        prefill_max_inflight_requests_per_worker: i32,
+        prefill_queue_size: Option<usize>,
+        prefill_queue_timeout_secs: Option<u64>,
     ) -> PyResult<Self> {
         let mut all_urls = worker_urls.clone();
 
@@ -1495,6 +1507,9 @@ impl Router {
             rdma_listen_ip,
             rdma_slot_ttl_s,
             log_mm_timing,
+            prefill_max_inflight_requests_per_worker,
+            prefill_queue_size,
+            prefill_queue_timeout_secs,
         })
     }
 

bindings/python/src/smg/router.py

@@ -251,6 +251,15 @@ class Router:
             only). If not specified, uses the main policy. Default: None
         decode_policy: Specific load balancing policy for decode nodes (PD mode only).
             If not specified, uses the main policy. Default: None
+        prefill_max_inflight_requests_per_worker: Maximum in-flight Prefill requests
+            per worker in PD or EPD mode. A non-positive value disables the limit.
+            Default: -1
+        prefill_queue_size: Maximum number of requests waiting for Prefill admission.
+            Defaults to 100 when Prefill admission is enabled. Set to 0 to disable
+            waiting. Default: None
+        prefill_queue_timeout_secs: Maximum time in seconds a request may wait for
+            Prefill admission. Defaults to 60 when Prefill admission is enabled.
+            Default: None
         request_id_headers: List of HTTP headers to check for request IDs. If not
             specified, uses common defaults: ['x-request-id', 'x-correlation-id',
             'x-trace-id', 'request-id']. Example: ['x-my-request-id',

bindings/python/src/smg/router_args.py

@@ -283,6 +283,9 @@ class RouterArgs:
     rdma_slot_ttl_s: int | None = None
     # Per-request multimodal timing at INFO; False falls back to SMG_LOG_MM_TIMING
     log_mm_timing: bool = False
+    prefill_max_inflight_requests_per_worker: int = -1
+    prefill_queue_size: int | None = None
+    prefill_queue_timeout_secs: int | None = None
 
     @staticmethod
     def add_cli_args(
@@ -874,6 +877,33 @@ def add_cli_args(
             metavar=("URL",),
             help="Decode server URL. Can be specified multiple times.",
         )
+        pd_group.add_argument(
+            f"--{prefix}prefill-max-inflight-requests-per-worker",
+            type=int,
+            default=RouterArgs.prefill_max_inflight_requests_per_worker,
+            help=(
+                "Maximum in-flight Prefill requests per worker in PD or EPD mode."
+                " A non-positive value disables the limit (default: -1)."
+            ),
+        )
+        pd_group.add_argument(
+            f"--{prefix}prefill-queue-size",
+            type=int,
+            default=RouterArgs.prefill_queue_size,
+            help=(
+                "Maximum number of requests waiting for Prefill admission."
+                " Defaults to 100 when Prefill admission is enabled; 0 disables waiting."
+            ),
+        )
+        pd_group.add_argument(
+            f"--{prefix}prefill-queue-timeout-secs",
+            type=int,
+            default=RouterArgs.prefill_queue_timeout_secs,
+            help=(
+                "Maximum time in seconds a request may wait for Prefill admission."
+                " Defaults to 60 when Prefill admission is enabled."
+            ),
+        )
         pd_group.add_argument(
             f"--{prefix}worker-startup-timeout-secs",
             type=int,

bindings/python/tests/test_arg_parser.py

@@ -838,6 +838,29 @@ def test_parse_pd_args_with_new_policies(self):
         assert router_args.prefill_policy == "consistent_hashing"
         assert router_args.decode_policy == "prefix_hash"
 
+    def test_parse_pd_prefill_admission_args(self):
+        """Test parsing explicit Prefill admission options."""
+        args = [
+            "--pd-disaggregation",
+            "--prefill",
+            "http://prefill1:8000",
+            "none",
+            "--decode",
+            "http://decode1:8001",
+            "--prefill-max-inflight-requests-per-worker",
+            "7",
+            "--prefill-queue-size",
+            "13",
+            "--prefill-queue-timeout-secs",
+            "17",
+        ]
+
+        router_args = parse_router_args(args)
+
+        assert router_args.prefill_max_inflight_requests_per_worker == 7
+        assert router_args.prefill_queue_size == 13
+        assert router_args.prefill_queue_timeout_secs == 17
+
     def test_parse_service_discovery_args(self):
         """Test parsing service discovery arguments."""
         args_a = [
@@ -1515,6 +1538,9 @@ class TestRouterArgsFieldOrder:
         "rdma_listen_ip",
         "rdma_slot_ttl_s",
         "log_mm_timing",
+        "prefill_max_inflight_requests_per_worker",
+        "prefill_queue_size",
+        "prefill_queue_timeout_secs",
     ]
 
     def test_complete_field_sequence_is_frozen(self):

bindings/python/tests/test_startup_sequence.py

@@ -697,6 +697,9 @@ def start(self):
         prefill_selector=None,
         decode_selector=None,
         cors_allowed_origins=None,
+        prefill_max_inflight_requests_per_worker=7,
+        prefill_queue_size=13,
+        prefill_queue_timeout_secs=17,
     )
 
     r = Router.from_args(args)
@@ -707,6 +710,9 @@ def start(self):
     assert captured["decode_selector"] is None
     assert captured["cors_allowed_origins"] is None
     assert captured["worker_urls"] == ["http://w1:8000"]
+    assert captured["prefill_max_inflight_requests_per_worker"] == 7
+    assert captured["prefill_queue_size"] == 13
+    assert captured["prefill_queue_timeout_secs"] == 17
     from smg.smg_rs import PolicyType
 
     assert captured["policy"] == PolicyType.RoundRobin

model_gateway/src/app_context.rs

@@ -11,6 +11,7 @@ use smg_data_connector::{
     StorageFactoryConfig,
 };
 use smg_mcp::McpOrchestrator;
+use tokio::sync::broadcast::error::RecvError;
 use tool_parser::ParserFactory as ToolParserFactory;
 use tracing::debug;
 
@@ -28,7 +29,10 @@ use crate::{
         grpc::multimodal::MultimodalConfigRegistry,
     },
     wasm::{config::WasmRuntimeConfig, module_manager::WasmModuleManager},
-    worker::{KvEventMonitor, WorkerHttpClientCache, WorkerMonitor, WorkerRegistry, WorkerService},
+    worker::{
+        KvEventMonitor, PrefillAdmission, WorkerHttpClientCache, WorkerMonitor, WorkerRegistry,
+        WorkerService,
+    },
     workflow::{JobQueue, WorkflowEngines},
 };
 
@@ -66,6 +70,7 @@ pub struct AppContext {
     pub reasoning_parser_factory: Option<ReasoningParserFactory>,
     pub tool_parser_factory: Option<ToolParserFactory>,
     pub worker_registry: Arc<WorkerRegistry>,
+    pub prefill_admission: Option<Arc<PrefillAdmission>>,
     pub policy_registry: Arc<PolicyRegistry>,
     pub gateway: Option<Arc<Gateway>>,
     pub response_storage: Arc<dyn ResponseStorage>,
@@ -102,6 +107,28 @@ impl std::fmt::Debug for AppContext {
     }
 }
 
+fn start_prefill_admission_notifier(
+    worker_registry: &WorkerRegistry,
+    admission: &Arc<PrefillAdmission>,
+) -> Result<(), AppContextBuildError> {
+    let mut events = worker_registry.subscribe_events();
+    let admission = Arc::downgrade(admission);
+    let runtime = tokio::runtime::Handle::try_current().map_err(|error| {
+        AppContextBuildError::InvalidConfig(format!(
+            "Prefill admission requires a Tokio runtime: {error}"
+        ))
+    })?;
+    runtime.spawn(async move {
+        while let Ok(_) | Err(RecvError::Lagged(_)) = events.recv().await {
+            let Some(admission) = admission.upgrade() else {
+                break;
+            };
+            admission.notify_capacity_changed();
+        }
+    });
+    Ok(())
+}
+
 pub struct AppContextBuilder {
     client: Option<Client>,
     router_config: Option<RouterConfig>,
@@ -360,6 +387,24 @@ impl AppContextBuilder {
             .worker_job_queue
             .ok_or(AppContextBuildError::MissingField("worker_job_queue"))?;
 
+        let prefill_admission =
+            usize::try_from(router_config.prefill_max_inflight_requests_per_worker)
+                .ok()
+                .filter(|max| *max > 0)
+                .map(|max| {
+                    Arc::new(PrefillAdmission::new(
+                        max,
+                        router_config.effective_prefill_queue_size(),
+                        Duration::from_secs(router_config.effective_prefill_queue_timeout_secs()),
+                    ))
+                });
+        if let Some(admission) = prefill_admission
+            .as_ref()
+            .filter(|_| router_config.effective_prefill_queue_size() > 0)
+        {
+            start_prefill_admission_notifier(&worker_registry, admission)?;
+        }
+
         // Create WorkerService from the already-built components
         let worker_service = Arc::new(WorkerService::new(
             worker_registry.clone(),
@@ -390,6 +435,7 @@ impl AppContextBuilder {
             reasoning_parser_factory: self.reasoning_parser_factory,
             tool_parser_factory: self.tool_parser_factory,
             worker_registry,
+            prefill_admission,
             policy_registry: self
                 .policy_registry
                 .ok_or(AppContextBuildError::MissingField("policy_registry"))?,
@@ -786,8 +832,15 @@ fn mcp_bootstrap_config(file: Option<&smg_mcp::McpConfig>) -> smg_mcp::McpConfig
 
 #[cfg(test)]
 mod tests {
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    use openai_protocol::worker::{HealthCheckConfig, WorkerStatus};
+
     use super::*;
-    use crate::config::types::PolicyConfig;
+    use crate::{
+        config::types::PolicyConfig,
+        worker::{BasicWorkerBuilder, PrefillAdmissionAttempt, Worker, WorkerType},
+    };
 
     /// Loopback echo server; axum::serve accepts HTTP/1.1 and prior-knowledge
     /// h2c on the same listener, mirroring a dual-protocol engine.
@@ -1117,4 +1170,95 @@ policy:
             Some("http://env-proxy.example:3128")
         );
     }
+
+    fn prefill_worker(url: &str) -> Arc<dyn Worker> {
+        Arc::new(
+            BasicWorkerBuilder::new(url)
+                .worker_type(WorkerType::Prefill)
+                .health_config(HealthCheckConfig {
+                    disable_health_check: true,
+                    ..Default::default()
+                })
+                .build(),
+        )
+    }
+
+    #[tokio::test]
+    async fn worker_status_change_wakes_prefill_admission() {
+        let registry = Arc::new(WorkerRegistry::new());
+        let admission = Arc::new(PrefillAdmission::new(1, 1, Duration::from_secs(1)));
+        start_prefill_admission_notifier(&registry, &admission).unwrap();
+
+        let first = prefill_worker("http://prefill-1:8000");
+        registry.register(Arc::clone(&first)).unwrap();
+        let occupied = admission
+            .admit(None, {
+                let first = Arc::clone(&first);
+                move |capacity| capacity.select(Arc::clone(&first), ())
+            })
+            .await
+            .unwrap();
+
+        let attempts = Arc::new(AtomicUsize::new(0));
+        #[expect(
+            clippy::disallowed_methods,
+            reason = "test waiter must run concurrently with worker status changes"
+        )]
+        let waiting = tokio::spawn({
+            let admission = Arc::clone(&admission);
+            let registry = Arc::clone(&registry);
+            let attempts = Arc::clone(&attempts);
+            async move {
+                admission
+                    .admit(None, |capacity| {
+                        attempts.fetch_add(1, Ordering::Relaxed);
+                        let workers = registry.get_by_type(WorkerType::Prefill);
+                        if workers.is_empty() {
+                            return PrefillAdmissionAttempt::Unavailable;
+                        }
+                        workers
+                            .iter()
+                            .find(|worker| worker.is_available() && capacity.has_capacity(worker))
+                            .map_or(PrefillAdmissionAttempt::AtCapacity, |worker| {
+                                capacity.select(Arc::clone(worker), Arc::clone(worker))
+                            })
+                    })
+                    .await
+            }
+        });
+
+        tokio::time::timeout(Duration::from_secs(1), async {
+            while admission.queued_requests() != 1 {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .unwrap();
+
+        let second: Arc<dyn Worker> = Arc::new(
+            BasicWorkerBuilder::new("http://prefill-2:8000")
+                .worker_type(WorkerType::Prefill)
+                .build(),
+        );
+        let second_id = registry.register(Arc::clone(&second)).unwrap();
+        tokio::time::timeout(Duration::from_secs(1), async {
+            while attempts.load(Ordering::Relaxed) < 2 {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .unwrap();
+        assert_eq!(admission.queued_requests(), 1);
+
+        registry.transition_status(&second_id, WorkerStatus::Ready);
+        let selected = tokio::time::timeout(Duration::from_secs(1), waiting)
+            .await
+            .unwrap()
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(selected.selected.url(), second.url());
+        drop(selected);
+        drop(occupied);
+    }
 }