import copy import gc import time import weakref from contextlib import contextmanager from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union, cast from vllm.model_executor.models.interfaces_base import VllmModelForPooling from vllm.utils.jsontree import json_map_leaves import numpy as np import torch import torch.distributed import torch.nn as nn from tqdm import tqdm from typing_extensions import TypeAlias import vllm.envs as envs from vllm.attention import AttentionType, get_attn_backend from vllm.attention.backends.abstract import AttentionBackend from vllm.attention.layer import Attention from vllm.attention import AttentionMetadata, get_attn_backend from vllm.compilation.counter import compilation_counter from vllm.config import (CompilationLevel, VllmConfig, get_layers_from_vllm_config) from vllm.distributed.eplb.eplb_state import EplbState from vllm.distributed.kv_transfer import (get_kv_transfer_group, has_kv_transfer_group) from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1 from vllm.distributed.parallel_state import ( get_pp_group, get_tp_group, graph_capture, is_global_first_rank, get_tensor_model_parallel_rank, prepare_communication_buffer_for_model) from vllm.forward_context import (BatchDescriptor, DPMetadata, set_forward_context) from vllm.logger import init_logger from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader from vllm.model_executor.models import supports_multimodal from vllm.model_executor.models.interfaces_base import ( VllmModelForPooling, is_pooling_model, is_text_generation_model) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange from vllm.multimodal.utils import group_mm_inputs_by_modality from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingType from vllm.sequence import IntermediateTensors from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler, GiB_bytes, LazyLoader, async_tensor_h2d, cdiv, check_use_alibi, get_dtype_size, is_pin_memory_available, round_up) from vllm.model_executor.model_loader import get_model from vllm.v1.attention.backends.utils import (CommonAttentionMetadata) from vllm.v1.attention.backends.utils import ( AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata, create_fast_prefill_custom_backend, reorder_batch_to_split_decodes_and_prefills, split_attn_metadata) from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm.v1.worker.ubatch_utils import UBatchSlice, UBatchSlices from vllm.v1.worker.cpu_model_runner import _torch_cuda_wrapper from vllm.v1.worker.ubatch_splitting import (check_ubatch_thresholds, ubatch_split) from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, DraftTokenIds, LogprobsLists, LogprobsTensors, ModelRunnerOutput, PoolerOutput, SamplerOutput) from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.spec_decode.medusa import MedusaProposer from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.spec_decode.ngram_proposer import NgramProposer from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext from vllm.v1.worker.block_table import BlockTable from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import (AttentionSpec, ChunkedLocalAttentionSpec, CrossAttentionSpec, EncoderOnlyAttentionSpec, FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, MambaSpec, SlidingWindowSpec) from vllm.utils.jsontree import json_map_leaves # from vllm.sample.logits_processor import LogitsProcessorManager # from vllm.worker.utils import (gather_mm_placeholders, initialize_kv_cache_for_kv_sharing, # sanity_check_mm_encoder_outputs, scatter_mm_placeholders) from vacc_tools.trace_logger import get_trace_api trace_time, register_module_trace, trace_autograd_function, register_optimizer_trace = ( get_trace_api("deepseek") ) logger = init_logger(__name__) AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata] # list when ubatching is enabled PerLayerAttnMetadata: TypeAlias = Union[list[AttnMetadataDict], AttnMetadataDict] def draft_model_eagle_load_model(draft_model, target_model: nn.Module) -> None: draft_model_config = \ draft_model.vllm_config.speculative_config.draft_model_config target_attn_layer_names = set( get_layers_from_vllm_config(draft_model.vllm_config, Attention).keys()) # from vllm.compilation.backends import set_model_tag # with set_model_tag("eagle_head"): draft_model.model = get_model(vllm_config=draft_model.vllm_config, model_config=draft_model_config) draft_attn_layer_names = ( get_layers_from_vllm_config(draft_model.vllm_config, Attention).keys() - target_attn_layer_names) draft_model.attn_layer_names = list(draft_attn_layer_names) if supports_multimodal(target_model): # handle multimodality draft_model.model.config.image_token_index = ( target_model.config.image_token_index) target_language_model = target_model.get_language_model() else: target_language_model = target_model # share embed_tokens with the target model if needed if get_pp_group().world_size == 1 \ and draft_model.model.model.embed_tokens.weight.shape \ == target_language_model.model.embed_tokens.weight.shape: logger.info( "Assuming the EAGLE head shares the same vocab embedding" \ " with the target model." ) del draft_model.model.model.embed_tokens draft_model.model.model.embed_tokens = ( target_language_model.model.embed_tokens) else: logger.info( "The EAGLE head's vocab embedding will be loaded separately" \ " from the target model." ) # share lm_head with the target model if needed # some model definition do not define lm_head explicitly # and reuse embed_tokens for lm_head, e.g., CohereForCausalLM if draft_model.vllm_config.speculative_config.method != "eagle3" and \ hasattr(target_language_model, "lm_head"): logger.info("Loading EAGLE LM head weights from the target model.") draft_model.model.lm_head = target_language_model.lm_head class VACCModelRunner(GPUModelRunner): def __init__(self, vllm_config: VllmConfig, device: torch.device): with _torch_cuda_wrapper(): super().__init__(vllm_config, device) # assert self.speculative_config is None, "spec decode is not supported." self.use_cuda_graph = False self.cascade_attn_enabled = False # remove _postprocess_tenosrs to avoid gpu->cpu # self._postprocess_tenosrs() self.positions = self._make_buffer(self.max_num_tokens, dtype=torch.int32) self.num_accepted_tokens = self._make_buffer(self.max_num_reqs, dtype=torch.int32) # self.slot_mapping = self.slot_mapping.to(torch.int32) self.seq_lens = [0 for i in range(self.max_num_reqs)] if self.speculative_config is not None: self.speculative_config.disable_padded_drafter_batch = True self.deepstack_input_embeds = None if not (self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder): if not self.enable_prompt_embeds: # --enable-prompt-embeds try: self.inputs_embeds = None torch.vacc.empty_cache() except Exception as e: print(f"remove self.inputs_embeds cache fail: {e}") if self.uses_mrope: self.mrope_positions = self._make_buffer( (3, self.max_num_tokens + 1), dtype=torch.int32) 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) # 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 mm_hash in scheduler_output.free_encoder_mm_hashes: self.encoder_cache.pop(mm_hash, 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) reqs_to_add: list[CachedRequestState] = [] # 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 self.is_pooling_model: assert pooling_params is not None task = pooling_params.task assert task is not None, "You did not set `task` in the API" model = cast(VllmModelForPooling, self.get_model()) to_update = model.pooler.get_pooling_updates(task) to_update.apply(pooling_params) req_state = CachedRequestState( req_id=req_id, prompt_token_ids=new_req_data.prompt_token_ids, prompt_embeds=new_req_data.prompt_embeds, deepstack_input_embeds=new_req_data.deepstack_input_embeds, mm_features=new_req_data.mm_features, 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, ) self.requests[req_id] = req_state # Only relevant for models using M-RoPE (e.g, Qwen2-VL) if self.uses_mrope: self._init_mrope_positions(req_state) reqs_to_add.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 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: if new_block_ids is not None: # 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: assert new_block_ids is not None # 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. reqs_to_add.append(req_state) continue # Update the persistent batch. self.input_batch.num_computed_tokens_cpu[req_index] = ( num_computed_tokens) if new_block_ids is not None: 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 request in reqs_to_add: self.input_batch.add_request(request) # 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() # @trace_time('_prepare_inputs') def _prepare_inputs( self, scheduler_output: "SchedulerOutput" ) -> tuple[PerLayerAttnMetadata, torch.Tensor, Optional[SpecDecodeMetadata], np.ndarray, Optional[CommonAttentionMetadata], int, Optional[UBatchSlices], Optional[torch.Tensor]]: """ :return: tuple[ attn_metadata: layer-to-attention_metadata mapping, logits_indices, spec_decode_metadata ] """ total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens assert total_num_scheduled_tokens > 0 num_reqs = self.input_batch.num_reqs assert num_reqs > 0 # OPTIMIZATION: Start copying the block table first. # This way, we can overlap the copy with the following CPU operations. self.input_batch.block_table.commit_block_table(num_reqs) # copy to gpu max_model_len//block_size # Get the number of scheduled tokens for each request. req_ids = self.input_batch.req_ids tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] num_scheduled_tokens = np.array(tokens, dtype=np.int32) max_num_scheduled_tokens = max(tokens) # Get request indices. # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2] req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens) # cu_num_tokens: [2, 5, 3] -> [2, 7, 10] # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] cu_num_tokens, arange = self._get_cumsum_and_arange( num_scheduled_tokens) # Get positions. positions_np = self.positions.np[:total_num_scheduled_tokens] np.add(self.input_batch.num_computed_tokens_cpu[req_indices], arange, out=positions_np) # Calculate M-RoPE positions. # Only relevant for models using M-RoPE (e.g, Qwen2-VL) if self.uses_mrope: self._calc_mrope_positions(scheduler_output) # Get token indices. # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2] # where M is the max_model_len. token_indices = (positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1]) token_indices_tensor = torch.from_numpy(token_indices) # NOTE(woosuk): We use torch.index_select instead of np.take here # because torch.index_select is much faster than np.take for large # tensors. torch.index_select(self.input_batch.token_ids_cpu_tensor.flatten(), 0, token_indices_tensor, out=self.input_ids.cpu[:total_num_scheduled_tokens]) if self.enable_prompt_embeds: is_token_ids = self.input_batch.is_token_ids.flatten() torch.index_select( is_token_ids, 0, token_indices_tensor, out=self.is_token_ids.cpu[:total_num_scheduled_tokens]) # Because we did not pre-allocate a massive prompt_embeds CPU tensor on # the InputBatch, we need to fill in the prompt embeds into the expected # spots in the GpuModelRunner's pre-allocated prompt_embeds tensor. if self.input_batch.req_prompt_embeds: output_idx = 0 deepstack_input_embeds_all_seqs = [] for req_idx in range(num_reqs): num_sched = num_scheduled_tokens[req_idx] # Skip if this request doesn't have embeddings if req_idx not in self.input_batch.req_prompt_embeds: output_idx += num_sched continue # Skip if no tokens scheduled if num_sched <= 0: output_idx += num_sched continue req_embeds = self.input_batch.req_prompt_embeds[req_idx] deepstack_input_embeds = None if self.input_batch.req_deepstack_input_embeds: deepstack_input_embeds = self.input_batch.req_deepstack_input_embeds[req_idx] start_pos = self.input_batch.num_computed_tokens_cpu[req_idx] # Skip if trying to read beyond available embeddings if start_pos >= req_embeds.shape[0]: output_idx += num_sched continue # Copy available embeddings end_pos = start_pos + num_sched actual_end = min(end_pos, req_embeds.shape[0]) actual_num_sched = actual_end - start_pos if actual_num_sched > 0: # for prefill stage self.inputs_embeds.cpu[output_idx:output_idx + actual_num_sched].copy_( req_embeds[start_pos:actual_end] ) if deepstack_input_embeds is not None: deepstack_input_embeds_all_seqs.append(deepstack_input_embeds) output_idx += num_sched if len(deepstack_input_embeds_all_seqs) > 0: # for multi-batch self.deepstack_input_embeds = torch.concatenate(deepstack_input_embeds_all_seqs, -2) self.input_batch.block_table.compute_slot_mapping( req_indices, positions_np) self.input_batch.block_table.commit_slot_mapping( total_num_scheduled_tokens) # copy to gpu # Prepare the attention metadata. self.query_start_loc.np[0] = 0 self.query_start_loc.np[1:num_reqs + 1] = cu_num_tokens # Note: pad query_start_loc to be non-decreasing, as kernels # like FlashAttention requires that self.query_start_loc.np[num_reqs + 1:].fill(cu_num_tokens[-1]) num_computed_tokens = scheduler_output.scheduled_cached_reqs.num_computed_tokens is_all_decode = len(num_computed_tokens) > 0 and all([i > 0 for i in num_computed_tokens]) if not is_all_decode: self.query_start_loc.copy_to_gpu() query_start_loc = self.query_start_loc.gpu[:num_reqs + 1] else: query_start_loc = self.query_start_loc.np[:num_reqs + 1].tolist() num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens num_tokens_padded = num_tokens_unpadded + self.get_local_padding( num_tokens_unpadded) uniform_decode = \ (max_num_scheduled_tokens == self.uniform_decode_query_len) and \ (total_num_scheduled_tokens == num_reqs * max_num_scheduled_tokens) ubatch_slices, num_tokens_after_padding = \ ubatch_split(num_scheduled_tokens, num_tokens_unpadded, num_tokens_padded, uniform_decode=uniform_decode, vllm_config=self.vllm_config) # patch here to seq_lengs to list v0.11.0 # self.seq_lens.np[:num_reqs] = ( # self.input_batch.num_computed_tokens_cpu[:num_reqs] + # num_scheduled_tokens) # # Fill unused with 0 for full cuda graph mode. # self.seq_lens.np[num_reqs:].fill(0) # self.seq_lens.copy_to_gpu() # seq_lens = self.seq_lens.gpu[:num_reqs] # max_seq_len = self.seq_lens.np[:num_reqs].max().item() # self.seq_lens.np[:num_reqs] = ( # self.input_batch.num_computed_tokens_cpu[:num_reqs] + # num_scheduled_tokens) self.seq_lens[:num_reqs] = [int(computed_token) + int(scheduled_token) for computed_token, scheduled_token in zip(self.input_batch.num_computed_tokens_cpu[:num_reqs], num_scheduled_tokens)] seq_lens = self.seq_lens[:num_reqs] max_seq_len = max(seq_lens) # seq_lens is list # num_tokens = [ # self.requests[r].num_tokens for r in self.input_batch.req_ids # ] # num_tokens_np = np.array(num_tokens, dtype=np.int32) # Record the index of requests that should not be sampled, # so that we could clear the sampled tokens before returning # discard_requests_mask = self.seq_lens.np[:num_reqs] < num_tokens_np # discard_request_indices = np.nonzero(discard_requests_mask)[0] # self.num_discarded_requests = len(discard_request_indices) # self.discard_request_indices.np[:self.num_discarded_requests] = ( # discard_request_indices) # self.discard_request_indices.copy_to_gpu(self.num_discarded_requests) # Copy the tensors to the GPU. self._prepare_input_ids(total_num_scheduled_tokens, cu_num_tokens) if self.uses_mrope: # Only relevant for models using M-RoPE (e.g, Qwen2-VL) self.mrope_positions.gpu[:, :total_num_scheduled_tokens].copy_( self.mrope_positions.cpu[:, :total_num_scheduled_tokens], non_blocking=True) else: # Common case (1D positions) self.positions.copy_to_gpu(total_num_scheduled_tokens) # copy position use_spec_decode = len( scheduler_output.scheduled_spec_decode_tokens) > 0 if not use_spec_decode: # NOTE(woosuk): Due to chunked prefills, the batch may contain # partial requests. While we should not sample any token # from these partial requests, we do so for simplicity. # We will ignore the sampled tokens from the partial requests. # TODO: Support prompt logprobs. # logits_indices = query_start_loc[1:] - 1 if not is_all_decode: logits_indices = query_start_loc[1:] - 1 else: logits_indices = None num_draft_tokens = None spec_decode_metadata = None else: # Get the number of draft tokens for each request. # Iterate over the dictionary rather than all requests since not all # requests have draft tokens. num_draft_tokens = np.zeros(num_reqs, dtype=np.int32) # For chunked prefills, use -1 as mask rather than 0, as guided # decoding may rollback speculative tokens. # num_decode_draft_tokens = np.full(num_reqs, -1, dtype=np.int32) for req_id, draft_token_ids in ( scheduler_output.scheduled_spec_decode_tokens.items()): req_idx = self.input_batch.req_id_to_index[req_id] num_draft_tokens[req_idx] = len(draft_token_ids) # num_decode_draft_tokens[req_idx] = (len(draft_token_ids) if ( # self.input_batch.num_computed_tokens_cpu[req_idx] # >= self.input_batch.num_prompt_tokens[req_idx]) else -1) spec_decode_metadata = self._calc_spec_decode_metadata( num_draft_tokens, cu_num_tokens) # copy logits_indices = spec_decode_metadata.logits_indices # For DECODE only cuda graph of some attention backends (e.g., GDN). # self.num_decode_draft_tokens.np[: # num_reqs] = num_decode_draft_tokens # self.num_decode_draft_tokens.np[num_reqs:].fill(-1) # self.num_decode_draft_tokens.copy_to_gpu() # # logits_indices_padded = None # if self.cache_config.kv_sharing_fast_prefill: # logits_indices_padded = self._prepare_kv_sharing_fast_prefill( # logits_indices) attn_metadata: PerLayerAttnMetadata = {} if ubatch_slices is not None: attn_metadata = [dict() for _ in range(len(ubatch_slices))] # Used in the below loop. query_start_loc_cpu = self.query_start_loc.cpu[:num_reqs + 1] seq_lens_cpu = self.seq_lens#.cpu[:num_reqs] num_computed_tokens_cpu = ( self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs]) spec_decode_common_attn_metadata = None if use_spec_decode: self.num_accepted_tokens.np[:num_reqs] = ( self.input_batch.num_accepted_tokens_cpu[:num_reqs]) # self.num_accepted_tokens.np[num_reqs:].fill(1) # self.num_accepted_tokens.copy_to_gpu() # Prepare the attention metadata for each KV cache group and make layers # in the same group share the same metadata. for kv_cache_group_id, kv_cache_group_spec in enumerate( self.kv_cache_config.kv_cache_groups): encoder_seq_lens = self._get_encoder_seq_lens( scheduler_output, kv_cache_group_spec.kv_cache_spec, num_reqs) if isinstance(kv_cache_group_spec.kv_cache_spec, EncoderOnlyAttentionSpec): # Encoder-only layers do not have KV cache, so we need to # create a dummy block table and slot mapping for them. blk_table_tensor = torch.zeros( (num_reqs, 1), dtype=torch.int32, device=self.device, ) slot_mapping = torch.zeros( (total_num_scheduled_tokens, ), dtype=torch.int32, #int64 to int32 device=self.device, ) num_common_prefix_blocks = 0 else: blk_table = self.input_batch.block_table[kv_cache_group_id] # blk_table_tensor = blk_table.get_device_tensor(num_reqs) blk_table_tensor = blk_table.block_table.gpu # slice first num_reqs move to metadata_build slot_mapping = blk_table.slot_mapping.gpu[: total_num_scheduled_tokens] # may copy # Fill unused with -1. Needed for reshape_and_cache in full cuda # graph mode. # blk_table.slot_mapping.gpu[total_num_scheduled_tokens:].fill_( # -1) num_common_prefix_blocks = ( scheduler_output. num_common_prefix_blocks[kv_cache_group_id]) common_attn_metadata = CommonAttentionMetadata( query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc_cpu, seq_lens=seq_lens, seq_lens_cpu=seq_lens_cpu, num_computed_tokens_cpu=num_computed_tokens_cpu, num_reqs=num_reqs, num_actual_tokens=total_num_scheduled_tokens, max_query_len=max_num_scheduled_tokens, max_seq_len=max_seq_len, block_table_tensor=blk_table_tensor, slot_mapping=slot_mapping, logits_indices_padded=None, #logits_indices_padded, num_logits_indices=None, #logits_indices.size(0), causal=True, encoder_seq_lens=encoder_seq_lens, ) if (self.speculative_config and spec_decode_common_attn_metadata is None): if isinstance(self.drafter, EagleProposer): if (self.drafter.attn_layer_names[0] in kv_cache_group_spec.layer_names): spec_decode_common_attn_metadata = common_attn_metadata else: spec_decode_common_attn_metadata = common_attn_metadata for attn_group in self.attn_groups[kv_cache_group_id]: # Prepare for cascade attention if enabled & beneficial. common_prefix_len = 0 builder = attn_group.get_metadata_builder() if self.cascade_attn_enabled: common_prefix_len = self._compute_cascade_attn_prefix_len( num_scheduled_tokens, num_common_prefix_blocks, attn_group.kv_cache_spec, builder, ) extra_attn_metadata_args = {} # if use_spec_decode and isinstance(builder, # GDNAttentionMetadataBuilder): # extra_attn_metadata_args = dict( # num_accepted_tokens=self.num_accepted_tokens. # gpu[:num_reqs], # num_decode_draft_tokens_cpu=self. # num_decode_draft_tokens.cpu[:num_reqs], # ) if ubatch_slices is not None: common_attn_metadata_list = split_attn_metadata( ubatch_slices, common_attn_metadata) for ubid, common_attn_metadata in enumerate( common_attn_metadata_list): attn_metadata_i = (attn_group.get_metadata_builder( ubatch_id=ubid).build( common_prefix_len=common_prefix_len, common_attn_metadata=common_attn_metadata)) for layer_name in kv_cache_group_spec.layer_names: assert type(attn_metadata) is list attn_metadata[ubid][layer_name] = attn_metadata_i else: assert isinstance(attn_metadata, dict) attn_metadata_i = builder.build( common_prefix_len=common_prefix_len, common_attn_metadata=common_attn_metadata, **extra_attn_metadata_args) for layer_name in attn_group.layer_names: attn_metadata[layer_name] = attn_metadata_i # Hot-Swap lora model if self.lora_config: self.set_active_loras(self.input_batch, num_scheduled_tokens) return (attn_metadata, logits_indices, spec_decode_metadata, num_scheduled_tokens, spec_decode_common_attn_metadata, max_num_scheduled_tokens, ubatch_slices, num_tokens_after_padding) def _preprocess( self, scheduler_output: "SchedulerOutput", intermediate_tensors: Optional[IntermediateTensors] = None, ubatch_slices: Optional[UBatchSlices] = None, num_tokens_after_padding: Optional[torch.Tensor] = None, ) -> tuple[int, int, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], torch.Tensor, Optional[IntermediateTensors], dict[str, Any]]: num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens if ubatch_slices: assert num_tokens_after_padding is not None num_input_tokens = int(num_tokens_after_padding[0].item() * 2) self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens) elif ubatch_slices is None: num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens) num_pad, num_tokens_after_padding = self.get_dp_padding( num_input_tokens) num_input_tokens += num_pad deepstack_input_embeds = None # _prepare_inputs may reorder the batch, so we must gather multi # modal outputs after that to ensure the correct order if (self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder): # Run the multimodal encoder if any. self._execute_mm_encoder(scheduler_output) mm_embeds = self._gather_mm_embeddings(scheduler_output) # NOTE(woosuk): To unify token ids and soft tokens (vision # embeddings), we always use embeddings (rather than token ids) # as input to the multimodal model, even when the input is text. inputs_embeds_scheduled = self.model.get_input_embeddings( input_ids=self.input_ids.gpu[:num_scheduled_tokens], multimodal_embeddings=mm_embeds or None, ) # # TODO(woosuk): Avoid the copy. Optimize. # self.inputs_embeds.gpu[:num_scheduled_tokens].copy_( # inputs_embeds_scheduled) input_ids = None # inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] # 完整vl模型，直接从 get_input_embeddings 拿到 device 上的 tensor 直接传, 不需要用self.inputs_embeds.gpu # enable_prompt_embeds 输入是 cpu 的 input_embeddings 才会需要拷贝 inputs_embeds = inputs_embeds_scheduled model_kwargs = { **self._init_model_kwargs(num_scheduled_tokens), **self._extract_mm_kwargs(scheduler_output), } elif self.enable_prompt_embeds and get_pp_group().is_first_rank: # Get the input embeddings for the tokens that are not input embeds, # then put them into the appropriate positions. # TODO(qthequartermasterman): Since even when prompt embeds are # enabled, (a) not all requests will use prompt embeds, and (b) # after the initial prompt is processed, the rest of the generated # tokens will be token ids, it is not desirable to have the # embedding layer outside of the CUDA graph all the time. The v0 # engine avoids this by "double compiling" the CUDA graph, once # with input_ids and again with inputs_embeds, for all num_tokens. # If a batch only has token ids, then including the embedding layer # in the CUDA graph will be more performant (like in the else case # below). # patch here to opt embedding model # token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens] \ # .nonzero(as_tuple=False) \ # .squeeze(1) # Some tokens ids may need to become embeds # if token_ids_idx.numel() > 0: # token_ids = self.input_ids.gpu[token_ids_idx] # tokens_to_embeds = self.model.get_input_embeddings( # input_ids=token_ids) # self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds # inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] token_ids_idx = self.is_token_ids.cpu[:num_scheduled_tokens] \ .nonzero(as_tuple=False) \ .squeeze(1) if token_ids_idx.numel() > 0: inputs_embeds = self.model.get_input_embeddings( input_ids=self.input_ids.gpu[:num_input_tokens]) else: inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] if self.deepstack_input_embeds is not None: # modelruner preprocess deepstack_input_embeds copy to vacc self.deepstack_input_embeds = self.deepstack_input_embeds.to(self.device) deepstack_input_embeds = self.deepstack_input_embeds #TODO for dict # for key in self.deepstack_input_embeds: # self.deepstack_input_embeds[key] = self.deepstack_input_embeds[key].to(self.device) model_kwargs = self._init_model_kwargs(num_input_tokens) input_ids = None else: # For text-only models, we use token ids as input. # While it is possible to use embeddings as input just like the # multimodal models, it is not desirable for performance since # then the embedding layer is not included in the CUDA graph. input_ids = self.input_ids.gpu[:num_input_tokens] inputs_embeds = None model_kwargs = self._init_model_kwargs(num_input_tokens) if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_input_tokens] else: positions = self.positions.gpu[:num_input_tokens] if get_pp_group().is_first_rank: intermediate_tensors = None else: self.intermediate_tensors = intermediate_tensors # intermediate_tensors = self.sync_and_slice_intermediate_tensors( # num_input_tokens, intermediate_tensors, True) if (self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs): encoder_inputs = self._extract_encoder_inputs(scheduler_output) model_kwargs.update(encoder_inputs) if deepstack_input_embeds is not None: model_kwargs['deepstack_input_embeds'] = deepstack_input_embeds return ( num_scheduled_tokens, num_input_tokens, num_tokens_after_padding, input_ids, inputs_embeds, positions, intermediate_tensors, model_kwargs, ) def load_model(self, eep_scale_up: bool = False) -> None: logger.info("Starting to load model %s...", self.model_config.model) with DeviceMemoryProfiler() as m: # noqa: SIM117 time_before_load = time.perf_counter() model_loader = get_model_loader(self.load_config) logger.info("Loading model from scratch...") self.model = model_loader.load_model( vllm_config=self.vllm_config, model_config=self.model_config) if self.lora_config: self.model = self.load_lora_model(self.model, self.model_config, self.scheduler_config, self.lora_config, self.device) if hasattr(self, "drafter"): logger.info("Loading drafter model...") # self.drafter.load_model(self.model) draft_model_eagle_load_model(self.drafter, self.model) if self.use_aux_hidden_state_outputs: self.model.set_aux_hidden_state_layers( self.model.get_eagle3_aux_hidden_state_layers()) time_after_load = time.perf_counter() self.model_memory_usage = m.consumed_memory logger.info("Model loading took %.4f GiB and %.6f seconds", self.model_memory_usage / GiB_bytes, time_after_load - time_before_load) register_module_trace(self.model) register_module_trace(self.sampler) def _sample( self, logits: Optional[torch.Tensor], spec_decode_metadata: Optional[SpecDecodeMetadata], is_first_calculate ) -> SamplerOutput: # 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, is_first_calculate=is_first_calculate, req_ids = self.input_batch.req_ids ) 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. # bonus_logits = logits[spec_decode_metadata.bonus_logits_indices] # # if bonus_logits_indices is int32, int32 => int64, copy + index_tensor_out # copy + index_tensor_out => new + index_tensor_out by use torch.ops.aten.index # logits = logits.to(torch.float32) bonus_logits_indices = spec_decode_metadata.bonus_logits_indices if isinstance(bonus_logits_indices, np.ndarray): if bonus_logits_indices.shape[0] == 1: bonus_logits = logits[bonus_logits_indices[0] : bonus_logits_indices[0] + 1] # only cpoy else: # bonus_logits = logits[bonus_logits_indices] # cpoy + silce bonus_logits_indices = torch.from_numpy(bonus_logits_indices) # new + slice size = len(spec_decode_metadata.bonus_logits_indices.shape) # bonus_logits = logits.new_empty(spec_decode_metadata.bonus_logits_indices.shape + logits.shape[size:] ) torch.ops.aten.index(logits, [bonus_logits_indices], out=bonus_logits) else: size = len(spec_decode_metadata.bonus_logits_indices.shape) bonus_logits = logits.new_empty(spec_decode_metadata.bonus_logits_indices.shape + logits.shape[size:] ) torch.ops.aten.index(logits, [spec_decode_metadata.bonus_logits_indices], out=bonus_logits) sampler_output = self.sampler( logits=bonus_logits, sampling_metadata=sampling_metadata, is_first_calculate=is_first_calculate, req_ids = self.input_batch.req_ids ) 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] target_logits_indices = spec_decode_metadata.target_logits_indices if isinstance(target_logits_indices, np.ndarray): if target_logits_indices.shape[0] == 1: target_logits = logits[target_logits_indices[0] : target_logits_indices[0] + 1] # only copy else: target_logits_indices = torch.from_numpy(target_logits_indices) size = len(spec_decode_metadata.target_logits_indices.shape) target_logits = logits.new_empty(spec_decode_metadata.target_logits_indices.shape + logits.shape[size:] ) torch.ops.aten.index(logits, [target_logits_indices], out=target_logits) # target_logits = logits[target_logits_indices] else: size = len(spec_decode_metadata.target_logits_indices.shape) target_logits = logits.new_empty(spec_decode_metadata.target_logits_indices.shape + logits.shape[size:] ) torch.ops.aten.index(logits, [spec_decode_metadata.target_logits_indices], out=target_logits) 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 self._update_states_after_model_execute(output_token_ids) return sampler_output def _pool( self, hidden_states: torch.Tensor, num_scheduled_tokens: int, num_scheduled_tokens_np: np.ndarray, ) -> ModelRunnerOutput: assert self.input_batch.num_reqs ==\ len(self.input_batch.pooling_params), \ "Either all or none of the requests in" \ " a batch must be pooling request" hidden_states = hidden_states[:num_scheduled_tokens] pooling_metadata = self.input_batch.get_pooling_metadata() pooling_metadata.build_pooling_cursor(num_scheduled_tokens_np.tolist(), device=hidden_states.device) seq_lens_cpu = self.seq_lens[:self.input_batch.num_reqs] model = cast(VllmModelForPooling, self.model) raw_pooler_output: PoolerOutput = model.pooler( hidden_states=hidden_states, pooling_metadata=pooling_metadata, ) raw_pooler_output = json_map_leaves( lambda x: x.to("cpu", non_blocking=True), raw_pooler_output, ) self._sync_device() pooler_output: list[Optional[torch.Tensor]] = [] for raw_output, seq_len, prompt_len in zip( raw_pooler_output, seq_lens_cpu, pooling_metadata.prompt_lens): output = raw_output if seq_len == prompt_len else None pooler_output.append(output) return ModelRunnerOutput( req_ids=self.input_batch.req_ids, req_id_to_index=self.input_batch.req_id_to_index, sampled_token_ids=[], logprobs=None, prompt_logprobs_dict={}, pooler_output=pooler_output, ) @torch.inference_mode() def execute_model( self, scheduler_output: "SchedulerOutput", intermediate_tensors: Optional[IntermediateTensors] = None, ) -> Union[ModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]: with record_function_or_nullcontext("Preprocess"): with self.synchronize_input_prep(): # Update persistent batch states. self._update_states(scheduler_output) if not scheduler_output.total_num_scheduled_tokens: if not has_kv_transfer_group(): # Return empty ModelRunnerOutput if no work to do. return EMPTY_MODEL_RUNNER_OUTPUT return self.kv_connector_no_forward( scheduler_output, self.vllm_config) if self.cache_config.kv_sharing_fast_prefill: assert not self.input_batch.num_prompt_logprobs, ( "--kv-sharing-fast-prefill produces incorrect " "logprobs for prompt tokens, tokens, please disable " "it when the requests need prompt logprobs") # Prepare the decoder inputs. (attn_metadata, logits_indices, spec_decode_metadata, num_scheduled_tokens_np, spec_decode_common_attn_metadata, max_query_len, ubatch_slices, num_tokens_after_padding ) = self._prepare_inputs(scheduler_output) if attn_metadata.items().__iter__().__next__()[1].prefill_metadata is not None: try: torch.vacc.empty_cache() except Exception as e: logger.warn("vacc empty cache skiping...") ( num_scheduled_tokens, num_input_tokens, num_tokens_across_dp, input_ids, inputs_embeds, positions, intermediate_tensors, model_kwargs, ) = self._preprocess(scheduler_output, intermediate_tensors, ubatch_slices, num_tokens_after_padding) uniform_decode = (max_query_len == self.uniform_decode_query_len) and ( num_scheduled_tokens == self.input_batch.num_reqs * max_query_len) batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens, uniform_decode=uniform_decode) cudagraph_runtime_mode, batch_descriptor = \ self.cudagraph_dispatcher.dispatch(batch_descriptor) # This is currently to get around the assert in the DPMetadata # where it wants `num_tokens_across_dp` to align with `num_tokens` if ubatch_slices is not None: num_input_tokens = ubatch_slices[0].num_tokens # 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, cudagraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_descriptor, ubatch_slices=ubatch_slices, ), record_function_or_nullcontext("Forward"), 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, **model_kwargs, ) with record_function_or_nullcontext("Postprocess"): if self.use_aux_hidden_state_outputs: # True when EAGLE 3 is used. hidden_states, aux_hidden_states = model_output else: # Common case. hidden_states = model_output aux_hidden_states = None if not self.broadcast_pp_output: # Common case. if not get_pp_group().is_last_rank: # Return the intermediate tensors. assert isinstance(hidden_states, IntermediateTensors) hidden_states.kv_connector_output = kv_connector_output return hidden_states if self.is_pooling_model: # Return the pooling output. output = self._pool(hidden_states, num_scheduled_tokens, num_scheduled_tokens_np) output.kv_connector_output = kv_connector_output return output if logits_indices is None or logits_indices.shape[0] == hidden_states.shape[0]: sample_hidden_states = hidden_states else: if isinstance(logits_indices, np.ndarray): if logits_indices.shape[0] == 1: sample_hidden_states = hidden_states[logits_indices[0] : logits_indices[0] + 1] # only cpoy else: logits_indices = torch.from_numpy(logits_indices) size = len(logits_indices.shape) new_shape = logits_indices.shape + hidden_states.shape[size:] sample_hidden_states = hidden_states.new_empty(new_shape) # copy + slice torch.ops.aten.index(hidden_states, [logits_indices], out=sample_hidden_states) else: # sample_hidden_states = hidden_states[logits_indices] size = len(logits_indices.shape) new_shape = logits_indices.shape + hidden_states.shape[size:] # new + slice sample_hidden_states = hidden_states.new_empty(new_shape) torch.ops.aten.index(hidden_states, [logits_indices], out=sample_hidden_states) logits = self.model.compute_logits(sample_hidden_states) else: # Rare case. assert not self.is_pooling_model if not get_pp_group().is_last_rank: all_gather_tensors = { "residual": not is_residual_scattered_for_sp( self.vllm_config, num_input_tokens) } get_pp_group().send_tensor_dict( hidden_states.tensors, all_gather_group=get_tp_group(), all_gather_tensors=all_gather_tensors) logits = None else: # sample_hidden_states = hidden_states[logits_indices] if logits_indices is None or logits_indices.shape[0] == hidden_states.shape[0]: sample_hidden_states = hidden_states else: if isinstance(logits_indices, np.ndarray): if logits_indices.shape[0] == 1: sample_hidden_states = hidden_states[logits_indices[0] : logits_indices[0] + 1] # only cpoy else: logits_indices = torch.from_numpy(logits_indices) # new + slice size = len(logits_indices.shape) new_shape = logits_indices.shape + hidden_states.shape[size:] sample_hidden_states = hidden_states.new_empty(new_shape) torch.ops.aten.index(hidden_states, [logits_indices], out=sample_hidden_states) else: # sample_hidden_states = hidden_states[logits_indices] # copy + slice size = len(logits_indices.shape) new_shape = logits_indices.shape + hidden_states.shape[size:] # new + slice sample_hidden_states = hidden_states.new_empty(new_shape) torch.ops.aten.index(hidden_states, [logits_indices], out=sample_hidden_states) logits = self.model.compute_logits(sample_hidden_states) model_output_broadcast_data = {} if logits is not None: model_output_broadcast_data["logits"] = logits.contiguous() 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: apply_grammar_bitmask(scheduler_output, self.input_batch, logits, self.device) with record_function_or_nullcontext("Sample"): is_first_calculate = attn_metadata[list(attn_metadata.keys())[0]].num_prefills sampler_output = self._sample(logits, spec_decode_metadata, is_first_calculate) def propose_draft_token_ids(sampled_token_ids): assert spec_decode_common_attn_metadata is not None with record_function_or_nullcontext("Draft"): self._draft_token_ids = self.propose_draft_token_ids( scheduler_output, sampled_token_ids, self.input_batch.sampling_metadata, hidden_states, sample_hidden_states, aux_hidden_states, spec_decode_metadata, spec_decode_common_attn_metadata, ) use_padded_batch_for_eagle = self.speculative_config and \ self.speculative_config.use_eagle() and \ not self.speculative_config.disable_padded_drafter_batch effective_drafter_max_model_len = self.max_model_len if effective_drafter_max_model_len is None: effective_drafter_max_model_len = self.model_config.max_model_len if (self.speculative_config and self.speculative_config.draft_model_config is not None and self.speculative_config.draft_model_config.max_model_len is not None): effective_drafter_max_model_len = ( self.speculative_config.draft_model_config.max_model_len) input_fits_in_drafter = spec_decode_common_attn_metadata and ( max(spec_decode_common_attn_metadata.seq_lens) + self.speculative_config.num_speculative_tokens <= effective_drafter_max_model_len) if use_padded_batch_for_eagle and input_fits_in_drafter: # EAGLE speculative decoding can use the GPU sampled tokens # as inputs, and does not need to wait for bookkeeping to finish. propose_draft_token_ids(sampler_output.sampled_token_ids) with record_function_or_nullcontext("Bookkeep"): ( num_nans_in_logits, logprobs_lists, valid_sampled_token_ids, prompt_logprobs_dict, req_ids_output_copy, req_id_to_index_output_copy, invalid_req_indices, ) = self._bookkeeping_sync(scheduler_output, sampler_output, logits, hidden_states, num_scheduled_tokens) if (self.speculative_config and not use_padded_batch_for_eagle and input_fits_in_drafter): # ngram and other speculative decoding methods use the sampled # tokens on the CPU, so they are run after bookkeeping. propose_draft_token_ids(valid_sampled_token_ids) with record_function_or_nullcontext("EPLB"): self.eplb_step() output = ModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, sampled_token_ids=valid_sampled_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, ) if not self.use_async_scheduling: return output return AsyncGPUModelRunnerOutput( model_runner_output=output, sampled_token_ids=sampler_output.sampled_token_ids, invalid_req_indices=invalid_req_indices, async_output_copy_stream=self.async_output_copy_stream, ) def propose_draft_token_ids( self, scheduler_output: "SchedulerOutput", sampled_token_ids: Union[torch.Tensor, list[list[int]]], sampling_metadata: SamplingMetadata, hidden_states: torch.Tensor, sample_hidden_states: torch.Tensor, aux_hidden_states: Optional[list[torch.Tensor]], spec_decode_metadata: Optional[SpecDecodeMetadata], common_attn_metadata: CommonAttentionMetadata, ) -> Union[list[list[int]], torch.Tensor]: num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens if self.speculative_config.method == "ngram": assert isinstance(sampled_token_ids, list) assert isinstance(self.drafter, NgramProposer) draft_token_ids = self.drafter.propose( sampled_token_ids, self.input_batch.req_ids, self.input_batch.num_tokens_no_spec, self.input_batch.token_ids_cpu, self.input_batch.spec_decode_unsupported_reqs) elif self.speculative_config.method == "medusa": assert isinstance(sampled_token_ids, list) assert isinstance(self.drafter, MedusaProposer) if sample_hidden_states.shape[0] == len(sampled_token_ids): # The input to the target model does not include draft tokens. hidden_states = sample_hidden_states else: indices = [] offset = 0 assert spec_decode_metadata is not None for num_draft, tokens in zip( spec_decode_metadata.num_draft_tokens, sampled_token_ids): indices.append(offset + len(tokens) - 1) offset += num_draft + 1 indices = torch.tensor(indices, device=self.device) hidden_states = sample_hidden_states[indices] draft_token_ids = self.drafter.propose( target_hidden_states=hidden_states, sampling_metadata=sampling_metadata, ) elif self.speculative_config.use_eagle(): assert isinstance(self.drafter, EagleProposer) if self.speculative_config.disable_padded_drafter_batch: # When padded-batch is disabled, the sampled_token_ids should be # the cpu-side list[list[int]] of valid sampled tokens for each # request, with invalid requests having empty lists. assert isinstance(sampled_token_ids, list), \ "sampled_token_ids should be a python list when" \ "padded-batch is disabled." next_token_ids = self.drafter.prepare_next_token_ids_cpu( sampled_token_ids, self.requests, self.input_batch, scheduler_output.num_scheduled_tokens) else: # When using padded-batch, the sampled_token_ids should be # the gpu tensor of sampled tokens for each request, of shape # (num_reqs, num_spec_tokens + 1) with rejected tokens having # value -1. assert isinstance(sampled_token_ids, torch.Tensor), \ "sampled_token_ids should be a torch.Tensor when" \ "padded-batch is enabled." next_token_ids, valid_sampled_tokens_count = \ self.drafter.prepare_next_token_ids_padded( common_attn_metadata, sampled_token_ids, self.requests, self.input_batch, self.discard_request_indices.gpu, self.num_discarded_requests ) if spec_decode_metadata is None: token_indices_to_sample = None # input_ids can be None for multimodal models. target_token_ids = self.input_ids.gpu[:num_scheduled_tokens] # TODO(woosuk): Support M-RoPE. target_positions = self.positions.gpu[:num_scheduled_tokens] if self.use_aux_hidden_state_outputs: assert aux_hidden_states is not None target_hidden_states = torch.cat( [h[:num_scheduled_tokens] for h in aux_hidden_states], dim=-1) else: target_hidden_states = hidden_states[:num_scheduled_tokens] else: if self.speculative_config.disable_padded_drafter_batch: token_indices_to_sample = None common_attn_metadata, token_indices =\ self.drafter.prepare_inputs( common_attn_metadata, sampled_token_ids, spec_decode_metadata.num_draft_tokens) else: common_attn_metadata, token_indices, \ token_indices_to_sample =\ self.drafter.prepare_inputs_padded( common_attn_metadata, spec_decode_metadata, valid_sampled_tokens_count) # target_token_ids = self.input_ids.gpu[token_indices] # # TODO(woosuk): Support M-RoPE. # target_positions = self.positions.gpu[token_indices] # print('token_indices', token_indices) assert len(token_indices) > 0, len(token_indices) if len(token_indices) > 1: token_indices_vacc = torch.tensor(token_indices, dtype=torch.int32).to(self.input_ids.gpu.device, non_blocking=True) target_token_ids = self.input_ids.gpu.new_empty(token_indices_vacc.shape + self.input_ids.gpu.shape[1:] ) torch.ops.aten.index(self.input_ids.gpu, [token_indices_vacc], out=target_token_ids) else: # len(token_indices) == 1: target_token_ids = self.input_ids.gpu[token_indices[0] : token_indices[0] + 1] # TODO(woosuk): Support M-RoPE. # target_positions = self.positions[token_indices] if len(token_indices) == 1: target_positions = self.positions.gpu[token_indices[0] : token_indices[0] + 1] else: target_positions = self.positions.gpu.new_empty(token_indices_vacc.shape + self.positions.gpu.shape[1:] ) torch.ops.aten.index(self.positions.gpu, [token_indices_vacc], out=target_positions) if self.use_aux_hidden_state_outputs: assert aux_hidden_states is not None target_hidden_states = torch.cat( [h[token_indices] for h in aux_hidden_states], dim=-1) else: # target_hidden_states = hidden_states[token_indices] if len(token_indices) == hidden_states.shape[0]: # all accepted no need slice target_hidden_states = hidden_states elif len(token_indices) == 1: # batch == 1 and noly 1 element: use copy instead of copy + slice target_hidden_states = hidden_states[token_indices[0] : token_indices[0] + 1] else: target_hidden_states = hidden_states.new_empty(token_indices_vacc.shape + hidden_states.shape[1:]) torch.ops.aten.index(hidden_states, [token_indices_vacc], out=target_hidden_states) mm_embeds = None if self.supports_mm_inputs: mm_embeds = self._gather_mm_embeddings(scheduler_output, shift_computed_tokens=1) draft_token_ids = self.drafter.propose( target_token_ids=target_token_ids, target_positions=target_positions, target_hidden_states=target_hidden_states, next_token_ids=next_token_ids, last_token_indices=token_indices_to_sample, sampling_metadata=sampling_metadata, common_attn_metadata=common_attn_metadata, mm_embeds=mm_embeds, ) return draft_token_ids def _calc_spec_decode_metadata( self, num_draft_tokens: np.ndarray, cu_num_scheduled_tokens: np.ndarray, ) -> SpecDecodeMetadata: # Inputs: # cu_num_scheduled_tokens: [ 4, 104, 107, 207, 209] # num_draft_tokens: [ 3, 0, 2, 0, 1] # Outputs: # cu_num_draft_tokens: [ 3, 3, 5, 5, 6] # logits_indices: [ 0, 1, 2, 3, 103, 104, 105, 106, # 206, 207, 208] # target_logits_indices: [ 0, 1, 2, 5, 6, 9] # bonus_logits_indices: [ 3, 4, 7, 8, 10] # Compute the logits indices. # [4, 1, 3, 1, 2] num_sampled_tokens = num_draft_tokens + 1 # Step 1. cu_num_sampled_tokens: [4, 5, 8, 9, 11] # arange: [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1] cu_num_sampled_tokens, arange = self._get_cumsum_and_arange( num_sampled_tokens, cumsum_dtype=np.int32) # Step 2. [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207] logits_indices = np.repeat( cu_num_scheduled_tokens - num_sampled_tokens, num_sampled_tokens) # Step 3. [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208] logits_indices += arange # Compute the bonus logits indices. bonus_logits_indices = cu_num_sampled_tokens - 1 # Compute the draft logits indices. # cu_num_draft_tokens: [3, 3, 5, 5, 6] # arange: [0, 1, 2, 0, 1, 0] cu_num_draft_tokens, arange = self._get_cumsum_and_arange( num_draft_tokens, cumsum_dtype=np.int32) # [0, 0, 0, 5, 5, 9] target_logits_indices = np.repeat( cu_num_sampled_tokens - num_sampled_tokens, num_draft_tokens) # [0, 1, 2, 5, 6, 9] target_logits_indices += arange # draft_token_ids = self.input_ids.gpu[logits_indices] # draft_token_ids = draft_token_ids[target_logits_indices + 1] indice = (target_logits_indices + 1).tolist() if len(indice) == 1: draft_token_ids = self.input_ids.gpu[indice[0] : indice[0] + 1] # vacc_copy else: draft_token_ids = self.input_ids.gpu[indice] # copy + copy + index_out # TODO: Optimize the CPU -> GPU copy. # cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to( # self.device, non_blocking=True) # logits_indices = torch.from_numpy(logits_indices).to(self.device, # non_blocking=True) # target_logits_indices = torch.from_numpy(target_logits_indices).to( # self.device, non_blocking=True) # bonus_logits_indices = torch.from_numpy(bonus_logits_indices).to( # self.device, non_blocking=True) # Compute the draft token ids. # draft_token_indices: [ 1, 2, 3, 105, 106, 208] metadata = SpecDecodeMetadata( draft_token_ids=draft_token_ids, num_draft_tokens=num_draft_tokens.tolist(), cu_num_draft_tokens=None, #cu_num_draft_tokens, target_logits_indices=target_logits_indices, bonus_logits_indices=bonus_logits_indices, logits_indices=logits_indices, ) return metadata def warming_up_model(self) -> None: logger.info("Warming up model for the compilation...") # Only generate graph for the generic shape with _set_global_compilation_settings(self.vllm_config): self._dummy_run(max(16, self.max_num_reqs)) logger.info("Warming up done.") def _init_device_properties(self) -> None: pass def _sync_device(self) -> None: pass def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None: """ Update the order of requests in the batch based on the attention backend's needs. For example, some attention backends (namely MLA) may want to separate requests based on if the attention computation will be compute-bound or memory-bound. Args: scheduler_output: The scheduler output. """ # Attention free models have zero kv_cache_groups, however models # like Mamba are also attention free but use the kv_cache for # keeping its internal state. This is why we check the number # of kv_cache groups instead of solely checking # for self.model_config.is_attention_free. if len(self.kv_cache_config.kv_cache_groups) == 0: return if len(self.kv_cache_config.kv_cache_groups) > 1: raise ValueError("Multiple KVCacheGroups is not" "currently supported with CPU model runner.") # Guard against encoder-only / pooling models where `attn_groups` # may be empty or lack the expected metadata_builder. # Without this check, accessing `attn_groups[0][0]` would trigger # an AssertionError on CPU backend. if not hasattr(self, "attn_groups") or not self.attn_groups: return if not self.attn_groups[0]: return mb = getattr(self.attn_groups[0][0], "metadata_builder", None) # from vllm.v1.attention.backends.cpu_attn import TorchSDPAMetadataBuilderV1 from vllm_vacc.vllm.v1.attention.backends.vacc_attn import VACCAttentionMetadataBuilderV1 mb = getattr(self.attn_groups[0][0], "metadata_builders", None) if isinstance(mb, list): if not isinstance(mb[0], VACCAttentionMetadataBuilderV1): return mb[0].reorder_batch(self.input_batch, scheduler_output) return elif not isinstance(mb, VACCAttentionMetadataBuilderV1): # Encoder-only / rerank models do not benefit from reordering, # so we safely skip here. return # Safe path for decoder/attention-heavy models mb.reorder_batch(self.input_batch, scheduler_output) @contextmanager def _set_global_compilation_settings(config: VllmConfig): import torch._inductor.config inductor_config = config.compilation_config.inductor_compile_config try: # Note: The MKLDNN and CPPGEMM backend requires freezing parameters. freezing_value = torch._inductor.config.freezing if inductor_config.get("max_autotune", False): torch._inductor.config.freezing = True yield finally: torch._inductor.config.freezing = freezing_value