feat(router): add load-aware generate paths

README.md

@@ -77,6 +77,21 @@ vllm-router \
     --intra-node-data-parallel-size 8
 ```
 
+#### Additional Generate Paths
+
+Register extension endpoints that use the `/inference/v1/generate` request
+schema. Configured paths get typed routing, active-load accounting, retries,
+and circuit breaking. Request JSON shape, backend authorization behavior,
+response status, headers, and body remain compatible with transparent proxying.
+Built-in routes keep their existing handlers. Workers must expose same paths.
+
+```bash
+vllm-router \
+  --worker-urls http://worker1:8000 http://worker2:8000 \
+  --policy cache_aware \
+  --extra-generate-paths /custom/v1/generate
+```
+
 #### Prefill-Decode Disaggregation
 ```bash
 # When vLLM runs the NIXL connector, prefill/decode URLs are required.

py_src/vllm_router/router.py

@@ -26,6 +26,8 @@ class Router:
     Args:
         worker_urls: List of URLs for worker nodes that will handle requests. Each URL should include
             the protocol, host, and port (e.g., ['http://worker1:8000', 'http://worker2:8000'])
+        extra_generate_paths: Additional HTTP paths that use the inference-generate request schema.
+            Requests on these paths use typed routing, retries, circuit breaking, and load tracking.
         policy: Load balancing policy to use. Options:
             - PolicyType.Random: Randomly select workers
             - PolicyType.RoundRobin: Distribute requests in round-robin fashion

py_src/vllm_router/router_args.py

@@ -89,6 +89,8 @@ class RouterArgs:
     cb_timeout_duration_secs: int = 60
     cb_window_duration_secs: int = 120
     disable_circuit_breaker: bool = False
