801 lines
35 KiB
Python
801 lines
35 KiB
Python
from typing import TYPE_CHECKING, Any, Optional, Union, cast
|
|
|
|
import numpy as np
|
|
import torch
|
|
import vllm.envs as envs
|
|
from vllm.distributed.kv_transfer import has_kv_transfer_group
|
|
from vllm.distributed.parallel_state import get_pp_group, get_tp_group
|
|
from vllm.forward_context import set_forward_context
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.models.interfaces_base import VllmModelForPooling
|
|
from vllm.multimodal.inputs import MultiModalKwargs
|
|
from vllm.sampling_params import SamplingType
|
|
from vllm.sequence import IntermediateTensors
|
|
from vllm.utils import LazyLoader, round_up
|
|
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
|
|
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
|
|
from vllm.v1.spec_decode.eagle import EagleProposer
|
|
from vllm.v1.worker.gpu_input_batch import CachedRequestState
|
|
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
|
|
|
|
from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
|
|
from vllm_omni.outputs import OmniModelRunnerOutput
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.v1.core.sched.output import SchedulerOutput
|
|
else:
|
|
xgr = LazyLoader("xgr", globals(), "xgrammar")
|
|
xgr_torch_compile = LazyLoader(
|
|
"xgr_torch_compile",
|
|
globals(),
|
|
"xgrammar.kernels.apply_token_bitmask_inplace_torch_compile",
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class OmniGPUModelRunner(GPUModelRunner):
|
|
def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
|
|
"""Update the cached states and the persistent batch with the scheduler
|
|
output.
|
|
|
|
The updated states are used by the `_prepare_inputs` function to create
|
|
the input GPU tensors for the model.
|
|
|
|
The SamplingMetadata is updated and copied to the GPU if there is a
|
|
new/resumed/paused/finished request in the batch.
|
|
"""
|
|
# Remove finished requests from the cached states.
|
|
for req_id in scheduler_output.finished_req_ids:
|
|
self.requests.pop(req_id, None)
|
|
self.encoder_cache.pop(req_id, None)
|
|
# 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
|
|
# then resubmitted with the same ID. In this case, we treat them as two
|
|
# distinct requests - clearing the cached states for the first request
|
|
# and handling the second as a new request.
|
|
for req_id in scheduler_output.finished_req_ids:
|
|
self.input_batch.remove_request(req_id)
|
|
|
|
# Free the cached encoder outputs.
|
|
for req_id, input_id in scheduler_output.free_encoder_input_ids:
|
|
encoder_outputs = self.encoder_cache.get(req_id)
|
|
if encoder_outputs is not None:
|
|
encoder_outputs.pop(input_id, None)
|
|
if not encoder_outputs:
|
|
self.encoder_cache.pop(req_id, None)
|
|
|
|
# Remove the unscheduled requests from the persistent batch.
|
|
# NOTE(woosuk): The unscheduled requests are either preempted requests
|
|
# or running requests that are not scheduled in this step. We remove
|
|
# them from the persistent batch but keep their cached states since
|
|
# they will be scheduled again sometime in the future.
|
|
scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
|
|
cached_req_ids = self.input_batch.req_id_to_index.keys()
|
|
unscheduled_req_ids = cached_req_ids - scheduled_req_ids
|
|
# NOTE(woosuk): The persistent batch optimization assumes that
|
|
# consecutive batches contain mostly the same requests. If batches
|
|
# have low request overlap (e.g., alternating between two distinct
|
|
# sets of requests), this optimization becomes very inefficient.
|
|
for req_id in unscheduled_req_ids:
|
|
self.input_batch.remove_request(req_id)
|
|
|
|
req_ids_to_add: list[str] = []
|
|
# Add new requests to the cached states.
|
|
for new_req_data in scheduler_output.scheduled_new_reqs:
|
|
req_id = new_req_data.req_id
|
|
sampling_params = new_req_data.sampling_params
|
|
pooling_params = new_req_data.pooling_params
|
|
|
|
if (
|
|
sampling_params
|
|
and sampling_params.sampling_type == SamplingType.RANDOM_SEED
|
|
):
|
|
generator = torch.Generator(device=self.device)
|
|
generator.manual_seed(sampling_params.seed)
|
|
else:
|
|
generator = None
|
|
|
|
if pooling_params:
|
|
assert (
|
|
task := pooling_params.task
|
|
) is not None, "You did not set `task` in the API"
|
|
|
|
model = cast(VllmModelForPooling, self.model)
|
|
to_update = model.pooler.get_pooling_updates(task)
|
|
to_update.apply(pooling_params)
|
|
|
|
self.requests[req_id] = CachedRequestState(
|
|
req_id=req_id,
|
|
prompt_token_ids=new_req_data.prompt_token_ids,
|
|
mm_inputs=new_req_data.mm_inputs,
|
|
mm_positions=new_req_data.mm_positions,
|
|
sampling_params=sampling_params,
|
|
pooling_params=pooling_params,
|
|
generator=generator,
|
|
block_ids=new_req_data.block_ids,
|
|
num_computed_tokens=new_req_data.num_computed_tokens,
|
|
output_token_ids=[],
|
|
lora_request=new_req_data.lora_request,
|
|
)
|
|
|
|
# If prompt embeddings are provided, decode and attach to inter_data
|
|
try:
|
|
if getattr(new_req_data, "prompt_embeds", None) is not None:
|
|
payload = new_req_data.prompt_embeds
|
|
dtype = getattr(np, payload.dtype)
|
|
arr = np.frombuffer(payload.data, dtype=dtype)
|
|
arr = arr.reshape(payload.shape)
|
|
pe_cpu = torch.from_numpy(arr)
|
|
# Store temporarily on CPU; later moved to device in builder
|
|
setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
|
|
# Also replace payload with Tensor for user visibility in
|
|
# scheduler_output
|
|
try:
|
|
new_req_data.prompt_embeds = pe_cpu # type: ignore[assignment]
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
# Decode additional_information payloads (dictionary)
|
|
try:
|
|
if getattr(new_req_data, "additional_information", None) is not None:
|
|
payload_info = new_req_data.additional_information
|
|
info_dict = {}
|
|
if isinstance(payload_info, dict):
|
|
info_dict = payload_info
|
|
else:
|
|
from vllm.v1.engine import AdditionalInformationPayload
|
|
|
|
if isinstance(payload_info, AdditionalInformationPayload):
|
|
for k, entry in payload_info.entries.items():
|
|
if entry.tensor_data is not None:
|
|
dt = np.dtype(
|
|
getattr(entry, "tensor_dtype", "float32")
|
|
)
|
|
arr = np.frombuffer(entry.tensor_data, dtype=dt)
|
|
arr = arr.reshape(entry.tensor_shape)
|
|
info_dict[k] = torch.from_numpy(arr)
|
|
else:
|
|
info_dict[k] = entry.list_data
|
|
if info_dict:
|
|
setattr(
|
|
self.requests[req_id],
|
|
"additional_information_cpu",
|
|
info_dict,
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
# Only relevant for models using M-RoPE (e.g, Qwen2-VL)
|
|
if self.uses_mrope:
|
|
image_grid_thw = []
|
|
video_grid_thw = []
|
|
second_per_grid_ts = []
|
|
audio_feature_lengths = []
|
|
use_audio_in_video = False
|
|
for mm_input in self.requests[req_id].mm_inputs:
|
|
if mm_input.get("image_grid_thw") is not None:
|
|
image_grid_thw.extend(mm_input["image_grid_thw"].tolist())
|
|
if mm_input.get("video_grid_thw") is not None:
|
|
video_grid_thw.extend(mm_input["video_grid_thw"].tolist())
|
|
if mm_input.get("second_per_grid_ts") is not None:
|
|
second_per_grid_ts.extend(mm_input["second_per_grid_ts"])
|
|
if mm_input.get("audio_feature_lengths") is not None:
|
|
audio_feature_lengths.extend(mm_input["audio_feature_lengths"])
|
|
if mm_input.get("use_audio_in_video") is True:
|
|
use_audio_in_video = True
|
|
|
|
hf_config = self.model_config.hf_config
|
|
|
|
(
|
|
self.requests[req_id].mrope_positions,
|
|
self.requests[req_id].mrope_position_delta,
|
|
) = MRotaryEmbedding.get_input_positions_tensor(
|
|
self.requests[req_id].prompt_token_ids,
|
|
hf_config=hf_config,
|
|
image_grid_thw=image_grid_thw,
|
|
video_grid_thw=video_grid_thw,
|
|
second_per_grid_ts=second_per_grid_ts,
|
|
audio_feature_lengths=audio_feature_lengths,
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
|
|
req_ids_to_add.append(req_id)
|
|
|
|
# Update the states of the running/resumed requests.
|
|
is_last_rank = get_pp_group().is_last_rank
|
|
req_data = scheduler_output.scheduled_cached_reqs
|
|
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]
|
|
new_block_ids = req_data.new_block_ids[i]
|
|
resumed_from_preemption = req_data.resumed_from_preemption[i]
|
|
|
|
# Update the cached states.
|
|
req_state.num_computed_tokens = num_computed_tokens
|
|
|
|
if not is_last_rank:
|
|
# When using PP, the scheduler sends the sampled tokens back,
|
|
# because there's no direct communication between the first-
|
|
# stage worker and the last-stage worker.
|
|
new_token_ids = req_data.new_token_ids[i]
|
|
# Add the sampled token(s) from the previous step (if any).
|
|
# This doesn't include "unverified" tokens like spec tokens.
|
|
num_new_tokens = (
|
|
num_computed_tokens + len(new_token_ids) - req_state.num_tokens
|
|
)
|
|
if num_new_tokens == 1:
|
|
# Avoid slicing list in most common case.
|
|
req_state.output_token_ids.append(new_token_ids[-1])
|
|
elif num_new_tokens > 0:
|
|
req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
|
|
|
|
# Update the block IDs.
|
|
if not resumed_from_preemption:
|
|
# Append the new blocks to the existing block IDs.
|
|
for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
|
|
block_ids.extend(new_ids)
|
|
else:
|
|
# The request is resumed from preemption.
|
|
# Replace the existing block IDs with the new ones.
|
|
req_state.block_ids = new_block_ids
|
|
|
|
req_index = self.input_batch.req_id_to_index.get(req_id)
|
|
if req_index is None:
|
|
# The request is not in the persistent batch.
|
|
# The request was either preempted and resumed later, or was not
|
|
# scheduled in the previous step and needs to be added again.
|
|
req_ids_to_add.append(req_id)
|
|
continue
|
|
|
|
# Update the persistent batch.
|
|
self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens
|
|
self.input_batch.block_table.append_row(new_block_ids, req_index)
|
|
|
|
# For the last rank, we don't need to update the token_ids_cpu
|
|
# because the sampled tokens are already cached.
|
|
if not is_last_rank:
|
|
# Add new_token_ids to token_ids_cpu.
|
|
start_token_index = num_computed_tokens
|
|
end_token_index = num_computed_tokens + len(new_token_ids)
|
|
self.input_batch.token_ids_cpu[
|
|
req_index, start_token_index:end_token_index
|
|
] = new_token_ids
|
|
self.input_batch.num_tokens_no_spec[req_index] = end_token_index
|
|
self.input_batch.num_tokens[req_index] = end_token_index
|
|
|
|
# Add spec_token_ids to token_ids_cpu.
|
|
spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(
|
|
req_id, ()
|
|
)
|
|
if spec_token_ids:
|
|
num_spec_tokens = len(spec_token_ids)
|
|
start_index = self.input_batch.num_tokens_no_spec[req_index]
|
|
end_token_index = start_index + num_spec_tokens
|
|
self.input_batch.token_ids_cpu[
|
|
req_index, start_index:end_token_index
|
|
] = spec_token_ids
|
|
# NOTE(woosuk): `num_tokens` here may include spec tokens.
|
|
self.input_batch.num_tokens[req_index] += num_spec_tokens
|
|
|
|
# Add the new or resumed requests to the persistent batch.
|
|
# The smaller empty indices are filled first.
|
|
for req_id in req_ids_to_add:
|
|
req_state = self.requests[req_id]
|
|
self.input_batch.add_request(req_state)
|
|
|
|
# Condense the batched states if there are gaps left by removed requests
|
|
self.input_batch.condense()
|
|
# Allow attention backend to reorder the batch, potentially
|
|
self._may_reorder_batch(scheduler_output)
|
|
# Refresh batch metadata with any pending updates.
|
|
self.input_batch.refresh_metadata()
|
|
|
|
@torch.inference_mode()
|
|
def execute_model(
|
|
self,
|
|
scheduler_output: "SchedulerOutput",
|
|
intermediate_tensors: Optional[IntermediateTensors] = None,
|
|
) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
|
|
self._update_states(scheduler_output)
|
|
if not scheduler_output.total_num_scheduled_tokens:
|
|
if not has_kv_transfer_group():
|
|
# Return empty ModelRunnerOutput if there's no work to do.
|
|
return EMPTY_MODEL_RUNNER_OUTPUT
|
|
|
|
return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
|
|
|
|
# Prepare the decoder inputs.
|
|
(
|
|
attn_metadata,
|
|
attention_cuda_graphs,
|
|
logits_indices,
|
|
spec_decode_metadata,
|
|
num_scheduled_tokens_np,
|
|
spec_decode_common_attn_metadata,
|
|
) = self._prepare_inputs(scheduler_output)
|
|
|
|
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
|
|
if (
|
|
self.use_cuda_graph
|
|
and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]
|
|
):
|
|
# Use piecewise CUDA graphs.
|
|
# Add padding to the batch size.
|
|
num_input_tokens = self.vllm_config.pad_for_cudagraph(num_scheduled_tokens)
|
|
else:
|
|
# Eager mode.
|
|
# Pad tokens to multiple of tensor_parallel_size when
|
|
# enabled collective fusion for SP
|
|
tp_size = self.vllm_config.parallel_config.tensor_parallel_size
|
|
if (
|
|
self.compilation_config.pass_config.enable_sequence_parallelism
|
|
and tp_size > 1
|
|
):
|
|
num_input_tokens = round_up(num_scheduled_tokens, tp_size)
|
|
else:
|
|
num_input_tokens = num_scheduled_tokens
|
|
|
|
# Padding for DP
|
|
num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
|
|
num_input_tokens += num_pad
|
|
|
|
# _prepare_inputs may reorder the batch, so we must gather multi
|
|
# modal outputs after that to ensure the correct order
|
|
if self.is_multimodal_model:
|
|
# Run the multimodal encoder if any.
|
|
self._execute_mm_encoder(scheduler_output)
|
|
mm_embeds = self._gather_mm_embeddings(scheduler_output)
|
|
else:
|
|
mm_embeds = []
|
|
|
|
if get_pp_group().is_first_rank:
|
|
# Always prepare inputs_embeds on the first PP rank so both
|
|
# multimodal and text-only models can consume embeddings.
|
|
inputs_embeds_scheduled = self.model.get_input_embeddings(
|
|
input_ids=self.input_ids[:num_scheduled_tokens],
|
|
multimodal_embeddings=(
|
|
(mm_embeds or None) if self.is_multimodal_model else None
|
|
),
|
|
)
|
|
|
|
# TODO(woosuk): Avoid the copy. Optimize.
|
|
self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
|
|
|
|
# Overlay per-request custom prompt_embeds for the prompt portion
|
|
for req_index, req_id in enumerate(self.input_batch.req_ids):
|
|
req_state = self.requests[req_id]
|
|
pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
|
|
if pe_cpu is None:
|
|
continue
|
|
num_computed_tokens = int(
|
|
self.input_batch.num_computed_tokens_cpu[req_index]
|
|
)
|
|
prompt_len = len(req_state.prompt_token_ids)
|
|
prompt_remaining = max(0, prompt_len - num_computed_tokens)
|
|
sched_tokens = int(num_scheduled_tokens_np[req_index])
|
|
overlay_len = min(sched_tokens, prompt_remaining)
|
|
if overlay_len <= 0:
|
|
continue
|
|
src = pe_cpu[
|
|
num_computed_tokens : num_computed_tokens + overlay_len
|
|
].to(dtype=self.dtype, device=self.device, non_blocking=True)
|
|
start_offset = int(self.query_start_loc_cpu[req_index])
|
|
self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
|
|
|
|
input_ids = self.input_ids[:num_input_tokens] # preserved for APIs
|
|
inputs_embeds = self.inputs_embeds[:num_input_tokens]
|
|
model_mm_kwargs = (
|
|
self._extract_mm_kwargs(scheduler_output)
|
|
if self.is_multimodal_model
|
|
else {}
|
|
)
|
|
else:
|
|
# For non-first PP ranks, use token ids as usual; embeddings are
|
|
# only consumed on the first rank.
|
|
input_ids = self.input_ids[:num_input_tokens]
|
|
inputs_embeds = None
|
|
model_mm_kwargs = {}
|
|
if self.uses_mrope:
|
|
positions = self.mrope_positions[:, :num_input_tokens]
|
|
else:
|
|
positions = self.positions[:num_input_tokens]
|
|
|
|
if get_pp_group().is_first_rank:
|
|
intermediate_tensors = None
|
|
else:
|
|
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
|
|
num_input_tokens, intermediate_tensors, True
|
|
)
|
|
|
|
# Some attention backends only support CUDA Graphs in pure decode.
|
|
# If attention doesn't support CUDA Graphs for this batch, but we
|
|
# compiled with full CUDA graphs, we have to skip them entirely.
|
|
skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
|
|
|
|
# Run the model.
|
|
# Use persistent buffers for CUDA graphs.
|
|
with set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=num_input_tokens,
|
|
num_tokens_across_dp=num_tokens_across_dp,
|
|
skip_cuda_graphs=skip_cuda_graphs,
|
|
), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output:
|
|
|
|
model_output = self.model(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
**MultiModalKwargs.as_kwargs(
|
|
model_mm_kwargs,
|
|
device=self.device,
|
|
),
|
|
sampling_metadata=self.input_batch.sampling_metadata,
|
|
logits_index=logits_indices,
|
|
sampler=self.sampler,
|
|
)
|
|
|
|
if self.use_aux_hidden_state_outputs:
|
|
hidden_states, aux_hidden_states = model_output
|
|
else:
|
|
hidden_states = model_output
|
|
aux_hidden_states = None
|
|
|
|
text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
|
|
hidden_states
|
|
)
|
|
|
|
# Broadcast PP output for external_launcher (torchrun)
|
|
# to make sure we are synced across pp ranks
|
|
# TODO: Support overlapping mirco-batches
|
|
# https://github.com/vllm-project/vllm/issues/18019
|
|
broadcast_pp_output = (
|
|
self.parallel_config.distributed_executor_backend == "external_launcher"
|
|
and len(get_pp_group().ranks) > 0
|
|
)
|
|
if not get_pp_group().is_last_rank:
|
|
# For mid-pipeline stages, return the hidden states.
|
|
assert isinstance(text_hidden_states, IntermediateTensors)
|
|
if not broadcast_pp_output:
|
|
text_hidden_states.kv_connector_output = kv_connector_output
|
|
return text_hidden_states
|
|
get_pp_group().send_tensor_dict(
|
|
text_hidden_states.tensors, all_gather_group=get_tp_group()
|
|
)
|
|
logits = None
|
|
else:
|
|
if self.input_batch.pooling_params:
|
|
return self._pool(
|
|
text_hidden_states,
|
|
num_scheduled_tokens,
|
|
num_scheduled_tokens_np,
|
|
kv_connector_output,
|
|
)
|
|
|
|
sample_hidden_states = text_hidden_states[logits_indices]
|
|
logits = self.model.compute_logits(sample_hidden_states, None)
|
|
if broadcast_pp_output:
|
|
model_output_broadcast_data = (
|
|
{
|
|
"logits": logits.contiguous(),
|
|
}
|
|
if logits is not None
|
|
else {}
|
|
)
|
|
model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
|
|
model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
|
|
)
|
|
assert model_output_broadcast_data is not None
|
|
logits = model_output_broadcast_data["logits"]
|
|
|
|
# Apply structured output bitmasks if present
|
|
if scheduler_output.grammar_bitmask is not None:
|
|
self.apply_grammar_bitmask(scheduler_output, logits)
|
|
|
|
# Sample the next token and get logprobs if needed.
|
|
sampling_metadata = self.input_batch.sampling_metadata
|
|
if spec_decode_metadata is None:
|
|
sampler_output = self.sampler(
|
|
logits=logits,
|
|
sampling_metadata=sampling_metadata,
|
|
)
|
|
else:
|
|
# When indexing with a tensor (bonus_logits_indices), PyTorch
|
|
# creates a new tensor with separate storage from the original
|
|
# logits tensor. This means any in-place operations on bonus_logits
|
|
# won't affect the original logits tensor.
|
|
assert logits is not None
|
|
bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
|
|
sampler_output = self.sampler(
|
|
logits=bonus_logits,
|
|
sampling_metadata=sampling_metadata,
|
|
)
|
|
bonus_token_ids = sampler_output.sampled_token_ids
|
|
|
|
# Just like `bonus_logits`, `target_logits` is a new tensor with
|
|
# separate storage from the original `logits` tensor. Therefore,
|
|
# it is safe to update `target_logits` in place.
|
|
target_logits = logits[spec_decode_metadata.target_logits_indices]
|
|
output_token_ids = self.rejection_sampler(
|
|
spec_decode_metadata,
|
|
None, # draft_probs
|
|
target_logits,
|
|
bonus_token_ids,
|
|
sampling_metadata,
|
|
)
|
|
sampler_output.sampled_token_ids = output_token_ids
|
|
|
|
num_nans_in_logits = {}
|
|
if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
|
|
num_nans_in_logits = self._get_nans_in_logits(logits)
|
|
|
|
# TODO(woosuk): The following loop can be slow since it iterates over
|
|
# the requests one by one. Optimize.
|
|
discard_sampled_tokens_req_indices = []
|
|
for i, req_id in enumerate(self.input_batch.req_ids):
|
|
req_state = self.requests[req_id]
|
|
seq_len = (
|
|
req_state.num_computed_tokens
|
|
+ scheduler_output.num_scheduled_tokens[req_id]
|
|
)
|
|
if seq_len < req_state.num_tokens:
|
|
# Ignore the sampled token for partial prefills.
|
|
# Rewind the generator state as if the token was not sampled.
|
|
# This relies on cuda-specific torch-internal impl details
|
|
generator = self.input_batch.generators.get(i)
|
|
if generator is not None:
|
|
generator.set_offset(generator.get_offset() - 4)
|
|
# Record the index of the request that should not be sampled,
|
|
# so that we could clear the sampled tokens before returning.
|
|
discard_sampled_tokens_req_indices.append(i)
|
|
|
|
# NOTE: GPU -> CPU Sync happens here.
|
|
# Move as many CPU operations as possible before this sync point.
|
|
logprobs_tensors = sampler_output.logprobs_tensors
|
|
logprobs_lists = (
|
|
logprobs_tensors.tolists() if logprobs_tensors is not None else None
|
|
)
|
|
|
|
# Compute prompt logprobs if needed.
|
|
prompt_logprobs_dict = self._get_prompt_logprobs_dict(
|
|
text_hidden_states[:num_scheduled_tokens],
|
|
scheduler_output,
|
|
)
|
|
|
|
# Get the valid generated tokens.
|
|
# import os # unused import
|
|
|
|
sampled_token_ids = sampler_output.sampled_token_ids
|
|
max_gen_len = sampled_token_ids.shape[-1]
|
|
if max_gen_len == 1:
|
|
# No spec decode tokens.
|
|
valid_sampled_token_ids = sampled_token_ids.tolist()
|
|
else:
|
|
# Includes spec decode tokens.
|
|
valid_sampled_token_ids = self.rejection_sampler.parse_output(
|
|
sampled_token_ids,
|
|
self.input_batch.vocab_size,
|
|
)
|
|
# Mask out the sampled tokens that should not be sampled.
|
|
for i in discard_sampled_tokens_req_indices:
|
|
valid_sampled_token_ids[i].clear()
|
|
|
|
# Cache the sampled tokens in the model runner, so that the scheduler
|
|
# doesn't need to send them back.
|
|
# NOTE(woosuk): As an exception, when using PP, the scheduler sends
|
|
# the sampled tokens back, because there's no direct communication
|
|
# between the first-stage worker and the last-stage worker.
|
|
for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):
|
|
if not sampled_ids:
|
|
continue
|
|
|
|
start_idx = self.input_batch.num_tokens_no_spec[req_idx]
|
|
end_idx = start_idx + len(sampled_ids)
|
|
assert end_idx <= self.max_model_len, (
|
|
"Sampled token IDs exceed the max model length. "
|
|
f"Total number of tokens: {end_idx} > max_model_len: "
|
|
f"{self.max_model_len}"
|
|
)
|
|
|
|
self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids
|
|
self.input_batch.num_tokens_no_spec[req_idx] = end_idx
|
|
self.input_batch.num_tokens[req_idx] = end_idx
|
|
req_id = self.input_batch.req_ids[req_idx]
|
|
req_state = self.requests[req_id]
|
|
req_state.output_token_ids.extend(sampled_ids)
|
|
|
|
if not self.speculative_config:
|
|
# Speculative decoding is not enabled.
|
|
spec_token_ids = None
|
|
else:
|
|
assert spec_decode_common_attn_metadata is not None
|
|
spec_token_ids = self.propose_draft_token_ids(
|
|
scheduler_output,
|
|
valid_sampled_token_ids,
|
|
sampling_metadata,
|
|
text_hidden_states,
|
|
sample_hidden_states,
|
|
aux_hidden_states,
|
|
spec_decode_metadata,
|
|
spec_decode_common_attn_metadata,
|
|
)
|
|
|
|
self.eplb_step()
|
|
|
|
return OmniModelRunnerOutput(
|
|
req_ids=self.input_batch.req_ids,
|
|
req_id_to_index=self.input_batch.req_id_to_index,
|
|
sampled_token_ids=valid_sampled_token_ids,
|
|
spec_token_ids=spec_token_ids,
|
|
logprobs=logprobs_lists,
|
|
prompt_logprobs_dict=prompt_logprobs_dict,
|
|
pooler_output=[],
|
|
kv_connector_output=kv_connector_output,
|
|
num_nans_in_logits=num_nans_in_logits,
|
|
multimodal_outputs=multimodal_outputs,
|
|
)
|
|
|
|
@torch.inference_mode()
|
|
def extract_multimodal_outputs(
|
|
self, hidden_states: torch.Tensor
|
|
) -> tuple[torch.Tensor, dict]:
|
|
if (
|
|
hasattr(self.model, "have_multimodal_outputs")
|
|
and self.model.have_multimodal_outputs
|
|
):
|
|
text_hidden_states = hidden_states.text_hidden_states
|
|
multimodal_outputs = hidden_states.multimodal_outputs
|
|
|
|
else:
|
|
text_hidden_states = hidden_states
|
|
multimodal_outputs = {}
|
|
return text_hidden_states, multimodal_outputs
|
|
|
|
@torch.inference_mode()
|
|
def _dummy_run(
|
|
self,
|
|
num_tokens: int,
|
|
capture_attn_cudagraph: bool = False,
|
|
skip_eplb: bool = False,
|
|
is_profile: bool = False,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
|
|
# Padding for DP
|
|
num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
|
|
num_tokens += num_pad
|
|
|
|
# Set num_scheduled_tokens based on num_tokens and max_num_seqs
|
|
# for dummy run with LoRA so that the num_reqs collectively
|
|
# has num_tokens in total.
|
|
assert num_tokens <= self.scheduler_config.max_num_batched_tokens
|
|
max_num_reqs = self.scheduler_config.max_num_seqs
|
|
num_reqs = min(num_tokens, max_num_reqs)
|
|
min_tokens_per_req = num_tokens // num_reqs
|
|
num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
|
|
num_scheduled_tokens_list[-1] += num_tokens % num_reqs
|
|
assert sum(num_scheduled_tokens_list) == num_tokens
|
|
assert len(num_scheduled_tokens_list) == num_reqs
|
|
num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
|
|
|
|
attn_metadata: Optional[dict[str, Any]] = None
|
|
if capture_attn_cudagraph:
|
|
attn_metadata = {}
|
|
|
|
# Make sure max_model_len is used at the graph capture time.
|
|
self.seq_lens_np[:num_reqs] = self.max_model_len
|
|
self.seq_lens_np[num_reqs:] = 0
|
|
self.seq_lens[:num_reqs].copy_(
|
|
self.seq_lens_cpu[:num_reqs], non_blocking=True
|
|
)
|
|
|
|
for kv_cache_group_id, kv_cache_group_spec in enumerate(
|
|
self.kv_cache_config.kv_cache_groups
|
|
):
|
|
common_attn_metadata = CommonAttentionMetadata(
|
|
query_start_loc=self.query_start_loc[: num_reqs + 1],
|
|
query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1],
|
|
seq_lens=self.seq_lens[:num_reqs],
|
|
seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
|
|
num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[ # noqa: E501
|
|
:num_reqs
|
|
],
|
|
num_reqs=num_reqs,
|
|
num_actual_tokens=num_tokens,
|
|
max_query_len=num_tokens,
|
|
block_table_tensor=self.input_batch.block_table[
|
|
kv_cache_group_id
|
|
].get_device_tensor()[:num_reqs],
|
|
slot_mapping=self.input_batch.block_table[
|
|
kv_cache_group_id
|
|
].slot_mapping[:num_tokens],
|
|
causal=True,
|
|
)
|
|
|
|
for attn_group in self.attn_groups[kv_cache_group_id]:
|
|
attn_metadata_i = (
|
|
attn_group.metadata_builder.build_for_cudagraph_capture(
|
|
common_attn_metadata
|
|
)
|
|
)
|
|
for layer_name in kv_cache_group_spec.layer_names:
|
|
attn_metadata[layer_name] = attn_metadata_i
|
|
|
|
with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens):
|
|
if self.is_multimodal_model:
|
|
input_ids = None
|
|
inputs_embeds = self.inputs_embeds[:num_tokens]
|
|
model_mm_kwargs = self._dummy_mm_kwargs(num_reqs)
|
|
else:
|
|
input_ids = self.input_ids[:num_tokens]
|
|
inputs_embeds = None
|
|
model_mm_kwargs = {}
|
|
|
|
if self.uses_mrope:
|
|
positions = self.mrope_positions[:, :num_tokens]
|
|
else:
|
|
positions = self.positions[:num_tokens]
|
|
|
|
if get_pp_group().is_first_rank:
|
|
intermediate_tensors = None
|
|
else:
|
|
if self.intermediate_tensors is None:
|
|
self.intermediate_tensors = (
|
|
self.model.make_empty_intermediate_tensors(
|
|
batch_size=self.max_num_tokens,
|
|
dtype=self.model_config.dtype,
|
|
device=self.device,
|
|
)
|
|
)
|
|
|
|
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
|
|
num_tokens, None, False
|
|
)
|
|
|
|
with self.maybe_randomize_inputs(input_ids), set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=num_tokens,
|
|
num_tokens_across_dp=num_tokens_across_dp,
|
|
):
|
|
outputs = self.model(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
**MultiModalKwargs.as_kwargs(
|
|
model_mm_kwargs,
|
|
device=self.device,
|
|
),
|
|
sampler=None,
|
|
)
|
|
|
|
if self.use_aux_hidden_state_outputs:
|
|
hidden_states, _ = outputs
|
|
else:
|
|
hidden_states = outputs
|
|
|
|
text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
|
|
hidden_states
|
|
)
|
|
|
|
if self.speculative_config and self.speculative_config.use_eagle():
|
|
assert isinstance(self.drafter, EagleProposer)
|
|
self.drafter.dummy_run(num_tokens)
|
|
|
|
# This is necessary to avoid blocking DP.
|
|
# For dummy runs, we typically skip EPLB since we don't have any real
|
|
# requests to process.
|
|
# However, in DP settings, there may be cases when some DP ranks do
|
|
# not have any requests to process, so they're executing dummy batches.
|
|
# In such cases, we still have to trigger EPLB to make sure
|
|
# ranks execute the rearrangement in synchronization.
|
|
if not skip_eplb:
|
|
self.eplb_step(is_dummy=True, is_profile=is_profile)
|
|
|
|
logit_indices = np.cumsum(num_scheduled_tokens) - 1
|
|
return text_hidden_states, text_hidden_states[logit_indices]
|