[Qwen3-TTS] float32 speaker embedding silently promotes the whole talker prefill + KV cache to float32, ~2.4x slower decode on bf16 checkpoints

[fwoeck]

Summary

With x-vector conditioning (ref_audio without ref_text), extract_speaker_embedding() in mlx_audio/tts/models/qwen3_tts/qwen3_tts.py returns a float32 embedding (the speaker encoder is excluded from quantization by model_quant_predicate, and no cast happens on return). In _prepare_generation_inputs, that float32 vector is concatenated with the talker's bfloat16 input embeddings:

codec_embed = mx.concatenate(
    [
        codec_embed,                      # bfloat16 (talker input embeddings)
        speaker_embed.reshape(1, 1, -1),  # float32 (speaker encoder output)
        codec_embed_suffix,               # bfloat16
    ],
    axis=1,
)

MLX's type promotion upcasts the result to float32 — and with it the entire talker prefill and, downstream, the whole KV cache, for the entire generation. On a bf16 checkpoint the model then decodes in float32 throughout.

Measured impact

mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16, M1 Pro (32 GB), paired A/B over the same seeds and sentences, streaming decode:

	unpatched	with cast
RTF	1.60	0.66
TTFB	0.73 s	0.40 s

No audible change in cloning quality. The unpatched RTF > 1.0 means real-time streaming playback falls behind and stutters on this hardware; with the cast it runs comfortably ahead.

Repro / verification

import mlx.core as mx
from mlx_audio.tts.utils import load_model
from mlx_audio.utils import load_audio

model = load_model("mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16")
audio = load_audio("reference.wav", sample_rate=24000)
print(model.extract_speaker_embedding(audio).dtype)  # float32, talker embeds are bfloat16

Timing the same generate(..., ref_audio=...) call with and without the workaround below shows the throughput difference directly.

Workaround we ship

inner = model.extract_speaker_embedding
model.extract_speaker_embedding = lambda *a, **k: inner(*a, **k).astype(mx.bfloat16)

Suggested fix

Cast the speaker embedding to the talker's embedding dtype before the concatenate, e.g. in _prepare_generation_inputs:

speaker_embed = speaker_embed.astype(codec_embed.dtype)

(or equivalently at the end of extract_speaker_embedding, casting to the talker input-embedding dtype). Keeping the speaker encoder itself in float32 is fine — it is only the handoff into the bf16 talker that needs the cast.

Possibly related: #827 reports Base-model cloning ~2.4-4x slower in a multi-model server — same magnitude, though the trigger described there is order-dependent, so it may be a different mechanism on top of this one.

Observed in 0.4.7; the code on current main is unchanged at both sites (no cast in extract_speaker_embedding, none at the concatenate).

[linear-code]

BLA-488

[lucasnewman]

@fwoeck Good catch, thanks! Fixing here: #879