100 lines
5.1 KiB
Python
100 lines
5.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Observe the real runner/model boundary without replacing inference."""
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner as LegacyRunner
|
|
from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner
|
|
|
|
|
|
class PhaseContractWorkerExtension:
|
|
model_runner: "LegacyRunner | OmniARModelRunner"
|
|
|
|
def start_phase_contract_probe(self):
|
|
runner = self.model_runner
|
|
self.phase_contract_events: list[tuple[int, int, int, bool]] = []
|
|
self.phase_contract_codes: list[list[list[int]]] = []
|
|
tail_embeddings = {}
|
|
original_preprocess = runner.model.preprocess
|
|
model_state = getattr(runner, "model_state", None)
|
|
is_v2 = model_state is not None
|
|
original_mtp = model_state._run_batched_mtp if model_state is not None else runner._talker_mtp_forward
|
|
|
|
def preprocess(input_ids, input_embeds, **info):
|
|
rid = info["req_id"] if is_v2 else info["request_id"]
|
|
if is_v2:
|
|
prompt_len = info["_omni_prompt_len"]
|
|
computed = info["_omni_num_computed_tokens"]
|
|
else:
|
|
index = runner.input_batch.req_id_to_index[rid]
|
|
prompt_len = len(runner.requests[rid].prompt_token_ids)
|
|
computed = int(runner.input_batch.num_computed_tokens_cpu[index])
|
|
assert info["_omni_prompt_len"] == prompt_len
|
|
assert info["_omni_num_computed_tokens"] == computed
|
|
assert info["_omni_is_prefill"] is (computed < prompt_len)
|
|
result = original_preprocess(input_ids=input_ids, input_embeds=input_embeds, **info)
|
|
span = input_ids.numel()
|
|
self.phase_contract_events.append((prompt_len, computed, span, info["_omni_is_prefill"]))
|
|
if computed == prompt_len - 1 and span == 1:
|
|
if is_v2:
|
|
tail_embeddings[rid] = result[1].detach().clone()
|
|
else:
|
|
offset = int(runner.query_start_loc.cpu[index])
|
|
tail_embeddings[offset] = result[1].detach().clone()
|
|
return result
|
|
|
|
def legacy_mtp(decode_req_ids, inputs_embeds, start_offsets=None):
|
|
for rid in decode_req_ids:
|
|
index = runner.input_batch.req_id_to_index[rid]
|
|
assert int(runner.input_batch.num_computed_tokens_cpu[index]) >= len(
|
|
runner.requests[rid].prompt_token_ids
|
|
), "prefill row was routed through MTP"
|
|
original_mtp(decode_req_ids, inputs_embeds, start_offsets)
|
|
for offset, expected in tail_embeddings.items():
|
|
torch.testing.assert_close(inputs_embeds[offset : offset + 1], expected, rtol=0, atol=0)
|
|
tail_embeddings.clear()
|
|
for rid in decode_req_ids:
|
|
self.phase_contract_codes.append(runner.model_intermediate_buffer[rid]["codes"]["audio"].tolist())
|
|
|
|
def v2_mtp(mtp_batches, input_ids, embeds, input_batch, gpu_keys, *args, **kwargs):
|
|
assert model_state is not None
|
|
for batch_index, _offset, _mtp_inputs in mtp_batches:
|
|
req_index = int(input_batch.idx_mapping_np[batch_index])
|
|
computed = model_state._get_input_batch_num_computed(input_batch, req_index, batch_index)
|
|
prompt_len = model_state._get_req_state_value(runner.req_states.prompt_len, req_index)
|
|
assert computed is not None and prompt_len is not None and computed >= prompt_len, (
|
|
"prefill row was routed through MTP"
|
|
)
|
|
original_mtp(mtp_batches, input_ids, embeds, input_batch, gpu_keys, *args, **kwargs)
|
|
for batch_index, _offset, _mtp_inputs in mtp_batches:
|
|
req_index = int(input_batch.idx_mapping_np[batch_index])
|
|
codes = model_state.intermediate_buffer.buffers[req_index]["codes"]["audio"]
|
|
self.phase_contract_codes.append(codes.tolist())
|
|
|
|
runner.model.preprocess = preprocess
|
|
if model_state is not None:
|
|
original_run_preprocess = model_state.run_preprocess
|
|
|
|
def run_preprocess(input_batch, model_inputs, *args, **kwargs):
|
|
original_run_preprocess(input_batch, model_inputs, *args, **kwargs)
|
|
embeds = model_inputs["inputs_embeds"]
|
|
for rid, expected in tail_embeddings.items():
|
|
batch_index = input_batch.req_ids.index(rid)
|
|
offset = int(input_batch.query_start_loc_np[batch_index])
|
|
torch.testing.assert_close(embeds[offset : offset + 1], expected, rtol=0, atol=0)
|
|
tail_embeddings.clear()
|
|
|
|
model_state.run_preprocess = run_preprocess
|
|
model_state._run_batched_mtp = v2_mtp
|
|
else:
|
|
runner._talker_mtp_forward = legacy_mtp
|
|
return True
|
|
|
|
def get_phase_contract_probe(self):
|
|
return {"events": self.phase_contract_events, "codes": self.phase_contract_codes}
|