fix: bind all inputs for all talker & local

- All inputs can be bound as shapes are
static.
- Using in_place to replace the values
of the buffer
- Codec decoder next, as there is still
a bug in it.

src/inference/qwen3_tts_inferencer_onnx.py

@@ -440,7 +440,8 @@ class Qwen3TTSInferencerONNX:
 
         # Non-graphed utility sessions
         log.info(f"  codec decoder   <- {codec_decoder_model_path}")
-        self._codec_decoder_sess = _sess(codec_decoder_model_path, graph=self._enable_cuda_graph)
+        # self._codec_decoder_sess = _sess(codec_decoder_model_path, graph=self._enable_cuda_graph)
+        self._codec_decoder_sess = _sess(codec_decoder_model_path, graph=False)
         log.info(f"  speaker encoder <- {speaker_encoder_model_path}")
         self._speaker_encoder = _sess(speaker_encoder_model_path, graph=False)
         log.info(f"  codec embed     <- {talker_codec_embed_model_path}")
@@ -560,6 +561,23 @@ class Qwen3TTSInferencerONNX:
         # Define cos and sin table for RoPE embedding
         self._cos_rope, self._sin_rope = self._build_rope_tables_numpy()
 
+        # Pre-allocate input OrtValues for Talker prefill
+        self._talker_prefill_inputs_embeds_ov = _make_device_ortvalue(
+            (1, _TALKER_PREFILL_LEN, self._hidden_size), np.float32, self._device, cuda_device_id
+        )
+        self._talker_prefill_cache_position_ov = _make_device_ortvalue(
+            (_TALKER_PREFILL_LEN), np.int64, self._device, cuda_device_id
+        )
+        self._talker_prefill_attention_mask_ov = _make_device_ortvalue(
+            (1, 1, _TALKER_PREFILL_LEN, _TALKER_MAX_SEQ_LEN), np.float32, self._device, cuda_device_id
+        )
+        self._talker_prefill_cos_rope_ov = _make_device_ortvalue(
+            (1, 1, _TALKER_PREFILL_LEN, self._head_dim), np.float32, self._device, cuda_device_id
+        )
+        self._talker_prefill_sin_rope_ov = _make_device_ortvalue(
+            (1, 1, _TALKER_PREFILL_LEN, self._head_dim), np.float32, self._device, cuda_device_id
+        )
+
         # Pre-allocate output OrtValues for Talker prefill
         # Prefill outputs logits over the full q_len=9 sequence; we only need
         # the last position for sampling, but we allocate for the full length.
@@ -570,23 +588,55 @@ class Qwen3TTSInferencerONNX:
             (1, 1, self._hidden_size), np.float32, self._device, cuda_device_id
         )
 
+        # Pre-allocate input OrtValues for Talker step
+        self._talker_inputs_embeds_ov = _make_device_ortvalue(
+            (1, 1, self._hidden_size), np.float32, self._device, cuda_device_id
+        )
+        self._talker_cache_position_ov = _make_device_ortvalue((1), np.int64, self._device, cuda_device_id)
+        self._talker_attention_mask_ov = _make_device_ortvalue(
+            (1, 1, 1, _TALKER_MAX_SEQ_LEN), np.float32, self._device, cuda_device_id
+        )
+        self._talker_cos_rope_ov = _make_device_ortvalue(
+            (1, 1, 1, self._head_dim), np.float32, self._device, cuda_device_id
+        )
+        self._talker_sin_rope_ov = _make_device_ortvalue(
+            (1, 1, 1, self._head_dim), np.float32, self._device, cuda_device_id
+        )
+
         # Pre-allocate output OrtValues for the Talker step
         self._talker_logits_ov = _make_device_ortvalue((1, self._vocab_size), np.float32, self._device, cuda_device_id)
         self._talker_hidden_ov = _make_device_ortvalue(
             (1, 1, self._hidden_size), np.float32, self._device, cuda_device_id
         )
 
+        # 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)
+
         # Pre-allocate output OrtValues for Local Talker prefill
         # Backbone outputs last_hidden [1, cp_hidden_size]; logits come from lm_head
         self._local_prefill_hidden_ov = _make_device_ortvalue(
             (1, self._local_hidden_size), np.float32, self._device, cuda_device_id
         )
 
