From ad271318d2c399b6a6747295b66cc937df0a6fc1 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Sun, 1 Nov 2026 10:15:30 -0700 Subject: [PATCH 1/3] Qwen3-Omni: fit the encoder captures on an 80 GB GPU #1072's vision block-loop buckets up to 32768 patches, four AuT capture batch sizes and the 512 MB default FlashInfer workspace per encoder slot leave Qwen3-Omni unable to start on an 80 GB GPU. Declare 64 MB workspaces for both encoder attentions, capture the AuT encoder at batch 8 only (smaller batches pad), and stop the vision buckets at 8192. Co-Authored-By: Claude Opus 5.5 --- mstar/model/qwen3_omni/qwen3_omni_model.py | 4 ++++ mstar/model/qwen3_omni/submodules.py | 8 +++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/mstar/model/qwen3_omni/qwen3_omni_model.py b/mstar/model/qwen3_omni/qwen3_omni_model.py index d1ea5739e7..ee112b6a84 100644 --- a/mstar/model/qwen3_omni/qwen3_omni_model.py +++ b/mstar/model/qwen3_omni/qwen3_omni_model.py @@ -342,6 +342,9 @@ 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), + # Split-KV partials of these shapes stay far below the 512 MB + # default, which costs 2 GB across both encoders and slots. + workspace_mb=64, ), ), RaggedAttentionSpec( @@ -354,6 +357,7 @@ 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, + workspace_mb=64, ), ), ] diff --git a/mstar/model/qwen3_omni/submodules.py b/mstar/model/qwen3_omni/submodules.py index 1edd67d00c..a64f0e3aa9 100644 --- a/mstar/model/qwen3_omni/submodules.py +++ b/mstar/model/qwen3_omni/submodules.py @@ -116,6 +116,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 @@ -152,7 +156,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. From 18e69bdce08eb625726e034317b170c8653aa062 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Sun, 1 Nov 2026 10:22:56 -0700 Subject: [PATCH 2/3] Qwen3-Omni: size the encoder workspaces for FA2's captured plans FA2 splits KV in every captured plan, at the bucket's segment ceiling; the AuT bs=8 buckets need ~230 MB and the 8192-patch vision bucket ~95 MB, so 64 MB would fail capture sizing off Hopper. Also update the bucket note in docs/models.rst. Co-Authored-By: Claude Opus 5.5 --- docs/models.rst | 3 ++- mstar/engine/resources/attn/ragged/config.py | 5 +++-- mstar/model/qwen3_omni/qwen3_omni_model.py | 10 ++++++---- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index 29989b95b3..b0b41b7f2b 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -500,7 +500,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. - Streamed audio starts after 12 codec frames (0.96 s of audio), then 13, then 25-frame chunks behind 25 frames of left context. Playback stays gapless while the Talker produces frames at about 1.5x real time over the first 50 diff --git a/mstar/engine/resources/attn/ragged/config.py b/mstar/engine/resources/attn/ragged/config.py index 3cf2f37055..175c79432f 100644 --- a/mstar/engine/resources/attn/ragged/config.py +++ b/mstar/engine/resources/attn/ragged/config.py @@ -94,8 +94,9 @@ class RaggedAttentionConfig: shard_over_sp: bool = True # Each label's FlashInfer workspace, per slot; None takes - # MSTAR_WORKSPACE_BUFFER_MB. A ragged prefill only spends it on split-KV - # partials (tens of MB), so a resource with many labels can declare less. + # MSTAR_WORKSPACE_BUFFER_MB. 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): diff --git a/mstar/model/qwen3_omni/qwen3_omni_model.py b/mstar/model/qwen3_omni/qwen3_omni_model.py index ee112b6a84..7124c6a700 100644 --- a/mstar/model/qwen3_omni/qwen3_omni_model.py +++ b/mstar/model/qwen3_omni/qwen3_omni_model.py @@ -342,9 +342,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), - # Split-KV partials of these shapes stay far below the 512 MB - # default, which costs 2 GB across both encoders and slots. - workspace_mb=64, + # 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( @@ -357,7 +358,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, - workspace_mb=64, + # FA2's split-KV partials at the 8192-patch bucket: ~95 MB. + workspace_mb=128, ), ), ] From de5272df75add1573c6bb709d792b0313e65080a Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Sun, 1 Nov 2026 12:52:19 -0700 Subject: [PATCH 3/3] Qwen3-Omni colocated: 176 Thinker KV pages so its graphs fit on 80 GB Co-Authored-By: Claude Opus 5.5 --- configs/qwen3omni_colocated.yaml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/configs/qwen3omni_colocated.yaml b/configs/qwen3omni_colocated.yaml index 97a080f3c3..e07af94747 100644 --- a/configs/qwen3omni_colocated.yaml +++ b/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: