fix: local talker & talker input structure

- Now, local talker follows the backbone talker
to use attention_mask and cos/sin RoPE as input.
- This reduces the amount of CPU fallback ops
from mask and cos/sin creations.
- Both still use scatter with indexing from cache_pos
for KV-cache update. This needs to be modified to
avoid CPU fallback.
- Also separate the .onnx file as _prefill and _step
as handling the KV-cache without scatter in the
next revision will need baked static flow.
- This is within the process of optimizing Ops from
CPU fallback.

qwen3-tts_onnx/{talker_local_model.onnx → talker_local_model_prefill.onnx}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8be4f6ddc1af78b7933ea9ae7c6ca787aa912d6a9aae54819d6bb5882d97b8f3
-size 314888365
+oid sha256:a64b5cc7bd8345a6a40c3cb0882d6ac699152817bd19e6f3a203214642c843c5
+size 314880101

qwen3-tts_onnx/{talker_model.onnx → talker_local_model_step.onnx}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8a3ad9936260c693af2eca5e8b22fb4a73fe1ca2d297a568e0f05ba520bd6033
-size 1775928943
+oid sha256:76e7a751f3d8038c5adf5e81c0b1f4f56a480b43be643928ff1e4f748ff4ce43
+size 314880101

qwen3-tts_onnx/talker_model_prefill.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:144ff25e52308068829932309acf2e8444965963a192b0649d07126b89983b4d
+size 1775928908

qwen3-tts_onnx/talker_model_step.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e0a94a3c6afafa67373fb841deef090bc432bdcf8c201bf86754db735da3932
+size 1775928908

src/inference/qwen3_tts_inferencer_onnx.py

@@ -349,10 +349,14 @@ class Qwen3TTSInferencerONNX:
 
     Parameters
     ----------
-    talker_model_path : str
-        Path to ``talker_model.onnx``.
-    talker_local_model_path : str
-        Path to ``talker_local_model.onnx``.
+    talker_model_prefill_path : str
+        Path to ``talker_model_prefill.onnx``.
+    talker_model_step_path : str
+        Path to ``talker_model_step.onnx``.
+    talker_local_model_prefill_path : str
+        Path to ``talker_local_model_prefill.onnx``.
+    talker_local_model_step_path : str
+        Path to ``talker_local_model_step.onnx``.
     talker_local_lm_head_model_path : str
         Path to ``talker_local_lm_head.onnx``  (batched lm_head for all 15 codebook groups).
     codec_decoder_model_path : str
