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

[BugFix]Fix keyError: num_processed_tokens_delta (#2213)

Browse Source 
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
This commit is contained in:
amy-why-3459 authored and GitHub committed 2026-03-26 15:16:25 +08:00
1 parent 8305477750
commit 029c85393a
2 files changed
+1 -127
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/model_executor/models/test_qwen3_omni_async_decode.py
-122
View File
@@ -1,122 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import torch
from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
def _make_model() -> Qwen3OmniMoeForConditionalGeneration:
model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
model.talker = SimpleNamespace(
num_code_groups=2,
text_projection=lambda x: x + 10,
)
model.tts_eos_embed = torch.tensor([[100.0, 101.0]], dtype=torch.bfloat16)
model.tts_pad_embed = torch.tensor([[200.0, 201.0]], dtype=torch.bfloat16)
return model
def test_talker_preprocess_decode_starts_from_prefill_consumed_boundary():
model = _make_model()
observed = {}
def fake_decode(input_ids, input_embeds, update_dict, **info_dict):
observed["num_processed_tokens"] = info_dict["num_processed_tokens"]
update_dict["num_processed_tokens_delta"] = 1
return torch.zeros((1, 2), dtype=torch.bfloat16), torch.ones((1, 2), dtype=torch.bfloat16), update_dict
model.talker_preprocess_decode = fake_decode
_, _, update_dict = model.talker_preprocess(
input_ids=torch.tensor([1], dtype=torch.long),
input_embeds=torch.ones((1, 2), dtype=torch.bfloat16),
decode_flag=False,
num_processed_tokens=0,
prefill_consumed_text_tokens=1,
)
assert observed["num_processed_tokens"] == 1
assert update_dict["decode_flag"] is True
assert update_dict["num_processed_tokens"] == 2
def test_async_decode_terminal_steps_do_not_advance_processed_tokens():
model = _make_model()
device = torch.device("cpu")
update_dict = {}
text_step = model._thinker_decode_to_talker_decode(
{
"cached_thinker_decode_embeddings": torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16),
"num_processed_tokens": 1,
"thinker_output_token_ids": [11, 12],
},
device,
update_dict,
)
assert torch.equal(text_step, model.tts_eos_embed)
assert update_dict["finished_flag"] is True
assert update_dict["num_processed_tokens_delta"] == 0
update_dict = {}
text_step = model._thinker_decode_to_talker_decode(
{
"cached_thinker_decode_embeddings": torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16),
"num_processed_tokens": 1,
"thinker_output_token_ids": [11, 12],
"finished_flag": True,
},
device,
update_dict,
)
assert torch.equal(text_step, model.tts_pad_embed)
assert update_dict["num_processed_tokens_delta"] == 0
def test_async_decode_consumes_cached_embedding_and_appends_new_runtime_embed():
model = _make_model()
device = torch.device("cpu")
update_dict = {}
text_step = model._thinker_decode_to_talker_decode(
{
"cached_thinker_decode_embeddings": torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16),
"thinker_decode_embeddings": torch.tensor([[3.0, 4.0]], dtype=torch.bfloat16),
"num_processed_tokens": 0,
"thinker_output_token_ids": [11, 12, 13],
},
device,
update_dict,
)
assert torch.equal(text_step, torch.tensor([11.0, 12.0], dtype=torch.bfloat16))
assert torch.equal(
update_dict["cached_thinker_decode_embeddings"],
torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.bfloat16),
)
assert update_dict["num_processed_tokens_delta"] == 1
def test_async_decode_consumes_runtime_embedding_when_cache_is_empty():
model = _make_model()
device = torch.device("cpu")
update_dict = {}
text_step = model._thinker_decode_to_talker_decode(
{
"thinker_decode_embeddings": torch.tensor([[5.0, 6.0]], dtype=torch.bfloat16),
"num_processed_tokens": 0,
"thinker_output_token_ids": [11, 12, 13],
},
device,
update_dict,
)
assert torch.equal(text_step, torch.tensor([15.0, 16.0], dtype=torch.bfloat16))
assert update_dict["thinker_decode_embeddings"] is None
assert update_dict["num_processed_tokens_delta"] == 1
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
+1 -5
View File
@@ -602,7 +602,6 @@ class Qwen3OmniMoeForConditionalGeneration(
dtype=torch.long,
)
update_dict["code_predictor_codes"] = code_predictor_codes
update_dict["num_processed_tokens_delta"] = span_len
else:
# decode
if not info_dict.get("decode_flag", False):
@@ -620,8 +619,7 @@ class Qwen3OmniMoeForConditionalGeneration(
)
update_dict["mtp_inputs"] = last_talker_hidden, text_step
processed_delta = update_dict.pop("num_processed_tokens_delta")
update_dict["num_processed_tokens"] = info_dict.get("num_processed_tokens", 0) + processed_delta
update_dict["num_processed_tokens"] = info_dict.get("num_processed_tokens", 0) + span_len
return input_ids, input_embeds, update_dict
def talker_mtp(
@@ -918,7 +916,6 @@ class Qwen3OmniMoeForConditionalGeneration(
start_index = info_dict.get("num_processed_tokens", 0)
thinker_output_token_ids = info_dict.get("thinker_output_token_ids", [])
if start_index >= len(thinker_output_token_ids) - 1:
update_dict["num_processed_tokens_delta"] = 0
if info_dict.get("finished_flag"):
return self.tts_pad_embed.to(device)
update_dict["finished_flag"] = True
@@ -937,7 +934,6 @@ class Qwen3OmniMoeForConditionalGeneration(
thinker_embed = thinker_embed.to(device)
update_dict["thinker_decode_embeddings"] = None
update_dict["num_processed_tokens_delta"] = 1
return self.talker.text_projection(thinker_embed).to(device)
def talker_preprocess_decode(