[Perf][Qwen3-TTS] Free unused decoder in Talker SpeechTokenizer to VRAM

[Sy0307]

Summary

The Talker stage lazily loads the full Qwen3TTSTokenizerV2Model (encoder + decoder) onto GPU when processing voice-clone requests (_ensure_speech_tokenizer_loaded). However, Talker only calls tok.encode() -> model.encoder.encode() -- the decoder is never used. This wastes ~331 MiB of GPU memory (bf16) for the duration of the session.

This PR moves only the encoder sub-module to GPU and immediately frees the decoder weights, saving ~331 MiB VRAM with zero functional impact.

Motivation

Qwen3-TTS uses a 2-stage pipeline where the SpeechTokenizer serves different roles in each stage:

Stage	Component	Uses encoder	Uses decoder
Stage 0 (Talker)	_ensure_speech_tokenizer_loaded()	Yes (tok.encode() for ref audio)	No
Stage 1 (Code2Wav)	_ensure_speech_tokenizer_loaded()	No	Yes (tok.model.decoder)

The existing code does tok.model.to(dev), which moves both encoder and decoder to GPU. Since Talker never touches the decoder, those weights sit idle on the GPU.

Changes

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py -- 3-line change in _ensure_speech_tokenizer_loaded():

• tok.model.to(dev) -> tok.model.encoder.to(dev) (only move encoder to GPU)
• Add tok.model.decoder = None (free decoder weights from CPU memory as well)
• Update error message to reflect the narrower scope

Why this is safe

1. tok.encode() only calls self.encoder.encode() -- verified in Qwen3TTSTokenizerV2Model.encode(). The decoder is not referenced anywhere in the encode path.

2. tok.model.dtype still works -- after setting decoder = None, the model still has encoder parameters, so nn.Module.dtype resolves correctly.

3. Code2Wav is unaffected -- it runs in a separate EngineCore process with its own _ensure_speech_tokenizer_loaded().

4. Encode output is bit-identical -- verified on H20 GPU with Qwen3-TTS-12Hz-1.7B-Base, 0/400 codes differ between baseline and optimized.

Test Plan

Verified on H20 (98 GB) with Qwen3-TTS-12Hz-1.7B-Base:

# Bit-exact comparison: baseline (full model on GPU) vs optimized (encoder only)
tok1 = Qwen3TTSTokenizer.from_pretrained(dir, torch_dtype=torch.bfloat16)
tok1.model.to(dev)  # baseline
codes1 = tok1.encode(wav, sr=sr).audio_codes[0]

tok2 = Qwen3TTSTokenizer.from_pretrained(dir, torch_dtype=torch.bfloat16)
tok2.model.encoder.to(dev); tok2.model.decoder = None  # optimized
codes2 = tok2.encode(wav, sr=sr).audio_codes[0]

torch.equal(codes1, codes2)  # True -- bit-identical

VRAM measurement

Config	VRAM	Saved
Full model on GPU (baseline)	504.1 MiB	--
Encoder only on GPU (this PR)	173.3 MiB	330.8 MiB

E2E tests:

python -m pytest tests/e2e/online_serving/test_qwen3_tts.py -v --timeout=300
python -m pytest tests/e2e/online_serving/test_qwen3_tts_base.py -v --timeout=300

[hsliuustc0106]

@linyueqian

[Copilot]

Pull request overview

Optimizes the Qwen3-TTS Talker stage’s SpeechTokenizer GPU memory usage by moving only the encoder to CUDA (since Talker only encodes reference audio) and dropping the unused decoder to reduce VRAM footprint during voice-clone requests.

Changes:

• Move tok.model.encoder to GPU instead of the full tokenizer model.
• Explicitly release the tokenizer decoder module in Talker (tok.model.decoder = None).
• Refine the error message to reflect encoder-only device movement.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The comment hard-codes a VRAM saving estimate (~218 MiB), but the PR description and measurements cite ~331 MiB saved. To avoid future confusion (and because this can vary by checkpoint/dtype), consider either updating the number to match the verified measurement or removing the specific MiB figure from the comment.

Suggested change

	# calls tok.encode()) and would waste ~218 MiB bf16 VRAM.
	# calls tok.encode()) and would otherwise consume additional bf16 VRAM.

[hsliuustc0106]

i test it in 0.6B model w/o this PR, it looks help save some memory during request. cc @linyueqian

[lishunyang12]

looks good, small and safe optimization. +1 to linyueqian's point about documenting the encode-only constraint on the cached instance.

[linyueqian]

LGTM