@@ -380,8 +384,10 @@ class Qwen3TTSInferencerONNX:
 
     def __init__(
         self,
-        talker_model_path: str,
-        talker_local_model_path: str,
+        talker_model_prefill_path: str,
+        talker_model_step_path: str,
+        talker_local_model_prefill_path: str,
+        talker_local_model_step_path: str,
         talker_local_lm_head_model_path: str,
         codec_decoder_model_path: str,
         speaker_encoder_model_path: str,
@@ -421,18 +427,18 @@ class Qwen3TTSInferencerONNX:
         log.info("Loading ONNX sessions...")
 
         # Prefill sessions: IOBinding + CUDA graph (q_len is always fixed)
-        log.info(f"  talker prefill  <- {talker_model_path}")
-        _talker_prefill_sess = _sess(talker_model_path, graph=self._enable_cuda_graph)
-        log.info(f"  local prefill   <- {talker_local_model_path}")
-        _local_prefill_sess = _sess(talker_local_model_path, graph=self._enable_cuda_graph)
+        log.info(f"  talker prefill  <- {talker_model_prefill_path}")
+        _talker_prefill_sess = _sess(talker_model_prefill_path, graph=self._enable_cuda_graph)
+        log.info(f"  local prefill   <- {talker_local_model_prefill_path}")
+        _local_prefill_sess = _sess(talker_local_model_prefill_path, graph=self._enable_cuda_graph)
 
         # Step sessions: IOBinding + CUDA graph
-        log.info(f"  talker step     <- {talker_model_path}")
-        _talker_step_sess = _sess(talker_model_path, graph=self._enable_cuda_graph)
+        log.info(f"  talker step     <- {talker_model_step_path}")
+        _talker_step_sess = _sess(talker_model_step_path, graph=self._enable_cuda_graph)
 
         # Unified Local Talker backbone step (shared for steps 2..15)
-        log.info(f"  local step      <- {talker_local_model_path}")
-        _local_step_sess = _sess(talker_local_model_path, graph=self._enable_cuda_graph)
+        log.info(f"  local step      <- {talker_local_model_step_path}")
+        _local_step_sess = _sess(talker_local_model_step_path, graph=self._enable_cuda_graph)
 
         # Batched lm_head (all 15 codebook groups in one session)
         log.info(f"  local lm_head   <- {talker_local_lm_head_model_path}")
@@ -492,13 +498,14 @@ class Qwen3TTSInferencerONNX:
         self._num_hidden_layers = self._talker_config.num_hidden_layers  # 28
         self._num_key_value_heads = self._talker_config.num_key_value_heads  # 8
         self._text_hidden_size = self._talker_config.text_hidden_size  # 2048
-        self._text_vocab_size = self._talker_config.text_vocab_size
+        self._text_vocab_size = self._talker_config.text_vocab_size  # 151936
         self._vocab_size = self._talker_config.vocab_size  # 3072
 
         # Local Talker dims
         self._local_head_dim = self._code_predictor_config.head_dim  # 128
         self._local_hidden_size = self._code_predictor_config.hidden_size  # 1024
         self._local_num_hidden_layers = self._code_predictor_config.num_hidden_layers  # 5
+        self._local_num_attention_heads = self._code_predictor_config.num_attention_heads  # 16
         self._local_num_key_value_heads = self._code_predictor_config.num_key_value_heads  # 8
         self._local_vocab_size = self._code_predictor_config.vocab_size  # 2048
 
@@ -608,11 +615,49 @@ class Qwen3TTSInferencerONNX:
             (1, 1, self._hidden_size), np.float32, self._device, cuda_device_id
         )
 
+        # Define causal_tril for attention_mask Local prefill and step
+        _causal_tril_local = np.tril(np.ones((self._num_code_groups, self._num_code_groups), dtype=bool))
+        self._causal_tril_local = _causal_tril_local.reshape(1, 1, self._num_code_groups, self._num_code_groups)
+
+        # Define cache_write mask for each position in the Local prefill and steps, used to update KV-cache
+        # self._local_cache_write_mask_list = self._make_local_cache_write_mask(batch_size=1)  # 15 cache write masks
+        # 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()
+
         # Pre-allocate input OrtValues for Local Talker prefill
         self._local_prefill_inputs_embeds_ov = _make_device_ortvalue(
             (1, 2, self._hidden_size), np.float32, self._device, cuda_device_id
         )
         self._local_prefill_cache_position_ov = _make_device_ortvalue((2), np.int64, self._device, cuda_device_id)
+        self._local_prefill_attention_mask_ov = _make_device_ortvalue(
+            (1, 1, 2, self._num_code_groups), np.float32, self._device, cuda_device_id
+        )
+        self._local_prefill_cos_rope_ov = _make_device_ortvalue(
+            (1, 2, self._head_dim), np.float32, self._device, cuda_device_id
+        )
+        self._local_prefill_sin_rope_ov = _make_device_ortvalue(
+            (1, 2, self._head_dim), np.float32, self._device, cuda_device_id
+        )
+        # self._local_prefill_cache_write_mask_ov = _make_device_ortvalue(
+        #     (1, self._local_num_key_value_heads, self._num_code_groups, self._local_head_dim),
+        #     np.bool,
+        #     self._device,
+        #     cuda_device_id,
+        # )
+        # self._local_prefill_key_gqa_expanded_ov = _make_device_ortvalue(
+        #     (1, self._local_num_attention_heads, self._num_code_groups, self._local_head_dim),
+        #     np.float32,
+        #     self._device,
+        #     cuda_device_id,
+        # )
+        # self._local_prefill_value_gqa_expanded_ov = _make_device_ortvalue(
+        #     (1, self._local_num_attention_heads, self._num_code_groups, self._local_head_dim),
+        #     np.float32,
+        #     self._device,
+        #     cuda_device_id,
+        # )
 
         # Pre-allocate output OrtValues for Local Talker prefill
         # Backbone outputs last_hidden [1, cp_hidden_size]; logits come from lm_head
