anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[BugFix] Fix the issue of thinker requests being preempted, causing shape mismatch. (#3147)

Browse Source 
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
This commit is contained in:
amy-why-3459 authored and GitHub committed 2026-05-15 17:07:42 +08:00
1 parent 440c718d2a
commit e7ee5de09f
5 files changed
+65 -74
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
+11 -10
View File
@@ -990,19 +990,10 @@ class Qwen3OmniMoeForConditionalGeneration(
"""
embed = payload.get("embed", {})
meta = payload.get("meta", {})
ids = payload.get("ids", {})
cached_thinker_decode_embeds = embed.get("cached_decode", None)
thinker_decode_embed = embed.get("decode", None)
start_index = meta.get("num_processed_tokens", 0)
thinker_output_token_ids = ids.get("output", [])
if start_index >= len(thinker_output_token_ids) - 1:
# When the tokens output by the thinker are exhausted, an EOS token needs to be appended.
# Use the finished_flag to mark that all tokens output by thinker have been consumed.
if meta.get("eos_emitted", False):
return self.tts_pad_embed.to(device)
update_dict.setdefault("meta", {})["eos_emitted"] = True
return self.tts_eos_embed.to(device)
if cached_thinker_decode_embeds is not None and start_index < cached_thinker_decode_embeds.shape[0]:
cached_thinker_decode_embeds = cached_thinker_decode_embeds.to(device)
@@ -1011,10 +1002,20 @@ class Qwen3OmniMoeForConditionalGeneration(
thinker_decode_embed = thinker_decode_embed.to(device)
cached_thinker_decode_embeds = torch.cat([cached_thinker_decode_embeds, thinker_decode_embed], dim=0)
update_dict.setdefault("embed", {})["cached_decode"] = cached_thinker_decode_embeds
else:
elif thinker_decode_embed is not None:
thinker_embed = thinker_decode_embed
if thinker_embed.device != device:
thinker_embed = thinker_embed.to(device)
else:
# When the tokens output by the thinker are exhausted, an EOS token needs to be appended.
# Use the finished_flag to mark that all tokens output by thinker have been consumed.
if meta.get("eos_emitted", False):
return self.tts_pad_embed.to(device)
update_dict.setdefault("meta", {})["eos_emitted"] = True
return self.tts_eos_embed.to(device)
update_dict.setdefault("embed", {})["decode"] = None
return self.talker.text_projection(thinker_embed).to(device)