Qwen3-Omni: fit the encoder captures on an 80 GB GPU

configs/qwen3omni_colocated.yaml

@@ -1,10 +1,12 @@
 model: "qwen3_omni"
 max_seq_len: 32768
 # The KV pools default to 2,048 pages. The Thinker's is 24 GiB at 12 MiB a page,
-# and every node's weights on one 80 GB GPU leave about 12 GiB.
+# and every node's weights on one 80 GB GPU leave about 12 GiB. At 256 Thinker
+# pages the encoder and Thinker graphs no longer fit beside them; 176 leaves
+# room for the Talker's and a third of the Thinker's.
 resources:
   thinker_kv:
-    max_num_pages: 256
+    max_num_pages: 176
   talker_kv:
     max_num_pages: 256
 node_groups:

docs/models.rst

@@ -646,7 +646,8 @@ Qwen3-Omni notes
   needed. A deployment can pick the kernel under ``resources:`` in its YAML,
   e.g. ``vision_attn: {flashinfer_backend: cutlass}``. Each encoder's block stack
   is captured as a CUDA graph per token-count bucket; a request past the
-  largest bucket (32768 patches, about 5 minutes of audio) runs it eagerly.
+  largest bucket (8192 patches, about a 1920x1080 image; 4096 audio tokens)
+  runs it eagerly.
 - A request with no audio output decodes on its own walk, ``thinker_text_decode``:
   its token stays on the device between steps and reaches the client inline, as
   Qwen3.5's does, and nothing goes to the Talker. A request with audio output

mstar/engine/resources/attn/ragged/config.py

@@ -99,8 +99,9 @@ class RaggedAttentionConfig:
 
     # Where each label's FlashInfer workspace starts, per slot; None takes
     # MSTAR_WORKSPACE_BUFFER_MB, which a declared size wins over. A floor: a plan that
-    # asks for more grows it. A ragged prefill only spends it on split-KV partials, so a
-    # resource with many labels can declare less.
+    # asks for more grows it. FA3 plans don't use it; FA2 spends it on split-KV
+    # partials, which a captured plan sizes from its bucket's tokens and segment
+    # ceiling, so a resource with many labels can declare less.
     workspace_mb: int | None = None
 
     def __post_init__(self):

mstar/model/qwen3_omni/qwen3_omni_model.py

@@ -369,6 +369,10 @@ def _encoder_resources(self) -> list[NodeResourceSpec]:
                     num_kv_heads=audio.encoder_attention_heads,
                     head_dim=audio.head_dim,
                     max_segments_per_request=AuTEncoderSubmodule.max_windows_per_replay(audio),
+                    # The 512 MB default costs 2 GB across both encoders and slots.
+                    # FA3 plans take none of it; FA2's captured plans split KV at
+                    # the bucket's segment ceiling, ~230 MB here at bs=8.
+                    workspace_mb=256,
                 ),
             ),
             RaggedAttentionSpec(
@@ -381,6 +385,8 @@ def _encoder_resources(self) -> list[NodeResourceSpec]:
                     # bucket at this many; a request with more frames runs the
                     # block loop eagerly (`VisionEncoderSubmodule._fits_capture`).
                     max_segments_per_request=VisionEncoderSubmodule.BLOCK_LOOP_PATCH_BUCKETS[-1] // 256,
+                    # FA2's split-KV partials at the 8192-patch bucket: ~95 MB.
+                    workspace_mb=128,
                 ),
             ),
         ]

mstar/model/qwen3_omni/submodules.py

@@ -134,6 +134,10 @@ class AudioEncoderSubmodule(AuTEncoderSubmodule):
     """The AuT encoder (see ``AuTEncoderSubmodule``) over the unpadded
     ``(num_mel_bins, frames)`` log-mel features ``process_prompt`` made."""
 
+    # One captured batch size: a smaller batch pads to 8 with empty rows, and
+    # each extra size costs a graph per token bucket beside the Thinker.
+    ENCODE_CAPTURE_BATCH_SIZES = [8]
+
     def __init__(self, audio_encoder: AuTEncoder, config: Qwen3OmniModelConfig):
         super().__init__(audio_encoder, attn_key=AUDIO_ATTN)
         self.config = config
@@ -170,7 +174,9 @@ class VisionEncoderSubmodule(NodeSubmodule):
     """
 
     # Patch-count buckets, one request per replay. A larger request runs eagerly.
-    BLOCK_LOOP_PATCH_BUCKETS = [256, 512, 1024, 2048, 4096, 8192, 16384, 32768]
+    # 8192 covers 1920x1080; sizing a 16384 or 32768 graph doesn't fit beside
+    # the Thinker on an 80 GB GPU.
+    BLOCK_LOOP_PATCH_BUCKETS = [256, 512, 1024, 2048, 4096, 8192]
     # Not compiled: the block loop is captured, and around it the forward
     # takes the frame lengths, so a compiled wrapper would recompile per
     # frame count while the prefix cache holds what earlier graphs computed.