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

[Rebase] Rebase to vllm v0.21.0 (#3530)

Browse Source 
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Signed-off-by: dsinghvi <divyanshsinghvi@gmail.com>
Signed-off-by: Dnoob <dxpouo@gmail.com>
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Signed-off-by: rein yang <ruiruyang2@gmail.com>
Signed-off-by: Nick Cao <ncao@redhat.com>
Signed-off-by: zhumingjue <zhumingjue@huawei.com>
Signed-off-by: Ricardo Noriega De Soto <rnoriega@redhat.com>
Signed-off-by: lyj-jjj <liuyingjun5@huawei.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: weizhoublue <weizhoublue@github.com>
Signed-off-by: weizhou.lan@daocloud.io <weizhou.lan@daocloud.io>
Signed-off-by: dengyunyang <584797741@qq.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Signed-off-by: David Chen <530634352@qq.com>
Signed-off-by: Jie Liu <33612777+keeper-jie@users.noreply.github.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Signed-off-by: bjf-frz <frz123db@gmail.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Signed-off-by: KexiongYu <yukexiong1@huawei.com>
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: NumberWan <wantszkin2003@gmail.com>
Co-authored-by: dsinghvi <divyanshsinghvi@gmail.com>
Co-authored-by: Dnoob <dxpouo@gmail.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: rein yang <73573651+R2-Y@users.noreply.github.com>
Co-authored-by: Nick Cao <ncao@redhat.com>
Co-authored-by: zhumingjue138 <zhumingjue@huawei.com>
Co-authored-by: Ricardo Noriega <rnoriega@redhat.com>
Co-authored-by: lyj-jjj <liuyingjun5@huawei.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com>
Co-authored-by: weizhoublue <45163302+weizhoublue@users.noreply.github.com>
Co-authored-by: weizhoublue <weizhoublue@github.com>
Co-authored-by: dengyunyang <584797741@qq.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Co-authored-by: Jie Liu <33612777+keeper-jie@users.noreply.github.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: NATURE <wzliu@connect.hku.hk>
Co-authored-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Y. Fisher <yukexiong1@huawei.com>
Co-authored-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
+20
Zhou TaichangClaude Opus 4.7NumberWandsinghviDnoobYueqian LinHongsheng Liuknlnguyen1802Samitrein yangNick Caozhumingjue138Ricardo Noriegalyj-jjjCursorgcanlinwangyuweizhoublueweizhoubluedengyunyang汪志鹏WeiQing ChenJie LiuYueqian LinNATUREbjf-frzamy-why-3459Y. FisherZeyu Huang | 黃澤宇Copilot Autofix powered by AI
 authored and GitHub committed 2026-05-16 18:18:43 +08:00
1 parent 5e1986206f
commit eb4e60ee64
46 files changed
+589 -258
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/worker/gpu_model_runner.py
+59 -10
View File
@@ -1,4 +1,5 @@
import contextlib
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, cast
import numpy as np
@@ -18,6 +19,11 @@ from vllm.v1.spec_decode.dflash import DFlashProposer
from vllm.v1.spec_decode.draft_model import DraftModelProposer
from vllm.v1.spec_decode.eagle import EagleProposer
from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
from vllm.v1.spec_decode.gemma4 import Gemma4Proposer
from vllm.v1.spec_decode.ngram_proposer_gpu import (
update_ngram_gpu_tensors_incremental,
update_scheduler_for_invalid_drafts,
)
from vllm.v1.worker.gpu_input_batch import CachedRequestState
from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata
from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
@@ -250,7 +256,7 @@ class OmniGPUModelRunner(GPUModelRunner):
mrope_pos_ptr += completion_part_len
def _update_states(self, scheduler_output: "SchedulerOutput"):
def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None:
"""Update the cached states and the persistent batch with the scheduler
output.
@@ -284,8 +290,7 @@ class OmniGPUModelRunner(GPUModelRunner):
if cleanup_finished_request is not None:
cleanup_finished_request(req_id)
if hasattr(self, "late_interaction_runner"):
self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids)
self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids)
# Remove the finished requests from the persistent batch.
# NOTE(woosuk): There could be an edge case where finished_req_ids and
# scheduled_req_ids overlap. This happens when a request is aborted and
@@ -326,8 +331,9 @@ class OmniGPUModelRunner(GPUModelRunner):
for req_id in unscheduled_req_ids:
self.input_batch.remove_request(req_id)
if self.use_async_spec_decode:
self.prev_num_draft_tokens.np.fill(0)
is_ngram_gpu = self.speculative_config is not None and self.speculative_config.use_ngram_gpu()
if is_ngram_gpu:
ngram_gpu_new_reqs: list[CachedRequestState] = []
reqs_to_add: list[CachedRequestState] = []
deferred_spec_decode_corrections = []
@@ -379,8 +385,7 @@ class OmniGPUModelRunner(GPUModelRunner):
lora_request=new_req_data.lora_request,
)
self.requests[req_id] = req_state
if hasattr(self, "late_interaction_runner"):
self.late_interaction_runner.register_request(req_id, pooling_params)
self.late_interaction_runner.register_request(req_id, pooling_params)
# If prompt embeddings are provided, decode and attach to inter_data
try:
@@ -429,12 +434,30 @@ class OmniGPUModelRunner(GPUModelRunner):
self._init_xdrope_positions(req_state)
reqs_to_add.append(self.requests[req_id])
# Track new requests for ngram_gpu full tensor copy
if is_ngram_gpu:
ngram_gpu_new_reqs.append(req_state)
# Update the states of the running/resumed requests.
is_last_rank = get_pp_group().is_last_rank
req_data = scheduler_output.scheduled_cached_reqs
scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
# Save scheduler-allocated spec lengths before trimming so
# prev_num_draft_len keeps the optimistic count for rejection correction.
original_num_spec_per_req: dict[str, int] = {}
if self.speculative_config is not None and self.speculative_config.use_ngram_gpu():
for req_id, toks in scheduled_spec_tokens.items():
original_num_spec_per_req[req_id] = len(toks)
update_scheduler_for_invalid_drafts(
self._num_valid_draft_tokens_event,
self._num_valid_draft_tokens_cpu,
scheduler_output,
self.input_batch.req_id_to_index,
)
if self.use_async_spec_decode:
self.prev_num_draft_tokens.np.fill(0)
for i, req_id in enumerate(req_data.req_ids):
req_state = self.requests[req_id]
num_computed_tokens = req_data.num_computed_tokens[i]
@@ -473,6 +496,9 @@ class OmniGPUModelRunner(GPUModelRunner):
if prev_req_index is not None:
self.prev_num_draft_tokens.np[prev_req_index] = optimistic_num_accepted
if is_ngram_gpu and optimistic_num_accepted > 0:
self.input_batch.num_tokens_no_spec[req_index] += optimistic_num_accepted
# Update the cached states.
req_state.num_computed_tokens = num_computed_tokens
@@ -521,6 +547,9 @@ class OmniGPUModelRunner(GPUModelRunner):
req_state.output_token_ids = resumed_token_ids[-num_output_tokens:]
reqs_to_add.append(req_state)
# Track resumed requests for ngram_gpu full tensor copy
if is_ngram_gpu:
ngram_gpu_new_reqs.append(req_state)
continue
# Update the persistent batch.
@@ -539,6 +568,11 @@ class OmniGPUModelRunner(GPUModelRunner):
# Add spec_token_ids to token_ids_cpu.
self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens)
# Restore scheduler-side draft count after ngram trimming.
if original_num_spec_per_req:
orig = original_num_spec_per_req.get(req_id, 0)
if orig != req_state.prev_num_draft_len:
req_state.prev_num_draft_len = orig
# Add the new or resumed requests to the persistent batch.
# The smaller empty indices are filled first.
@@ -553,6 +587,18 @@ class OmniGPUModelRunner(GPUModelRunner):
# Refresh batch metadata with any pending updates.
self.input_batch.refresh_metadata()
# Incrementally update ngram_gpu tensors after batch is stable
if is_ngram_gpu:
update_ngram_gpu_tensors_incremental(
self.input_batch,
self.token_ids_gpu_tensor,
self.num_tokens_no_spec_gpu,
ngram_gpu_new_reqs,
self.device,
_pinned_idx_buf=self._ngram_pinned_idx_buf,
_pinned_val_buf=self._ngram_pinned_val_buf,
)
if deferred_spec_decode_corrections:
def correct_spec_decode_token_counts():
@@ -577,6 +623,9 @@ class OmniGPUModelRunner(GPUModelRunner):
if cur_req_index is None:
continue
self.input_batch.num_computed_tokens_cpu[cur_req_index] -= correction
if is_ngram_gpu and correction > 0:
self.input_batch.num_tokens_no_spec[cur_req_index] -= correction
self.num_tokens_no_spec_gpu[cur_req_index] -= correction
return correct_spec_decode_token_counts
else:
@@ -846,7 +895,7 @@ class OmniGPUModelRunner(GPUModelRunner):
device=self.device,
)
intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens_padded, None, False)
intermediate_tensors = self.sync_and_gather_intermediate_tensors(num_tokens_padded, None, False)
if ubatch_slices_padded is not None:
# Adjust values to reflect a single ubatch.
@@ -900,7 +949,7 @@ class OmniGPUModelRunner(GPUModelRunner):
):
assert isinstance(
self.drafter,
EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer,
EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer | Gemma4Proposer,
)
assert self.speculative_config is not None
# Eagle currently only supports PIECEWISE cudagraphs.
@@ -1300,7 +1349,7 @@ class OmniGPUModelRunner(GPUModelRunner):
intermediate_tensors = None
else:
assert intermediate_tensors is not None
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
intermediate_tensors = self.sync_and_gather_intermediate_tensors(
num_input_tokens, intermediate_tensors, True
)