From 5ac3be017d41abb78492eb09d3d36528970bf0a4 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:51 -0700 Subject: [PATCH 001/110] sdk: reference audio for client.tts voice cloning --- mstar/client/client.py | 32 +++++++++++++++++++++++++++++--- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/mstar/client/client.py b/mstar/client/client.py index 1fdc9f3a9..0abb05dc1 100644 --- a/mstar/client/client.py +++ b/mstar/client/client.py @@ -139,9 +139,35 @@ def generate_image(self, prompt: str, **model_kwargs) -> bytes: raise RuntimeError("Server returned no image output") return res.images[0] - def tts(self, text: str, *, voice: str | None = None, **model_kwargs) -> AudioBuffer: - """Text-to-speech. Returns an :class:`AudioBuffer` (``.to_wav(path)``).""" - res = self.generate(text=text, output_modalities=("audio",), voice=voice, **model_kwargs) + def tts( + self, + text: str, + *, + voice: str | None = None, + reference_audio=None, + **model_kwargs, + ) -> AudioBuffer: + """Text-to-speech. Returns an :class:`AudioBuffer` (``.to_wav(path)``). + + ``voice`` names a built-in speaker. ``reference_audio`` (a path, raw + ``bytes`` or a ``(filename, bytes)`` tuple) clones a voice from a clip + on models that support it; the clip travels as the request's audio + input, and model-specific knobs (``ref_text``, ``instruct``, + ``language``, ...) are forwarded verbatim as ``model_kwargs``. + """ + audio = None + input_modalities = None + if reference_audio is not None: + audio = [reference_audio] + input_modalities = ("audio", "text") + res = self.generate( + text=text, + audio=audio, + input_modalities=input_modalities, + output_modalities=("audio",), + voice=voice, + **model_kwargs, + ) if res.audio is None: raise RuntimeError("Server returned no audio output") return res.audio From 48c06d936a891399c61100e61467fbc1d9ecb263 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:51 -0700 Subject: [PATCH 002/110] test: client.tts forwards voice and reference audio --- test/modular/test_client_sdk.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/test/modular/test_client_sdk.py b/test/modular/test_client_sdk.py index 00cb64fc3..7dfa0f692 100644 --- a/test/modular/test_client_sdk.py +++ b/test/modular/test_client_sdk.py @@ -80,3 +80,24 @@ def test_audiobuffer_wav_bytes(): pcm = np.array([0, 16000, -16000], dtype=" Date: Fri, 18 Sep 2026 02:09:51 -0700 Subject: [PATCH 003/110] qwen3_tts: variant-aware config with speaker-encoder settings --- mstar/model/qwen3_tts/config.py | 137 +++++++++++++++++++++++++++++++- 1 file changed, 134 insertions(+), 3 deletions(-) diff --git a/mstar/model/qwen3_tts/config.py b/mstar/model/qwen3_tts/config.py index fe4c4518a..f45072a12 100644 --- a/mstar/model/qwen3_tts/config.py +++ b/mstar/model/qwen3_tts/config.py @@ -1,9 +1,16 @@ -"""Checkpoint-backed configuration for Qwen3-TTS 12 Hz CustomVoice. +"""Checkpoint-backed configuration for the Qwen3-TTS 12 Hz family. + +One dataclass tree serves every published 12 Hz checkpoint: the 0.6B and +1.7B CustomVoice models (built-in speakers, optional style instruction on +1.7B), VoiceDesign (voice described by an instruction) and Base (voice +cloned from reference audio through an ECAPA-TDNN speaker encoder). Which +paths a checkpoint supports is read from ``config.json``, never hard-coded. Qwen publishes configuration across three files rather than one monolithic object: -* ``config.json``: Talker architecture, special IDs, speakers, and languages +* ``config.json``: Talker architecture, special IDs, speakers, languages and + (Base only) the speaker encoder * ``generation_config.json``: main Talker and residual sampling defaults * ``speech_tokenizer/config.json``: neural audio decoder architecture/rates @@ -25,6 +32,15 @@ TALKER_SAMPLER = "talker_sampler" CODE_PRED_SAMPLER = "code_predictor" +# --------------------------------------------------------------------------- +# Fixed ChatML wrapper of the assistant turn, as the Qwen2 tokenizer emits it: +# ``<|im_start|>assistant\n`` ... ``<|im_end|>\n<|im_start|>assistant\n``. +# The reference slices the text span as ``input_id[:, 3:-5]``; the API side +# uses these to size the span and the Talker verifies them before prefill. +# --------------------------------------------------------------------------- +CHATML_ASSISTANT_PREFIX_TOKEN_IDS = (151644, 77091, 198) +CHATML_ASSISTANT_SUFFIX_TOKEN_IDS = (151645, 198, 151644, 77091, 198) + def _read_json(path: Path) -> dict[str, Any]: """Read optional checkpoint metadata, leaving dataclass defaults intact.""" @@ -152,6 +168,46 @@ def from_dict(cls, data: dict[str, Any]) -> "Qwen3TTSTalkerConfig": return cls(**values) +@dataclass +class Qwen3TTSSpeakerEncoderConfig: + """ECAPA-TDNN speaker encoder shipped with the Base checkpoint. + + Defaults mirror ``Qwen3TTSSpeakerEncoderConfig`` in the reference + implementation; ``config.json`` only pins ``enc_dim`` (the Talker hidden + size the x-vector is added into) and the sample rate of the reference + audio. + """ + + mel_dim: int = 128 + enc_dim: int = 1024 + enc_channels: tuple[int, ...] = (512, 512, 512, 512, 1536) + enc_kernel_sizes: tuple[int, ...] = (5, 3, 3, 3, 1) + enc_dilations: tuple[int, ...] = (1, 2, 3, 4, 1) + enc_attention_channels: int = 128 + enc_res2net_scale: int = 8 + enc_se_channels: int = 128 + sample_rate: int = 24000 + + # Mel front end used by the reference ``extract_speaker_embedding``. + n_fft: int = 1024 + hop_size: int = 256 + win_size: int = 1024 + fmin: int = 0 + fmax: int = 12000 + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "Qwen3TTSSpeakerEncoderConfig": + values = { + name: data[name] + for name in cls.__dataclass_fields__ + if name in data + } + for name in ("enc_channels", "enc_kernel_sizes", "enc_dilations"): + if name in values: + values[name] = tuple(values[name]) + return cls(**values) + + @dataclass class Qwen3TTSCodecConfig: """Official speech-tokenizer decoder plus M* streaming chunk controls.""" @@ -268,13 +324,28 @@ class Qwen3TTSModelConfig: tts_bos_token_id: int = 151672 tts_eos_token_id: int = 151673 - default_speaker: str = "vivian" default_language: str = "auto" talker: Qwen3TTSTalkerConfig = field(default_factory=Qwen3TTSTalkerConfig) codec: Qwen3TTSCodecConfig = field(default_factory=Qwen3TTSCodecConfig) generation: Qwen3TTSGenerationConfig = field( default_factory=Qwen3TTSGenerationConfig ) + # Present only on Base checkpoints (``speaker_encoder_config`` in + # config.json); CustomVoice and VoiceDesign carry no speaker encoder. + speaker_encoder: Qwen3TTSSpeakerEncoderConfig | None = None + + SUPPORTED_MODEL_TYPES = ("custom_voice", "voice_design", "base") + + def __post_init__(self) -> None: + if self.tts_model_type not in self.SUPPORTED_MODEL_TYPES: + raise ValueError( + f"Unsupported Qwen3-TTS tts_model_type {self.tts_model_type!r}; " + f"supported: {', '.join(self.SUPPORTED_MODEL_TYPES)}" + ) + if self.tts_model_type == "base" and self.speaker_encoder is None: + self.speaker_encoder = Qwen3TTSSpeakerEncoderConfig( + enc_dim=self.talker.hidden_size + ) @property def code_predictor(self) -> Qwen3TTSCodePredictorConfig: @@ -284,6 +355,62 @@ def code_predictor(self) -> Qwen3TTSCodePredictorConfig: def num_code_groups(self) -> int: return self.talker.num_code_groups + # -- Variant capabilities (all derived from the checkpoint metadata) ---- + + @property + def is_custom_voice(self) -> bool: + return self.tts_model_type == "custom_voice" + + @property + def is_voice_design(self) -> bool: + return self.tts_model_type == "voice_design" + + @property + def is_base(self) -> bool: + return self.tts_model_type == "base" + + @property + def has_builtin_speakers(self) -> bool: + """CustomVoice ships named speakers; VoiceDesign and Base do not.""" + return bool(self.talker.spk_id) + + @property + def default_speaker(self) -> str | None: + """Speaker used when a request names none (``None`` = no speaker tag).""" + if not self.has_builtin_speakers: + return None + return "vivian" if "vivian" in self.talker.spk_id else sorted(self.talker.spk_id)[0] + + @property + def supports_instruct(self) -> bool: + """Instruction text (style or voice description) in the prefill. + + VoiceDesign is driven by it; the 1.7B CustomVoice accepts it for + style/emotion control. The reference silently drops instructions for + the 0.6B CustomVoice, which was not trained with them, so M* rejects + them there instead of ignoring the request field. + """ + if self.is_voice_design: + return True + return self.is_custom_voice and self.tts_model_size != "0b6" + + @property + def requires_instruct(self) -> bool: + return self.is_voice_design + + @property + def supports_reference_audio(self) -> bool: + return self.is_base + + @property + def default_non_streaming_mode(self) -> bool: + """Reference/vLLM-Omni/SGLang-Omni default text layout per variant. + + CustomVoice and VoiceDesign place the whole text in the prefill; + Base (voice clone) feeds text one token per generated frame. + """ + return not self.is_base + @classmethod def from_pretrained(cls, model_dir: str | Path) -> "Qwen3TTSModelConfig": """Compose Talker, Codec, and generation configs from a local snapshot.""" @@ -315,4 +442,8 @@ def from_pretrained(cls, model_dir: str | Path) -> "Qwen3TTSModelConfig": codec=Qwen3TTSCodecConfig.from_dict(codec_data), generation=Qwen3TTSGenerationConfig.from_dict(generation_data), ) + if "speaker_encoder_config" in model_data: + values["speaker_encoder"] = Qwen3TTSSpeakerEncoderConfig.from_dict( + model_data["speaker_encoder_config"] + ) return cls(**values) From 28a34d9ad7236c0a88953e49f272bfb0c0a61e77 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:51 -0700 Subject: [PATCH 004/110] qwen3_tts: project Talker-width inputs into the code predictor --- mstar/model/qwen3_tts/components/talker.py | 24 ++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/mstar/model/qwen3_tts/components/talker.py b/mstar/model/qwen3_tts/components/talker.py index adbe24eff..cc2dc4f11 100644 --- a/mstar/model/qwen3_tts/components/talker.py +++ b/mstar/model/qwen3_tts/components/talker.py @@ -6,6 +6,9 @@ M*'s paged KV cache. Its attention/MLP projections are tensor-parallel. * Within one time step, the 5-layer CodePredictor walks codec groups 1-15. This short depth axis uses a fixed local KV tensor and is kept replicated. + Its inputs (the Talker hidden state and every codec embedding) live in the + Talker width; ``small_to_mtp_projection`` maps them into the predictor + width when the two differ (1.7B: 2048 -> 1024; 0.6B: identity). Class/module names intentionally follow Hugging Face checkpoint namespaces so ``load_hf_weights`` can stream parameters without a model-specific state-dict @@ -227,18 +230,24 @@ class Qwen3TTSCodePredictor(nn.Module): Callers first write the Talker hidden state at position 0, then repeatedly feed the preceding codec embedding at positions 1-15. Keeping this as tensor-only code allows the complete depth loop to be CUDA-graph captured. + + Every depth input arrives in the Talker width and passes through + ``small_to_mtp_projection`` (a biased linear layer on the 1.7B + checkpoints, identity on the 0.6B where both widths are 1024), exactly as + the reference ``Qwen3TTSTalkerCodePredictorModelForConditionalGeneration`` + does before its decoder. """ def __init__(self, config: Qwen3TTSModelConfig) -> None: super().__init__() cp = config.code_predictor - if cp.hidden_size != config.talker.hidden_size: - raise ValueError( - "M* currently requires equal Talker and CodePredictor hidden " - "sizes; the supported 0.6B checkpoint uses 1024 for both" - ) self.config = cp self.model = Qwen3TTSCodePredictorInnerModel(config) + self.small_to_mtp_projection: nn.Module = ( + nn.Linear(config.talker.hidden_size, cp.hidden_size, bias=True) + if cp.hidden_size != config.talker.hidden_size + else nn.Identity() + ) self.lm_head = nn.ModuleList([ nn.Linear(cp.hidden_size, cp.vocab_size, bias=False) for _ in range(config.num_code_groups - 1) @@ -274,8 +283,11 @@ def forward_depth_unrolled( ``kv_cache`` layout is ``[layers, batch, K/V, groups, kv_heads, head_dim]``. Unlike Talker cache, it is frame-local scratch space: every generated frame starts at ``cache_pos=0`` and overwrites it. + + ``inputs_embeds`` is ``[batch, 1, talker_hidden]``; the returned + hidden state is ``[batch, 1, predictor_hidden]``. """ - hidden_states = inputs_embeds + hidden_states = self.small_to_mtp_projection(inputs_embeds) batch_size, seq_len, _ = hidden_states.shape if seq_len != 1: raise ValueError("CodePredictor decode expects exactly one token") From 3259e936221ec0d5122d7dcc6a348746868e8715 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:51 -0700 Subject: [PATCH 005/110] qwen3_tts: build prefill with instruct, optional speaker and both text layouts --- mstar/model/qwen3_tts/submodules.py | 134 +++++++++++++++++----------- 1 file changed, 83 insertions(+), 51 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index ad3b99bbc..f7e293185 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -4,7 +4,10 @@ # # Two submodules cover the complete text-to-speech streaming pipeline: # 1. TalkerSubmodule (KV_CACHE engine) -# - Builds the official text/voice prefill sequence. +# - Builds the official prefill sequence for every 12 Hz variant: +# optional instruction turn, ChatML assistant role, codec language / +# speaker tags, then the text either whole (non-streaming layout) or +# one token per generated frame (streaming layout). # - Maintains the Talker paged KV cache across 12 Hz decode steps. # - Predicts codec group 0 with the Talker and groups 1-15 with the # depth-wise CodePredictor. @@ -50,6 +53,8 @@ Qwen3TTSTalkerModel, ) from mstar.model.qwen3_tts.config import ( + CHATML_ASSISTANT_PREFIX_TOKEN_IDS, + CHATML_ASSISTANT_SUFFIX_TOKEN_IDS, CODE_PRED_SAMPLER, TALKER_ATTN, TALKER_KV, @@ -89,8 +94,8 @@ class TalkerSubmodule(ARNodeSubmodule): disable_torch_compile = True MAX_BATCH_SIZE = 32 DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] - CHATML_ASSISTANT_PREFIX_TOKEN_IDS = (151644, 77091, 198) - CHATML_ASSISTANT_SUFFIX_TOKEN_IDS = (151645, 198, 151644, 77091, 198) + CHATML_ASSISTANT_PREFIX_TOKEN_IDS = CHATML_ASSISTANT_PREFIX_TOKEN_IDS + CHATML_ASSISTANT_SUFFIX_TOKEN_IDS = CHATML_ASSISTANT_SUFFIX_TOKEN_IDS def __init__( self, @@ -200,27 +205,10 @@ def _special_text_embeds( bos, eos, pad = self._project_text(token_ids).to(dtype).chunk(3, dim=1) return bos, eos, pad - def _build_prefill( - self, - request_id: str, - text_ids: torch.Tensor, - speaker_id: int, - language_id: int, - ) -> torch.Tensor: - """Build the official mixed text/codec prefill embedding sequence. - - The assistant-role prefix and codec conditioning tags enter the - one-shot prefill. Remaining prompt text is retained in per-request - state and added one token at a time to later recurrent codec embeds. - This aligns text progress with the 12 Hz acoustic generation steps. - """ - text_ids = text_ids.to(device=self.get_device(), dtype=torch.long).view(1, -1) - expected_prefix = text_ids.new_tensor( - self.CHATML_ASSISTANT_PREFIX_TOKEN_IDS - ) - expected_suffix = text_ids.new_tensor( - self.CHATML_ASSISTANT_SUFFIX_TOKEN_IDS - ) + def _validate_chatml(self, text_ids: torch.Tensor) -> tuple[int, int]: + """Check the fixed assistant-turn wrapper and return its (prefix, suffix) lengths.""" + expected_prefix = text_ids.new_tensor(self.CHATML_ASSISTANT_PREFIX_TOKEN_IDS) + expected_suffix = text_ids.new_tensor(self.CHATML_ASSISTANT_SUFFIX_TOKEN_IDS) prefix_len = expected_prefix.numel() suffix_len = expected_suffix.numel() if text_ids.shape[1] < prefix_len + 1 + suffix_len: @@ -240,6 +228,48 @@ def _build_prefill( f"{expected_suffix.tolist()}, got " f"{text_ids[0, -suffix_len:].tolist()}" ) + return prefix_len, suffix_len + + def _build_prefill( + self, + request_id: str, + text_ids: torch.Tensor, + prompt_layout: torch.Tensor, + speaker_id: int, + language_id: int, + ) -> torch.Tensor: + """Build the official mixed text/codec prefill embedding sequence. + + ``text_ids`` concatenates an optional instruction turn with the + assistant turn; ``prompt_layout`` is ``[instruct_len, text_len, + stream_text]`` and says where the split is and how the text is fed. + The layout mirrors ``Qwen3TTSForConditionalGeneration.generate``: + + * instruction (``<|im_start|>user ... <|im_end|>``) as plain projected + text embeddings (VoiceDesign, 1.7B CustomVoice style control); + * the assistant role, then the codec think/language tags, the speaker + tag when the checkpoint has built-in speakers (``speaker_id >= 0``) + and the codec PAD, all summed with TTS PAD / BOS text embeddings; + * ``stream_text == 0`` (the reference default for CustomVoice and + VoiceDesign): every text token plus TTS EOS enters the prefill over + codec PADs, closed by TTS PAD + codec BOS. Decode then adds TTS PAD + to each frame. + * ``stream_text == 1`` (the reference default for Base): only the first + text token enters the prefill over codec BOS; the remaining tokens + plus TTS EOS are kept in per-request state and added one per frame. + """ + text_ids = text_ids.to(device=self.get_device(), dtype=torch.long).view(1, -1) + instruct_len, text_len, stream_text = (int(v) for v in prompt_layout.tolist()) + instruct_ids = text_ids[:, :instruct_len] + assistant_ids = text_ids[:, instruct_len:] + prefix_len, suffix_len = self._validate_chatml(assistant_ids) + if assistant_ids.shape[1] != prefix_len + text_len + suffix_len: + raise ValueError( + "Qwen3-TTS prompt layout disagrees with the token stream: " + f"expected {prefix_len + text_len + suffix_len} assistant tokens, " + f"got {assistant_ids.shape[1]}" + ) + text_tokens = assistant_ids[:, prefix_len:prefix_len + text_len] codec = self.talker_config codec_prefix = ( @@ -252,43 +282,44 @@ def _build_prefill( codec.codec_think_eos_id, ] ) + speaker_tag = [speaker_id] if speaker_id >= 0 else [] codec_ids = torch.tensor( - [[*codec_prefix, speaker_id, codec.codec_pad_id, codec.codec_bos_id]], + [[*codec_prefix, *speaker_tag, codec.codec_pad_id, codec.codec_bos_id]], dtype=torch.long, device=self.get_device(), ) codec_embeds = self.model.model.codec_embedding(codec_ids) - bos_embed, eos_embed, pad_embed = self._special_text_embeds( - codec_embeds.dtype - ) + dtype = codec_embeds.dtype + bos_embed, eos_embed, pad_embed = self._special_text_embeds(dtype) + + def project(ids: torch.Tensor) -> torch.Tensor: + return self._project_text(ids).to(dtype) - # Prefix layout mirrors the official CustomVoice generation helper: - # assistant role, language/voice codec tags, then first text token. - role_embed = self._project_text( - text_ids[:, :prefix_len] - ).to(codec_embeds.dtype) + # Tags: TTS PAD over every codec tag but the last, TTS BOS over the + # codec PAD; the codec BOS pairs with text below. + role_embed = project(assistant_ids[:, :prefix_len]) tag_text = torch.cat([ pad_embed.expand(-1, codec_embeds.shape[1] - 2, -1), bos_embed, ], dim=1) - tag_embed = tag_text + codec_embeds[:, :-1] - first_text = ( - self._project_text( - text_ids[:, prefix_len:prefix_len + 1] - ).to(codec_embeds.dtype) - + codec_embeds[:, -1:] - ) - prefill = torch.cat([role_embed, tag_embed, first_text], dim=1) - - # The fixed five-token ChatML suffix is replaced by projected TTS EOS. - # Decode consumes this tensor by ``generation_step`` and uses PAD once - # the text condition has been exhausted. - trailing = torch.cat([ - self._project_text( - text_ids[:, prefix_len + 1:-suffix_len] - ).to(codec_embeds.dtype), - eos_embed, - ], dim=1) + pieces = [role_embed, tag_text + codec_embeds[:, :-1]] + if instruct_len: + pieces.insert(0, project(instruct_ids)) + + if stream_text: + pieces.append(project(text_tokens[:, :1]) + codec_embeds[:, -1:]) + trailing = torch.cat([project(text_tokens[:, 1:]), eos_embed], dim=1) + else: + text_embed = torch.cat([project(text_tokens), eos_embed], dim=1) + codec_pads = self.model.model.codec_embedding( + codec_ids.new_full((1, text_embed.shape[1]), codec.codec_pad_id) + ) + pieces.append(text_embed + codec_pads) + pieces.append(pad_embed + codec_embeds[:, -1:]) + # Nothing left to feed: every frame adds TTS PAD (empty stream). + trailing = eos_embed[:, :0] + prefill = torch.cat(pieces, dim=1) + self.request_state(request_id).add_all( trailing_text_hidden=trailing.squeeze(0), tts_pad_embed=pad_embed[0, 0], @@ -316,6 +347,7 @@ def prepare_inputs( input_embeds = self._build_prefill( fwd_info.request_id, inputs["text_inputs"][0], + inputs["prompt_layout"][0], int(inputs["speaker_id"][0].item()), int(inputs["language_id"][0].item()), ) From cbb27df4231bc994e51897cdc92fb20d7ea74fe1 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 006/110] qwen3_tts: serve CustomVoice, VoiceDesign and Base from one class --- mstar/model/qwen3_tts/qwen3_tts_model.py | 294 ++++++++++++++++------- 1 file changed, 213 insertions(+), 81 deletions(-) diff --git a/mstar/model/qwen3_tts/qwen3_tts_model.py b/mstar/model/qwen3_tts/qwen3_tts_model.py index bd26a574c..d90472446 100644 --- a/mstar/model/qwen3_tts/qwen3_tts_model.py +++ b/mstar/model/qwen3_tts/qwen3_tts_model.py @@ -1,10 +1,14 @@ """Qwen3-TTS model contract and two-partition streaming topology. -The 0.6B CustomVoice checkpoint is a text-to-speech model without the large -multimodal Thinker used by Qwen3-Omni. Its autoregressive Talker predicts one -12 Hz codec frame per step: group 0 comes from the Talker language model and -groups 1-15 come from a small depth-wise CodePredictor. The speech-tokenizer -decoder turns those frames into 24 kHz PCM. +One class serves every 12 Hz checkpoint: 0.6B/1.7B CustomVoice (built-in +speakers, style instructions on 1.7B), 1.7B VoiceDesign (voice described by +an instruction) and 1.7B Base (voice cloned from reference audio). They share +one architecture: an autoregressive Talker predicts one 12 Hz codec frame per +step, group 0 from the Talker language model and groups 1-15 from a small +depth-wise CodePredictor, and the speech-tokenizer decoder turns those frames +into 24 kHz PCM. What differs is only the prefill conditioning, which is +derived from ``config.json`` (``Qwen3TTSModelConfig``), never from the +registry key. Architecture (two asynchronous partitions): Talker - text/voice prefill, then autoregressive 16-group codec frames @@ -60,6 +64,8 @@ from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION from mstar.model.base import ForwardPassArgs, Model from mstar.model.qwen3_tts.config import ( + CHATML_ASSISTANT_PREFIX_TOKEN_IDS, + CHATML_ASSISTANT_SUFFIX_TOKEN_IDS, CODE_PRED_SAMPLER, TALKER_ATTN, TALKER_KV, @@ -173,13 +179,87 @@ def load_private_module(name: str, path: Path): ) +# --------------------------------------------------------------------------- +# Checkpoint completeness +# --------------------------------------------------------------------------- + +# Fused M* parameters and the per-shard checkpoint keys that feed them +# (mirrors ``LLAMA_STACKED_PARAMS`` in the loader). +_FUSED_SOURCES = { + "qkv_proj": ("q_proj", "k_proj", "v_proj"), + "gate_up_proj": ("gate_proj", "up_proj"), +} + + +def _checkpoint_keys(checkpoint_dir: str | Path, prefix: str) -> set[str]: + """Tensor names under ``prefix`` in a (possibly sharded) safetensors checkpoint.""" + import json + + from safetensors import safe_open + + root = Path(checkpoint_dir) + index = root / "model.safetensors.index.json" + if index.is_file(): + with index.open(encoding="utf-8") as f: + names = json.load(f)["weight_map"].keys() + else: + with safe_open(str(root / "model.safetensors"), framework="pt") as f: + names = list(f.keys()) + return {name.removeprefix(prefix) for name in names if name.startswith(prefix)} + + +def _expected_checkpoint_keys(module: torch.nn.Module) -> set[str]: + """Checkpoint keys an M* module consumes, expanding fused projections.""" + expected: set[str] = set() + for name in dict(module.named_parameters()): + for fused, sources in _FUSED_SOURCES.items(): + if f".{fused}." in name: + expected.update(name.replace(f".{fused}.", f".{source}.") for source in sources) + break + else: + expected.add(name) + return expected + + +def _verify_checkpoint_coverage( + module: torch.nn.Module, + loaded: set[str], + checkpoint_keys: set[str], + component: str, +) -> None: + """Fail startup on any parameter left uninitialized or any key left unused. + + Both directions matter: a missing key means random weights would serve + requests; an unused key means the checkpoint carries a component this + port silently ignores (the 1.7B code predictor projection, for example). + """ + expected = set(dict(module.named_parameters())) + missing = sorted(expected - loaded) + if missing: + preview = ", ".join(missing[:8]) + raise RuntimeError( + f"{component} checkpoint did not initialize {len(missing)} " + f"parameters: {preview}" + ) + unused = sorted( + key for key in checkpoint_keys - _expected_checkpoint_keys(module) + if "rotary_emb" not in key + ) + if unused: + preview = ", ".join(unused[:8]) + raise RuntimeError( + f"{component} checkpoint has {len(unused)} tensors this port does " + f"not load: {preview}" + ) + + # --------------------------------------------------------------------------- # Model contract # --------------------------------------------------------------------------- class Qwen3TTSModel(Model): - """Qwen3-TTS 12 Hz CustomVoice model contract. + """Qwen3-TTS 12 Hz model contract (CustomVoice, VoiceDesign, Base). GPU computation is split into an autoregressive Talker partition and a streaming Codec partition. This class owns only model-level scheduling, @@ -197,12 +277,9 @@ def __init__( # The lightweight API-side object needs config and tokenizer only. self.local_dir = _resolve_model_metadata(model_path_hf, cache_dir) + # Rejects unknown ``tts_model_type`` values; every supported variant + # is handled below through the config's capability properties. self.config = Qwen3TTSModelConfig.from_pretrained(self.local_dir) - if self.config.tts_model_type != "custom_voice": - raise ValueError( - "The first Qwen3-TTS integration supports only CustomVoice " - f"checkpoints, got {self.config.tts_model_type!r}" - ) self.tokenizer = AutoTokenizer.from_pretrained( self.local_dir, @@ -286,7 +363,7 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: # Talker-to-Codec stream. talker_prefill = GraphNode( name="Talker", - input_names=["text_inputs", "speaker_id", "language_id"], + input_names=list(self.PREFILL_INPUTS), outputs=[ GraphEdge( next_node=EMPTY_DESTINATION, @@ -392,45 +469,56 @@ def get_partition_topology(self) -> PartitionTopology: # API preprocessing # ----------------------------------------------------------------------- - def process_prompt( - self, - prompt: str | None, - input_modalities: list[str], - output_modalities: list[str], - tensors: NameToTensorList | None = None, - **kwargs: Any, - ) -> NameToTensorList: - """Validate a CustomVoice request and build Talker input tensors. - - Qwen3-TTS expects an assistant ChatML turn rather than a generic user - turn. Speaker and language are separate codec-side conditioning IDs; - ``-1`` means automatic language selection. + # Prompt templates of the reference ``qwen_tts`` inference wrapper. The + # assistant wrapper is fixed at 3 + 5 tokens (see ``_validate_chatml``), + # which is how the text span is located inside the tokenized turn. + ASSISTANT_TEMPLATE = "<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" + INSTRUCT_TEMPLATE = "<|im_start|>user\n{instruct}<|im_end|>\n" + # Tensors ``process_prompt`` produces and the Talker prefill consumes. + PREFILL_INPUTS = ("text_inputs", "prompt_layout", "speaker_id", "language_id") + + def _tokenize(self, text: str) -> torch.Tensor: + encoded = self.tokenizer(text, return_tensors="pt", padding=True) + ids = encoded["input_ids"] + if ids.ndim == 2: + ids = ids[0] + return ids.to(dtype=torch.long) + + def _resolve_speaker(self, kwargs: dict[str, Any]) -> tuple[str | None, int]: + """Map ``voice``/``speaker`` onto a codec speaker tag (``-1`` = none). + + CustomVoice checkpoints carry named speakers and fall back to the + default one. VoiceDesign and Base have none: the voice comes from the + instruction or the reference audio, so naming one is an error rather + than something to ignore silently. """ - del tensors - if not prompt: - raise ValueError("Qwen3-TTS requires a non-empty text prompt") - if set(input_modalities) != {"text"}: - raise ValueError( - "Qwen3-TTS CustomVoice currently supports text input only" - ) - if set(output_modalities) != {"audio"}: - raise ValueError("Qwen3-TTS CustomVoice supports audio output only") - if kwargs.get("instruct"): + requested = kwargs.get("speaker", kwargs.get("voice")) + if requested is None or requested == "": + requested = self.config.default_speaker + if requested is None: + return None, -1 + elif not self.config.has_builtin_speakers: raise ValueError( - "Qwen3-TTS 0.6B CustomVoice does not support instructions" + f"Qwen3-TTS {self.config.tts_model_type} checkpoints have no " + "built-in speakers; describe the voice with 'instruct' " + "(VoiceDesign) or supply reference audio (Base) instead of 'voice'" ) - - speaker = str( - kwargs.get("speaker", kwargs.get("voice", self.config.default_speaker)) - ).lower() + speaker = str(requested).lower() if speaker not in self.config.talker.spk_id: supported = ", ".join(sorted(self.config.talker.spk_id)) raise ValueError( f"Unsupported Qwen3-TTS speaker {speaker!r}; supported: {supported}" ) + return speaker, self.config.talker.spk_id[speaker] + def _resolve_language(self, kwargs: dict[str, Any], speaker: str | None) -> int: + """Map ``language`` onto a codec language tag (``-1`` = automatic).""" language = str(kwargs.get("language", self.config.default_language)).lower() - dialect = self.config.talker.spk_is_dialect.get(speaker, False) + dialect = ( + self.config.talker.spk_is_dialect.get(speaker, False) + if speaker is not None + else False + ) if dialect and language in {"auto", "chinese"}: language = str(dialect).lower() @@ -444,29 +532,80 @@ def process_prompt( raise ValueError( f"Unsupported Qwen3-TTS language {language!r}; supported: {supported}" ) + return self.config.talker.codec_language_id.get(language, -1) - # Match the official processor template exactly. `_build_prefill` - # relies on the fixed assistant suffix when separating prompt tokens - # into the initial prefill and per-frame text conditioning stream. - formatted = ( - f"<|im_start|>assistant\n{prompt}<|im_end|>\n" - "<|im_start|>assistant\n" - ) - encoded = self.tokenizer( - formatted, - return_tensors="pt", - padding=True, + def _resolve_instruct(self, kwargs: dict[str, Any]) -> str: + """Style/voice instruction; ``instructions`` is the OpenAI field name.""" + instruct = kwargs.get("instruct", kwargs.get("instructions")) or "" + instruct = str(instruct).strip() + if instruct and not self.config.supports_instruct: + raise ValueError( + f"Qwen3-TTS {self.config.tts_model_size} " + f"{self.config.tts_model_type} does not support instructions" + ) + if not instruct and self.config.requires_instruct: + raise ValueError( + "Qwen3-TTS VoiceDesign requires an 'instruct' describing the voice" + ) + return instruct + + def process_prompt( + self, + prompt: str | None, + input_modalities: list[str], + output_modalities: list[str], + tensors: NameToTensorList | None = None, + **kwargs: Any, + ) -> NameToTensorList: + """Validate a request against the checkpoint variant and tokenize it. + + Produces the four ``PREFILL_INPUTS`` tensors. ``text_inputs`` is the + optional instruction turn followed by the assistant turn (both in the + reference ChatML templates); ``prompt_layout`` is + ``[instruct_len, text_len, stream_text]``. ``stream_text`` follows the + reference default per variant (whole text in the prefill for + CustomVoice/VoiceDesign, one token per frame for Base) unless the + request sets ``non_streaming_mode``. + """ + del tensors + if not prompt: + raise ValueError("Qwen3-TTS requires a non-empty text prompt") + if set(input_modalities) != {"text"}: + raise ValueError("Qwen3-TTS currently supports text input only") + if set(output_modalities) != {"audio"}: + raise ValueError("Qwen3-TTS supports audio output only") + if self.config.is_base: + raise ValueError( + "Qwen3-TTS Base clones a voice from reference audio, which " + "this build does not accept yet; use a CustomVoice or " + "VoiceDesign checkpoint for text-only requests" + ) + + speaker, speaker_id = self._resolve_speaker(kwargs) + language_id = self._resolve_language(kwargs, speaker) + instruct = self._resolve_instruct(kwargs) + stream_text = not bool( + kwargs.get("non_streaming_mode", self.config.default_non_streaming_mode) ) - text_inputs = encoded["input_ids"] - if text_inputs.ndim == 2: - text_inputs = text_inputs[0] - language_id = self.config.talker.codec_language_id.get(language, -1) + assistant_ids = self._tokenize(self.ASSISTANT_TEMPLATE.format(text=prompt)) + wrapper_len = len(CHATML_ASSISTANT_PREFIX_TOKEN_IDS) + len( + CHATML_ASSISTANT_SUFFIX_TOKEN_IDS + ) + text_len = assistant_ids.numel() - wrapper_len + if text_len < 1: + raise ValueError("Qwen3-TTS prompt tokenized to no text tokens") + instruct_ids = ( + self._tokenize(self.INSTRUCT_TEMPLATE.format(instruct=instruct)) + if instruct + else assistant_ids.new_empty(0) + ) return { - "text_inputs": [text_inputs.to(dtype=torch.long)], - "speaker_id": [torch.tensor( - [self.config.talker.spk_id[speaker]], dtype=torch.long + "text_inputs": [torch.cat([instruct_ids, assistant_ids])], + "prompt_layout": [torch.tensor( + [instruct_ids.numel(), text_len, int(stream_text)], dtype=torch.long )], + "speaker_id": [torch.tensor([speaker_id], dtype=torch.long)], "language_id": [torch.tensor([language_id], dtype=torch.long)], } @@ -500,7 +639,7 @@ def get_initial_forward_pass_args( }, ) inputs = [] - for name in ("text_inputs", "speaker_id", "language_id"): + for name in self.PREFILL_INPUTS: edge = GraphEdge(next_node="Talker", name=name) edge.tensor_info = input_signals.get(name, []) inputs.append(edge) @@ -694,22 +833,6 @@ def get_submodule( self._submodule_cache[node_name] = submodule return submodule - @staticmethod - def _verify_loaded( - module: torch.nn.Module, - loaded: set[str], - component: str, - ) -> None: - """Fail startup if checkpoint filtering left any parameter uninitialized.""" - expected = set(dict(module.named_parameters())) - missing = sorted(expected - loaded) - if missing: - preview = ", ".join(missing[:8]) - raise RuntimeError( - f"{component} checkpoint did not initialize {len(missing)} " - f"parameters: {preview}" - ) - def _create_talker_submodule( self, device: str, @@ -746,7 +869,12 @@ def talker_weights(): talker_weights(), stacked_params=LLAMA_STACKED_PARAMS, ) - self._verify_loaded(talker, loaded, "Qwen3-TTS Talker") + cp_prefix = "talker.code_predictor." + talker_keys = { + key for key in _checkpoint_keys(self.local_dir, "talker.") + if not key.startswith("code_predictor.") + } + _verify_checkpoint_coverage(talker, loaded, talker_keys, "Qwen3-TTS Talker") talker.eval() # CodePredictor is small and depth-wise. It is loaded separately from @@ -756,7 +884,6 @@ def talker_weights(): if autocast_dtype is not None: code_predictor = code_predictor.to(autocast_dtype) code_predictor.to_empty(device=device) - cp_prefix = "talker.code_predictor." cp_weights = ( (name.removeprefix(cp_prefix), tensor) for name, tensor in iter_safetensors_shards( @@ -768,8 +895,11 @@ def talker_weights(): cp_weights, stacked_params=LLAMA_STACKED_PARAMS, ) - self._verify_loaded( - code_predictor, loaded, "Qwen3-TTS CodePredictor" + _verify_checkpoint_coverage( + code_predictor, + loaded, + _checkpoint_keys(self.local_dir, cp_prefix), + "Qwen3-TTS CodePredictor", ) # The captured depth loop indexes all residual LM heads as one tensor; # consolidate after the individual checkpoint heads are loaded. @@ -812,7 +942,9 @@ def _create_codec_submodule(self, device: str) -> NodeSubmodule: ) ) loaded = load_hf_weights(decoder, weights) - self._verify_loaded(decoder, loaded, "Qwen3-TTS Codec") + _verify_checkpoint_coverage( + decoder, loaded, _checkpoint_keys(codec_dir, prefix), "Qwen3-TTS Codec" + ) decoder.eval() return CodecSubmodule(decoder, self.config) From 83c160b30fbf02d9f3e616566b9c933d142aaeaa Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 007/110] Register the Qwen3-TTS 1.7B variants --- mstar/model/registry.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/mstar/model/registry.py b/mstar/model/registry.py index 61276ebf9..63e045ce0 100644 --- a/mstar/model/registry.py +++ b/mstar/model/registry.py @@ -13,6 +13,9 @@ "pi05": ("mstar.model.pi05.pi05_model", "Pi05Model"), "qwen3_omni": ("mstar.model.qwen3_omni.qwen3_omni_model", "Qwen3OmniModel"), "qwen3_tts": ("mstar.model.qwen3_tts.qwen3_tts_model", "Qwen3TTSModel"), + "qwen3_tts_1p7b": ("mstar.model.qwen3_tts.qwen3_tts_model", "Qwen3TTSModel"), + "qwen3_tts_voicedesign": ("mstar.model.qwen3_tts.qwen3_tts_model", "Qwen3TTSModel"), + "qwen3_tts_base": ("mstar.model.qwen3_tts.qwen3_tts_model", "Qwen3TTSModel"), "vjepa2": ("mstar.model.vjepa2.vjepa2_model", "VJepa2Model"), "vjepa2_ac": ("mstar.model.vjepa2.vjepa2_model", "VJepa2ACModel"), "wan22": ("mstar.model.wan22.wan22_model", "Wan22Model"), @@ -45,7 +48,14 @@ # state-dict remap inside Pi05Model.get_submodule(). "pi05": {"model_path_hf": "lerobot/pi05_base"}, "qwen3_omni": {"model_path_hf": "Qwen/Qwen3-Omni-30B-A3B-Instruct"}, + # Qwen3-TTS 12 Hz family: one class, the variant is read from config.json. + # CustomVoice = built-in speakers (1.7B also takes style instructions), + # VoiceDesign = voice described by an instruction, Base = voice cloned + # from reference audio. "qwen3_tts": {"model_path_hf": "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"}, + "qwen3_tts_1p7b": {"model_path_hf": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"}, + "qwen3_tts_voicedesign": {"model_path_hf": "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"}, + "qwen3_tts_base": {"model_path_hf": "Qwen/Qwen3-TTS-12Hz-1.7B-Base"}, # V-JEPA 2 standard (encoder + masked predictor). Default is ViT-L @ 256 # (~300M); the same class loads vitl/h/g at 256 or 384 by reading # config.json. From 5092cd4ae9d9e0945fc22c8ca30719e109759cc2 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 008/110] cli: default configs and hints for the Qwen3-TTS 1.7B variants --- mstar/cli/main.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/mstar/cli/main.py b/mstar/cli/main.py index 7784fe0c6..dafcba9d1 100644 --- a/mstar/cli/main.py +++ b/mstar/cli/main.py @@ -32,6 +32,9 @@ "qwen3_omni": "qwen3omni_2gpu.yaml", "qwen3_tts": "qwen3tts.yaml", "omnivoice": "omnivoice.yaml", + "qwen3_tts_1p7b": "qwen3tts_1p7b.yaml", + "qwen3_tts_voicedesign": "qwen3tts_voicedesign.yaml", + "qwen3_tts_base": "qwen3tts_base.yaml", "pi05": "pi05.yaml", "vjepa2": "vjepa2.yaml", "vjepa2_ac": "vjepa2_ac.yaml", @@ -112,13 +115,19 @@ def _next_steps(model: str, host: str, port: int) -> str: if model == "omnivoice": lines.append(" client.tts(\"Xin chào\", language=\"Vietnamese\").to_wav(\"out.wav\")") lines.append(" # clone a voice: ref_audio=\"ref.wav\", ref_text=\"\"") - if model in ("orpheus", "qwen3_omni", "qwen3_tts"): + if model in ("orpheus", "qwen3_omni", "qwen3_tts", "qwen3_tts_1p7b"): voice = { "orpheus": "tara", "qwen3_omni": "Ethan", "qwen3_tts": "Vivian", + "qwen3_tts_1p7b": "Vivian", }[model] lines.append(f" client.tts(\"Hello there\", voice=\"{voice}\").to_wav(\"out.wav\")") + if model == "qwen3_tts_voicedesign": + lines.append(" client.tts(\"Hello there\", instruct=\"A calm, warm female voice\").to_wav(\"out.wav\")") + if model == "qwen3_tts_base": + lines.append(" client.tts(\"Hello there\", reference_audio=\"ref.wav\", " + "ref_text=\"...\").to_wav(\"out.wav\")") if model in ("pi05", "vjepa2", "vjepa2_ac"): lines.append(" res = client.generate(text=\"...\", output_modalities=(\"" + ("action" if model == "pi05" else "video") + "\",))") From d63c7c939a2af1a43960fb96bde58300d192d058 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 009/110] Add the Qwen3-TTS 1.7B CustomVoice deployment config --- configs/qwen3tts_1p7b.yaml | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 configs/qwen3tts_1p7b.yaml diff --git a/configs/qwen3tts_1p7b.yaml b/configs/qwen3tts_1p7b.yaml new file mode 100644 index 000000000..c3d244e1f --- /dev/null +++ b/configs/qwen3tts_1p7b.yaml @@ -0,0 +1,18 @@ +# Qwen3-TTS-12Hz-1.7B-CustomVoice: built-in speakers, optional style instruction (instruct=...) +# Same graph, resources and single-GPU placement as the 0.6B deployment +# (configs/qwen3tts.yaml); only the checkpoint differs. KV geometry is +# identical across sizes (28 layers, 8 KV heads, head_dim 128). +model: "qwen3_tts_1p7b" +max_seq_len: 32768 +# The supported deployment image cannot build FlashInfer's Hopper FA3 JIT +# kernels yet. Keep the model default on "auto" and pin this deployment to FA2. +resources: + talker_attn: + flashinfer_backend: fa2 +node_groups: + - node_names: [Talker] + ranks: [0] + graph_walks: [talker_prefill, talker_decode] + - node_names: [Codec] + ranks: [0] + graph_walks: [codec_chunk] From 297ebedba94ab27e2000a77ef42e0a29ac569154 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 010/110] Add the Qwen3-TTS VoiceDesign deployment config --- configs/qwen3tts_voicedesign.yaml | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 configs/qwen3tts_voicedesign.yaml diff --git a/configs/qwen3tts_voicedesign.yaml b/configs/qwen3tts_voicedesign.yaml new file mode 100644 index 000000000..3d2477a4c --- /dev/null +++ b/configs/qwen3tts_voicedesign.yaml @@ -0,0 +1,18 @@ +# Qwen3-TTS-12Hz-1.7B-VoiceDesign: the voice is described by the request's instruct text +# Same graph, resources and single-GPU placement as the 0.6B deployment +# (configs/qwen3tts.yaml); only the checkpoint differs. KV geometry is +# identical across sizes (28 layers, 8 KV heads, head_dim 128). +model: "qwen3_tts_voicedesign" +max_seq_len: 32768 +# The supported deployment image cannot build FlashInfer's Hopper FA3 JIT +# kernels yet. Keep the model default on "auto" and pin this deployment to FA2. +resources: + talker_attn: + flashinfer_backend: fa2 +node_groups: + - node_names: [Talker] + ranks: [0] + graph_walks: [talker_prefill, talker_decode] + - node_names: [Codec] + ranks: [0] + graph_walks: [codec_chunk] From db2af3d72f516003ffc235b52df347ead3fd3a1a Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 011/110] Add the Qwen3-TTS Base deployment config --- configs/qwen3tts_base.yaml | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 configs/qwen3tts_base.yaml diff --git a/configs/qwen3tts_base.yaml b/configs/qwen3tts_base.yaml new file mode 100644 index 000000000..621e9cdd5 --- /dev/null +++ b/configs/qwen3tts_base.yaml @@ -0,0 +1,18 @@ +# Qwen3-TTS-12Hz-1.7B-Base: zero-shot voice clone from reference audio (speaker encoder + optional ICL) +# Same graph, resources and single-GPU placement as the 0.6B deployment +# (configs/qwen3tts.yaml); only the checkpoint differs. KV geometry is +# identical across sizes (28 layers, 8 KV heads, head_dim 128). +model: "qwen3_tts_base" +max_seq_len: 32768 +# The supported deployment image cannot build FlashInfer's Hopper FA3 JIT +# kernels yet. Keep the model default on "auto" and pin this deployment to FA2. +resources: + talker_attn: + flashinfer_backend: fa2 +node_groups: + - node_names: [Talker] + ranks: [0] + graph_walks: [talker_prefill, talker_decode] + - node_names: [Codec] + ranks: [0] + graph_walks: [codec_chunk] From 13022b72eea5a8169424374b68e15b4d465f10c9 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 012/110] api: OpenAI speech adapter for Qwen3-TTS --- mstar/api_server/openai/adapters.py | 35 ++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/mstar/api_server/openai/adapters.py b/mstar/api_server/openai/adapters.py index 5a715e273..e979339c2 100644 --- a/mstar/api_server/openai/adapters.py +++ b/mstar/api_server/openai/adapters.py @@ -325,7 +325,6 @@ def speech_to_request(self, req: SpeechRequest, upload_dir: Path) -> SubmitArgs: ) - class OmniVoiceAdapter(OpenAIAdapter): """OmniVoice: zero-shot TTS in three modes over one endpoint. @@ -382,6 +381,36 @@ def speech_to_request(self, req: SpeechRequest, upload_dir: Path) -> SubmitArgs: ) +class Qwen3TTSAdapter(OpenAIAdapter): + """Qwen3-TTS (CustomVoice / VoiceDesign / Base): text-to-speech, audio only. + + ``voice`` selects a built-in speaker (CustomVoice). ``instructions`` (the + OpenAI field; ``instruct`` also accepted) carries the style or voice + description. Non-standard knobs travel through ``extra_body``: + ``language``, ``non_streaming_mode``, ``top_k``, ``repetition_penalty``, + the residual-group ``subtalker_*`` sampling, ``max_new_tokens``. + ``temperature`` / ``top_p`` / ``seed`` map onto the Talker sampler. + """ + + supports_speech = True + + def speech_to_request(self, req: SpeechRequest, upload_dir: Path) -> SubmitArgs: # noqa: ARG002 + mk = _passthrough(req) + if getattr(req, "voice", None): + mk["voice"] = req.voice + # OpenAI's field is ``instructions``; the model reads ``instruct``. + instructions = mk.pop("instructions", None) + if instructions: + mk.setdefault("instruct", instructions) + _apply_sampling(req, mk, temperature_key="temperature", top_p_key="top_p", max_tokens_key=None) + return SubmitArgs( + text=req.input, + input_modalities=["text"], + output_modalities=["audio"], + model_kwargs=mk, + ) + + class Cosmos3Adapter(OpenAIAdapter): """NVIDIA Cosmos3: text-to-image and text/image-to-video generation. @@ -513,6 +542,10 @@ def video_to_request(self, req: VideoGenerationRequest, upload_dir: Path) -> Sub "qwen3_omni": Qwen3OmniAdapter(), "omnivoice": OmniVoiceAdapter(), "orpheus": OrpheusAdapter(), + "qwen3_tts": Qwen3TTSAdapter(), + "qwen3_tts_1p7b": Qwen3TTSAdapter(), + "qwen3_tts_voicedesign": Qwen3TTSAdapter(), + "qwen3_tts_base": Qwen3TTSAdapter(), "cosmos3": Cosmos3Adapter(), "cosmos3_droid": Cosmos3Adapter(), "cosmos3_super": Cosmos3Adapter(), From f6410a304adff5cc79a2b3b21a81b165cf364b2e Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 013/110] benchmark: Qwen3-TTS 1.7B and VoiceDesign model entries --- benchmark/base.py | 35 +++++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/benchmark/base.py b/benchmark/base.py index 7a9b644d0..9e4bbe30b 100644 --- a/benchmark/base.py +++ b/benchmark/base.py @@ -233,15 +233,40 @@ def get_supported_modalities(self): class Qwen3TTS(Model): - """Qwen3-TTS CustomVoice benchmark metadata for native M* requests.""" + """Qwen3-TTS CustomVoice benchmark metadata (0.6B by default). + + ``/v1/audio/speech`` requests carry the same ``voice`` and ``language`` + for every engine so the Talker prefill is identical across systems. + """ + + HF_URL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" def get_hf_url(self): - return "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" + return self.HF_URL + + def get_model_kwargs(self, request_type: RequestType): + return {"voice": "vivian", "language": "English"} def get_supported_modalities(self): return {RequestType.T2S} +class Qwen3TTS1p7B(Qwen3TTS): + HF_URL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" + + +class Qwen3TTSVoiceDesign(Qwen3TTS): + """VoiceDesign has no built-in speakers; the voice is the instruction.""" + + HF_URL = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign" + + def get_model_kwargs(self, request_type: RequestType): + return { + "language": "English", + "instructions": "A clear, friendly adult female voice with a neutral accent.", + } + + class Pi05(Model): """Physical Intelligence Pi0.5 VLA model. @@ -328,6 +353,8 @@ class ModelType(Enum): ORPHEUS = "orpheus" QWEN3OMNI = "qwen3omni" QWEN3TTS = "qwen3_tts" + QWEN3TTS_1P7B = "qwen3_tts_1p7b" + QWEN3TTS_VOICEDESIGN = "qwen3_tts_voicedesign" PI05 = "pi05" VJEPA2AC = "vjepa2ac" WHISPER_LARGE = "whisper_large" @@ -342,6 +369,10 @@ def inst(self, **kwargs) -> Model: return Qwen3Omni(**kwargs) if self == ModelType.QWEN3TTS: return Qwen3TTS(**kwargs) + if self == ModelType.QWEN3TTS_1P7B: + return Qwen3TTS1p7B(**kwargs) + if self == ModelType.QWEN3TTS_VOICEDESIGN: + return Qwen3TTSVoiceDesign(**kwargs) if self == ModelType.PI05: return Pi05(**kwargs) if self == ModelType.VJEPA2AC: From 3a87ece6d6e61cb144f62e8bfb655083f4431b8c Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 014/110] benchmark: route Qwen3-TTS through /v1/audio/speech --- benchmark/request.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmark/request.py b/benchmark/request.py index 478ddda7c..ff3585eae 100644 --- a/benchmark/request.py +++ b/benchmark/request.py @@ -14,7 +14,7 @@ import aiohttp import numpy as np -from benchmark.base import Bagel, Model, Orpheus, RequestType, Status +from benchmark.base import Bagel, Model, Orpheus, Qwen3TTS, RequestType, Status from benchmark.utils import _write_wav @@ -1494,7 +1494,7 @@ async def send_request( metrics=metrics, additional_model_kwargs=additional_model_kwargs, ) - if req_type.get_output_modalities() == "audio" and isinstance(model, Orpheus): + if req_type.get_output_modalities() == "audio" and isinstance(model, (Orpheus, Qwen3TTS)): metrics = RequestMetrics( request_id=request_id, type=req_type, From 8f32c650d098b4ea5533009048717518514305a6 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 015/110] test: Qwen3-TTS variants, prefill layouts and predictor projection --- test/modular/test_qwen3_tts_model.py | 254 ++++++++++++++++++++++++++- 1 file changed, 249 insertions(+), 5 deletions(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index eed6b957d..7ba27bbbd 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -42,14 +42,34 @@ CONFIG_PATH = Path(__file__).resolve().parents[2] / "configs" / "qwen3tts.yaml" +ASSISTANT_PREFIX = [151644, 77091, 198] +ASSISTANT_SUFFIX = [151645, 198, 151644, 77091, 198] +USER_PREFIX = [151644, 872, 198] +USER_SUFFIX = [151645, 198] + + class _TokenizerStub: + """Tokenizes the two reference templates: fixed ChatML wrappers, one id per word.""" + def __init__(self): - self.last_text = None + self.texts = [] + + @property + def last_text(self): + return self.texts[-1] if self.texts else None def __call__(self, text, **kwargs): - self.last_text = text + self.texts.append(text) assert kwargs == {"return_tensors": "pt", "padding": True} - return {"input_ids": torch.tensor([[1, 2, 3]])} + if text.startswith("<|im_start|>assistant\n"): + body = text[len("<|im_start|>assistant\n"):-len("<|im_end|>\n<|im_start|>assistant\n")] + prefix, suffix = ASSISTANT_PREFIX, ASSISTANT_SUFFIX + else: + assert text.startswith("<|im_start|>user\n") + body = text[len("<|im_start|>user\n"):-len("<|im_end|>\n")] + prefix, suffix = USER_PREFIX, USER_SUFFIX + words = [1000 + i for i, _ in enumerate(body.split())] + return {"input_ids": torch.tensor([prefix + words + suffix])} def _make_model() -> Qwen3TTSModel: @@ -186,6 +206,32 @@ def test_qwen3_tts_registry_engines_cache_and_yaml_are_consistent(): assert by_walk["codec_chunk"].consumes_stream is True +QWEN3_TTS_VARIANTS = { + "qwen3_tts_1p7b": ("Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", "qwen3tts_1p7b.yaml"), + "qwen3_tts_voicedesign": ("Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign", "qwen3tts_voicedesign.yaml"), + "qwen3_tts_base": ("Qwen/Qwen3-TTS-12Hz-1.7B-Base", "qwen3tts_base.yaml"), +} + + +def test_qwen3_tts_1p7b_variants_share_class_configs_and_adapter(): + from mstar.api_server.openai.adapters import Qwen3TTSAdapter, get_adapter + from mstar.cli.main import DEFAULT_CONFIGS + + base_yaml = yaml.safe_load(CONFIG_PATH.read_text(encoding="utf-8")) + for key, (hf_id, yaml_name) in QWEN3_TTS_VARIANTS.items(): + assert get_model_class(key) is Qwen3TTSModel + assert HF_MODELS[key] == {"model_path_hf": hf_id} + assert DEFAULT_CONFIGS[key] == yaml_name + deployment = yaml.safe_load( + (CONFIG_PATH.parent / yaml_name).read_text(encoding="utf-8") + ) + assert deployment["model"] == key + assert deployment["node_groups"] == base_yaml["node_groups"] + assert deployment["resources"] == base_yaml["resources"] + assert isinstance(get_adapter(key), Qwen3TTSAdapter) + assert isinstance(get_adapter("qwen3_tts"), Qwen3TTSAdapter) + + def test_qwen3_tts_cli_and_benchmark_entries_are_registered(): repo_root = str(Path(__file__).resolve().parents[2]) sys.path.insert(0, repo_root) @@ -284,11 +330,85 @@ def test_qwen3_tts_process_prompt_matches_official_template(): "<|im_start|>assistant\n你好<|im_end|>\n" "<|im_start|>assistant\n" ) - assert tensors["text_inputs"][0].tolist() == [1, 2, 3] + assert tensors["text_inputs"][0].tolist() == ASSISTANT_PREFIX + [1000] + ASSISTANT_SUFFIX + # CustomVoice default: whole text in the prefill (stream_text = 0). + assert tensors["prompt_layout"][0].tolist() == [0, 1, 0] assert tensors["speaker_id"][0].item() == 3065 assert tensors["language_id"][0].item() == 2055 +def _variant_model(tts_model_type: str, tts_model_size: str = "1b7") -> Qwen3TTSModel: + model = _make_model() + talker = Qwen3TTSTalkerConfig(hidden_size=2048, intermediate_size=6144) + if tts_model_type != "custom_voice": + talker.spk_id = {} + talker.spk_is_dialect = {} + model.config = Qwen3TTSModelConfig( + tts_model_type=tts_model_type, tts_model_size=tts_model_size, talker=talker + ) + return model + + +def test_qwen3_tts_1p7b_custom_voice_prepends_instruction_turn(): + model = _variant_model("custom_voice") + assert model.config.supports_instruct and not model.config.requires_instruct + + tensors = model.process_prompt( + "hello big world", + input_modalities=["text"], + output_modalities=["audio"], + voice="Ryan", + instructions="speak slowly", + non_streaming_mode=False, + ) + + instruct_ids = USER_PREFIX + [1000, 1001] + USER_SUFFIX + assistant_ids = ASSISTANT_PREFIX + [1000, 1001, 1002] + ASSISTANT_SUFFIX + assert model.tokenizer.texts[-1] == "<|im_start|>user\nspeak slowly<|im_end|>\n" + assert tensors["text_inputs"][0].tolist() == instruct_ids + assistant_ids + assert tensors["prompt_layout"][0].tolist() == [len(instruct_ids), 3, 1] + assert tensors["speaker_id"][0].item() == 3061 + + +def test_qwen3_tts_voice_design_requires_instruct_and_has_no_speakers(): + model = _variant_model("voice_design") + assert model.config.default_speaker is None + assert model.config.requires_instruct + + tensors = model.process_prompt( + "hello", + input_modalities=["text"], + output_modalities=["audio"], + instruct="A deep, calm male voice", + ) + assert tensors["speaker_id"][0].item() == -1 + assert tensors["prompt_layout"][0].tolist() == [3 + 5 + 2, 1, 0] + + with pytest.raises(ValueError, match="requires an 'instruct'"): + model.process_prompt("hello", input_modalities=["text"], output_modalities=["audio"]) + with pytest.raises(ValueError, match="no built-in speakers"): + model.process_prompt( + "hello", input_modalities=["text"], output_modalities=["audio"], + voice="vivian", instruct="x", + ) + + +def test_qwen3_tts_base_config_declares_speaker_encoder(): + model = _variant_model("base") + assert model.config.supports_reference_audio + assert model.config.speaker_encoder is not None + assert model.config.speaker_encoder.enc_dim == 2048 + # Base feeds text one token per frame by default (reference default). + assert model.config.default_non_streaming_mode is False + with pytest.raises(ValueError, match="reference audio"): + model.process_prompt("hello", input_modalities=["text"], output_modalities=["audio"]) + + +def test_qwen3_tts_config_rejects_unknown_variant(): + with pytest.raises(ValueError, match="tts_model_type"): + Qwen3TTSModelConfig(tts_model_type="duplex") + + def test_qwen3_tts_validates_speaker_dialect_after_language_override(): model = _make_model() @@ -352,8 +472,11 @@ def test_qwen3_tts_initial_partition_args_route_expected_inputs(): model = _make_model() pointers = { name: [SimpleNamespace(name=name)] - for name in ("text_inputs", "speaker_id", "language_id") + for name in Qwen3TTSModel.PREFILL_INPUTS } + assert Qwen3TTSModel.PREFILL_INPUTS == ( + "text_inputs", "prompt_layout", "speaker_id", "language_id", + ) talker = model.get_initial_forward_pass_args( "Talker", @@ -491,13 +614,16 @@ def test_qwen3_tts_talker_builds_official_streaming_prefill(): submodule.CHATML_ASSISTANT_PREFIX_TOKEN_IDS = (1, 2, 3) submodule.CHATML_ASSISTANT_SUFFIX_TOKEN_IDS = (8, 9, 10, 11, 12) + # 3 prefix + 4 text + 5 suffix tokens, streaming text layout. embeds = submodule._build_prefill( request_id="request", text_ids=torch.arange(1, 13), + prompt_layout=torch.tensor([0, 4, 1]), speaker_id=40, language_id=-1, ) + # role(3) + [nothink, think_bos, think_eos, speaker, pad](5) + first text token assert embeds.shape == (9, 16) state = submodule.request_state("request") assert state["trailing_text_hidden"].shape == (4, 16) @@ -505,6 +631,68 @@ def test_qwen3_tts_talker_builds_official_streaming_prefill(): assert state["generation_step"] == 0 +def test_qwen3_tts_talker_builds_official_non_streaming_prefill(): + config = _tiny_model_config() + submodule = TalkerSubmodule( + Qwen3TTSTalkerModel(config), Qwen3TTSCodePredictor(config), config + ) + submodule.CHATML_ASSISTANT_PREFIX_TOKEN_IDS = (1, 2, 3) + submodule.CHATML_ASSISTANT_SUFFIX_TOKEN_IDS = (8, 9, 10, 11, 12) + + embeds = submodule._build_prefill( + request_id="request", + text_ids=torch.arange(1, 13), + prompt_layout=torch.tensor([0, 4, 0]), + speaker_id=40, + language_id=41, + ) + + # role(3) + [think, think_bos, lang, think_eos, speaker, pad](6) + # + (4 text + tts_eos) over codec pads (5) + (tts_pad + codec_bos)(1) + assert embeds.shape == (15, 16) + state = submodule.request_state("request") + # Nothing streams: every decode frame adds the TTS PAD embedding. + assert state["trailing_text_hidden"].shape == (0, 16) + prepared = submodule.prepare_inputs( + "talker_decode", + SimpleNamespace(request_id="request"), + {"talker_input_embeds": [torch.zeros(1, 16)]}, + ) + assert torch.equal(prepared.input_embeds[0], state["tts_pad_embed"]) + + +def test_qwen3_tts_talker_prefill_prepends_instruction_without_speaker(): + config = _tiny_model_config() + submodule = TalkerSubmodule( + Qwen3TTSTalkerModel(config), Qwen3TTSCodePredictor(config), config + ) + submodule.CHATML_ASSISTANT_PREFIX_TOKEN_IDS = (1, 2, 3) + submodule.CHATML_ASSISTANT_SUFFIX_TOKEN_IDS = (8, 9, 10, 11, 12) + instruct = torch.tensor([20, 21, 22, 23, 24, 25]) + text_ids = torch.cat([instruct, torch.arange(1, 13)]) + + embeds = submodule._build_prefill( + request_id="request", + text_ids=text_ids, + prompt_layout=torch.tensor([6, 4, 1]), + speaker_id=-1, + language_id=-1, + ) + + # instruct(6) + role(3) + [nothink, think_bos, think_eos, pad](4) + first text + assert embeds.shape == (14, 16) + # A layout whose text span disagrees with the token stream is rejected + # even when the ChatML wrapper itself still lines up. + with pytest.raises(ValueError, match="prompt layout disagrees"): + submodule._build_prefill( + request_id="request", + text_ids=text_ids, + prompt_layout=torch.tensor([6, 3, 1]), + speaker_id=-1, + language_id=-1, + ) + + def test_qwen3_tts_talker_rejects_changed_chatml_layout(): config = _tiny_model_config() submodule = TalkerSubmodule( @@ -521,6 +709,7 @@ def test_qwen3_tts_talker_rejects_changed_chatml_layout(): submodule._build_prefill( request_id="request", text_ids=text_ids, + prompt_layout=torch.tensor([0, 4, 1]), speaker_id=40, language_id=-1, ) @@ -713,6 +902,61 @@ def test_qwen3_tts_talker_batches_and_captures_decode(): assert submodule.can_use_cuda_graphs(batch, model_inputs) +def test_qwen3_tts_code_predictor_projects_wider_talker_inputs(): + """1.7B: Talker width 2048 vs predictor width 1024 -> biased projection on + every depth input; 0.6B (equal widths) -> identity, no extra parameters.""" + narrow = _tiny_model_config() + assert isinstance( + Qwen3TTSCodePredictor(narrow).small_to_mtp_projection, torch.nn.Identity + ) + + wide = _tiny_model_config() + wide.talker.hidden_size = 32 + predictor = Qwen3TTSCodePredictor(wide) + projection = predictor.small_to_mtp_projection + assert isinstance(projection, torch.nn.Linear) + assert projection.weight.shape == (16, 32) + assert projection.bias.shape == (16,) + # Residual embedding tables stay in the Talker width: their sum feeds the + # next Talker step, only the predictor input is projected. + assert predictor.model.codec_embedding[0].weight.shape == (32, 32) + assert {"small_to_mtp_projection.weight", "small_to_mtp_projection.bias"} <= set( + dict(predictor.named_parameters()) + ) + + for layer in predictor.model.layers: + layer.input_layernorm = torch.nn.Identity() + layer.post_attention_layernorm = torch.nn.Identity() + layer.self_attn.q_norm = torch.nn.Identity() + layer.self_attn.k_norm = torch.nn.Identity() + predictor.model.norm = torch.nn.Identity() + import mstar.model.qwen3_tts.components.talker as talker_module + original_rope = talker_module.apply_rope_pos_ids + original_attn = talker_module.decode_attn_nhd + talker_module.apply_rope_pos_ids = lambda q, k, pos, theta: (q, k) + talker_module.decode_attn_nhd = lambda q, k_cache, v_cache, n: ( + torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), k_cache[:, :n].transpose(1, 2), + v_cache[:, :n].transpose(1, 2), enable_gqa=True, + ).transpose(1, 2) + ) + try: + cp = wide.talker.code_predictor + out = predictor.forward_depth_unrolled( + inputs_embeds=torch.randn(2, 1, 32), + position_ids=torch.zeros(2, 1, dtype=torch.long), + kv_cache=torch.zeros( + cp.num_hidden_layers, 2, 2, wide.talker.num_code_groups, + cp.num_key_value_heads, cp.head_dim, + ), + cache_pos=0, + ) + finally: + talker_module.apply_rope_pos_ids = original_rope + talker_module.decode_attn_nhd = original_attn + assert out.shape == (2, 1, 16) + + def test_qwen3_tts_code_predictor_uses_decode_attn_nhd(monkeypatch): config = _tiny_model_config() predictor = Qwen3TTSCodePredictor(config) From 20db9eeecdb3f68317fa3e476c54d95f51ca4e7f Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 016/110] test: Qwen3-TTS speech adapter --- test/modular/test_openai_adapters.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/test/modular/test_openai_adapters.py b/test/modular/test_openai_adapters.py index 1ecddbfad..ff18a8421 100644 --- a/test/modular/test_openai_adapters.py +++ b/test/modular/test_openai_adapters.py @@ -74,6 +74,25 @@ def test_qwen3_speech_maps_talker_sampling(tmp_path): assert sa.model_kwargs["voice"] == "Ethan" +def test_qwen3_tts_speech_maps_voice_instructions_and_extra_body(tmp_path): + req = SpeechRequest( + input="hello", voice="Vivian", instructions="speak slowly", + temperature=0.7, top_p=0.9, seed=3, language="English", non_streaming_mode=False, + ) + sa = adapters.Qwen3TTSAdapter().speech_to_request(req, tmp_path) + assert sa.text == "hello" + assert sa.input_modalities == ["text"] and sa.output_modalities == ["audio"] + mk = sa.model_kwargs + assert mk["voice"] == "Vivian" + # OpenAI's ``instructions`` becomes the model's ``instruct`` (and only that). + assert mk["instruct"] == "speak slowly" and "instructions" not in mk + assert (mk["temperature"], mk["top_p"], mk["seed"]) == (0.7, 0.9, 3) + assert mk["language"] == "English" and mk["non_streaming_mode"] is False + assert "max_output_tokens" not in mk + for key in ("qwen3_tts", "qwen3_tts_1p7b", "qwen3_tts_voicedesign", "qwen3_tts_base"): + assert isinstance(adapters.get_adapter(key), adapters.Qwen3TTSAdapter) + + def test_chat_and_image_honor_seed(tmp_path): chat = ChatCompletionRequest(model="bagel", messages=[{"role": "user", "content": "x"}], seed=7) assert adapters.BagelAdapter().chat_to_request(chat, tmp_path).model_kwargs["seed"] == 7 From 1b66b1e495a7bb311bcd09988504cbf48d81b885 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 017/110] test: Qwen3-TTS 1.7B real-weight loading checks --- .../test_qwen3_tts_1p7b_real_weights.py | 134 ++++++++++++++++++ 1 file changed, 134 insertions(+) create mode 100644 test/integration/test_qwen3_tts_1p7b_real_weights.py diff --git a/test/integration/test_qwen3_tts_1p7b_real_weights.py b/test/integration/test_qwen3_tts_1p7b_real_weights.py new file mode 100644 index 000000000..28346ad6c --- /dev/null +++ b/test/integration/test_qwen3_tts_1p7b_real_weights.py @@ -0,0 +1,134 @@ +"""Real-weight loading checks for the Qwen3-TTS 1.7B family. + +Never downloads. Each variant is exercised only when its checkpoint is already +in the local Hugging Face cache and CUDA is available. The parity harness +(``test/qwen3-tts/parity_qwen3_tts.py``) is the functional check; this file +pins what loading must get right for every variant: complete checkpoint +coverage, the Talker-to-predictor projection, and the prefill layouts. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from mstar.model.qwen3_tts.qwen3_tts_model import Qwen3TTSModel + +VARIANTS = { + "custom_voice": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", + "voice_design": "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign", + "base": "Qwen/Qwen3-TTS-12Hz-1.7B-Base", +} +# talker.* minus code_predictor.* | code_predictor.* | speech_tokenizer decoder.* +TALKER_PARAMS = 1_741_550_592 +CODE_PREDICTOR_PARAMS = 175_125_760 +CODEC_PARAMS = 114_323_137 + + +def _find_cached_snapshot(repo: str) -> Path | None: + repo_dir = f"models--{repo.replace('/', '--')}" + roots = [] + if os.environ.get("HF_HUB_CACHE"): + roots.append(Path(os.environ["HF_HUB_CACHE"])) + if os.environ.get("HF_HOME"): + roots.append(Path(os.environ["HF_HOME"]) / "hub") + roots.append(Path.home() / ".cache" / "huggingface" / "hub") + for root in roots: + snapshots = root / repo_dir / "snapshots" + if not snapshots.is_dir(): + continue + for snapshot in snapshots.iterdir(): + if ( + (snapshot / "model.safetensors").is_file() + and (snapshot / "speech_tokenizer" / "model.safetensors").is_file() + ): + return snapshot + return None + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.fixture(scope="module", params=sorted(VARIANTS)) +def loaded(request): + variant = request.param + snapshot = _find_cached_snapshot(VARIANTS[variant]) + if snapshot is None: + pytest.skip(f"{VARIANTS[variant]} is not in the local Hugging Face cache") + model = Qwen3TTSModel(model_path_hf=str(snapshot)) + talker = model.get_submodule("Talker", device="cuda:0", autocast_dtype=torch.bfloat16) + codec = model.get_submodule("Codec", device="cuda:0") + yield variant, model, talker, codec + del talker, codec + model._submodule_cache.clear() + torch.cuda.empty_cache() + + +def test_variant_metadata_and_weights_load_completely(loaded): + variant, model, talker, codec = loaded + assert model.config.tts_model_type == variant + assert model.config.tts_model_size == "1b7" + assert model.config.talker.hidden_size == 2048 + assert model.config.code_predictor.hidden_size == 1024 + # The coverage check in get_submodule already failed loudly on any missing + # or unused tensor; the counts pin the architecture the checkpoint carries. + assert sum(p.numel() for p in talker.model.parameters()) == TALKER_PARAMS + assert sum(p.numel() for p in talker.code_predictor.parameters()) == CODE_PREDICTOR_PARAMS + assert sum(p.numel() for p in codec.decoder.parameters()) == CODEC_PARAMS + projection = talker.code_predictor.small_to_mtp_projection + assert isinstance(projection, torch.nn.Linear) + assert projection.weight.shape == (1024, 2048) + assert next(talker.model.parameters()).dtype == torch.bfloat16 + assert next(codec.decoder.parameters()).dtype == torch.float32 + + +def test_prefill_layout_matches_variant(loaded): + variant, model, talker, _ = loaded + kwargs = {"input_modalities": ["text"], "output_modalities": ["audio"]} + if variant == "base": + with pytest.raises(ValueError, match="reference audio"): + model.process_prompt("Testing the base model.", **kwargs) + return + if variant == "custom_voice": + tensors = model.process_prompt( + "Testing Qwen three TTS.", voice="Vivian", language="English", + instruct="Speak slowly and clearly.", **kwargs, + ) + assert tensors["speaker_id"][0].item() == model.config.talker.spk_id["vivian"] + else: + tensors = model.process_prompt( + "Testing Qwen three TTS.", instruct="A calm male voice.", **kwargs, + ) + assert tensors["speaker_id"][0].item() == -1 + instruct_len, text_len, stream_text = tensors["prompt_layout"][0].tolist() + assert instruct_len > 0 and text_len > 0 and stream_text == 0 + + prepared = talker.prepare_inputs( + "talker_prefill", SimpleNamespace(request_id=f"prefill-{variant}"), tensors, + ) + assert prepared.input_embeds.shape[1] == model.config.talker.hidden_size + # instruct + role(3) + codec tags + (text + eos) + closing pad/bos + tags = 3 + 1 + (1 if variant == "custom_voice" else 0) + 1 # think..., [speaker], pad + assert prepared.input_seq_len == instruct_len + 3 + tags + (text_len + 1) + 1 + state = talker.request_state(f"prefill-{variant}") + assert state["trailing_text_hidden"].shape == (0, model.config.talker.hidden_size) + + +def test_depth_loop_runs_through_projection(loaded): + variant, model, talker, _ = loaded + del variant + batch = 2 + hidden = torch.randn( + batch, model.config.talker.hidden_size, device="cuda:0", dtype=torch.bfloat16 + ) + layer0 = torch.randint(0, 2048, (batch,), device="cuda:0") + codes, embed_sum = talker._depth_loop(hidden, layer0, lambda logits: logits.argmax(-1)) + torch.cuda.synchronize() + assert codes.shape == (batch, model.config.num_code_groups) + assert (codes[:, 0] == layer0).all() + assert (codes[:, 1:] < model.config.code_predictor.vocab_size).all() + assert embed_sum.shape == (batch, model.config.talker.hidden_size) From 6ba7e924b581dcb76cf0fe31e0c5ffe0101dd442 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:52 -0700 Subject: [PATCH 018/110] benchmark: streaming /v1/audio/speech client shared by all TTS engines --- benchmark/tts_speech_bench.py | 335 ++++++++++++++++++++++++++++++++++ 1 file changed, 335 insertions(+) create mode 100644 benchmark/tts_speech_bench.py diff --git a/benchmark/tts_speech_bench.py b/benchmark/tts_speech_bench.py new file mode 100644 index 000000000..eed1b2ae5 --- /dev/null +++ b/benchmark/tts_speech_bench.py @@ -0,0 +1,335 @@ +#!/usr/bin/env python3 +"""Streaming ``/v1/audio/speech`` benchmark shared by M*, vLLM-Omni and SGLang-Omni. + +One client, one metric definition, every engine (BENCHMARK_PROTOCOL.md, TTS row): + +* time-to-first-audio (TTFA): request start to the first PCM byte (WAV header + excluded), p50 / p95 over the requests of a repeat; +* end-to-end latency and RTF = wall time / seconds of audio produced; +* audio-seconds generated per wall-clock second at the given concurrency; +* the PCM of every request can be written out for a WER check + (``benchmark/tts_wer.py``). + +Requests are closed-loop: ``--concurrency`` requests in flight at all times until +every sentence of the input file has been synthesized once per repeat. Warmup +requests are excluded. Repeats are reported individually and as the median. + +Engine specifics are limited to how audio is streamed: + +* ``mstar``: ``response_format=wav``, one open-ended WAV (44-byte header, then PCM16); +* ``vllm-omni``: ``response_format=pcm`` + ``stream_format=audio`` -> raw PCM16; +* ``sglang-omni``: ``response_format=pcm`` -> raw PCM16. + +Example (same node, back to back, 200 sentences, 3 repeats):: + + python -m benchmark.tts_speech_bench --engine mstar --url http://127.0.0.1:8000 \\ + --model qwen3_tts_1p7b --sentences $BENCH/tts/sentences_200.txt \\ + --voice vivian --language English --concurrency 8 --repeats 3 \\ + --out results/mstar_c8.json --save-audio-dir results/mstar_c8_wav + python -m benchmark.tts_speech_bench --engine vllm-omni --url http://127.0.0.1:8002 \\ + --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice ... +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import statistics +import struct +import subprocess +import sys +import time +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +import aiohttp + +WAV_HEADER_BYTES = 44 +BYTES_PER_SAMPLE = 2 # PCM16 mono + + +@dataclass +class RequestResult: + sentence_id: int + words: int + ttfa_s: float | None + e2e_s: float | None + audio_s: float + pcm_bytes: int + error: str | None = None + + @property + def rtf(self) -> float | None: + if self.e2e_s is None or self.audio_s <= 0: + return None + return self.e2e_s / self.audio_s + + +@dataclass +class RepeatSummary: + repeat: int + requests: int + errors: int + wall_s: float + ttfa_p50_ms: float + ttfa_p95_ms: float + e2e_p50_s: float + e2e_p95_s: float + rtf_mean: float + rtf_p95: float + audio_s_total: float + audio_s_per_wall_s: float + results: list[dict[str, Any]] = field(default_factory=list) + + +def _percentile(values: list[float], p: float) -> float: + if not values: + return float("nan") + ordered = sorted(values) + idx = min(len(ordered) - 1, max(0, round(p * (len(ordered) - 1)))) + return ordered[idx] + + +def _payload(args: argparse.Namespace, text: str) -> dict[str, Any]: + payload: dict[str, Any] = {"model": args.model, "input": text, "stream": True} + if args.engine == "mstar": + payload["response_format"] = "wav" + else: + payload["response_format"] = "pcm" + if args.engine == "vllm-omni": + payload["stream_format"] = "audio" + if args.voice: + payload["voice"] = args.voice + if args.language: + payload["language"] = args.language + if args.instructions: + payload["instructions"] = args.instructions + if args.seed is not None: + payload["seed"] = args.seed + for item in args.extra: + key, _, raw = item.partition("=") + try: + payload[key] = json.loads(raw) + except json.JSONDecodeError: + payload[key] = raw + return payload + + +def _write_wav(path: Path, pcm: bytes, sample_rate: int) -> None: + header = struct.pack( + "<4sI4s4sIHHIIHH4sI", + b"RIFF", 36 + len(pcm), b"WAVE", b"fmt ", 16, 1, 1, sample_rate, + sample_rate * BYTES_PER_SAMPLE, BYTES_PER_SAMPLE, 8 * BYTES_PER_SAMPLE, + b"data", len(pcm), + ) + path.write_bytes(header + pcm) + + +async def _one_request( + session: aiohttp.ClientSession, + args: argparse.Namespace, + sentence_id: int, + text: str, + save_dir: Path | None, +) -> RequestResult: + payload = _payload(args, text) + skip = WAV_HEADER_BYTES if args.engine == "mstar" else 0 + pcm = bytearray() + ttfa = None + start = time.perf_counter() + try: + async with session.post( + f"{args.url}/v1/audio/speech", + json=payload, + headers={"Authorization": "Bearer EMPTY"}, + timeout=aiohttp.ClientTimeout(total=args.timeout, sock_read=args.timeout), + ) as resp: + if resp.status != 200: + body = await resp.text() + return RequestResult(sentence_id, len(text.split()), None, None, 0.0, 0, + error=f"HTTP {resp.status}: {body[:200]}") + async for chunk in resp.content.iter_any(): + if not chunk: + continue + if skip: + drop = min(skip, len(chunk)) + chunk = chunk[drop:] + skip -= drop + if not chunk: + continue + if ttfa is None: + ttfa = time.perf_counter() - start + pcm.extend(chunk) + except Exception as exc: # noqa: BLE001 - every failure is a benchmark error row + return RequestResult(sentence_id, len(text.split()), None, None, 0.0, 0, error=str(exc)[:200]) + e2e = time.perf_counter() - start + if not pcm: + return RequestResult(sentence_id, len(text.split()), ttfa, e2e, 0.0, 0, error="empty audio") + if save_dir is not None: + _write_wav(save_dir / f"{sentence_id:04d}.wav", bytes(pcm), args.sample_rate) + audio_s = len(pcm) / (BYTES_PER_SAMPLE * args.sample_rate) + return RequestResult(sentence_id, len(text.split()), ttfa, e2e, audio_s, len(pcm)) + + +async def _run_repeat( + args: argparse.Namespace, sentences: list[tuple[int, str]], repeat: int, save_dir: Path | None, +) -> RepeatSummary: + semaphore = asyncio.Semaphore(args.concurrency) + connector = aiohttp.TCPConnector(limit=0) + + async def limited(session, sid, text): + async with semaphore: + return await _one_request(session, args, sid, text, save_dir) + + async with aiohttp.ClientSession(connector=connector) as session: + wall_start = time.perf_counter() + results = await asyncio.gather(*(limited(session, sid, text) for sid, text in sentences)) + wall = time.perf_counter() - wall_start + + ok = [r for r in results if r.error is None] + ttfa = [r.ttfa_s for r in ok if r.ttfa_s is not None] + e2e = [r.e2e_s for r in ok if r.e2e_s is not None] + rtf = [r.rtf for r in ok if r.rtf is not None] + audio_total = sum(r.audio_s for r in ok) + return RepeatSummary( + repeat=repeat, + requests=len(results), + errors=len(results) - len(ok), + wall_s=wall, + ttfa_p50_ms=1000 * _percentile(ttfa, 0.50), + ttfa_p95_ms=1000 * _percentile(ttfa, 0.95), + e2e_p50_s=_percentile(e2e, 0.50), + e2e_p95_s=_percentile(e2e, 0.95), + rtf_mean=statistics.fmean(rtf) if rtf else float("nan"), + rtf_p95=_percentile(rtf, 0.95), + audio_s_total=audio_total, + audio_s_per_wall_s=audio_total / wall if wall > 0 else float("nan"), + results=[asdict(r) | {"rtf": r.rtf} for r in results], + ) + + +def _gpu_info() -> dict[str, str]: + try: + out = subprocess.run( + ["nvidia-smi", "--query-gpu=name,driver_version,clocks.sm,clocks.max.sm,memory.total", + "--format=csv,noheader"], + capture_output=True, text=True, check=True, timeout=10, + ).stdout.strip().splitlines() + except Exception as exc: # noqa: BLE001 - informational only + return {"error": str(exc)} + return {"gpus": out} + + +def _median_summary(repeats: list[RepeatSummary]) -> dict[str, float]: + keys = ("ttfa_p50_ms", "ttfa_p95_ms", "e2e_p50_s", "e2e_p95_s", "rtf_mean", "rtf_p95", + "audio_s_per_wall_s", "wall_s") + return {key: statistics.median(getattr(r, key) for r in repeats) for key in keys} + + +def _markdown_row(args: argparse.Namespace, med: dict[str, float], errors: int) -> str: + return ( + f"| {args.label or args.engine} | c={args.concurrency} | " + f"TTFA p50 {med['ttfa_p50_ms']:.0f} ms / p95 {med['ttfa_p95_ms']:.0f} ms | " + f"RTF {med['rtf_mean']:.3f} | {med['audio_s_per_wall_s']:.1f} audio-s/s | " + f"e2e p50 {med['e2e_p50_s']:.2f} s | errors {errors} |" + ) + + +async def _main_async(args: argparse.Namespace) -> None: + lines = [ln.strip() for ln in Path(args.sentences).read_text(encoding="utf-8").splitlines()] + sentences = [(i + 1, ln) for i, ln in enumerate(lines) if ln] + if args.num_sentences: + sentences = sentences[: args.num_sentences] + if not sentences: + sys.exit("no sentences to synthesize") + + save_dir = Path(args.save_audio_dir) if args.save_audio_dir else None + if save_dir is not None: + save_dir.mkdir(parents=True, exist_ok=True) + + if args.warmup: + warm = [sentences[i % len(sentences)] for i in range(args.warmup)] + print(f"warmup: {len(warm)} requests", file=sys.stderr) + await _run_repeat(args, warm, repeat=-1, save_dir=None) + + repeats: list[RepeatSummary] = [] + for rep in range(args.repeats): + # Only the last repeat keeps audio: identical inputs, and one copy is + # all the WER check needs. + summary = await _run_repeat( + args, sentences, rep, save_dir if rep == args.repeats - 1 else None, + ) + repeats.append(summary) + print( + f"repeat {rep}: {summary.requests} req, {summary.errors} errors, " + f"TTFA p50 {summary.ttfa_p50_ms:.0f} ms p95 {summary.ttfa_p95_ms:.0f} ms, " + f"RTF {summary.rtf_mean:.3f}, {summary.audio_s_per_wall_s:.1f} audio-s/s, " + f"wall {summary.wall_s:.1f} s", + file=sys.stderr, + ) + + med = _median_summary(repeats) + errors = sum(r.errors for r in repeats) + report = { + "engine": args.engine, + "label": args.label, + "url": args.url, + "model": args.model, + "engine_version": args.engine_version, + "sentences_file": str(Path(args.sentences).resolve()), + "num_sentences": len(sentences), + "concurrency": args.concurrency, + "repeats": args.repeats, + "warmup": args.warmup, + "request_fields": {k: v for k, v in _payload(args, "").items() if k != "input"}, + "sample_rate": args.sample_rate, + "gpu": _gpu_info(), + "started_at": args.started_at, + "median_over_repeats": med, + "errors_total": errors, + "repeats_detail": [asdict(r) for r in repeats], + "markdown_row": _markdown_row(args, med, errors), + } + if args.out: + Path(args.out).parent.mkdir(parents=True, exist_ok=True) + Path(args.out).write_text(json.dumps(report, indent=2), encoding="utf-8") + print(report["markdown_row"]) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--engine", choices=("mstar", "vllm-omni", "sglang-omni"), required=True) + parser.add_argument("--url", required=True, help="server base URL, e.g. http://127.0.0.1:8000") + parser.add_argument("--model", required=True, help="model name sent in the request") + parser.add_argument("--sentences", required=True, help="text file, one sentence per line") + parser.add_argument("--num-sentences", type=int, default=0, help="use only the first N sentences") + parser.add_argument("--concurrency", type=int, default=1) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--voice", default=None) + parser.add_argument("--language", default=None) + parser.add_argument("--instructions", default=None) + parser.add_argument("--seed", type=int, default=None) + parser.add_argument("--extra", action="append", default=[], metavar="KEY=JSON", + help="extra request field, e.g. --extra non_streaming_mode=false") + parser.add_argument("--sample-rate", type=int, default=24000) + parser.add_argument("--timeout", type=float, default=300.0) + parser.add_argument("--label", default=None, help="row label for the markdown table") + parser.add_argument("--engine-version", default=None, help="recorded verbatim in the report") + parser.add_argument("--out", default=None, help="JSON report path") + parser.add_argument("--save-audio-dir", default=None, help="write .wav per request (last repeat)") + args = parser.parse_args(argv) + if args.concurrency < 1 or args.repeats < 1: + parser.error("concurrency and repeats must be positive") + args.started_at = time.strftime("%Y-%m-%dT%H:%M:%S%z") + return args + + +def main(argv: list[str] | None = None) -> None: + asyncio.run(_main_async(parse_args(argv))) + + +if __name__ == "__main__": + main() From b4c4fa6b6255c017b0e1fb1a28535c53daa3eb8d Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 019/110] benchmark: WER scorer for synthesized speech --- benchmark/tts_wer.py | 108 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 benchmark/tts_wer.py diff --git a/benchmark/tts_wer.py b/benchmark/tts_wer.py new file mode 100644 index 000000000..5c16e3e43 --- /dev/null +++ b/benchmark/tts_wer.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Word error rate of synthesized speech, transcribed with whisper-large-v3-turbo. + +Quality guard for the TTS benchmark (BENCHMARK_PROTOCOL.md): every engine's WAVs +from ``benchmark/tts_speech_bench.py --save-audio-dir`` are transcribed with the +same ASR model and scored against the same input sentences, after Whisper's +English text normalizer. Engines must land within one WER point of each other. + + python -m benchmark.tts_wer --audio-dir results/mstar_c8_wav \\ + --sentences $BENCH/tts/sentences_200.txt --out results/mstar_c8_wer.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +def word_errors(reference: list[str], hypothesis: list[str]) -> int: + """Levenshtein distance over words (substitutions + insertions + deletions).""" + previous = list(range(len(hypothesis) + 1)) + for i, ref_word in enumerate(reference, start=1): + current = [i] + for j, hyp_word in enumerate(hypothesis, start=1): + current.append(min( + previous[j] + 1, + current[j - 1] + 1, + previous[j - 1] + (ref_word != hyp_word), + )) + previous = current + return previous[-1] + + +def load_sentences(path: str) -> dict[int, str]: + lines = Path(path).read_text(encoding="utf-8").splitlines() + return {i + 1: ln.strip() for i, ln in enumerate(lines) if ln.strip()} + + +def transcribe(audio_paths: list[Path], model_id: str, device: str, batch_size: int) -> list[str]: + import torch + from transformers import pipeline + + asr = pipeline( + "automatic-speech-recognition", + model=model_id, + torch_dtype=torch.float16 if device.startswith("cuda") else torch.float32, + device=device, + ) + outputs = asr( + [str(p) for p in audio_paths], + batch_size=batch_size, + generate_kwargs={"language": "en", "task": "transcribe"}, + return_timestamps=False, + ) + return [o["text"] for o in outputs] + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--audio-dir", required=True, help="directory of .wav files") + parser.add_argument("--sentences", required=True, help="text file, one sentence per line (id = line number)") + parser.add_argument("--asr-model", default="openai/whisper-large-v3-turbo") + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--batch-size", type=int, default=16) + parser.add_argument("--out", default=None) + args = parser.parse_args(argv) + + from transformers.models.whisper.english_normalizer import EnglishTextNormalizer + + sentences = load_sentences(args.sentences) + audio_paths = sorted(Path(args.audio_dir).glob("*.wav")) + if not audio_paths: + sys.exit(f"no WAV files in {args.audio_dir}") + hypotheses = transcribe(audio_paths, args.asr_model, args.device, args.batch_size) + normalize = EnglishTextNormalizer({}) + + rows = [] + total_errors = total_words = 0 + for path, hypothesis in zip(audio_paths, hypotheses, strict=True): + sid = int(path.stem) + reference = sentences[sid] + ref_words = normalize(reference).split() + hyp_words = normalize(hypothesis).split() + errors = word_errors(ref_words, hyp_words) + total_errors += errors + total_words += len(ref_words) + rows.append({"id": sid, "reference": reference, "hypothesis": hypothesis.strip(), + "errors": errors, "words": len(ref_words)}) + wer = 100.0 * total_errors / max(total_words, 1) + report = { + "audio_dir": str(Path(args.audio_dir).resolve()), + "asr_model": args.asr_model, + "files": len(rows), + "wer_percent": wer, + "total_words": total_words, + "total_errors": total_errors, + "rows": rows, + } + if args.out: + Path(args.out).parent.mkdir(parents=True, exist_ok=True) + Path(args.out).write_text(json.dumps(report, indent=2), encoding="utf-8") + print(f"WER {wer:.2f}% over {len(rows)} files ({total_errors}/{total_words} words)") + + +if __name__ == "__main__": + main() From ce08c3d946b30e8cd6d38343f7351d27b3a22a47 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 020/110] test: Qwen3-TTS parity harness against the qwen-tts reference --- test/qwen3-tts/parity_qwen3_tts.py | 453 +++++++++++++++++++++++++++++ 1 file changed, 453 insertions(+) create mode 100644 test/qwen3-tts/parity_qwen3_tts.py diff --git a/test/qwen3-tts/parity_qwen3_tts.py b/test/qwen3-tts/parity_qwen3_tts.py new file mode 100644 index 000000000..3c6a2c49c --- /dev/null +++ b/test/qwen3-tts/parity_qwen3_tts.py @@ -0,0 +1,453 @@ +#!/usr/bin/env python3 +"""Parity of the M* Qwen3-TTS port against the ``qwen-tts`` reference (GPU). + +Three checks, all on real weights, all against the reference package that the +checkpoint ships with (``qwen_tts`` 0.1.1, bf16 Talker, fp32 codec): + +1. **Teacher-forced Talker.** The reference generates ``--frames`` greedy codec + frames and returns the Talker hidden state that produced each one. M* is + driven through its served path (process_prompt -> prefill -> paged + FlashInfer decode, one step per frame) while being fed the reference's + frames, and its group-0 logits at every frame are compared with the + reference's. The two sides build their prefill independently, so the first + frame validates the prompt construction as well as the backbone. A second + reference pass over M*'s own prefill (``backbone_only``) isolates the + backbone from the prompt if the first comparison ever fails. +2. **Teacher-forced CodePredictor.** For every frame, M*'s depth loop receives + the reference Talker hidden state and the reference codes and its 15 group + logits are compared with the reference ``forward_finetune`` logits. +3. **Greedy end to end.** M* generates the same number of frames on its own + (temperature 0 on both samplers, the checkpoint's repetition penalty) and + the codes are compared frame by frame; both code sequences are decoded to + audio by the codec on each side and the waveform max-abs-diff is reported. + +Run inside the GPU allocation (weights must already be in the HF cache):: + + python test/qwen3-tts/parity_qwen3_tts.py --repo Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \\ + --frames 64 --voice vivian --language English --json results/parity_1p7b.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path +from typing import Any + +import torch + +from mstar.communication.tensors import LocalTransferEngine +from mstar.conductor.request_info import CurrentForwardPassInfo +from mstar.distributed.communication import CommGroup, JointGroups +from mstar.engine.resources import StepContext, StepRunner, resolve_spec_dependencies +from mstar.engine.resources.base import EngineResourceInfo, build_resource +from mstar.engine.resources.kv.transfer import TransferEngineInfo +from mstar.model.qwen3_tts.qwen3_tts_model import Qwen3TTSModel +from mstar.model.submodule_base import ModelInputsFromEngine + +GREEDY_KWARGS = {"do_sample": False, "subtalker_dosample": False} + + +# --------------------------------------------------------------------------- +# Checkpoint + reference +# --------------------------------------------------------------------------- + + +def resolve_snapshot(repo: str) -> str: + if Path(repo).is_dir(): + return repo + from huggingface_hub import snapshot_download + + return snapshot_download(repo, local_files_only=True) + + +def load_reference(snapshot: str, device: str): + """The reference stack: ``Qwen3TTSForConditionalGeneration`` + processor + codec.""" + from qwen_tts import Qwen3TTSModel as ReferenceModel + + return ReferenceModel.from_pretrained( + snapshot, device_map=device, dtype=torch.bfloat16, attn_implementation="sdpa", + ) + + +def reference_generate(ref, args) -> tuple[torch.Tensor, list[torch.Tensor]]: + """Greedy reference codes ``[frames, groups]`` and the embeds the Talker saw. + + Uses the low-level ``generate`` of the reference so the prompt layout + (speaker, language, instruct, streaming vs non-streaming text) is exactly + the reference's own, independent of M*'s ``process_prompt``. + """ + model = ref.model + device = model.device + input_ids = ref._tokenize_texts([ref._build_assistant_text(args.text)]) + instruct_ids = None + if args.instruct: + instruct_ids = ref._tokenize_texts([ref._build_instruct_text(args.instruct)]) + speakers = [args.voice] if args.voice else None + non_streaming = model.tts_model_type in ("custom_voice", "voice_design") + if args.non_streaming_mode is not None: + non_streaming = args.non_streaming_mode + codes_list, hidden_list = model.generate( + input_ids=input_ids, + instruct_ids=instruct_ids, + languages=[args.language or "auto"], + speakers=speakers, + non_streaming_mode=non_streaming, + max_new_tokens=args.frames, + do_sample=False, + subtalker_dosample=False, + repetition_penalty=args.repetition_penalty, + ) + codes = codes_list[0].to(device) + if codes.shape[0] < args.frames: + print(f"reference stopped at EOS after {codes.shape[0]} frames", file=sys.stderr) + return codes, hidden_list[0] + + +def reference_frame_embeds(ref, codes: torch.Tensor) -> torch.Tensor: + """Sum of the 16 codec embeddings per frame, in the Talker width.""" + talker = ref.model.talker + embeds = talker.get_input_embeddings()(codes[:, 0]) + residual_tables = talker.code_predictor.get_input_embeddings() + for group in range(1, codes.shape[1]): + embeds = embeds + residual_tables[group - 1](codes[:, group]) + return embeds + + +def reference_teacher_forced(ref, prefill: torch.Tensor, frame_embeds: torch.Tensor, trailing: torch.Tensor, + tts_pad: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Reference group-0 logits and normed hidden per frame over one full sequence.""" + talker = ref.model.talker + steps = frame_embeds.shape[0] + text_cond = torch.stack([ + trailing[t] if t < trailing.shape[0] else tts_pad for t in range(steps) + ]) + # Frame t's logits come from the position holding frame t-1's input; the + # prefill's last position predicts frame 0, so the last frame input is + # not needed as an input at all. + inputs = torch.cat([prefill, (frame_embeds + text_cond)[:-1]], dim=0).unsqueeze(0) + out = talker.model(inputs_embeds=inputs, use_cache=False) + hidden = out.last_hidden_state[0, prefill.shape[0] - 1:] + return talker.codec_head(hidden), hidden + + +def reference_code_predictor_logits(ref, hidden: torch.Tensor, codes: torch.Tensor) -> torch.Tensor: + """Reference residual-group logits ``[frames, groups-1, vocab]`` (teacher forced).""" + talker = ref.model.talker + cp = talker.code_predictor + first = talker.get_input_embeddings()(codes[:, :1]) + residual = [cp.get_input_embeddings()[g - 1](codes[:, g:g + 1]) for g in range(1, codes.shape[1] - 1)] + inputs = torch.cat([hidden.unsqueeze(1), first, *residual], dim=1) + return cp.forward_finetune(inputs_embeds=inputs).logits + + +# --------------------------------------------------------------------------- +# M* side: the served path without the worker around it +# --------------------------------------------------------------------------- + + +class MStarTalkerDriver: + """Drive ``TalkerSubmodule`` through declare -> admit -> plan -> forward -> commit. + + Mirrors ``Engine._drive_step`` for one eager request, on resources built + from the model's own ``get_node_resources`` declaration. + """ + + def __init__(self, model: Qwen3TTSModel, talker, device: str, max_num_pages: int = 64): + self.model = model + self.talker = talker + self.device = torch.device(device) + specs = model.get_node_resources() + by_key = resolve_spec_dependencies(specs) + for spec in specs: + if hasattr(spec, "apply_yaml_overrides") and hasattr(spec.config, "max_num_pages"): + spec.apply_yaml_overrides(max_num_pages=max_num_pages) + groups = JointGroups(tp_group=CommGroup.trivial(), sp_group=CommGroup.trivial()) + transfer = TransferEngineInfo("h", "h", LocalTransferEngine("h")) + self.resources = { + spec.resource_key: build_resource( + spec, + EngineResourceInfo( + device=self.device, + joint_comm_group=groups, + transfer_engine_info=transfer, + kv_dtype=torch.bfloat16, + dependencies={key: by_key[key] for key in spec.depends_on()}, + ), + ) + for spec in specs + } + talker.bind_node_resources(self.resources) + self.runner = StepRunner(self.resources) + self.fwd_index = 0 + + def open_request(self, rid: str, model_kwargs: dict[str, Any]) -> CurrentForwardPassInfo: + configs = self.model.get_request_resource_configs({}, model_kwargs) + for config in configs.values(): + config.apply_conductor_config(seed=1234) + self.runner.ingest_request(rid, configs) + return CurrentForwardPassInfo( + request_id=rid, graph_walk="talker_prefill", fwd_index=0, random_seed=1234, + max_tokens=8192, resource_configs=configs, + step_metadata={"talker_max_tokens": 8192, "is_prefill": True}, + ) + + def close_request(self, rid: str) -> None: + self.runner.remove_request(rid) + self.talker.cleanup_request(rid) + + def step(self, walk: str, fwd: CurrentForwardPassInfo, inputs: dict, forward): + """One step; ``forward(engine_inputs, **preprocessed)`` runs the compute.""" + rid = fwd.request_id + fwd.graph_walk = walk + prepared = self.talker.prepare_inputs(walk, fwd, inputs) + step = self.talker.declare_step(walk, [rid], [prepared]) + step.set_ctx(StepContext(request_ids=(rid,), graph_walk=walk, slot=0, capture=False)) + outcome = self.runner.admit(step) + assert outcome.ok, f"admit failed: {outcome.reason}" + self.runner.plan(step) + engine_inputs = ModelInputsFromEngine( + request_ids=[rid], per_request_info={rid: fwd}, resources=dict(self.resources), + per_request_states={rid: self.talker.request_state(rid)}, step=step, + ) + preprocessed = self.talker.preprocess(walk, engine_inputs, [prepared]) + out = forward(engine_inputs, **preprocessed) + self.runner.commit(step) + return out + + +@torch.no_grad() +def mstar_teacher_forced(driver: MStarTalkerDriver, tensors: dict, frame_embeds: torch.Tensor, + greedy_kwargs: dict) -> tuple[torch.Tensor, torch.Tensor]: + """M* group-0 logits and normed hidden per frame, fed the reference frames.""" + talker = driver.talker + rid = "teacher-forced" + fwd = driver.open_request(rid, greedy_kwargs) + logits, hiddens = [], [] + + def backbone(engine_inputs, input_embeds, last_token_indices, suppress_eos): + del engine_inputs, suppress_eos + hidden = talker.model(input_embeds, label="main") + last = hidden.index_select(0, last_token_indices) + hiddens.append(last[0]) + logits.append(talker.model.codec_head(last)[0]) + + driver.step("talker_prefill", fwd, tensors, backbone) + for t in range(frame_embeds.shape[0] - 1): + # prepare_inputs adds the text condition for this step itself. + driver.step("talker_decode", fwd, {"talker_input_embeds": [frame_embeds[t:t + 1]]}, backbone) + driver.close_request(rid) + return torch.stack(logits), torch.stack(hiddens) + + +@torch.no_grad() +def mstar_code_predictor_logits(talker, hidden: torch.Tensor, codes: torch.Tensor) -> torch.Tensor: + """M* residual-group logits per frame, fed the reference hidden and codes.""" + out = [] + for t in range(codes.shape[0]): + collected = [] + target = iter(codes[t, 1:].tolist()) + + def teacher(cp_logits, collected=collected, target=target): + collected.append(cp_logits[0]) + return cp_logits.new_tensor([next(target)], dtype=torch.long) + + talker._depth_loop(hidden[t:t + 1], codes[t:t + 1, 0], teacher) + out.append(torch.stack(collected)) + return torch.stack(out) + + +@torch.no_grad() +def mstar_greedy(driver: MStarTalkerDriver, tensors: dict, frames: int, greedy_kwargs: dict) -> torch.Tensor: + """M*'s own greedy generation through ``TalkerSubmodule.forward``.""" + talker = driver.talker + rid = "greedy" + fwd = driver.open_request(rid, greedy_kwargs) + codes = [] + + def forward(engine_inputs, **kw): + return talker.forward(fwd.graph_walk, engine_inputs, **kw) + + out = driver.step("talker_prefill", fwd, tensors, forward) + codes.append(out["codec_tokens"][0][0]) + talker.postprocess(rid, fwd, out) + eos = talker.talker_config.codec_eos_token_id + while len(codes) < frames and int(codes[-1][0]) != eos: + out = driver.step("talker_decode", fwd, {"talker_input_embeds": out["talker_input_embeds"]}, forward) + codes.append(out["codec_tokens"][0][0]) + talker.postprocess(rid, fwd, out) + driver.close_request(rid) + codes = torch.stack(codes) + return codes[codes[:, 0] != eos] + + +# --------------------------------------------------------------------------- +# Comparison +# --------------------------------------------------------------------------- + + +def compare_logits(name: str, ours: torch.Tensor, theirs: torch.Tensor) -> dict[str, Any]: + ours = ours.float() + theirs = theirs.float() + diff = (ours - theirs).abs() + agree = (ours.argmax(-1) == theirs.argmax(-1)).float() + top2 = theirs.topk(2, dim=-1).values + margin = (top2[..., 0] - top2[..., 1]) + return { + "name": name, + "positions": int(agree.numel()), + "argmax_agreement": float(agree.mean()), + "max_abs_diff": float(diff.max()), + "mean_abs_diff": float(diff.mean()), + "ref_logit_scale": float(theirs.abs().mean()), + # disagreements should sit on near-ties: report the reference top-2 + # margin where the argmax differs + "disagreement_margin_median": ( + float(margin[agree == 0].median()) if (agree == 0).any() else None + ), + } + + +def compare_codes(ours: torch.Tensor, theirs: torch.Tensor) -> dict[str, Any]: + n = min(ours.shape[0], theirs.shape[0]) + equal_frames = (ours[:n] == theirs[:n]).all(dim=1) + first_div = int(equal_frames.logical_not().nonzero()[0]) if not equal_frames.all() else n + return { + "frames_mstar": int(ours.shape[0]), + "frames_reference": int(theirs.shape[0]), + "frames_compared": n, + "identical_frames_before_divergence": first_div, + "group0_agreement": float((ours[:n, 0] == theirs[:n, 0]).float().mean()), + "all_groups_agreement": float(equal_frames.float().mean()), + } + + +@torch.no_grad() +def decode_audio(codec, codes: torch.Tensor) -> torch.Tensor: + """M* codec: ``[frames, groups]`` -> float waveform in [-1, 1].""" + wav = codec.decoder(codes.t().unsqueeze(0).contiguous()) + return wav.squeeze().float() + + +@torch.no_grad() +def reference_decode_audio(ref, codes: torch.Tensor) -> torch.Tensor: + wavs, _ = ref.model.speech_tokenizer.decode([{"audio_codes": codes}]) + return torch.as_tensor(wavs[0]).float() + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--repo", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice") + parser.add_argument( + "--text", + default="The quick brown fox jumps over the lazy dog while the sun sets behind the hills.", + ) + parser.add_argument("--voice", default="vivian") + parser.add_argument("--language", default="English") + parser.add_argument("--instruct", default=None) + parser.add_argument("--non-streaming-mode", type=lambda s: s.lower() == "true", default=None) + parser.add_argument("--frames", type=int, default=64) + parser.add_argument("--repetition-penalty", type=float, default=1.05) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--json", default=None) + args = parser.parse_args(argv) + if args.voice == "": + args.voice = None + + torch.manual_seed(0) + snapshot = resolve_snapshot(args.repo) + t0 = time.perf_counter() + ref = load_reference(snapshot, args.device) + print(f"reference loaded in {time.perf_counter() - t0:.1f}s", file=sys.stderr) + + model = Qwen3TTSModel(model_path_hf=snapshot) + talker = model.get_submodule("Talker", device=args.device, autocast_dtype=torch.bfloat16) + codec = model.get_submodule("Codec", device=args.device) + driver = MStarTalkerDriver(model, talker, args.device) + print(f"M* loaded in {time.perf_counter() - t0:.1f}s", file=sys.stderr) + + request_kwargs = {"language": args.language, **GREEDY_KWARGS, + "repetition_penalty": args.repetition_penalty} + if args.voice: + request_kwargs["voice"] = args.voice + if args.instruct: + request_kwargs["instruct"] = args.instruct + if args.non_streaming_mode is not None: + request_kwargs["non_streaming_mode"] = args.non_streaming_mode + tensors = model.process_prompt(args.text, ["text"], ["audio"], **request_kwargs) + + # 1 + 2: teacher forced against the reference's greedy frames. ``ref_hidden`` + # is the hidden state the reference's own generation used for each frame. + ref_codes, ref_hidden = reference_generate(ref, args) + ref_hidden = ref_hidden.to(torch.bfloat16) + frames = ref_codes.shape[0] + theirs_logits = ref.model.talker.codec_head(ref_hidden) + frame_embeds = reference_frame_embeds(ref, ref_codes) + ours_logits, ours_hidden = mstar_teacher_forced(driver, tensors, frame_embeds, request_kwargs) + talker_report = compare_logits("talker_group0_logits", ours_logits, theirs_logits) + hidden_report = { + "name": "talker_hidden", + "max_abs_diff": float((ours_hidden.float() - ref_hidden.float()).abs().max()), + "rel_diff": float((ours_hidden.float() - ref_hidden.float()).norm() / ref_hidden.float().norm()), + } + # Diagnostic: the reference backbone over M*'s own prefill embeddings. + prefill = talker._build_prefill( + "layout-probe", tensors["text_inputs"][0], tensors["prompt_layout"][0], + int(tensors["speaker_id"][0]), int(tensors["language_id"][0]), + ) + probe_state = talker.request_state("layout-probe") + backbone_logits, _ = reference_teacher_forced( + ref, prefill, frame_embeds, probe_state["trailing_text_hidden"], probe_state["tts_pad_embed"], + ) + talker.cleanup_request("layout-probe") + backbone_report = compare_logits("backbone_only_on_mstar_prefill", ours_logits, backbone_logits) + cp_ours = mstar_code_predictor_logits(talker, ref_hidden, ref_codes) + cp_theirs = reference_code_predictor_logits(ref, ref_hidden, ref_codes) + cp_report = compare_logits("code_predictor_logits", cp_ours, cp_theirs) + + # 3: greedy end to end + audio. + ours_codes = mstar_greedy(driver, tensors, frames, request_kwargs) + codes_report = compare_codes(ours_codes, ref_codes) + n = codes_report["frames_compared"] + audio_ref_codes_mstar = decode_audio(codec, ref_codes[:n]) + audio_ref_codes_ref = reference_decode_audio(ref, ref_codes[:n]).to(audio_ref_codes_mstar.device) + m = min(audio_ref_codes_mstar.numel(), audio_ref_codes_ref.numel()) + codec_report = { + "name": "codec_same_codes", + "samples": m, + "max_abs_diff": float((audio_ref_codes_mstar[:m] - audio_ref_codes_ref[:m]).abs().max()), + "length_mismatch": int(audio_ref_codes_mstar.numel() - audio_ref_codes_ref.numel()), + } + audio_ours = decode_audio(codec, ours_codes[:n]) + k = min(audio_ours.numel(), audio_ref_codes_ref.numel()) + e2e_audio = { + "name": "audio_greedy_e2e", + "samples": k, + "max_abs_diff": float((audio_ours[:k] - audio_ref_codes_ref[:k]).abs().max()), + "snr_db": float(10 * torch.log10( + audio_ref_codes_ref[:k].pow(2).mean() / ((audio_ours[:k] - audio_ref_codes_ref[:k]).pow(2).mean() + 1e-12) + )), + } + + report = { + "repo": args.repo, "snapshot": snapshot, "text": args.text, "voice": args.voice, + "language": args.language, "instruct": args.instruct, "frames": int(ref_codes.shape[0]), + "talker": talker_report, "talker_hidden": hidden_report, "backbone_only": backbone_report, + "code_predictor": cp_report, + "greedy_codes": codes_report, "codec": codec_report, "audio": e2e_audio, + "torch": torch.__version__, "gpu": torch.cuda.get_device_name(0), + } + print(json.dumps(report, indent=2)) + if args.json: + Path(args.json).parent.mkdir(parents=True, exist_ok=True) + Path(args.json).write_text(json.dumps(report, indent=2), encoding="utf-8") + ok = talker_report["argmax_agreement"] >= 0.99 and cp_report["argmax_agreement"] >= 0.99 \ + and codec_report["max_abs_diff"] < 1e-3 + print("PARITY OK" if ok else "PARITY FAILED", file=sys.stderr) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() From 02799bff0a7b501c51f9ae917fb74151684f4d59 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 021/110] qwen3_tts: ECAPA-TDNN speaker encoder and mel front end components --- .../qwen3_tts/components/speaker_encoder.py | 263 ++++++++++++++++++ 1 file changed, 263 insertions(+) create mode 100644 mstar/model/qwen3_tts/components/speaker_encoder.py diff --git a/mstar/model/qwen3_tts/components/speaker_encoder.py b/mstar/model/qwen3_tts/components/speaker_encoder.py new file mode 100644 index 000000000..f936f0adc --- /dev/null +++ b/mstar/model/qwen3_tts/components/speaker_encoder.py @@ -0,0 +1,263 @@ +"""ECAPA-TDNN speaker encoder and its mel front end (Qwen3-TTS Base). + +The Base checkpoint clones a voice from reference audio: a 128-bin log-mel +spectrogram of the 24 kHz reference goes through an ECAPA-TDNN encoder whose +2048-wide output (the "x-vector") takes the place of the built-in speaker tag +in the Talker prefill. Module names follow the ``speaker_encoder.*`` keys of +the checkpoint so weights stream in without remapping. + +Both pieces are plain, batch-friendly PyTorch: the encoder is a stack of 1-D +convolutions with "same" reflect padding, and the front end is an STFT plus a +Slaney-normalized mel filterbank (the filterbank ``librosa.filters.mel`` +produces, computed here without the librosa dependency). +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from torch import nn + +from mstar.model.qwen3_tts.config import Qwen3TTSSpeakerEncoderConfig + +# --------------------------------------------------------------------------- +# Mel front end +# --------------------------------------------------------------------------- + + +def _hz_to_mel_slaney(freq: torch.Tensor) -> torch.Tensor: + """Slaney mel scale: linear below 1 kHz, logarithmic above (librosa default).""" + f_sp = 200.0 / 3 + min_log_hz = 1000.0 + min_log_mel = min_log_hz / f_sp + logstep = math.log(6.4) / 27.0 + mel = freq / f_sp + log_region = freq >= min_log_hz + mel = torch.where( + log_region, + min_log_mel + torch.log(freq.clamp(min=min_log_hz) / min_log_hz) / logstep, + mel, + ) + return mel + + +def _mel_to_hz_slaney(mel: torch.Tensor) -> torch.Tensor: + f_sp = 200.0 / 3 + min_log_hz = 1000.0 + min_log_mel = min_log_hz / f_sp + logstep = math.log(6.4) / 27.0 + freq = f_sp * mel + log_region = mel >= min_log_mel + return torch.where( + log_region, min_log_hz * torch.exp(logstep * (mel - min_log_mel)), freq + ) + + +def slaney_mel_filterbank( + sample_rate: int, n_fft: int, n_mels: int, fmin: float, fmax: float +) -> torch.Tensor: + """``librosa.filters.mel(sr, n_fft, n_mels, fmin, fmax)``: ``[n_mels, n_fft // 2 + 1]``.""" + fft_freqs = torch.linspace(0.0, sample_rate / 2, n_fft // 2 + 1, dtype=torch.float64) + mel_edges = torch.linspace( + _hz_to_mel_slaney(torch.tensor(float(fmin), dtype=torch.float64)).item(), + _hz_to_mel_slaney(torch.tensor(float(fmax), dtype=torch.float64)).item(), + n_mels + 2, + dtype=torch.float64, + ) + hz_edges = _mel_to_hz_slaney(mel_edges) + fdiff = hz_edges[1:] - hz_edges[:-1] + ramps = hz_edges.unsqueeze(1) - fft_freqs.unsqueeze(0) # [n_mels + 2, bins] + lower = -ramps[:-2] / fdiff[:-1].unsqueeze(1) + upper = ramps[2:] / fdiff[1:].unsqueeze(1) + weights = torch.clamp(torch.minimum(lower, upper), min=0.0) + # Slaney normalization: each filter integrates to about the same area. + enorm = 2.0 / (hz_edges[2:] - hz_edges[:-2]) + return (weights * enorm.unsqueeze(1)).to(torch.float32) + + +class Qwen3TTSMelFrontEnd(nn.Module): + """Log-mel features of 24 kHz reference audio, as ``extract_speaker_embedding`` computes them. + + ``waveform`` is ``[batch, samples]`` in ``[-1, 1]``; the output is + ``[batch, frames, mel_dim]`` in float32, ready for the encoder. Windows are + Hann, frames are not centered but the signal is reflect-padded by + ``(n_fft - hop) // 2`` on both sides, magnitudes get a ``1e-9`` floor and + the mel energies are log-compressed with a ``1e-5`` clamp. + """ + + def __init__(self, config: Qwen3TTSSpeakerEncoderConfig) -> None: + super().__init__() + self.n_fft = config.n_fft + self.hop_size = config.hop_size + self.win_size = config.win_size + self.register_buffer( + "mel_basis", + slaney_mel_filterbank(config.sample_rate, config.n_fft, config.mel_dim, config.fmin, config.fmax), + persistent=False, + ) + self.register_buffer("window", torch.hann_window(config.win_size), persistent=False) + + def forward(self, waveform: torch.Tensor) -> torch.Tensor: + waveform = waveform.to(torch.float32) + if waveform.ndim == 1: + waveform = waveform.unsqueeze(0) + padding = (self.n_fft - self.hop_size) // 2 + waveform = F.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1) + spec = torch.stft( + waveform, + self.n_fft, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.window, + center=False, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + magnitude = torch.sqrt(spec.real.pow(2) + spec.imag.pow(2) + 1e-9) + mel = torch.matmul(self.mel_basis, magnitude) + return torch.log(torch.clamp(mel, min=1e-5)).transpose(1, 2) + + +# --------------------------------------------------------------------------- +# ECAPA-TDNN +# --------------------------------------------------------------------------- + + +class _SameReflectConv1d(nn.Module): + """``nn.Conv1d(padding="same", padding_mode="reflect")`` with the checkpoint's ``conv`` name.""" + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int, dilation: int = 1) -> None: + super().__init__() + self.conv = nn.Conv1d( + in_channels, out_channels, kernel_size, dilation=dilation, + padding="same", padding_mode="reflect", + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.conv(hidden_states) + + +class TimeDelayNetBlock(_SameReflectConv1d): + """Conv1d + ReLU.""" + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return F.relu(self.conv(hidden_states)) + + +class Res2NetBlock(nn.Module): + """Hierarchical residual convolutions over ``scale`` channel groups.""" + + def __init__(self, in_channels: int, out_channels: int, scale: int, kernel_size: int, dilation: int) -> None: + super().__init__() + self.scale = scale + self.blocks = nn.ModuleList([ + TimeDelayNetBlock(in_channels // scale, out_channels // scale, kernel_size, dilation) + for _ in range(scale - 1) + ]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + outputs = [] + previous = None + for index, part in enumerate(hidden_states.chunk(self.scale, dim=1)): + if index == 0: + previous = part + elif index == 1: + previous = self.blocks[0](part) + else: + previous = self.blocks[index - 1](part + previous) + outputs.append(previous) + return torch.cat(outputs, dim=1) + + +class SqueezeExcitationBlock(nn.Module): + def __init__(self, in_channels: int, se_channels: int, out_channels: int) -> None: + super().__init__() + self.conv1 = nn.Conv1d(in_channels, se_channels, 1) + self.conv2 = nn.Conv1d(se_channels, out_channels, 1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + pooled = hidden_states.mean(dim=2, keepdim=True) + gate = torch.sigmoid(self.conv2(F.relu(self.conv1(pooled)))) + return hidden_states * gate + + +class SqueezeExcitationRes2NetBlock(nn.Module): + """TDNN -> Res2Net -> TDNN -> squeeze-excitation, with a residual connection.""" + + def __init__(self, in_channels: int, out_channels: int, res2net_scale: int, se_channels: int, + kernel_size: int, dilation: int) -> None: + super().__init__() + self.tdnn1 = TimeDelayNetBlock(in_channels, out_channels, 1, 1) + self.res2net_block = Res2NetBlock(out_channels, out_channels, res2net_scale, kernel_size, dilation) + self.tdnn2 = TimeDelayNetBlock(out_channels, out_channels, 1, 1) + self.se_block = SqueezeExcitationBlock(out_channels, se_channels, out_channels) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual = hidden_states + hidden_states = self.tdnn1(hidden_states) + hidden_states = self.res2net_block(hidden_states) + hidden_states = self.tdnn2(hidden_states) + return self.se_block(hidden_states) + residual + + +class AttentiveStatisticsPooling(nn.Module): + """Attention-weighted mean and standard deviation over time: ``[B, C, T] -> [B, 2C, 1]``.""" + + def __init__(self, channels: int, attention_channels: int) -> None: + super().__init__() + self.eps = 1e-12 + self.tdnn = TimeDelayNetBlock(channels * 3, attention_channels, 1, 1) + self.conv = nn.Conv1d(attention_channels, channels, 1) + + def _statistics(self, hidden_states: torch.Tensor, weights: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + mean = (weights * hidden_states).sum(dim=2) + variance = (weights * (hidden_states - mean.unsqueeze(2)).pow(2)).sum(dim=2) + return mean, torch.sqrt(variance.clamp(min=self.eps)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + frames = hidden_states.shape[-1] + uniform = torch.full_like(hidden_states[:, :1, :], 1.0 / frames) + mean, std = self._statistics(hidden_states, uniform) + context = torch.cat([ + hidden_states, + mean.unsqueeze(2).expand(-1, -1, frames), + std.unsqueeze(2).expand(-1, -1, frames), + ], dim=1) + attention = torch.softmax(self.conv(torch.tanh(self.tdnn(context))), dim=2) + mean, std = self._statistics(hidden_states, attention) + return torch.cat([mean, std], dim=1).unsqueeze(2) + + +class Qwen3TTSSpeakerEncoder(nn.Module): + """ECAPA-TDNN x-vector extractor: ``[batch, frames, mel_dim] -> [batch, enc_dim]``.""" + + def __init__(self, config: Qwen3TTSSpeakerEncoderConfig) -> None: + super().__init__() + channels = config.enc_channels + kernels = config.enc_kernel_sizes + dilations = config.enc_dilations + if not len(channels) == len(kernels) == len(dilations): + raise ValueError("enc_channels, enc_kernel_sizes and enc_dilations must have the same length") + self.blocks = nn.ModuleList([TimeDelayNetBlock(config.mel_dim, channels[0], kernels[0], dilations[0])]) + for index in range(1, len(channels) - 1): + self.blocks.append(SqueezeExcitationRes2NetBlock( + channels[index - 1], channels[index], config.enc_res2net_scale, config.enc_se_channels, + kernels[index], dilations[index], + )) + self.mfa = TimeDelayNetBlock(channels[-1], channels[-1], kernels[-1], dilations[-1]) + self.asp = AttentiveStatisticsPooling(channels[-1], config.enc_attention_channels) + self.fc = nn.Conv1d(channels[-1] * 2, config.enc_dim, 1) + + def forward(self, mels: torch.Tensor) -> torch.Tensor: + hidden_states = mels.transpose(1, 2) + features = [] + for block in self.blocks: + hidden_states = block(hidden_states) + features.append(hidden_states) + # Multi-layer feature aggregation over the SE-Res2Net outputs only. + hidden_states = self.mfa(torch.cat(features[1:], dim=1)) + return self.fc(self.asp(hidden_states)).squeeze(-1) From 12e5f87278aa3ecdc5ac8799bd81d3518de938be Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 022/110] test: Qwen3-TTS speaker encoder parity against qwen-tts --- .../modular/test_qwen3_tts_speaker_encoder.py | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) create mode 100644 test/modular/test_qwen3_tts_speaker_encoder.py diff --git a/test/modular/test_qwen3_tts_speaker_encoder.py b/test/modular/test_qwen3_tts_speaker_encoder.py new file mode 100644 index 000000000..4036ac2c2 --- /dev/null +++ b/test/modular/test_qwen3_tts_speaker_encoder.py @@ -0,0 +1,82 @@ +"""CPU parity of the M* ECAPA-TDNN speaker encoder and mel front end (Qwen3-TTS Base). + +Oracle: the ``qwen_tts`` reference package. Both encoders start from the same +random weights (the M* module loads the reference state dict verbatim, which +also pins the checkpoint key layout), so any drift is an implementation +difference, not initialization noise. +""" + +from __future__ import annotations + +import importlib.util +import warnings + +import pytest +import torch + +from mstar.model.qwen3_tts.components.speaker_encoder import ( + Qwen3TTSMelFrontEnd, + Qwen3TTSSpeakerEncoder, + slaney_mel_filterbank, +) +from mstar.model.qwen3_tts.config import Qwen3TTSSpeakerEncoderConfig + +pytestmark = pytest.mark.skipif( + importlib.util.find_spec("qwen_tts") is None, reason="qwen-tts reference not installed" +) + + +def _reference_modules(): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # pysox probes for a SoX binary on import + from qwen_tts.core.models.configuration_qwen3_tts import ( + Qwen3TTSSpeakerEncoderConfig as RefConfig, + ) + from qwen_tts.core.models.modeling_qwen3_tts import ( + Qwen3TTSSpeakerEncoder as RefEncoder, + ) + from qwen_tts.core.models.modeling_qwen3_tts import mel_spectrogram + return RefConfig, RefEncoder, mel_spectrogram + + +def test_speaker_encoder_matches_reference_on_shared_weights(): + RefConfig, RefEncoder, _ = _reference_modules() + torch.manual_seed(0) + config = Qwen3TTSSpeakerEncoderConfig(enc_dim=64, enc_channels=(32, 32, 32, 32, 96), enc_se_channels=16, + enc_attention_channels=16) + reference = RefEncoder(RefConfig( + enc_dim=64, enc_channels=[32, 32, 32, 32, 96], enc_se_channels=16, enc_attention_channels=16, + )).eval() + ours = Qwen3TTSSpeakerEncoder(config).eval() + # Identical parameter names: the checkpoint's ``speaker_encoder.*`` keys + # load without a remap. + assert set(dict(ours.named_parameters())) == set(dict(reference.named_parameters())) + ours.load_state_dict(reference.state_dict()) + + mels = torch.randn(3, 97, config.mel_dim) + with torch.no_grad(): + expected = reference(mels) + actual = ours(mels) + assert actual.shape == (3, 64) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +def test_mel_front_end_matches_reference(): + _, _, mel_spectrogram = _reference_modules() + config = Qwen3TTSSpeakerEncoderConfig(enc_dim=2048) + torch.manual_seed(1) + wave = (torch.randn(2, 24000) * 0.3).clamp(-1, 1) + expected = mel_spectrogram( + wave, n_fft=config.n_fft, num_mels=config.mel_dim, sampling_rate=config.sample_rate, + hop_size=config.hop_size, win_size=config.win_size, fmin=config.fmin, fmax=config.fmax, + ).transpose(1, 2) + actual = Qwen3TTSMelFrontEnd(config)(wave) + assert actual.shape == expected.shape == (2, 24000 // config.hop_size, config.mel_dim) + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4) + + +def test_slaney_filterbank_matches_librosa(): + librosa = pytest.importorskip("librosa") + expected = torch.from_numpy(librosa.filters.mel(sr=24000, n_fft=1024, n_mels=128, fmin=0, fmax=12000)) + actual = slaney_mel_filterbank(24000, 1024, 128, 0, 12000) + torch.testing.assert_close(actual, expected.to(torch.float32), rtol=1e-6, atol=1e-7) From 124fb138b74f1993810d1d65801a27244c255f98 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 023/110] qwen3_tts: codec encoder settings and reference frame arithmetic --- mstar/model/qwen3_tts/config.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/mstar/model/qwen3_tts/config.py b/mstar/model/qwen3_tts/config.py index f45072a12..ad260f443 100644 --- a/mstar/model/qwen3_tts/config.py +++ b/mstar/model/qwen3_tts/config.py @@ -239,6 +239,12 @@ class Qwen3TTSCodecConfig: input_sample_rate: int = 24000 output_sample_rate: int = 24000 decode_upsample_rate: int = 1920 + encode_downsample_rate: int = 1920 + encoder_valid_num_quantizers: int = 16 + # Raw Mimi encoder configuration (``encoder_config`` in the speech + # tokenizer's config.json); it is handed verbatim to the encoder that + # turns reference audio into codec frames for voice cloning. + encoder_config: dict[str, Any] = field(default_factory=dict) # M* stream policy: 300 new 12 Hz frames with 25 frames of overlap. chunk_frames: int = 300 @@ -258,17 +264,27 @@ def from_dict(cls, data: dict[str, Any]) -> "Qwen3TTSCodecConfig": "input_sample_rate", "output_sample_rate", "decode_upsample_rate", + "encode_downsample_rate", + "encoder_valid_num_quantizers", + "encoder_config", ) if name in data }) return cls(**values) + def frames_for_samples(self, num_samples: int) -> int: + """Codec frames the encoder emits for ``num_samples`` of input audio.""" + return -(-int(num_samples) // self.encode_downsample_rate) + def decoder_kwargs(self) -> dict[str, Any]: """Arguments accepted by the official 12 Hz decoder config.""" excluded = { "input_sample_rate", "output_sample_rate", "decode_upsample_rate", + "encode_downsample_rate", + "encoder_valid_num_quantizers", + "encoder_config", "chunk_frames", "left_context_frames", } From b708a577ff46367f67579054a3946e8a6bb9b7cc Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 024/110] qwen3_tts: reference encoder node, in-context clone prefill, codec trimming --- mstar/model/qwen3_tts/submodules.py | 245 ++++++++++++++++++++++++---- 1 file changed, 217 insertions(+), 28 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index f7e293185..6b4588bca 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -17,8 +17,11 @@ # 2. CodecSubmodule (STATELESS engine) # - Receives buffered codec frames from the Talker partition. # - Pads variable final tails to fixed CUDA Graph capture shapes. -# - Runs the official speech-tokenizer decoder and trims overlap before -# emitting 24 kHz PCM. +# - Runs the official speech-tokenizer decoder and trims overlap (and, for +# voice clones, the reference frames) before emitting 24 kHz PCM. +# 3. RefEncoderSubmodule (no resources; Base voice clone only) +# - Turns one reference clip into the ECAPA x-vector that replaces the +# speaker tag and, for in-context cloning, into its codec frames. # # Engine-facing lifecycle: # prepare_inputs -> preprocess -> forward/forward_batched @@ -48,6 +51,10 @@ from mstar.engine.engine import ExecutingBatch from mstar.engine.resources import AttentionStep, KVStep, PositionStep, SamplerStep, Segment, SlotLease, SubmoduleStep from mstar.engine.resources.sampler.resource import SamplerResource +from mstar.model.qwen3_tts.components.speaker_encoder import ( + Qwen3TTSMelFrontEnd, + Qwen3TTSSpeakerEncoder, +) from mstar.model.qwen3_tts.components.talker import ( Qwen3TTSCodePredictor, Qwen3TTSTalkerModel, @@ -67,6 +74,7 @@ ARNodeSubmodule, ModelInputsFromEngine, NodeInputs, + NodeSubmodule, ) # the CodePredictor depth loop, as a piecewise capture region @@ -230,6 +238,14 @@ def _validate_chatml(self, text_ids: torch.Tensor) -> tuple[int, int]: ) return prefix_len, suffix_len + def _frame_embeds(self, codes: torch.Tensor) -> torch.Tensor: + """Sum of the 16 codec embeddings per frame, in the Talker width: ``[frames, D]``.""" + embeds = self.model.model.codec_embedding(codes[:, 0]) + tables = self.code_predictor.model.codec_embedding + for group in range(1, codes.shape[1]): + embeds = embeds + tables[group - 1](codes[:, group]) + return embeds + def _build_prefill( self, request_id: str, @@ -237,31 +253,42 @@ def _build_prefill( prompt_layout: torch.Tensor, speaker_id: int, language_id: int, + speaker_embed: torch.Tensor | None = None, + ref_codes: torch.Tensor | None = None, ) -> torch.Tensor: """Build the official mixed text/codec prefill embedding sequence. - ``text_ids`` concatenates an optional instruction turn with the - assistant turn; ``prompt_layout`` is ``[instruct_len, text_len, - stream_text]`` and says where the split is and how the text is fed. - The layout mirrors ``Qwen3TTSForConditionalGeneration.generate``: + ``text_ids`` concatenates an optional instruction turn, the assistant + turn and (voice clone) the reference transcript turn; ``prompt_layout`` + is ``[instruct_len, text_len, stream_text, ref_text_len, ref_frames]`` + and says where the splits are and how the text is fed. The layout + mirrors ``Qwen3TTSForConditionalGeneration.generate``: * instruction (``<|im_start|>user ... <|im_end|>``) as plain projected text embeddings (VoiceDesign, 1.7B CustomVoice style control); * the assistant role, then the codec think/language tags, the speaker - tag when the checkpoint has built-in speakers (``speaker_id >= 0``) - and the codec PAD, all summed with TTS PAD / BOS text embeddings; - * ``stream_text == 0`` (the reference default for CustomVoice and - VoiceDesign): every text token plus TTS EOS enters the prefill over - codec PADs, closed by TTS PAD + codec BOS. Decode then adds TTS PAD - to each frame. + (a built-in tag when ``speaker_id >= 0``, the reference x-vector + when ``speaker_embed`` is given, nothing for VoiceDesign) and the + codec PAD, all summed with TTS PAD / BOS text embeddings; + * ``ref_frames > 0`` (in-context clone): reference transcript + text + + TTS EOS over codec BOS + the reference frames' codec embeddings, + laid out per ``generate_icl_prompt``; + * otherwise ``stream_text == 0`` (the reference default for + CustomVoice and VoiceDesign): every text token plus TTS EOS enters + the prefill over codec PADs, closed by TTS PAD + codec BOS; decode + then adds TTS PAD to each frame; * ``stream_text == 1`` (the reference default for Base): only the first text token enters the prefill over codec BOS; the remaining tokens plus TTS EOS are kept in per-request state and added one per frame. """ text_ids = text_ids.to(device=self.get_device(), dtype=torch.long).view(1, -1) - instruct_len, text_len, stream_text = (int(v) for v in prompt_layout.tolist()) + layout = [int(v) for v in prompt_layout.tolist()] + instruct_len, text_len, stream_text = layout[:3] + ref_text_len, ref_frames = (layout[3], layout[4]) if len(layout) >= 5 else (0, 0) instruct_ids = text_ids[:, :instruct_len] - assistant_ids = text_ids[:, instruct_len:] + assistant_end = text_ids.shape[1] - ref_text_len + assistant_ids = text_ids[:, instruct_len:assistant_end] + ref_text_ids = text_ids[:, assistant_end:] prefix_len, suffix_len = self._validate_chatml(assistant_ids) if assistant_ids.shape[1] != prefix_len + text_len + suffix_len: raise ValueError( @@ -269,6 +296,11 @@ def _build_prefill( f"expected {prefix_len + text_len + suffix_len} assistant tokens, " f"got {assistant_ids.shape[1]}" ) + if ref_frames > 0 and (ref_codes is None or ref_codes.shape[0] < ref_frames): + raise ValueError( + f"Qwen3-TTS in-context clone needs {ref_frames} reference frames, got " + f"{0 if ref_codes is None else ref_codes.shape[0]}" + ) text_tokens = assistant_ids[:, prefix_len:prefix_len + text_len] codec = self.talker_config @@ -282,21 +314,32 @@ def _build_prefill( codec.codec_think_eos_id, ] ) - speaker_tag = [speaker_id] if speaker_id >= 0 else [] codec_ids = torch.tensor( - [[*codec_prefix, *speaker_tag, codec.codec_pad_id, codec.codec_bos_id]], + [[*codec_prefix, codec.codec_pad_id, codec.codec_bos_id]], dtype=torch.long, device=self.get_device(), ) codec_embeds = self.model.model.codec_embedding(codec_ids) dtype = codec_embeds.dtype + # The speaker slot sits between the language tags and the codec PAD: + # a built-in speaker's codec embedding, or the reference x-vector. + speaker_vector = None + if speaker_id >= 0: + speaker_vector = self.model.model.codec_embedding(codec_ids.new_tensor([[speaker_id]])) + elif speaker_embed is not None: + speaker_vector = speaker_embed.to(device=self.get_device(), dtype=dtype).reshape(1, 1, -1) + if speaker_vector is not None: + split = len(codec_prefix) + codec_embeds = torch.cat( + [codec_embeds[:, :split], speaker_vector, codec_embeds[:, split:]], dim=1 + ) bos_embed, eos_embed, pad_embed = self._special_text_embeds(dtype) def project(ids: torch.Tensor) -> torch.Tensor: return self._project_text(ids).to(dtype) # Tags: TTS PAD over every codec tag but the last, TTS BOS over the - # codec PAD; the codec BOS pairs with text below. + # codec PAD; the codec BOS pairs with text (or the reference) below. role_embed = project(assistant_ids[:, :prefix_len]) tag_text = torch.cat([ pad_embed.expand(-1, codec_embeds.shape[1] - 2, -1), @@ -305,8 +348,33 @@ def project(ids: torch.Tensor) -> torch.Tensor: pieces = [role_embed, tag_text + codec_embeds[:, :-1]] if instruct_len: pieces.insert(0, project(instruct_ids)) - - if stream_text: + empty_trailing = eos_embed[:, :0] + + if ref_frames > 0: + # ``<|im_start|>assistant\n{ref}<|im_end|>\n`` -> transcript tokens only. + ref_tokens = ref_text_ids[:, prefix_len:ref_text_ids.shape[1] - 2] + text_embed = torch.cat([project(torch.cat([ref_tokens, text_tokens], dim=1)), eos_embed], dim=1) + reference = ref_codes.to(device=self.get_device(), dtype=torch.long)[:ref_frames] + codec_embed = torch.cat( + [codec_embeds[:, -1:], self._frame_embeds(reference).to(dtype).unsqueeze(0)], dim=1 + ) + if not stream_text: + codec_pads = self.model.model.codec_embedding( + codec_ids.new_full((1, text_embed.shape[1]), codec.codec_pad_id) + ) + pieces += [text_embed + codec_pads, codec_embed + pad_embed] + trailing = empty_trailing + elif text_embed.shape[1] > codec_embed.shape[1]: + pieces.append(text_embed[:, :codec_embed.shape[1]] + codec_embed) + trailing = text_embed[:, codec_embed.shape[1]:] + else: + padded = torch.cat([ + text_embed, + pad_embed.expand(-1, codec_embed.shape[1] - text_embed.shape[1], -1), + ], dim=1) + pieces.append(padded + codec_embed) + trailing = empty_trailing + elif stream_text: pieces.append(project(text_tokens[:, :1]) + codec_embeds[:, -1:]) trailing = torch.cat([project(text_tokens[:, 1:]), eos_embed], dim=1) else: @@ -317,15 +385,20 @@ def project(ids: torch.Tensor) -> torch.Tensor: pieces.append(text_embed + codec_pads) pieces.append(pad_embed + codec_embeds[:, -1:]) # Nothing left to feed: every frame adds TTS PAD (empty stream). - trailing = eos_embed[:, :0] + trailing = empty_trailing prefill = torch.cat(pieces, dim=1) - self.request_state(request_id).add_all( + state = self.request_state(request_id) + state.add_all( trailing_text_hidden=trailing.squeeze(0), tts_pad_embed=pad_embed[0, 0], generation_step=0, generated_frames=0, ) + if ref_frames > 0: + # The codec decodes the reference frames ahead of the generated + # ones (their audio is trimmed), exactly as the reference does. + state.add("reference_frames", reference) return prefill.squeeze(0) def prepare_inputs( @@ -343,13 +416,15 @@ def prepare_inputs( therefore prepare requests before admitting them to a micro-batch. """ del kwargs - if graph_walk == "talker_prefill": + if graph_walk in ("talker_prefill", "talker_prefill_clone"): input_embeds = self._build_prefill( fwd_info.request_id, inputs["text_inputs"][0], inputs["prompt_layout"][0], int(inputs["speaker_id"][0].item()), int(inputs["language_id"][0].item()), + speaker_embed=inputs["speaker_embed"][0] if "speaker_embed" in inputs else None, + ref_codes=inputs["ref_codes"][0] if "ref_codes" in inputs else None, ) state = self.request_state(fwd_info.request_id) elif graph_walk == "talker_decode": @@ -600,6 +675,20 @@ def _run_depth_loop_piecewise( ) return output["all_codes"], output["codec_embed_sum"] + def _codec_stream_items(self, graph_walk: str, request_id: str, frame: torch.Tensor) -> list[torch.Tensor]: + """Frames this step pushes into the codec stream, one item per frame. + + The clone prefill leads with the reference clip's frames so the codec + warms up on the voice being cloned; ``CodecSubmodule`` trims their + audio. Decode steps (the captured path) always push exactly one frame. + """ + if graph_walk != "talker_prefill_clone": + return [frame] + reference = self.request_state(request_id).get("reference_frames") + if reference is None: + return [frame] + return [*reference.unbind(0), frame] + def forward( self, graph_walk: str, @@ -609,11 +698,16 @@ def forward( suppress_eos: torch.Tensor, **kwargs: Any, ) -> NameToTensorList: - del graph_walk, kwargs + del kwargs output = self._run_frame( engine_inputs, input_embeds, last_token_indices, suppress_eos ) - return {name: [tensor] for name, tensor in output.items()} + request_id = engine_inputs.request_ids[0] + return { + "talker_input_embeds": [output["talker_input_embeds"]], + "codec_tokens": self._codec_stream_items(graph_walk, request_id, output["codec_tokens"][0]), + "new_token": [output["new_token"]], + } def forward_batched( self, @@ -624,14 +718,14 @@ def forward_batched( suppress_eos: torch.Tensor, **kwargs: Any, ) -> dict[str, NameToTensorList]: - del graph_walk, kwargs + del kwargs output = self._run_frame( engine_inputs, input_embeds, last_token_indices, suppress_eos ) return { request_id: { "talker_input_embeds": [output["talker_input_embeds"][i:i + 1]], - "codec_tokens": [output["codec_tokens"][i]], + "codec_tokens": self._codec_stream_items(graph_walk, request_id, output["codec_tokens"][i]), "new_token": [output["new_token"][i]], } for i, request_id in enumerate(engine_inputs.request_ids) @@ -655,7 +749,7 @@ def postprocess( if "new_token" in outputs: outputs["layer0_codes"] = outputs.pop("new_token") elif "layer0_codes" not in outputs and "codec_tokens" in outputs: - codec_tokens = outputs["codec_tokens"][0] + codec_tokens = outputs["codec_tokens"][-1] outputs["layer0_codes"] = [codec_tokens.reshape(-1)[0]] if "layer0_codes" in outputs: state = self.request_state(request_id) @@ -700,7 +794,7 @@ def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bo together freely. """ return ( - batch.graph_walk in {"talker_prefill", "talker_decode"} + batch.graph_walk in {"talker_prefill", "talker_prefill_clone", "talker_decode"} and bool(model_inputs) and len(model_inputs) <= self.MAX_BATCH_SIZE ) @@ -868,6 +962,14 @@ def prepare_inputs( so differently sized final tails can reuse the same CUDA Graph. """ del graph_walk, kwargs + state = self.request_state(fwd_info.request_id) + if "ref_frames" in inputs and "skip_samples" not in state: + # Voice clone: the stream leads with the reference clip's frames, + # whose audio the client must not hear. + state.add( + "skip_samples", + int(inputs["ref_frames"][0].reshape(-1)[0].item()) * self.total_upsample, + ) codes = inputs["codec_tokens"][0].to( device=self.get_device(), dtype=torch.long ) @@ -963,6 +1065,11 @@ def postprocess( left_context = self.config.codec.left_context_frames if emitted else 0 start = left_context * self.total_upsample end = frames * self.total_upsample + skip = int(state.get("skip_samples", 0)) + if skip: + dropped = min(skip, max(end - start, 0)) + start += dropped + state.add("skip_samples", skip - dropped) outputs["audio_chunk"][0] = outputs["audio_chunk"][0][start:end] state.add("codec_chunk_emitted", True) @@ -1017,3 +1124,85 @@ def can_use_cuda_graphs( ) and super().can_use_cuda_graphs(batch, model_inputs) ) + + +# =========================================================================== +# 3. RefEncoderSubmodule - reference audio -> speaker conditioning (Base) +# =========================================================================== + + +class RefEncoderSubmodule(NodeSubmodule): + """Encode one reference clip into the Talker's voice conditioning. + + Runs once per request, before the clone prefill, and owns no resources. + The x-vector comes from the ECAPA-TDNN encoder over a log-mel spectrogram + of the 24 kHz clip; for in-context cloning the codec encoder also turns + the clip into ``ref_frames`` 16-group frames. The mel front end and the + codec encoder run in float32 regardless of the engine's autocast dtype; + the x-vector is produced in the encoder's own (Talker) dtype. + """ + + disable_torch_compile = True + + def __init__( + self, + speaker_encoder: Qwen3TTSSpeakerEncoder, + mel_front_end: Qwen3TTSMelFrontEnd, + codec_encoder: torch.nn.Module, + config: Qwen3TTSModelConfig, + ) -> None: + super().__init__() + self.speaker_encoder = speaker_encoder + self.mel_front_end = mel_front_end + self.codec_encoder = codec_encoder + self.config = config + + def prepare_inputs( + self, + graph_walk: str, + fwd_info: CurrentForwardPassInfo, + inputs: NameToTensorList, + **kwargs: Any, + ) -> NodeInputs: + del graph_walk, fwd_info, kwargs + waveform = inputs["audio_inputs"][0].to(device=self.get_device(), dtype=torch.float32) + if waveform.ndim > 1: + waveform = waveform.mean(dim=0) if waveform.shape[0] < waveform.shape[-1] else waveform.mean(dim=-1) + layout = inputs["prompt_layout"][0].tolist() + ref_frames = int(layout[4]) if len(layout) >= 5 else 0 + return NodeInputs( + tensor_inputs={"waveform": waveform.reshape(-1)}, + kwargs={"ref_frames": ref_frames}, + ) + + def forward( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + waveform: torch.Tensor, + ref_frames: int = 0, + **kwargs: Any, + ) -> NameToTensorList: + del graph_walk, engine_inputs, kwargs + device_type = waveform.device.type + with torch.autocast(device_type=device_type, enabled=False): + mels = self.mel_front_end(waveform.unsqueeze(0)) + encoder_dtype = next(self.speaker_encoder.parameters()).dtype + speaker_embed = self.speaker_encoder(mels.to(encoder_dtype))[0] + + num_quantizers = self.config.codec.num_quantizers + if ref_frames > 0: + with torch.autocast(device_type=device_type, enabled=False): + encoded = self.codec_encoder.encode( + input_values=waveform.view(1, 1, -1).float(), return_dict=True + ) + codes = encoded.audio_codes[0, :num_quantizers, :ref_frames].t().contiguous().long() + if codes.shape[0] < ref_frames: + raise ValueError( + f"codec encoder produced {codes.shape[0]} frames for a clip declared as {ref_frames}" + ) + else: + # x-vector-only clone: no in-context frames. The edge still needs + # a tensor; the Talker ignores it because the layout says 0 frames. + codes = torch.zeros(1, num_quantizers, dtype=torch.long, device=waveform.device) + return {"speaker_embed": [speaker_embed], "ref_codes": [codes]} From 6439f219bba375ca9f32c4eb594f3fc3f829ca3b Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 025/110] qwen3_tts: Base voice-clone walks and CPU-built codec modules --- mstar/model/qwen3_tts/qwen3_tts_model.py | 350 +++++++++++++++++------ 1 file changed, 262 insertions(+), 88 deletions(-) diff --git a/mstar/model/qwen3_tts/qwen3_tts_model.py b/mstar/model/qwen3_tts/qwen3_tts_model.py index d90472446..9f4daf853 100644 --- a/mstar/model/qwen3_tts/qwen3_tts_model.py +++ b/mstar/model/qwen3_tts/qwen3_tts_model.py @@ -11,15 +11,16 @@ registry key. Architecture (two asynchronous partitions): - Talker - text/voice prefill, then autoregressive 16-group codec frames - Codec - stateless speech-tokenizer decoder producing PCM chunks + Talker - text/voice prefill, then autoregressive 16-group codec frames + RefEncoder - (Base) reference clip -> x-vector + codec frames, feeds the prefill + Codec - stateless speech-tokenizer decoder producing PCM chunks Streaming topology: Talker --[codec_tokens, LeftContextChunkPolicy(300, 25)]--> Codec Request state machine: - Talker: talker_prefill -> talker_decode loop -> done on EOS/token limit - Codec: waits for streamed frames -> codec_chunk -> emits audio -> waits + Talker: talker_prefill | talker_prefill_clone -> talker_decode loop -> done on EOS/token limit + Codec: waits for streamed frames -> codec_chunk | codec_chunk_clone -> emits audio -> waits This class runs in the API/conductor side. It owns request validation, graph and partition declarations, state-machine transitions, sampling defaults, and @@ -59,10 +60,11 @@ GraphNode, GraphSection, Loop, + Sequential, TensorPointerInfo, ) from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION -from mstar.model.base import ForwardPassArgs, Model +from mstar.model.base import ForwardPassArgs, Model, TensorAndMetadata from mstar.model.qwen3_tts.config import ( CHATML_ASSISTANT_PREFIX_TOKEN_IDS, CHATML_ASSISTANT_SUFFIX_TOKEN_IDS, @@ -112,8 +114,11 @@ def _resolve_model_metadata(repo_id: str, cache_dir: str | None) -> str: @lru_cache(maxsize=1) -def _load_qwen3_tts_decoder_classes() -> tuple[type, type]: - """Load only qwen-tts' 12 Hz decoder modules. +def _load_qwen3_tts_codec_classes() -> tuple[type, type, type]: + """Load only qwen-tts' 12 Hz speech-tokenizer modules. + + Returns the decoder config class, the decoder (codes -> waveform) and the + encoder (waveform -> codes, used for voice-clone reference audio). ``qwen_tts.__init__`` eagerly imports its high-level inference package, which in turn imports the unrelated 25 Hz tokenizer and pysox. Pysox @@ -176,6 +181,7 @@ def load_private_module(name: str, path: Path): return ( config_module.Qwen3TTSTokenizerV2DecoderConfig, model_module.Qwen3TTSTokenizerV2Decoder, + model_module.Qwen3TTSTokenizerV2Encoder, ) @@ -356,27 +362,45 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: ``talker_input_embeds`` is the recurrent Talker edge. ``codec_tokens`` crosses the asynchronous partition boundary and is buffered according - to ``get_partition_topology`` before Codec is scheduled. + to ``get_partition_topology`` before Codec is scheduled. Base + checkpoints add a clone prefill (RefEncoder -> Talker) and a codec walk + that also receives the reference frame count to trim. """ - # Prefill seeds both recurrent paths: the embedding for the next - # Talker step is persisted, while the first codec frame starts the - # Talker-to-Codec stream. - talker_prefill = GraphNode( - name="Talker", - input_names=list(self.PREFILL_INPUTS), - outputs=[ - GraphEdge( - next_node=EMPTY_DESTINATION, - name="talker_input_embeds", - persist=True, - ), - StreamingGraphEdge( - next_node="Codec", - name="codec_tokens", - target_partition="Codec", - ), - ], - ) + def talker_prefill_node(input_names: list[str]) -> GraphNode: + # Prefill seeds both recurrent paths: the embedding for the next + # Talker step is persisted, while the first codec frame(s) start + # the Talker-to-Codec stream. + return GraphNode( + name="Talker", + input_names=input_names, + outputs=[ + GraphEdge( + next_node=EMPTY_DESTINATION, + name="talker_input_embeds", + persist=True, + ), + StreamingGraphEdge( + next_node="Codec", + name="codec_tokens", + target_partition="Codec", + ), + ], + ) + + def codec_node(input_names: list[str]) -> GraphNode: + # Codec is deliberately a separate walk/engine so waveform decoding + # can overlap with subsequent Talker steps. + return GraphNode( + name="Codec", + input_names=input_names, + outputs=[ + GraphEdge( + next_node=EMIT_TO_CLIENT, + name="audio_chunk", + output_modality="audio", + ), + ], + ) # Each loop iteration predicts one complete 16-group codec frame and # feeds the summed codec embedding back into the next Talker step. @@ -400,25 +424,26 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: max_iters=self.get_max_output_tokens(), outputs=[], ) - - # Codec is deliberately a separate walk/engine so waveform decoding - # can overlap with subsequent Talker steps. - codec_chunk = GraphNode( - name="Codec", - input_names=["codec_tokens"], - outputs=[ - GraphEdge( - next_node=EMIT_TO_CLIENT, - name="audio_chunk", - output_modality="audio", - ), - ], - ) - return { - "talker_prefill": talker_prefill, + walks = { + "talker_prefill": talker_prefill_node(list(self.PREFILL_INPUTS)), "talker_decode": talker_decode, - "codec_chunk": codec_chunk, + "codec_chunk": codec_node(["codec_tokens"]), } + if self.config.supports_reference_audio: + ref_encoder = GraphNode( + name="RefEncoder", + input_names=list(self.REF_ENCODER_INPUTS), + outputs=[ + GraphEdge(next_node="Talker", name="speaker_embed"), + GraphEdge(next_node="Talker", name="ref_codes"), + ], + ) + walks["talker_prefill_clone"] = Sequential([ + ref_encoder, + talker_prefill_node([*self.PREFILL_INPUTS, "speaker_embed", "ref_codes"]), + ]) + walks["codec_chunk_clone"] = codec_node(["codec_tokens", "ref_frames"]) + return walks # ----------------------------------------------------------------------- # Asynchronous partitions and stream buffering @@ -426,16 +451,21 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: def get_partitions(self) -> list[PartitionDefinition]: """Split autoregressive generation from independently scheduled audio.""" + talker_walks = {"talker_prefill", "talker_decode"} + codec_walks = {"codec_chunk"} + if self.config.supports_reference_audio: + talker_walks.add("talker_prefill_clone") + codec_walks.add("codec_chunk_clone") return [ PartitionDefinition( name="Talker", - graph_walks={"talker_prefill", "talker_decode"}, + graph_walks=talker_walks, initial_walk="talker_prefill", producer_partitions=[], ), PartitionDefinition( name="Codec", - graph_walks={"codec_chunk"}, + graph_walks=codec_walks, initial_walk=None, producer_partitions=["Talker"], ), @@ -474,8 +504,26 @@ def get_partition_topology(self) -> PartitionTopology: # which is how the text span is located inside the tokenized turn. ASSISTANT_TEMPLATE = "<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" INSTRUCT_TEMPLATE = "<|im_start|>user\n{instruct}<|im_end|>\n" - # Tensors ``process_prompt`` produces and the Talker prefill consumes. + REFERENCE_TEMPLATE = "<|im_start|>assistant\n{text}<|im_end|>\n" + # API tensors the Talker prefill consumes; the clone prefill adds the + # RefEncoder's ``speaker_embed`` and ``ref_codes`` to them. PREFILL_INPUTS = ("text_inputs", "prompt_layout", "speaker_id", "language_id") + # API tensors the RefEncoder consumes (reference clip + layout). + REF_ENCODER_INPUTS = ("audio_inputs", "prompt_layout") + + def load_audio(self, filepath: str, device: str) -> TensorAndMetadata: + """Decode a reference clip to 24 kHz mono float32 (speaker encoder + codec rate).""" + import soundfile + import torchaudio.functional as audio_functional + + waveform, sample_rate = soundfile.read(filepath, dtype="float32", always_2d=True) + audio = torch.from_numpy(waveform).mean(dim=1) + target_rate = self.config.codec.input_sample_rate + if sample_rate != target_rate: + audio = audio_functional.resample(audio, sample_rate, target_rate) + return TensorAndMetadata( + data=audio.to(device), metadata={"sample_rate": target_rate, "num_channels": 1} + ) def _tokenize(self, text: str) -> torch.Tensor: encoded = self.tokenizer(text, return_tensors="pt", padding=True) @@ -549,6 +597,43 @@ def _resolve_instruct(self, kwargs: dict[str, Any]) -> str: ) return instruct + def _resolve_reference( + self, kwargs: dict[str, Any], tensors: NameToTensorList | None, input_modalities: list[str], + ) -> tuple[str, int]: + """Voice-clone reference: (transcript, frames). ``frames == 0`` means x-vector only. + + Base needs exactly one reference clip. In-context cloning (the + default) also needs the clip's transcript; ``x_vector_only_mode`` + drops both the transcript and the frames and conditions on the + x-vector alone. + """ + clips = (tensors or {}).get("audio_inputs", []) + if not self.config.supports_reference_audio: + if clips or "audio" in input_modalities: + raise ValueError( + f"Qwen3-TTS {self.config.tts_model_type} does not take reference " + "audio; voice cloning needs a Base checkpoint" + ) + return "", 0 + if len(clips) != 1: + raise ValueError( + "Qwen3-TTS Base clones a voice from exactly one reference clip " + f"(got {len(clips)}); pass it as the request's audio input" + ) + ref_text = str(kwargs.get("ref_text") or "").strip() + if kwargs.get("x_vector_only_mode", False): + return "", 0 + if not ref_text: + raise ValueError( + "Qwen3-TTS Base needs 'ref_text' (the reference clip's transcript) " + "unless 'x_vector_only_mode' is set" + ) + num_samples = int(clips[0].reshape(-1).shape[0]) + frames = self.config.codec.frames_for_samples(num_samples) + if frames < 1: + raise ValueError("Qwen3-TTS reference clip is empty") + return ref_text, frames + def process_prompt( self, prompt: str | None, @@ -559,31 +644,36 @@ def process_prompt( ) -> NameToTensorList: """Validate a request against the checkpoint variant and tokenize it. - Produces the four ``PREFILL_INPUTS`` tensors. ``text_inputs`` is the - optional instruction turn followed by the assistant turn (both in the - reference ChatML templates); ``prompt_layout`` is - ``[instruct_len, text_len, stream_text]``. ``stream_text`` follows the - reference default per variant (whole text in the prefill for - CustomVoice/VoiceDesign, one token per frame for Base) unless the - request sets ``non_streaming_mode``. + Produces the ``PREFILL_INPUTS`` tensors. ``text_inputs`` is the + optional instruction turn, the assistant turn and (in-context clone) + the reference transcript turn, all in the reference ChatML templates; + ``prompt_layout`` is ``[instruct_len, text_len, stream_text, + ref_text_len, ref_frames]``. ``stream_text`` follows the reference + default per variant (whole text in the prefill for CustomVoice and + VoiceDesign, one token per frame for Base) unless the request sets + ``non_streaming_mode``. Base requests also get ``ref_frames`` as its + own tensor for the codec's trimming. """ - del tensors if not prompt: raise ValueError("Qwen3-TTS requires a non-empty text prompt") - if set(input_modalities) != {"text"}: - raise ValueError("Qwen3-TTS currently supports text input only") - if set(output_modalities) != {"audio"}: - raise ValueError("Qwen3-TTS supports audio output only") - if self.config.is_base: + if "audio" in input_modalities and not self.config.supports_reference_audio: + raise ValueError( + f"Qwen3-TTS {self.config.tts_model_type} does not take reference " + "audio; voice cloning needs a Base checkpoint" + ) + allowed = {"text", "audio"} if self.config.supports_reference_audio else {"text"} + if "text" not in input_modalities or not set(input_modalities) <= allowed: raise ValueError( - "Qwen3-TTS Base clones a voice from reference audio, which " - "this build does not accept yet; use a CustomVoice or " - "VoiceDesign checkpoint for text-only requests" + "Qwen3-TTS currently supports text input only" + + (" (plus one reference clip for Base)" if self.config.supports_reference_audio else "") ) + if set(output_modalities) != {"audio"}: + raise ValueError("Qwen3-TTS supports audio output only") speaker, speaker_id = self._resolve_speaker(kwargs) language_id = self._resolve_language(kwargs, speaker) instruct = self._resolve_instruct(kwargs) + ref_text, ref_frames = self._resolve_reference(kwargs, tensors, input_modalities) stream_text = not bool( kwargs.get("non_streaming_mode", self.config.default_non_streaming_mode) ) @@ -600,14 +690,23 @@ def process_prompt( if instruct else assistant_ids.new_empty(0) ) - return { - "text_inputs": [torch.cat([instruct_ids, assistant_ids])], + ref_ids = ( + self._tokenize(self.REFERENCE_TEMPLATE.format(text=ref_text)) + if ref_frames + else assistant_ids.new_empty(0) + ) + outputs = { + "text_inputs": [torch.cat([instruct_ids, assistant_ids, ref_ids])], "prompt_layout": [torch.tensor( - [instruct_ids.numel(), text_len, int(stream_text)], dtype=torch.long + [instruct_ids.numel(), text_len, int(stream_text), ref_ids.numel(), ref_frames], + dtype=torch.long, )], "speaker_id": [torch.tensor([speaker_id], dtype=torch.long)], "language_id": [torch.tensor([language_id], dtype=torch.long)], } + if self.config.supports_reference_audio: + outputs["ref_frames"] = [torch.tensor([ref_frames], dtype=torch.long)] + return outputs # ----------------------------------------------------------------------- # Conductor partition state machine @@ -623,24 +722,30 @@ def get_initial_forward_pass_args( ) -> ForwardPassArgs: """Create each partition's initial state. - Talker starts immediately from API tensors. Codec has no direct API - inputs and remains dormant until its incoming streaming connection has - enough frames to schedule ``codec_chunk``. + Talker starts immediately from API tensors; a request that carries a + reference clip takes the clone prefill, whose RefEncoder runs first. + Codec has no direct API inputs and remains dormant until its incoming + streaming connection has enough frames to schedule its chunk walk. """ model_kwargs = model_kwargs or {} + clone = "audio" in input_modalities if partition_name == "Talker": + walk = "talker_prefill_clone" if clone else "talker_prefill" metadata = CurrentForwardConductorMetadata( input_modalities=input_modalities, output_modalities=output_modalities, - graph_walk="talker_prefill", + graph_walk=walk, is_prefill=True, kwargs={ "talker_max_tokens": self.get_max_output_tokens(**model_kwargs), }, ) + routes = [(name, "Talker") for name in self.PREFILL_INPUTS] + if clone: + routes += [(name, "RefEncoder") for name in self.REF_ENCODER_INPUTS] inputs = [] - for name in self.PREFILL_INPUTS: - edge = GraphEdge(next_node="Talker", name=name) + for name, node in routes: + edge = GraphEdge(next_node=node, name=name) edge.tensor_info = input_signals.get(name, []) inputs.append(edge) return ForwardPassArgs( @@ -657,7 +762,7 @@ def get_initial_forward_pass_args( metadata = CurrentForwardConductorMetadata( input_modalities=input_modalities, output_modalities=output_modalities, - graph_walk="codec_chunk", + graph_walk="codec_chunk_clone" if clone else "codec_chunk", is_prefill=False, ) return ForwardPassArgs( @@ -684,7 +789,7 @@ def get_partition_forward_pass_args( """ del incoming_connections if partition_name == "Talker": - if partition_metadata.graph_walk == "talker_prefill": + if partition_metadata.graph_walk in ("talker_prefill", "talker_prefill_clone"): partition_metadata.graph_walk = "talker_decode" partition_metadata.is_prefill = False edge = GraphEdge(next_node="Talker", name="talker_input_embeds") @@ -708,10 +813,18 @@ def get_partition_forward_pass_args( ) if partition_name == "Codec": - partition_metadata.graph_walk = "codec_chunk" + inputs = [] + if partition_metadata.graph_walk == "codec_chunk_clone": + # The reference frame count is an API tensor; every codec + # invocation of a clone request re-reads it (cheap, one int). + edge = GraphEdge(next_node="Codec", name="ref_frames") + edge.tensor_info = persist_signals.get("ref_frames", []) + inputs.append(edge) + else: + partition_metadata.graph_walk = "codec_chunk" return ForwardPassArgs( full_metadata=partition_metadata, - inputs=[], + inputs=inputs, unpersist_tensors=[], step_metadata={ "codec_chunk_frames": self.config.codec.chunk_frames, @@ -816,7 +929,7 @@ def get_submodule( ) -> NodeSubmodule | None: """Build only the node assigned to this worker and cache the wrapper.""" del sp_group - if node_name not in ("Talker", "Codec"): + if node_name not in ("Talker", "Codec", "RefEncoder"): raise ValueError(f"Unknown Qwen3-TTS node: {node_name!r}") if node_name in self._submodule_cache: return self._submodule_cache[node_name] @@ -828,6 +941,10 @@ def get_submodule( tp_group=tp_group, autocast_dtype=autocast_dtype, ) + elif node_name == "RefEncoder": + submodule = self._create_ref_encoder_submodule( + device=device, autocast_dtype=autocast_dtype + ) else: submodule = self._create_codec_submodule(device=device) self._submodule_cache[node_name] = submodule @@ -910,10 +1027,7 @@ def talker_weights(): def _create_codec_submodule(self, device: str) -> NodeSubmodule: """Build the official speech-tokenizer decoder from its sub-checkpoint.""" try: - ( - Qwen3TTSTokenizerV2DecoderConfig, - Qwen3TTSTokenizerV2Decoder, - ) = _load_qwen3_tts_decoder_classes() + decoder_config_cls, decoder_cls, _ = _load_qwen3_tts_codec_classes() except ImportError as exc: raise ImportError( "Qwen3-TTS Codec requires the 'qwen-tts' package; install " @@ -925,13 +1039,12 @@ def _create_codec_submodule(self, device: str) -> NodeSubmodule: from mstar.model.qwen3_tts.submodules import CodecSubmodule # Reuse the official decoder implementation, but keep graph scheduling, - # chunk padding, overlap trimming, and output transport in M*. - decoder_config = Qwen3TTSTokenizerV2DecoderConfig( - **self.config.codec.decoder_kwargs() - ) - with torch.device("meta"): - decoder = Qwen3TTSTokenizerV2Decoder(decoder_config) - decoder.to_empty(device=device) + # chunk padding, overlap trimming, and output transport in M*. The + # 114M-parameter module is built on the CPU rather than on ``meta``: + # its rotary tables are non-persistent buffers that ``to_empty`` would + # leave uninitialized (no checkpoint tensor refills them). + decoder_config = decoder_config_cls(**self.config.codec.decoder_kwargs()) + decoder = decoder_cls(decoder_config).to(device=device) codec_dir = Path(self.local_dir) / "speech_tokenizer" prefix = "decoder." @@ -948,6 +1061,67 @@ def _create_codec_submodule(self, device: str) -> NodeSubmodule: decoder.eval() return CodecSubmodule(decoder, self.config) + def _create_ref_encoder_submodule( + self, device: str, autocast_dtype: torch.dtype | None = None, + ) -> NodeSubmodule: + """Speaker encoder (main checkpoint) + codec encoder (speech tokenizer) for Base.""" + if not self.config.supports_reference_audio or self.config.speaker_encoder is None: + raise ValueError( + f"Qwen3-TTS {self.config.tts_model_type} has no reference-audio encoder" + ) + from transformers import MimiConfig + + from mstar.model.loader import load_hf_weights + from mstar.model.loader.iterators import iter_safetensors_shards + from mstar.model.qwen3_tts.components.speaker_encoder import ( + Qwen3TTSMelFrontEnd, + Qwen3TTSSpeakerEncoder, + ) + from mstar.model.qwen3_tts.submodules import RefEncoderSubmodule + + # The x-vector is computed in the Talker's dtype, as the reference does. + speaker_encoder = Qwen3TTSSpeakerEncoder(self.config.speaker_encoder) + if autocast_dtype is not None: + speaker_encoder = speaker_encoder.to(autocast_dtype) + speaker_encoder = speaker_encoder.to(device=device) + prefix = "speaker_encoder." + loaded = load_hf_weights( + speaker_encoder, + ( + (name.removeprefix(prefix), tensor) + for name, tensor in iter_safetensors_shards(self.local_dir, device=device, prefix=prefix) + ), + ) + _verify_checkpoint_coverage( + speaker_encoder, loaded, _checkpoint_keys(self.local_dir, prefix), "Qwen3-TTS speaker encoder" + ) + speaker_encoder.eval() + + # Mimi encoder of the speech tokenizer: reference clip -> codec frames. + # Float32 like the decoder; built on the CPU for its non-persistent + # buffers (rotary tables, convolution geometry). + _, _, encoder_cls = _load_qwen3_tts_codec_classes() + codec_encoder = encoder_cls(MimiConfig(**self.config.codec.encoder_config)).to(device=device) + codec_dir = Path(self.local_dir) / "speech_tokenizer" + prefix = "encoder." + loaded = load_hf_weights( + codec_encoder, + ( + (name.removeprefix(prefix), tensor) + for name, tensor in iter_safetensors_shards(codec_dir, device=device, prefix=prefix) + ), + ) + _verify_checkpoint_coverage( + codec_encoder, loaded, _checkpoint_keys(codec_dir, prefix), "Qwen3-TTS codec encoder" + ) + codec_encoder.eval() + return RefEncoderSubmodule( + speaker_encoder, + Qwen3TTSMelFrontEnd(self.config.speaker_encoder).to(device=device), + codec_encoder, + self.config, + ) + @staticmethod def _talker_step_metadata( metadata: CurrentForwardConductorMetadata, From 9b3453e1b40fb9b5001728b4fd952b165ab7e08c Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 026/110] api: reference audio for Qwen3-TTS speech requests --- mstar/api_server/openai/adapters.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/mstar/api_server/openai/adapters.py b/mstar/api_server/openai/adapters.py index e979339c2..394413a73 100644 --- a/mstar/api_server/openai/adapters.py +++ b/mstar/api_server/openai/adapters.py @@ -394,7 +394,7 @@ class Qwen3TTSAdapter(OpenAIAdapter): supports_speech = True - def speech_to_request(self, req: SpeechRequest, upload_dir: Path) -> SubmitArgs: # noqa: ARG002 + def speech_to_request(self, req: SpeechRequest, upload_dir: Path) -> SubmitArgs: mk = _passthrough(req) if getattr(req, "voice", None): mk["voice"] = req.voice @@ -403,6 +403,24 @@ def speech_to_request(self, req: SpeechRequest, upload_dir: Path) -> SubmitArgs: if instructions: mk.setdefault("instruct", instructions) _apply_sampling(req, mk, temperature_key="temperature", top_p_key="top_p", max_tokens_key=None) + # Voice clone (Base): ``ref_audio`` is a data URL, an http(s) URL, a + # local path or bare base64 (vLLM-Omni's field). It becomes the + # request's audio input for one request; ``ref_text`` and + # ``x_vector_only_mode`` ride along as model kwargs. Named, persisted + # voices come with the shared voice registry (engine/voice-registry). + ref_audio = mk.pop("ref_audio", None) + if ref_audio: + if ref_audio.startswith(("data:", "http://", "https://")) or Path(ref_audio).suffix: + _, path = media_io.resolve_media_ref(ref_audio, upload_dir) + else: + _, path = media_io.save_base64(ref_audio, "wav", "audio", upload_dir) + return SubmitArgs( + text=req.input, + file_paths={"audio": [path]}, + input_modalities=["audio", "text"], + output_modalities=["audio"], + model_kwargs=mk, + ) return SubmitArgs( text=req.input, input_modalities=["text"], From 6d3e80b3cfb63fc702dbd54fb577ada26eda53ea Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 027/110] configs: map the Qwen3-TTS RefEncoder node --- configs/qwen3tts_base.yaml | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/configs/qwen3tts_base.yaml b/configs/qwen3tts_base.yaml index 621e9cdd5..abd22d29e 100644 --- a/configs/qwen3tts_base.yaml +++ b/configs/qwen3tts_base.yaml @@ -1,7 +1,8 @@ # Qwen3-TTS-12Hz-1.7B-Base: zero-shot voice clone from reference audio (speaker encoder + optional ICL) -# Same graph, resources and single-GPU placement as the 0.6B deployment -# (configs/qwen3tts.yaml); only the checkpoint differs. KV geometry is -# identical across sizes (28 layers, 8 KV heads, head_dim 128). +# Same Talker/Codec graph and resources as the 0.6B deployment +# (configs/qwen3tts.yaml) plus the RefEncoder node, which turns the request's +# reference clip into the x-vector and codec frames the clone prefill needs. +# KV geometry is identical across sizes (28 layers, 8 KV heads, head_dim 128). model: "qwen3_tts_base" max_seq_len: 32768 # The supported deployment image cannot build FlashInfer's Hopper FA3 JIT @@ -10,9 +11,9 @@ resources: talker_attn: flashinfer_backend: fa2 node_groups: - - node_names: [Talker] + - node_names: [RefEncoder, Talker] ranks: [0] - graph_walks: [talker_prefill, talker_decode] + graph_walks: [talker_prefill, talker_prefill_clone, talker_decode] - node_names: [Codec] ranks: [0] - graph_walks: [codec_chunk] + graph_walks: [codec_chunk, codec_chunk_clone] From 29722ccf82c65071bb28e6591f6627c30e2b7044 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 028/110] test: Qwen3-TTS voice-clone prompt, prefill and codec trimming --- test/modular/test_qwen3_tts_model.py | 283 +++++++++++++++++++++++++-- 1 file changed, 272 insertions(+), 11 deletions(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index 7ba27bbbd..10b5bece8 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -61,9 +61,14 @@ def last_text(self): def __call__(self, text, **kwargs): self.texts.append(text) assert kwargs == {"return_tensors": "pt", "padding": True} - if text.startswith("<|im_start|>assistant\n"): + if text.endswith("<|im_end|>\n<|im_start|>assistant\n"): + # the assistant turn to synthesize body = text[len("<|im_start|>assistant\n"):-len("<|im_end|>\n<|im_start|>assistant\n")] prefix, suffix = ASSISTANT_PREFIX, ASSISTANT_SUFFIX + elif text.startswith("<|im_start|>assistant\n"): + # the reference transcript turn (voice clone) + body = text[len("<|im_start|>assistant\n"):-len("<|im_end|>\n")] + prefix, suffix = ASSISTANT_PREFIX, USER_SUFFIX else: assert text.startswith("<|im_start|>user\n") body = text[len("<|im_start|>user\n"):-len("<|im_end|>\n")] @@ -226,8 +231,14 @@ def test_qwen3_tts_1p7b_variants_share_class_configs_and_adapter(): (CONFIG_PATH.parent / yaml_name).read_text(encoding="utf-8") ) assert deployment["model"] == key - assert deployment["node_groups"] == base_yaml["node_groups"] assert deployment["resources"] == base_yaml["resources"] + ranks = {name: group["ranks"] for group in deployment["node_groups"] for name in group["node_names"]} + base_ranks = {name: group["ranks"] for group in base_yaml["node_groups"] for name in group["node_names"]} + assert ranks["Talker"] == base_ranks["Talker"] and ranks["Codec"] == base_ranks["Codec"] + if key == "qwen3_tts_base": + assert ranks["RefEncoder"] == ranks["Talker"] # the clone prefill runs on the Talker's GPU + else: + assert deployment["node_groups"] == base_yaml["node_groups"] assert isinstance(get_adapter(key), Qwen3TTSAdapter) assert isinstance(get_adapter("qwen3_tts"), Qwen3TTSAdapter) @@ -255,9 +266,9 @@ def test_qwen3_tts_decoder_import_does_not_probe_sox(): script = """ import sys import importlib.util -from mstar.model.qwen3_tts.qwen3_tts_model import _load_qwen3_tts_decoder_classes -config_cls, decoder_cls = _load_qwen3_tts_decoder_classes() -print(config_cls.__name__, decoder_cls.__name__) +from mstar.model.qwen3_tts.qwen3_tts_model import _load_qwen3_tts_codec_classes +config_cls, decoder_cls, encoder_cls = _load_qwen3_tts_codec_classes() +print(config_cls.__name__, decoder_cls.__name__, encoder_cls.__name__) print('sox_loaded=' + str('sox' in sys.modules)) print('public_qwen_tts_loaded=' + str(any( name == 'qwen_tts' or name.startswith('qwen_tts.') for name in sys.modules @@ -271,7 +282,7 @@ def test_qwen3_tts_decoder_import_does_not_probe_sox(): capture_output=True, text=True, ) - assert "Qwen3TTSTokenizerV2DecoderConfig Qwen3TTSTokenizerV2Decoder" in ( + assert "Qwen3TTSTokenizerV2DecoderConfig Qwen3TTSTokenizerV2Decoder Qwen3TTSTokenizerV2Encoder" in ( result.stdout ) assert "sox_loaded=False" in result.stdout @@ -331,8 +342,10 @@ def test_qwen3_tts_process_prompt_matches_official_template(): "<|im_start|>assistant\n" ) assert tensors["text_inputs"][0].tolist() == ASSISTANT_PREFIX + [1000] + ASSISTANT_SUFFIX - # CustomVoice default: whole text in the prefill (stream_text = 0). - assert tensors["prompt_layout"][0].tolist() == [0, 1, 0] + # CustomVoice default: whole text in the prefill (stream_text = 0); no + # reference transcript or frames. + assert tensors["prompt_layout"][0].tolist() == [0, 1, 0, 0, 0] + assert "ref_frames" not in tensors assert tensors["speaker_id"][0].item() == 3065 assert tensors["language_id"][0].item() == 2055 @@ -366,7 +379,7 @@ def test_qwen3_tts_1p7b_custom_voice_prepends_instruction_turn(): assistant_ids = ASSISTANT_PREFIX + [1000, 1001, 1002] + ASSISTANT_SUFFIX assert model.tokenizer.texts[-1] == "<|im_start|>user\nspeak slowly<|im_end|>\n" assert tensors["text_inputs"][0].tolist() == instruct_ids + assistant_ids - assert tensors["prompt_layout"][0].tolist() == [len(instruct_ids), 3, 1] + assert tensors["prompt_layout"][0].tolist() == [len(instruct_ids), 3, 1, 0, 0] assert tensors["speaker_id"][0].item() == 3061 @@ -382,7 +395,7 @@ def test_qwen3_tts_voice_design_requires_instruct_and_has_no_speakers(): instruct="A deep, calm male voice", ) assert tensors["speaker_id"][0].item() == -1 - assert tensors["prompt_layout"][0].tolist() == [3 + 5 + 2, 1, 0] + assert tensors["prompt_layout"][0].tolist() == [3 + 5 + 2, 1, 0, 0, 0] with pytest.raises(ValueError, match="requires an 'instruct'"): model.process_prompt("hello", input_modalities=["text"], output_modalities=["audio"]) @@ -400,8 +413,100 @@ def test_qwen3_tts_base_config_declares_speaker_encoder(): assert model.config.speaker_encoder.enc_dim == 2048 # Base feeds text one token per frame by default (reference default). assert model.config.default_non_streaming_mode is False - with pytest.raises(ValueError, match="reference audio"): + + +def test_qwen3_tts_base_process_prompt_builds_in_context_clone(): + model = _variant_model("base") + clip = torch.zeros(24000 + 1) # 1 s + 1 sample -> 13 codec frames at 1920 samples/frame + tensors = model.process_prompt( + "hello big world", + input_modalities=["audio", "text"], + output_modalities=["audio"], + tensors={"audio_inputs": [clip]}, + ref_text="the reference says", + language="English", + ) + ref_ids = ASSISTANT_PREFIX + [1000, 1001, 1002] + USER_SUFFIX # same <|im_end|>\n tail + assistant_ids = ASSISTANT_PREFIX + [1000, 1001, 1002] + ASSISTANT_SUFFIX + assert model.tokenizer.texts[-1] == "<|im_start|>assistant\nthe reference says<|im_end|>\n" + assert tensors["text_inputs"][0].tolist() == assistant_ids + ref_ids + # [instruct_len, text_len, stream_text (Base default: streaming), ref_text_len, ref_frames] + assert tensors["prompt_layout"][0].tolist() == [0, 3, 1, len(ref_ids), 13] + assert tensors["ref_frames"][0].item() == 13 + assert tensors["speaker_id"][0].item() == -1 + + xvec = model.process_prompt( + "hello", input_modalities=["audio", "text"], output_modalities=["audio"], + tensors={"audio_inputs": [clip]}, x_vector_only_mode=True, + ) + assert xvec["prompt_layout"][0].tolist() == [0, 1, 1, 0, 0] + assert xvec["ref_frames"][0].item() == 0 + + with pytest.raises(ValueError, match="ref_text"): + model.process_prompt( + "hello", input_modalities=["audio", "text"], output_modalities=["audio"], + tensors={"audio_inputs": [clip]}, + ) + with pytest.raises(ValueError, match="exactly one reference clip"): model.process_prompt("hello", input_modalities=["text"], output_modalities=["audio"]) + with pytest.raises(ValueError, match="no built-in speakers"): + model.process_prompt( + "hello", input_modalities=["audio", "text"], output_modalities=["audio"], + tensors={"audio_inputs": [clip]}, ref_text="x", voice="vivian", + ) + # Other variants refuse reference audio instead of ignoring it. + with pytest.raises(ValueError, match="does not take reference audio"): + _variant_model("custom_voice").process_prompt( + "hello", input_modalities=["audio", "text"], output_modalities=["audio"], + tensors={"audio_inputs": [clip]}, + ) + + +def test_qwen3_tts_base_declares_clone_walks_and_routes_reference_audio(): + model = _variant_model("base") + walks = model.get_graph_walk_graphs() + assert {"talker_prefill_clone", "codec_chunk_clone"} <= set(walks) + assert "RefEncoder" in model.nodes + partitions = {part.name: part for part in model.get_partitions()} + assert "talker_prefill_clone" in partitions["Talker"].graph_walks + assert "codec_chunk_clone" in partitions["Codec"].graph_walks + # Non-Base variants do not even declare the clone walks (config-driven). + assert "talker_prefill_clone" not in _variant_model("voice_design").get_graph_walk_graphs() + + pointers = { + name: [SimpleNamespace(name=name)] + for name in (*Qwen3TTSModel.PREFILL_INPUTS, "audio_inputs", "ref_frames") + } + talker = model.get_initial_forward_pass_args( + "Talker", input_modalities=["audio", "text"], output_modalities=["audio"], + input_signals=pointers, + ) + assert talker.full_metadata.graph_walk == "talker_prefill_clone" + routes = {(edge.name, edge.next_node) for edge in talker.inputs} + assert ("audio_inputs", "RefEncoder") in routes and ("prompt_layout", "RefEncoder") in routes + assert ("text_inputs", "Talker") in routes + codec = model.get_initial_forward_pass_args( + "Codec", input_modalities=["audio", "text"], output_modalities=["audio"], + input_signals=pointers, + ) + assert codec.full_metadata.graph_walk == "codec_chunk_clone" + rearmed = model.get_partition_forward_pass_args( + "Codec", codec.full_metadata, persist_signals={"ref_frames": pointers["ref_frames"]}, + ) + assert rearmed.full_metadata.graph_walk == "codec_chunk_clone" + assert [edge.name for edge in rearmed.inputs] == ["ref_frames"] + assert rearmed.inputs[0].tensor_info == pointers["ref_frames"] + + # The Base deployment maps the extra node and walks. + deployment = yaml.safe_load((CONFIG_PATH.parent / "qwen3tts_base.yaml").read_text(encoding="utf-8")) + groups = {name: group for group in deployment["node_groups"] for name in group["node_names"]} + assert "RefEncoder" in groups + assert "talker_prefill_clone" in groups["RefEncoder"]["graph_walks"] + assert "codec_chunk_clone" in groups["Codec"]["graph_walks"] + by_walk = {} + for worker_graph in model.get_worker_graphs(str(CONFIG_PATH.parent / "qwen3tts_base.yaml")): + by_walk.setdefault(next(iter(worker_graph.graph_walks)), worker_graph) + assert {"talker_prefill_clone", "codec_chunk_clone"} <= set(by_walk) def test_qwen3_tts_config_rejects_unknown_variant(): @@ -693,6 +798,78 @@ def test_qwen3_tts_talker_prefill_prepends_instruction_without_speaker(): ) +def test_qwen3_tts_talker_builds_in_context_clone_prefill(): + config = _tiny_model_config() + submodule = TalkerSubmodule( + Qwen3TTSTalkerModel(config), Qwen3TTSCodePredictor(config), config + ) + submodule.CHATML_ASSISTANT_PREFIX_TOKEN_IDS = (1, 2, 3) + submodule.CHATML_ASSISTANT_SUFFIX_TOKEN_IDS = (8, 9, 10, 11, 12) + assistant = torch.arange(1, 13) # 4 text tokens + reference = torch.tensor([1, 2, 3, 50, 51, 8, 9]) # 2 transcript tokens + <|im_end|>\n + text_ids = torch.cat([assistant, reference]) + ref_codes = torch.randint(0, 32, (5, config.talker.num_code_groups)) + speaker_embed = torch.randn(config.talker.hidden_size) + + def build(layout): + return submodule._build_prefill( + request_id="clone", text_ids=text_ids, prompt_layout=torch.tensor(layout), + speaker_id=-1, language_id=-1, speaker_embed=speaker_embed, ref_codes=ref_codes, + ) + + # Streaming text (Base default): text (2 ref + 4 + eos = 7) longer than + # codec (bos + 5 frames = 6) -> 6 in the prefill, 1 trailing. + embeds = build([0, 4, 1, 7, 5]) + # role(3) + [nothink, think_bos, think_eos, xvec, pad](5) + icl(6) + assert embeds.shape == (3 + 5 + 6, 16) + state = submodule.request_state("clone") + assert state["trailing_text_hidden"].shape == (1, 16) + assert torch.equal(state["reference_frames"], ref_codes) + # The x-vector occupies the speaker slot right after the three think tags. + tags = embeds[3:3 + 5] + assert torch.allclose(tags[3], speaker_embed.to(tags.dtype) + state["tts_pad_embed"], atol=1e-5) + + # Non-streaming: (text + eos) over codec pads, then (bos + frames) + tts pad. + embeds = build([0, 4, 0, 7, 5]) + assert embeds.shape == (3 + 5 + 7 + 6, 16) + assert submodule.request_state("clone")["trailing_text_hidden"].shape == (0, 16) + + # Text shorter than the codec span: padded with TTS PAD, nothing trails. + short = torch.cat([torch.tensor([1, 2, 3, 4, 8, 9, 10, 11, 12]), reference]) + embeds = submodule._build_prefill( + request_id="clone", text_ids=short, prompt_layout=torch.tensor([0, 1, 1, 7, 5]), + speaker_id=-1, language_id=-1, speaker_embed=speaker_embed, ref_codes=ref_codes, + ) + assert embeds.shape == (3 + 5 + 6, 16) + assert submodule.request_state("clone")["trailing_text_hidden"].shape == (0, 16) + + # x-vector only: standard layout with the x-vector in the speaker slot. + embeds = submodule._build_prefill( + request_id="xvec", text_ids=assistant, prompt_layout=torch.tensor([0, 4, 1, 0, 0]), + speaker_id=-1, language_id=-1, speaker_embed=speaker_embed, ref_codes=None, + ) + assert embeds.shape == (3 + 5 + 1, 16) + assert "reference_frames" not in submodule.request_state("xvec") + + with pytest.raises(ValueError, match="reference frames"): + build([0, 4, 1, 7, 9]) + + +def test_qwen3_tts_clone_prefill_streams_reference_frames_first(): + config = _tiny_model_config() + submodule = TalkerSubmodule( + Qwen3TTSTalkerModel(config), Qwen3TTSCodePredictor(config), config + ) + reference = torch.arange(8).view(2, 4) + submodule.request_state("clone").add("reference_frames", reference) + frame = torch.tensor([9, 9, 9, 9]) + items = submodule._codec_stream_items("talker_prefill_clone", "clone", frame) + assert [item.tolist() for item in items] == [[0, 1, 2, 3], [4, 5, 6, 7], [9, 9, 9, 9]] + # Only the clone prefill leads with the reference; decode never does. + assert submodule._codec_stream_items("talker_decode", "clone", frame) == [frame] + assert submodule._codec_stream_items("talker_prefill_clone", "other", frame) == [frame] + + def test_qwen3_tts_talker_rejects_changed_chatml_layout(): config = _tiny_model_config() submodule = TalkerSubmodule( @@ -1089,6 +1266,30 @@ def test_qwen3_tts_codec_trims_overlap_after_first_chunk(): assert second["audio_chunk"][0].tolist() == list(range(8, 20)) +def test_qwen3_tts_codec_trims_reference_audio_from_clone_streams(): + config = _tiny_model_config() # upsample 4 samples per frame, chunk 3, left context 2 + submodule = CodecSubmodule(_FakeCodecDecoder(4), config) + codes = torch.ones(3, 4, dtype=torch.long) + submodule.prepare_inputs( + "codec_chunk_clone", SimpleNamespace(request_id="clone"), + {"codec_tokens": [codes], "ref_frames": [torch.tensor([4])]}, + ) + state = submodule.request_state("clone") + assert state["skip_samples"] == 16 + + # First chunk: 3 frames = 12 samples, all reference -> nothing emitted. + first = {"audio_chunk": [torch.arange(20)]} + submodule.postprocess("clone", None, first) + assert first["audio_chunk"][0].numel() == 0 + assert state["skip_samples"] == 4 + # Second chunk: 2 context + 3 new frames; 4 more samples belong to the reference. + state.add("latest_codec_frames", 5) + second = {"audio_chunk": [torch.arange(20)]} + submodule.postprocess("clone", None, second) + assert second["audio_chunk"][0].tolist() == list(range(12, 20)) + assert state["skip_samples"] == 0 + + def test_qwen3_tts_codec_filters_eos_and_pads_to_capture_shape(): config = _tiny_model_config() submodule = CodecSubmodule(_FakeCodecDecoder(4), config) @@ -1180,3 +1381,63 @@ def test_qwen3_tts_codec_batches_and_declares_cuda_graphs(): oversized = model_inputs * 5 assert len(oversized) == 10 assert not submodule.can_batch(batch, oversized) + + +class _FakeCodecEncoder(torch.nn.Module): + """Stands in for the Mimi encoder: deterministic codes, one frame per 4 samples.""" + + def __init__(self, num_quantizers: int): + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros(())) + self.num_quantizers = num_quantizers + self.calls = 0 + + def encode(self, input_values, return_dict=True): + del return_dict + self.calls += 1 + frames = input_values.shape[-1] // 4 + 2 # the real encoder pads a little + codes = torch.arange(frames).repeat(self.num_quantizers + 1, 1).unsqueeze(0) + return SimpleNamespace(audio_codes=codes) + + +def test_qwen3_tts_ref_encoder_emits_xvector_and_reference_frames(): + from mstar.model.qwen3_tts.components.speaker_encoder import ( + Qwen3TTSMelFrontEnd, + Qwen3TTSSpeakerEncoder, + ) + from mstar.model.qwen3_tts.config import Qwen3TTSSpeakerEncoderConfig + from mstar.model.qwen3_tts.submodules import RefEncoderSubmodule + + config = _tiny_model_config() + speaker_config = Qwen3TTSSpeakerEncoderConfig( + enc_dim=config.talker.hidden_size, enc_channels=(16, 16, 16, 16, 48), + enc_se_channels=8, enc_attention_channels=8, + ) + encoder = _FakeCodecEncoder(config.codec.num_quantizers) + submodule = RefEncoderSubmodule( + Qwen3TTSSpeakerEncoder(speaker_config), Qwen3TTSMelFrontEnd(speaker_config), encoder, config, + ) + clip = torch.randn(2, 4000) * 0.1 # stereo, averaged to mono + prepared = submodule.prepare_inputs( + "talker_prefill_clone", SimpleNamespace(request_id="clone"), + {"audio_inputs": [clip], "prompt_layout": [torch.tensor([0, 2, 1, 5, 7])]}, + ) + assert prepared.tensor_inputs["waveform"].shape == (4000,) + assert prepared.kwargs == {"ref_frames": 7} + engine_inputs = ModelInputsFromEngine(request_ids=["clone"], per_request_info={}) + out = submodule.forward("talker_prefill_clone", engine_inputs, **submodule.preprocess( + "talker_prefill_clone", engine_inputs, [prepared])) + assert out["speaker_embed"][0].shape == (config.talker.hidden_size,) + assert out["ref_codes"][0].shape == (7, config.codec.num_quantizers) + assert out["ref_codes"][0][:, 0].tolist() == list(range(7)) + assert encoder.calls == 1 + + # x-vector only: the codec encoder is skipped and a placeholder frame rides the edge. + prepared = submodule.prepare_inputs( + "talker_prefill_clone", SimpleNamespace(request_id="xvec"), + {"audio_inputs": [clip[0]], "prompt_layout": [torch.tensor([0, 2, 1, 0, 0])]}, + ) + out = submodule.forward("talker_prefill_clone", engine_inputs, **submodule.preprocess( + "talker_prefill_clone", engine_inputs, [prepared])) + assert out["ref_codes"][0].shape == (1, config.codec.num_quantizers) + assert encoder.calls == 1 From 7cc58722396238ed17090515d694eaebe5516b4c Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 029/110] test: Qwen3-TTS reference audio through the speech adapter --- test/modular/test_openai_adapters.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/test/modular/test_openai_adapters.py b/test/modular/test_openai_adapters.py index ff18a8421..3e3ee617c 100644 --- a/test/modular/test_openai_adapters.py +++ b/test/modular/test_openai_adapters.py @@ -93,6 +93,28 @@ def test_qwen3_tts_speech_maps_voice_instructions_and_extra_body(tmp_path): assert isinstance(adapters.get_adapter(key), adapters.Qwen3TTSAdapter) +def test_qwen3_tts_speech_reference_audio_becomes_audio_input(tmp_path): + wav = base64.b64encode(b"RIFF....WAVEfmt ").decode() + req = SpeechRequest( + input="clone me", ref_audio=f"data:audio/wav;base64,{wav}", ref_text="reference words", + x_vector_only_mode=False, + ) + sa = adapters.Qwen3TTSAdapter().speech_to_request(req, tmp_path) + assert sa.input_modalities == ["audio", "text"] and sa.output_modalities == ["audio"] + (path,) = sa.file_paths["audio"] + assert Path(path).is_file() and Path(path).read_bytes() == b"RIFF....WAVEfmt " + assert "ref_audio" not in sa.model_kwargs + assert sa.model_kwargs["ref_text"] == "reference words" + assert sa.model_kwargs["x_vector_only_mode"] is False + + # Bare base64 (vLLM-Omni style) is accepted too. + sa = adapters.Qwen3TTSAdapter().speech_to_request(SpeechRequest(input="x", ref_audio=wav), tmp_path) + assert sa.input_modalities == ["audio", "text"] and sa.file_paths["audio"] + # No reference -> plain text request. + sa = adapters.Qwen3TTSAdapter().speech_to_request(SpeechRequest(input="x", voice="vivian"), tmp_path) + assert sa.input_modalities == ["text"] and sa.file_paths is None + + def test_chat_and_image_honor_seed(tmp_path): chat = ChatCompletionRequest(model="bagel", messages=[{"role": "user", "content": "x"}], seed=7) assert adapters.BagelAdapter().chat_to_request(chat, tmp_path).model_kwargs["seed"] == 7 From c0bdd17ee6b0c88febacc0e685a5dcbd3db17d30 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:53 -0700 Subject: [PATCH 030/110] test: float32 reference codec in the Qwen3-TTS parity harness --- test/qwen3-tts/parity_qwen3_tts.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/test/qwen3-tts/parity_qwen3_tts.py b/test/qwen3-tts/parity_qwen3_tts.py index 3c6a2c49c..9c1ddbf71 100644 --- a/test/qwen3-tts/parity_qwen3_tts.py +++ b/test/qwen3-tts/parity_qwen3_tts.py @@ -332,8 +332,15 @@ def decode_audio(codec, codes: torch.Tensor) -> torch.Tensor: @torch.no_grad() -def reference_decode_audio(ref, codes: torch.Tensor) -> torch.Tensor: - wavs, _ = ref.model.speech_tokenizer.decode([{"audio_codes": codes}]) +def reference_decode_audio(snapshot: str, device: str, codes: torch.Tensor) -> torch.Tensor: + """Reference codec in float32 (M* runs its codec in float32; the reference + wrapper would otherwise inherit the Talker's bf16).""" + from qwen_tts import Qwen3TTSTokenizer + + tokenizer = Qwen3TTSTokenizer.from_pretrained( + str(Path(snapshot) / "speech_tokenizer"), device_map=device, dtype=torch.float32, + ) + wavs, _ = tokenizer.decode([{"audio_codes": codes}]) return torch.as_tensor(wavs[0]).float() @@ -412,7 +419,7 @@ def main(argv: list[str] | None = None) -> None: codes_report = compare_codes(ours_codes, ref_codes) n = codes_report["frames_compared"] audio_ref_codes_mstar = decode_audio(codec, ref_codes[:n]) - audio_ref_codes_ref = reference_decode_audio(ref, ref_codes[:n]).to(audio_ref_codes_mstar.device) + audio_ref_codes_ref = reference_decode_audio(snapshot, args.device, ref_codes[:n]).to(audio_ref_codes_mstar.device) m = min(audio_ref_codes_mstar.numel(), audio_ref_codes_ref.numel()) codec_report = { "name": "codec_same_codes", From 82cc57657a4d3ceabbfe5931954e713f3a802375 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 031/110] docs: Qwen3-TTS 1.7B variants, voice cloning and benchmark client --- docs/models.rst | 51 ++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 42 insertions(+), 9 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index 9cb3b131a..2523378b8 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -44,6 +44,15 @@ Registry keys live in ``mstar/model/registry.py`` (``MODEL_REGISTRY`` / ``HF_MOD * - ``qwen3_tts`` - ``Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`` - Streaming text-to-speech with built-in speakers: Talker + 12 Hz speech codec. + * - ``qwen3_tts_1p7b`` + - ``Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice`` + - 1.7B CustomVoice: built-in speakers plus style/emotion ``instruct`` control. + * - ``qwen3_tts_voicedesign`` + - ``Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign`` + - 1.7B VoiceDesign: the voice is described by the request's ``instruct`` text. + * - ``qwen3_tts_base`` + - ``Qwen/Qwen3-TTS-12Hz-1.7B-Base`` + - 1.7B Base: zero-shot voice clone from one reference clip (x-vector, optional in-context transcript). * - ``vjepa2`` - ``facebook/vjepa2-vitl-fpc64-256`` - V-JEPA 2 video encoder + masked predictor. @@ -90,11 +99,23 @@ Qwen3-TTS notes --------------- - Install the model-specific dependencies with ``pip install -e '.[qwen3_tts]'`` - and launch the default single-GPU deployment with - ``mstar serve qwen3_tts --gpus 0``. -- The first integration supports the CustomVoice checkpoint and text-to-audio - requests. ``voice`` selects one of the checkpoint's built-in speakers and - ``language`` defaults to automatic detection. + and launch a single-GPU deployment with ``mstar serve --gpus 0`` where + ```` is one of ``qwen3_tts`` (0.6B CustomVoice), ``qwen3_tts_1p7b``, + ``qwen3_tts_voicedesign`` or ``qwen3_tts_base``. One model class serves every + 12 Hz checkpoint; the variant (speakers, instruction support, reference + audio) is read from the checkpoint's ``config.json``. +- Requests: ``voice`` selects a built-in speaker (CustomVoice), ``language`` + defaults to automatic detection, ``instruct`` (OpenAI: ``instructions``) + carries a style instruction (1.7B CustomVoice) or the voice description + (VoiceDesign, required). Base clones a voice from one reference clip: on + ``/v1/audio/speech`` pass ``ref_audio`` (data URL, URL, path or base64) plus + ``ref_text`` (its transcript) or ``x_vector_only_mode: true``; with the SDK, + ``client.tts(text, reference_audio="ref.wav", ref_text="...")``. The clip is + used for that request only; named, persisted voices arrive with the shared + voice registry. +- Text layout follows the reference defaults: CustomVoice and VoiceDesign put + the whole text in the prefill; Base feeds it one token per frame. Override + per request with ``non_streaming_mode``. - Codec CUDA graphs are captured through batch size 8. The upstream decoder's batch-16 capture can exhaust an H100 after Talker weights and CodePredictor graphs are resident; larger Codec batches therefore use the scheduler's safe @@ -105,10 +126,13 @@ Qwen3-TTS notes carried as a graph tensor input so replay does not consult capture-slot dummy request state. Residual ``subtalker_*`` sampling is per-request through the ``code_predictor`` aux sampler, so custom values neither block batching nor - fall off the graph. -- The 12 Hz decoder does not require the system SoX executable. M* imports only - the exact upstream decoder modules, avoiding qwen-tts's unrelated 25 Hz SoX - probe during worker startup. + fall off the graph. On the 1.7B checkpoints the CodePredictor projects the + Talker-width inputs through ``small_to_mtp_projection`` before its depth loop. +- Weight loading checks coverage in both directions: a parameter the checkpoint + does not fill, or a checkpoint tensor the port does not load, fails startup. +- The 12 Hz codec does not require the system SoX executable. M* imports only + the exact upstream speech-tokenizer modules, avoiding qwen-tts's unrelated + 25 Hz SoX probe during worker startup. For throughput/latency validation, run the native serving benchmark with the Qwen3-TTS model metadata rather than the Orpheus compatibility entry:: @@ -131,6 +155,15 @@ fixed-length decode throughput rather than end-user latency. The first process-local request can include eager FlashInfer kernel JIT, so keep the warmup requests enabled when reporting steady-state latency. +For cross-engine comparisons (M*, vLLM-Omni, SGLang-Omni) use the shared +streaming ``/v1/audio/speech`` client, which measures time-to-first-audio, RTF +and audio-seconds per second with the same request for every engine, and score +the saved WAVs with ``benchmark/tts_wer.py``:: + + python -m benchmark.tts_speech_bench --engine mstar --url http://127.0.0.1:8000 \ + --model qwen3_tts_1p7b --sentences sentences_200.txt --voice vivian \ + --language English --concurrency 8 --repeats 3 --out results/mstar_c8.json + Cosmos3 environment requirements -------------------------------- From f327647823448657717c1d0496e344e5df22674b Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 032/110] examples: Qwen3-TTS SDK usage for the three variants --- examples/sdk_tts_qwen3.py | 40 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 examples/sdk_tts_qwen3.py diff --git a/examples/sdk_tts_qwen3.py b/examples/sdk_tts_qwen3.py new file mode 100644 index 000000000..9122aab13 --- /dev/null +++ b/examples/sdk_tts_qwen3.py @@ -0,0 +1,40 @@ +"""Qwen3-TTS through the Python SDK: built-in voices, voice design and voice cloning. + +Start one server per checkpoint, e.g.: + mstar serve qwen3_tts_1p7b # built-in speakers + style instructions + mstar serve qwen3_tts_voicedesign # voice described by an instruction + mstar serve qwen3_tts_base # voice cloned from a reference clip +""" + +import sys + +from mstar import MStarClient + +client = MStarClient("http://localhost:8000") +variant = sys.argv[1] if len(sys.argv) > 1 else "custom_voice" + +if variant == "custom_voice": + audio = client.tts( + "Hello from M star! It is a beautiful day for a walk.", + voice="vivian", + language="English", + instruct="Speak with great enthusiasm.", # 1.7B only; drop it for the 0.6B checkpoint + ) +elif variant == "voice_design": + audio = client.tts( + "Hello from M star! It is a beautiful day for a walk.", + language="English", + instruct="A calm, warm adult female voice with a slight British accent.", + ) +elif variant == "voice_clone": + audio = client.tts( + "Hello from M star! It is a beautiful day for a walk.", + language="English", + reference_audio="reference.wav", # 3-10 s clip of the target speaker + ref_text="Transcript of the reference clip.", # or x_vector_only_mode=True + ) +else: + sys.exit(f"unknown variant {variant!r}: custom_voice | voice_design | voice_clone") + +audio.to_wav("out.wav") +print(f"wrote out.wav — {len(audio)} samples @ {audio.sample_rate} Hz") From 6e65cb5d392e32b7d7cf447d15a4e46c56f94c19 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 033/110] test: Qwen3-TTS rejects reference audio on non-Base checkpoints --- test/modular/test_qwen3_tts_model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index 10b5bece8..01fe03ca0 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -541,7 +541,8 @@ def test_qwen3_tts_validates_speaker_dialect_after_language_override(): ("prompt", "inputs", "outputs", "kwargs", "message"), [ ("", ["text"], ["audio"], {}, "non-empty"), - ("hello", ["audio"], ["audio"], {}, "text input only"), + ("hello", ["audio"], ["audio"], {}, "does not take reference audio"), + ("hello", ["video", "text"], ["audio"], {}, "text input only"), ("hello", ["text"], ["text"], {}, "audio output only"), ("hello", ["text"], ["audio", "text"], {}, "audio output only"), ("hello", ["text"], ["audio"], {"voice": "unknown"}, "speaker"), From f60d139cd4abe5cae838791dd54679386fe87a22 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 034/110] streaming: scheduled left-context chunk policy for early first audio --- mstar/streaming/chunk_policy.py | 64 +++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/mstar/streaming/chunk_policy.py b/mstar/streaming/chunk_policy.py index 6c527fe46..3e945d771 100644 --- a/mstar/streaming/chunk_policy.py +++ b/mstar/streaming/chunk_policy.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from collections.abc import Sequence class ChunkPolicy(ABC): @@ -150,3 +151,66 @@ def window_size(self) -> int: def continue_after_producer_done(self) -> bool: return self._continue_after_done + + +class ScheduledLeftContextChunkPolicy(ChunkPolicy): + """Left-context chunking whose chunk sizes follow a ramp. + + Streaming vocoders want the first audio out as early as possible and + larger chunks once the stream is running. Chunk ``k`` delivers + ``schedule[k]`` new items (``chunk`` once the schedule is exhausted) with + up to ``left_context`` already-delivered items in front of them, so a + causal decoder can warm up on frames it has processed before. Unlike + ``LeftContextChunkPolicy`` the first chunk may be smaller than the + context: the context is whatever has been delivered so far, capped. + + Example (Qwen3-TTS, 12 Hz frames): ``schedule=(4, 8, 16)``, ``chunk=25``, + ``left_context=25`` pops windows of 4, 4+8, 12+16, 25+25, 25+25, ... + items and the first audio leaves after four frames instead of 300. + + The consumer learns how many leading items of a window are context from + ``StreamChunk.context_items`` (the worker passes it along as + ``step_metadata["stream_chunks"][edge]["context_items"]``), so it can trim + the duplicated output without re-deriving this schedule. + """ + + def __init__(self, schedule: Sequence[int], chunk: int, left_context: int): + super().__init__() + if chunk <= 0 or any(size <= 0 for size in schedule) or left_context < 0: + raise ValueError("chunk sizes must be positive and left_context non-negative") + self._schedule = tuple(int(size) for size in schedule) + self._chunk = int(chunk) + self._left_context = int(left_context) + self._chunks_popped = 0 + self._delivered = 0 # new items handed to the consumer so far + + def _new_items(self) -> int: + if self._chunks_popped < len(self._schedule): + return self._schedule[self._chunks_popped] + return self._chunk + + def _context(self) -> int: + return min(self._left_context, self._delivered) + + def is_ready(self, buffer_len: int) -> bool: + return buffer_len >= self.window_size() + + def window_size(self) -> int: + return self._context() + self._new_items() + + def next_chunk_size(self, buffer_len: int) -> int: + # The buffer pointer sits ``context`` items before the first new item. + # After this pop it must sit ``next context`` items before the next + # chunk's first new item. + delivered_after = self._delivered + self._new_items() + next_context = min(self._left_context, delivered_after) + return (delivered_after - next_context) - (self._delivered - self._context()) + + def register_chunk(self, chunk_size: int): + super().register_chunk(chunk_size) + # A regular pop delivers exactly this chunk's new items. The only + # other caller is the terminal flush (producer done, window not + # full), after which no data-carrying chunk follows, so treating it + # the same keeps the bookkeeping trivially correct where it matters. + self._delivered += self._new_items() + self._chunks_popped += 1 From c5a123e6ff92739015ea1d7b421ae344046803ec Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 035/110] streaming: report each chunk's already-delivered context items --- mstar/streaming/stream_buffer.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/mstar/streaming/stream_buffer.py b/mstar/streaming/stream_buffer.py index d2ce728c0..3375f7f6b 100644 --- a/mstar/streaming/stream_buffer.py +++ b/mstar/streaming/stream_buffer.py @@ -14,6 +14,9 @@ class StreamChunk: chunk_index: int start_offset: int = 0 # global position of the first item in this chunk is_final: bool = False + # leading items of this chunk that an earlier chunk already delivered + # (sliding-window overlap / left context); the consumer trims their output + context_items: int = 0 @dataclass @@ -39,6 +42,9 @@ class StreamBuffer: _id_to_tensor: dict = field(default_factory=dict) _consumed: int = 0 _chunks_popped: int = 0 + # global position just past the last item ever handed out; everything + # before it in a later window is context, not new data + _delivered_end: int = 0 producer_done: bool = False # Set once a chunk has been popped with ``is_final=True`` (the terminal # flush). Guards the empty-buffer final flush below so it fires exactly @@ -138,7 +144,9 @@ def pop_chunk(self) -> StreamChunk: chunk_index=self._chunks_popped, start_offset=offset, is_final=is_final, + context_items=min(max(self._delivered_end - offset, 0), len(items)), ) + self._delivered_end = max(self._delivered_end, offset + len(items)) self._chunks_popped += 1 return chunk From cd61dfefb213050aa303ab5d711f53db9c037e32 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 036/110] graph: carry stream chunk offset and context on synthetic edges --- mstar/graph/base.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/mstar/graph/base.py b/mstar/graph/base.py index 17c38de34..7b9524848 100644 --- a/mstar/graph/base.py +++ b/mstar/graph/base.py @@ -74,6 +74,11 @@ class GraphEdge: # set on a synthetic streaming-input edge carrying the final chunk, so the # consuming pass (not the earlier ingest) reports the partition done _final_stream_chunk: bool = field(default=False) + # set on a synthetic streaming-input edge: where the chunk starts in the + # stream and how many of its leading items were delivered before (context); + # None on every other edge + _stream_chunk_offset: int | None = field(default=None) + _stream_chunk_context: int | None = field(default=None) # Set for sharded configurations _total_fanin: int = 1 @@ -90,6 +95,8 @@ def clone(self): output_modality=self.output_modality, _persist_for_loop=self._persist_for_loop, _final_stream_chunk=self._final_stream_chunk, + _stream_chunk_offset=self._stream_chunk_offset, + _stream_chunk_context=self._stream_chunk_context, ) From 3534f50c5cb7824b6e34e4db70569a745c3e940a Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 037/110] worker: expose stream chunk geometry to consumers via step_metadata --- mstar/worker/worker.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 8a9ad0523..0dd09f4f4 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -789,6 +789,8 @@ def _pop_streaming_edge( name=edge_name, tensor_info=[], _final_stream_chunk=chunk.is_final, + _stream_chunk_offset=chunk.start_offset, + _stream_chunk_context=chunk.context_items, ) else: # Normal chunk — store tensor and create edge with tensor_info. @@ -807,6 +809,8 @@ def _pop_streaming_edge( name=edge_name, tensor_info=tensor_infos.get(edge_name, []), _final_stream_chunk=chunk.is_final, + _stream_chunk_offset=chunk.start_offset, + _stream_chunk_context=chunk.context_items, ) return synthetic_edge @@ -958,6 +962,7 @@ def _build_executing_batch(self, batch: ScheduledBatch) -> ExecutingBatch: for request_id, node in batch.node_objects.items(): tensors = {} + stream_chunks = {} ready_inputs = node.ready_signals.ready_inputs for input_name, edge in ready_inputs.items(): tensors[input_name] = [ @@ -967,8 +972,20 @@ def _build_executing_batch(self, batch: ScheduledBatch) -> ExecutingBatch: ] if edge._final_stream_chunk: final_stream_rids.add(request_id) + if edge._stream_chunk_context is not None: + stream_chunks[input_name] = { + "start_offset": edge._stream_chunk_offset, + "context_items": edge._stream_chunk_context, + "is_final": edge._final_stream_chunk, + } per_request_inputs[request_id] = tensors - per_request_info[request_id] = self.worker_graphs_manager.get_fwd_info(request_id, batch_partition) + fwd_info = self.worker_graphs_manager.get_fwd_info(request_id, batch_partition) + if stream_chunks: + # Where each streamed input sits in its stream and how many of + # its leading items are repeated context, for the consumer's + # ``prepare_inputs`` (a vocoder trims that context's audio). + fwd_info.step_metadata["stream_chunks"] = stream_chunks + per_request_info[request_id] = fwd_info return self._make_executing_batch( node_name=batch.node_name, From 796651bb48cd2b4d0ed1ae598fea96461e1ebf17 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 038/110] test: scheduled chunk policy and chunk context geometry --- test/modular/test_stream_chunk_schedule.py | 126 +++++++++++++++++++++ 1 file changed, 126 insertions(+) create mode 100644 test/modular/test_stream_chunk_schedule.py diff --git a/test/modular/test_stream_chunk_schedule.py b/test/modular/test_stream_chunk_schedule.py new file mode 100644 index 000000000..c35d6f528 --- /dev/null +++ b/test/modular/test_stream_chunk_schedule.py @@ -0,0 +1,126 @@ +"""``ScheduledLeftContextChunkPolicy`` and the chunk geometry a StreamBuffer reports. + +A streaming vocoder decodes each popped window and must drop the audio of the +leading ``context_items`` (frames an earlier chunk already delivered). These +tests pin (a) the window / context sequence of the ramped policy, (b) that every +item reaches the consumer exactly once as new data whatever the ramp, and (c) +that ``StreamChunk.context_items`` is right for every policy in the tree, so a +consumer can rely on it instead of re-deriving the policy's schedule. +""" + +import pytest +import torch + +from mstar.graph.base import GraphEdge +from mstar.streaming.chunk_policy import ( + FixedChunkPolicy, + LeftContextChunkPolicy, + ScheduledLeftContextChunkPolicy, + SlidingWindowChunkPolicy, +) +from mstar.streaming.stream_buffer import StreamBuffer + + +def _drive(policy, total_items, drain_before_done=True): + """Feed ``total_items`` one-row items; return (chunks, is_final flags).""" + buffer = StreamBuffer(request_id="r", edge_name="codec_tokens", from_partition="Talker", policy=policy) + chunks = [] + + def poll(): + for _ in range(total_items + 50): + if not buffer.has_chunk_ready(): + return + chunks.append(buffer.pop_chunk()) + raise AssertionError("has_chunk_ready never went False") + + for i in range(total_items): + buffer.pre_read_register(f"t{i}") + buffer.put(f"t{i}", torch.tensor([i])) + if drain_before_done: + poll() + buffer.signal_done() + poll() + return chunks + + +def _geometry(chunks): + """(window, context, start_offset) per data-carrying chunk.""" + out = [] + for chunk in chunks: + data = chunk.data["data"] + if data is None: + continue + items = data.reshape(-1).tolist() + out.append((len(items), chunk.context_items, chunk.start_offset, items)) + return out + + +def _new_items(chunks): + delivered = [] + for window, context, _, items in _geometry(chunks): + delivered.extend(items[context:window]) + return delivered + + +def test_scheduled_policy_ramps_windows_and_reports_context(): + policy = ScheduledLeftContextChunkPolicy(schedule=(4, 8, 16), chunk=25, left_context=25) + chunks = _drive(policy, total_items=120) + geometry = [(w, c, o) for w, c, o, _ in _geometry(chunks)] + # window = context + new: 4 | 4+8 | 12+16 | 25+25 | 25+25 | flush + assert geometry[:5] == [(4, 0, 0), (12, 4, 0), (28, 12, 0), (50, 25, 3), (50, 25, 28)] + # The terminal flush hands over the remaining 17 new items behind 25 of context. + assert geometry[-1] == (42, 25, 78) + assert _new_items(chunks) == list(range(120)) + assert sum(chunk.is_final for chunk in chunks) == 1 and chunks[-1].is_final + # Every window's leading context is exactly the items the previous chunk ended with. + geometry_all = _geometry(chunks) + for prev, curr in zip(geometry_all, geometry_all[1:], strict=False): + assert curr[3][:curr[1]] == prev[3][len(prev[3]) - curr[1]:] + + +@pytest.mark.parametrize("drain_before_done", [True, False]) +@pytest.mark.parametrize("total_items", [0, 1, 3, 4, 5, 12, 28, 53, 78, 100, 153]) +def test_scheduled_policy_delivers_everything_once_with_one_final(total_items, drain_before_done): + policy = ScheduledLeftContextChunkPolicy(schedule=(4, 8, 16), chunk=25, left_context=25) + chunks = _drive(policy, total_items, drain_before_done) + assert _new_items(chunks) == list(range(total_items)) + assert sum(chunk.is_final for chunk in chunks) == 1 and chunks[-1].is_final + + +def test_scheduled_policy_context_smaller_than_first_chunks(): + # Left context shorter than the ramp steps: context saturates at 2. + policy = ScheduledLeftContextChunkPolicy(schedule=(1, 3), chunk=5, left_context=2) + chunks = _drive(policy, total_items=14) + assert [(w, c, o) for w, c, o, _ in _geometry(chunks)] == [ + (1, 0, 0), (4, 1, 0), (7, 2, 2), (7, 2, 7), (2, 2, 12), + ] + assert _new_items(chunks) == list(range(14)) + + +def test_scheduled_policy_rejects_bad_sizes(): + with pytest.raises(ValueError): + ScheduledLeftContextChunkPolicy(schedule=(0,), chunk=5, left_context=1) + with pytest.raises(ValueError): + ScheduledLeftContextChunkPolicy(schedule=(), chunk=5, left_context=-1) + + +@pytest.mark.parametrize( + ("policy", "expected"), + [ + (FixedChunkPolicy(chunk_size=3), [(3, 0), (3, 0), (3, 0), (1, 0)]), + (SlidingWindowChunkPolicy(window=4, stride=2), [(4, 0), (4, 2), (4, 2), (4, 2), (2, 2)]), + (LeftContextChunkPolicy(chunk=4, left_context=1), [(4, 0), (5, 1), (3, 1)]), + ], +) +def test_existing_policies_report_context_items(policy, expected): + chunks = _drive(policy, total_items=10) + assert [(w, c) for w, c, _, _ in _geometry(chunks)] == expected + assert _new_items(chunks) == list(range(10)) + + +def test_graph_edge_clone_keeps_stream_chunk_geometry(): + edge = GraphEdge(next_node="Codec", name="codec_tokens", _final_stream_chunk=True, + _stream_chunk_offset=7, _stream_chunk_context=3) + clone = edge.clone() + assert (clone._stream_chunk_offset, clone._stream_chunk_context, clone._final_stream_chunk) == (7, 3, True) + assert GraphEdge(next_node="x", name="y")._stream_chunk_context is None From e1b8584ce36ef9a03c3bed50b62f3f49fa88b03d Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 039/110] qwen3_tts: ramped codec chunk schedule with 25-frame left context --- mstar/model/qwen3_tts/config.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/mstar/model/qwen3_tts/config.py b/mstar/model/qwen3_tts/config.py index ad260f443..e64c63e14 100644 --- a/mstar/model/qwen3_tts/config.py +++ b/mstar/model/qwen3_tts/config.py @@ -246,8 +246,13 @@ class Qwen3TTSCodecConfig: # turns reference audio into codec frames for voice cloning. encoder_config: dict[str, Any] = field(default_factory=dict) - # M* stream policy: 300 new 12 Hz frames with 25 frames of overlap. - chunk_frames: int = 300 + # M* stream policy: the codec pops a ramp of small chunks first (first + # audio after 4 frames = 320 ms of speech), then ``chunk_frames`` new + # frames per call, each preceded by up to ``left_context_frames`` already + # decoded frames so the causal decoder warms up (the reference's own + # ``chunked_decode`` uses 25 frames of left context). + chunk_schedule: tuple[int, ...] = (4, 8, 16) + chunk_frames: int = 25 left_context_frames: int = 25 @classmethod @@ -272,6 +277,19 @@ def from_dict(cls, data: dict[str, Any]) -> "Qwen3TTSCodecConfig": }) return cls(**values) + def codec_windows(self) -> list[int]: + """Distinct window sizes (context + new frames) the chunk schedule produces. + + These are the shapes the codec captures CUDA graphs for; a terminal + flush shorter than a window is padded up to the next one. + """ + windows = set() + delivered = 0 + for size in (*self.chunk_schedule, self.chunk_frames): + windows.add(min(self.left_context_frames, delivered) + size) + delivered += size + return sorted(windows) + def frames_for_samples(self, num_samples: int) -> int: """Codec frames the encoder emits for ``num_samples`` of input audio.""" return -(-int(num_samples) // self.encode_downsample_rate) @@ -285,6 +303,7 @@ def decoder_kwargs(self) -> dict[str, Any]: "encode_downsample_rate", "encoder_valid_num_quantizers", "encoder_config", + "chunk_schedule", "chunk_frames", "left_context_frames", } From f4940f7410eac30072c7fd3652647e508810e0e7 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 040/110] qwen3_tts: stream codec frames through the scheduled chunk policy --- mstar/model/qwen3_tts/qwen3_tts_model.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/mstar/model/qwen3_tts/qwen3_tts_model.py b/mstar/model/qwen3_tts/qwen3_tts_model.py index 9f4daf853..ddb33342c 100644 --- a/mstar/model/qwen3_tts/qwen3_tts_model.py +++ b/mstar/model/qwen3_tts/qwen3_tts_model.py @@ -16,7 +16,7 @@ Codec - stateless speech-tokenizer decoder producing PCM chunks Streaming topology: - Talker --[codec_tokens, LeftContextChunkPolicy(300, 25)]--> Codec + Talker --[codec_tokens, ScheduledLeftContextChunkPolicy((4, 8, 16), 25, 25)]--> Codec Request state machine: Talker: talker_prefill | talker_prefill_clone -> talker_decode loop -> done on EOS/token limit @@ -76,7 +76,7 @@ Qwen3TTSModelConfig, ) from mstar.model.submodule_base import NodeSubmodule -from mstar.streaming.chunk_policy import LeftContextChunkPolicy +from mstar.streaming.chunk_policy import ScheduledLeftContextChunkPolicy from mstar.streaming.topology import Connection, PartitionTopology, StreamingGraphEdge # --------------------------------------------------------------------------- @@ -472,12 +472,14 @@ def get_partitions(self) -> list[PartitionDefinition]: ] def get_partition_topology(self) -> PartitionTopology: - """Buffer codec frames with the decoder's required left context. - - The first Codec invocation receives up to ``chunk_frames`` new frames. - Later invocations prepend ``left_context_frames`` old frames to avoid - convolution boundary artifacts; ``CodecSubmodule.postprocess`` removes - the duplicated PCM prefix before emission. + """Buffer codec frames in a ramp of chunks with left context. + + The first Codec invocation runs after ``chunk_schedule[0]`` frames so + audio starts flowing early; chunks then grow to ``chunk_frames``. Every + window after the first is preceded by up to ``left_context_frames`` + already decoded frames to avoid boundary artifacts; + ``CodecSubmodule.postprocess`` removes the duplicated PCM prefix using + the context count the stream buffer reports for each window. """ codec = self.config.codec return PartitionTopology( @@ -487,7 +489,8 @@ def get_partition_topology(self) -> PartitionTopology: from_partition="Talker", to_partition="Codec", edge_name="codec_tokens", - chunk_policy_factory=lambda: LeftContextChunkPolicy( + chunk_policy_factory=lambda: ScheduledLeftContextChunkPolicy( + schedule=codec.chunk_schedule, chunk=codec.chunk_frames, left_context=codec.left_context_frames, ), @@ -831,6 +834,7 @@ def get_partition_forward_pass_args( "codec_left_context_frames": ( self.config.codec.left_context_frames ), + "codec_chunk_schedule": list(self.config.codec.chunk_schedule), }, ) raise ValueError(f"Unknown Qwen3-TTS partition: {partition_name!r}") From 616ab0f0a42a31183afc3458e77f3befa6382ea3 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 041/110] qwen3_tts: codec windows padded to captured buckets, context from the stream --- mstar/model/qwen3_tts/submodules.py | 154 ++++++++++++++++------------ 1 file changed, 91 insertions(+), 63 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index 6b4588bca..e88e07cc5 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -16,7 +16,7 @@ # piecewise graph covering that loop on eager paths like prefill. # 2. CodecSubmodule (STATELESS engine) # - Receives buffered codec frames from the Talker partition. -# - Pads variable final tails to fixed CUDA Graph capture shapes. +# - Pads each window of the chunk ramp to its CUDA Graph capture bucket. # - Runs the official speech-tokenizer decoder and trims overlap (and, for # voice clones, the reference frames) before emitting 24 kHz PCM. # 3. RefEncoderSubmodule (no resources; Base voice clone only) @@ -28,7 +28,7 @@ # -> postprocess -> check_stop (Talker only) # # Streaming topology: -# Talker --[codec_tokens, LeftContextChunkPolicy(300, 25)]--> Codec +# Talker --[codec_tokens, ScheduledLeftContextChunkPolicy((4, 8, 16), 25, 25)]--> Codec # --------------------------------------------------------------------------- from __future__ import annotations @@ -917,8 +917,15 @@ class CodecSubmodule(ARNodeSubmodule): The node runs on a stateless engine, but ``ARNodeInputs`` is reused as the typed container for fixed-length codec tensors. Per-request state stores - only how many non-padding frames arrived and whether a prior chunk was - emitted; the neural decoder itself has no cross-call state. + only the geometry of the latest window (how many frames are real, how + many of them are repeated context) and, for voice clones, how much + reference audio is still to be dropped; the neural decoder itself has no + cross-call state. + + Windows follow the model's ``ScheduledLeftContextChunkPolicy``: a ramp of + small chunks, then ``chunk_frames`` new frames behind ``left_context_frames`` + of context. Each distinct window size is a CUDA-graph bucket; a shorter + terminal flush is zero-padded up to the next bucket and trimmed after. """ # fp32, uncompiled: what ``get_stateless_flavor`` used to buy on the old @@ -926,21 +933,18 @@ class CodecSubmodule(ARNodeSubmodule): disable_torch_compile = True disable_autocast = True - # The official 114M-parameter decoder materializes large fixed-shape - # activations while CUDA graphs are captured. Capturing bs=16 exhausts an - # H100 once Talker weights and the CodePredictor graphs are resident, so - # keep the safe ceiling at 8 until the decoder is ported to M*'s - # lighter-weight codec components. - MAX_BATCH_SIZE = 8 - CAPTURE_BATCH_SIZES = [1, 2, 4, 8] + # Windows are at most chunk + left_context frames (50 by default, 4 s of + # audio), so the decoder's activations stay small enough to capture + # batches of 16 next to the Talker; ``can_batch`` keeps the ceiling. + MAX_BATCH_SIZE = 16 + CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16] def __init__(self, decoder: torch.nn.Module, config: Qwen3TTSModelConfig): super().__init__() self.decoder = decoder self.config = config - self.full_seq_len = ( - config.codec.chunk_frames + config.codec.left_context_frames - ) + self.windows = config.codec.codec_windows() + self.max_window = self.windows[-1] self.total_upsample = 1 for factor in ( *config.codec.upsample_rates, @@ -948,6 +952,21 @@ def __init__(self, decoder: torch.nn.Module, config: Qwen3TTSModelConfig): ): self.total_upsample *= factor + def _bucket(self, frames: int) -> int: + """Smallest captured window that holds ``frames`` (the terminal flush is shorter).""" + for window in self.windows: + if frames <= window: + return window + raise ValueError( + f"Codec chunk has {frames} frames, maximum is {self.max_window}" + ) + + @staticmethod + def _stream_chunk_meta(fwd_info: CurrentForwardPassInfo) -> dict[str, Any]: + """Geometry of this window from the stream buffer (offset, context, final).""" + step_metadata = getattr(fwd_info, "step_metadata", None) or {} + return step_metadata.get("stream_chunks", {}).get("codec_tokens", {}) + def prepare_inputs( self, graph_walk: str, @@ -955,11 +974,12 @@ def prepare_inputs( inputs: NameToTensorList, **kwargs: Any, ) -> ARNodeInputs: - """Remove EOS frames and pad one stream chunk to its capture shape. + """Remove EOS frames and pad one stream window to its capture bucket. Input arrives as ``[frames, code_groups]``. The official decoder wants - ``[quantizers, frames]``; every request is padded to ``chunk + context`` - so differently sized final tails can reuse the same CUDA Graph. + ``[quantizers, frames]``; every request is padded to the smallest + captured window that fits so the ramp's windows and the terminal tail + all replay CUDA graphs. """ del graph_walk, kwargs state = self.request_state(fwd_info.request_id) @@ -980,24 +1000,21 @@ def prepare_inputs( f"Expected codec tokens with shape (frames, groups), got {codes.shape}" ) # EOS belongs to Talker loop control and is not a valid codec codebook - # index for waveform reconstruction. + # index for waveform reconstruction. It is only ever the last frame of + # a stream, so the leading context count is unaffected. codes = codes[ codes[:, 0] != self.config.talker.codec_eos_token_id, :self.config.codec.num_quantizers, ] - original_frames = codes.shape[0] - if original_frames > self.full_seq_len: - raise ValueError( - f"Codec chunk has {original_frames} frames, maximum is " - f"{self.full_seq_len}" - ) - if original_frames < self.full_seq_len: - codes = torch.nn.functional.pad( - codes, - (0, 0, 0, self.full_seq_len - original_frames), - ) - self.request_state(fwd_info.request_id).add( - "latest_codec_frames", original_frames + frames = codes.shape[0] + context = int(self._stream_chunk_meta(fwd_info).get("context_items", 0)) + bucket = self._bucket(max(frames, 1)) + if frames < bucket: + codes = torch.nn.functional.pad(codes, (0, 0, 0, bucket - frames)) + state.add_all( + latest_codec_frames=frames, + latest_context_frames=min(context, frames), + codec_bucket=bucket, ) return ARNodeInputs( tensor_inputs={"codec_tokens": codes.t().contiguous()}, @@ -1009,7 +1026,7 @@ def preprocess( engine_inputs: ModelInputsFromEngine, inputs: list[ARNodeInputs], ) -> dict[str, torch.Tensor]: - """Stack equal fixed-shape codec chunks into one continuous batch.""" + """Stack equal fixed-shape codec windows into one continuous batch.""" del graph_walk, engine_inputs return { "codec_tokens": torch.stack([ @@ -1053,17 +1070,14 @@ def postprocess( outputs: dict[str, list[torch.Tensor]], **kwargs: Any, ) -> None: - """Remove padded tail and duplicated left-context PCM before emission.""" + """Drop padding, repeated-context audio and (clone) reference audio before emission.""" del request_info, kwargs if "audio_chunk" not in outputs: return state = self.request_state(request_id) frames = int(state.get("latest_codec_frames", 0)) - emitted = bool(state.get("codec_chunk_emitted", False)) - # The first chunk has no overlap. Later stream chunks include old codec - # frames at the front, whose decoded samples must not be emitted twice. - left_context = self.config.codec.left_context_frames if emitted else 0 - start = left_context * self.total_upsample + context = int(state.get("latest_context_frames", 0)) + start = context * self.total_upsample end = frames * self.total_upsample skip = int(state.get("skip_samples", 0)) if skip: @@ -1071,7 +1085,6 @@ def postprocess( start += dropped state.add("skip_samples", skip - dropped) outputs["audio_chunk"][0] = outputs["audio_chunk"][0][start:end] - state.add("codec_chunk_emitted", True) def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: """Batch codec requests only when their decoder input shapes match.""" @@ -1084,42 +1097,57 @@ def max_batch_size(self, graph_walk: str) -> int: del graph_walk return self.MAX_BATCH_SIZE + def cg_key_info( + self, + graph_walk: str, + per_request_info: Mapping[str, CurrentForwardPassInfo], + ) -> Any: + """The window bucket this batch was padded to (``can_batch`` keeps it uniform).""" + del graph_walk + buckets = { + self.request_state(request_id).get("codec_bucket") + for request_id in per_request_info + } + return buckets.pop() if len(buckets) == 1 else None + def get_cuda_graph_configs( self, device: torch.device, tp_world_size: int = 1 ) -> list[CudaGraphConfig]: - """Capture fixed-length Codec batches for all scheduler buckets.""" + """Capture every window of the chunk schedule for all batch buckets.""" del tp_world_size - return [BatchedCudaGraphConfig( - capture_graph_walk="codec_chunk", - single_request_inputs=ARNodeInputs( - # 1, not full_seq_len: batched buckets match on bs, and this - # keeps the intern seq_len from aliasing the fixed trailing dims. - input_seq_len=1, - tensor_inputs={ - "codec_tokens": torch.zeros( - self.config.codec.num_quantizers, - self.full_seq_len, - dtype=torch.long, - device=device, - ) - }, - ), - capture_batch_sizes=self.CAPTURE_BATCH_SIZES, - compile=False, - )] + return [ + BatchedCudaGraphConfig( + capture_graph_walk="codec_chunk", + replay_graph_walks=["codec_chunk", "codec_chunk_clone"], + single_request_inputs=ARNodeInputs( + # 1, not the window: batched buckets match on bs, and this + # keeps the intern seq_len from aliasing the trailing dims. + input_seq_len=1, + tensor_inputs={ + "codec_tokens": torch.zeros( + self.config.codec.num_quantizers, + window, + dtype=torch.long, + device=device, + ) + }, + ), + additional_key_info=window, + capture_batch_sizes=self.CAPTURE_BATCH_SIZES, + compile=False, + ) + for window in self.windows + ] def can_use_cuda_graphs( self, batch: ExecutingBatch, model_inputs: list[NodeInputs] ) -> bool: return ( - batch.graph_walk == "codec_chunk" + batch.graph_walk in ("codec_chunk", "codec_chunk_clone") and self.can_batch(batch, model_inputs) and all( item.tensor_inputs["codec_tokens"].shape - == ( - self.config.codec.num_quantizers, - self.full_seq_len, - ) + in {(self.config.codec.num_quantizers, window) for window in self.windows} for item in model_inputs ) and super().can_use_cuda_graphs(batch, model_inputs) From 9b36195966eb4fc220aa1ad3d458802c852ee7da Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 042/110] test: Qwen3-TTS codec ramp, bucket padding and context trimming --- test/modular/test_qwen3_tts_model.py | 117 ++++++++++++++++++++------- 1 file changed, 86 insertions(+), 31 deletions(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index 01fe03ca0..b58d5193c 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -36,7 +36,7 @@ from mstar.model.qwen3_tts.submodules import CodecSubmodule, TalkerSubmodule from mstar.model.registry import HF_MODELS, get_model_class from mstar.model.submodule_base import ARNodeInputs, ModelInputsFromEngine -from mstar.streaming.chunk_policy import LeftContextChunkPolicy +from mstar.streaming.chunk_policy import ScheduledLeftContextChunkPolicy from mstar.streaming.stream_buffer import StreamBuffer CONFIG_PATH = Path(__file__).resolve().parents[2] / "configs" / "qwen3tts.yaml" @@ -703,6 +703,7 @@ def _tiny_model_config() -> Qwen3TTSModelConfig: talker=talker, codec=Qwen3TTSCodecConfig( num_quantizers=4, + chunk_schedule=(1,), chunk_frames=3, left_context_frames=2, upsample_rates=(2,), @@ -1252,20 +1253,30 @@ def forward(self, codes): return torch.zeros(codes.shape[0], 1, length, dtype=torch.float32) -def test_qwen3_tts_codec_trims_overlap_after_first_chunk(): +def test_qwen3_tts_codec_trims_reported_context_audio(): config = _tiny_model_config() submodule = CodecSubmodule(_FakeCodecDecoder(4), config) + assert submodule.windows == [1, 4] and submodule.max_window == 4 state = submodule.request_state("request") - state.add("latest_codec_frames", 5) + # The stream buffer reports how many leading frames are repeated context; + # the first window has none, later ones up to left_context (2). + state.add_all(latest_codec_frames=5, latest_context_frames=0) first = {"audio_chunk": [torch.arange(20)]} submodule.postprocess("request", None, first) assert first["audio_chunk"][0].tolist() == list(range(20)) + state.add_all(latest_codec_frames=5, latest_context_frames=2) second = {"audio_chunk": [torch.arange(20)]} submodule.postprocess("request", None, second) assert second["audio_chunk"][0].tolist() == list(range(8, 20)) + # Padding frames of a bucket never reach the client. + state.add_all(latest_codec_frames=3, latest_context_frames=1) + padded = {"audio_chunk": [torch.arange(16)]} + submodule.postprocess("request", None, padded) + assert padded["audio_chunk"][0].tolist() == list(range(4, 12)) + def test_qwen3_tts_codec_trims_reference_audio_from_clone_streams(): config = _tiny_model_config() # upsample 4 samples per frame, chunk 3, left context 2 @@ -1279,12 +1290,13 @@ def test_qwen3_tts_codec_trims_reference_audio_from_clone_streams(): assert state["skip_samples"] == 16 # First chunk: 3 frames = 12 samples, all reference -> nothing emitted. + assert state["latest_codec_frames"] == 3 and state["latest_context_frames"] == 0 first = {"audio_chunk": [torch.arange(20)]} submodule.postprocess("clone", None, first) assert first["audio_chunk"][0].numel() == 0 assert state["skip_samples"] == 4 # Second chunk: 2 context + 3 new frames; 4 more samples belong to the reference. - state.add("latest_codec_frames", 5) + state.add_all(latest_codec_frames=5, latest_context_frames=2) second = {"audio_chunk": [torch.arange(20)]} submodule.postprocess("clone", None, second) assert second["audio_chunk"][0].tolist() == list(range(12, 20)) @@ -1301,50 +1313,80 @@ def test_qwen3_tts_codec_filters_eos_and_pads_to_capture_shape(): [5, 6, 7, 8], ]) - prepared = submodule.prepare_inputs( - "codec_chunk", - SimpleNamespace(request_id="request"), - {"codec_tokens": [codes]}, + fwd_info = SimpleNamespace( + request_id="request", + step_metadata={"stream_chunks": {"codec_tokens": { + "start_offset": 1, "context_items": 1, "is_final": False, + }}}, ) + prepared = submodule.prepare_inputs("codec_chunk", fwd_info, {"codec_tokens": [codes]}) + # Two real frames pad up to the smallest captured window (4), not the largest. packed = prepared.tensor_inputs["codec_tokens"] - assert packed.shape == (4, 5) + assert packed.shape == (4, 4) assert packed[:, :2].t().tolist() == [[1, 2, 3, 4], [5, 6, 7, 8]] assert packed[:, 2:].count_nonzero().item() == 0 - assert submodule.request_state("request")["latest_codec_frames"] == 2 + state = submodule.request_state("request") + assert state["latest_codec_frames"] == 2 + assert state["latest_context_frames"] == 1 + assert state["codec_bucket"] == 4 + + # A single frame lands in the first ramp bucket; too many frames is an error. + one = submodule.prepare_inputs( + "codec_chunk", SimpleNamespace(request_id="one"), {"codec_tokens": [codes[:1]]}, + ) + assert one.tensor_inputs["codec_tokens"].shape == (4, 1) + with pytest.raises(ValueError, match="maximum is 4"): + submodule.prepare_inputs( + "codec_chunk", SimpleNamespace(request_id="big"), + {"codec_tokens": [torch.ones(5, 4, dtype=torch.long)]}, + ) -def test_qwen3_tts_streaming_policy_flushes_only_new_tail_audio(): +def test_qwen3_tts_streaming_policy_ramps_and_flushes_only_new_tail_audio(): config = _tiny_model_config() stream = StreamBuffer( request_id="request", edge_name="codec_tokens", from_partition="Talker", - policy=LeftContextChunkPolicy( + policy=ScheduledLeftContextChunkPolicy( + schedule=config.codec.chunk_schedule, chunk=config.codec.chunk_frames, left_context=config.codec.left_context_frames, ), ) + chunks = [] for i in range(5): tensor_id = str(i) stream.pre_read_register(tensor_id) stream.put(tensor_id, torch.tensor([i])) - if i == 2: - first = stream.pop_chunk() - assert first.data["data"].flatten().tolist() == [0, 1, 2] + while stream.has_chunk_ready(): + chunks.append(stream.pop_chunk()) + # First audio after a single frame, then 1 context + 3 new frames. + assert [c.data["data"].flatten().tolist() for c in chunks] == [[0], [0, 1, 2, 3]] + assert [c.context_items for c in chunks] == [0, 1] stream.signal_done() assert stream.has_chunk_ready() tail = stream.pop_chunk() - assert tail.data["data"].flatten().tolist() == [1, 2, 3, 4] + assert tail.data["data"].flatten().tolist() == [2, 3, 4] + assert tail.context_items == 2 assert tail.is_final is True + # The codec trims exactly the context frames the buffer reported (tail: 1 new frame). codec = CodecSubmodule(_FakeCodecDecoder(4), config) - state = codec.request_state("request") - state.add_all(latest_codec_frames=4, codec_chunk_emitted=True) + fwd_info = SimpleNamespace( + request_id="request", + step_metadata={"stream_chunks": {"codec_tokens": { + "start_offset": tail.start_offset, "context_items": tail.context_items, "is_final": True, + }}}, + ) + tail_codes = tail.data["data"].view(3, 1).expand(3, 4) + prepared = codec.prepare_inputs("codec_chunk", fwd_info, {"codec_tokens": [tail_codes]}) + assert prepared.tensor_inputs["codec_tokens"].shape == (4, 4) # 3 frames padded to the 4-frame bucket outputs = {"audio_chunk": [torch.arange(16)]} codec.postprocess("request", None, outputs) - assert outputs["audio_chunk"][0].tolist() == list(range(8, 16)) + assert outputs["audio_chunk"][0].tolist() == list(range(8, 12)) def test_qwen3_tts_codec_batches_and_declares_cuda_graphs(): @@ -1352,7 +1394,7 @@ def test_qwen3_tts_codec_batches_and_declares_cuda_graphs(): submodule = CodecSubmodule(_FakeCodecDecoder(4), config) model_inputs = [ ARNodeInputs(tensor_inputs={ - "codec_tokens": torch.zeros(4, 5, dtype=torch.long) + "codec_tokens": torch.zeros(4, 4, dtype=torch.long) }) for _ in range(2) ] @@ -1370,17 +1412,30 @@ def test_qwen3_tts_codec_batches_and_declares_cuda_graphs(): ModelInputsFromEngine(request_ids=["a", "b"], per_request_info={}), model_inputs, ) - assert packed["codec_tokens"].shape == (2, 4, 5) - graph_config = submodule.get_cuda_graph_configs(torch.device("cpu"))[0] - assert graph_config.capture_graph_walk == "codec_chunk" - assert submodule.max_batch_size("codec_chunk") == 8 - assert graph_config.capture_batch_sizes == [1, 2, 4, 8] - assert graph_config.single_request_inputs.tensor_inputs[ - "codec_tokens" - ].shape == (4, 5) - - oversized = model_inputs * 5 - assert len(oversized) == 10 + assert packed["codec_tokens"].shape == (2, 4, 4) + # One capture per window of the chunk ramp, keyed by the window, replayed + # by both codec walks. + graph_configs = submodule.get_cuda_graph_configs(torch.device("cpu")) + assert [c.additional_key_info for c in graph_configs] == [1, 4] + for graph_config in graph_configs: + assert graph_config.capture_graph_walk == "codec_chunk" + assert set(graph_config.replay_graph_walks) == {"codec_chunk", "codec_chunk_clone"} + assert graph_config.capture_batch_sizes == [1, 2, 4, 8, 16] + assert graph_config.single_request_inputs.tensor_inputs["codec_tokens"].shape == ( + 4, graph_config.additional_key_info, + ) + assert submodule.max_batch_size("codec_chunk") == 16 + # The batch's capture key is the bucket its requests were padded to. + for rid in ("a", "b"): + submodule.request_state(rid).add("codec_bucket", 4) + assert submodule.cg_key_info("codec_chunk", {"a": None, "b": None}) == 4 + submodule.request_state("b").add("codec_bucket", 1) + assert submodule.cg_key_info("codec_chunk", {"a": None, "b": None}) is None + + mixed = model_inputs + [ARNodeInputs(tensor_inputs={"codec_tokens": torch.zeros(4, 1, dtype=torch.long)})] + assert not submodule.can_batch(batch, mixed) + oversized = model_inputs * 9 + assert len(oversized) == 18 assert not submodule.can_batch(batch, oversized) From f27ce586e300d4a39c47f598dad2003dd7642380 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:54 -0700 Subject: [PATCH 043/110] docs: Qwen3-TTS codec chunk ramp and capture buckets --- docs/models.rst | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index 2523378b8..3da48cb03 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -116,10 +116,12 @@ Qwen3-TTS notes - Text layout follows the reference defaults: CustomVoice and VoiceDesign put the whole text in the prefill; Base feeds it one token per frame. Override per request with ``non_streaming_mode``. -- Codec CUDA graphs are captured through batch size 8. The upstream decoder's - batch-16 capture can exhaust an H100 after Talker weights and CodePredictor - graphs are resident; larger Codec batches therefore use the scheduler's safe - ceiling. +- Audio streams in a ramp of codec chunks: the first window is decoded after 4 + frames (320 ms of speech), later windows grow to 25 new frames behind 25 + frames of already decoded left context (the reference's own + ``chunked_decode`` context). Each window size is a CUDA-graph bucket + captured for batch sizes 1 to 16; the stream buffer reports how many leading + frames of a window are repeated context, and the codec trims their audio. - Talker prefill remains eager because it runs once with variable sequence lengths. Decode always uses the whole-walk CUDA Graph, with the 15-step CodePredictor loop captured inside it; request-local EOS suppression is From ae116719deed3c421ca3c899809a95010f5a4cf1 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 044/110] qwen3_tts: reference encoder keeps its float32 front ends under autocast --- mstar/model/qwen3_tts/submodules.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index e88e07cc5..4070dd0b0 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -1165,12 +1165,14 @@ class RefEncoderSubmodule(NodeSubmodule): Runs once per request, before the clone prefill, and owns no resources. The x-vector comes from the ECAPA-TDNN encoder over a log-mel spectrogram of the 24 kHz clip; for in-context cloning the codec encoder also turns - the clip into ``ref_frames`` 16-group frames. The mel front end and the - codec encoder run in float32 regardless of the engine's autocast dtype; - the x-vector is produced in the encoder's own (Talker) dtype. + the clip into ``ref_frames`` 16-group frames. The node keeps the dtypes it + was built with (``disable_autocast``): the mel front end and the codec + encoder stay in float32, the speaker encoder runs in the Talker's dtype + as the reference does. """ disable_torch_compile = True + disable_autocast = True def __init__( self, From 7c227969247a4add71b2a23df52174f2899b07f2 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 045/110] test: voice-clone mode in the Qwen3-TTS parity harness --- test/qwen3-tts/parity_qwen3_tts.py | 82 +++++++++++++++++++++++++++--- 1 file changed, 76 insertions(+), 6 deletions(-) diff --git a/test/qwen3-tts/parity_qwen3_tts.py b/test/qwen3-tts/parity_qwen3_tts.py index 9c1ddbf71..9c7c10164 100644 --- a/test/qwen3-tts/parity_qwen3_tts.py +++ b/test/qwen3-tts/parity_qwen3_tts.py @@ -72,12 +72,23 @@ def load_reference(snapshot: str, device: str): ) -def reference_generate(ref, args) -> tuple[torch.Tensor, list[torch.Tensor]]: - """Greedy reference codes ``[frames, groups]`` and the embeds the Talker saw. +def reference_clone_prompt(ref, args): + """The reference's voice-clone prompt (x-vector + codes) for ``--ref-audio``, or None.""" + if not args.ref_audio: + return None + items = ref.create_voice_clone_prompt( + ref_audio=args.ref_audio, ref_text=args.ref_text, x_vector_only_mode=args.x_vector_only, + ) + return ref._prompt_items_to_voice_clone_prompt(items), items[0] + + +def reference_generate(ref, args, clone=None) -> tuple[torch.Tensor, list[torch.Tensor]]: + """Greedy reference codes ``[frames, groups]`` and the hidden states behind them. Uses the low-level ``generate`` of the reference so the prompt layout - (speaker, language, instruct, streaming vs non-streaming text) is exactly - the reference's own, independent of M*'s ``process_prompt``. + (speaker, language, instruct, reference clip, streaming vs non-streaming + text) is exactly the reference's own, independent of M*'s + ``process_prompt``. """ model = ref.model device = model.device @@ -89,6 +100,13 @@ def reference_generate(ref, args) -> tuple[torch.Tensor, list[torch.Tensor]]: non_streaming = model.tts_model_type in ("custom_voice", "voice_design") if args.non_streaming_mode is not None: non_streaming = args.non_streaming_mode + extra = {} + if clone is not None: + prompt_dict, item = clone + extra = { + "voice_clone_prompt": prompt_dict, + "ref_ids": [ref._tokenize_texts([ref._build_ref_text(item.ref_text)])[0]] if item.ref_text else None, + } codes_list, hidden_list = model.generate( input_ids=input_ids, instruct_ids=instruct_ids, @@ -99,6 +117,7 @@ def reference_generate(ref, args) -> tuple[torch.Tensor, list[torch.Tensor]]: do_sample=False, subtalker_dosample=False, repetition_penalty=args.repetition_penalty, + **extra, ) codes = codes_list[0].to(device) if codes.shape[0] < args.frames: @@ -201,6 +220,8 @@ def close_request(self, rid: str) -> None: def step(self, walk: str, fwd: CurrentForwardPassInfo, inputs: dict, forward): """One step; ``forward(engine_inputs, **preprocessed)`` runs the compute.""" rid = fwd.request_id + if walk == "talker_prefill" and ("speaker_embed" in inputs or "ref_codes" in inputs): + walk = "talker_prefill_clone" fwd.graph_walk = walk prepared = self.talker.prepare_inputs(walk, fwd, inputs) step = self.talker.declare_step(walk, [rid], [prepared]) @@ -355,6 +376,9 @@ def main(argv: list[str] | None = None) -> None: parser.add_argument("--language", default="English") parser.add_argument("--instruct", default=None) parser.add_argument("--non-streaming-mode", type=lambda s: s.lower() == "true", default=None) + parser.add_argument("--ref-audio", default=None, help="Base: reference clip (voice clone)") + parser.add_argument("--ref-text", default=None, help="Base: transcript of the reference clip") + parser.add_argument("--x-vector-only", action="store_true", help="Base: skip in-context frames") parser.add_argument("--frames", type=int, default=64) parser.add_argument("--repetition-penalty", type=float, default=1.05) parser.add_argument("--device", default="cuda:0") @@ -383,11 +407,54 @@ def main(argv: list[str] | None = None) -> None: request_kwargs["instruct"] = args.instruct if args.non_streaming_mode is not None: request_kwargs["non_streaming_mode"] = args.non_streaming_mode - tensors = model.process_prompt(args.text, ["text"], ["audio"], **request_kwargs) + clone = reference_clone_prompt(ref, args) + clone_report = None + if clone is None: + tensors = model.process_prompt(args.text, ["text"], ["audio"], **request_kwargs) + else: + # Voice clone: M*'s load_audio -> process_prompt -> RefEncoder, compared + # with the reference's own x-vector and codec frames for the same clip. + clip = model.load_audio(args.ref_audio, args.device) + request_kwargs.update({"ref_text": args.ref_text, "x_vector_only_mode": args.x_vector_only}) + tensors = model.process_prompt( + args.text, ["audio", "text"], ["audio"], tensors={"audio_inputs": [clip.data]}, **request_kwargs, + ) + ref_encoder = model.get_submodule("RefEncoder", device=args.device, autocast_dtype=torch.bfloat16) + prepared = ref_encoder.prepare_inputs( + "talker_prefill_clone", CurrentForwardPassInfo( + request_id="clone", graph_walk="talker_prefill_clone", fwd_index=0, random_seed=0, max_tokens=0, + ), {"audio_inputs": [clip.data], "prompt_layout": tensors["prompt_layout"]}, + ) + with torch.no_grad(): + encoded = ref_encoder.forward( + "talker_prefill_clone", ModelInputsFromEngine(request_ids=["clone"], per_request_info={}), + **ref_encoder.preprocess("talker_prefill_clone", None, [prepared]), + ) + tensors["speaker_embed"] = encoded["speaker_embed"] + tensors["ref_codes"] = encoded["ref_codes"] + _, item = clone + their_xvec = item.ref_spk_embedding.to(args.device).float() + our_xvec = encoded["speaker_embed"][0].float() + clone_report = { + "xvector_cosine": float(torch.nn.functional.cosine_similarity(our_xvec, their_xvec, dim=0)), + "xvector_max_abs_diff": float((our_xvec - their_xvec).abs().max()), + "xvector_scale": float(their_xvec.abs().mean()), + } + if item.ref_code is not None: + their_codes = item.ref_code.to(args.device) + our_codes = encoded["ref_codes"][0] + n = min(their_codes.shape[0], our_codes.shape[0]) + clone_report.update({ + "ref_frames_mstar": int(our_codes.shape[0]), + "ref_frames_reference": int(their_codes.shape[0]), + "ref_code_agreement": float((our_codes[:n] == their_codes[:n]).float().mean()), + }) + del ref_encoder + torch.cuda.empty_cache() # 1 + 2: teacher forced against the reference's greedy frames. ``ref_hidden`` # is the hidden state the reference's own generation used for each frame. - ref_codes, ref_hidden = reference_generate(ref, args) + ref_codes, ref_hidden = reference_generate(ref, args, clone) ref_hidden = ref_hidden.to(torch.bfloat16) frames = ref_codes.shape[0] theirs_logits = ref.model.talker.codec_head(ref_hidden) @@ -403,6 +470,8 @@ def main(argv: list[str] | None = None) -> None: prefill = talker._build_prefill( "layout-probe", tensors["text_inputs"][0], tensors["prompt_layout"][0], int(tensors["speaker_id"][0]), int(tensors["language_id"][0]), + speaker_embed=tensors.get("speaker_embed", [None])[0], + ref_codes=tensors.get("ref_codes", [None])[0], ) probe_state = talker.request_state("layout-probe") backbone_logits, _ = reference_teacher_forced( @@ -441,6 +510,7 @@ def main(argv: list[str] | None = None) -> None: report = { "repo": args.repo, "snapshot": snapshot, "text": args.text, "voice": args.voice, "language": args.language, "instruct": args.instruct, "frames": int(ref_codes.shape[0]), + "clone": clone_report, "talker": talker_report, "talker_hidden": hidden_report, "backbone_only": backbone_report, "code_predictor": cp_report, "greedy_codes": codes_report, "codec": codec_report, "audio": e2e_audio, From 29691b57bd9946a92be5f11492086c34540bbee8 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 046/110] qwen3_tts: advertise the codec clone walk only on Base checkpoints --- mstar/model/qwen3_tts/submodules.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index 4070dd0b0..46cc2864a 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -952,6 +952,13 @@ def __init__(self, decoder: torch.nn.Module, config: Qwen3TTSModelConfig): ): self.total_upsample *= factor + def _codec_walks(self) -> list[str]: + """Walks this node runs: the clone walk exists only on Base checkpoints.""" + walks = ["codec_chunk"] + if self.config.supports_reference_audio: + walks.append("codec_chunk_clone") + return walks + def _bucket(self, frames: int) -> int: """Smallest captured window that holds ``frames`` (the terminal flush is shorter).""" for window in self.windows: @@ -1118,7 +1125,7 @@ def get_cuda_graph_configs( return [ BatchedCudaGraphConfig( capture_graph_walk="codec_chunk", - replay_graph_walks=["codec_chunk", "codec_chunk_clone"], + replay_graph_walks=self._codec_walks(), single_request_inputs=ARNodeInputs( # 1, not the window: batched buckets match on bs, and this # keeps the intern seq_len from aliasing the trailing dims. @@ -1143,7 +1150,7 @@ def can_use_cuda_graphs( self, batch: ExecutingBatch, model_inputs: list[NodeInputs] ) -> bool: return ( - batch.graph_walk in ("codec_chunk", "codec_chunk_clone") + batch.graph_walk in self._codec_walks() and self.can_batch(batch, model_inputs) and all( item.tensor_inputs["codec_tokens"].shape From 81c05147e50187988c8abbb8d594dbe41ec826d3 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 047/110] test: codec graph configs replay the clone walk only for Base --- test/modular/test_qwen3_tts_model.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index b58d5193c..ebac9c1ad 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -1419,11 +1419,18 @@ def test_qwen3_tts_codec_batches_and_declares_cuda_graphs(): assert [c.additional_key_info for c in graph_configs] == [1, 4] for graph_config in graph_configs: assert graph_config.capture_graph_walk == "codec_chunk" - assert set(graph_config.replay_graph_walks) == {"codec_chunk", "codec_chunk_clone"} + # CustomVoice has no clone walk; a Base config would add codec_chunk_clone. + assert set(graph_config.replay_graph_walks) == {"codec_chunk"} assert graph_config.capture_batch_sizes == [1, 2, 4, 8, 16] assert graph_config.single_request_inputs.tensor_inputs["codec_tokens"].shape == ( 4, graph_config.additional_key_info, ) + base_config = _tiny_model_config() + base_config.tts_model_type = "base" + base_codec = CodecSubmodule(_FakeCodecDecoder(4), base_config) + assert set(base_codec.get_cuda_graph_configs(torch.device("cpu"))[0].replay_graph_walks) == { + "codec_chunk", "codec_chunk_clone", + } assert submodule.max_batch_size("codec_chunk") == 16 # The batch's capture key is the bucket its requests were padded to. for rid in ("a", "b"): From 52b047eebe3ecb6b7a4390cf63280b77d9fc1b34 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 048/110] qwen3_tts: pick the codec capture bucket from stream metadata --- mstar/model/qwen3_tts/submodules.py | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index 46cc2864a..de2a4ef64 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -1006,6 +1006,7 @@ def prepare_inputs( raise ValueError( f"Expected codec tokens with shape (frames, groups), got {codes.shape}" ) + num_items = codes.shape[0] # EOS belongs to Talker loop control and is not a valid codec codebook # index for waveform reconstruction. It is only ever the last frame of # a stream, so the leading context count is unaffected. @@ -1014,8 +1015,11 @@ def prepare_inputs( :self.config.codec.num_quantizers, ] frames = codes.shape[0] - context = int(self._stream_chunk_meta(fwd_info).get("context_items", 0)) - bucket = self._bucket(max(frames, 1)) + meta = self._stream_chunk_meta(fwd_info) + context = int(meta.get("context_items", 0)) + # The bucket is chosen from the window's item count (EOS included) so + # that ``cg_key_info``, which only sees the stream metadata, agrees. + bucket = self._bucket(max(int(meta.get("num_items", num_items)), 1)) if frames < bucket: codes = torch.nn.functional.pad(codes, (0, 0, 0, bucket - frames)) state.add_all( @@ -1109,12 +1113,21 @@ def cg_key_info( graph_walk: str, per_request_info: Mapping[str, CurrentForwardPassInfo], ) -> Any: - """The window bucket this batch was padded to (``can_batch`` keeps it uniform).""" + """The window bucket this batch pads to (``can_batch`` keeps it uniform). + + Derived from the stream metadata the worker attaches to each request + (available before ``prepare_inputs`` runs, so a pre-planned lease can + find its capture); the state written by ``prepare_inputs`` is the + fallback for callers without that metadata. + """ del graph_walk - buckets = { - self.request_state(request_id).get("codec_bucket") - for request_id in per_request_info - } + buckets = set() + for request_id, fwd_info in per_request_info.items(): + num_items = self._stream_chunk_meta(fwd_info).get("num_items") + if num_items is not None: + buckets.add(self._bucket(max(int(num_items), 1))) + else: + buckets.add(self.request_state(request_id).get("codec_bucket")) return buckets.pop() if len(buckets) == 1 else None def get_cuda_graph_configs( From f4ac4321bb542ddbe66676252b18b04fb5ef413c Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 049/110] test: codec capture key from stream metadata --- test/modular/test_qwen3_tts_model.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index ebac9c1ad..52f8d058d 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -1316,12 +1316,12 @@ def test_qwen3_tts_codec_filters_eos_and_pads_to_capture_shape(): fwd_info = SimpleNamespace( request_id="request", step_metadata={"stream_chunks": {"codec_tokens": { - "start_offset": 1, "context_items": 1, "is_final": False, + "start_offset": 1, "context_items": 1, "num_items": 3, "is_final": False, }}}, ) prepared = submodule.prepare_inputs("codec_chunk", fwd_info, {"codec_tokens": [codes]}) - # Two real frames pad up to the smallest captured window (4), not the largest. + # Three items (one of them EOS) pad up to the smallest captured window (4). packed = prepared.tensor_inputs["codec_tokens"] assert packed.shape == (4, 4) assert packed[:, :2].t().tolist() == [[1, 2, 3, 4], [5, 6, 7, 8]] @@ -1432,7 +1432,15 @@ def test_qwen3_tts_codec_batches_and_declares_cuda_graphs(): "codec_chunk", "codec_chunk_clone", } assert submodule.max_batch_size("codec_chunk") == 16 - # The batch's capture key is the bucket its requests were padded to. + # The batch's capture key is the bucket its requests pad to: read off the + # stream metadata when present (before prepare_inputs), else off the state. + def meta(num_items): + return SimpleNamespace(step_metadata={"stream_chunks": {"codec_tokens": { + "num_items": num_items, "context_items": 0, "start_offset": 0, "is_final": False, + }}}) + + assert submodule.cg_key_info("codec_chunk", {"a": meta(3), "b": meta(4)}) == 4 + assert submodule.cg_key_info("codec_chunk", {"a": meta(1), "b": meta(4)}) is None for rid in ("a", "b"): submodule.request_state(rid).add("codec_bucket", 4) assert submodule.cg_key_info("codec_chunk", {"a": None, "b": None}) == 4 From 22013c4e5d64bf1ed51e56b0763d5d2f55d263dd Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 050/110] streaming: report each chunk's item count --- mstar/streaming/stream_buffer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mstar/streaming/stream_buffer.py b/mstar/streaming/stream_buffer.py index 3375f7f6b..aa8dec71a 100644 --- a/mstar/streaming/stream_buffer.py +++ b/mstar/streaming/stream_buffer.py @@ -17,6 +17,9 @@ class StreamChunk: # leading items of this chunk that an earlier chunk already delivered # (sliding-window overlap / left context); the consumer trims their output context_items: int = 0 + # items in this chunk (context included); lets a consumer pick its + # capture bucket before the chunk tensor is unpacked + num_items: int = 0 @dataclass @@ -145,6 +148,7 @@ def pop_chunk(self) -> StreamChunk: start_offset=offset, is_final=is_final, context_items=min(max(self._delivered_end - offset, 0), len(items)), + num_items=len(items), ) self._delivered_end = max(self._delivered_end, offset + len(items)) self._chunks_popped += 1 From f7d58e8c03bceabd0972ff5c424e00328895d30a Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 051/110] graph: carry the stream chunk item count on synthetic edges --- mstar/graph/base.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mstar/graph/base.py b/mstar/graph/base.py index 7b9524848..a4e0b8b62 100644 --- a/mstar/graph/base.py +++ b/mstar/graph/base.py @@ -79,6 +79,7 @@ class GraphEdge: # None on every other edge _stream_chunk_offset: int | None = field(default=None) _stream_chunk_context: int | None = field(default=None) + _stream_chunk_items: int | None = field(default=None) # Set for sharded configurations _total_fanin: int = 1 @@ -97,6 +98,7 @@ def clone(self): _final_stream_chunk=self._final_stream_chunk, _stream_chunk_offset=self._stream_chunk_offset, _stream_chunk_context=self._stream_chunk_context, + _stream_chunk_items=self._stream_chunk_items, ) From 3972df81811af0aeb27e5c4aeb20b4069efb0e90 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 052/110] worker: include the chunk item count in stream_chunks metadata --- mstar/worker/worker.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 0dd09f4f4..569541e30 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -791,6 +791,7 @@ def _pop_streaming_edge( _final_stream_chunk=chunk.is_final, _stream_chunk_offset=chunk.start_offset, _stream_chunk_context=chunk.context_items, + _stream_chunk_items=chunk.num_items, ) else: # Normal chunk — store tensor and create edge with tensor_info. @@ -811,6 +812,7 @@ def _pop_streaming_edge( _final_stream_chunk=chunk.is_final, _stream_chunk_offset=chunk.start_offset, _stream_chunk_context=chunk.context_items, + _stream_chunk_items=chunk.num_items, ) return synthetic_edge @@ -976,6 +978,7 @@ def _build_executing_batch(self, batch: ScheduledBatch) -> ExecutingBatch: stream_chunks[input_name] = { "start_offset": edge._stream_chunk_offset, "context_items": edge._stream_chunk_context, + "num_items": edge._stream_chunk_items, "is_final": edge._final_stream_chunk, } per_request_inputs[request_id] = tensors From 381826344c9ab8df98b22ab2bf3a608f8e8b9767 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 053/110] test: stream chunk item counts --- test/modular/test_stream_chunk_schedule.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/test/modular/test_stream_chunk_schedule.py b/test/modular/test_stream_chunk_schedule.py index c35d6f528..7abc6aa05 100644 --- a/test/modular/test_stream_chunk_schedule.py +++ b/test/modular/test_stream_chunk_schedule.py @@ -49,8 +49,10 @@ def _geometry(chunks): for chunk in chunks: data = chunk.data["data"] if data is None: + assert chunk.num_items == 0 continue items = data.reshape(-1).tolist() + assert chunk.num_items == len(items) out.append((len(items), chunk.context_items, chunk.start_offset, items)) return out @@ -120,7 +122,9 @@ def test_existing_policies_report_context_items(policy, expected): def test_graph_edge_clone_keeps_stream_chunk_geometry(): edge = GraphEdge(next_node="Codec", name="codec_tokens", _final_stream_chunk=True, - _stream_chunk_offset=7, _stream_chunk_context=3) + _stream_chunk_offset=7, _stream_chunk_context=3, _stream_chunk_items=12) clone = edge.clone() - assert (clone._stream_chunk_offset, clone._stream_chunk_context, clone._final_stream_chunk) == (7, 3, True) - assert GraphEdge(next_node="x", name="y")._stream_chunk_context is None + assert (clone._stream_chunk_offset, clone._stream_chunk_context, clone._stream_chunk_items, + clone._final_stream_chunk) == (7, 3, 12, True) + plain = GraphEdge(next_node="x", name="y") + assert plain._stream_chunk_context is None and plain._stream_chunk_items is None From c1ef190d42c98af30015b315f430d13e180d1a59 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 054/110] deps: soundfile in the qwen3_tts extra for reference audio --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 4682e1193..7ae9fea08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -155,6 +155,7 @@ qwen3_tts = [ "mooncake-transfer-engine", "qwen-tts==0.1.1", "safetensors", + "soundfile", # reference-audio decoding for Base voice cloning "transformers>=4.57.3", ] From 36efac18b5cf1a7834469d1bf6389ab8d3b0a611 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 055/110] qwen3_tts: hand the codec its reference frame count from the first clone chunk --- mstar/model/qwen3_tts/qwen3_tts_model.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/mstar/model/qwen3_tts/qwen3_tts_model.py b/mstar/model/qwen3_tts/qwen3_tts_model.py index ddb33342c..b7ed98a72 100644 --- a/mstar/model/qwen3_tts/qwen3_tts_model.py +++ b/mstar/model/qwen3_tts/qwen3_tts_model.py @@ -770,12 +770,23 @@ def get_initial_forward_pass_args( ) return ForwardPassArgs( full_metadata=metadata, - inputs=[], + # The clone walk needs the reference frame count next to the + # stream from its very first chunk; the API tensor stays + # persisted so every later re-arm can read it again. + inputs=self._codec_ref_frames_edges(input_signals) if clone else [], unpersist_tensors=[], request_done="audio" not in output_modalities, ) raise ValueError(f"Unknown Qwen3-TTS partition: {partition_name!r}") + @staticmethod + def _codec_ref_frames_edges( + signals: dict[str, list[TensorPointerInfo]], + ) -> list[GraphEdge]: + edge = GraphEdge(next_node="Codec", name="ref_frames") + edge.tensor_info = signals.get("ref_frames", []) + return [edge] + def get_partition_forward_pass_args( self, partition_name: str, @@ -820,9 +831,7 @@ def get_partition_forward_pass_args( if partition_metadata.graph_walk == "codec_chunk_clone": # The reference frame count is an API tensor; every codec # invocation of a clone request re-reads it (cheap, one int). - edge = GraphEdge(next_node="Codec", name="ref_frames") - edge.tensor_info = persist_signals.get("ref_frames", []) - inputs.append(edge) + inputs = self._codec_ref_frames_edges(persist_signals) else: partition_metadata.graph_walk = "codec_chunk" return ForwardPassArgs( From bc5b6d38023805747d78ae04c88303904d42ab7a Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:09:55 -0700 Subject: [PATCH 056/110] test: codec clone walk receives ref_frames on its initial inputs --- test/modular/test_qwen3_tts_model.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index 52f8d058d..d09989a24 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -490,6 +490,11 @@ def test_qwen3_tts_base_declares_clone_walks_and_routes_reference_audio(): input_signals=pointers, ) assert codec.full_metadata.graph_walk == "codec_chunk_clone" + # The very first codec chunk already needs the reference frame count: it + # rides the initial inputs (and stays persisted for every later chunk). + assert [edge.name for edge in codec.inputs] == ["ref_frames"] + assert codec.inputs[0].tensor_info == pointers["ref_frames"] + assert codec.unpersist_tensors == [] rearmed = model.get_partition_forward_pass_args( "Codec", codec.full_metadata, persist_signals={"ref_frames": pointers["ref_frames"]}, ) From 752611ed2aef1d64f2dfb9bc4dd9953c2ec64052 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:46:53 -0700 Subject: [PATCH 057/110] api: sentence splitter for long text-to-speech inputs --- mstar/api_server/openai/speech_chunking.py | 100 +++++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 mstar/api_server/openai/speech_chunking.py diff --git a/mstar/api_server/openai/speech_chunking.py b/mstar/api_server/openai/speech_chunking.py new file mode 100644 index 000000000..e21cadf93 --- /dev/null +++ b/mstar/api_server/openai/speech_chunking.py @@ -0,0 +1,100 @@ +"""Sentence chunking for ``/v1/audio/speech``. + +Long inputs are split into sentence groups that are synthesized as separate +requests and played back in order. This keeps every autoregressive +text-to-speech request inside the length its model was trained for, lets the +engine batch the pieces of one long input like independent requests, and +starts the second piece while the first is still streaming (see +``serving_speech``). Splitting is purely textual and model-agnostic. +""" + +from __future__ import annotations + +import re + +# One sentence: a lazy body up to a Latin terminator (with any closing quotes +# or brackets) that precedes whitespace or the end, a CJK terminator, a +# paragraph break, or the end of the text. +_SENTENCE = re.compile( + r""".+?(?: + [.!?;]["'”’)\]]*(?=\s|\Z) + | [。！？；] + | (?=\n[ \t]*\n) + | \Z + )""", + re.VERBOSE | re.DOTALL, +) +# Soft break points inside an over-long sentence, most preferred first. +_SOFT_BREAKS = ( + re.compile(r"(?<=[,;:，；：])\s*"), + re.compile(r"\s+"), +) + + +def _hard_wrap(sentence: str, max_chars: int) -> list[str]: + """Split one over-long sentence at clause boundaries, then at spaces.""" + pieces = [sentence] + for pattern in _SOFT_BREAKS: + wrapped: list[str] = [] + for piece in pieces: + if len(piece) <= max_chars: + wrapped.append(piece) + continue + current = "" + for part in (p.strip() for p in pattern.split(piece)): + if not part: + continue + candidate = f"{current} {part}" if current else part + if current and len(candidate) > max_chars: + wrapped.append(current) + current = part + else: + current = candidate + if current: + wrapped.append(current) + pieces = wrapped + return pieces + + +def split_sentences(text: str, max_chars: int = 400, min_chars: int = 24) -> list[str]: + """Group ``text`` into sentence chunks of at most ``max_chars`` characters. + + Sentences are never cut unless one alone exceeds ``max_chars`` (then it is + wrapped at clause boundaries, or spaces as a last resort). A trailing + fragment shorter than ``min_chars`` is merged into its predecessor so the + model is not asked to voice a lone "Okay." Returns ``[text]`` when nothing + needs splitting and ``[]`` for blank input. + """ + if max_chars <= 0: + raise ValueError("max_chars must be positive") + text = text.strip() + if not text: + return [] + if len(text) <= max_chars: + return [text] + + sentences: list[str] = [] + for match in _SENTENCE.finditer(text): + piece = " ".join(match.group(0).split()) + if not piece: + continue + sentences.extend(_hard_wrap(piece, max_chars) if len(piece) > max_chars else [piece]) + + chunks: list[str] = [] + current = "" + for sentence in sentences: + if current and len(current) + 1 + len(sentence) > max_chars: + chunks.append(current) + current = sentence + else: + current = f"{current} {sentence}".strip() + if current: + chunks.append(current) + + if ( + len(chunks) > 1 + and len(chunks[-1]) < min_chars + and len(chunks[-2]) + 1 + len(chunks[-1]) <= max_chars * 5 // 4 + ): + chunks[-2:] = [f"{chunks[-2]} {chunks[-1]}"] + return chunks From e234e948fd587fb939b7c4fd8873810ea28db6a8 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:46:53 -0700 Subject: [PATCH 058/110] api: per-adapter sentence-chunking thresholds for /v1/audio/speech --- mstar/api_server/openai/adapters.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/mstar/api_server/openai/adapters.py b/mstar/api_server/openai/adapters.py index 394413a73..ec24716c1 100644 --- a/mstar/api_server/openai/adapters.py +++ b/mstar/api_server/openai/adapters.py @@ -185,6 +185,16 @@ class OpenAIAdapter: supports_videos: bool = False # POST /v1/videos/generations supports_realtime: bool = False # /v1/realtime (bidirectional speech WebSocket) + # ``/v1/audio/speech`` sentence chunking (``serving_speech``): inputs of at + # least ``speech_chunk_min_chars`` characters are split into sentence + # groups of about ``speech_chunk_max_chars`` and synthesized as ordered + # sub-requests. ``None`` keeps the whole text in one request unless the + # client sends ``sentence_chunking: true``. + speech_chunk_min_chars: int | None = None + speech_chunk_max_chars: int = 400 + # sub-requests kept in flight ahead of the one being streamed + speech_chunk_lookahead: int = 2 + def chat_to_request(self, req: ChatCompletionRequest, upload_dir: Path) -> SubmitArgs: # noqa: ARG002 # Output modalities vary by model: e.g. Qwen3-Omni speech output also # emits text, whereas BAGEL chat is text-only. From afc57df05f6a11dae27d89ca8833977c591bfe06 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:46:53 -0700 Subject: [PATCH 059/110] api: synthesize long speech inputs as ordered sentence chunks --- mstar/api_server/openai/serving_speech.py | 82 ++++++++++++++++++----- 1 file changed, 65 insertions(+), 17 deletions(-) diff --git a/mstar/api_server/openai/serving_speech.py b/mstar/api_server/openai/serving_speech.py index 335f491f6..883318675 100644 --- a/mstar/api_server/openai/serving_speech.py +++ b/mstar/api_server/openai/serving_speech.py @@ -3,14 +3,48 @@ Non-streaming returns the full audio as a container blob (WAV by default). Streaming returns a single open-ended WAV response (header + PCM16 frames) as the audio is produced. + +Long inputs can be synthesized as ordered sentence chunks (one engine request +per chunk, ``speech_chunking.split_sentences``): the adapter's +``speech_chunk_min_chars`` turns this on for long texts, and a client can force +or suppress it per request with ``sentence_chunking: true|false``. The next +chunks are submitted while the current one streams, so the engine batches them +and playback never waits for a prefill. """ from __future__ import annotations +from collections.abc import Callable + from fastapi.responses import Response, StreamingResponse from mstar.api_server import media_io from mstar.api_server.openai._util import rid +from mstar.api_server.openai.speech_chunking import split_sentences + + +def _plan_chunks(req, adapter, text: str) -> list[str]: + """The texts to synthesize, in playback order (``[text]`` when unchunked).""" + # Extra request fields live in pydantic's ``model_extra``; plain objects + # (tests, other frontends) may carry the attribute directly. + requested = (getattr(req, "model_extra", None) or {}).get("sentence_chunking") + if requested is None: + requested = getattr(req, "sentence_chunking", None) + min_chars = getattr(adapter, "speech_chunk_min_chars", None) + if requested is False or (requested is None and (min_chars is None or len(text) < min_chars)): + return [text] + chunks = split_sentences(text, max_chars=getattr(adapter, "speech_chunk_max_chars", 400)) + return chunks or [text] + + +def _chunk_kwargs(model_kwargs: dict, index: int) -> dict: + """Per-chunk model kwargs: identical, except a client seed advances per chunk.""" + kwargs = dict(model_kwargs) + kwargs.pop("sentence_chunking", None) + seed = kwargs.get("seed") + if isinstance(seed, int) and not isinstance(seed, bool): + kwargs["seed"] = seed + index + return kwargs async def create_speech(api, model_name, adapter, req, raw_request=None): # noqa: ARG001 @@ -18,32 +52,46 @@ async def create_speech(api, model_name, adapter, req, raw_request=None): # noq request_id = rid("speech") sample_rate = api.model.get_output_sample_rate("audio") if api.model is not None else 24000 fmt = (req.response_format or "wav").lower() + chunks = _plan_chunks(req, adapter, args.text or "") - api.submit_request( - text=args.text, - file_paths=args.file_paths, - input_modalities=args.input_modalities, - output_modalities=args.output_modalities, - model_kwargs=args.model_kwargs, - streaming=bool(req.stream), - request_id=request_id, - ) + def submit(index: int) -> str: + chunk_id = request_id if len(chunks) == 1 else f"{request_id}-{index}" + return api.submit_request( + text=chunks[index], + file_paths=args.file_paths, + input_modalities=args.input_modalities, + output_modalities=args.output_modalities, + model_kwargs=_chunk_kwargs(args.model_kwargs, index), + streaming=bool(req.stream), + request_id=chunk_id, + ) + lookahead = max(1, int(getattr(adapter, "speech_chunk_lookahead", 2))) if req.stream: return StreamingResponse( - _stream_wav(api, request_id, sample_rate), + _stream_wav(api, submit, len(chunks), lookahead, sample_rate), media_type="audio/wav", headers={"Cache-Control": "no-cache"}, ) - chunks = await api.collect_results(request_id, raw_request) - pcm = b"".join(c.data for c in chunks if c.modality == "audio") - audio_bytes, mime = media_io.pcm16_to_container(pcm, sample_rate, fmt) + pcm_parts: list[bytes] = [] + pending: list[str] = [submit(i) for i in range(min(lookahead, len(chunks)))] + for index in range(len(chunks)): + if len(pending) < len(chunks): + pending.append(submit(len(pending))) + results = await api.collect_results(pending[index], raw_request) + pcm_parts.append(b"".join(c.data for c in results if c.modality == "audio")) + audio_bytes, mime = media_io.pcm16_to_container(b"".join(pcm_parts), sample_rate, fmt) return Response(content=audio_bytes, media_type=mime) -async def _stream_wav(api, request_id, sample_rate): +async def _stream_wav(api, submit: Callable[[int], str], num_chunks: int, lookahead: int, sample_rate: int): yield media_io.wav_stream_header(sample_rate) - async for c in api.iter_result_chunks(request_id): - if c.modality == "audio" and c.data: - yield c.data + pending: list[str] = [submit(i) for i in range(min(lookahead, num_chunks))] + for index in range(num_chunks): + if len(pending) < num_chunks: + # Keep the next chunk generating while this one plays. + pending.append(submit(len(pending))) + async for c in api.iter_result_chunks(pending[index]): + if c.modality == "audio" and c.data: + yield c.data From 0b778386a76938bc81a5ae0c45863985f254dcfd Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:46:53 -0700 Subject: [PATCH 060/110] test: sentence chunking on /v1/audio/speech --- test/modular/test_speech_chunking.py | 209 +++++++++++++++++++++++++++ 1 file changed, 209 insertions(+) create mode 100644 test/modular/test_speech_chunking.py diff --git a/test/modular/test_speech_chunking.py b/test/modular/test_speech_chunking.py new file mode 100644 index 000000000..37a80e329 --- /dev/null +++ b/test/modular/test_speech_chunking.py @@ -0,0 +1,209 @@ +"""Sentence chunking on ``/v1/audio/speech``: the splitter and the ordered sub-requests. + +The router is mounted on a FastAPI app with a stubbed APIServer (as in +``test_openai_router.py``); every sub-request the handler submits is recorded +so ordering, ids, seeds and playback concatenation can be checked without an +engine. +""" + +import sys +import tempfile +import types +from pathlib import Path + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("pydantic") +pytest.importorskip("httpx") +np = pytest.importorskip("numpy") + +from fastapi import FastAPI # noqa: E402 +from fastapi.testclient import TestClient # noqa: E402 + +from mstar.api_server.openai import adapters # noqa: E402 +from mstar.api_server.openai.speech_chunking import split_sentences # noqa: E402 + +LONG_TEXT = ( + "The train to the coast leaves at seven tomorrow morning, so please pack your bag tonight. " + "She opened the window and let the cool evening air drift into the kitchen. " + "Our meeting has been moved to Thursday afternoon because the room is being repainted! " + "A gentle rain fell over the harbor while the fishing boats returned one by one. " + "Remember to water the tomatoes twice a week during the hottest part of the summer? " + "The museum's new exhibit traces the history of printing from wooden blocks to modern presses." +) + + +# --------------------------------------------------------------------------- +# splitter +# --------------------------------------------------------------------------- + + +def test_short_text_is_one_chunk(): + assert split_sentences("Hello there. How are you?", max_chars=400) == ["Hello there. How are you?"] + assert split_sentences(" ", max_chars=400) == [] + + +def test_sentences_are_grouped_up_to_max_chars_and_never_cut(): + chunks = split_sentences(LONG_TEXT, max_chars=200) + assert len(chunks) == 3 + assert all(len(c) <= 200 for c in chunks) + assert " ".join(chunks) == " ".join(LONG_TEXT.split()) + # every chunk ends where a sentence ends + assert all(c[-1] in ".!?" for c in chunks) + + +def test_cjk_terminators_and_paragraph_breaks_split(): + text = "今天天气很好。我们去公园散步吧！你觉得怎么样？\n\nSecond paragraph here." + chunks = split_sentences(text, max_chars=8, min_chars=1) + assert chunks[:3] == ["今天天气很好。", "我们去公园散步吧！", "你觉得怎么样？"] + assert " ".join(chunks[3:]) == "Second paragraph here." + + paragraphs = split_sentences("First paragraph no period\n\nsecond paragraph no period", max_chars=30, min_chars=1) + assert paragraphs == ["First paragraph no period", "second paragraph no period"] + + +def test_quotes_after_terminators_stay_with_their_sentence(): + text = 'He said "Wait!" Then she left. "Really?" she asked. Yes.' + chunks = split_sentences(text, max_chars=20, min_chars=1) + assert chunks[0] == 'He said "Wait!"' + assert chunks[1] == "Then she left." + assert chunks[2] == '"Really?" she asked.' + + +def test_overlong_sentence_is_wrapped_at_clauses_then_spaces(): + sentence = "alpha beta gamma, delta epsilon zeta, eta theta iota, kappa lambda mu nu xi omicron" + chunks = split_sentences(sentence, max_chars=30, min_chars=1) + assert all(len(c) <= 30 for c in chunks) + assert " ".join(chunks).replace(" ,", ",") == sentence + assert chunks[0] == "alpha beta gamma," + + +def test_tiny_trailing_fragment_merges_into_previous_chunk(): + text = ("A sentence that is fairly long and goes on for a while to fill the chunk nicely. " + "Yes.") + chunks = split_sentences(text, max_chars=90, min_chars=24) + assert chunks == [" ".join(text.split())] + + +def test_max_chars_must_be_positive(): + with pytest.raises(ValueError): + split_sentences("a. b.", max_chars=0) + + +# --------------------------------------------------------------------------- +# handler +# --------------------------------------------------------------------------- + + +class _Chunk: + def __init__(self, modality, data, metadata=None): + self.modality = modality + self.data = data + self.metadata = metadata or {} + + +class _StubModel: + def get_output_sample_rate(self, modality="audio"): + return 24000 + + +def _pcm(*vals): + return np.array(vals, dtype=" bytes: + return content[44:] + + +def test_long_input_is_synthesized_as_ordered_sentence_chunks(client_and_stub): + client, stub = client_and_stub + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": LONG_TEXT, "voice": "tara", "seed": 10}) + assert r.status_code == 200 and r.headers["content-type"] == "audio/wav" + texts = [s["text"] for s in stub.submits] + assert texts == split_sentences(LONG_TEXT, max_chars=200) + # One id per chunk, derived from the request id; seeds advance per chunk; + # the model kwargs are otherwise identical and carry no chunking flag. + ids = [s["request_id"] for s in stub.submits] + assert [i.rsplit("-", 1)[1] for i in ids] == ["0", "1", "2"] and len({i.rsplit("-", 1)[0] for i in ids}) == 1 + assert [s["model_kwargs"]["seed"] for s in stub.submits] == [10, 11, 12] + for submit in stub.submits: + assert submit["model_kwargs"]["voice"] == "tara" + assert "sentence_chunking" not in submit["model_kwargs"] + # Playback order == submission order. + assert _wav_pcm(r.content) == _pcm(0, 0) + _pcm(1, 1) + _pcm(2, 2) + + +def test_streaming_chunks_are_concatenated_in_order(client_and_stub): + client, stub = client_and_stub + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": LONG_TEXT, "stream": True}) + assert r.status_code == 200 and r.content[:4] == b"RIFF" + assert len(stub.submits) == 3 and all(s["streaming"] is True for s in stub.submits) + assert _wav_pcm(r.content) == _pcm(0, 0) + _pcm(1, 1) + _pcm(2, 2) + + +def test_short_input_and_opt_out_keep_a_single_request(client_and_stub): + client, stub = client_and_stub + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": "Hi there. All good?"}) + assert r.status_code == 200 and len(stub.submits) == 1 + assert stub.submits[0]["request_id"].startswith("speech-") and "-" not in stub.submits[0]["request_id"][7:] + + stub.submits.clear() + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": LONG_TEXT, "sentence_chunking": False}) + assert r.status_code == 200 and len(stub.submits) == 1 + assert stub.submits[0]["text"] == LONG_TEXT and "sentence_chunking" not in stub.submits[0]["model_kwargs"] + + +def test_client_can_force_chunking_below_the_threshold(client_and_stub, monkeypatch): + client, stub = client_and_stub + monkeypatch.setattr(adapters.OrpheusAdapter, "speech_chunk_min_chars", None) + monkeypatch.setattr(adapters.OrpheusAdapter, "speech_chunk_max_chars", 60) + text = LONG_TEXT[:150] + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": text, "sentence_chunking": True}) + assert r.status_code == 200 and len(stub.submits) >= 2 + assert " ".join(s["text"] for s in stub.submits) == text + assert all(len(s["text"]) <= 60 for s in stub.submits) From bf73d675bc2b0fe1c583d1b4c40d9e4f48b24c1f Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:48:02 -0700 Subject: [PATCH 061/110] api: sentence-chunk Qwen3-TTS speech inputs from 600 characters --- mstar/api_server/openai/adapters.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/mstar/api_server/openai/adapters.py b/mstar/api_server/openai/adapters.py index ec24716c1..7ac404cfa 100644 --- a/mstar/api_server/openai/adapters.py +++ b/mstar/api_server/openai/adapters.py @@ -400,9 +400,16 @@ class Qwen3TTSAdapter(OpenAIAdapter): ``language``, ``non_streaming_mode``, ``top_k``, ``repetition_penalty``, the residual-group ``subtalker_*`` sampling, ``max_new_tokens``. ``temperature`` / ``top_p`` / ``seed`` map onto the Talker sampler. + + Inputs of 600+ characters are synthesized as ordered sentence chunks + (``serving_speech``): each Talker request stays near the lengths the + model was trained on and the chunks batch like independent requests. + Benchmark sentences (5-40 words) never reach the threshold. """ supports_speech = True + speech_chunk_min_chars = 600 + speech_chunk_max_chars = 400 def speech_to_request(self, req: SpeechRequest, upload_dir: Path) -> SubmitArgs: mk = _passthrough(req) From dacf98b73f418d02e2a3f5437ac1e2686569ce9a Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:48:02 -0700 Subject: [PATCH 062/110] test: Qwen3-TTS speech chunking thresholds --- test/modular/test_openai_adapters.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/modular/test_openai_adapters.py b/test/modular/test_openai_adapters.py index 3e3ee617c..8e40e8c8f 100644 --- a/test/modular/test_openai_adapters.py +++ b/test/modular/test_openai_adapters.py @@ -91,6 +91,10 @@ def test_qwen3_tts_speech_maps_voice_instructions_and_extra_body(tmp_path): assert "max_output_tokens" not in mk for key in ("qwen3_tts", "qwen3_tts_1p7b", "qwen3_tts_voicedesign", "qwen3_tts_base"): assert isinstance(adapters.get_adapter(key), adapters.Qwen3TTSAdapter) + # Long inputs are sentence-chunked by the speech handler; short ones are not. + assert adapters.Qwen3TTSAdapter.speech_chunk_min_chars == 600 + assert adapters.Qwen3TTSAdapter.speech_chunk_max_chars == 400 + assert adapters.OrpheusAdapter.speech_chunk_min_chars is None def test_qwen3_tts_speech_reference_audio_becomes_audio_input(tmp_path): From 46562e697b3fcce6ce4f06bbadaa023cb125fa45 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:48:02 -0700 Subject: [PATCH 063/110] docs: Qwen3-TTS sentence chunking on /v1/audio/speech --- docs/models.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/models.rst b/docs/models.rst index 3da48cb03..0c9f1aa0f 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -116,6 +116,10 @@ Qwen3-TTS notes - Text layout follows the reference defaults: CustomVoice and VoiceDesign put the whole text in the prefill; Base feeds it one token per frame. Override per request with ``non_streaming_mode``. +- ``/v1/audio/speech`` inputs of 600 or more characters are synthesized as + ordered sentence chunks of about 400 characters (one Talker request each, + two kept in flight while the current one streams); set + ``sentence_chunking: false`` (or ``true`` for shorter texts) per request. - Audio streams in a ramp of codec chunks: the first window is decoded after 4 frames (320 ms of speech), later windows grow to 25 new frames behind 25 frames of already decoded left context (the reference's own From bc4aa0fe53ea1295a55d287bd6fd7c3e73d65beb Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:54:05 -0700 Subject: [PATCH 064/110] test: parity harness applies the deployment's resource overrides --- test/qwen3-tts/parity_qwen3_tts.py | 30 +++++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/test/qwen3-tts/parity_qwen3_tts.py b/test/qwen3-tts/parity_qwen3_tts.py index 9c7c10164..a86813e88 100644 --- a/test/qwen3-tts/parity_qwen3_tts.py +++ b/test/qwen3-tts/parity_qwen3_tts.py @@ -41,7 +41,7 @@ from mstar.communication.tensors import LocalTransferEngine from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.distributed.communication import CommGroup, JointGroups -from mstar.engine.resources import StepContext, StepRunner, resolve_spec_dependencies +from mstar.engine.resources import StepContext, StepRunner, apply_yaml_overrides, resolve_spec_dependencies from mstar.engine.resources.base import EngineResourceInfo, build_resource from mstar.engine.resources.kv.transfer import TransferEngineInfo from mstar.model.qwen3_tts.qwen3_tts_model import Qwen3TTSModel @@ -55,6 +55,19 @@ # --------------------------------------------------------------------------- +def default_deployment(config) -> str | None: + """The deployment YAML ``mstar serve`` would use for this checkpoint variant.""" + root = Path(__file__).resolve().parents[2] / "configs" + if config.is_base: + name = "qwen3tts_base.yaml" + elif config.is_voice_design: + name = "qwen3tts_voicedesign.yaml" + else: + name = "qwen3tts.yaml" if config.tts_model_size == "0b6" else "qwen3tts_1p7b.yaml" + path = root / name + return str(path) if path.is_file() else None + + def resolve_snapshot(repo: str) -> str: if Path(repo).is_dir(): return repo @@ -174,12 +187,19 @@ class MStarTalkerDriver: from the model's own ``get_node_resources`` declaration. """ - def __init__(self, model: Qwen3TTSModel, talker, device: str, max_num_pages: int = 64): + def __init__(self, model: Qwen3TTSModel, talker, device: str, deployment: str | None = None, + max_num_pages: int = 64): self.model = model self.talker = talker self.device = torch.device(device) specs = model.get_node_resources() by_key = resolve_spec_dependencies(specs) + if deployment is not None: + # The served deployment's resource overrides (e.g. the FA2 pin), so + # the harness runs the same kernels as ``mstar serve``. + import yaml + + apply_yaml_overrides(specs, yaml.safe_load(Path(deployment).read_text(encoding="utf-8"))) for spec in specs: if hasattr(spec, "apply_yaml_overrides") and hasattr(spec.config, "max_num_pages"): spec.apply_yaml_overrides(max_num_pages=max_num_pages) @@ -382,6 +402,8 @@ def main(argv: list[str] | None = None) -> None: parser.add_argument("--frames", type=int, default=64) parser.add_argument("--repetition-penalty", type=float, default=1.05) parser.add_argument("--device", default="cuda:0") + parser.add_argument("--config", default=None, + help="deployment YAML for resource overrides (default: the variant's)") parser.add_argument("--json", default=None) args = parser.parse_args(argv) if args.voice == "": @@ -396,7 +418,9 @@ def main(argv: list[str] | None = None) -> None: model = Qwen3TTSModel(model_path_hf=snapshot) talker = model.get_submodule("Talker", device=args.device, autocast_dtype=torch.bfloat16) codec = model.get_submodule("Codec", device=args.device) - driver = MStarTalkerDriver(model, talker, args.device) + deployment = args.config or default_deployment(model.config) + print(f"resource overrides from {deployment}", file=sys.stderr) + driver = MStarTalkerDriver(model, talker, args.device, deployment=deployment) print(f"M* loaded in {time.perf_counter() - t0:.1f}s", file=sys.stderr) request_kwargs = {"language": args.language, **GREEDY_KWARGS, From d3b8e60bb556d32af693755dab4176796fc7ba87 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:54:05 -0700 Subject: [PATCH 065/110] benchmark: render the TTS protocol tables from result JSONs --- benchmark/tts_report.py | 88 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 benchmark/tts_report.py diff --git a/benchmark/tts_report.py b/benchmark/tts_report.py new file mode 100644 index 000000000..ff383bbf1 --- /dev/null +++ b/benchmark/tts_report.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +"""Assemble the BENCHMARK_PROTOCOL.md TTS table from ``tts_speech_bench`` / ``tts_wer`` / parity JSONs. + + python -m benchmark.tts_report --results results/2026-09-18 --out results/2026-09-18/REPORT.md + +Every ``*_c.json`` written by ``benchmark/tts_speech_bench.py`` becomes one row +(label, concurrency, TTFA p50/p95, RTF, audio-seconds per second, errors); a +sibling ``*_c_wer.json`` from ``benchmark/tts_wer.py`` fills the WER column +and ``parity_*.json`` files from ``test/qwen3-tts/parity_qwen3_tts.py`` become +the parity table. The markdown is printed and optionally written to ``--out``. +""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path + + +def _load(path: Path) -> dict: + return json.loads(path.read_text(encoding="utf-8")) + + +def benchmark_rows(results: Path) -> list[str]: + rows = [] + for path in sorted(results.glob("*_c[0-9]*.json")): + if path.name.endswith("_wer.json"): + continue + report = _load(path) + med = report["median_over_repeats"] + wer_path = path.with_name(path.stem + "_wer.json") + wer = f"{_load(wer_path)['wer_percent']:.2f}" if wer_path.is_file() else "n/a" + version = report.get("engine_version") or "" + label = report.get("label") or report["engine"] + rows.append( + f"| {label} {version} | {report['concurrency']} | " + f"{med['ttfa_p50_ms']:.0f} / {med['ttfa_p95_ms']:.0f} | {med['rtf_mean']:.3f} | " + f"{med['audio_s_per_wall_s']:.1f} | {wer} | {report['errors_total']} | " + f"{report['num_sentences']} sentences x {report['repeats']} repeats |" + ) + return rows + + +def parity_rows(results: Path) -> list[str]: + rows = [] + for path in sorted(results.glob("parity_*.json")): + r = _load(path) + clone = r.get("clone") or {} + clone_cell = ( + f"cos {clone['xvector_cosine']:.4f}, codes {clone.get('ref_code_agreement', float('nan')):.3f}" + if clone else "-" + ) + rows.append( + f"| {r['repo'].split('/')[-1]} | {re.sub(r'^parity_', '', path.stem)} | {r['frames']} | " + f"{r['talker']['argmax_agreement']:.4f} | {r['code_predictor']['argmax_agreement']:.4f} | " + f"{r['greedy_codes']['identical_frames_before_divergence']}/{r['greedy_codes']['frames_compared']} | " + f"{r['codec']['max_abs_diff']:.2e} | {r['audio']['max_abs_diff']:.3f} | {clone_cell} |" + ) + return rows + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--results", required=True, help="directory with the benchmark / parity JSON files") + parser.add_argument("--out", default=None) + args = parser.parse_args(argv) + results = Path(args.results) + + lines = ["## Benchmarks (H100, back to back, warmup excluded, median over repeats)", "", + "| System (version) | concurrency | TTFA p50 / p95 ms | RTF | audio-s / s | WER % | errors | notes |", + "|---|---|---|---|---|---|---|---|", *benchmark_rows(results), "", + "## Parity vs qwen-tts (greedy, bf16 Talker, fp32 codec)", "", + "| checkpoint | mode | frames | Talker argmax agreement | CodePredictor argmax agreement | " + "identical greedy frames | codec max-abs-diff | greedy audio max-abs-diff | " + "clone (x-vector cosine, ref-code agreement) |", + "|---|---|---|---|---|---|---|---|---|", *parity_rows(results)] + env = results / "environment.txt" + if env.is_file(): + lines += ["", "## Environment", "", "```", env.read_text(encoding="utf-8").strip(), "```"] + text = "\n".join(lines) + "\n" + print(text) + if args.out: + Path(args.out).write_text(text, encoding="utf-8") + + +if __name__ == "__main__": + main() From 9a8d8e2ce4788d9994fad992f68ab366f9f714ad Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:59:41 -0700 Subject: [PATCH 066/110] test: end-to-end smoke test for a served Qwen3-TTS variant --- test/qwen3-tts/smoke_qwen3_tts.py | 176 ++++++++++++++++++++++++++++++ 1 file changed, 176 insertions(+) create mode 100644 test/qwen3-tts/smoke_qwen3_tts.py diff --git a/test/qwen3-tts/smoke_qwen3_tts.py b/test/qwen3-tts/smoke_qwen3_tts.py new file mode 100644 index 000000000..4a137f360 --- /dev/null +++ b/test/qwen3-tts/smoke_qwen3_tts.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +"""End-to-end smoke test of a running Qwen3-TTS server through ``/v1/audio/speech``. + +Exercises the served path (API -> conductor -> Talker/RefEncoder -> codec +stream -> WAV) for whichever variant the server hosts, saving every WAV and +checking that audio is present, non-silent, of plausible length, and that a +greedy request is repeatable. Run inside the GPU allocation against +``mstar serve ``:: + + python test/qwen3-tts/smoke_qwen3_tts.py --url http://127.0.0.1:8000 --variant custom_voice --out results/smoke + python test/qwen3-tts/smoke_qwen3_tts.py --url http://127.0.0.1:8000 --variant base \\ + --ref-audio $BENCH/tts/ref/clone_2.wav --ref-text "Okay. Yeah. I resent you. ..." +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import struct +import sys +import time +from pathlib import Path + +import numpy as np +import requests + +SAMPLE_RATE = 24000 +TEXT = "The train to the coast leaves at seven tomorrow morning, so please pack your bag tonight." +LONG_TEXT = " ".join([ + "The train to the coast leaves at seven tomorrow morning, so please pack your bag tonight.", + "She opened the window and let the cool evening air drift into the kitchen.", + "Our meeting has been moved to Thursday afternoon because the conference room is being repainted.", + "A gentle rain fell over the harbor while the fishing boats returned one by one.", + "Remember to water the tomatoes twice a week during the hottest part of the summer.", + "The museum's new exhibit traces the history of printing from wooden blocks to modern digital presses.", + "He laughed so hard at the joke that he spilled coffee all over his notes.", +]) + + +def speech(url: str, payload: dict, stream: bool, timeout: float) -> tuple[bytes, float, float]: + """Return (pcm16 bytes, time to first audio, total time).""" + start = time.perf_counter() + first = None + pcm = bytearray() + with requests.post(f"{url}/v1/audio/speech", json={**payload, "stream": stream, "response_format": "wav"}, + stream=True, timeout=timeout) as resp: + resp.raise_for_status() + skip = 44 + for chunk in resp.iter_content(chunk_size=None): + if not chunk: + continue + if skip: + drop = min(skip, len(chunk)) + chunk = chunk[drop:] + skip -= drop + if not chunk: + continue + if first is None: + first = time.perf_counter() - start + pcm.extend(chunk) + if not stream: + # non-streaming: the WAV may carry a full header with the data length; PCM after byte 44 + pass + return bytes(pcm), first if first is not None else float("nan"), time.perf_counter() - start + + +def write_wav(path: Path, pcm: bytes) -> None: + header = struct.pack("<4sI4s4sIHHIIHH4sI", b"RIFF", 36 + len(pcm), b"WAVE", b"fmt ", 16, 1, 1, + SAMPLE_RATE, SAMPLE_RATE * 2, 2, 16, b"data", len(pcm)) + path.write_bytes(header + pcm) + + +def stats(pcm: bytes) -> dict: + audio = np.frombuffer(pcm, dtype=" None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--url", default="http://127.0.0.1:8000") + parser.add_argument("--variant", choices=("custom_voice", "voice_design", "base"), required=True) + parser.add_argument("--voice", default="vivian") + parser.add_argument("--ref-audio", default=None) + parser.add_argument("--ref-text", default=None) + parser.add_argument("--out", default="results/smoke") + parser.add_argument("--timeout", type=float, default=300.0) + parser.add_argument("--no-instruct", action="store_true", help="skip the instruction case (0.6B CustomVoice)") + args = parser.parse_args(argv) + out = Path(args.out) + out.mkdir(parents=True, exist_ok=True) + + base: dict = {"model": args.variant, "language": "English"} + if args.variant == "custom_voice": + base["voice"] = args.voice + elif args.variant == "voice_design": + base["instructions"] = "A clear, friendly adult female voice with a neutral accent." + else: + if not args.ref_audio: + sys.exit("--ref-audio is required for the base variant") + base["ref_audio"] = "data:audio/wav;base64," + base64.b64encode(Path(args.ref_audio).read_bytes()).decode() + if args.ref_text: + base["ref_text"] = args.ref_text + else: + base["x_vector_only_mode"] = True + + report: dict = {"variant": args.variant, "cases": {}} + failures = [] + + def run(name: str, payload: dict, stream: bool, min_seconds: float, max_seconds: float) -> bytes: + try: + pcm, ttfa, total = speech(args.url, payload, stream, args.timeout) + except requests.RequestException as exc: + failures.append(f"{name}: {exc}") + report["cases"][name] = {"error": str(exc)} + print(f"{name:28s} FAILED: {exc}") + return b"" + info = {**stats(pcm), "ttfa_s": ttfa, "total_s": total, "stream": stream} + report["cases"][name] = info + write_wav(out / f"{args.variant}_{name}.wav", pcm) + problems = [] + if not (min_seconds <= info["seconds"] <= max_seconds): + problems.append(f"length {info['seconds']:.2f}s outside [{min_seconds}, {max_seconds}]") + if info["peak"] < 0.05 or info["rms"] < 0.005: + problems.append(f"near-silent audio (peak {info['peak']:.3f}, rms {info['rms']:.4f})") + if problems: + failures.append(f"{name}: " + "; ".join(problems)) + print( + f"{name:28s} {info['seconds']:6.2f}s audio peak {info['peak']:.2f} " + f"ttfa {ttfa * 1000:6.0f} ms total {total:.2f}s" + ) + return pcm + + # 1. streamed sentence: 5-10 s of speech, first audio quickly. + run("stream", {**base, "input": TEXT}, stream=True, min_seconds=3.0, max_seconds=12.0) + # 2. non-streaming container response. + run("blob", {**base, "input": TEXT}, stream=False, min_seconds=3.0, max_seconds=12.0) + # 3. greedy is repeatable byte for byte (same seed). + greedy = {**base, "input": TEXT, "do_sample": False, "subtalker_dosample": False, "seed": 7} + first = run("greedy_a", greedy, stream=True, min_seconds=3.0, max_seconds=12.0) + second = run("greedy_b", greedy, stream=True, min_seconds=3.0, max_seconds=12.0) + report["greedy_repeatable"] = first == second + if first != second: + failures.append("greedy runs with the same seed differ") + # 4. long input goes through sentence chunking (server side) and stays continuous. + run("long_chunked", {**base, "input": LONG_TEXT}, stream=True, min_seconds=25.0, max_seconds=90.0) + # 5. instruction control (1.7B CustomVoice style, VoiceDesign voice description). + if args.variant in ("custom_voice", "voice_design") and not args.no_instruct: + run("instruct", {**base, "input": TEXT, "instructions": "Whisper, very quietly."}, stream=True, + min_seconds=2.0, max_seconds=15.0) + # 6. error paths surface as 4xx, not hangs. + for name, bad in (("bad_voice", {**base, "input": TEXT, "voice": "nobody"}), + ("empty_input", {**base, "input": ""})): + try: + r = requests.post(f"{args.url}/v1/audio/speech", json=bad, timeout=60) + report["cases"][name] = {"status": r.status_code} + print(f"{name:28s} HTTP {r.status_code}") + if args.variant == "custom_voice" and name == "bad_voice" and r.status_code // 100 != 4: + failures.append(f"{name}: expected 4xx, got {r.status_code}") + if name == "empty_input" and r.status_code // 100 != 4: + failures.append(f"{name}: expected 4xx, got {r.status_code}") + except requests.RequestException as exc: + failures.append(f"{name}: {exc}") + + report["failures"] = failures + (out / f"{args.variant}_smoke.json").write_text(json.dumps(report, indent=2), encoding="utf-8") + print("SMOKE OK" if not failures else "SMOKE FAILED:\n " + "\n ".join(failures)) + sys.exit(0 if not failures else 1) + + +if __name__ == "__main__": + main() From 5b27bcd9ed934441be01cb44bcf50f939b18691e Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:12:09 -0700 Subject: [PATCH 067/110] test: parity harness tolerates specs without a config --- test/qwen3-tts/parity_qwen3_tts.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/qwen3-tts/parity_qwen3_tts.py b/test/qwen3-tts/parity_qwen3_tts.py index a86813e88..ef522f437 100644 --- a/test/qwen3-tts/parity_qwen3_tts.py +++ b/test/qwen3-tts/parity_qwen3_tts.py @@ -201,7 +201,8 @@ def __init__(self, model: Qwen3TTSModel, talker, device: str, deployment: str | apply_yaml_overrides(specs, yaml.safe_load(Path(deployment).read_text(encoding="utf-8"))) for spec in specs: - if hasattr(spec, "apply_yaml_overrides") and hasattr(spec.config, "max_num_pages"): + config = getattr(spec, "config", None) # sampler specs carry none + if config is not None and hasattr(config, "max_num_pages"): spec.apply_yaml_overrides(max_num_pages=max_num_pages) groups = JointGroups(tp_group=CommGroup.trivial(), sp_group=CommGroup.trivial()) transfer = TransferEngineInfo("h", "h", LocalTransferEngine("h")) From f5ac4e22b5005f772f74f4e0849274ad8e98ec5e Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:12:09 -0700 Subject: [PATCH 068/110] test: 1.7B real-weight layouts for the five-field prompt layout and a clone prefill check --- .../test_qwen3_tts_1p7b_real_weights.py | 63 ++++++++++++++++++- 1 file changed, 61 insertions(+), 2 deletions(-) diff --git a/test/integration/test_qwen3_tts_1p7b_real_weights.py b/test/integration/test_qwen3_tts_1p7b_real_weights.py index 28346ad6c..cc0d9dbcc 100644 --- a/test/integration/test_qwen3_tts_1p7b_real_weights.py +++ b/test/integration/test_qwen3_tts_1p7b_real_weights.py @@ -90,7 +90,8 @@ def test_prefill_layout_matches_variant(loaded): variant, model, talker, _ = loaded kwargs = {"input_modalities": ["text"], "output_modalities": ["audio"]} if variant == "base": - with pytest.raises(ValueError, match="reference audio"): + # Base clones a voice: a text-only request has no reference clip. + with pytest.raises(ValueError, match="reference clip"): model.process_prompt("Testing the base model.", **kwargs) return if variant == "custom_voice": @@ -104,8 +105,9 @@ def test_prefill_layout_matches_variant(loaded): "Testing Qwen three TTS.", instruct="A calm male voice.", **kwargs, ) assert tensors["speaker_id"][0].item() == -1 - instruct_len, text_len, stream_text = tensors["prompt_layout"][0].tolist() + instruct_len, text_len, stream_text, ref_text_len, ref_frames = tensors["prompt_layout"][0].tolist() assert instruct_len > 0 and text_len > 0 and stream_text == 0 + assert ref_text_len == 0 and ref_frames == 0 prepared = talker.prepare_inputs( "talker_prefill", SimpleNamespace(request_id=f"prefill-{variant}"), tensors, @@ -132,3 +134,60 @@ def test_depth_loop_runs_through_projection(loaded): assert (codes[:, 0] == layer0).all() assert (codes[:, 1:] < model.config.code_predictor.vocab_size).all() assert embed_sum.shape == (batch, model.config.talker.hidden_size) + + +def test_base_reference_clip_becomes_xvector_frames_and_clone_prefill(loaded): + """Base only: the shared reference clip runs through load_audio -> RefEncoder -> clone prefill. + + ``QWEN3_TTS_REF_AUDIO`` names a 24 kHz mono clip (the benchmark's + ``clone_2.wav``); its transcript is fixed here. + """ + variant, model, talker, _ = loaded + if variant != "base": + pytest.skip("voice cloning is a Base feature") + clip_path = os.environ.get("QWEN3_TTS_REF_AUDIO") + if not clip_path or not Path(clip_path).is_file(): + pytest.skip("set QWEN3_TTS_REF_AUDIO to a reference clip") + from mstar.model.submodule_base import ModelInputsFromEngine + + clip = model.load_audio(clip_path, "cuda:0") + assert clip.metadata["sample_rate"] == 24000 and clip.data.ndim == 1 + frames = model.config.codec.frames_for_samples(clip.data.shape[0]) + tensors = model.process_prompt( + "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye.", + input_modalities=["audio", "text"], output_modalities=["audio"], + tensors={"audio_inputs": [clip.data]}, + ref_text="Okay. Yeah. I resent you. I love you. I respect you. " + "But you know what? You blew it! And thanks to you.", + ) + assert tensors["prompt_layout"][0].tolist()[4] == frames and tensors["ref_frames"][0].item() == frames + + ref_encoder = model.get_submodule("RefEncoder", device="cuda:0", autocast_dtype=torch.bfloat16) + prepared = ref_encoder.prepare_inputs( + "talker_prefill_clone", SimpleNamespace(request_id="clone"), + {"audio_inputs": [clip.data], "prompt_layout": tensors["prompt_layout"]}, + ) + engine_inputs = ModelInputsFromEngine(request_ids=["clone"], per_request_info={}) + with torch.no_grad(): + encoded = ref_encoder.forward( + "talker_prefill_clone", engine_inputs, + **ref_encoder.preprocess("talker_prefill_clone", engine_inputs, [prepared]), + ) + torch.cuda.synchronize() + xvec, codes = encoded["speaker_embed"][0], encoded["ref_codes"][0] + assert xvec.shape == (model.config.talker.hidden_size,) and torch.isfinite(xvec.float()).all() + assert codes.shape == (frames, model.config.num_code_groups) + assert (codes >= 0).all() and (codes < model.config.codec.codebook_size).all() + + prefill = talker.prepare_inputs( + "talker_prefill_clone", SimpleNamespace(request_id="clone-prefill"), + {**tensors, "speaker_embed": [xvec], "ref_codes": [codes]}, + ) + # role(3) + [think, think_bos, lang?, think_eos, x-vector, pad] + in-context span (streaming text default) + assert prefill.input_embeds.shape[1] == model.config.talker.hidden_size + assert prefill.input_seq_len > 3 + 5 + frames + state = talker.request_state("clone-prefill") + assert torch.equal(state["reference_frames"], codes) + model._submodule_cache.pop("RefEncoder", None) + del ref_encoder + torch.cuda.empty_cache() From 5396fdb65f781d4f7fb4dc22bc5945e5cb403651 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:14:59 -0700 Subject: [PATCH 069/110] qwen3_tts: load codec quantizer buffers and check coverage against the state dict --- mstar/model/qwen3_tts/qwen3_tts_model.py | 86 +++++++++++++++--------- 1 file changed, 54 insertions(+), 32 deletions(-) diff --git a/mstar/model/qwen3_tts/qwen3_tts_model.py b/mstar/model/qwen3_tts/qwen3_tts_model.py index b7ed98a72..b5f9dd5f4 100644 --- a/mstar/model/qwen3_tts/qwen3_tts_model.py +++ b/mstar/model/qwen3_tts/qwen3_tts_model.py @@ -215,9 +215,10 @@ def _checkpoint_keys(checkpoint_dir: str | Path, prefix: str) -> set[str]: def _expected_checkpoint_keys(module: torch.nn.Module) -> set[str]: - """Checkpoint keys an M* module consumes, expanding fused projections.""" + """Checkpoint keys an M* module consumes: its state dict (parameters and + persistent buffers), with fused projections expanded to their shards.""" expected: set[str] = set() - for name in dict(module.named_parameters()): + for name in module.state_dict(): for fused, sources in _FUSED_SOURCES.items(): if f".{fused}." in name: expected.update(name.replace(f".{fused}.", f".{source}.") for source in sources) @@ -227,25 +228,48 @@ def _expected_checkpoint_keys(module: torch.nn.Module) -> set[str]: return expected +def _load_buffers( + module: torch.nn.Module, + weights, +) -> set[str]: + """Copy checkpoint tensors into the module's persistent buffers. + + ``load_hf_weights`` only fills parameters. Codec quantizers keep their + codebooks in buffers (``embed_sum`` / ``cluster_usage``), so a checkpoint + must be able to refill those too. Returns the buffer names it filled. + """ + persistent = set(module.state_dict()) - set(dict(module.named_parameters())) + buffers = {name: buf for name, buf in module.named_buffers() if name in persistent} + loaded: set[str] = set() + for name, tensor in weights: + target = buffers.get(name) + if target is None: + continue + target.copy_(tensor.to(device=target.device, dtype=target.dtype)) + loaded.add(name) + return loaded + + def _verify_checkpoint_coverage( module: torch.nn.Module, loaded: set[str], checkpoint_keys: set[str], component: str, ) -> None: - """Fail startup on any parameter left uninitialized or any key left unused. + """Fail startup on any state left uninitialized or any key left unused. - Both directions matter: a missing key means random weights would serve - requests; an unused key means the checkpoint carries a component this - port silently ignores (the 1.7B code predictor projection, for example). + Both directions matter: a missing key means random weights (or default + buffers) would serve requests; an unused key means the checkpoint carries + a component this port silently ignores (the 1.7B code predictor + projection, or a quantizer codebook kept in buffers, for example). """ - expected = set(dict(module.named_parameters())) + expected = set(module.state_dict()) missing = sorted(expected - loaded) if missing: preview = ", ".join(missing[:8]) raise RuntimeError( f"{component} checkpoint did not initialize {len(missing)} " - f"parameters: {preview}" + f"tensors: {preview}" ) unused = sorted( key for key in checkpoint_keys - _expected_checkpoint_keys(module) @@ -1061,13 +1085,12 @@ def _create_codec_submodule(self, device: str) -> NodeSubmodule: codec_dir = Path(self.local_dir) / "speech_tokenizer" prefix = "decoder." - weights = ( - (name.removeprefix(prefix), tensor) - for name, tensor in iter_safetensors_shards( - codec_dir, device=device, prefix=prefix - ) - ) - loaded = load_hf_weights(decoder, weights) + + def weights(): + for name, tensor in iter_safetensors_shards(codec_dir, device=device, prefix=prefix): + yield name.removeprefix(prefix), tensor + + loaded = load_hf_weights(decoder, weights()) | _load_buffers(decoder, weights()) _verify_checkpoint_coverage( decoder, loaded, _checkpoint_keys(codec_dir, prefix), "Qwen3-TTS Codec" ) @@ -1098,13 +1121,12 @@ def _create_ref_encoder_submodule( speaker_encoder = speaker_encoder.to(autocast_dtype) speaker_encoder = speaker_encoder.to(device=device) prefix = "speaker_encoder." - loaded = load_hf_weights( - speaker_encoder, - ( - (name.removeprefix(prefix), tensor) - for name, tensor in iter_safetensors_shards(self.local_dir, device=device, prefix=prefix) - ), - ) + + def speaker_weights(): + for name, tensor in iter_safetensors_shards(self.local_dir, device=device, prefix=prefix): + yield name.removeprefix(prefix), tensor + + loaded = load_hf_weights(speaker_encoder, speaker_weights()) _verify_checkpoint_coverage( speaker_encoder, loaded, _checkpoint_keys(self.local_dir, prefix), "Qwen3-TTS speaker encoder" ) @@ -1112,20 +1134,20 @@ def _create_ref_encoder_submodule( # Mimi encoder of the speech tokenizer: reference clip -> codec frames. # Float32 like the decoder; built on the CPU for its non-persistent - # buffers (rotary tables, convolution geometry). + # buffers (rotary tables, convolution geometry). Its quantizer keeps + # the codebooks in persistent buffers, which the checkpoint refills. _, _, encoder_cls = _load_qwen3_tts_codec_classes() codec_encoder = encoder_cls(MimiConfig(**self.config.codec.encoder_config)).to(device=device) codec_dir = Path(self.local_dir) / "speech_tokenizer" - prefix = "encoder." - loaded = load_hf_weights( - codec_encoder, - ( - (name.removeprefix(prefix), tensor) - for name, tensor in iter_safetensors_shards(codec_dir, device=device, prefix=prefix) - ), - ) + encoder_prefix = "encoder." + + def encoder_weights(): + for name, tensor in iter_safetensors_shards(codec_dir, device=device, prefix=encoder_prefix): + yield name.removeprefix(encoder_prefix), tensor + + loaded = load_hf_weights(codec_encoder, encoder_weights()) | _load_buffers(codec_encoder, encoder_weights()) _verify_checkpoint_coverage( - codec_encoder, loaded, _checkpoint_keys(codec_dir, prefix), "Qwen3-TTS codec encoder" + codec_encoder, loaded, _checkpoint_keys(codec_dir, encoder_prefix), "Qwen3-TTS codec encoder" ) codec_encoder.eval() return RefEncoderSubmodule( From 316ac5a816285694ea81fc2f71ea3ef24345eee1 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:16:11 -0700 Subject: [PATCH 070/110] test: VoiceDesign real-weight layout check with an explicit language --- test/integration/test_qwen3_tts_1p7b_real_weights.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/test/integration/test_qwen3_tts_1p7b_real_weights.py b/test/integration/test_qwen3_tts_1p7b_real_weights.py index cc0d9dbcc..fa6b1ed01 100644 --- a/test/integration/test_qwen3_tts_1p7b_real_weights.py +++ b/test/integration/test_qwen3_tts_1p7b_real_weights.py @@ -102,7 +102,7 @@ def test_prefill_layout_matches_variant(loaded): assert tensors["speaker_id"][0].item() == model.config.talker.spk_id["vivian"] else: tensors = model.process_prompt( - "Testing Qwen three TTS.", instruct="A calm male voice.", **kwargs, + "Testing Qwen three TTS.", language="English", instruct="A calm male voice.", **kwargs, ) assert tensors["speaker_id"][0].item() == -1 instruct_len, text_len, stream_text, ref_text_len, ref_frames = tensors["prompt_layout"][0].tolist() @@ -113,8 +113,10 @@ def test_prefill_layout_matches_variant(loaded): "talker_prefill", SimpleNamespace(request_id=f"prefill-{variant}"), tensors, ) assert prepared.input_embeds.shape[1] == model.config.talker.hidden_size - # instruct + role(3) + codec tags + (text + eos) + closing pad/bos - tags = 3 + 1 + (1 if variant == "custom_voice" else 0) + 1 # think..., [speaker], pad + # instruct + role(3) + codec tags + (text + eos) + closing pad/bos; with an + # explicit language the tags are think, think_bos, language, think_eos, + # [speaker], pad. + tags = 3 + 1 + (1 if variant == "custom_voice" else 0) + 1 assert prepared.input_seq_len == instruct_len + 3 + tags + (text_len + 1) + 1 state = talker.request_state(f"prefill-{variant}") assert state["trailing_text_hidden"].shape == (0, model.config.talker.hidden_size) From 3a9c156a212a44f7c0f6d12f678428c54bcf3609 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:20:06 -0700 Subject: [PATCH 071/110] test: parity harness keeps whole frames in the greedy loop --- test/qwen3-tts/parity_qwen3_tts.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/test/qwen3-tts/parity_qwen3_tts.py b/test/qwen3-tts/parity_qwen3_tts.py index ef522f437..6460e0b8e 100644 --- a/test/qwen3-tts/parity_qwen3_tts.py +++ b/test/qwen3-tts/parity_qwen3_tts.py @@ -313,12 +313,14 @@ def forward(engine_inputs, **kw): return talker.forward(fwd.graph_walk, engine_inputs, **kw) out = driver.step("talker_prefill", fwd, tensors, forward) - codes.append(out["codec_tokens"][0][0]) + # ``codec_tokens`` is the list of stream items (one ``[groups]`` frame each; + # the clone prefill leads with the reference frames): keep the new frame. + codes.append(out["codec_tokens"][-1]) talker.postprocess(rid, fwd, out) eos = talker.talker_config.codec_eos_token_id while len(codes) < frames and int(codes[-1][0]) != eos: out = driver.step("talker_decode", fwd, {"talker_input_embeds": out["talker_input_embeds"]}, forward) - codes.append(out["codec_tokens"][0][0]) + codes.append(out["codec_tokens"][-1]) talker.postprocess(rid, fwd, out) driver.close_request(rid) codes = torch.stack(codes) @@ -411,6 +413,7 @@ def main(argv: list[str] | None = None) -> None: args.voice = None torch.manual_seed(0) + torch.set_grad_enabled(False) snapshot = resolve_snapshot(args.repo) t0 = time.perf_counter() ref = load_reference(snapshot, args.device) From 33ada0c34715b7312b04aef72c1afd1b8dfb5498 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:24:51 -0700 Subject: [PATCH 072/110] test: parity verdict on confident-position agreement and relative logit error --- test/qwen3-tts/parity_qwen3_tts.py | 42 +++++++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/test/qwen3-tts/parity_qwen3_tts.py b/test/qwen3-tts/parity_qwen3_tts.py index 6460e0b8e..010f8a810 100644 --- a/test/qwen3-tts/parity_qwen3_tts.py +++ b/test/qwen3-tts/parity_qwen3_tts.py @@ -21,6 +21,14 @@ the codes are compared frame by frame; both code sequences are decoded to audio by the codec on each side and the waveform max-abs-diff is reported. +Verdict: the Talker and CodePredictor logits must differ from the reference +by less than 2% of the logit scale and agree on every confident position +(reference top-2 margin above 1.0 logit), and the codec must decode the same +codes to the same waveform (max-abs-diff below 1e-3). Raw argmax agreement +and the greedy divergence point are reported but not gated: bf16 with +different attention kernels flips near-ties, and one flipped code changes +every later frame of an autoregressive run. + Run inside the GPU allocation (weights must already be in the HF cache):: python test/qwen3-tts/parity_qwen3_tts.py --repo Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \\ @@ -332,25 +340,43 @@ def forward(engine_inputs, **kw): # --------------------------------------------------------------------------- +CONFIDENT_MARGIN = 1.0 # reference top-2 logit gap above which bf16 noise cannot flip the argmax + + def compare_logits(name: str, ours: torch.Tensor, theirs: torch.Tensor) -> dict[str, Any]: + """Logit-level agreement between M* and the reference over the same inputs. + + Both sides run bf16 with different attention kernels and matmul orders, + so raw argmax agreement is bounded by near-ties. The decisive numbers are + the mean difference relative to the logit scale and the agreement on + *confident* positions (reference top-2 margin above ``CONFIDENT_MARGIN``), + which an implementation bug would break and numerical noise cannot. + """ ours = ours.float() theirs = theirs.float() diff = (ours - theirs).abs() agree = (ours.argmax(-1) == theirs.argmax(-1)).float() top2 = theirs.topk(2, dim=-1).values margin = (top2[..., 0] - top2[..., 1]) + confident = margin > CONFIDENT_MARGIN return { "name": name, "positions": int(agree.numel()), "argmax_agreement": float(agree.mean()), + "confident_positions": int(confident.sum()), + "confident_agreement": float(agree[confident].mean()) if confident.any() else None, "max_abs_diff": float(diff.max()), "mean_abs_diff": float(diff.mean()), "ref_logit_scale": float(theirs.abs().mean()), + "rel_mean_diff": float(diff.mean() / theirs.abs().mean()), # disagreements should sit on near-ties: report the reference top-2 # margin where the argmax differs "disagreement_margin_median": ( float(margin[agree == 0].median()) if (agree == 0).any() else None ), + "disagreement_margin_max": ( + float(margin[agree == 0].max()) if (agree == 0).any() else None + ), } @@ -548,9 +574,19 @@ def main(argv: list[str] | None = None) -> None: if args.json: Path(args.json).parent.mkdir(parents=True, exist_ok=True) Path(args.json).write_text(json.dumps(report, indent=2), encoding="utf-8") - ok = talker_report["argmax_agreement"] >= 0.99 and cp_report["argmax_agreement"] >= 0.99 \ - and codec_report["max_abs_diff"] < 1e-3 - print("PARITY OK" if ok else "PARITY FAILED", file=sys.stderr) + def faithful(report: dict[str, Any]) -> bool: + confident = report["confident_agreement"] + return report["rel_mean_diff"] < 0.02 and (confident is None or confident >= 0.995) + + ok = faithful(talker_report) and faithful(cp_report) and codec_report["max_abs_diff"] < 1e-3 + verdict = "PARITY OK" if ok else "PARITY FAILED" + print( + f"{verdict}: talker rel diff {talker_report['rel_mean_diff']:.4f}, confident agreement " + f"{talker_report['confident_agreement']}; code predictor rel diff {cp_report['rel_mean_diff']:.4f}, " + f"confident agreement {cp_report['confident_agreement']}; " + f"codec max-abs-diff {codec_report['max_abs_diff']:.2e}", + file=sys.stderr, + ) sys.exit(0 if ok else 1) From e7866d5db10921f33d562d8685abe117c6bf3e44 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:24:51 -0700 Subject: [PATCH 073/110] benchmark: report confident agreement and relative logit error in the parity table --- benchmark/tts_report.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/benchmark/tts_report.py b/benchmark/tts_report.py index ff383bbf1..08bf3542a 100644 --- a/benchmark/tts_report.py +++ b/benchmark/tts_report.py @@ -51,9 +51,15 @@ def parity_rows(results: Path) -> list[str]: f"cos {clone['xvector_cosine']:.4f}, codes {clone.get('ref_code_agreement', float('nan')):.3f}" if clone else "-" ) + def agreement(report: dict) -> str: + confident = report.get("confident_agreement") + confident_cell = f"{confident:.4f}" if confident is not None else "n/a" + rel = report.get("rel_mean_diff", float("nan")) + return f"{report['argmax_agreement']:.4f} / {confident_cell} / {rel:.4f}" + rows.append( f"| {r['repo'].split('/')[-1]} | {re.sub(r'^parity_', '', path.stem)} | {r['frames']} | " - f"{r['talker']['argmax_agreement']:.4f} | {r['code_predictor']['argmax_agreement']:.4f} | " + f"{agreement(r['talker'])} | {agreement(r['code_predictor'])} | " f"{r['greedy_codes']['identical_frames_before_divergence']}/{r['greedy_codes']['frames_compared']} | " f"{r['codec']['max_abs_diff']:.2e} | {r['audio']['max_abs_diff']:.3f} | {clone_cell} |" ) @@ -71,7 +77,8 @@ def main(argv: list[str] | None = None) -> None: "| System (version) | concurrency | TTFA p50 / p95 ms | RTF | audio-s / s | WER % | errors | notes |", "|---|---|---|---|---|---|---|---|", *benchmark_rows(results), "", "## Parity vs qwen-tts (greedy, bf16 Talker, fp32 codec)", "", - "| checkpoint | mode | frames | Talker argmax agreement | CodePredictor argmax agreement | " + "| checkpoint | mode | frames | Talker argmax / confident agreement / rel. logit diff | " + "CodePredictor argmax / confident agreement / rel. logit diff | " "identical greedy frames | codec max-abs-diff | greedy audio max-abs-diff | " "clone (x-vector cosine, ref-code agreement) |", "|---|---|---|---|---|---|---|---|---|", *parity_rows(results)] From 61d6bfaf3c4915f9e31d3e7996a4cd56d63b8902 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:29:07 -0700 Subject: [PATCH 074/110] test: clone parity compares codec encoders in both dtypes and shares reference codes --- test/qwen3-tts/parity_qwen3_tts.py | 60 +++++++++++++++++++++++------- 1 file changed, 47 insertions(+), 13 deletions(-) diff --git a/test/qwen3-tts/parity_qwen3_tts.py b/test/qwen3-tts/parity_qwen3_tts.py index 010f8a810..4b208fe30 100644 --- a/test/qwen3-tts/parity_qwen3_tts.py +++ b/test/qwen3-tts/parity_qwen3_tts.py @@ -401,19 +401,45 @@ def decode_audio(codec, codes: torch.Tensor) -> torch.Tensor: return wav.squeeze().float() +_FP32_TOKENIZER = {} + + +def reference_fp32_tokenizer(snapshot: str, device: str): + """The reference speech tokenizer in float32 (M* runs its codec in float32; + the reference wrapper would otherwise inherit the Talker's bf16).""" + key = (snapshot, device) + if key not in _FP32_TOKENIZER: + from qwen_tts import Qwen3TTSTokenizer + + _FP32_TOKENIZER[key] = Qwen3TTSTokenizer.from_pretrained( + str(Path(snapshot) / "speech_tokenizer"), device_map=device, dtype=torch.float32, + ) + return _FP32_TOKENIZER[key] + + @torch.no_grad() def reference_decode_audio(snapshot: str, device: str, codes: torch.Tensor) -> torch.Tensor: - """Reference codec in float32 (M* runs its codec in float32; the reference - wrapper would otherwise inherit the Talker's bf16).""" - from qwen_tts import Qwen3TTSTokenizer - - tokenizer = Qwen3TTSTokenizer.from_pretrained( - str(Path(snapshot) / "speech_tokenizer"), device_map=device, dtype=torch.float32, - ) - wavs, _ = tokenizer.decode([{"audio_codes": codes}]) + wavs, _ = reference_fp32_tokenizer(snapshot, device).decode([{"audio_codes": codes}]) return torch.as_tensor(wavs[0]).float() +@torch.no_grad() +def reference_encode_fp32(snapshot: str, device: str, clip_path: str) -> torch.Tensor: + """Reference encoder (float32) codes ``[frames, groups]`` for the clip.""" + tokenizer = reference_fp32_tokenizer(snapshot, device) + return tokenizer.encode(clip_path).audio_codes[0].to(device) + + +def code_agreement(ours: torch.Tensor, theirs: torch.Tensor) -> dict[str, Any]: + n = min(ours.shape[0], theirs.shape[0]) + same = (ours[:n] == theirs[:n]) + return { + "frames": n, + "all_groups": float(same.float().mean()), + "per_group": [round(float(v), 3) for v in same.float().mean(dim=0)], + } + + def main(argv: list[str] | None = None) -> None: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--repo", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice") @@ -484,8 +510,6 @@ def main(argv: list[str] | None = None) -> None: "talker_prefill_clone", ModelInputsFromEngine(request_ids=["clone"], per_request_info={}), **ref_encoder.preprocess("talker_prefill_clone", None, [prepared]), ) - tensors["speaker_embed"] = encoded["speaker_embed"] - tensors["ref_codes"] = encoded["ref_codes"] _, item = clone their_xvec = item.ref_spk_embedding.to(args.device).float() our_xvec = encoded["speaker_embed"][0].float() @@ -494,15 +518,25 @@ def main(argv: list[str] | None = None) -> None: "xvector_max_abs_diff": float((our_xvec - their_xvec).abs().max()), "xvector_scale": float(their_xvec.abs().mean()), } + tensors["speaker_embed"] = encoded["speaker_embed"] + tensors["ref_codes"] = encoded["ref_codes"] if item.ref_code is not None: - their_codes = item.ref_code.to(args.device) + # M*'s float32 Mimi encoder vs the reference's own codes (bf16, the + # default dtype it inherits) and vs the reference encoder in float32. our_codes = encoded["ref_codes"][0] - n = min(their_codes.shape[0], our_codes.shape[0]) + their_codes = item.ref_code.to(args.device) + fp32_codes = reference_encode_fp32(snapshot, args.device, args.ref_audio) clone_report.update({ "ref_frames_mstar": int(our_codes.shape[0]), "ref_frames_reference": int(their_codes.shape[0]), - "ref_code_agreement": float((our_codes[:n] == their_codes[:n]).float().mean()), + "ref_code_agreement": code_agreement(our_codes, their_codes)["all_groups"], + "ref_codes_vs_reference_bf16": code_agreement(our_codes, their_codes), + "ref_codes_vs_reference_fp32": code_agreement(our_codes, fp32_codes), + "reference_fp32_vs_bf16": code_agreement(fp32_codes, their_codes), }) + # The Talker comparison must see the same in-context frames on both + # sides, so M* is fed the reference's own codes from here on. + tensors["ref_codes"] = [their_codes] del ref_encoder torch.cuda.empty_cache() From 0c4d7c03da77cda39cfb91aa9182eef2e60e7a7d Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:38:54 -0700 Subject: [PATCH 075/110] api: streaming speech surfaces an up-front engine error as its HTTP status --- mstar/api_server/openai/serving_speech.py | 38 ++++++++++++++++++++--- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/mstar/api_server/openai/serving_speech.py b/mstar/api_server/openai/serving_speech.py index 883318675..c56471130 100644 --- a/mstar/api_server/openai/serving_speech.py +++ b/mstar/api_server/openai/serving_speech.py @@ -10,12 +10,17 @@ or suppress it per request with ``sentence_chunking: true|false``. The next chunks are submitted while the current one streams, so the engine batches them and playback never waits for a prefill. + +A streaming request only commits to HTTP 200 once its first result chunk has +arrived and is not an error; an error chunk before that becomes the HTTP error +it carries (the non-streaming path gets the same from ``collect_results``). """ from __future__ import annotations from collections.abc import Callable +from fastapi import HTTPException from fastapi.responses import Response, StreamingResponse from mstar.api_server import media_io @@ -68,8 +73,15 @@ def submit(index: int) -> str: lookahead = max(1, int(getattr(adapter, "speech_chunk_lookahead", 2))) if req.stream: + pending = [submit(i) for i in range(min(lookahead, len(chunks)))] + # Look at the first result before committing to a 200: a request the + # engine rejects (bad voice, dead worker, ...) must surface as an HTTP + # error, not as an empty WAV. + first_iter = api.iter_result_chunks(pending[0]) + first = await anext(first_iter, None) + _raise_if_error(first) return StreamingResponse( - _stream_wav(api, submit, len(chunks), lookahead, sample_rate), + _stream_wav(api, submit, len(chunks), pending, first_iter, first, sample_rate), media_type="audio/wav", headers={"Cache-Control": "no-cache"}, ) @@ -85,13 +97,31 @@ def submit(index: int) -> str: return Response(content=audio_bytes, media_type=mime) -async def _stream_wav(api, submit: Callable[[int], str], num_chunks: int, lookahead: int, sample_rate: int): +def _raise_if_error(chunk) -> None: + """A data-worker failure arrives as an ``error`` chunk; turn it into the HTTP error it carries.""" + if chunk is not None and chunk.modality == "error": + raise HTTPException( + status_code=int((chunk.metadata or {}).get("status", 500)), + detail=chunk.data.decode("utf-8", "replace") if isinstance(chunk.data, bytes) else str(chunk.data), + ) + + +async def _stream_wav(api, submit: Callable[[int], str], num_chunks: int, pending: list[str], + first_iter, first, sample_rate: int): yield media_io.wav_stream_header(sample_rate) - pending: list[str] = [submit(i) for i in range(min(lookahead, num_chunks))] for index in range(num_chunks): if len(pending) < num_chunks: # Keep the next chunk generating while this one plays. pending.append(submit(len(pending))) - async for c in api.iter_result_chunks(pending[index]): + if index == 0: + iterator, head = first_iter, first + else: + iterator, head = api.iter_result_chunks(pending[index]), None + if head is not None and head.modality == "audio" and head.data: + yield head.data + async for c in iterator: + # Mid-stream the status is already sent; closing the stream is the + # only honest signal left, so raise rather than end quietly. + _raise_if_error(c) if c.modality == "audio" and c.data: yield c.data From edd23dccfff9b16ea8dff9eeb6589236b2068787 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:38:54 -0700 Subject: [PATCH 076/110] test: streaming speech returns the engine error status --- test/modular/test_speech_chunking.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/test/modular/test_speech_chunking.py b/test/modular/test_speech_chunking.py index 37a80e329..cc30b1760 100644 --- a/test/modular/test_speech_chunking.py +++ b/test/modular/test_speech_chunking.py @@ -207,3 +207,30 @@ def test_client_can_force_chunking_below_the_threshold(client_and_stub, monkeypa assert r.status_code == 200 and len(stub.submits) >= 2 assert " ".join(s["text"] for s in stub.submits) == text assert all(len(s["text"]) <= 60 for s in stub.submits) + + +def test_streaming_request_that_fails_up_front_returns_the_error_status(client_and_stub): + client, stub = client_and_stub + + def failing_submit(**kw): + stub.submits.append(kw) + stub._chunks[kw["request_id"]] = [ + _Chunk("error", b"Unsupported Qwen3-TTS speaker 'nobody'", {"status": 400}), + ] + return kw["request_id"] + + stub.submit_request = failing_submit + payload = {"model": "orpheus", "input": "hi there", "voice": "nobody", "stream": True} + r = client.post("/v1/audio/speech", json=payload) + assert r.status_code == 400 + assert "nobody" in r.json()["error"]["message"] + # and the non-streaming path keeps returning the error too + stub._chunks.clear() + + async def collect(request_id, raw_request=None): + from fastapi import HTTPException + raise HTTPException(status_code=400, detail="Unsupported Qwen3-TTS speaker 'nobody'") + + stub.collect_results = collect + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": "hi there", "voice": "nobody"}) + assert r.status_code == 400 From b1e29132438bc0e2f2b72a80b89807a5386eff69 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:51:23 -0700 Subject: [PATCH 077/110] test: parity reports token diversity and loudness of each side's greedy run --- test/qwen3-tts/parity_qwen3_tts.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/test/qwen3-tts/parity_qwen3_tts.py b/test/qwen3-tts/parity_qwen3_tts.py index 4b208fe30..87a87c74b 100644 --- a/test/qwen3-tts/parity_qwen3_tts.py +++ b/test/qwen3-tts/parity_qwen3_tts.py @@ -394,6 +394,22 @@ def compare_codes(ours: torch.Tensor, theirs: torch.Tensor) -> dict[str, Any]: } +def describe_codes(codes: torch.Tensor, audio: torch.Tensor, eos: int) -> dict[str, Any]: + """What one greedy run produced: token diversity and the loudness of its + audio. Greedy decoding of a codec LM can lock onto a repeated (silent) + frame; this tells silence apart from speech on each side independently.""" + group0 = codes[:, 0] + counts = torch.bincount(group0, minlength=1) + return { + "frames": int(codes.shape[0]), + "reached_eos": bool((group0 == eos).any()), + "unique_group0": int((counts > 0).sum()), + "top_group0_share": float(counts.max() / max(1, group0.numel())), + "audio_peak": float(audio.abs().max()) if audio.numel() else 0.0, + "audio_rms": float(audio.pow(2).mean().sqrt()) if audio.numel() else 0.0, + } + + @torch.no_grad() def decode_audio(codec, codes: torch.Tensor) -> torch.Tensor: """M* codec: ``[frames, groups]`` -> float waveform in [-1, 1].""" @@ -585,6 +601,9 @@ def main(argv: list[str] | None = None) -> None: "length_mismatch": int(audio_ref_codes_mstar.numel() - audio_ref_codes_ref.numel()), } audio_ours = decode_audio(codec, ours_codes[:n]) + eos = talker.talker_config.codec_eos_token_id + codes_report["mstar"] = describe_codes(ours_codes[:n], audio_ours, eos) + codes_report["reference"] = describe_codes(ref_codes[:n], audio_ref_codes_ref, eos) k = min(audio_ours.numel(), audio_ref_codes_ref.numel()) e2e_audio = { "name": "audio_greedy_e2e", From 60e596174c12cd431964381c0f72cedc05dda07a Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:55:24 -0700 Subject: [PATCH 078/110] qwen3_tts: memoise reference-clip conditioning by content --- mstar/model/qwen3_tts/submodules.py | 30 +++++++++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index de2a4ef64..b7089bbd1 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -33,6 +33,8 @@ from __future__ import annotations +import hashlib +from collections import OrderedDict from collections.abc import Callable, Mapping from typing import Any @@ -1189,10 +1191,16 @@ class RefEncoderSubmodule(NodeSubmodule): was built with (``disable_autocast``): the mel front end and the codec encoder stay in float32, the speaker encoder runs in the Talker's dtype as the reference does. + + The same clip recurs: a voice is reused across a session, and every + sentence chunk of a long clone request carries it. The conditioning is + memoised by clip content, so the encoders run once per distinct clip and a + repeat costs the prefill alone. """ disable_torch_compile = True disable_autocast = True + CONDITIONING_CACHE_SIZE = 64 def __init__( self, @@ -1206,6 +1214,14 @@ def __init__( self.mel_front_end = mel_front_end self.codec_encoder = codec_encoder self.config = config + # (clip digest, ref_frames) -> (speaker_embed, ref_codes), oldest first. + self._conditioning: OrderedDict[tuple[str, int], tuple[torch.Tensor, torch.Tensor]] = OrderedDict() + + @staticmethod + def clip_digest(clip: torch.Tensor) -> str: + """Content key of a reference clip (its float32 samples).""" + samples = clip.detach().to(device="cpu", dtype=torch.float32).contiguous() + return hashlib.blake2b(samples.numpy().tobytes(), digest_size=16).hexdigest() def prepare_inputs( self, @@ -1215,14 +1231,15 @@ def prepare_inputs( **kwargs: Any, ) -> NodeInputs: del graph_walk, fwd_info, kwargs - waveform = inputs["audio_inputs"][0].to(device=self.get_device(), dtype=torch.float32) + clip = inputs["audio_inputs"][0] + waveform = clip.to(device=self.get_device(), dtype=torch.float32) if waveform.ndim > 1: waveform = waveform.mean(dim=0) if waveform.shape[0] < waveform.shape[-1] else waveform.mean(dim=-1) layout = inputs["prompt_layout"][0].tolist() ref_frames = int(layout[4]) if len(layout) >= 5 else 0 return NodeInputs( tensor_inputs={"waveform": waveform.reshape(-1)}, - kwargs={"ref_frames": ref_frames}, + kwargs={"ref_frames": ref_frames, "clip_key": (self.clip_digest(clip), ref_frames)}, ) def forward( @@ -1231,9 +1248,14 @@ def forward( engine_inputs: ModelInputsFromEngine, waveform: torch.Tensor, ref_frames: int = 0, + clip_key: tuple[str, int] | None = None, **kwargs: Any, ) -> NameToTensorList: del graph_walk, engine_inputs, kwargs + if clip_key is not None and clip_key in self._conditioning: + self._conditioning.move_to_end(clip_key) + speaker_embed, codes = self._conditioning[clip_key] + return {"speaker_embed": [speaker_embed], "ref_codes": [codes]} device_type = waveform.device.type with torch.autocast(device_type=device_type, enabled=False): mels = self.mel_front_end(waveform.unsqueeze(0)) @@ -1255,4 +1277,8 @@ def forward( # x-vector-only clone: no in-context frames. The edge still needs # a tensor; the Talker ignores it because the layout says 0 frames. codes = torch.zeros(1, num_quantizers, dtype=torch.long, device=waveform.device) + if clip_key is not None: + self._conditioning[clip_key] = (speaker_embed, codes) + while len(self._conditioning) > self.CONDITIONING_CACHE_SIZE: + self._conditioning.popitem(last=False) return {"speaker_embed": [speaker_embed], "ref_codes": [codes]} From 9befd9bb3ab4403f6a23659c141fe7c627af6c0c Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:55:24 -0700 Subject: [PATCH 079/110] test: reference-clip conditioning memo --- test/modular/test_qwen3_tts_model.py | 52 +++++++++++++++++++++++++++- 1 file changed, 51 insertions(+), 1 deletion(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index d09989a24..f6ab6c8f2 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -1499,7 +1499,7 @@ def test_qwen3_tts_ref_encoder_emits_xvector_and_reference_frames(): {"audio_inputs": [clip], "prompt_layout": [torch.tensor([0, 2, 1, 5, 7])]}, ) assert prepared.tensor_inputs["waveform"].shape == (4000,) - assert prepared.kwargs == {"ref_frames": 7} + assert prepared.kwargs["ref_frames"] == 7 engine_inputs = ModelInputsFromEngine(request_ids=["clone"], per_request_info={}) out = submodule.forward("talker_prefill_clone", engine_inputs, **submodule.preprocess( "talker_prefill_clone", engine_inputs, [prepared])) @@ -1517,3 +1517,53 @@ def test_qwen3_tts_ref_encoder_emits_xvector_and_reference_frames(): "talker_prefill_clone", engine_inputs, [prepared])) assert out["ref_codes"][0].shape == (1, config.codec.num_quantizers) assert encoder.calls == 1 + + +def test_qwen3_tts_ref_encoder_memoises_conditioning_by_clip_content(): + from mstar.model.qwen3_tts.components.speaker_encoder import ( + Qwen3TTSMelFrontEnd, + Qwen3TTSSpeakerEncoder, + ) + from mstar.model.qwen3_tts.config import Qwen3TTSSpeakerEncoderConfig + from mstar.model.qwen3_tts.submodules import RefEncoderSubmodule + + config = _tiny_model_config() + speaker_config = Qwen3TTSSpeakerEncoderConfig( + enc_dim=config.talker.hidden_size, enc_channels=(16, 16, 16, 16, 48), + enc_se_channels=8, enc_attention_channels=8, + ) + encoder = _FakeCodecEncoder(config.codec.num_quantizers) + submodule = RefEncoderSubmodule( + Qwen3TTSSpeakerEncoder(speaker_config), Qwen3TTSMelFrontEnd(speaker_config), encoder, config, + ) + engine_inputs = ModelInputsFromEngine(request_ids=["clone"], per_request_info={}) + + def encode(rid: str, clip: torch.Tensor, ref_frames: int): + prepared = submodule.prepare_inputs( + "talker_prefill_clone", SimpleNamespace(request_id=rid), + {"audio_inputs": [clip], "prompt_layout": [torch.tensor([0, 2, 1, ref_frames, ref_frames])]}, + ) + return submodule.forward("talker_prefill_clone", engine_inputs, **submodule.preprocess( + "talker_prefill_clone", engine_inputs, [prepared])) + + clip_a = torch.randn(4000) * 0.1 + first = encode("a1", clip_a, 7) + assert encoder.calls == 1 + # The same samples again (a fresh tensor, as an upload produces) cost no encoder pass + # and yield the same conditioning. + again = encode("a2", clip_a.clone(), 7) + assert encoder.calls == 1 + assert again["speaker_embed"][0] is first["speaker_embed"][0] + assert torch.equal(again["ref_codes"][0], first["ref_codes"][0]) + # Other content, or the same clip used x-vector-only, is a different entry. + encode("b1", torch.randn(4000) * 0.1, 7) + assert encoder.calls == 2 + xvec = encode("a3", clip_a, 0) + assert encoder.calls == 2 and xvec["ref_codes"][0].shape == (1, config.codec.num_quantizers) + assert len(submodule._conditioning) == 3 + # The memo is bounded, oldest first. + submodule.CONDITIONING_CACHE_SIZE = 2 + encode("c1", torch.randn(4000) * 0.1, 7) + assert encoder.calls == 3 and len(submodule._conditioning) == 2 + encode("a4", clip_a, 7) # evicted, so encoded again + assert encoder.calls == 4 From 7dc9746ea62f401482fe1e121734e147d4ecc4b5 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:55:58 -0700 Subject: [PATCH 080/110] docs: reference clips are encoded once per distinct clip --- docs/models.rst | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index 0c9f1aa0f..274c934b6 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -110,9 +110,10 @@ Qwen3-TTS notes (VoiceDesign, required). Base clones a voice from one reference clip: on ``/v1/audio/speech`` pass ``ref_audio`` (data URL, URL, path or base64) plus ``ref_text`` (its transcript) or ``x_vector_only_mode: true``; with the SDK, - ``client.tts(text, reference_audio="ref.wav", ref_text="...")``. The clip is - used for that request only; named, persisted voices arrive with the shared - voice registry. + ``client.tts(text, reference_audio="ref.wav", ref_text="...")``. The clip's + conditioning (x-vector and codec frames) is memoised by content, so a voice + reused across requests, or by the sentence chunks of one long request, is + encoded once; named, persisted voices arrive with the shared voice registry. - Text layout follows the reference defaults: CustomVoice and VoiceDesign put the whole text in the prefill; Base feeds it one token per frame. Override per request with ``non_streaming_mode``. From 3171700cfc3176efce65aa02c5af0fc2903e9a6f Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:12:01 -0700 Subject: [PATCH 081/110] qwen3_tts: codec window geometry travels with the pass, not request state --- mstar/model/qwen3_tts/submodules.py | 32 ++++++++++++++++++++++------- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index b7089bbd1..a931b6eee 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -34,6 +34,7 @@ from __future__ import annotations import hashlib +import logging from collections import OrderedDict from collections.abc import Callable, Mapping from typing import Any @@ -86,6 +87,8 @@ # 1. TalkerSubmodule - autoregressive 12 Hz codec-frame generation # =========================================================================== +logger = logging.getLogger(__name__) + class TalkerSubmodule(ARNodeSubmodule): """Run text/voice prefill and produce one 16-code frame per AR step. @@ -785,6 +788,7 @@ def check_stop( not ignore_eos and token == self.talker_config.codec_eos_token_id ) if reached_eos or generated >= max_tokens: + logger.debug("talker %s: stop after %d frames (eos=%s)", request_id, generated, reached_eos) return {"talker_decode_loop"} return set() @@ -1024,13 +1028,17 @@ def prepare_inputs( bucket = self._bucket(max(int(meta.get("num_items", num_items)), 1)) if frames < bucket: codes = torch.nn.functional.pad(codes, (0, 0, 0, bucket - frames)) - state.add_all( - latest_codec_frames=frames, - latest_context_frames=min(context, frames), - codec_bucket=bucket, + logger.debug( + "codec %s: window items=%d context=%d frames=%d bucket=%d final=%s", + fwd_info.request_id, num_items, context, frames, bucket, meta.get("is_final"), ) + state.add("codec_bucket", bucket) # graph-key fallback when a caller has no stream metadata + # The window's geometry rides with this pass: under speculative + # scheduling the next window of the same request is prepared before + # this one is postprocessed, so request state would be overwritten. return ARNodeInputs( tensor_inputs={"codec_tokens": codes.t().contiguous()}, + kwargs={"frames": frames, "context": min(context, frames)}, ) def preprocess( @@ -1081,15 +1089,22 @@ def postprocess( request_id: str, request_info: CurrentForwardPassInfo, outputs: dict[str, list[torch.Tensor]], + inputs: ARNodeInputs | None = None, **kwargs: Any, ) -> None: - """Drop padding, repeated-context audio and (clone) reference audio before emission.""" + """Drop padding, repeated-context audio and (clone) reference audio before emission. + + ``inputs`` is this pass's ``prepare_inputs`` result (the engine hands it + back), which carries the window's frame and context counts. + """ del request_info, kwargs if "audio_chunk" not in outputs: return + if inputs is None: + raise ValueError("codec postprocess needs the pass's inputs for its window geometry") state = self.request_state(request_id) - frames = int(state.get("latest_codec_frames", 0)) - context = int(state.get("latest_context_frames", 0)) + frames = int(inputs.kwargs["frames"]) + context = int(inputs.kwargs["context"]) start = context * self.total_upsample end = frames * self.total_upsample skip = int(state.get("skip_samples", 0)) @@ -1098,6 +1113,9 @@ def postprocess( start += dropped state.add("skip_samples", skip - dropped) outputs["audio_chunk"][0] = outputs["audio_chunk"][0][start:end] + logger.debug( + "codec %s: emit frames %d..%d (%d samples)", request_id, context, frames, max(end - start, 0) + ) def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: """Batch codec requests only when their decoder input shapes match.""" From 9262ccbabde0180a05523528c8ca23707c07629c Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:12:01 -0700 Subject: [PATCH 082/110] test: codec trims each pass with its own window geometry --- test/modular/test_qwen3_tts_model.py | 53 +++++++++++++++++++--------- 1 file changed, 36 insertions(+), 17 deletions(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index f6ab6c8f2..314018cf1 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -1262,27 +1262,50 @@ def test_qwen3_tts_codec_trims_reported_context_audio(): config = _tiny_model_config() submodule = CodecSubmodule(_FakeCodecDecoder(4), config) assert submodule.windows == [1, 4] and submodule.max_window == 4 - state = submodule.request_state("request") # The stream buffer reports how many leading frames are repeated context; - # the first window has none, later ones up to left_context (2). - state.add_all(latest_codec_frames=5, latest_context_frames=0) + # the first window has none, later ones up to left_context (2). The + # geometry travels with the pass's inputs. first = {"audio_chunk": [torch.arange(20)]} - submodule.postprocess("request", None, first) + submodule.postprocess("request", None, first, inputs=_geometry(frames=5, context=0)) assert first["audio_chunk"][0].tolist() == list(range(20)) - state.add_all(latest_codec_frames=5, latest_context_frames=2) second = {"audio_chunk": [torch.arange(20)]} - submodule.postprocess("request", None, second) + submodule.postprocess("request", None, second, inputs=_geometry(frames=5, context=2)) assert second["audio_chunk"][0].tolist() == list(range(8, 20)) # Padding frames of a bucket never reach the client. - state.add_all(latest_codec_frames=3, latest_context_frames=1) padded = {"audio_chunk": [torch.arange(16)]} - submodule.postprocess("request", None, padded) + submodule.postprocess("request", None, padded, inputs=_geometry(frames=3, context=1)) assert padded["audio_chunk"][0].tolist() == list(range(4, 12)) +def _geometry(frames: int, context: int) -> ARNodeInputs: + return ARNodeInputs(kwargs={"frames": frames, "context": context}) + + +def test_qwen3_tts_codec_postprocess_uses_its_own_pass_geometry(): + """Speculative scheduling prepares a request's next window before the + current one is postprocessed; the trim must follow the pass, not the + request's latest state.""" + config = _tiny_model_config() # upsample 4, windows [1, 4] + submodule = CodecSubmodule(_FakeCodecDecoder(4), config) + meta = lambda context: SimpleNamespace( # noqa: E731 + request_id="request", + step_metadata={"stream_chunks": {"codec_tokens": {"context_items": context, "is_final": False}}}, + ) + first = submodule.prepare_inputs("codec_chunk", meta(0), {"codec_tokens": [torch.ones(1, 4, dtype=torch.long)]}) + second = submodule.prepare_inputs("codec_chunk", meta(1), {"codec_tokens": [torch.ones(4, 4, dtype=torch.long)]}) + assert (first.kwargs, second.kwargs) == ({"frames": 1, "context": 0}, {"frames": 4, "context": 1}) + + out_first = {"audio_chunk": [torch.arange(4)]} + submodule.postprocess("request", None, out_first, inputs=first) + assert out_first["audio_chunk"][0].tolist() == [0, 1, 2, 3] + out_second = {"audio_chunk": [torch.arange(16)]} + submodule.postprocess("request", None, out_second, inputs=second) + assert out_second["audio_chunk"][0].tolist() == list(range(4, 16)) + + def test_qwen3_tts_codec_trims_reference_audio_from_clone_streams(): config = _tiny_model_config() # upsample 4 samples per frame, chunk 3, left context 2 submodule = CodecSubmodule(_FakeCodecDecoder(4), config) @@ -1295,15 +1318,13 @@ def test_qwen3_tts_codec_trims_reference_audio_from_clone_streams(): assert state["skip_samples"] == 16 # First chunk: 3 frames = 12 samples, all reference -> nothing emitted. - assert state["latest_codec_frames"] == 3 and state["latest_context_frames"] == 0 first = {"audio_chunk": [torch.arange(20)]} - submodule.postprocess("clone", None, first) + submodule.postprocess("clone", None, first, inputs=_geometry(frames=3, context=0)) assert first["audio_chunk"][0].numel() == 0 assert state["skip_samples"] == 4 # Second chunk: 2 context + 3 new frames; 4 more samples belong to the reference. - state.add_all(latest_codec_frames=5, latest_context_frames=2) second = {"audio_chunk": [torch.arange(20)]} - submodule.postprocess("clone", None, second) + submodule.postprocess("clone", None, second, inputs=_geometry(frames=5, context=2)) assert second["audio_chunk"][0].tolist() == list(range(12, 20)) assert state["skip_samples"] == 0 @@ -1331,10 +1352,8 @@ def test_qwen3_tts_codec_filters_eos_and_pads_to_capture_shape(): assert packed.shape == (4, 4) assert packed[:, :2].t().tolist() == [[1, 2, 3, 4], [5, 6, 7, 8]] assert packed[:, 2:].count_nonzero().item() == 0 - state = submodule.request_state("request") - assert state["latest_codec_frames"] == 2 - assert state["latest_context_frames"] == 1 - assert state["codec_bucket"] == 4 + assert prepared.kwargs == {"frames": 2, "context": 1} + assert submodule.request_state("request")["codec_bucket"] == 4 # A single frame lands in the first ramp bucket; too many frames is an error. one = submodule.prepare_inputs( @@ -1390,7 +1409,7 @@ def test_qwen3_tts_streaming_policy_ramps_and_flushes_only_new_tail_audio(): prepared = codec.prepare_inputs("codec_chunk", fwd_info, {"codec_tokens": [tail_codes]}) assert prepared.tensor_inputs["codec_tokens"].shape == (4, 4) # 3 frames padded to the 4-frame bucket outputs = {"audio_chunk": [torch.arange(16)]} - codec.postprocess("request", None, outputs) + codec.postprocess("request", None, outputs, inputs=prepared) assert outputs["audio_chunk"][0].tolist() == list(range(8, 12)) From 5affa75184c67d6a059a4bbc1bd3ed9db936485b Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:14:11 -0700 Subject: [PATCH 083/110] benchmark: order report rows by system and concurrency, trim the environment dump --- benchmark/tts_report.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/benchmark/tts_report.py b/benchmark/tts_report.py index 08bf3542a..72671724d 100644 --- a/benchmark/tts_report.py +++ b/benchmark/tts_report.py @@ -23,11 +23,14 @@ def _load(path: Path) -> dict: def benchmark_rows(results: Path) -> list[str]: + reports = [ + _load(path) for path in sorted(results.glob("*_c[0-9]*.json")) + if not path.name.endswith("_wer.json") + ] rows = [] - for path in sorted(results.glob("*_c[0-9]*.json")): - if path.name.endswith("_wer.json"): - continue - report = _load(path) + # one system's rows in concurrency order, systems alphabetically + for report in sorted(reports, key=lambda r: (r.get("label") or r["engine"], int(r["concurrency"]))): + path = results / f"{report.get('label') or report['engine']}_c{report['concurrency']}.json" med = report["median_over_repeats"] wer_path = path.with_name(path.stem + "_wer.json") wer = f"{_load(wer_path)['wer_percent']:.2f}" if wer_path.is_file() else "n/a" @@ -84,7 +87,9 @@ def main(argv: list[str] | None = None) -> None: "|---|---|---|---|---|---|---|---|---|", *parity_rows(results)] env = results / "environment.txt" if env.is_file(): - lines += ["", "## Environment", "", "```", env.read_text(encoding="utf-8").strip(), "```"] + # versions, GPU and clocks; the raw nvidia-smi clock dump that follows is left out + summary = env.read_text(encoding="utf-8").split("==============NVSMI LOG")[0].strip() + lines += ["", "## Environment", "", "```", summary, "```"] text = "\n".join(lines) + "\n" print(text) if args.out: From 22de57de9ebc86e974a92e9cc038c6379cd4b63f Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:35:26 -0700 Subject: [PATCH 084/110] qwen3_tts: codec batches mixed windows and returns per-request samples from a single forward --- mstar/model/qwen3_tts/submodules.py | 35 ++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index a931b6eee..be7db75fd 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -1047,11 +1047,20 @@ def preprocess( engine_inputs: ModelInputsFromEngine, inputs: list[ARNodeInputs], ) -> dict[str, torch.Tensor]: - """Stack equal fixed-shape codec windows into one continuous batch.""" + """Stack the batch's windows, padded to its widest bucket. + + Requests reach the codec at different points of the chunk ramp, so one + batch mixes windows. The decoder is causal, so trailing padding never + changes the frames in front of it; each request's ``postprocess`` keeps + only its own frames. + """ del graph_walk, engine_inputs + windows = [item.tensor_inputs["codec_tokens"] for item in inputs] + width = max(window.shape[-1] for window in windows) return { "codec_tokens": torch.stack([ - item.tensor_inputs["codec_tokens"] for item in inputs + torch.nn.functional.pad(window, (0, width - window.shape[-1])) + for window in windows ]) } @@ -1067,8 +1076,9 @@ def forward( codec_tokens: torch.Tensor, **kwargs: Any, ) -> NameToTensorList: + """One request's window (``[1, quantizers, frames]``) -> its ``[samples]`` audio.""" del graph_walk, engine_inputs, kwargs - return {"audio_chunk": [self._decode(codec_tokens)]} + return {"audio_chunk": [self._decode(codec_tokens)[0]]} def forward_batched( self, @@ -1112,17 +1122,21 @@ def postprocess( dropped = min(skip, max(end - start, 0)) start += dropped state.add("skip_samples", skip - dropped) - outputs["audio_chunk"][0] = outputs["audio_chunk"][0][start:end] + audio = outputs["audio_chunk"][0].reshape(-1) # one request's samples, whatever the batch shape + if audio.numel() < end: + raise ValueError( + f"codec produced {audio.numel()} samples for a window of {frames} frames " + f"({end} expected)" + ) + outputs["audio_chunk"][0] = audio[start:end] logger.debug( "codec %s: emit frames %d..%d (%d samples)", request_id, context, frames, max(end - start, 0) ) def can_batch(self, batch: ExecutingBatch, model_inputs: list[NodeInputs]) -> bool: - """Batch codec requests only when their decoder input shapes match.""" + """Any mix of windows batches (``preprocess`` pads to the widest).""" del batch - return 0 < len(model_inputs) <= self.MAX_BATCH_SIZE and len({ - item.tensor_inputs["codec_tokens"].shape for item in model_inputs - }) == 1 + return 0 < len(model_inputs) <= self.MAX_BATCH_SIZE def max_batch_size(self, graph_walk: str) -> int: del graph_walk @@ -1133,7 +1147,7 @@ def cg_key_info( graph_walk: str, per_request_info: Mapping[str, CurrentForwardPassInfo], ) -> Any: - """The window bucket this batch pads to (``can_batch`` keeps it uniform). + """The widest window bucket in this batch, which ``preprocess`` pads to. Derived from the stream metadata the worker attaches to each request (available before ``prepare_inputs`` runs, so a pre-planned lease can @@ -1148,7 +1162,8 @@ def cg_key_info( buckets.add(self._bucket(max(int(num_items), 1))) else: buckets.add(self.request_state(request_id).get("codec_bucket")) - return buckets.pop() if len(buckets) == 1 else None + buckets.discard(None) + return max(buckets) if buckets else None def get_cuda_graph_configs( self, device: torch.device, tp_world_size: int = 1 From 10b293bd4c6869420a28224f150960c69d6edff4 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:35:26 -0700 Subject: [PATCH 085/110] test: codec mixed-window batching and single-forward output shape --- test/modular/test_qwen3_tts_model.py | 42 +++++++++++++++++++++++++--- 1 file changed, 38 insertions(+), 4 deletions(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index 314018cf1..9445a44c5 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -1464,20 +1464,54 @@ def meta(num_items): }}}) assert submodule.cg_key_info("codec_chunk", {"a": meta(3), "b": meta(4)}) == 4 - assert submodule.cg_key_info("codec_chunk", {"a": meta(1), "b": meta(4)}) is None + # Requests at different points of the ramp share a batch: the key (and the + # padding in preprocess) is the widest window among them. + assert submodule.cg_key_info("codec_chunk", {"a": meta(1), "b": meta(4)}) == 4 for rid in ("a", "b"): submodule.request_state(rid).add("codec_bucket", 4) assert submodule.cg_key_info("codec_chunk", {"a": None, "b": None}) == 4 submodule.request_state("b").add("codec_bucket", 1) - assert submodule.cg_key_info("codec_chunk", {"a": None, "b": None}) is None + assert submodule.cg_key_info("codec_chunk", {"a": None, "b": None}) == 4 - mixed = model_inputs + [ARNodeInputs(tensor_inputs={"codec_tokens": torch.zeros(4, 1, dtype=torch.long)})] - assert not submodule.can_batch(batch, mixed) + mixed = model_inputs + [ARNodeInputs(tensor_inputs={"codec_tokens": torch.ones(4, 1, dtype=torch.long)})] + assert submodule.can_batch(batch, mixed) + assert submodule.can_use_cuda_graphs(batch, mixed) + packed = submodule.preprocess( + "codec_chunk", ModelInputsFromEngine(request_ids=["a", "b", "c"], per_request_info={}), mixed, + ) + assert packed["codec_tokens"].shape == (3, 4, 4) + assert packed["codec_tokens"][2].tolist() == [[1, 0, 0, 0]] * 4 # 1-frame window padded on the right oversized = model_inputs * 9 assert len(oversized) == 18 assert not submodule.can_batch(batch, oversized) +def test_qwen3_tts_codec_single_forward_yields_one_request_samples(): + """The eager single-request path must hand postprocess a ``[samples]`` + tensor: slicing a ``[1, samples]`` batch on its first axis emitted empty + chunks whenever a window carried context (the truncation seen at c=8).""" + config = _tiny_model_config() # upsample 4, windows [1, 4] + submodule = CodecSubmodule(_FakeCodecDecoder(4), config) + engine_inputs = ModelInputsFromEngine(request_ids=["r"], per_request_info={}) + window = ARNodeInputs( + tensor_inputs={"codec_tokens": torch.ones(4, 4, dtype=torch.long)}, + kwargs={"frames": 4, "context": 2}, + ) + packed = submodule.preprocess("codec_chunk", engine_inputs, [window]) + out = submodule.forward("codec_chunk", engine_inputs, **packed) + assert out["audio_chunk"][0].shape == (16,) + submodule.postprocess("r", None, out, inputs=window) + assert out["audio_chunk"][0].shape == (8,) # frames 2..4 of 4 + + # A batch-shaped chunk is flattened rather than sliced on the batch axis; + # too few samples for the window is an error, never a silent cut. + batched = {"audio_chunk": [torch.arange(16).view(1, 16)]} + submodule.postprocess("r", None, batched, inputs=window) + assert batched["audio_chunk"][0].tolist() == list(range(8, 16)) + with pytest.raises(ValueError, match="samples"): + submodule.postprocess("r", None, {"audio_chunk": [torch.arange(8)]}, inputs=window) + + class _FakeCodecEncoder(torch.nn.Module): """Stands in for the Mimi encoder: deterministic codes, one frame per 4 samples.""" From 0ba9b8f1b62fecee75b880e7e43b7d347fe47133 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:38:54 -0700 Subject: [PATCH 086/110] qwen3_tts: codec batches up to 32 windows --- mstar/model/qwen3_tts/submodules.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index be7db75fd..a3862e496 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -942,8 +942,8 @@ class CodecSubmodule(ARNodeSubmodule): # Windows are at most chunk + left_context frames (50 by default, 4 s of # audio), so the decoder's activations stay small enough to capture # batches of 16 next to the Talker; ``can_batch`` keeps the ceiling. - MAX_BATCH_SIZE = 16 - CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16] + MAX_BATCH_SIZE = 32 + CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] def __init__(self, decoder: torch.nn.Module, config: Qwen3TTSModelConfig): super().__init__() From 36afc51926207f10ab8925e1dc71c9e73726aa23 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:38:54 -0700 Subject: [PATCH 087/110] test: codec batch limit of 32 --- test/modular/test_qwen3_tts_model.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index 9445a44c5..ed2b446ac 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -1445,7 +1445,7 @@ def test_qwen3_tts_codec_batches_and_declares_cuda_graphs(): assert graph_config.capture_graph_walk == "codec_chunk" # CustomVoice has no clone walk; a Base config would add codec_chunk_clone. assert set(graph_config.replay_graph_walks) == {"codec_chunk"} - assert graph_config.capture_batch_sizes == [1, 2, 4, 8, 16] + assert graph_config.capture_batch_sizes == [1, 2, 4, 8, 16, 32] assert graph_config.single_request_inputs.tensor_inputs["codec_tokens"].shape == ( 4, graph_config.additional_key_info, ) @@ -1455,7 +1455,7 @@ def test_qwen3_tts_codec_batches_and_declares_cuda_graphs(): assert set(base_codec.get_cuda_graph_configs(torch.device("cpu"))[0].replay_graph_walks) == { "codec_chunk", "codec_chunk_clone", } - assert submodule.max_batch_size("codec_chunk") == 16 + assert submodule.max_batch_size("codec_chunk") == 32 # The batch's capture key is the bucket its requests pad to: read off the # stream metadata when present (before prepare_inputs), else off the state. def meta(num_items): @@ -1481,8 +1481,8 @@ def meta(num_items): ) assert packed["codec_tokens"].shape == (3, 4, 4) assert packed["codec_tokens"][2].tolist() == [[1, 0, 0, 0]] * 4 # 1-frame window padded on the right - oversized = model_inputs * 9 - assert len(oversized) == 18 + oversized = model_inputs * 17 + assert len(oversized) == 34 assert not submodule.can_batch(batch, oversized) From fc540678e23a48186a9c116377be65dc22f7ddde Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:38:54 -0700 Subject: [PATCH 088/110] docs: codec graphs captured for batch sizes up to 32 --- docs/models.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/models.rst b/docs/models.rst index 274c934b6..bcd6813ad 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -125,7 +125,7 @@ Qwen3-TTS notes frames (320 ms of speech), later windows grow to 25 new frames behind 25 frames of already decoded left context (the reference's own ``chunked_decode`` context). Each window size is a CUDA-graph bucket - captured for batch sizes 1 to 16; the stream buffer reports how many leading + captured for batch sizes 1 to 32; the stream buffer reports how many leading frames of a window are repeated context, and the codec trims their audio. - Talker prefill remains eager because it runs once with variable sequence lengths. Decode always uses the whole-walk CUDA Graph, with the 15-step From f55e16ed236b4a52bf725741731bfb0cf9d3d65a Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:43:34 -0700 Subject: [PATCH 089/110] qwen3_tts: first codec window is a single frame --- mstar/model/qwen3_tts/config.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/mstar/model/qwen3_tts/config.py b/mstar/model/qwen3_tts/config.py index e64c63e14..2aa822ed8 100644 --- a/mstar/model/qwen3_tts/config.py +++ b/mstar/model/qwen3_tts/config.py @@ -246,12 +246,14 @@ class Qwen3TTSCodecConfig: # turns reference audio into codec frames for voice cloning. encoder_config: dict[str, Any] = field(default_factory=dict) - # M* stream policy: the codec pops a ramp of small chunks first (first - # audio after 4 frames = 320 ms of speech), then ``chunk_frames`` new - # frames per call, each preceded by up to ``left_context_frames`` already - # decoded frames so the causal decoder warms up (the reference's own - # ``chunked_decode`` uses 25 frames of left context). - chunk_schedule: tuple[int, ...] = (4, 8, 16) + # M* stream policy: the codec pops a ramp of small chunks first (the first + # frame alone, so first audio leaves one Talker step after prefill; the + # decoder is causal, so a frame's audio does not depend on how it was + # chunked), then ``chunk_frames`` new frames per call, each preceded by up + # to ``left_context_frames`` already decoded frames so the causal decoder + # warms up (the reference's own ``chunked_decode`` uses 25 frames of left + # context). + chunk_schedule: tuple[int, ...] = (1, 3, 8, 16) chunk_frames: int = 25 left_context_frames: int = 25 From 7b3381e9691231b1ff9b2d7570fc2c0aff96efe4 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:43:34 -0700 Subject: [PATCH 090/110] qwen3_tts: graph docstring follows the 1-3-8-16 chunk ramp --- mstar/model/qwen3_tts/qwen3_tts_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mstar/model/qwen3_tts/qwen3_tts_model.py b/mstar/model/qwen3_tts/qwen3_tts_model.py index b5f9dd5f4..96e5bde48 100644 --- a/mstar/model/qwen3_tts/qwen3_tts_model.py +++ b/mstar/model/qwen3_tts/qwen3_tts_model.py @@ -16,7 +16,7 @@ Codec - stateless speech-tokenizer decoder producing PCM chunks Streaming topology: - Talker --[codec_tokens, ScheduledLeftContextChunkPolicy((4, 8, 16), 25, 25)]--> Codec + Talker --[codec_tokens, ScheduledLeftContextChunkPolicy((1, 3, 8, 16), 25, 25)]--> Codec Request state machine: Talker: talker_prefill | talker_prefill_clone -> talker_decode loop -> done on EOS/token limit From f342ff57cf2ddfedbe451b2d23d1db91b9fe4d9e Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:43:34 -0700 Subject: [PATCH 091/110] qwen3_tts: submodule header follows the 1-3-8-16 chunk ramp --- mstar/model/qwen3_tts/submodules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index a3862e496..1ec0aff61 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -28,7 +28,7 @@ # -> postprocess -> check_stop (Talker only) # # Streaming topology: -# Talker --[codec_tokens, ScheduledLeftContextChunkPolicy((4, 8, 16), 25, 25)]--> Codec +# Talker --[codec_tokens, ScheduledLeftContextChunkPolicy((1, 3, 8, 16), 25, 25)]--> Codec # --------------------------------------------------------------------------- from __future__ import annotations From e8b7fca401419db7c6beac8be4541e14af2c00ec Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:43:34 -0700 Subject: [PATCH 092/110] docs: first audio after a single codec frame --- docs/models.rst | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index bcd6813ad..ecd80b6fd 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -121,10 +121,11 @@ Qwen3-TTS notes ordered sentence chunks of about 400 characters (one Talker request each, two kept in flight while the current one streams); set ``sentence_chunking: false`` (or ``true`` for shorter texts) per request. -- Audio streams in a ramp of codec chunks: the first window is decoded after 4 - frames (320 ms of speech), later windows grow to 25 new frames behind 25 - frames of already decoded left context (the reference's own - ``chunked_decode`` context). Each window size is a CUDA-graph bucket +- Audio streams in a ramp of codec chunks: the first frame is decoded on its + own (first audio one Talker step after prefill), the next windows add 3, 8 + and 16 frames, then 25 new frames behind 25 frames of already decoded left + context (the reference's own ``chunked_decode`` context; the decoder is + causal, so chunking does not change the audio). Each window size is a CUDA-graph bucket captured for batch sizes 1 to 32; the stream buffer reports how many leading frames of a window are repeated context, and the codec trims their audio. - Talker prefill remains eager because it runs once with variable sequence From b4fff0b34dbaf205e1312e32919efddd9d9f0fde Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:45:42 -0700 Subject: [PATCH 093/110] test: smoke repeats a seeded sampled run for Base, greedy for the others --- test/qwen3-tts/smoke_qwen3_tts.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/test/qwen3-tts/smoke_qwen3_tts.py b/test/qwen3-tts/smoke_qwen3_tts.py index 4a137f360..d5411f05a 100644 --- a/test/qwen3-tts/smoke_qwen3_tts.py +++ b/test/qwen3-tts/smoke_qwen3_tts.py @@ -139,13 +139,18 @@ def run(name: str, payload: dict, stream: bool, min_seconds: float, max_seconds: run("stream", {**base, "input": TEXT}, stream=True, min_seconds=3.0, max_seconds=12.0) # 2. non-streaming container response. run("blob", {**base, "input": TEXT}, stream=False, min_seconds=3.0, max_seconds=12.0) - # 3. greedy is repeatable byte for byte (same seed). - greedy = {**base, "input": TEXT, "do_sample": False, "subtalker_dosample": False, "seed": 7} - first = run("greedy_a", greedy, stream=True, min_seconds=3.0, max_seconds=12.0) - second = run("greedy_b", greedy, stream=True, min_seconds=3.0, max_seconds=12.0) - report["greedy_repeatable"] = first == second + # 3. the same seed is repeatable byte for byte. CustomVoice/VoiceDesign run + # greedy; Base clone prompts degenerate under greedy decoding in the + # reference implementation too (repeated near-silent frames, no EOS), so + # there the seeded default sampling is what must repeat. + repeat = {**base, "input": TEXT, "seed": 7} + if args.variant != "base": + repeat.update(do_sample=False, subtalker_dosample=False) + first = run("repeat_a", repeat, stream=True, min_seconds=3.0, max_seconds=12.0) + second = run("repeat_b", repeat, stream=True, min_seconds=3.0, max_seconds=12.0) + report["seed_repeatable"] = first == second if first != second: - failures.append("greedy runs with the same seed differ") + failures.append("runs with the same seed differ") # 4. long input goes through sentence chunking (server side) and stays continuous. run("long_chunked", {**base, "input": LONG_TEXT}, stream=True, min_seconds=25.0, max_seconds=90.0) # 5. instruction control (1.7B CustomVoice style, VoiceDesign voice description). From 45c7b15d48b1d3242d496a2bcf0cf357bd5abb93 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:45:42 -0700 Subject: [PATCH 094/110] benchmark: WER loads audio with soundfile, no ffmpeg needed --- benchmark/tts_wer.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/benchmark/tts_wer.py b/benchmark/tts_wer.py index 5c16e3e43..9f3bbbd0a 100644 --- a/benchmark/tts_wer.py +++ b/benchmark/tts_wer.py @@ -38,6 +38,26 @@ def load_sentences(path: str) -> dict[int, str]: return {i + 1: ln.strip() for i, ln in enumerate(lines) if ln.strip()} +ASR_SAMPLE_RATE = 16000 + + +def load_for_asr(path: Path) -> dict: + """A WAV file as the pipeline's raw-audio input (mono float32 at 16 kHz). + + Decoded with soundfile so the benchmark hosts need no ffmpeg binary. + """ + import numpy as np + import soundfile as sf + import torch + import torchaudio.functional as taf + + audio, sample_rate = sf.read(str(path), dtype="float32", always_2d=True) + mono = torch.from_numpy(np.ascontiguousarray(audio.mean(axis=1))) + if sample_rate != ASR_SAMPLE_RATE: + mono = taf.resample(mono, sample_rate, ASR_SAMPLE_RATE) + return {"raw": mono.numpy(), "sampling_rate": ASR_SAMPLE_RATE} + + def transcribe(audio_paths: list[Path], model_id: str, device: str, batch_size: int) -> list[str]: import torch from transformers import pipeline @@ -49,7 +69,7 @@ def transcribe(audio_paths: list[Path], model_id: str, device: str, batch_size: device=device, ) outputs = asr( - [str(p) for p in audio_paths], + [load_for_asr(p) for p in audio_paths], batch_size=batch_size, generate_kwargs={"language": "en", "task": "transcribe"}, return_timestamps=False, From 9042454148e1bd5c6a1c0f6d8b33b21f15c91daf Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:48:29 -0700 Subject: [PATCH 095/110] conductor: optional rank_devices mapping places worker ranks on a shared device --- mstar/conductor/conductor.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index c68c066bc..d156785ed 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -111,6 +111,16 @@ def _exit_when_orphaned(worker_id: str, parent=None, poll_s: float = 0.5) -> Non ) +def worker_device(device_type: str, rank: int, rank_devices: dict[int, int] | None = None) -> str: + """The device string a worker rank runs on: its own number unless the + deployment's ``rank_devices`` maps it elsewhere (several ranks may share a + GPU); CPU deployments ignore the mapping.""" + if device_type == "cpu": + return "cpu" + index = (rank_devices or {}).get(rank, rank) + return f"{device_type}:{index}" + + def _worker_process_target( worker_id: str, worker_ids: list[str], @@ -317,6 +327,14 @@ def __init__( ) assert "max_seq_len" in self.model_config assert "node_groups" in self.model_config + # Optional ``rank_devices: {rank: device_index}`` places a worker rank + # on a device other than the one its number implies, e.g. two workers + # sharing one GPU so a light node's steps stop interleaving with a + # heavy node's on the same worker loop. + self.rank_devices = { + int(rank): int(index) + for rank, index in (self.model_config.get("rank_devices") or {}).items() + } self.default_sharding_config = model.get_sharding_config(model_config_file) self.worker_graphs = { @@ -491,10 +509,7 @@ def _launch_workers(self): "model": self.model, "enable_nvtx": self.enable_nvtx, "enable_prof": self.enable_prof, - "device": ( - f"{self.device_type}:{rank}" - if self.device_type != "cpu" else "cpu" - ), + "device": worker_device(self.device_type, rank, self.rank_devices), "log_level": self.log_level, "tensor_comm_protocol": self.tensor_comm_protocol, "tcp_transfer_device": self.tcp_transfer_device From 03cd27700a40947809ef8af5f34609b5255b68ac Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:48:29 -0700 Subject: [PATCH 096/110] docs: rank_devices for several workers on one GPU --- docs/serving.rst | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/serving.rst b/docs/serving.rst index 1bc969983..17b497d48 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -217,6 +217,18 @@ no longer read. It raises an error with a message describing the migration. - {node_names: [vae_encoder, vae_decoder], ranks: [0]} - {node_names: [LLM], ranks: [0]} +**Several workers on one GPU.** A worker is one process per rank, and by default rank +``n`` runs on device ``n``. ``rank_devices`` maps ranks onto devices explicitly, so two +node groups can run as separate workers on the same GPU: a heavy autoregressive node then +never waits for a light node's steps on its worker loop (a TTS talker and its codec): + +.. code-block:: yaml + + node_groups: + - {node_names: [Talker], ranks: [0]} + - {node_names: [Codec], ranks: [1]} + rank_devices: {1: 0} + **Disaggregation.** The same node can live on different GPUs *per graph walk* — e.g. prefill, decode, and image generation on three GPUs: From 7b4a979546132c502a084b2f5f14face8e1a2aa0 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:48:29 -0700 Subject: [PATCH 097/110] test: worker device mapping --- test/modular/test_worker_device_map.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 test/modular/test_worker_device_map.py diff --git a/test/modular/test_worker_device_map.py b/test/modular/test_worker_device_map.py new file mode 100644 index 000000000..e461f3054 --- /dev/null +++ b/test/modular/test_worker_device_map.py @@ -0,0 +1,20 @@ +"""Deployment ``rank_devices``: which device string each worker rank gets.""" + +from mstar.conductor.conductor import worker_device + + +def test_rank_is_the_device_by_default(): + assert worker_device("cuda", 0) == "cuda:0" + assert worker_device("cuda", 3, {}) == "cuda:3" + assert worker_device("xpu", 1) == "xpu:1" + + +def test_rank_devices_places_ranks_on_a_shared_device(): + mapping = {1: 0, 2: 0} + assert [worker_device("cuda", rank, mapping) for rank in (0, 1, 2, 3)] == [ + "cuda:0", "cuda:0", "cuda:0", "cuda:3", + ] + + +def test_cpu_deployments_ignore_the_mapping(): + assert worker_device("cpu", 1, {1: 0}) == "cpu" From 16f33205e0f56f3843eaab9c7333c3acadcb488d Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:49:07 -0700 Subject: [PATCH 098/110] qwen3_tts: deployment with the codec on a second worker sharing the GPU --- configs/qwen3tts_1p7b_split.yaml | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 configs/qwen3tts_1p7b_split.yaml diff --git a/configs/qwen3tts_1p7b_split.yaml b/configs/qwen3tts_1p7b_split.yaml new file mode 100644 index 000000000..94fb0578b --- /dev/null +++ b/configs/qwen3tts_1p7b_split.yaml @@ -0,0 +1,18 @@ +# Qwen3-TTS-12Hz-1.7B-CustomVoice with the codec on its own worker. +# Both workers run on GPU 0 (rank_devices), so the Talker's decode steps no +# longer alternate with codec windows on one worker loop: the Talker runs +# back to back and the codec batches the windows that become ready while it +# does. Same graph, resources and checkpoint as configs/qwen3tts_1p7b.yaml. +model: "qwen3_tts_1p7b" +max_seq_len: 32768 +resources: + talker_attn: + flashinfer_backend: fa2 +node_groups: + - node_names: [Talker] + ranks: [0] + graph_walks: [talker_prefill, talker_decode] + - node_names: [Codec] + ranks: [1] + graph_walks: [codec_chunk] +rank_devices: {1: 0} From a44e180a82c2547fd5fd160a9ffc7fd0e486a6d5 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:50:00 -0700 Subject: [PATCH 099/110] benchmark: WER drives Whisper directly from 16 kHz arrays --- benchmark/tts_wer.py | 36 +++++++++++++++++++++--------------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/benchmark/tts_wer.py b/benchmark/tts_wer.py index 9f3bbbd0a..1e7e025bb 100644 --- a/benchmark/tts_wer.py +++ b/benchmark/tts_wer.py @@ -59,22 +59,28 @@ def load_for_asr(path: Path) -> dict: def transcribe(audio_paths: list[Path], model_id: str, device: str, batch_size: int) -> list[str]: + """Whisper transcripts, one per file, in the order given. + + Drives the model directly rather than through ``pipeline(...)``: the + pipeline decodes files with ffmpeg/torchcodec, which benchmark hosts may + not have, while the processor only needs 16 kHz arrays. + """ import torch - from transformers import pipeline - - asr = pipeline( - "automatic-speech-recognition", - model=model_id, - torch_dtype=torch.float16 if device.startswith("cuda") else torch.float32, - device=device, - ) - outputs = asr( - [load_for_asr(p) for p in audio_paths], - batch_size=batch_size, - generate_kwargs={"language": "en", "task": "transcribe"}, - return_timestamps=False, - ) - return [o["text"] for o in outputs] + from transformers import WhisperForConditionalGeneration, WhisperProcessor + + dtype = torch.float16 if device.startswith("cuda") else torch.float32 + processor = WhisperProcessor.from_pretrained(model_id) + model = WhisperForConditionalGeneration.from_pretrained(model_id, torch_dtype=dtype).to(device).eval() + texts: list[str] = [] + for start in range(0, len(audio_paths), batch_size): + clips = [load_for_asr(p)["raw"] for p in audio_paths[start:start + batch_size]] + features = processor( + clips, sampling_rate=ASR_SAMPLE_RATE, return_tensors="pt", + ).input_features.to(device=device, dtype=dtype) + with torch.inference_mode(): + generated = model.generate(features, language="en", task="transcribe", max_new_tokens=440) + texts.extend(processor.batch_decode(generated, skip_special_tokens=True)) + return [t.strip() for t in texts] def main(argv: list[str] | None = None) -> None: From 35ed0405bd6c732b3a8885e397ccc7af95db8be2 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:54:30 -0700 Subject: [PATCH 100/110] qwen3_tts: talker prefill replays a packed CUDA graph --- mstar/model/qwen3_tts/submodules.py | 38 +++++++++++++++++++++++------ 1 file changed, 31 insertions(+), 7 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index 1ec0aff61..fc76bcb97 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -46,6 +46,7 @@ from mstar.engine.cuda_graph_config import ( BatchedCudaGraphConfig, CudaGraphConfig, + PackedCudaGraphConfig, PiecewiseBatchedConfig, PiecewiseCallInputs, PiecewiseCaptureShape, @@ -107,6 +108,11 @@ class TalkerSubmodule(ARNodeSubmodule): disable_torch_compile = True MAX_BATCH_SIZE = 32 DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + # Prefill prompts are short (a sentence plus the ChatML frame, an optional + # instruction and speaker slot): a few token buckets cover them, and the + # eager alternative costs ~20 ms of launch overhead per request. + PREFILL_TOKEN_BUCKETS = [32, 64, 128, 256, 512, 1024] + PREFILL_CAPTURE_BATCH_SIZES = [1, 2, 4, 8] CHATML_ASSISTANT_PREFIX_TOKEN_IDS = CHATML_ASSISTANT_PREFIX_TOKEN_IDS CHATML_ASSISTANT_SUFFIX_TOKEN_IDS = CHATML_ASSISTANT_SUFFIX_TOKEN_IDS @@ -812,16 +818,34 @@ def max_batch_size(self, graph_walk: str) -> int: def get_cuda_graph_configs( self, device: torch.device, tp_world_size: int = 1 ) -> list[CudaGraphConfig]: - """Capture fixed one-token Talker decode batches, including sampling. + """Capture fixed one-token Talker decode batches and packed prefills, + sampling included. Dynamic EOS suppression is carried by ``ARNodeInputs`` and packed into a graph input, so replay never consults capture-slot dummy request state. - Prefill remains eager because it is variable-length and runs once per - request. + Prefill replays a packed capture of the smallest token bucket that + holds the batch; the clone prefill stays eager because it also pushes + the reference clip's frames into the codec stream. """ del tp_world_size dtype = self.model.model.codec_embedding.weight.dtype - return [BatchedCudaGraphConfig( + hidden = self.talker_config.hidden_size + + def prefill_input(num_tokens: int) -> ARNodeInputs: + return ARNodeInputs( + input_embeds=torch.zeros(num_tokens, hidden, dtype=dtype, device=device), + input_seq_len=num_tokens, + tensor_inputs={"suppress_eos": torch.ones(1, dtype=torch.bool, device=device)}, + ) + + prefill = PackedCudaGraphConfig( + capture_graph_walk="talker_prefill", + capture_token_lengths=self.PREFILL_TOKEN_BUCKETS, + make_node_input=prefill_input, + capture_batch_sizes=self.PREFILL_CAPTURE_BATCH_SIZES, + compile=False, + ) + return [prefill, BatchedCudaGraphConfig( capture_graph_walk="talker_decode", single_request_inputs=ARNodeInputs( input_embeds=torch.zeros( @@ -904,9 +928,9 @@ def declare_step( def can_use_cuda_graphs( self, batch: ExecutingBatch, model_inputs: list[NodeInputs] ) -> bool: - """Replay the whole decode graph; sampling params are read from buffers, - so no request's settings can disqualify it.""" - if batch.graph_walk != "talker_decode" or not self.can_batch( + """Replay the whole decode graph or a packed prefill; sampling params + are read from buffers, so no request's settings can disqualify it.""" + if batch.graph_walk not in ("talker_decode", "talker_prefill") or not self.can_batch( batch, model_inputs ): return False From 27ebc0dc3a8effb84263fda34ae1077aa434b6bc Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:54:30 -0700 Subject: [PATCH 101/110] test: talker prefill graph capture config --- test/modular/test_qwen3_tts_model.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index ed2b446ac..18e2844c9 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -1072,12 +1072,31 @@ def test_qwen3_tts_talker_batches_and_captures_decode(): assert packed["input_embeds"].shape == (2, 16) assert packed["last_token_indices"].tolist() == [0, 1] assert packed["suppress_eos"].tolist() == [True, True] - graph_config = submodule.get_cuda_graph_configs(torch.device("cpu"))[0] - assert graph_config.capture_graph_walk == "talker_decode" + configs = {c.capture_graph_walk: c for c in submodule.get_cuda_graph_configs(torch.device("cpu"))} + graph_config = configs["talker_decode"] assert graph_config.capture_batch_sizes == [1, 2, 4, 8, 16, 32] assert graph_config.single_request_inputs.tensor_inputs[ "suppress_eos" ].item() is True + # Prefill replays a packed capture: token buckets, padding rows shaped like + # a prepared prefill (embeds + the dynamic EOS-suppression key). + prefill_config = configs["talker_prefill"] + assert prefill_config.capture_token_lengths == [32, 64, 128, 256, 512, 1024] + assert prefill_config.capture_batch_sizes == [1, 2, 4, 8] + padding = prefill_config.make_node_input(7) + assert padding.input_embeds.shape == (7, 16) and padding.input_seq_len == 7 + assert padding.tensor_inputs["suppress_eos"].item() is True + assert prefill_config.replay_graph_walks == ["talker_prefill"] # the clone prefill stays eager + prefill_batch = ExecutingBatch( + node_name="Talker", step_context=_step_context("talker_prefill", ["a", "b"]), + per_request_input_tensors={}, per_request_info=info, + ) + assert submodule.can_use_cuda_graphs(prefill_batch, model_inputs) + clone_batch = ExecutingBatch( + node_name="Talker", step_context=_step_context("talker_prefill_clone", ["a", "b"]), + per_request_input_tensors={}, per_request_info=info, + ) + assert not submodule.can_use_cuda_graphs(clone_batch, model_inputs) # Residual sampling params live in per-request sampler buffers, so requests # that disagree about them still batch AND still replay the decode graph. # (They used to fall out of both.) From 88479ddf32306b7055cae4a5bd865abe794dce2d Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 04:54:30 -0700 Subject: [PATCH 102/110] docs: talker prefill is graph-captured --- docs/models.rst | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index ecd80b6fd..7641ba8c5 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -128,8 +128,10 @@ Qwen3-TTS notes causal, so chunking does not change the audio). Each window size is a CUDA-graph bucket captured for batch sizes 1 to 32; the stream buffer reports how many leading frames of a window are repeated context, and the codec trims their audio. -- Talker prefill remains eager because it runs once with variable sequence - lengths. Decode always uses the whole-walk CUDA Graph, with the 15-step +- Talker prefill replays a packed CUDA Graph for the smallest token bucket + (32 to 1024 tokens) that holds the batch; only the clone prefill, which also + pushes the reference clip's frames into the codec stream, runs eager. Decode + always uses the whole-walk CUDA Graph, with the 15-step CodePredictor loop captured inside it; request-local EOS suppression is carried as a graph tensor input so replay does not consult capture-slot dummy request state. Residual ``subtalker_*`` sampling is per-request through the From 58eca9a5b13e77da6cad0308c75719914a6407a9 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 05:06:58 -0700 Subject: [PATCH 103/110] qwen3_tts: codec_dtype deployment option (float32 default, bfloat16 opt-in) --- mstar/model/qwen3_tts/qwen3_tts_model.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/mstar/model/qwen3_tts/qwen3_tts_model.py b/mstar/model/qwen3_tts/qwen3_tts_model.py index 96e5bde48..9182bd612 100644 --- a/mstar/model/qwen3_tts/qwen3_tts_model.py +++ b/mstar/model/qwen3_tts/qwen3_tts_model.py @@ -288,6 +288,21 @@ def _verify_checkpoint_coverage( # --------------------------------------------------------------------------- +_CODEC_DTYPES = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16} + + +def codec_dtype(value: str | torch.dtype) -> torch.dtype: + """The codec decoder dtype named by a deployment's ``codec_dtype``.""" + if isinstance(value, torch.dtype): + return value + try: + return _CODEC_DTYPES[str(value)] + except KeyError: + raise ValueError( + f"codec_dtype must be one of {sorted(_CODEC_DTYPES)}, got {value!r}" + ) from None + + class Qwen3TTSModel(Model): """Qwen3-TTS 12 Hz model contract (CustomVoice, VoiceDesign, Base). @@ -304,6 +319,10 @@ def __init__( ) -> None: self.model_path_hf = model_path_hf self.cache_dir = cache_dir + # Server-init knob (deployment YAML ``model_kwargs``): the speech + # codec decoder's dtype. float32 reproduces the reference decoder + # bit for bit; bfloat16 roughly halves the codec's GPU time. + self.codec_dtype = codec_dtype(kwargs.get("codec_dtype", "float32")) # The lightweight API-side object needs config and tokenizer only. self.local_dir = _resolve_model_metadata(model_path_hf, cache_dir) @@ -1095,6 +1114,8 @@ def weights(): decoder, loaded, _checkpoint_keys(codec_dir, prefix), "Qwen3-TTS Codec" ) decoder.eval() + if self.codec_dtype != torch.float32: + decoder.to(dtype=self.codec_dtype) return CodecSubmodule(decoder, self.config) def _create_ref_encoder_submodule( From 44cdb5c49e0a236cc96e9e0a06e8c05bb44e7cfe Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 05:06:58 -0700 Subject: [PATCH 104/110] qwen3_tts: codec batch comment matches the 32-request ceiling --- mstar/model/qwen3_tts/submodules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index fc76bcb97..b635dab50 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -965,7 +965,7 @@ class CodecSubmodule(ARNodeSubmodule): # Windows are at most chunk + left_context frames (50 by default, 4 s of # audio), so the decoder's activations stay small enough to capture - # batches of 16 next to the Talker; ``can_batch`` keeps the ceiling. + # batches of 32 next to the Talker; ``can_batch`` keeps the ceiling. MAX_BATCH_SIZE = 32 CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] From fe6d9ad1eb6de3e3770b1317ad8466946dc97cb4 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 05:06:58 -0700 Subject: [PATCH 105/110] docs: codec_dtype option --- docs/models.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/models.rst b/docs/models.rst index 7641ba8c5..38933ce12 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -138,6 +138,10 @@ Qwen3-TTS notes ``code_predictor`` aux sampler, so custom values neither block batching nor fall off the graph. On the 1.7B checkpoints the CodePredictor projects the Talker-width inputs through ``small_to_mtp_projection`` before its depth loop. +- The codec decoder runs in float32 by default, reproducing the reference + decoder bit for bit. ``model_kwargs: {codec_dtype: bfloat16}`` in the + deployment YAML halves its GPU time when throughput matters more than + bit-exactness (the Talker is bf16 either way). - Weight loading checks coverage in both directions: a parameter the checkpoint does not fill, or a checkpoint tensor the port does not load, fails startup. - The 12 Hz codec does not require the system SoX executable. M* imports only From ac890efb25bec1903c0512b488880088d4e86d8b Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 05:06:58 -0700 Subject: [PATCH 106/110] test: codec_dtype parsing --- test/modular/test_qwen3_tts_model.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index 18e2844c9..cafdd51be 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -1639,3 +1639,13 @@ def encode(rid: str, clip: torch.Tensor, ref_frames: int): assert encoder.calls == 3 and len(submodule._conditioning) == 2 encode("a4", clip_a, 7) # evicted, so encoded again assert encoder.calls == 4 + + +def test_qwen3_tts_codec_dtype_option(): + from mstar.model.qwen3_tts.qwen3_tts_model import codec_dtype + + assert codec_dtype("float32") is torch.float32 + assert codec_dtype("bfloat16") is torch.bfloat16 + assert codec_dtype(torch.float16) is torch.float16 + with pytest.raises(ValueError, match="codec_dtype"): + codec_dtype("int8") From 5fca373df8ef19a24384c04232f6fb9650336932 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 05:48:36 -0700 Subject: [PATCH 107/110] qwen3_tts: clone stream carries only the reference tail the codec can use as context --- mstar/model/qwen3_tts/submodules.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index b635dab50..c334ed178 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -407,8 +407,9 @@ def project(ids: torch.Tensor) -> torch.Tensor: generated_frames=0, ) if ref_frames > 0: - # The codec decodes the reference frames ahead of the generated - # ones (their audio is trimmed), exactly as the reference does. + # The codec warms up on the tail of the reference clip ahead of + # the generated frames (their audio is trimmed); see + # ``_codec_stream_items``. state.add("reference_frames", reference) return prefill.squeeze(0) @@ -689,16 +690,19 @@ def _run_depth_loop_piecewise( def _codec_stream_items(self, graph_walk: str, request_id: str, frame: torch.Tensor) -> list[torch.Tensor]: """Frames this step pushes into the codec stream, one item per frame. - The clone prefill leads with the reference clip's frames so the codec - warms up on the voice being cloned; ``CodecSubmodule`` trims their - audio. Decode steps (the captured path) always push exactly one frame. + The clone prefill leads with the tail of the reference clip's frames + so the codec warms up on the voice being cloned: only the last + ``left_context_frames`` matter, since that is all the context a + window ever carries, and every reference frame the stream carries + costs a codec pass whose audio ``CodecSubmodule`` then trims. Decode + steps (the captured path) always push exactly one frame. """ if graph_walk != "talker_prefill_clone": return [frame] reference = self.request_state(request_id).get("reference_frames") if reference is None: return [frame] - return [*reference.unbind(0), frame] + return [*reference[-self.config.codec.left_context_frames:].unbind(0), frame] def forward( self, @@ -1021,11 +1025,13 @@ def prepare_inputs( del graph_walk, kwargs state = self.request_state(fwd_info.request_id) if "ref_frames" in inputs and "skip_samples" not in state: - # Voice clone: the stream leads with the reference clip's frames, - # whose audio the client must not hear. + # Voice clone: the stream leads with the last ``left_context_frames`` + # of the reference clip (all of it when shorter), whose audio the + # client must not hear. + ref_frames = int(inputs["ref_frames"][0].reshape(-1)[0].item()) state.add( "skip_samples", - int(inputs["ref_frames"][0].reshape(-1)[0].item()) * self.total_upsample, + min(ref_frames, self.config.codec.left_context_frames) * self.total_upsample, ) codes = inputs["codec_tokens"][0].to( device=self.get_device(), dtype=torch.long From 8c9072172ba14060f1af661e097d4c8ef7424ceb Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 05:48:36 -0700 Subject: [PATCH 108/110] test: clone stream reference tail and trimming --- test/modular/test_qwen3_tts_model.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index cafdd51be..3489b7230 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -867,11 +867,17 @@ def test_qwen3_tts_clone_prefill_streams_reference_frames_first(): submodule = TalkerSubmodule( Qwen3TTSTalkerModel(config), Qwen3TTSCodePredictor(config), config ) - reference = torch.arange(8).view(2, 4) + # Three reference frames, left context 2: the stream carries the last two + # (all a window can use as context) ahead of the first generated frame. + reference = torch.arange(12).view(3, 4) submodule.request_state("clone").add("reference_frames", reference) frame = torch.tensor([9, 9, 9, 9]) items = submodule._codec_stream_items("talker_prefill_clone", "clone", frame) - assert [item.tolist() for item in items] == [[0, 1, 2, 3], [4, 5, 6, 7], [9, 9, 9, 9]] + assert [item.tolist() for item in items] == [[4, 5, 6, 7], [8, 9, 10, 11], [9, 9, 9, 9]] + submodule.request_state("short").add("reference_frames", reference[:1]) + assert [item.tolist() for item in submodule._codec_stream_items("talker_prefill_clone", "short", frame)] == [ + [0, 1, 2, 3], [9, 9, 9, 9], + ] # Only the clone prefill leads with the reference; decode never does. assert submodule._codec_stream_items("talker_decode", "clone", frame) == [frame] assert submodule._codec_stream_items("talker_prefill_clone", "other", frame) == [frame] @@ -1334,18 +1340,18 @@ def test_qwen3_tts_codec_trims_reference_audio_from_clone_streams(): {"codec_tokens": [codes], "ref_frames": [torch.tensor([4])]}, ) state = submodule.request_state("clone") - assert state["skip_samples"] == 16 + # A 4-frame clip with left context 2: the stream carried its last 2 frames (8 samples). + assert state["skip_samples"] == 8 - # First chunk: 3 frames = 12 samples, all reference -> nothing emitted. + # First chunk: 3 frames = 12 samples, the first 8 are reference -> one frame emitted. first = {"audio_chunk": [torch.arange(20)]} submodule.postprocess("clone", None, first, inputs=_geometry(frames=3, context=0)) - assert first["audio_chunk"][0].numel() == 0 - assert state["skip_samples"] == 4 - # Second chunk: 2 context + 3 new frames; 4 more samples belong to the reference. + assert first["audio_chunk"][0].tolist() == [8, 9, 10, 11] + assert state["skip_samples"] == 0 + # Second chunk: 2 context + 3 new frames, nothing left to drop. second = {"audio_chunk": [torch.arange(20)]} submodule.postprocess("clone", None, second, inputs=_geometry(frames=5, context=2)) - assert second["audio_chunk"][0].tolist() == list(range(12, 20)) - assert state["skip_samples"] == 0 + assert second["audio_chunk"][0].tolist() == list(range(8, 20)) def test_qwen3_tts_codec_filters_eos_and_pads_to_capture_shape(): From 2a8fa504db6fe49ac70a8f16a2c0d5ff4fc530bc Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 06:07:47 -0700 Subject: [PATCH 109/110] qwen3_tts: split deployment caps the KV pool for two workers on one GPU --- configs/qwen3tts_1p7b_split.yaml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/configs/qwen3tts_1p7b_split.yaml b/configs/qwen3tts_1p7b_split.yaml index 94fb0578b..c6245254c 100644 --- a/configs/qwen3tts_1p7b_split.yaml +++ b/configs/qwen3tts_1p7b_split.yaml @@ -8,6 +8,11 @@ max_seq_len: 32768 resources: talker_attn: flashinfer_backend: fa2 + # Two workers share the GPU: cap the KV pool at 1024 pages x 128 tokens + # (16 GB; 32 concurrent requests of a few hundred tokens use ~300 pages) + # so the Talker's weights, graphs and pool fit next to the codec worker's. + talker_kv: + max_num_pages: 1024 node_groups: - node_names: [Talker] ranks: [0] From 708f72d05a5e238f920fb4a60c6d009945ab4d51 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 07:28:35 -0700 Subject: [PATCH 110/110] docs: split deployment and CUDA MPS note --- docs/models.rst | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/models.rst b/docs/models.rst index 38933ce12..017656bdf 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -138,6 +138,12 @@ Qwen3-TTS notes ``code_predictor`` aux sampler, so custom values neither block batching nor fall off the graph. On the 1.7B checkpoints the CodePredictor projects the Talker-width inputs through ``small_to_mtp_projection`` before its depth loop. +- ``configs/qwen3tts_1p7b_split.yaml`` runs the Codec on a second worker that + shares GPU 0 (``rank_devices``) and caps the Talker KV pool so both fit. + On its own it lowers first-audio latency at high concurrency; with a + user-level CUDA MPS daemon (``nvidia-cuda-mps-control -d`` before + ``mstar serve``) the two workers' kernels also overlap, which raised c=32 + throughput by about a fifth on an H100. - The codec decoder runs in float32 by default, reproducing the reference decoder bit for bit. ``model_kwargs: {codec_dtype: bfloat16}`` in the deployment YAML halves its GPU time when throughput matters more than