feat: gplang-omni reserves 70% of gpu 0 memory, preventing oom errors on consumer cards

10 files changed, 254 insertions, 3 deletions

diff --git a/app/backends/sglomni/catalog.py b/app/backends/sglomni/catalog.py
index 160eb8b..f20509f 100644
--- a/app/backends/sglomni/catalog.py
+++ b/app/backends/sglomni/catalog.py
@@ -104,10 +104,16 @@ _QWEN_EXTRAS: Tuple[Extra, ...] = (
 _SOX_HINT = ("install the sox system package (e.g. sudo pacman -S sox, "
              "sudo apt install sox, brew install sox)")
 # The Fish Audio and ZONOS2 pipelines use the Descript DAC codec, which
-# upstream installs WITH dependencies (nothing conflicts).
+# upstream installs WITH dependencies — but descript-audiotools carries a
+# vestigial 2021-era pin, protobuf<3.20 (its code never imports protobuf),
+# that downgrades the protobuf 6.x the sglang-omni stack itself needs
+# (smg-grpc-proto, grpcio health/reflection, cutlass-dsl, onnxruntime,
+# s3prl all demand >=4). The final extra re-pins protobuf AFTER
+# descript-audiotools so the downgrade never survives the install.
 _DAC_EXTRAS: Tuple[Extra, ...] = (
     ("descript-audiotools==0.7.2", False),
-    ("descript-audio-codec==1.0.0", False))
+    ("descript-audio-codec==1.0.0", False),
+    ("protobuf==6.33.6", False))
 
 # Companion distributions whose top-level import name differs from the pip
 # name's plain dash-to-underscore normalization (verified against their