+        # Pre-allocate input OrtValues for Local Talker step
+        self._local_step_inputs_embeds_ov = _make_device_ortvalue(
+            (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)
+
         # Pre-allocate output OrtValue for the unified Local Talker backbone step
         self._local_step_hidden_ov = _make_device_ortvalue(
             (1, self._local_hidden_size), np.float32, self._device, cuda_device_id
         )
 
+        # Pre-allocate input OrtValue for the batched lm_head
+        self._local_lm_head_hidden_states_ov = _make_device_ortvalue(
+            (1, self._local_hidden_size), np.float32, self._device, cuda_device_id
+        )
+
         # Pre-allocate output OrtValue for the batched lm_head
         # Shape: [15, local_vocab_size] — all codebook group logits at once
         self._local_lm_head_logits_ov = _make_device_ortvalue(
@@ -622,6 +672,14 @@ class Qwen3TTSInferencerONNX:
         self._codec_decoder_bound = _BoundSession(self._codec_decoder_sess, self._device, cuda_device_id)
         self._codec_step_idx = 0
 
+        # 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
+        )
+
         # Pre-allocate OrtValues for Codec Decoder
         self._codec_wav_ov = _make_device_ortvalue(
             (1, 1, self._codec_chunk_frames * self._speech_tokenizer_decoder_total_upsample),
@@ -707,6 +765,8 @@ class Qwen3TTSInferencerONNX:
         b.bind_output_device("current_hidden_state_cache", self._codec_hidden_cache_ov)
         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)
 
@@ -729,6 +789,11 @@ class Qwen3TTSInferencerONNX:
         here; they are re-bound on every call in ``_run_talker_prefill``.
         """
         b = self._talker_prefill_bound
+        b.bind_input_device("inputs_embeds", self._talker_prefill_inputs_embeds_ov)
+        b.bind_input_device("cache_position", self._talker_prefill_cache_position_ov)
+        b.bind_input_device("attention_mask", self._talker_prefill_attention_mask_ov)
+        b.bind_input_device("cos_rope", self._talker_prefill_cos_rope_ov)
+        b.bind_input_device("sin_rope", self._talker_prefill_sin_rope_ov)
         for i in range(self._num_hidden_layers):
             b.bind_input_device(f"past_key_{i}", self._kv_talker_ov[2 * i])
             b.bind_input_device(f"past_value_{i}", self._kv_talker_ov[2 * i + 1])
@@ -746,6 +811,8 @@ class Qwen3TTSInferencerONNX:
         Dynamic inputs are re-bound on every call in ``_run_local_prefill``.
         """
         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)
         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])
@@ -765,6 +832,11 @@ class Qwen3TTSInferencerONNX:
         here; they are re-bound on every call in ``_run_talker_step``.
         """
         b = self._talker_step_bound
+        b.bind_input_device("inputs_embeds", self._talker_inputs_embeds_ov)
+        b.bind_input_device("cache_position", self._talker_cache_position_ov)
+        b.bind_input_device("attention_mask", self._talker_attention_mask_ov)
+        b.bind_input_device("cos_rope", self._talker_cos_rope_ov)
+        b.bind_input_device("sin_rope", self._talker_sin_rope_ov)
         for i in range(self._num_hidden_layers):
             b.bind_input_device(f"past_key_{i}", self._kv_talker_ov[2 * i])
             b.bind_input_device(f"past_value_{i}", self._kv_talker_ov[2 * i + 1])
@@ -781,6 +853,8 @@ class Qwen3TTSInferencerONNX:
         is applied separately by ``_wire_local_lm_head_bindings``.
         """
         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)
         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])
@@ -796,6 +870,7 @@ class Qwen3TTSInferencerONNX:
         bound to the pre-allocated device OrtValue.
         """
         b = self._local_lm_head_bound
+        b.bind_input_device("hidden_states", self._local_lm_head_hidden_states_ov)
         b.bind_output_device("logits", self._local_lm_head_logits_ov)
 
     # ── Prefill KV snapshot helpers ───────────────────────────────────────────
@@ -1152,11 +1227,11 @@ class Qwen3TTSInferencerONNX:
         and ``_talker_prefill_hidden_ov`` respectively.
         """
         b = self._talker_prefill_bound
-        b.bind_input_cpu("inputs_embeds", inputs_embeds)
-        b.bind_input_cpu("cache_position", cache_position)
-        b.bind_input_cpu("attention_mask", attention_mask)
-        b.bind_input_cpu("cos_rope", cos_rope)
-        b.bind_input_cpu("sin_rope", sin_rope)
+        _copy_numpy_to_ortvalue(inputs_embeds, self._talker_prefill_inputs_embeds_ov)
+        _copy_numpy_to_ortvalue(cache_position, self._talker_prefill_cache_position_ov)
+        _copy_numpy_to_ortvalue(attention_mask, self._talker_prefill_attention_mask_ov)
+        _copy_numpy_to_ortvalue(cos_rope, self._talker_prefill_cos_rope_ov)
+        _copy_numpy_to_ortvalue(sin_rope, self._talker_prefill_sin_rope_ov)
         b.run()
 
     # ── Talker step (IOBinding + CUDA graph) ──────────────────────────────────