+    # Additional typed inference-generate routes
+    extra_generate_paths: List[str] = dataclasses.field(default_factory=list)
 
     @staticmethod
     def add_cli_args(
@@ -128,6 +130,13 @@ def add_cli_args(
             default=[],
             help="List of worker URLs (e.g., http://worker1:8000 http://worker2:8000)",
         )
+        parser.add_argument(
+            f"--{prefix}extra-generate-paths",
+            type=str,
+            nargs="*",
+            default=[],
+            help="Additional HTTP paths using the inference-generate request schema",
+        )
 
         # Routing policy configuration
         parser.add_argument(
@@ -518,6 +527,16 @@ def from_cli_args(
         return cls(**args_dict)
 
     def _validate_router_args(self):
+        if len(set(self.extra_generate_paths)) != len(self.extra_generate_paths):
+            raise ValueError("extra_generate_paths must not contain duplicates")
+        for path in self.extra_generate_paths:
+            if (
+                len(path) < 2
+                or not path.startswith("/")
+                or any(character in path for character in "?#{}*")
+            ):
+                raise ValueError(f"invalid extra generate path: {path}")
+
         # Validate configuration based on mode
         if self.vllm_pd_disaggregation:
             # Validate PD configuration - skip URL requirements if using service discovery

py_test/unit/test_arg_parser.py

@@ -24,6 +24,7 @@ def test_default_values(self):
         assert args.port == 30000
         assert args.policy == "cache_aware"
         assert args.worker_urls == []
+        assert args.extra_generate_paths == []
         assert args.vllm_pd_disaggregation is False
         assert args.prefill_urls == []
         assert args.decode_urls == []
@@ -440,6 +441,22 @@ def test_parse_basic_args(self):
         assert router_args.worker_urls == ["http://worker1:8000", "http://worker2:8000"]
         assert router_args.policy == "round_robin"
 
+    def test_parse_extra_generate_paths(self):
+        router_args = parse_router_args(
+            [
+                "--worker-urls",
+                "http://worker1:8000",
+                "--extra-generate-paths",
+                "/custom/v1/generate",
+                "/extension/generate",
+            ]
+        )
+
+        assert router_args.extra_generate_paths == [
+            "/custom/v1/generate",
+            "/extension/generate",
+        ]
+
     def test_parse_pd_args(self):
         """Test parsing PD disaggregated mode arguments."""
         args = [

py_test/unit/test_validation.py

@@ -399,6 +399,29 @@ def test_request_id_headers_validation(self):
 class TestLaunchValidation:
     """Test launch-time validation logic."""
 
+    def test_extra_generate_paths_accept_distinct_extension_routes(self):
+        args = RouterArgs(
+            extra_generate_paths=["/custom/v1/generate", "/extension/generate"]
+        )
+
+        args._validate_router_args()
+
+    @pytest.mark.parametrize(
+        "paths",
+        [
+            ["custom/v1/generate"],
+            ["/custom/{path}"],
+            ["/custom/v1/generate", "/custom/v1/generate"],
+        ],
+    )
+    def test_extra_generate_paths_reject_invalid_or_duplicate_routes(self, paths):
+        args = RouterArgs(extra_generate_paths=paths)
+
+        with pytest.raises(
+            ValueError, match="extra_generate_paths|extra generate path"
+        ):
+            args._validate_router_args()
+
     def test_pd_mode_requires_urls(self):
         """Test that PD mode requires prefill and decode URLs."""
         # PD mode without URLs should fail

src/lib.rs

@@ -34,6 +34,7 @@ struct Router {
     host: String,
     port: u16,
     worker_urls: Vec<String>,
+    extra_generate_paths: Vec<String>,
     policy: PolicyType,
     worker_startup_timeout_secs: u64,
     worker_startup_check_interval: u64,
@@ -310,6 +311,7 @@ impl Router {
         otlp_traces_endpoint = None,
         // KV connector default (PD disaggregation)
         kv_connector = String::from("nixl"),
+        extra_generate_paths = vec![],
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -372,11 +374,13 @@ impl Router {
         enable_trace: bool,
         otlp_traces_endpoint: Option<String>,
         kv_connector: String,
+        extra_generate_paths: Vec<String>,
     ) -> PyResult<Self> {
         Ok(Router {
             host,
             port,
             worker_urls,
+            extra_generate_paths,
             policy,
             worker_startup_timeout_secs,
             worker_startup_check_interval,
@@ -502,6 +506,7 @@ impl Router {
                 } else {
                     None
                 },
+                extra_generate_paths: self.extra_generate_paths.clone(),
             })
             .await
             .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))

src/main.rs

@@ -106,6 +106,10 @@ struct CliArgs {
     #[arg(long, num_args = 0..)]
     worker_urls: Vec<String>,
 
+    /// Additional HTTP paths using the inference-generate request schema
+    #[arg(long, num_args = 0..)]
+    extra_generate_paths: Vec<String>,
+
     /// Load balancing policy to use
     #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash", "rendezvous_hash"])]
     policy: String,
@@ -614,6 +618,7 @@ impl CliArgs {
             } else {
                 None
             },
+            extra_generate_paths: self.extra_generate_paths.clone(),
         }
     }
 }

src/middleware.rs

@@ -3,6 +3,7 @@ use axum::{
     response::IntoResponse, response::Response,
 };
 use rand::Rng;
+use std::future::Future;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
@@ -496,17 +497,30 @@ pub async fn concurrency_limit_middleware(
     request: Request<axum::body::Body>,
     next: Next,
 ) -> Response {
+    let path = request.uri().path().to_string();
+    run_with_concurrency_limit(&app_state, &path, || next.run(request)).await
+}
+
+pub async fn run_with_concurrency_limit<F, Fut>(
+    app_state: &Arc<AppState>,
+    path: &str,
+    operation: F,
+) -> Response
+where
+    F: FnOnce() -> Fut,
+    Fut: Future<Output = Response>,
+{
     // Static counter for embeddings queue size
     static EMBEDDINGS_QUEUE_SIZE: AtomicU64 = AtomicU64::new(0);
 
     // Identify if this is an embeddings request based on path
-    let is_embeddings = request.uri().path().contains("/v1/embeddings");
+    let is_embeddings = path.contains("/v1/embeddings");
     let token_bucket = app_state.context.rate_limiter.clone();
 
     // Try to acquire token immediately
     if token_bucket.try_acquire(1.0).await.is_ok() {
         debug!("Acquired token immediately");
-        let response = next.run(request).await;
+        let response = operation().await;
 
         // Return the token to the bucket
         token_bucket.return_tokens(1.0).await;
@@ -545,7 +559,7 @@ pub async fn concurrency_limit_middleware(
                                 RouterMetrics::set_embeddings_queue_size(new_val as usize);
                             }
 
-                            let response = next.run(request).await;
+                            let response = operation().await;
 
                             // Return the token to the bucket
                             token_bucket.return_tokens(1.0).await;

src/protocols/spec.rs

@@ -1,4 +1,4 @@
-use serde::{Deserialize, Serialize};
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use serde_json::Value;
 use std::collections::HashMap;
 
@@ -2009,6 +2009,51 @@ impl GenerationRequest for InferenceGenerateRequest {
     }
 }
 
+/// Typed routing metadata paired with the original JSON object.
+///
+/// Additional generate paths use the typed view for routing while serializing
+/// the original value unchanged. This preserves the transparent proxy's JSON
+/// contract, including whether optional fields were omitted.
+#[derive(Debug, Clone)]
+pub struct PassthroughInferenceGenerateRequest {
+    routing: InferenceGenerateRequest,
+    original: Value,
+}
+
+impl GenerationRequest for PassthroughInferenceGenerateRequest {
+    fn is_stream(&self) -> bool {
+        self.routing.is_stream()
+    }
+
+    fn get_model(&self) -> Option<&str> {
+        self.routing.get_model()
+    }
+
+    fn extract_text_for_routing(&self) -> String {
+        self.routing.extract_text_for_routing()
+    }
+}
+
+impl Serialize for PassthroughInferenceGenerateRequest {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        self.original.serialize(serializer)
+    }
+}
+
+impl<'de> Deserialize<'de> for PassthroughInferenceGenerateRequest {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let original = Value::deserialize(deserializer)?;
+        let routing = serde_json::from_value(original.clone()).map_err(serde::de::Error::custom)?;
+        Ok(Self { routing, original })
+    }
+}
+
 // ==================================================================
 // =            VLLM SPEC - RERANK API                            =
 // ==================================================================
@@ -2426,7 +2471,7 @@ mod tests {
     }
 
     #[test]
-    fn test_inference_generate_lossless_passthrough() {
+    fn test_inference_generate_preserves_extension_fields() {
         let body = serde_json::json!({
             "token_ids": [1, 2, 3],
             "model": "qwen3-30b",
@@ -2448,6 +2493,27 @@ mod tests {
         assert_eq!(forwarded["model"], "qwen3-30b");
     }
 
+    #[test]
+    fn test_passthrough_inference_generate_preserves_complete_json_shape() {
+        let body = serde_json::json!({
+            "token_ids": [1, 2, 3],
+            "model": "model-a",
+            "sampling_params": {
+                "max_tokens": 4096,
+                "nested_extension": {"enabled": true}
+            },
+            "priority": 7,
+            "custom_request_field": "preserved"
+        });
+        let req: PassthroughInferenceGenerateRequest =
+            serde_json::from_value(body.clone()).unwrap();
+
+        assert_eq!(req.extract_text_for_routing(), "1 2 3");
+        assert_eq!(req.get_model(), Some("model-a"));
+        assert!(!req.is_stream());
+        assert_eq!(serde_json::to_value(&req).unwrap(), body);
+    }
+
     #[test]
     fn test_inference_generate_same_tokens_same_key_regardless_of_sampling() {
         let make = |seed: i64| -> InferenceGenerateRequest {

src/routers/http/openai_router.rs

@@ -5,7 +5,7 @@ use crate::core::{CircuitBreaker, CircuitBreakerConfig as CoreCircuitBreakerConf
 use crate::otel_http::{self, ClientRequestOptions};
 use crate::protocols::spec::{
     ChatCompletionRequest, CompletionRequest, GenerateRequest, InferenceGenerateRequest,
-    RerankRequest,
+    PassthroughInferenceGenerateRequest, RerankRequest,
 };
 use async_trait::async_trait;
 use axum::{
@@ -223,6 +223,20 @@ impl super::super::RouterTrait for OpenAIRouter {
             .into_response()
     }
 
+    async fn route_inference_generate_path(
+        &self,
+        _headers: Option<&HeaderMap>,
+        _body: &PassthroughInferenceGenerateRequest,
+        _path: &str,
+        _model_id: Option<&str>,
+    ) -> Response {
+        (
+            StatusCode::NOT_IMPLEMENTED,
+            "Generate endpoint not supported for OpenAI backend",
+        )
+            .into_response()
+    }
+
     async fn route_chat(
         &self,
         headers: Option<&HeaderMap>,