@@ -625,6 +670,33 @@ class Qwen3TTSInferencerONNX:
             (1, 1, self._hidden_size), np.float32, self._device, cuda_device_id
         )
         self._local_step_cache_position_ov = _make_device_ortvalue((1), np.int64, self._device, cuda_device_id)
+        self._local_step_attention_mask_ov = _make_device_ortvalue(
+            (1, 1, 1, self._num_code_groups), np.float32, self._device, cuda_device_id
+        )
+        self._local_step_cos_rope_ov = _make_device_ortvalue(
+            (1, 1, self._head_dim), np.float32, self._device, cuda_device_id
+        )
+        self._local_step_sin_rope_ov = _make_device_ortvalue(
+            (1, 1, self._head_dim), np.float32, self._device, cuda_device_id
+        )
+        # self._local_step_cache_write_mask_ov = _make_device_ortvalue(
+        #     (1, self._local_num_key_value_heads, self._num_code_groups, self._local_head_dim),
+        #     np.bool,
+        #     self._device,
+        #     cuda_device_id,
+        # )
+        # self._local_step_key_gqa_expanded_ov = _make_device_ortvalue(
+        #     (1, self._local_num_attention_heads, self._num_code_groups, self._local_head_dim),
+        #     np.float32,
+        #     self._device,
+        #     cuda_device_id,
+        # )
+        # self._local_step_value_gqa_expanded_ov = _make_device_ortvalue(
+        #     (1, self._local_num_attention_heads, self._num_code_groups, self._local_head_dim),
+        #     np.float32,
+        #     self._device,
+        #     cuda_device_id,
+        # )
 
         # Pre-allocate output OrtValue for the unified Local Talker backbone step
         self._local_step_hidden_ov = _make_device_ortvalue(
@@ -812,6 +884,12 @@ class Qwen3TTSInferencerONNX:
         b = self._local_prefill_bound
         b.bind_input_device("inputs_embeds", self._local_prefill_inputs_embeds_ov)
         b.bind_input_device("cache_position", self._local_prefill_cache_position_ov)
+        b.bind_input_device("attention_mask", self._local_prefill_attention_mask_ov)
+        b.bind_input_device("cos_rope", self._local_prefill_cos_rope_ov)
+        b.bind_input_device("sin_rope", self._local_prefill_sin_rope_ov)
+        # b.bind_input_device("cache_write_mask", self._local_prefill_cache_write_mask_ov)
+        # b.bind_input_device("key_gqa_expanded", self._local_prefill_key_gqa_expanded_ov)
+        # b.bind_input_device("value_gqa_expanded", self._local_prefill_value_gqa_expanded_ov)
         for i in range(self._local_num_hidden_layers):
             b.bind_input_device(f"past_key_{i}", self._kv_local_ov[2 * i])
             b.bind_input_device(f"past_value_{i}", self._kv_local_ov[2 * i + 1])
@@ -854,6 +932,12 @@ class Qwen3TTSInferencerONNX:
         b = self._local_step_bound
         b.bind_input_device("inputs_embeds", self._local_step_inputs_embeds_ov)
         b.bind_input_device("cache_position", self._local_step_cache_position_ov)
+        b.bind_input_device("attention_mask", self._local_step_attention_mask_ov)
+        b.bind_input_device("cos_rope", self._local_step_cos_rope_ov)
+        b.bind_input_device("sin_rope", self._local_step_sin_rope_ov)
+        # b.bind_input_device("cache_write_mask", self._local_step_cache_write_mask_ov)
+        # b.bind_input_device("key_gqa_expanded", self._local_step_key_gqa_expanded_ov)
+        # b.bind_input_device("value_gqa_expanded", self._local_step_value_gqa_expanded_ov)
         for i in range(self._local_num_hidden_layers):
             b.bind_input_device(f"past_key_{i}", self._kv_local_ov[2 * i])
             b.bind_input_device(f"past_value_{i}", self._kv_local_ov[2 * i + 1])
@@ -936,6 +1020,7 @@ class Qwen3TTSInferencerONNX:
         return langs
 
     # ── Helper to build RoPE cos and sin tables for the talker ────────────────
+
     def _build_rope_tables_numpy(self):
         """
         Build NumPy RoPE cosine and sine lookup tables with interleaved multimodal layout.
@@ -1039,9 +1124,97 @@ class Qwen3TTSInferencerONNX:
         # Add singleton dimension to match expected broadcast shape.
         cos = np.expand_dims(cos, axis=unsqueeze_dim)
         sin = np.expand_dims(sin, axis=unsqueeze_dim)
-        # each cos and sin with shape (bs, 1, positions, rope_dim*2)
+        # each cos and sin with shape (bs, 1, positions, rope_dim)
 
-        return cos, sin
+        return cos.astype(np.float32), sin.astype(np.float32)
+
+    def _build_local_rope_tables_numpy(self):
+        """
+        Build NumPy RoPE cosine and sine lookup tables for the local talker.
+
+        This function reproduces the logic of the provided Torch code in NumPy:
+        1. Compute inverse RoPE frequencies from `rope_theta`.
+        2. Build position-angle frequencies for `max_steps`.
+        3. Form full rotary embeddings by duplicating the frequency half.
+        4. Return `cos` and `sin` tables with an added singleton axis.
+
+        Parameters obtained from self.
+        ----------
+        _hidden_size : int
+            Total hidden size of the model.
+        _num_attention_heads : int
+            Number of attention heads.
+        _max_steps : int
+            Maximum number of positions to precompute.
+        base (rope_theta) : float
+            RoPE base frequency parameter.
+        mrope_section : sequence of int
+            Section boundaries used to interleave multimodal rotary dimensions.
+
+        Returns
+        -------
+        cos : np.ndarray
+            Cosine table with shape matching the Torch version, including the
+            extra singleton dimension inserted at axis 1.
+        sin : np.ndarray
+            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)
+
+        # Compute position-dependent phase angles.
+        inv_freq = (1.0 / (base ** (inv_idx / rope_dim))).astype(np.float32)
+        # Match the original broadcasted matmul behavior.
+        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_expanded = np.broadcast_to(
+            position_ids[None, None, :], (1, 1, len(position_ids))
+        )  # shape (bs, 1, positions)
+        freqs = np.matmul(inv_freq_expanded, position_ids_expanded).transpose(
+            0, 2, 1
+        )  # shape (bs, positions, rope_dim)
+
+        # Duplicate the half-dimension to form full rotary embeddings.
+        emb = np.concatenate([freqs, freqs], axis=-1)
+        cos = np.cos(emb)
+        sin = np.sin(emb)
+        # each cos and sin with shape (bs, positions, rope_dim)
+
+        return cos.astype(np.float32), sin.astype(np.float32)
+
+    # ── Other helper functions ────────────────────────────────────────────────
+
+    # def _make_local_cache_write_mask(self, batch_size=1):
+    #     t = np.arange(self._num_code_groups, dtype=np.int64).reshape(batch_size, 1, 1, self._num_code_groups)
+
+    #     cache_position = np.asarray([0, 1], dtype=np.int64).reshape(batch_size, 1, -1, 1)
+    #     mask = (t == cache_position).any(axis=2)  # shape [1, 1, max_seq_len]
+    #     mask = mask[..., None]  # shape [1, 1, max_seq_len, 1]
+    #     mask_list = [
+    #         np.broadcast_to(
+    #             mask, (batch_size, self._local_num_key_value_heads, self._num_code_groups, self._local_head_dim)
+    #         )
+    #     ]
+
+    #     # cache write mask for each codebook from 1 to 15 (0 is generated by talker)
+    #     for pos in range(2, self._num_code_groups):
+    #         cache_position = np.asarray([pos], dtype=np.int64).reshape(1, 1, -1, 1)
+    #         mask = (t == cache_position).any(axis=2)  # shape [1, 1, max_seq_len]
+    #         mask = mask[..., None]  # shape [1, 1, max_seq_len, 1]
+    #         mask_list.append(
+    #             np.broadcast_to(
+    #                 mask, (batch_size, self._local_num_key_value_heads, self._num_code_groups, self._local_head_dim)
+    #             )
+    #         )
+    #     # this is used for torch.where inside attention function to update KV-cache buffer in-place
+
+    #     return mask_list
 
     # ── CPU sampling ──────────────────────────────────────────────────────────
 