@@ -1172,11 +1247,11 @@ class Qwen3TTSInferencerONNX:
         """Execute one Talker step via IOBinding and return logits as NumPy."""
         b = self._talker_step_bound
         # Re-bind the two small dynamic inputs (H→D copy, ~4 KB total)
-        b.bind_input_cpu("inputs_embeds", inputs_embeds)
-        b.bind_input_cpu("cache_position", cache_position)
-        b.bind_input_cpu("attention_mask", attention_mask)
-        b.bind_input_cpu("cos_rope", cos_rope)
-        b.bind_input_cpu("sin_rope", sin_rope)
+        _copy_numpy_to_ortvalue(inputs_embeds, self._talker_inputs_embeds_ov)
+        _copy_numpy_to_ortvalue(cache_position, self._talker_cache_position_ov)
+        _copy_numpy_to_ortvalue(attention_mask, self._talker_attention_mask_ov)
+        _copy_numpy_to_ortvalue(cos_rope, self._talker_cos_rope_ov)
+        _copy_numpy_to_ortvalue(sin_rope, self._talker_sin_rope_ov)
         # All KV and output bindings are already wired statically
         b.run()
         # Copy logits to CPU for sampling (~12 KB for vocab_size=3072)
@@ -1196,8 +1271,8 @@ class Qwen3TTSInferencerONNX:
         Logits are obtained by calling ``_run_local_lm_head`` with this hidden.
         """
         b = self._local_prefill_bound
-        b.bind_input_cpu("inputs_embeds", inputs_embeds)
-        b.bind_input_cpu("cache_position", cache_position)
+        _copy_numpy_to_ortvalue(inputs_embeds, self._local_prefill_inputs_embeds_ov)
+        _copy_numpy_to_ortvalue(cache_position, self._local_prefill_cache_position_ov)
         b.run()
         return self._local_prefill_hidden_ov.numpy()  # [1, cp_hidden_size]
 
@@ -1214,8 +1289,8 @@ class Qwen3TTSInferencerONNX:
         Logits are obtained by calling ``_run_local_lm_head`` with this hidden.
         """
         b = self._local_step_bound
-        b.bind_input_cpu("inputs_embeds", inputs_embeds)
-        b.bind_input_cpu("cache_position", cache_position)
+        _copy_numpy_to_ortvalue(inputs_embeds, self._local_step_inputs_embeds_ov)
+        _copy_numpy_to_ortvalue(cache_position, self._local_step_cache_position_ov)
         b.run()
         return self._local_step_hidden_ov.numpy()  # [1, cp_hidden_size]
 
@@ -1231,7 +1306,7 @@ class Qwen3TTSInferencerONNX:
         ``logits[head_idx]`` on the CPU for sampling.
         """
         b = self._local_lm_head_bound
-        b.bind_input_cpu("hidden_states", hidden)
+        _copy_numpy_to_ortvalue(hidden, self._local_lm_head_hidden_states_ov)
         b.run()
         return self._local_lm_head_logits_ov.numpy()  # [15, local_vocab_size]
 
@@ -1246,11 +1321,16 @@ class Qwen3TTSInferencerONNX:
         b = self._codec_decoder_bound
         # Re-bind the two small dynamic inputs (H→D copy, ~4 KB total)
         b.bind_input_cpu("codes", chunk_tokens)
-        b.bind_input_cpu("cache_position", pos)
+        # log.info(f"run_codec_decoder chunk_tokens {chunk_tokens.shape} {chunk_tokens.dtype}")
+        # b.bind_input_cpu("cache_position", pos)
+        # log.info(f"run_codec_decoder pos {pos.shape} {pos.dtype}")
+        # _copy_numpy_to_ortvalue(chunk_tokens, self._codec_codes_ov)
+        _copy_numpy_to_ortvalue(pos, self._codec_cache_position_ov)
         # All KV and output bindings are already wired statically
         b.run()
         # Copy wav to CPU
-        return b.get_output_numpy("wav")
+        # return b.get_output_numpy("wav")
+        return self._codec_wav_ov.numpy()
 
     # ── Main autoregressive step ──────────────────────────────────────────────
 
@@ -1285,9 +1365,14 @@ class Qwen3TTSInferencerONNX:
         cos_rope = self._cos_rope[:, :, cache_position]
         sin_rope = self._sin_rope[:, :, cache_position]
         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}")
 
         # 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}"
+        )
         self._talker_seq_len += 1
 
         # ── Sample first token (CPU) ──────────────────────────────────────────