fix: codec decoder input structure and bugs

- Codec decoder now receives pre-computed
attention_mask, cos/sin RoPE, similar as talker
and local talker.
- The computation of attention_mask previously
also bugged because new KV is assigned from the
right, and the attention_mask has to ignore
left pads that have not yet reached sliding_window
length.
- The above point caused initial sound bugs in the
previous commits, which has now been fixed.
- All input axes for codec decoder are also now made
static in the export to make the graph capture better.
It still can be made dynamic in the export with consistent
length at runtime, but graph capture will perform better
if the exported model already contains pure static axes.
- This includes batch and chunk_frames, batch=1,
chunk_frames=4 (320 ms).
- Next will be to make all exported axes static
for talker, local_talker, and lm_head,
as these are also captured by CUDA graphs.

qwen3-tts_onnx/codec_decoder_model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:246c19dba19adf0ad2620b82fece12e17e874a5d9bd52a3b089fd8b8107b57de
-size 456761251
+oid sha256:5db3b88ca51c04b39ad2603ab781887fef300a65b5614f7cbeff1be86d0da51c
+size 456699494

src/inference/qwen3_tts_inferencer_onnx.py

@@ -372,7 +372,6 @@ class Qwen3TTSInferencerONNX:
     cuda_device_id : int
     enable_cuda_graph : bool
     num_threads : int
-    chunk_frames : int
     temperature : float
     top_p : float
     top_k : int
@@ -401,8 +400,8 @@ class Qwen3TTSInferencerONNX:
         use_cuda: bool = True,
         cuda_device_id: int = 0,
         enable_cuda_graph: bool = True,
-        num_threads: int = 4,
         chunk_frames: int = 4,
+        num_threads: int = 4,
         temperature: float = 0.85,
         top_p: float = 0.8,
         top_k: int = 50,
@@ -472,13 +471,14 @@ class Qwen3TTSInferencerONNX:
         self._speech_tokenizer_latent_dim = dec_cfg.latent_dim
         self._speech_tokenizer_codebook_dim = dec_cfg.codebook_dim
         self._speech_tokenizer_head_dim = dec_cfg.head_dim
+        self._speech_tokenizer_rope_theta = dec_cfg.rope_theta
         self._speech_tokenizer_num_attention_heads = dec_cfg.num_attention_heads
         self._speech_tokenizer_num_hidden_layers = dec_cfg.num_hidden_layers
         self._speech_tokenizer_num_key_value_heads = dec_cfg.num_key_value_heads
         self._speech_tokenizer_sliding_window = dec_cfg.sliding_window
         self._speech_tokenizer_decoder_left_context_size = 25
         self._speech_tokenizer_decoder_total_upsample = 1920
-        self._codec_chunk_frames = 4
+        self._codec_chunk_frames = chunk_frames
         self.output_sample_rate = self._speech_tokenizer_config.output_sample_rate
 
         with open(model_config_path, "r") as f:
@@ -546,7 +546,7 @@ class Qwen3TTSInferencerONNX:
         # Streaming / buffering
         self.text_buffer_size = 32
         self.min_text_chunk_chars = 8
-        self.chunk_frames = chunk_frames
+        self.chunk_frames = self._codec_chunk_frames
         self.overlap_frames = 0
         self._max_steps = _TALKER_MAX_SEQ_LEN - _TALKER_PREFILL_LEN  # 751
 
@@ -565,7 +565,7 @@ class Qwen3TTSInferencerONNX:
         self._causal_tril = _causal_tril.reshape(1, 1, _TALKER_MAX_SEQ_LEN, _TALKER_MAX_SEQ_LEN)
 
         # Define cos and sin table for RoPE embedding