@@ -1200,8 +1373,8 @@ class Qwen3TTSInferencerONNX:
         cache_position = np.arange(_TALKER_PREFILL_LEN, dtype=np.int64)
         causal_rows = self._causal_tril[:, :, cache_position, :]  # [1,1,q_len,MAX]
         attn_mask = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, MAX_SEQ_LEN]
-        cos_rope = self._cos_rope[:, :, cache_position]
-        sin_rope = self._sin_rope[:, :, cache_position]
+        cos_rope = self._cos_rope[:, :, cache_position].copy()
+        sin_rope = self._sin_rope[:, :, cache_position].copy()
         self._run_talker_prefill(inputs_embeds, cache_position, attn_mask, cos_rope, sin_rope)
         self._talker_seq_len = _TALKER_PREFILL_LEN
         # Snapshot the post-prefill KV state for turn resets
@@ -1262,7 +1435,11 @@ class Qwen3TTSInferencerONNX:
         self,
         inputs_embeds: np.ndarray,  # [1, 2, talker_hidden_size]
         cache_position: np.ndarray,  # [2]
+        attention_mask: np.ndarray,  # [1, 1, 2, max_seq_len]
+        cos_rope: np.ndarray,  # [1, 1, 2, 128]
+        sin_rope: np.ndarray,  # [1, 1, 2, 128]
     ) -> np.ndarray:
+        # cache_write_mask: np.ndarray,  # [1, 8, 16, D]
         """Execute the Local Talker prefill backbone via IOBinding (CUDA graph).
 
         The local KV-cache device OrtValues are updated in-place.  Returns
@@ -1272,6 +1449,10 @@ class Qwen3TTSInferencerONNX:
         b = self._local_prefill_bound
         _copy_numpy_to_ortvalue(inputs_embeds, self._local_prefill_inputs_embeds_ov)
         _copy_numpy_to_ortvalue(cache_position, self._local_prefill_cache_position_ov)
+        _copy_numpy_to_ortvalue(attention_mask, self._local_prefill_attention_mask_ov)
+        _copy_numpy_to_ortvalue(cos_rope, self._local_prefill_cos_rope_ov)
+        _copy_numpy_to_ortvalue(sin_rope, self._local_prefill_sin_rope_ov)
+        # _copy_numpy_to_ortvalue(cache_write_mask, self._local_prefill_cache_write_mask_ov)
         b.run()
         return self._local_prefill_hidden_ov.numpy()  # [1, cp_hidden_size]
 
@@ -1281,7 +1462,11 @@ class Qwen3TTSInferencerONNX:
         self,
         inputs_embeds: np.ndarray,  # [1, 1, talker_hidden_size]
         cache_position: np.ndarray,  # [1]
+        attention_mask: np.ndarray,  # [1, 1, 1, max_seq_len]
+        cos_rope: np.ndarray,  # [1, 1, 1, 128]
+        sin_rope: np.ndarray,  # [1, 1, 1, 128]
     ) -> np.ndarray:
+        # cache_write_mask: np.ndarray,  # [1, 8, 16, 128]
         """Execute one unified Local Talker backbone step via IOBinding.
 
         Returns the last-token hidden state [1, cp_hidden_size] as NumPy.