@@ -115,6 +121,11 @@ _DAC_EXTRAS: Tuple[Extra, ...] = (
 _EXTRA_IMPORT_OVERRIDES = {
     "descript-audiotools": "audiotools",
     "descript-audio-codec": "dac",
+    # Not a naming quirk: a version marker. `google.protobuf` imports just
+    # as happily at the downgraded 3.19.6 as at the required 6.33.6, but
+    # this submodule only exists from protobuf 5.27 — so the venv probe
+    # fails (and the extra re-installs) whenever the audiotools pin won.
+    "protobuf": "google.protobuf.runtime_version",
 }
 
 
diff --git a/app/backends/sglomni/configs/moss_tts.yaml b/app/backends/sglomni/configs/moss_tts.yaml
index ef8d8fa..41bddd5 100644
--- a/app/backends/sglomni/configs/moss_tts.yaml
+++ b/app/backends/sglomni/configs/moss_tts.yaml
@@ -1,2 +1,24 @@
+# MOSS-TTS v1.5 with a pinned AR-engine memory budget.
+#
+# The upstream pipeline (MossTTSPipelineConfig) colocates its three stages
+# (preprocessing -> tts_engine -> vocoder) in one process on GPU 0 and leaves
+# the engine's sglang mem_fraction_static unset: the static pool (weights +
+# KV cache) is auto-sized to nearly all free VRAM at boot. On a 24 GB card
+# that leaves only tens of MiB free once the engine's CUDA graphs and the
+# colocated vocoder are resident — the first /v1/audio/speech request aborts
+# with "CUDA out of memory. Tried to allocate ~100 MiB" and every retry fails
+# identically (the same failure the qwen3_tts/higgs/zonos2 vendored configs
+# fix; note moss_tts_local does NOT need this — MossTTSLocalPipelineConfig
+# already budgets its colocated stages explicitly: 0.15/0.67/0.18).
+#
+# 0.70 pins the static pool at ~70% of the card (~16.5 GB on a 24 GB GPU —
+# a KV pool far larger than any narration request needs) and leaves ~7 GB
+# for the vocoder, CUDA graphs, transient allocations, and other GPU
+# processes. Precedent: dots_tts.yaml pins this same knob.
 config_cls: MossTTSPipelineConfig
 model_path: OpenMOSS-Team/MOSS-TTS-v1.5
+
+stages:
+  tts_engine:
+    engine:
+      mem_fraction_static: 0.70
diff --git a/app/backends/sglomni/configs/qwen3_tts_0_6b.yaml b/app/backends/sglomni/configs/qwen3_tts_0_6b.yaml
index a712ef9..bbb88bb 100644
--- a/app/backends/sglomni/configs/qwen3_tts_0_6b.yaml
+++ b/app/backends/sglomni/configs/qwen3_tts_0_6b.yaml
@@ -1,2 +1,23 @@
+# Qwen3-TTS 0.6B Base with a pinned AR-engine memory budget.
+#
+# The upstream pipeline (Qwen3TTSPipelineConfig) colocates its three stages
+# (preprocessing -> tts_engine -> vocoder) in one process on GPU 0 and leaves
+# the engine's sglang mem_fraction_static unset: the static pool (weights +
+# KV cache) is auto-sized to nearly all free VRAM at boot. On a 24 GB card
+# that leaves only tens of MiB free once the engine's CUDA graphs (~1.5 GB)
+# and the colocated vocoder are resident — the first /v1/audio/speech request
+# aborts with "CUDA out of memory. Tried to allocate ~100 MiB" and every
+# retry fails identically (the weights themselves are tiny: 0.6B bf16).
+#
+# 0.70 pins the static pool at ~70% of the card (~16.5 GB on a 24 GB GPU —
+# a KV pool far larger than any narration request needs) and leaves ~7 GB
+# for the vocoder, CUDA graphs, transient allocations, and other GPU
+# processes. Precedent: dots_tts.yaml pins this same knob; zonos2_bf16.yaml
+# and higgs_audio_v3_tts.yaml fix the same 24 GB OOM class their way.
 config_cls: Qwen3TTSPipelineConfig
 model_path: Qwen/Qwen3-TTS-12Hz-0.6B-Base
+
+stages:
+  tts_engine:
+    engine:
+      mem_fraction_static: 0.70
diff --git a/app/backends/sglomni/configs/qwen3_tts_0_6b_customvoice.yaml b/app/backends/sglomni/configs/qwen3_tts_0_6b_customvoice.yaml
index 6b284da..7481058 100644
--- a/app/backends/sglomni/configs/qwen3_tts_0_6b_customvoice.yaml
+++ b/app/backends/sglomni/configs/qwen3_tts_0_6b_customvoice.yaml
@@ -1,2 +1,23 @@
+# Qwen3-TTS 0.6B CustomVoice with a pinned AR-engine memory budget.
+#
+# The upstream pipeline (Qwen3TTSPipelineConfig) colocates its three stages
+# (preprocessing -> tts_engine -> vocoder) in one process on GPU 0 and leaves
+# the engine's sglang mem_fraction_static unset: the static pool (weights +
+# KV cache) is auto-sized to nearly all free VRAM at boot. On a 24 GB card
+# that leaves only tens of MiB free once the engine's CUDA graphs (~1.5 GB)
+# and the colocated vocoder are resident — the first /v1/audio/speech request
+# aborts with "CUDA out of memory. Tried to allocate ~100 MiB" and every
+# retry fails identically (the weights themselves are tiny: 0.6B bf16).
+#
+# 0.70 pins the static pool at ~70% of the card (~16.5 GB on a 24 GB GPU —
+# a KV pool far larger than any narration request needs) and leaves ~7 GB
+# for the vocoder, CUDA graphs, transient allocations, and other GPU
+# processes. Precedent: dots_tts.yaml pins this same knob; zonos2_bf16.yaml
+# and higgs_audio_v3_tts.yaml fix the same 24 GB OOM class their way.
 config_cls: Qwen3TTSPipelineConfig
 model_path: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+
+stages:
+  tts_engine:
+    engine:
+      mem_fraction_static: 0.70
diff --git a/app/backends/sglomni/configs/qwen3_tts_1_7b.yaml b/app/backends/sglomni/configs/qwen3_tts_1_7b.yaml
index 4f7706d..b9cb5ec 100644
--- a/app/backends/sglomni/configs/qwen3_tts_1_7b.yaml
+++ b/app/backends/sglomni/configs/qwen3_tts_1_7b.yaml
@@ -1,2 +1,23 @@
+# Qwen3-TTS 1.7B Base with a pinned AR-engine memory budget.
+#
+# The upstream pipeline (Qwen3TTSPipelineConfig) colocates its three stages
+# (preprocessing -> tts_engine -> vocoder) in one process on GPU 0 and leaves
+# the engine's sglang mem_fraction_static unset: the static pool (weights +
+# KV cache) is auto-sized to nearly all free VRAM at boot. On a 24 GB card
+# that leaves only tens of MiB free once the engine's CUDA graphs (~1.5 GB)
+# and the colocated vocoder are resident — the first /v1/audio/speech request
+# aborts with "CUDA out of memory. Tried to allocate ~100 MiB" and every
+# retry fails identically (the weights themselves are small: 1.7B bf16).
+#
+# 0.70 pins the static pool at ~70% of the card (~16.5 GB on a 24 GB GPU —
+# a KV pool far larger than any narration request needs) and leaves ~7 GB
+# for the vocoder, CUDA graphs, transient allocations, and other GPU
+# processes. Precedent: dots_tts.yaml pins this same knob; zonos2_bf16.yaml
+# and higgs_audio_v3_tts.yaml fix the same 24 GB OOM class their way.
 config_cls: Qwen3TTSPipelineConfig
 model_path: Qwen/Qwen3-TTS-12Hz-1.7B-Base
+
+stages:
+  tts_engine:
+    engine:
+      mem_fraction_static: 0.70
diff --git a/app/backends/sglomni/configs/qwen3_tts_1_7b_voicedesign.yaml b/app/backends/sglomni/configs/qwen3_tts_1_7b_voicedesign.yaml
index 20ae0b8..ade6a24 100644
--- a/app/backends/sglomni/configs/qwen3_tts_1_7b_voicedesign.yaml
+++ b/app/backends/sglomni/configs/qwen3_tts_1_7b_voicedesign.yaml
@@ -1,2 +1,23 @@
+# Qwen3-TTS 1.7B VoiceDesign with a pinned AR-engine memory budget.
+#
+# The upstream pipeline (Qwen3TTSPipelineConfig) colocates its three stages
+# (preprocessing -> tts_engine -> vocoder) in one process on GPU 0 and leaves
+# the engine's sglang mem_fraction_static unset: the static pool (weights +
+# KV cache) is auto-sized to nearly all free VRAM at boot. On a 24 GB card
+# that leaves only tens of MiB free once the engine's CUDA graphs (~1.5 GB)
+# and the colocated vocoder are resident — the first /v1/audio/speech request
+# aborts with "CUDA out of memory. Tried to allocate ~100 MiB" and every
+# retry fails identically (the weights themselves are small: 1.7B bf16).
+#
+# 0.70 pins the static pool at ~70% of the card (~16.5 GB on a 24 GB GPU —
+# a KV pool far larger than any narration request needs) and leaves ~7 GB
+# for the vocoder, CUDA graphs, transient allocations, and other GPU
+# processes. Precedent: dots_tts.yaml pins this same knob; zonos2_bf16.yaml
+# and higgs_audio_v3_tts.yaml fix the same 24 GB OOM class their way.
 config_cls: Qwen3TTSPipelineConfig
 model_path: Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign
+
+stages:
+  tts_engine:
+    engine:
+      mem_fraction_static: 0.70
diff --git a/app/backends/sglomni/configs/voxtral_tts.yaml b/app/backends/sglomni/configs/voxtral_tts.yaml
index 450cbce..b53f83a 100644
--- a/app/backends/sglomni/configs/voxtral_tts.yaml
+++ b/app/backends/sglomni/configs/voxtral_tts.yaml
@@ -1,2 +1,25 @@
+# Voxtral TTS 4B with a pinned AR-engine memory budget.
+#
+# The upstream pipeline (VoxtralTTSPipelineConfig) colocates its three stages
+# (preprocessing -> tts_generation -> vocoder) in one process on GPU 0 and
+# leaves the engine's sglang mem_fraction_static unset: the static pool
+# (weights + KV cache) is auto-sized to nearly all free VRAM at boot. On a
+# 24 GB card that leaves only tens of MiB free once the engine's CUDA graphs
+# and the colocated vocoder are resident — the first /v1/audio/speech request
+# aborts with "CUDA out of memory. Tried to allocate ~100 MiB" and every
+# retry fails identically (the same failure the qwen3_tts/higgs/zonos2
+# vendored configs fix).
+#
+# 0.70 pins the static pool at ~70% of the card (~16.5 GB on a 24 GB GPU —
+# ~8 GB of weights leaves a KV pool far larger than any narration request
+# needs) and leaves ~7 GB for the vocoder, CUDA graphs, transient
+# allocations, and other GPU processes. Precedent: dots_tts.yaml pins this
+# same knob. NOTE the stage name: Voxtral's engine stage is tts_generation,
+# not tts_engine.
 config_cls: VoxtralTTSPipelineConfig
 model_path: mistralai/Voxtral-4B-TTS-2603
+
+stages:
+  tts_generation:
+    engine:
+      mem_fraction_static: 0.70
diff --git a/app/backends/sglomni/wizard.py b/app/backends/sglomni/wizard.py
index 2083ec8..96679d6 100644
--- a/app/backends/sglomni/wizard.py
+++ b/app/backends/sglomni/wizard.py
@@ -101,7 +101,7 @@ def _wizard(stdscr, args: argparse.Namespace) -> Optional[dict]:
         installed = set(sg_models.installed_keys())
         picked = tui.checkbox_tree(
             stdscr, "Select SGLang-Omni Models to Install", families,
-            expand_all=True, back_value=_GO_BACK,
+            back_value=_GO_BACK,
             checked={(index, option["key"])
                      for index, family in enumerate(families)
                      for option in family["options"]
diff --git a/app/docs/backend-sglomni.md b/app/docs/backend-sglomni.md
index 58e5ce5..9091ebd 100644
--- a/app/docs/backend-sglomni.md
+++ b/app/docs/backend-sglomni.md
@@ -170,6 +170,33 @@ Find and stop the process holding the port
 (`ss -tlnp 'sport = :8100'`), or move this server to a free port
 (Settings → SGLang-Omni port), then start again.
 
+**The first request dies with `HTTP 500: CUDA out of memory. Tried to
+allocate ~100 MiB` (and every retry fails identically) on a 24 GB card,
+with the model far smaller than the card.** The model size is not the
+problem: upstream's Qwen3-TTS, MOSS-TTS and Voxtral pipelines colocate
+their AR engine and vocoder on GPU 0 with the engine's sglang
+`mem_fraction_static` unset, so the static pool (weights + KV cache) is
+auto-sized to nearly all free VRAM at boot. By the time a real request
+runs, the engine's CUDA graphs (~1.5 GB) and the colocated vocoder are
+resident too and only tens of MiB are left — the transient allocations
+generation needs do not fit. The vendored configs for those models pin
+the pool at `mem_fraction_static: 0.70` (~7 GB of headroom on a 24 GB
+card) exactly for this; Higgs (`gpu_memory_fraction: 0.80`) and ZONOS2's
+bf16 fallback (`0.70`) carry the same treatment. If you hit the failure
+on a server started by hand or an older checkout, pin the budget on the
+command line and let this tool use the running server as-is:
+
+```bash
+app/envs/sglomni/bin/sgl-omni serve \
+  --model-path Qwen/Qwen3-TTS-12Hz-0.6B-Base \
+  --config app/backends/sglomni/configs/qwen3_tts_0_6b.yaml \
+  --port 8100 --mem-fraction-static 0.70
+```
+
+(MOSS-TTS Local needs none of this: its upstream config already budgets
+its colocated stages explicitly. Fish Speech S2-Pro's OOM is a different
+problem — the model itself needs ~24 GB, upstream issue #359.)
+
 ## Manual setup
 
 ```bash
diff --git a/app/tests/test_backends_sglomni.py b/app/tests/test_backends_sglomni.py
index 8cb86cb..e61d51e 100644
--- a/app/tests/test_backends_sglomni.py
+++ b/app/tests/test_backends_sglomni.py
@@ -203,6 +203,9 @@ class CompanionPackageTests(unittest.TestCase):
             "audiotools")
         self.assertEqual(
             catalog.extra_import_name("descript-audio-codec==1.0.0"), "dac")
+        self.assertEqual(
+            catalog.extra_import_name("protobuf==6.33.6"),
+            "google.protobuf.runtime_version")
 
     def test_missing_companions_reports_absent_modules(self):
         entry = entry_by_key("qwen3_tts_1_7b_base")
@@ -227,6 +230,28 @@ class CompanionPackageTests(unittest.TestCase):
         with patch.object(envs, "env_exists", return_value=False):
             self.assertEqual(models.missing_companions(entry), [])
 
+    def test_dac_heal_reinstalls_protobuf_after_the_downgrade(self):
+        # descript-audiotools carries a vestigial protobuf<3.20 pin: its
+        # with-deps install downgrades the protobuf 6.x the sglang-omni
+        # stack itself needs. The restore extra must therefore re-enter
+        # missing_companions whenever that downgrade happened — probed via
+        # google.protobuf.runtime_version, which 3.19.6 does not provide.
+        entry = entry_by_key("zonos2")
+        self.assertIn(("protobuf==6.33.6", False), entry.extras)
+        with patch.object(envs, "env_exists", return_value=True), \
+                patch.object(envs, "module_available",
+                             side_effect=lambda name, env_dir:
+                             name != "google.protobuf.runtime_version"):
+            missing = models.missing_companions(entry)
+        self.assertEqual(missing, [("protobuf==6.33.6", False)])
+
+    def test_dac_protobuf_restore_installs_after_descript_audiotools(self):
+        # The restore only re-pins protobuf if it pip-installs AFTER the
+        # package whose dependency resolution did the downgrading.
+        specs = [spec for spec, _no_deps in entry_by_key("fish_s2_pro").extras]
+        self.assertLess(specs.index("descript-audiotools==0.7.2"),
+                        specs.index("protobuf==6.33.6"))
+
     def test_install_companions_installs_only_missing_with_no_deps(self):
         entry = entry_by_key("qwen3_tts_1_7b_base")
         calls = []
@@ -415,6 +440,58 @@ class HiggsConfigTests(unittest.TestCase):
         self.assertIsNone(entry_by_key("zonos2").chunk_words)
 
 
+class EngineMemoryBudgetTests(unittest.TestCase):
+    """Colocated AR pipelines with an unpinned engine budget OOM on 24 GB.
+
+    Upstream auto-sizes the engine's sglang static pool to nearly all free
+    VRAM when mem_fraction_static is unset (qwen3_tts, moss_tts, voxtral):
+    on a 24 GB card the first /v1/audio/speech request aborts with
+    "CUDA out of memory. Tried to allocate ~100 MiB" and every retry fails
+    identically. These vendored configs pin the pool at 0.70, leaving ~7 GB
+    for the colocated vocoder, CUDA graphs, and other GPU processes.
+    """
+
+    # (catalog key, engine stage name). Voxtral's engine stage is
+    # tts_generation; every other engine-bearing pipeline names it
+    # tts_engine. moss_tts_local is intentionally absent — its upstream
+    # config class already budgets its colocated stages explicitly
+    # (0.15 preprocessing / 0.67 AR / 0.18 vocoder), and pinning would
+    # fight that logic (the codec reserve derives from the fractions).
+    # s2-pro's OOM is model size, not budgeting (upstream issue #359).
+    PINNED = (
+        ("qwen3_tts_0_6b_base", "tts_engine"),
+        ("qwen3_tts_0_6b_customvoice", "tts_engine"),
+        ("qwen3_tts_1_7b_base", "tts_engine"),
+        ("qwen3_tts_1_7b_voicedesign", "tts_engine"),
+        ("moss_tts", "tts_engine"),
+        ("voxtral_tts", "tts_generation"),
+    )
+    UNPINNED = ("moss_tts_local", "fish_s2_pro")
+
+    def test_engine_budget_is_pinned_under_the_engine_stage(self):
+        for key, stage in self.PINNED:
+            with self.subTest(key=key, stage=stage):
+                path = config_path(entry_by_key(key))
+                self.assertIsNotNone(path)
+                self.assertTrue(path.is_file(), f"missing {path}")
+                text = path.read_text(encoding="utf-8")
+                self.assertIn(
+                    "stages:\n"
+                    f"  {stage}:\n"
+                    "    engine:\n"
+                    "      mem_fraction_static: 0.70\n",
+                    text)
+
+    def test_already_budgeted_and_oversized_models_stay_unpinned(self):
+        for key in self.UNPINNED:
+            with self.subTest(key=key):
+                path = config_path(entry_by_key(key))
+                self.assertIsNotNone(path)
+                self.assertTrue(path.is_file(), f"missing {path}")
+                text = path.read_text(encoding="utf-8")
+                self.assertNotIn("mem_fraction_static", text)
+
+
 class Fp8FallbackTests(unittest.TestCase):
     """FP8-only pipelines fall back to a vendored bf16 config on old GPUs."""
 
@@ -908,6 +985,13 @@ class SetupWizardTests(unittest.TestCase):
         self.assertEqual(settings["uninstall_keys"], [])
         self.assertEqual(confirms, [])
 
+    def test_setup_tree_starts_minimized(self):
+        """Like audio.cpp: no expand_all — every family starts collapsed."""
+        from backends.sglomni import wizard
+        _settings, trees, _confirms = self._wizard([wizard._GO_BACK])
+        self.assertEqual(trees[0][0], "Select SGLang-Omni Models to Install")
+        self.assertFalse(trees[0][1].get("expand_all", False))
+
     def test_steps_remove_before_downloading(self):
         from backends.sglomni import wizard
         gone, added = ENTRIES[0], ENTRIES[1]