[Router] Bound the e2e worker memory budget so prefill graph capture stops OOMing (#39713)

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

experimental/sgl-router/tests/e2e/infra/model_specs.py

@@ -33,6 +33,25 @@ def _resolve_model_path(hf_path: str) -> str:
     return hf_path
 
 
+# Every worker in this suite exists to answer a handful of 8–1024-token router
+# assertions, so the engine's serving defaults are actively wrong here. At
+# ``--mem-fraction-static 0.83`` the KV pool claims whatever the card has minus
+# a slack of ``free_memory_at_dist_init * (1 - mem_fraction_static)`` — on an
+# idle 80 GB H100 that is ~60 GB of KV (556K tokens for a 0.6B model) against
+# ~13 GB of slack, and the activation working set for the default
+# ``chunked_prefill_size=8192`` already accounts for nearly all of it. The
+# prefill CUDA graph then has nothing left to capture its 58 num-token buckets
+# into, dies part-way through with ``CUDA error: out of memory`` inside
+# ``graph.capture_end()``, and the worker exits during startup — surfacing as a
+# router e2e failure that has nothing to do with the router. Bounding the pool
+# and skipping the prefill graph removes the whole class; it also cuts the
+# per-spawn capture time, which this suite pays once per test.
+CI_WORKER_ARGS: list[str] = [
+    "--mem-fraction-static=0.6",
+    "--cuda-graph-backend-prefill=disabled",
+]
+
+
 MODEL_SPECS: dict[str, dict] = {
     # Fast-start tiny model for convergence / decode-affinity / stale-request
     # tests. Single GPU, ~2 GB weights, sub-30s start on a warm cache.
@@ -41,6 +60,7 @@ MODEL_SPECS: dict[str, dict] = {
         "memory_gb": 4,
         "tp": 1,
         "features": ["chat", "streaming"],
+        "worker_args": CI_WORKER_ARGS,
     },
     # Standard small chat model — matches SMG's `llama-1b` entry.
     "llama-1b": {
@@ -48,6 +68,7 @@ MODEL_SPECS: dict[str, dict] = {
         "memory_gb": 4,
         "tp": 1,
         "features": ["chat", "streaming"],
+        "worker_args": CI_WORKER_ARGS,
     },
     # Primary 8B chat model — matches SMG's `llama-8b`.
     "llama-8b": {
@@ -55,6 +76,7 @@ MODEL_SPECS: dict[str, dict] = {
         "memory_gb": 16,
         "tp": 1,
         "features": ["chat", "streaming"],
+        "worker_args": CI_WORKER_ARGS,
     },
 }