@@ -1290,6 +1475,10 @@ class Qwen3TTSInferencerONNX:
         b = self._local_step_bound
         _copy_numpy_to_ortvalue(inputs_embeds, self._local_step_inputs_embeds_ov)
         _copy_numpy_to_ortvalue(cache_position, self._local_step_cache_position_ov)
+        _copy_numpy_to_ortvalue(attention_mask, self._local_step_attention_mask_ov)
+        _copy_numpy_to_ortvalue(cos_rope, self._local_step_cos_rope_ov)
+        _copy_numpy_to_ortvalue(sin_rope, self._local_step_sin_rope_ov)
+        # _copy_numpy_to_ortvalue(cache_write_mask, self._local_step_cache_write_mask_ov)
         b.run()
         return self._local_step_hidden_ov.numpy()  # [1, cp_hidden_size]
 
@@ -1356,8 +1545,8 @@ class Qwen3TTSInferencerONNX:
         cache_position = np.array([self._talker_seq_len], dtype=np.int64)
         causal_rows = self._causal_tril[:, :, cache_position, :]  # [1,1,q_len,MAX]
         attn_mask = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, MAX_SEQ_LEN]
-        cos_rope = self._cos_rope[:, :, cache_position]
-        sin_rope = self._sin_rope[:, :, cache_position]
+        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}")
 
@@ -1418,8 +1607,21 @@ class Qwen3TTSInferencerONNX:
         )  # [1, 2, talker_hidden_size]
         cache_position_prefill = np.array([0, 1], dtype=np.int64)
 
+        causal_rows = self._causal_tril_local[:, :, cache_position_prefill, :]  # [1,1,q_len,MAX]
+        attn_mask_prefill = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, MAX_SEQ_LEN]
+        cos_rope_prefill = self._cos_rope_local[:, cache_position_prefill].copy()
+        sin_rope_prefill = self._sin_rope_local[:, cache_position_prefill].copy()
+        # cache_write_mask_prefill = self._local_cache_write_mask_list[0]
+
         # Backbone prefill → last_hidden [1, cp_hidden_size]
-        hidden = self._run_local_prefill(inputs_embeds_prefill, cache_position_prefill)
+        hidden = self._run_local_prefill(
+            inputs_embeds_prefill,
+            cache_position_prefill,
+            attn_mask_prefill,
+            cos_rope_prefill,
+            sin_rope_prefill,
+        )
+        # cache_write_mask_prefill,
         # Batched lm_head → [15, local_vocab_size]; select group 0 (codebook group 1)
         all_logits = self._run_local_lm_head(hidden)  # [15, local_vocab_size]
         logits_g1 = all_logits[0:1, :]  # [1, local_vocab_size]
@@ -1443,11 +1645,21 @@ class Qwen3TTSInferencerONNX:
             head_idx = step_i - 1  # 1..14 (0-indexed into the 15-row logits tensor)
             cache_position_step = np.array([step_i], dtype=np.int64)
 
+            causal_rows = self._causal_tril_local[:, :, cache_position_step, :]  # [1,1,q_len,MAX]
+            attn_mask_step = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, MAX_SEQ_LEN]
+            cos_rope_step = self._cos_rope_local[:, cache_position_step].copy()
+            sin_rope_step = self._sin_rope_local[:, cache_position_step].copy()
+            # cache_write_mask_step = self._local_cache_write_mask_list[step_i - 1]
+
             # Unified backbone step → last_hidden [1, cp_hidden_size]
             hidden = self._run_local_step(
                 prev_embed.astype(np.float32),
                 cache_position_step,
+                attn_mask_step,
+                cos_rope_step,
+                sin_rope_step,
             )
