fix: warn before using a likely too-big chunk size for sglang-omni models

2 files changed, 37 insertions, 14 deletions

diff --git a/app/backends/sglomni/catalog.py b/app/backends/sglomni/catalog.py
index 1e4ca2d..160eb8b 100644
--- a/app/backends/sglomni/catalog.py
+++ b/app/backends/sglomni/catalog.py
@@ -69,11 +69,20 @@ class ModelEntry:
     # 86.13 fps, ~12 s) and Higgs's to 2048 frames (75 fps, ~27 s, and
     # per-request values are clamped to the engine cap, so its vendored
     # config raises the cap too) — and silently truncate longer text.
-    # 12288 frames covers a full 250-word CHUNK_SIZE sub-chunk on both and
-    # stays under the KV-pool admission check on every GPU that can host
-    # the model (the scheduler rejects prompt + max_new_tokens above it,
-    # and these requests are not auto-clamped).
+    # 12288 frames covers a full 250-word CHUNK_SIZE sub-chunk on ZONOS2
+    # (whose KV pools are >= 28083 tokens on every GPU that can host it)
+    # but NOT on Higgs: its thinker engine pins context_length at 4096,
+    # and the scheduler rejects any request whose prompt tokens plus
+    # max_new_tokens exceed that window (kv_capacity=4095, on every GPU —
+    # the pool side is never the binding constraint there).
     max_new_tokens: Optional[int] = None
+    # The largest sub-request (words) the model can narrate within its
+    # admission window, for models whose engine cannot cover a full
+    # CHUNK_SIZE sub-chunk (None = no cap; config.CHUNK_SIZE stands).
+    # Higgs: 3000 frames ≈ 40 s at 75 fps ≈ 80 words of narration, and
+    # the run pre-flight offers to clamp CHUNK_SIZE for the run (the
+    # client's adaptive retry still rescues requests the window rejects).
+    chunk_words: Optional[int] = None
 
 
 # The Qwen3-TTS CustomVoice speaker table — the same built-in speakers the
@@ -177,9 +186,18 @@ ENTRIES: Tuple[ModelEntry, ...] = (
         requires_reference=False,
         # The engine's 2048-frame default is ~27 s of speech at the codec's
         # 75 fps; requests are clamped to the engine cap server-side, so the
-        # yaml raises the cap and every request carries 12288 frames
-        # (~164 s) — enough for a full CHUNK_SIZE sub-chunk.
-        max_new_tokens=12288,
+        # yaml raises the cap and every request carries 3000 frames (~40 s)
+        # — the most the admission window allows: upstream pins the thinker
+        # engine's context_length at 4096, and the scheduler rejects any
+        # request whose prompt (including the reference-audio tokens) plus
+        # max_new_tokens exceeds it. 3000 frames leaves ~1095 tokens of
+        # prompt headroom (an 80-word chunk with a 20.5 s reference
+        # measured 684). Sub-requests cap at 80 words so the text fits the
+        # window too — the pre-flight offers to clamp CHUNK_SIZE for the
+        # run, and the client refits rejected requests to whatever the
+        # server reports as its capacity.
+        max_new_tokens=3000,
+        chunk_words=80,
         notes="zero-shot narration, cloning from a reference clip",
     ),
     ModelEntry(
diff --git a/app/backends/sglomni/configs/higgs_audio_v3_tts.yaml b/app/backends/sglomni/configs/higgs_audio_v3_tts.yaml
index 74736fa..8a42771 100644
--- a/app/backends/sglomni/configs/higgs_audio_v3_tts.yaml
+++ b/app/backends/sglomni/configs/higgs_audio_v3_tts.yaml
@@ -9,22 +9,27 @@
 # "CUDA out of memory. Tried to allocate 14.00 MiB".
 #
 # 0.80 trims the engine's static pool by ~1.2 GB per 24 GB of VRAM while
-# leaving a KV cache pool (~10 GB on a 24 GB card) far larger than any
-# narration request needs. Cards with heavy other-GPU-process usage can go
-# lower (e.g. 0.75).
+# leaving a KV cache pool far larger than any narration request needs.
+# Cards with heavy other-GPU-process usage can go lower (e.g. 0.75).
 #
 # The tts_engine factory also caps every request at max_new_tokens=2048
 # audio frames, and per-request values are clamped to that cap server-side
 # (make_higgs_scheduler_adapters) — the Higgs codec runs 75 frames per
 # second (24 kHz / 320 downsample), so the default is ~27 s of speech, which
 # silently truncates this tool's full 250-word sub-chunks (~100 s). Raising
-# the factory cap is the only way past it; the catalog also sends
-# max_new_tokens=12288 per request (the same value ZONOS2 uses) so a request
-# may use the room: 12288 frames ≈ 164 s.
+# the factory cap is the only way past it, but the ceiling is hard: upstream
+# pins the thinker engine's context_length at 4096 (HiggsTtsEngineBuilder —
+# not overridable), and the scheduler rejects any request whose prompt
+# tokens (including the reference-audio tokens) plus max_new_tokens exceed
+# that window ("Request requires more tokens than the thinker KV cache can
+# hold", kv_capacity=4095, on every GPU). The cap therefore lands at 3000
+# frames ≈ 40 s — the most the window allows with prompt headroom (an
+# 80-word chunk with a 20.5 s reference measured 684 prompt tokens) — and
+# the catalog caps sub-requests at 80 words to match (chunk_words).
 config_cls: HiggsTtsPipelineConfig
 model_path: bosonai/higgs-audio-v3-tts-4b
 stages:
   tts_engine:
     gpu_memory_fraction: 0.80
     factory:
-      max_new_tokens: 12288
+      max_new_tokens: 3000