From 61f1d29f95827b70587bcd73e1284ac3dd2f30f7 Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Mon, 3 Aug 2026 03:53:14 -0700 Subject: [PATCH 1/2] Add CosyVoice3 + DSpark speculative decoding via vllm V1 API Integrates DSpark speculative decoding into the CosyVoice3 inference pipeline without requiring any manual export scripts. cosyvoice/utils/file_utils.py - export_cosyvoice3_vllm_spec(): exports CosyVoice3LM to HF format with extended embed_tokens/lm_head for DSpark compatibility (speech tokens at offset text_vocab_size=151924). Skips if output dir exists, mirroring export_cosyvoice2_vllm(). Saves cosyvoice3_metadata.json with speech_token_offset for downstream mapping. cosyvoice/cli/model.py - CosyVoice3Model.load_vllm_spec(model_dir, draft_model_dir): auto-exports via export_cosyvoice3_vllm_spec() into /hf_spec, then creates a vllm V1 LLM with DSpark speculative_config (probabilistic sampling + draft_apply_repetition_penalty=True). Frees PyTorch layers after loading. cosyvoice/llm/llm.py - CosyVoice3LM._inference_spec(): token-based speculative generation using the Qwen2 chat template. Decodes text tokens back to string, encodes prompt speech tokens as <|s_N|>, calls vllm_spec.generate(), and maps output token IDs back to speech IDs via speech_token_offset. - Qwen2LM.inference(): branches to _inference_spec() when vllm_spec is set. vllm_example.py - cosyvoice3_spec_example(): two-line usage matching load_vllm() style. Usage: cosyvoice.model.load_vllm_spec( 'pretrained_models/Fun-CosyVoice3-0.5B', 'yuekai/cosyvoice3_llm_dspark', ) Requires the speculative vllm fork (yuekaizhang/vllm branch dspark-draft-sampling-mirrors) with vllm-omni 0.25.1 compiled extensions via PYTHONPATH (enable_prompt_embeds is incompatible with DSpark's V2 model runner requirement; token-based path avoids this constraint). Co-Authored-By: Claude Sonnet 4.6 (1M context) Signed-off-by: Yuekai Zhang --- cosyvoice/cli/model.py | 43 +++++++++++ cosyvoice/llm/llm.py | 39 +++++++++- cosyvoice/utils/file_utils.py | 135 ++++++++++++++++++++++++++++++++++ vllm_example.py | 33 ++++++++- 4 files changed, 248 insertions(+), 2 deletions(-) diff --git a/cosyvoice/cli/model.py b/cosyvoice/cli/model.py index 92a15d985..7713d9a10 100644 --- a/cosyvoice/cli/model.py +++ b/cosyvoice/cli/model.py @@ -422,6 +422,49 @@ def __init__(self, # FSQ silent and breath token self.silent_tokens = [1, 2, 28, 29, 55, 248, 494, 2241, 2242, 2322, 2323] + def load_vllm_spec(self, model_dir, draft_model_dir, gpu_memory_utilization=0.6): + """Load CosyVoice3 + DSpark speculative decoding (vllm V1 API, token-based). + + `model_dir` is the CosyVoice3 pretrained model directory (same as used to load + the model, e.g. 'pretrained_models/Fun-CosyVoice3-0.5B'). The HuggingFace export + is created automatically under /hf_spec the first time (like + export_cosyvoice2_vllm() for load_vllm()). + `draft_model_dir` is a speculators-format DSpark checkpoint (local path or HF repo). + + After loading, self.llm.vllm_spec is used automatically inside tts() calls. + """ + import json + from transformers import AutoTokenizer + from vllm import LLM + from cosyvoice.utils.file_utils import export_cosyvoice3_vllm_spec + + hf_llm_dir = os.path.join(model_dir, 'CosyVoice-BlankEN') + hf_model_dir = os.path.join(model_dir, 'hf_spec') + export_cosyvoice3_vllm_spec(self.llm, hf_model_dir, hf_llm_dir, self.device) + + with open(os.path.join(hf_model_dir, 'cosyvoice3_metadata.json')) as f: + metadata = json.load(f) + with open(os.path.join(draft_model_dir, 'config.json')) as f: + draft_cfg = json.load(f) + method = draft_cfg.get('speculators_model_type', 'dspark') + num_spec_tokens = draft_cfg.get('block_size', 8) - 1 + + self.llm.vllm_spec_tokenizer = AutoTokenizer.from_pretrained(hf_model_dir) + self.llm.vllm_spec_speech_offset = metadata['speech_token_offset'] + self.llm.vllm_spec = LLM( + model=hf_model_dir, + speculative_config={ + 'model': draft_model_dir, + 'method': method, + 'num_speculative_tokens': num_spec_tokens, + 'draft_sample_method': 'probabilistic', + 'draft_apply_repetition_penalty': True, + }, + gpu_memory_utilization=gpu_memory_utilization, + ) + # free PyTorch model weights; vllm manages GPU memory from here + del self.llm.llm.model.model.layers + def token2wav(self, token, prompt_token, prompt_feat, embedding, token_offset, uuid, stream=False, finalize=False, speed=1.0): with torch.cuda.amp.autocast(self.fp16): tts_mel, _ = self.flow.inference(token=token.to(self.device, dtype=torch.int32), diff --git a/cosyvoice/llm/llm.py b/cosyvoice/llm/llm.py index e8e81d942..bbc11f981 100644 --- a/cosyvoice/llm/llm.py +++ b/cosyvoice/llm/llm.py @@ -498,6 +498,10 @@ def inference( max_len = int((text_len - prompt_text_len) * max_token_text_ratio) # 5. step by step decode + # speculative vllm path: token-based V1 API (CosyVoice3 only) + if hasattr(self, 'vllm_spec'): + yield from self._inference_spec(text, prompt_speech_token, min_len, max_len) + return for token in self.inference_wrapper(lm_input, sampling, min_len, max_len, uuid): yield token @@ -703,4 +707,37 @@ def __init__( self.stop_token_ids = [speech_token_size + i for i in range(200)] self.vllm_output_queue = {} if online_feature is True: - self.speech_token_extractor = SpeechTokenExtractor(model_path=os.path.join(onnx_path, 'speech_tokenizer_v3.batch.onnx')) \ No newline at end of file + self.speech_token_extractor = SpeechTokenExtractor(model_path=os.path.join(onnx_path, 'speech_tokenizer_v3.batch.onnx')) + + def _inference_spec(self, text: torch.Tensor, prompt_speech_token: torch.Tensor, + min_len: int, max_len: int): + """Token-based speculative decoding via vllm V1 LLM (set by load_vllm_spec). + + `text` is the concatenated [prompt_text, target_text] as Qwen2 token IDs. + `prompt_speech_token` holds voice-clone speech token IDs (0..speech_token_size). + Output: yields integer speech token IDs. + """ + from vllm import SamplingParams + offset = self.vllm_spec_speech_offset + # only pass the primary eos stop token (vllm caps stop_token_ids at 128); + # the yield loop below also filters any token outside [0, speech_token_size) + stop_ids = [offset + self.eos_token] + + text_str = self.vllm_spec_tokenizer.decode(text[0].tolist(), skip_special_tokens=False) + speech_str = ''.join(f'<|s_{t}|>' for t in prompt_speech_token[0].tolist()) if prompt_speech_token.shape[1] > 0 else '' + + chat = [{'role': 'user', 'content': text_str}] + if speech_str: + chat.append({'role': 'assistant', 'content': speech_str}) + prompt = self.vllm_spec_tokenizer.apply_chat_template( + chat, tokenize=False, continue_final_message=True) + + params = SamplingParams( + temperature=0.8, top_p=0.95, top_k=15, repetition_penalty=1.1, + min_tokens=min_len, max_tokens=max_len, stop_token_ids=stop_ids) + output = self.vllm_spec.generate([prompt], params, use_tqdm=False)[0] + for token_id in output.outputs[0].token_ids: + speech_id = token_id - offset + if speech_id < 0 or speech_id >= self.speech_token_size: + break + yield speech_id \ No newline at end of file diff --git a/cosyvoice/utils/file_utils.py b/cosyvoice/utils/file_utils.py index b173ef201..bab30e1a3 100644 --- a/cosyvoice/utils/file_utils.py +++ b/cosyvoice/utils/file_utils.py @@ -89,6 +89,141 @@ def convert_onnx_to_trt(trt_model, trt_kwargs, onnx_model, fp16): # NOTE do not support bistream inference as only speech token embedding/head is kept +_COSYVOICE3_TEXT_SPECIAL_TOKENS = [ + '<|im_start|>', '<|im_end|>', '<|endofprompt|>', + '[breath]', '', '', '[noise]', '[laughter]', '[cough]', + '[clucking]', '[accent]', '[quick_breath]', '', '', + '[hissing]', '[sigh]', '[vocalized-noise]', '[lipsmack]', '[mn]', '<|endofsystem|>', + '[AA]', '[AA0]', '[AA1]', '[AA2]', '[AE]', '[AE0]', '[AE1]', '[AE2]', + '[AH]', '[AH0]', '[AH1]', '[AH2]', '[AO]', '[AO0]', '[AO1]', '[AO2]', + '[AW]', '[AW0]', '[AW1]', '[AW2]', '[AY]', '[AY0]', '[AY1]', '[AY2]', + '[B]', '[CH]', '[D]', '[DH]', '[EH]', '[EH0]', '[EH1]', '[EH2]', + '[ER]', '[ER0]', '[ER1]', '[ER2]', '[EY]', '[EY0]', '[EY1]', '[EY2]', + '[F]', '[G]', '[HH]', '[IH]', '[IH0]', '[IH1]', '[IH2]', + '[IY]', '[IY0]', '[IY1]', '[IY2]', '[JH]', '[K]', '[L]', '[M]', '[N]', + '[NG]', '[OW]', '[OW0]', '[OW1]', '[OW2]', '[OY]', '[OY0]', '[OY1]', '[OY2]', + '[P]', '[R]', '[S]', '[SH]', '[T]', '[TH]', '[UH]', '[UH0]', '[UH1]', '[UH2]', + '[UW]', '[UW0]', '[UW1]', '[UW2]', '[V]', '[W]', '[Y]', '[Z]', '[ZH]', + '[a]', '[ai]', '[an]', '[ang]', '[ao]', '[b]', '[c]', '[ch]', '[d]', '[e]', + '[ei]', '[en]', '[eng]', '[f]', '[g]', '[h]', '[i]', '[ian]', '[in]', '[ing]', + '[iu]', '[ià]', '[iàn]', '[iàng]', '[iào]', '[iá]', '[ián]', '[iáng]', '[iáo]', + '[iè]', '[ié]', '[iòng]', '[ióng]', '[iù]', '[iú]', '[iā]', '[iān]', '[iāng]', + '[iāo]', '[iē]', '[iě]', '[iōng]', '[iū]', '[iǎ]', '[iǎn]', '[iǎng]', '[iǎo]', + '[iǒng]', '[iǔ]', '[j]', '[k]', '[l]', '[m]', '[n]', '[o]', '[ong]', '[ou]', + '[p]', '[q]', '[r]', '[s]', '[sh]', '[t]', '[u]', '[uang]', '[ue]', '[un]', + '[uo]', '[uà]', '[uài]', '[uàn]', '[uàng]', '[uá]', '[uái]', '[uán]', '[uáng]', + '[uè]', '[ué]', '[uì]', '[uí]', '[uò]', '[uó]', '[uā]', '[uāi]', '[uān]', + '[uāng]', '[uē]', '[uě]', '[uī]', '[uō]', '[uǎ]', '[uǎi]', '[uǎn]', '[uǎng]', + '[uǐ]', '[uǒ]', '[vè]', '[w]', '[x]', '[y]', '[z]', '[zh]', + '[à]', '[ài]', '[àn]', '[àng]', '[ào]', '[á]', '[ái]', '[án]', '[áng]', '[áo]', + '[è]', '[èi]', '[èn]', '[èng]', '[èr]', '[é]', '[éi]', '[én]', '[éng]', '[ér]', + '[ì]', '[ìn]', '[ìng]', '[í]', '[ín]', '[íng]', '[ò]', '[òng]', '[òu]', + '[ó]', '[óng]', '[óu]', '[ù]', '[ùn]', '[ú]', '[ún]', '[ā]', '[āi]', '[ān]', + '[āng]', '[āo]', '[ē]', '[ēi]', '[ēn]', '[ēng]', '[ě]', '[ěi]', '[ěn]', + '[ěng]', '[ěr]', '[ī]', '[īn]', '[īng]', '[ō]', '[ōng]', '[ōu]', '[ū]', + '[ūn]', '[ǎ]', '[ǎi]', '[ǎn]', '[ǎng]', '[ǎo]', '[ǐ]', '[ǐn]', '[ǐng]', + '[ǒ]', '[ǒng]', '[ǒu]', '[ǔ]', '[ǔn]', '[ǘ]', '[ǚ]', '[ǜ]', +] + +# Chat template for CosyVoice3 speculative vllm: <|sos|>{text}<|task_id|>{speech_tokens} +_COSYVOICE3_CHAT_TEMPLATE = ( + "{%- for message in messages %}" + "{%- if message['role'] == 'user' %}{{- '<|sos|>' + message['content'] + '<|task_id|>' }}" + "{%- elif message['role'] == 'assistant' %}{{- message['content']}}" + "{%- endif %}{%- endfor %}" +) + + +def export_cosyvoice3_vllm_spec(model, model_path, hf_llm_dir, device): + """Export CosyVoice3LM to HuggingFace format compatible with DSpark speculative decoding. + + Extends the Qwen2 tokenizer with CosyVoice3 text special tokens and speech tokens, + then splices speech_embedding / llm_decoder into the model's embed_tokens / lm_head + at offset `text_vocab_size`. This produces the same token-ID layout expected by the + DSpark draft model (yuekai/cosyvoice3_llm_dspark): speech token N lives at token ID + `text_vocab_size + N`. + + Mirrors export_cosyvoice2_vllm() in style: skips if model_path already exists. + """ + if os.path.exists(model_path): + return + + from transformers import AutoTokenizer + import math + + dtype = torch.bfloat16 + tokenizer = AutoTokenizer.from_pretrained(hf_llm_dir, trust_remote_code=True) + tokenizer.add_special_tokens({ + 'eos_token': '<|endoftext|>', + 'pad_token': '<|endoftext|>', + 'additional_special_tokens': _COSYVOICE3_TEXT_SPECIAL_TOKENS, + }) + text_vocab_size = len(tokenizer) + + speech_token_size = model.speech_embedding.num_embeddings # includes 200 special + base_speech_token_size = model.speech_token_size # e.g. 2512 or 6561 + + speech_tokens = [f'<|s_{i}|>' for i in range(speech_token_size)] + speech_tokens[base_speech_token_size + 0] = '<|sos|>' + speech_tokens[base_speech_token_size + 1] = '<|eos1|>' + speech_tokens[base_speech_token_size + 2] = '<|task_id|>' + speech_tokens[base_speech_token_size + 3] = '<|fill|>' + tokenizer.add_tokens(speech_tokens) + tokenizer.chat_template = _COSYVOICE3_CHAT_TEMPLATE + + # Pad vocab to multiple of 128 for efficiency + new_vocab = len(tokenizer) + padded_vocab = math.ceil(new_vocab / 128) * 128 + + qwen_model = model.llm.model + qwen_model.resize_token_embeddings(padded_vocab) + qwen_model.to(dtype).to(device) + + with torch.no_grad(): + # Splice speech_embedding into embed_tokens at [text_vocab_size:] + src = min(model.speech_embedding.weight.shape[0], speech_token_size) + qwen_model.get_input_embeddings().weight[text_vocab_size:text_vocab_size + src] = \ + model.speech_embedding.weight[:src].to(dtype) + + # Build new lm_head: text part zeroed (not generated), speech part = llm_decoder + has_bias = model.llm_decoder.bias is not None + new_lm_head = torch.nn.Linear( + qwen_model.config.hidden_size, padded_vocab, bias=has_bias, device=device, dtype=dtype) + new_lm_head.weight.data.zero_() + if has_bias: + new_lm_head.bias.data.fill_(float('-inf')) + # Copy text part from original lm_head so text logits stay valid (optional but clean) + orig_head = qwen_model.lm_head + copy_text = min(orig_head.weight.shape[0], text_vocab_size) + new_lm_head.weight[:copy_text] = orig_head.weight[:copy_text].to(dtype) + # Copy llm_decoder into speech part + dec_size = min(model.llm_decoder.weight.shape[0], speech_token_size) + new_lm_head.weight[text_vocab_size:text_vocab_size + dec_size] = \ + model.llm_decoder.weight[:dec_size].to(dtype) + if has_bias: + new_lm_head.bias[text_vocab_size:text_vocab_size + dec_size] = \ + model.llm_decoder.bias[:dec_size].to(dtype) + qwen_model.lm_head = new_lm_head + + eos_id = text_vocab_size + base_speech_token_size + 1 + qwen_model.config.vocab_size = padded_vocab + qwen_model.config.tie_word_embeddings = False + qwen_model.generation_config.eos_token_id = eos_id + qwen_model.generation_config.pad_token_id = eos_id + + os.makedirs(model_path, exist_ok=True) + qwen_model.save_pretrained(model_path) + tokenizer.save_pretrained(model_path) + json.dump({ + 'text_vocab_size': text_vocab_size, + 'base_speech_token_size': base_speech_token_size, + 'embedding_size': speech_token_size, + 'padded_vocab_size': padded_vocab, + 'eos_token_id': eos_id, + 'speech_token_offset': text_vocab_size, + }, open(os.path.join(model_path, 'cosyvoice3_metadata.json'), 'w'), indent=2) + + def export_cosyvoice2_vllm(model, model_path, device): if os.path.exists(model_path): return diff --git a/vllm_example.py b/vllm_example.py index 697d7a9a3..842444db6 100644 --- a/vllm_example.py +++ b/vllm_example.py @@ -30,9 +30,40 @@ def cosyvoice3_example(): continue +def cosyvoice3_spec_example(): + """CosyVoice3 + DSpark speculative decoding. + + Requires: + 1. A HuggingFace-format export of the CosyVoice3 LLM (produced by + runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py). + 2. A speculators-format DSpark drafter checkpoint (local path or HF repo). + 3. The speculative vllm fork (yuekaizhang/vllm, branch + dspark-draft-sampling-mirrors) loaded via PYTHONPATH with compiled + extensions from a vllm-omni 0.25.1 build. + + Example launch (set PYTHONPATH before starting Python): + PYTHONPATH=/path/to/speculative/vllm:/path/to/vllm025_venv/lib/python3.12/site-packages \\ + python vllm_example.py + """ + DRAFT_MODEL = 'yuekai/cosyvoice3_llm_dspark' # or local path to checkpoint_best/ + + cosyvoice = AutoModel(model_dir='pretrained_models/Fun-CosyVoice3-0.5B', load_trt=True, fp16=False) + # HF export is created automatically under pretrained_models/Fun-CosyVoice3-0.5B/hf_spec + cosyvoice.model.load_vllm_spec('pretrained_models/Fun-CosyVoice3-0.5B', DRAFT_MODEL) + + for i in tqdm(range(10)): + set_all_random_seed(i) + for j, result in enumerate(cosyvoice.inference_zero_shot( + '收到好友从远方寄来的生日礼物，那份意外的惊喜与深深的祝福让我心中充满了甜蜜的快乐，笑容如花儿般绽放。', + 'You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。', + './asset/zero_shot_prompt.wav', stream=False)): + result['tts_speech'].numpy().tofile(f'spec_{i}_{j}.pcm') + + def main(): # cosyvoice2_example() - cosyvoice3_example() + # cosyvoice3_example() + cosyvoice3_spec_example() if __name__ == '__main__': From ee734e0d24f371e107ed37bec112997e9557475d Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Mon, 3 Aug 2026 04:04:29 -0700 Subject: [PATCH 2/2] Simplify cosyvoice3_spec_example docstring; restore cosyvoice3_example as default Co-Authored-By: Claude Sonnet 4.6 (1M context) Signed-off-by: Yuekai Zhang --- vllm_example.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/vllm_example.py b/vllm_example.py index 842444db6..9a7bae537 100644 --- a/vllm_example.py +++ b/vllm_example.py @@ -33,17 +33,13 @@ def cosyvoice3_example(): def cosyvoice3_spec_example(): """CosyVoice3 + DSpark speculative decoding. - Requires: - 1. A HuggingFace-format export of the CosyVoice3 LLM (produced by - runtime/triton_trtllm/scripts/convert_cosyvoice3_to_hf.py). - 2. A speculators-format DSpark drafter checkpoint (local path or HF repo). - 3. The speculative vllm fork (yuekaizhang/vllm, branch - dspark-draft-sampling-mirrors) loaded via PYTHONPATH with compiled - extensions from a vllm-omni 0.25.1 build. + Setup (one-time): + pip install vllm-omni==0.25.1 + git clone -b dspark-draft-sampling-mirrors https://github.com/yuekaizhang/vllm + export PYTHONPATH=/path/to/yuekaizhang/vllm:$PYTHONPATH - Example launch (set PYTHONPATH before starting Python): - PYTHONPATH=/path/to/speculative/vllm:/path/to/vllm025_venv/lib/python3.12/site-packages \\ - python vllm_example.py + Then run: + python vllm_example.py """ DRAFT_MODEL = 'yuekai/cosyvoice3_llm_dspark' # or local path to checkpoint_best/ @@ -62,8 +58,8 @@ def cosyvoice3_spec_example(): def main(): # cosyvoice2_example() - # cosyvoice3_example() - cosyvoice3_spec_example() + cosyvoice3_example() + # cosyvoice3_spec_example() if __name__ == '__main__':