+            # cache_write_mask_step,
             # Batched lm_head → [15, local_vocab_size]; select the current group
             all_logits = self._run_local_lm_head(hidden)
             logits_gi = all_logits[head_idx : head_idx + 1, :]  # [1, local_vocab_size]
@@ -1524,8 +1736,8 @@ class Qwen3TTSInferencerONNX:
         causal_rows = self._causal_tril[:, :, cache_pos, :]  # [1,1,q_len,MAX]
         attn_mask = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, MAX_SEQ_LEN]
         for _ in range(n_iter):
-            cos_rope = self._cos_rope[:, :, cache_pos]
-            sin_rope = self._sin_rope[:, :, cache_pos]
+            cos_rope = self._cos_rope[:, :, cache_pos].copy()
+            sin_rope = self._sin_rope[:, :, cache_pos].copy()
             t0 = time.perf_counter()
             self._run_talker_prefill(dummy_embeds, cache_pos, attn_mask, cos_rope, sin_rope)
             times["talker_prefill"].append((time.perf_counter() - t0) * 1000)
@@ -1540,8 +1752,8 @@ class Qwen3TTSInferencerONNX:
             cache_pos_step = np.array([_TALKER_PREFILL_LEN + n], dtype=np.int64)
             causal_rows = self._causal_tril[:, :, cache_pos_step, :]  # [1,1,q_len,MAX]
             attn_mask = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, MAX_SEQ_LEN]
-            cos_rope = self._cos_rope[:, :, cache_pos_step]
-            sin_rope = self._sin_rope[:, :, cache_pos_step]
+            cos_rope = self._cos_rope[:, :, cache_pos_step].copy()
+            sin_rope = self._sin_rope[:, :, cache_pos_step].copy()
             t0 = time.perf_counter()
             self._run_talker_step(dummy_step_embeds, cache_pos_step, attn_mask, cos_rope, sin_rope)
             times["talker_step"].append((time.perf_counter() - t0) * 1000)
@@ -1553,9 +1765,15 @@ class Qwen3TTSInferencerONNX:
         log.info("Warming-up local talker prefill...")
         dummy_lp_embeds = np.random.randn(1, _LOCAL_PREFILL_LEN, self._local_hidden_size).astype(np.float32)
         cache_pos_lp = np.array([0, 1], dtype=np.int64)
+        causal_rows = self._causal_tril_local[:, :, cache_pos_lp, :]  # [1,1,q_len,MAX]
+        attn_mask_lp = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, MAX_SEQ_LEN]
+        cos_rope_lp = self._cos_rope_local[:, cache_pos_lp].copy()
+        sin_rope_lp = self._sin_rope_local[:, cache_pos_lp].copy()
+        # cache_write_mask_lp = self._local_cache_write_mask_list[0]
         for _ in range(n_iter):
             t0 = time.perf_counter()
-            self._run_local_prefill(dummy_lp_embeds, cache_pos_lp)
+            self._run_local_prefill(dummy_lp_embeds, cache_pos_lp, attn_mask_lp, cos_rope_lp, sin_rope_lp)
+            # cache_write_mask_lp
             times["local_prefill"].append((time.perf_counter() - t0) * 1000)
         exec_time_tot = sum(times["local_prefill"]) / 1000
         avg_exec_time = exec_time_tot / len(times["local_prefill"]) * 1000
@@ -1568,9 +1786,15 @@ class Qwen3TTSInferencerONNX:
         log.info("Warming-up local talker step...")
         dummy_ls_embeds = np.random.randn(1, 1, self._hidden_size).astype(np.float32)
         cache_pos_ls = np.array([2], dtype=np.int64)
