def decode_step(
        self,
        tgt_ids_Bx1xC: mx.array,  # [B, 1, C]
        tgt_pos_Bx1: mx.array,  # [B, 1]
        encoder_out: mx.array,  # [B, S, E]
        self_attn_mask: Any,  # None
        cross_attn_mask: mx.array,  # [B, 1, 1, S]
        self_attention_cache: List[KVCache],
        cross_attention_cache: List[KVCache],
    ) -> mx.array:
        """
        Performs a single decoding step, managing KV caches layer by layer.

        Returns:
            A tuple containing:
            - logits_Bx1xCV: The final output logits for the current step (B, 1, C*V), cast to float32.
            - new_cache: The updated KV cache for the next decoding step.
        """
        assert (
            self_attn_mask is None
        ), "Self-attention mask should be None, kept for pattern"

        x = None
        for i in range(self.num_channels):
            channel_tokens = tgt_ids_Bx1xC[..., i]
            channel_embed = self.embeddings[i](channel_tokens)
            x = channel_embed if x is None else x + channel_embed

        for i, layer in enumerate(self.layers):
            self_cache = self_attention_cache[i]
            cross_cache = cross_attention_cache[i]
            x = layer(
                x,  # (2, 1, D)
                encoder_out,  # (2, S, E)
                src_positions=None,  # CA KV is already computed
                tgt_positions=tgt_pos_Bx1,  # (2, 1)
                deterministic=True,
                self_attn_mask=None,
                cross_attn_mask=cross_attn_mask,
                self_attn_cache=self_cache,
                cross_attn_cache=cross_cache,
            )

        x = self.norm(x)
        logits_Bx1xCxV = self.logits_dense(x)

        # Convert to float32 if needed
        if logits_Bx1xCxV.dtype != mx.float32:
            logits_Bx1xCxV = logits_Bx1xCxV.astype(mx.float32)

        return logits_Bx1xCxV

        Performs a single decoding step, managing KV caches layer by layer.

        Returns:
            A tuple containing:
            - logits_Bx1xCV: The final output logits for the current step (B, 1, C*V), cast to float32.
            - new_cache: The updated KV cache for the next decoding step.
        
mlx_audio/tts/models/dia/layers.py