-        self._cos_rope, self._sin_rope = self._build_rope_tables_numpy()
+        self._cos_rope, self._sin_rope = self._build_multimodal_rope_tables_numpy()
 
         # Pre-allocate input OrtValues for Talker prefill
         self._talker_prefill_inputs_embeds_ov = _make_device_ortvalue(
@@ -624,7 +624,11 @@ class Qwen3TTSInferencerONNX:
         # For each, [B, H, 16, D], with varying 1/0 depending on the position
 
         # Define cos and sin table for Local RoPE embedding
-        self._cos_rope_local, self._sin_rope_local = self._build_local_rope_tables_numpy()
+        self._cos_rope_local, self._sin_rope_local = self._build_rope_tables_numpy(
+            rope_dim=self._code_predictor_config.head_dim,
+            base=self._code_predictor_config.rope_theta,
+            max_pos=self._num_code_groups,
+        )
 
         # Pre-allocate input OrtValues for Local Talker prefill
         self._local_prefill_inputs_embeds_ov = _make_device_ortvalue(
@@ -743,12 +747,39 @@ class Qwen3TTSInferencerONNX:
         self._codec_decoder_bound = _BoundSession(self._codec_decoder_sess, self._device, cuda_device_id)
         self._codec_step_idx = 0
 
+        # Define causal_tril for attention_mask Local prefill and step
+        _causal_tril_codec = np.tril(
+            np.ones((self._speech_tokenizer_sliding_window, self._speech_tokenizer_sliding_window), dtype=bool)
+        )
+        self._causal_tril_codec = _causal_tril_codec.reshape(
+            1, 1, self._speech_tokenizer_sliding_window, self._speech_tokenizer_sliding_window
+        )
+        # Define cos and sin table for Codec RoPE embedding
+        self._cos_rope_codec, self._sin_rope_codec = self._build_rope_tables_numpy(
+            rope_dim=self._speech_tokenizer_head_dim,
+            base=self._speech_tokenizer_rope_theta,
+            max_pos=(
+                self._max_steps
+                if self._max_steps % self.chunk_frames == 0
+                else self._max_steps + self.chunk_frames - (self._max_steps % self.chunk_frames)
+            ),
+        )
+
         # Pre-allocate OrtValues input for Codec Decoder
         self._codec_codes_ov = _make_device_ortvalue(
             (1, self._num_code_groups, self.chunk_frames), np.int64, self._device, cuda_device_id
         )
-        self._codec_cache_position_ov = _make_device_ortvalue(
-            (self.chunk_frames), np.int64, self._device, cuda_device_id
+        # self._codec_cache_position_ov = _make_device_ortvalue(
+        #     (self.chunk_frames), np.int64, self._device, cuda_device_id
+        # )
+        self._codec_attention_mask_ov = _make_device_ortvalue(
+            (1, 1, self.chunk_frames, self._speech_tokenizer_sliding_window), np.float32, self._device, cuda_device_id
+        )
+        self._codec_cos_rope_ov = _make_device_ortvalue(
+            (1, self.chunk_frames, self._speech_tokenizer_head_dim), np.float32, self._device, cuda_device_id
+        )
+        self._codec_sin_rope_ov = _make_device_ortvalue(
+            (1, self.chunk_frames, self._speech_tokenizer_head_dim), np.float32, self._device, cuda_device_id
         )
 
         # Pre-allocate OrtValues for Codec Decoder
@@ -837,9 +868,12 @@ class Qwen3TTSInferencerONNX:
         b.bind_output_device("current_pre_conv_hidden_state_cache", self._codec_pre_conv_cache_ov)
 
         b.bind_input_device("codes", self._codec_codes_ov)
-        b.bind_input_device("cache_position", self._codec_cache_position_ov)
         b.bind_input_device("hidden_state_cache", self._codec_hidden_cache_ov)
         b.bind_input_device("pre_conv_hidden_state_cache", self._codec_pre_conv_cache_ov)
+        # b.bind_input_device("cache_position", self._codec_cache_position_ov)
+        b.bind_input_device("attention_mask", self._codec_attention_mask_ov)
+        b.bind_input_device("cos_rope", self._codec_cos_rope_ov)
+        b.bind_input_device("sin_rope", self._codec_sin_rope_ov)
 
         for i in range(self._speech_tokenizer_num_hidden_layers):
             b.bind_input_device(f"past_key_{i}", self._codec_kv_ovs[2 * i])
@@ -1021,7 +1055,7 @@ class Qwen3TTSInferencerONNX:
 
     # ── Helper to build RoPE cos and sin tables for the talker ────────────────
 
-    def _build_rope_tables_numpy(self):
+    def _build_multimodal_rope_tables_numpy(self):
         """
         Build NumPy RoPE cosine and sine lookup tables with interleaved multimodal layout.
 
@@ -1128,7 +1162,7 @@ class Qwen3TTSInferencerONNX:
 
         return cos.astype(np.float32), sin.astype(np.float32)
 
-    def _build_local_rope_tables_numpy(self):
+    def _build_rope_tables_numpy(self, rope_dim: int, base: float, max_pos: int):
         """
         Build NumPy RoPE cosine and sine lookup tables for the local talker.
 
@@ -1160,9 +1194,6 @@ class Qwen3TTSInferencerONNX:
             Sine table with shape matching the Torch version, including the
             extra singleton dimension inserted at axis 1.
         """
-        rope_dim = self._code_predictor_config.head_dim
-        base = self._code_predictor_config.rope_theta
-
         # Build inverse frequencies for RoPE.
         inv_idx = np.arange(0, rope_dim, 2, dtype=np.float32)
 
@@ -1172,7 +1203,7 @@ class Qwen3TTSInferencerONNX:
         inv_freq_expanded = np.broadcast_to(
             inv_freq[None, :, None], (1, len(inv_idx), 1)
         )  # shape (bs, 1, rope_dim, 1)
-        position_ids = np.arange(self._num_code_groups, dtype=np.float32)
+        position_ids = np.arange(max_pos, dtype=np.float32)
         position_ids_expanded = np.broadcast_to(
             position_ids[None, None, :], (1, 1, len(position_ids))
         )  # shape (bs, 1, positions)
@@ -1503,13 +1534,19 @@ class Qwen3TTSInferencerONNX:
     def _run_codec_decoder(
         self,
         chunk_tokens: np.ndarray,  # [1, 16, chunk_length]
-        pos: np.ndarray,  # [chunk_length]
+        attention_mask: np.ndarray,  # [1, 1, chunk_length, 72]
+        cos_rope: np.ndarray,  # [1, chunk_length, 64]
+        sin_rope: np.ndarray,  # [1, chunk_length, 64]
     ) -> np.ndarray:
+        # pos: np.ndarray,  # [chunk_length]
         """Execute one Codec Decoder step via IOBinding and return wav as NumPy."""
         b = self._codec_decoder_bound
         # Copy numpy inputs to buffer in-place
         _copy_numpy_to_ortvalue(chunk_tokens, self._codec_codes_ov)
-        _copy_numpy_to_ortvalue(pos, self._codec_cache_position_ov)
+        # _copy_numpy_to_ortvalue(pos, self._codec_cache_position_ov)
+        _copy_numpy_to_ortvalue(attention_mask, self._codec_attention_mask_ov)
+        _copy_numpy_to_ortvalue(cos_rope, self._codec_cos_rope_ov)
+        _copy_numpy_to_ortvalue(sin_rope, self._codec_sin_rope_ov)
         # All KV and output bindings are already wired statically
         b.run()
         # Copy wav to CPU
@@ -1548,19 +1585,20 @@ class Qwen3TTSInferencerONNX:
         cos_rope = self._cos_rope[:, :, cache_position].copy()
         sin_rope = self._sin_rope[:, :, cache_position].copy()
         logits = self._run_talker_step(inputs_embeds, cache_position, attn_mask, cos_rope, sin_rope)  # [1, vocab]
-        log.info(f"step-idx-{self._step_idx} logits {logits} {logits.shape} {logits.dtype}")
+        # log.info(f"step-idx-{self._step_idx} logits {logits} {logits.shape} {logits.dtype}")
 
         # Retrieve hidden states for Local Talker prefill
         self._last_hidden_states_np = self._talker_hidden_ov.numpy()  # [1, 1, H]
-        log.info(
-            f"step-idx-{self._step_idx} self._last_hidden_states_np {self._last_hidden_states_np} "
-            f"{self._last_hidden_states_np.shape} {self._last_hidden_states_np.dtype}"
-        )
+        # log.info(
+        #     f"step-idx-{self._step_idx} self._last_hidden_states_np {self._last_hidden_states_np} "
+        #     f"{self._last_hidden_states_np.shape} {self._last_hidden_states_np.dtype}"
+        # )
         self._talker_seq_len += 1
 
         # ── Sample first token (CPU) ───────────────────────────────────���──────
         history = self._generated_tokens[:, :, 0]
         first_token = self._sample(logits, history, suppress=self._suppress_tokens)
+        log.info(f"step-idx-{self._step_idx} vq {first_token} {first_token.shape} {first_token.dtype}")
 
         self._last_first_token = first_token
         self._is_stopping = bool(first_token[0] == self._codec_eos_token_id)
@@ -1686,6 +1724,7 @@ class Qwen3TTSInferencerONNX:
         )[
             None, :, :
         ]  # [1, 1, 16]
+        log.info(f"local-step-idx-{self._step_idx} vq-rvq {audio_tokens} {audio_tokens.shape} {audio_tokens.dtype}")
 
         self._last_audio_tokens = audio_tokens.copy()
         self._generated_tokens = np.concatenate([self._generated_tokens, audio_tokens.copy()], axis=1)
@@ -1822,10 +1861,29 @@ class Qwen3TTSInferencerONNX:
         n_iter = n_iter // self.chunk_frames
         codec_step_idx = 0
         for n in range(n_iter):
-            pos = np.arange(codec_step_idx, codec_step_idx + chunk_tokens.shape[-1], dtype=np.int64)
-            codec_step_idx += chunk_tokens.shape[-1]
+            pos = np.arange(codec_step_idx, codec_step_idx + self.chunk_frames, dtype=np.int64)
+            cos_rope = self._cos_rope_codec[:, pos].copy()
+            sin_rope = self._sin_rope_codec[:, pos].copy()
+            if codec_step_idx + self.chunk_frames < self._speech_tokenizer_sliding_window:
+                # attention mask goes from the right to the left
+                # we pad left False if the current + chunks still less than sliding window left
+                pad_len = self._speech_tokenizer_sliding_window - (codec_step_idx + self.chunk_frames)
+                # here, basically, we need to pad left because we start from the right
+                # and there still remaining amounts before it attends full sliding window length
+                # so we must not attend to those previous remaining amounts, which is less than pos. 0
+                causal_rows_right = self._causal_tril_codec[:, :, pos, :-pad_len]  # [1,1,q_len,72-accum_len]
+                causal_rows = np.pad(
+                    causal_rows_right, ((0, 0), (0, 0), (0, 0), (pad_len, 0)), mode="constant", constant_values=False
+                )
+            else:
+                # this case where the current + chunks is exactly sliding window length
+                # or more than sliding window length
+                # simply take the last chunks from the causal tril
+                causal_rows = self._causal_tril_codec[:, :, -self.chunk_frames :, :]  # [1,1,q_len,72]
+            attn_mask = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, 72]
+            codec_step_idx += self.chunk_frames
             t0 = time.perf_counter()
-            self._run_codec_decoder(chunk_tokens, pos)
+            self._run_codec_decoder(chunk_tokens, attn_mask, cos_rope, sin_rope)
             times["codec_decoder"].append((time.perf_counter() - t0) * 1000)  # ms
         exec_time_tot = sum(times["codec_decoder"]) / 1000
         avg_exec_time = exec_time_tot / len(times["codec_decoder"]) * 1000
@@ -1897,8 +1955,10 @@ class Qwen3TTSInferencerONNX:
         if not self._prefilled:
             self.prefill()
             return outputs
+        log.info(f"self._pending_tokens {self._pending_tokens}")
         while self._pending_tokens and not self.is_finished:
             token = self._pending_tokens.pop(0)
+            log.info(f"pop token {token} -> pending_tokens {self._pending_tokens}")
             output = self.step(token)
             if output is not None:
                 outputs.append(output)
@@ -1906,8 +1966,10 @@ class Qwen3TTSInferencerONNX:
 
     def push_text(self, text_fragment: str) -> List[NDArrayInt]:
         self._text_cache += text_fragment
+        log.info(f"push text_fragment {text_fragment} -> text_cache {self._text_cache}")
         for segment in self._extract_text_segments(force=False):
             self._pending_tokens.extend(self._tokenize_texts([segment]))
+            log.info(f"push segment {segment} -> pending_tokens {self._pending_tokens}")
         return self._drain_pending_tokens()
 
     def end_text(self) -> List[NDArrayInt]:
@@ -1977,12 +2039,30 @@ class Qwen3TTSInferencerONNX:
 
         # cache_position for the decoder transformer
         # It starts from 0 and grows. The model handles the sliding window internally.
-        pos = np.arange(self._codec_step_idx, self._codec_step_idx + chunk_tokens.shape[-1], dtype=np.int64)
-        self._codec_step_idx += chunk_tokens.shape[-1]
+        pos = np.arange(self._codec_step_idx, self._codec_step_idx + self.chunk_frames, dtype=np.int64)
+        cos_rope = self._cos_rope_codec[:, pos].copy()
+        sin_rope = self._sin_rope_codec[:, pos].copy()
+        if self._codec_step_idx + self.chunk_frames < self._speech_tokenizer_sliding_window:
+            # attention mask goes from the right to the left
+            # we pad left False if the current + chunks still less than sliding window left
+            pad_len = self._speech_tokenizer_sliding_window - (self._codec_step_idx + self.chunk_frames)
+            # here, basically, we need to pad left because we start from the right
+            # and there still remaining amounts before it attends full sliding window length
+            # so we must not attend to those previous remaining amounts, which is less than pos. 0
+            causal_rows_right = self._causal_tril_codec[:, :, pos, :-pad_len]  # [1,1,q_len,72-accum_len]
+            causal_rows = np.pad(
+                causal_rows_right, ((0, 0), (0, 0), (0, 0), (pad_len, 0)), mode="constant", constant_values=False
+            )
+        else:
+            # this case where the current + chunks is exactly sliding window length
+            # or more than sliding window length
+            # simply take the last chunks from the causal tril
+            causal_rows = self._causal_tril_codec[:, :, -self.chunk_frames :, :]  # [1,1,q_len,72]
+        attn_mask = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, 72]
 
         # Run the wired graph
         # chunk_tokens need to be .copy(), otherwise it messed up if we bind input and/or use CUDA graph
-        wav = self._run_codec_decoder(chunk_tokens.copy(), pos).copy()
+        wav = self._run_codec_decoder(chunk_tokens.copy(), attn_mask, cos_rope, sin_rope).copy()
         # new memory address, so that output wav not only the last repeated
 
         # The decoder output already discards the upsampled left-context frames (25 frames)
@@ -1994,6 +2074,7 @@ class Qwen3TTSInferencerONNX:
             f"{np.min(wav)} {np.mean(wav)} {np.std(wav)} {np.max(wav)} {wav.shape}"
         )
 
+        self._codec_step_idx += self.chunk_frames
         return wav
 
     def _overlap_samples(self, wav: NDArrayFloat) -> int:

test_qwen3-tts-streaming_onnx.py

@@ -34,7 +34,6 @@ Usage:
         --repetition_penalty 1.9 \
         --repetition_window 50 \
         --num_threads 4 \
-        --chunk_frames 4 \
         --prompt_wav audio_ref/speaker.[wav|flac|mp3] \
         --out_wav output.wav \
         --text "Text to be synthesized" \
@@ -237,7 +236,6 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--output_wav_path", default=None)
     p.add_argument("--language", default=_LANGUAGE)
     p.add_argument("--text", default=_TEXT, help="Text to synthesise.")
-    p.add_argument("--chunk_frames", type=int, default=_CHUNK_FRAMES)
     p.add_argument("--temperature", type=float, default=_TEMPERATURE)
     p.add_argument("--top_p", type=float, default=_TOP_P)
     p.add_argument("--top_k", type=int, default=_TOP_K)
@@ -329,7 +327,7 @@ def main() -> None:
         cuda_device_id=args.cuda_device_id,
         enable_cuda_graph=args.enable_cuda_graph,
         num_threads=args.num_threads,
-        chunk_frames=args.chunk_frames,
+        chunk_frames=_CHUNK_FRAMES,
         temperature=args.temperature,
         top_p=args.top_p,
         top_k=args.top_k,