+        causal_rows = self._causal_tril_local[:, :, cache_pos_ls, :]  # [1,1,q_len,MAX]
+        attn_mask_ls = np.where(causal_rows, 0.0, -np.inf).astype(np.float32)  # [1, 1, q_len, MAX_SEQ_LEN]
+        cos_rope_ls = self._cos_rope_local[:, cache_pos_ls].copy()
+        sin_rope_ls = self._sin_rope_local[:, cache_pos_ls].copy()
+        # cache_write_mask_ls = self._local_cache_write_mask_list[1]
         for _ in range(n_iter):
             t0 = time.perf_counter()
-            self._run_local_step(dummy_ls_embeds, cache_pos_ls)
+            self._run_local_step(dummy_ls_embeds, cache_pos_ls, attn_mask_ls, cos_rope_ls, sin_rope_ls)
+            # cache_write_mask_ls
             times["local_step"].append((time.perf_counter() - t0) * 1000)
         exec_time_tot = sum(times["local_step"]) / 1000
         avg_exec_time = exec_time_tot / len(times["local_step"]) * 1000

test_qwen3-tts-streaming_onnx.py

@@ -24,12 +24,7 @@ This script demonstrates the full Qwen3-TTS-Streaming-ONNX pipeline by:
 5. Writing the concatenated audio to a WAV file.
 Usage:
     python test_qwen3-tts-streaming_onnx.py \
-        --talker_model_path qwen3-tts_onnx/talker_model.onnx \
-        --talker_local_model_path qwen3-tts_onnx/talker_local_model.onnx \
-        --codec_decoder_model_path qwen3-tts_onnx/codec_decoder_model.onnx \
-        --speaker_encoder_model_path qwen3-tts_onnx/speaker_encoder_model.onnx \
-        --talker_codec_embed_model_path qwen3-tts_onnx/talker_codec_embed_model.onnx \
-        --text_embed_proj_model_path qwen3-tts_onnx/text_embed_proj_model.onnx \
+        --onnx_dir qwen3-tts_onnx/ \
         --model_config_path configs/config.json \
         --codec_config_path configs/tokenizer_config.json \
         --preprocessor_config_dir configs/ \
@@ -281,8 +276,10 @@ def main() -> None:
     suffix = "_int8.onnx" if args.use_int8 else ".onnx"
 
     # Build model paths
-    talker_path = _onnx_path(args.onnx_dir, f"talker_model{suffix}")
-    talker_local_path = _onnx_path(args.onnx_dir, f"talker_local_model{suffix}")
+    talker_prefill_path = _onnx_path(args.onnx_dir, f"talker_model_prefill{suffix}")
+    talker_step_path = _onnx_path(args.onnx_dir, f"talker_model_step{suffix}")
+    talker_local_prefill_path = _onnx_path(args.onnx_dir, f"talker_local_model_prefill{suffix}")
+    talker_local_step_path = _onnx_path(args.onnx_dir, f"talker_local_model_step{suffix}")
     # Unified backbone step (shared for steps 2..15) and batched lm_head
     talker_local_lm_head_path = _onnx_path(args.onnx_dir, f"talker_local_lm_head{suffix}")
 
@@ -294,8 +291,10 @@ def main() -> None:
 
     # Validate required paths
     required = [
-        talker_path,
-        talker_local_path,
+        talker_prefill_path,
+        talker_step_path,
+        talker_local_prefill_path,
+        talker_local_step_path,
         talker_local_lm_head_path,
         codec_decoder_path,
         speaker_encoder_path,
@@ -312,8 +311,10 @@ def main() -> None:
 
     log.info("Building inferencer...")
     inferencer = Qwen3TTSInferencerONNX(
-        talker_model_path=talker_path,
-        talker_local_model_path=talker_local_path,
+        talker_model_prefill_path=talker_prefill_path,
+        talker_model_step_path=talker_step_path,
+        talker_local_model_prefill_path=talker_local_prefill_path,
+        talker_local_model_step_path=talker_local_step_path,
         talker_local_lm_head_model_path=talker_local_lm_head_path,
         codec_decoder_model_path=codec_decoder_path,
         speaker_encoder_model_path=speaker_encoder_path,