From e2d56bbbfc621a40994ee1216dfab1aa0dfac3d6 Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Wed, 16 Sep 2026 00:13:09 -0700 Subject: [PATCH] [Router] Bound the e2e worker memory budget so prefill graph capture stops OOMing (#39713) Co-authored-by: Kangyan Zhou Co-authored-by: Claude Opus 5 (1M context) --- .../sgl-router/tests/e2e/infra/model_specs.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/experimental/sgl-router/tests/e2e/infra/model_specs.py b/experimental/sgl-router/tests/e2e/infra/model_specs.py index 3363c8463..65a52a942 100644 --- a/experimental/sgl-router/tests/e2e/infra/model_specs.py +++ b/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, }, }