FoundationVision/LlamaGen
64
1import contextlib2import time3from enum import IntEnum4from typing import Dict, List, NamedTuple, Optional, Set, Tuple5 6import numpy as np7import torch8import torch.nn as nn9 10from vllm.attention import (AttentionMetadata, AttentionMetadataPerStage,11 get_attn_backend)12from vllm.config import (DeviceConfig, LoadConfig, LoRAConfig, ModelConfig,13 ParallelConfig, SchedulerConfig, VisionLanguageConfig)14from vllm.distributed import broadcast_tensor_dict, with_pynccl_for_all_reduce15from vllm.distributed.device_communicators import (custom_all_reduce,16 pynccl_utils)17from vllm.logger import init_logger18from vllm.lora.layers import LoRAMapping19from vllm.lora.request import LoRARequest20from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager21from vllm.model_executor import SamplingMetadata22from vllm.model_executor.model_loader import get_model23from vllm.sampling_params import SamplingParams, SamplingType24from vllm.sequence import (MultiModalData, SamplerOutput, SequenceData,25 SequenceGroupMetadata)26from vllm.utils import (CudaMemoryProfiler, async_tensor_h2d, is_hip,27 is_pin_memory_available, make_tensor_with_pad,28 maybe_expand_dim)29from serve.gpt_model import GPT_models30 31logger = init_logger(__name__)32 33_PAD_SLOT_ID = -134LORA_WARMUP_RANK = 835_BATCH_SIZE_ALIGNMENT = 836# Capture graphs for token size 1, 2, 4, 8, 16, 24, 32, 40, ..., 256.37# NOTE: _get_graph_batch_size needs to be updated if this list is changed.38_BATCH_SIZES_TO_CAPTURE = [1, 2, 4] + [39 _BATCH_SIZE_ALIGNMENT * i for i in range(1, 33)40]41 42 43class PreparePromptMetadata(NamedTuple):44 input_tokens: List[int]45 input_positions: List[int]46 attn_metadata: Optional[AttentionMetadataPerStage]47 prompt_lens: List[int]48 subquery_lens: List[int]49 lora_index_mapping: List[int]50 lora_prompt_mapping: List[int]51 lora_requests: Set[LoRARequest]52 multi_modal_input: Optional[torch.Tensor]53 slot_mapping: List[int]54 55 @classmethod56 def empty(cls):57 return PreparePromptMetadata(58 input_tokens=[],59 input_positions=[],60 attn_metadata=None,61 prompt_lens=[],62 subquery_lens=[],63 lora_index_mapping=[],64 lora_prompt_mapping=[],65 lora_requests=set(),66 multi_modal_input=None,67 slot_mapping=[],68 )69 70 71class PrepareDecodeMetadata(NamedTuple):72 input_tokens: List[int]73 input_positions: List[int]74 attn_metadata: Optional[AttentionMetadata]75 lora_index_mapping: List[int]76 lora_prompt_mapping: List[int]77 lora_requests: Set[LoRARequest]78 slot_mapping: List[int]79 80 @classmethod81 def empty(cls):82 return PrepareDecodeMetadata(83 input_tokens=[],84 input_positions=[],85 attn_metadata=None,86 lora_index_mapping=[],87 lora_prompt_mapping=[],88 lora_requests=set(),89 slot_mapping=[],90 )91 92 93# How batches are constructed.94class BatchType(IntEnum):95 # Every batch is prefill.96 PREFILL = 097 # Every batch is decode.98 DECODE = 199 # Batch is a mixture of prefill and decode.100 MIXED = 2101 102 103class ModelRunner:104 105 def __init__(106 self,107 model_config: ModelConfig,108 parallel_config: ParallelConfig,109 scheduler_config: SchedulerConfig,110 device_config: DeviceConfig,111 load_config: LoadConfig,112 lora_config: Optional[LoRAConfig],113 kv_cache_dtype: Optional[str] = "auto",114 is_driver_worker: bool = False,115 vision_language_config: Optional[VisionLanguageConfig] = None,116 ):117 self.model_config = model_config118 self.parallel_config = parallel_config119 self.scheduler_config = scheduler_config120 self.lora_config = lora_config121 self.load_config = load_config122 self.is_driver_worker = is_driver_worker123 124 # model_config can be None in tests/samplers/test_sampler.py.125 # FIXME(woosuk): This is a hack to make the tests work. Refactor this.126 self.sliding_window = (model_config.get_sliding_window()127 if model_config is not None else None)128 self.device_config = (device_config129 if device_config is not None else DeviceConfig())130 self.device = self.device_config.device131 132 # Set after load_model.133 self.lora_manager: LRUCacheWorkerLoRAManager = None134 135 self.graph_runners: Dict[int, CUDAGraphRunner] = {}136 self.graph_memory_pool: Optional[Tuple[137 int, int]] = None # Set during graph capture.138 139 self.max_context_len_to_capture = (140 self.model_config.max_context_len_to_capture141 if self.model_config is not None else 0)142 143 self.pin_memory = is_pin_memory_available()144 self.kv_cache_dtype = kv_cache_dtype145 self.vision_language_config = vision_language_config146 147 self.attn_backend = get_attn_backend(148 self.model_config.dtype if model_config is not None else None)149 150 # Lazy initialization151 self.model: torch.nn.Module # Set after load_model152 self.block_size: int # Set after initial profiling.153 # When using CUDA graph, the input block tables must be padded to154 # max_context_len_to_capture. However, creating the block table in155 # Python can be expensive. To optimize this, we cache the block table156 # in numpy and only copy the actual input content at every iteration.157 # The shape of the cached block table will be158 # (max batch size to capture, max context len to capture / block size).159 self.graph_block_tables: torch.Tensor # Set after initial profiling.160 161 def load_model(self, args) -> None:162 with CudaMemoryProfiler() as m:163 precision = {'none': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16}[args.precision]164 latent_size = args.image_size // args.downsample_size 165 gpt_model = GPT_models[args.gpt_model](166 vocab_size=args.codebook_size,167 block_size=latent_size ** 2,168 num_classes=args.num_classes,169 cls_token_num=args.cls_token_num,170 model_type=args.gpt_type,171 cfg_scale=args.cfg_scale,172 ).to(device='cuda', dtype=precision) # TODO: make device configurable173 174 checkpoint = torch.load(args.gpt_ckpt, map_location="cpu")175 if args.from_fsdp: # fspd176 model_weight = checkpoint177 elif "model" in checkpoint: # ddp178 model_weight = checkpoint["model"]179 elif "state_dict" in checkpoint:180 model_weight = checkpoint["state_dict"]181 else:182 raise Exception("please check model weight")183 gpt_model.custom_load_state_dict(model_weight)184 gpt_model.eval()185 del checkpoint186 self.model = gpt_model187 188 self.model_memory_usage = m.consumed_memory189 logger.info(f"Loading model weights took "190 f"{self.model_memory_usage / float(2**30):.4f} GB")191 192 if self.lora_config:193 assert hasattr(self.model, "supported_lora_modules"194 ) and self.model.supported_lora_modules, (195 "Model does not support LoRA")196 assert hasattr(197 self.model,198 "embedding_modules"), "Model does not have embedding_modules"199 assert hasattr(self.model, "embedding_padding_modules"200 ), "Model does not have embedding_padding_modules"201 self.lora_manager = LRUCacheWorkerLoRAManager(202 self.scheduler_config.max_num_seqs,203 self.scheduler_config.max_num_batched_tokens, self.vocab_size,204 self.lora_config, self.device, self.model.embedding_modules,205 self.model.embedding_padding_modules)206 self.model = self.lora_manager.create_lora_manager(self.model)207 208 if self.kv_cache_dtype == "fp8" and is_hip():209 # Currently scaled KV cache is only enabled on ROCm210 if self.model_config.quantization_param_path is not None:211 if callable(getattr(self.model, "load_kv_cache_scales", None)):212 self.model.load_kv_cache_scales(213 self.model_config.quantization_param_path)214 else:215 raise RuntimeError("Using FP8 KV cache and scaling "216 "factors provided but model "217 f"{self.model.__class__} does not "218 "support loading scaling factors.")219 else:220 logger.warn("Using FP8 KV cache but no scaling factors "221 "provided. Defaulting to scaling factors of 1.0. "222 "This may lead to less accurate results!")223 elif self.model_config.quantization_param_path is not None:224 logger.warn("KV cache scaling factors provided, "225 "but the KV cache data type is not FP8. "226 "KV cache scaling factors will not be used.")227 228 def set_block_size(self, block_size: int) -> None:229 self.block_size = block_size230 231 self.graph_block_tables = np.zeros(232 (max(_BATCH_SIZES_TO_CAPTURE), self.get_max_block_per_batch()),233 dtype=np.int32)234 235 def get_max_block_per_batch(self) -> int:236 block_size = self.block_size237 return (self.max_context_len_to_capture + block_size - 1) // block_size238 239 def _prepare_prompt(240 self,241 seq_group_metadata_list: List[SequenceGroupMetadata],242 ) -> PreparePromptMetadata:243 input_tokens: List[int] = []244 input_positions: List[int] = []245 slot_mapping: List[int] = []246 lora_index_mapping: List[int] = []247 lora_prompt_mapping: List[int] = []248 lora_requests: Set[LoRARequest] = set()249 250 prompt_lens: List[int] = []251 context_lens: List[int] = []252 subquery_lens: List[int] = []253 prefix_block_tables: List[List[int]] = []254 multi_modal_input_list: List[torch.Tensor] = []255 256 if len(seq_group_metadata_list) == 0:257 return PreparePromptMetadata.empty()258 259 for seq_group_metadata in seq_group_metadata_list:260 assert seq_group_metadata.is_prompt261 seq_ids = list(seq_group_metadata.seq_data.keys())262 assert len(seq_ids) == 1263 seq_id = seq_ids[0]264 265 computed_block_nums = seq_group_metadata.computed_block_nums266 if (self.scheduler_config is not None267 and self.scheduler_config.chunked_prefill_enabled268 and not (computed_block_nums is None269 or computed_block_nums == [])):270 raise RuntimeError(271 "chunked prefill cannot be used with prefix caching "272 "now.")273 274 token_chunk_size = seq_group_metadata.token_chunk_size275 seq_data = seq_group_metadata.seq_data[seq_id]276 computed_len = seq_data.get_num_computed_tokens()277 # We should use get_len here because in case of preemption278 # it contains output tokens.279 prefill_end = min(seq_data.get_len(),280 computed_len + token_chunk_size)281 prompt_tokens = seq_data.get_token_ids()[computed_len:prefill_end]282 prompt_len = prefill_end283 prompt_lens.append(prompt_len)284 285 # NOTE: This only works for oooooooxxx style attention.286 if computed_block_nums is not None and len(287 computed_block_nums) > 0 and self.sliding_window is None:288 # Prefix is not supported with sliding_window289 computed_len = len(computed_block_nums) * self.block_size290 prompt_tokens = prompt_tokens[computed_len:]291 prefix_block_tables.append(computed_block_nums)292 elif self.scheduler_config.chunked_prefill_enabled:293 if seq_group_metadata.block_tables is not None:294 # Prefill has chunked before.295 block_table = seq_group_metadata.block_tables[seq_id]296 prefix_block_tables.append(block_table)297 else:298 # The first prefill.299 prefix_block_tables.append([])300 else:301 prefix_block_tables.append([])302 # Right now, prefill start is always 0. However, this303 # assumption can be changed once chunked prefill is introduced.304 assert computed_len == 0305 306 # actual prompt lens307 context_lens.append(computed_len)308 subquery_lens.append(prompt_len - computed_len)309 310 input_tokens.extend(prompt_tokens)311 # NOTE(woosuk): Here we assume that the first token in the prompt312 # is always the first token in the sequence.313 input_positions.extend(list(range(computed_len, prefill_end)))314 lora_id = seq_group_metadata.lora_int_id315 316 if lora_id > 0:317 lora_requests.add(seq_group_metadata.lora_request)318 319 lora_index_mapping += [lora_id] * (prompt_len - computed_len)320 lora_prompt_mapping.extend(321 [lora_id] *322 (prompt_len - computed_len323 if seq_group_metadata.sampling_params.prompt_logprobs else 1))324 325 if seq_group_metadata.multi_modal_data:326 multi_modal_input_list.append(327 seq_group_metadata.multi_modal_data.data)328 329 if seq_group_metadata.block_tables is None:330 # During memory profiling, the block tables are not initialized331 # yet. In this case, we just use a dummy slot mapping.332 slot_mapping.extend([_PAD_SLOT_ID] * prompt_len)333 continue334 335 # Compute the slot mapping.336 block_table = seq_group_metadata.block_tables[seq_id]337 # Mask the [0, start_idx) tokens of the prompt with _PAD_SLOT_ID,338 # where start_idx is max(0, prompt_len - sliding_window).339 # For example, if the prompt len is 10, sliding window is 8, and340 # block size is 4, the first two tokens are masked and the slot341 # mapping will be [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1].342 start_idx = 0343 if self.sliding_window is not None:344 assert computed_len == 0, (345 "Prefix caching is currently not supported with "346 "sliding window attention")347 start_idx = max(0, prompt_len - self.sliding_window)348 349 for i in range(computed_len, prefill_end):350 if i < start_idx:351 slot_mapping.append(_PAD_SLOT_ID)352 continue353 354 block_number = block_table[i // self.block_size]355 block_offset = i % self.block_size356 slot = block_number * self.block_size + block_offset357 slot_mapping.append(slot)358 359 max_subquery_len = max(subquery_lens)360 max_prompt_len = max(prompt_lens)361 assert max_subquery_len > 0362 363 context_lens_tensor = torch.tensor(context_lens,364 dtype=torch.int,365 device=self.device)366 367 if multi_modal_input_list:368 assert self.vision_language_config, (369 "Multi-modal inputs are only supported by "370 "vision language models.")371 multi_modal_input = torch.cat(multi_modal_input_list,372 dim=0).to(self.device)373 else:374 multi_modal_input = None375 376 # Prepare prefix block tables377 max_prompt_block_table_len = max(len(t) for t in prefix_block_tables)378 block_tables = make_tensor_with_pad(379 prefix_block_tables,380 max_len=max_prompt_block_table_len,381 pad=0,382 dtype=torch.int,383 device=self.device,384 )385 386 # Query length can be shorter than key (i.e., prompt) when prefill387 # is chunked or prefix cached.388 subquery_lens_tensor = torch.tensor(subquery_lens,389 dtype=torch.long,390 device=self.device)391 subquery_start_loc = torch.zeros(subquery_lens_tensor.shape[0] + 1,392 dtype=torch.int32,393 device=self.device)394 395 prompt_lens_tensor = torch.tensor(prompt_lens,396 dtype=torch.long,397 device=self.device)398 seq_start_loc = torch.zeros(prompt_lens_tensor.shape[0] + 1,399 dtype=torch.int32,400 device=self.device)401 402 torch.cumsum(subquery_lens_tensor,403 dim=0,404 dtype=subquery_start_loc.dtype,405 out=subquery_start_loc[1:])406 407 torch.cumsum(prompt_lens_tensor,408 dim=0,409 dtype=seq_start_loc.dtype,410 out=seq_start_loc[1:])411 412 attn_metadata = self.attn_backend.make_metadata(413 is_prompt=True,414 prompt_lens=prompt_lens,415 prompt_lens_tensor=prompt_lens_tensor,416 max_subquery_len=max_subquery_len,417 max_context_len=None,418 max_prompt_len=max_prompt_len,419 subquery_start_loc=subquery_start_loc,420 seq_start_loc=seq_start_loc,421 context_lens=context_lens_tensor,422 block_tables=block_tables,423 use_cuda_graph=False,424 )425 426 return PreparePromptMetadata(427 input_tokens=input_tokens,428 input_positions=input_positions,429 attn_metadata=attn_metadata,430 prompt_lens=prompt_lens,431 subquery_lens=subquery_lens,432 lora_index_mapping=lora_index_mapping,433 lora_prompt_mapping=lora_prompt_mapping,434 lora_requests=lora_requests,435 multi_modal_input=multi_modal_input,436 slot_mapping=slot_mapping,437 )438 439 def _prepare_decode(440 self,441 seq_group_metadata_list: List[SequenceGroupMetadata],442 ) -> PrepareDecodeMetadata:443 input_tokens: List[int] = []444 input_positions: List[int] = []445 slot_mapping: List[int] = []446 context_lens: List[int] = []447 block_tables: List[List[int]] = []448 lora_index_mapping: List[int] = []449 lora_prompt_mapping: List[int] = []450 lora_requests: Set[LoRARequest] = set()451 452 if len(seq_group_metadata_list) == 0:453 return PrepareDecodeMetadata.empty()454 455 for seq_group_metadata in seq_group_metadata_list:456 assert not seq_group_metadata.is_prompt457 assert seq_group_metadata.token_chunk_size == 1458 459 seq_ids = list(seq_group_metadata.seq_data.keys())460 lora_id = seq_group_metadata.lora_int_id461 462 if lora_id > 0:463 lora_requests.add(seq_group_metadata.lora_request)464 465 for seq_id in seq_ids:466 seq_data = seq_group_metadata.seq_data[seq_id]467 generation_token = seq_data.get_last_token_id()468 input_tokens.append(generation_token)469 470 seq_len = seq_data.get_len()471 position = seq_len - 1472 input_positions.append(position)473 474 context_len = seq_len if self.sliding_window is None else min(475 seq_len, self.sliding_window)476 context_lens.append(context_len)477 478 block_table = seq_group_metadata.block_tables[seq_id]479 block_number = block_table[position // self.block_size]480 block_offset = position % self.block_size481 slot = block_number * self.block_size + block_offset482 slot_mapping.append(slot)483 lora_index_mapping.append(lora_id)484 lora_prompt_mapping.append(lora_id)485 486 if self.sliding_window is not None:487 sliding_window_blocks = (self.sliding_window //488 self.block_size)489 block_table = block_table[-sliding_window_blocks:]490 block_tables.append(block_table)491 492 # vLLM uses cuda graph only for decoding requests.493 # See `capture_model` API for more details.494 # For decoding requests, batch_size == input_tokens.495 batch_size = len(input_tokens)496 max_context_len = max(context_lens)497 use_captured_graph = (498 not self.model_config.enforce_eager499 and batch_size <= _BATCH_SIZES_TO_CAPTURE[-1]500 and max_context_len <= self.max_context_len_to_capture)501 if use_captured_graph:502 graph_batch_size = _get_graph_batch_size(batch_size)503 assert graph_batch_size >= batch_size504 for _ in range(graph_batch_size - batch_size):505 input_tokens.append(0)506 input_positions.append(0)507 slot_mapping.append(_PAD_SLOT_ID)508 context_lens.append(1)509 block_tables.append([])510 lora_index_mapping.append(0)511 batch_size = graph_batch_size512 513 context_lens_tensor = torch.tensor(context_lens,514 dtype=torch.int,515 device=self.device)516 517 if use_captured_graph:518 # When using cuda-graph all these tensors should be519 # padded.520 assert context_lens_tensor.shape[0] == len(input_tokens)521 assert context_lens_tensor.shape[0] == len(input_positions)522 assert context_lens_tensor.shape[0] == len(slot_mapping)523 524 # The shape of graph_block_tables is525 # [max batch size, max context len // block size].526 input_block_tables = self.graph_block_tables[:batch_size]527 for i, block_table in enumerate(block_tables):528 if block_table:529 input_block_tables[i, :len(block_table)] = block_table530 block_tables = torch.tensor(input_block_tables, device=self.device)531 else:532 max_block_table_len = max(533 len(block_table) for block_table in block_tables)534 block_tables = make_tensor_with_pad(535 block_tables,536 max_len=max_block_table_len,537 pad=0,538 dtype=torch.int,539 device=self.device,540 )541 542 attn_metadata = self.attn_backend.make_metadata(543 is_prompt=False,544 prompt_lens=None,545 prompt_lens_tensor=None,546 max_subquery_len=None,547 max_context_len=max_context_len,548 max_prompt_len=None,549 subquery_start_loc=None,550 seq_start_loc=None,551 context_lens=context_lens_tensor,552 block_tables=block_tables,553 use_cuda_graph=use_captured_graph,554 )555 return PrepareDecodeMetadata(556 input_tokens=input_tokens,557 input_positions=input_positions,558 attn_metadata=attn_metadata,559 lora_index_mapping=lora_index_mapping,560 lora_prompt_mapping=lora_prompt_mapping,561 lora_requests=lora_requests,562 slot_mapping=slot_mapping,563 )564 565 def _prepare_sample(566 self,567 seq_group_metadata_list: List[SequenceGroupMetadata],568 prompt_lens: List[int],569 subquery_lens: Optional[List[int]],570 ) -> SamplingMetadata:571 seq_groups: List[Tuple[List[int], SamplingParams]] = []572 selected_token_indices: List[int] = []573 generators: List[torch.Generator] = []574 selected_token_start_idx = 0575 categorized_sample_indices: Dict[SamplingType,576 List[Tuple[int, int]]] = {577 t: []578 for t in SamplingType579 }580 categorized_sample_indices_start_idx = 0581 categorized_sampled_token_indices_start_idx = 0582 583 for i, seq_group_metadata in enumerate(seq_group_metadata_list):584 seq_ids = list(seq_group_metadata.seq_data.keys())585 sampling_params = seq_group_metadata.sampling_params586 seq_groups.append((seq_ids, sampling_params))587 588 if seq_group_metadata.is_prompt:589 assert len(seq_ids) == 1590 assert subquery_lens is not None591 subquery_len = subquery_lens[i]592 if sampling_params.prompt_logprobs is not None:593 # NOTE: prompt token positions do not need sample, skip594 categorized_sample_indices_start_idx += subquery_len - 1595 596 categorized_sample_indices[597 sampling_params.sampling_type].append(598 (categorized_sample_indices_start_idx,599 categorized_sampled_token_indices_start_idx))600 categorized_sample_indices_start_idx += 1601 categorized_sampled_token_indices_start_idx += 1602 603 if sampling_params.prompt_logprobs is not None:604 selected_token_indices.extend(605 range(selected_token_start_idx,606 selected_token_start_idx + subquery_len - 1))607 selected_token_indices.append(selected_token_start_idx +608 subquery_len - 1)609 selected_token_start_idx += subquery_len610 611 if sampling_params.seed is not None:612 seq_group_metadata.state.generator = torch.Generator(613 device=self.device).manual_seed(sampling_params.seed)614 else:615 num_seqs = len(seq_ids)616 selected_token_indices.extend(617 range(selected_token_start_idx,618 selected_token_start_idx + num_seqs))619 selected_token_start_idx += num_seqs620 621 categorized_sample_indices[622 sampling_params.sampling_type].extend(623 list(624 zip(625 range(626 categorized_sample_indices_start_idx,627 categorized_sample_indices_start_idx +628 num_seqs),629 range(630 categorized_sampled_token_indices_start_idx,631 categorized_sampled_token_indices_start_idx632 + num_seqs))))633 categorized_sample_indices_start_idx += num_seqs634 categorized_sampled_token_indices_start_idx += num_seqs635 636 if sampling_params.seed is not None:637 generators.append(seq_group_metadata.state.generator)638 639 selected_token_indices = async_tensor_h2d(selected_token_indices,640 dtype=torch.long,641 target_device=self.device,642 pin_memory=self.pin_memory)643 644 categorized_sample_indices = {645 t: maybe_expand_dim(646 async_tensor_h2d(seq_ids,647 dtype=torch.int,648 target_device=self.device,649 pin_memory=self.pin_memory), 2, 2)650 for t, seq_ids in categorized_sample_indices.items()651 }652 653 seq_data: Dict[int, SequenceData] = {}654 for seq_group_metadata in seq_group_metadata_list:655 seq_data.update(seq_group_metadata.seq_data)656 657 sampling_metadata = SamplingMetadata(658 seq_groups=seq_groups,659 seq_data=seq_data,660 prompt_lens=prompt_lens,661 selected_token_indices=selected_token_indices,662 categorized_sample_indices=categorized_sample_indices,663 generators=generators,664 )665 return sampling_metadata666 667 def prepare_input_tensors(668 self,669 seq_group_metadata_list: List[SequenceGroupMetadata],670 ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata, SamplingMetadata,671 Set[LoRARequest], LoRAMapping, torch.Tensor]:672 if self.is_driver_worker:673 prefill_reqs = []674 decode_reqs = []675 for seq_group_meta in seq_group_metadata_list:676 if seq_group_meta.is_prompt:677 prefill_reqs.append(seq_group_meta)678 else:679 decode_reqs.append(seq_group_meta)680 681 # Prepare input tensors.682 (683 input_tokens,684 input_positions,685 prefill_attn_metadata,686 prompt_lens,687 subquery_lens,688 lora_index_mapping,689 lora_prompt_mapping,690 lora_requests,691 multi_modal_input,692 slot_mapping,693 ) = self._prepare_prompt(prefill_reqs)694 (695 decode_input_tokens,696 decode_input_positions,697 decode_attn_metadata,698 decode_lora_index_mapping,699 decode_lora_prompt_mapping,700 decode_lora_requests,701 decode_slot_mapping,702 ) = self._prepare_decode(decode_reqs)703 sampling_metadata = self._prepare_sample(seq_group_metadata_list,704 prompt_lens,705 subquery_lens)706 707 if not self.scheduler_config.chunked_prefill_enabled:708 assert (len(prefill_reqs) and len(decode_reqs)) == 0709 710 num_prefills = len(prompt_lens)711 num_prefill_tokens = len(input_tokens)712 num_decode_tokens = len(decode_input_tokens)713 714 # Coalesce tensors. Note that attn_metadata is currently not715 # coalesced for simplicity.716 input_tokens.extend(decode_input_tokens)717 input_positions.extend(decode_input_positions)718 slot_mapping.extend(decode_slot_mapping)719 lora_index_mapping.extend(decode_lora_index_mapping)720 lora_prompt_mapping.extend(decode_lora_prompt_mapping)721 lora_requests.update(decode_lora_requests)722 723 input_tokens = torch.tensor(input_tokens,724 dtype=torch.long,725 device=self.device)726 input_positions = torch.tensor(input_positions,727 dtype=torch.long,728 device=self.device)729 slot_mapping = torch.tensor(slot_mapping,730 dtype=torch.long,731 device=self.device)732 733 if self.lora_config:734 lora_mapping = LoRAMapping(735 lora_index_mapping,736 lora_prompt_mapping,737 )738 else:739 lora_mapping = None740 741 # Broadcast the metadata.742 # If batch contains both prefill and decode, it sends 2 broadcasts.743 # If it only contains 1 type, it triggers a single broadcast.744 if (prefill_attn_metadata is not None745 and decode_attn_metadata is not None):746 batch_type = BatchType.MIXED747 elif prefill_attn_metadata is not None:748 batch_type = BatchType.PREFILL749 else:750 batch_type = BatchType.DECODE751 752 metadata_dict = {753 "input_tokens": input_tokens,754 "input_positions": input_positions,755 "selected_token_indices":756 sampling_metadata.selected_token_indices,757 "lora_requests": lora_requests,758 "lora_mapping": lora_mapping,759 "multi_modal_input": multi_modal_input,760 "num_prefill_tokens": num_prefill_tokens,761 "num_decode_tokens": num_decode_tokens,762 "slot_mapping": slot_mapping,763 "num_prefills": num_prefills,764 "batch_type": batch_type,765 }766 if prefill_attn_metadata is not None:767 metadata_dict.update(prefill_attn_metadata.asdict_zerocopy())768 else:769 assert decode_attn_metadata is not None770 metadata_dict.update(decode_attn_metadata.asdict_zerocopy())771 broadcast_tensor_dict(metadata_dict, src=0)772 773 # Broadcast decode attn metadata for mixed batch type.774 # The additional broadcast costs 300us overhead on 4 A10 GPUs.775 # We can potentially reduce the overhead by coelescing tensors.776 if batch_type == BatchType.MIXED:777 assert decode_attn_metadata is not None778 metadata_dict = decode_attn_metadata.asdict_zerocopy()779 broadcast_tensor_dict(metadata_dict, src=0)780 else:781 metadata_dict = broadcast_tensor_dict(src=0)782 input_tokens = metadata_dict.pop("input_tokens")783 input_positions = metadata_dict.pop("input_positions")784 slot_mapping = metadata_dict.pop("slot_mapping")785 num_prefills = metadata_dict.pop("num_prefills")786 selected_token_indices = metadata_dict.pop(787 "selected_token_indices")788 lora_mapping = metadata_dict.pop("lora_mapping")789 lora_requests = metadata_dict.pop("lora_requests")790 multi_modal_input = metadata_dict.pop("multi_modal_input")791 num_prefill_tokens = metadata_dict.pop("num_prefill_tokens")792 num_decode_tokens = metadata_dict.pop("num_decode_tokens")793 batch_type = metadata_dict.pop("batch_type")794 795 # Create an attention metadata.796 prefill_attn_metadata = None797 decode_attn_metadata = None798 if batch_type == BatchType.PREFILL or batch_type == BatchType.MIXED:799 prefill_attn_metadata = self.attn_backend.make_metadata(800 **metadata_dict)801 else:802 decode_attn_metadata = self.attn_backend.make_metadata(803 **metadata_dict)804 sampling_metadata = SamplingMetadata(805 seq_groups=None,806 seq_data=None,807 prompt_lens=None,808 selected_token_indices=selected_token_indices,809 categorized_sample_indices=None,810 generators=None,811 perform_sampling=False,812 )813 814 # if it is a mixed batch, decode attn_metadata is broadcasted815 # separately.816 if batch_type == BatchType.MIXED:817 metadata_dict = broadcast_tensor_dict(src=0)818 decode_attn_metadata = self.attn_backend.make_metadata(819 **metadata_dict)820 821 attn_metadata = AttentionMetadata(822 num_prefills=num_prefills,823 slot_mapping=slot_mapping,824 num_prefill_tokens=num_prefill_tokens,825 num_decode_tokens=num_decode_tokens,826 prefill_metadata=prefill_attn_metadata,827 decode_metadata=decode_attn_metadata,828 kv_cache_dtype=self.kv_cache_dtype,829 )830 831 return (input_tokens, input_positions, attn_metadata,832 sampling_metadata, lora_requests, lora_mapping,833 multi_modal_input)834 835 @torch.inference_mode()836 def execute_model(837 self,838 seq_group_metadata_list: List[SequenceGroupMetadata],839 kv_caches: List[torch.Tensor],840 ) -> Optional[SamplerOutput]:841 (input_tokens, input_positions, attn_metadata, sampling_metadata,842 lora_requests, lora_mapping, multi_modal_input843 ) = self.prepare_input_tensors(seq_group_metadata_list)844 if self.lora_config:845 self.set_active_loras(lora_requests, lora_mapping)846 847 # Currently cuda graph is only supported by the decode phase.848 prefill_meta = attn_metadata.prefill_metadata849 decode_meta = attn_metadata.decode_metadata850 if prefill_meta is None and decode_meta.use_cuda_graph:851 graph_batch_size = input_tokens.shape[0]852 model_executable = self.graph_runners[graph_batch_size]853 else:854 model_executable = self.model855 execute_model_kwargs = {856 "input_ids": input_tokens,857 "positions": input_positions,858 "kv_caches": kv_caches,859 "attn_metadata": attn_metadata,860 }861 if self.vision_language_config:862 execute_model_kwargs.update({"image_input": multi_modal_input})863 hidden_states = model_executable(**execute_model_kwargs)864 865 # Compute the logits.866 logits = self.model.compute_logits(hidden_states, sampling_metadata)867 868 # Only perform sampling in the driver worker.869 if not sampling_metadata.perform_sampling:870 return None871 872 # Sample the next token.873 output = self.model.sample(874 logits=logits,875 sampling_metadata=sampling_metadata,876 )877 return output878 879 @torch.inference_mode()880 def profile_run(self) -> None:881 # Enable top-k sampling to reflect the accurate memory usage.882 sampling_params = SamplingParams(top_p=0.99, top_k=self.vocab_size - 1)883 max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens884 max_num_seqs = self.scheduler_config.max_num_seqs885 886 # This represents the maximum number of different requests887 # that will have unique loras, an therefore the max amount of memory888 # consumption create dummy lora request copies from the lora request889 # passed in, which contains a lora from the lora warmup path.890 dummy_lora_requests = []891 dummy_lora_requests_per_seq = []892 if self.lora_config:893 for idx in range(self.lora_config.max_loras):894 lora_id = idx + 1895 dummy_lora_request = LoRARequest(896 lora_name=f"warmup_{lora_id}",897 lora_int_id=lora_id,898 lora_local_path="/not/a/real/path",899 )900 self.lora_manager.add_dummy_lora(dummy_lora_request,901 rank=LORA_WARMUP_RANK)902 dummy_lora_requests.append(dummy_lora_request)903 dummy_lora_requests_per_seq = [904 dummy_lora_requests[idx % len(dummy_lora_requests)]905 for idx in range(max_num_seqs)906 ]907 908 # Profile memory usage with max_num_sequences sequences and the total909 # number of tokens equal to max_num_batched_tokens.910 seqs: List[SequenceGroupMetadata] = []911 # Additional GPU memory may be needed for vision encoding, which needs912 # to be accounted for when calculating the GPU blocks for913 # vLLM blocker manager.914 # To exercise the worst scenario for GPU memory consumption,915 # the number of seqs (batch_size) is chosen to maximize the number916 # of images processed.917 if self.vision_language_config:918 max_num_seqs = min(919 max_num_seqs,920 int(max_num_batched_tokens /921 self.vision_language_config.image_feature_size))922 for group_id in range(max_num_seqs):923 seq_len = (max_num_batched_tokens // max_num_seqs +924 (group_id < max_num_batched_tokens % max_num_seqs))925 seq_data, fake_multi_modal_input = _prepare_fake_inputs(926 seq_len, self.vision_language_config)927 seq = SequenceGroupMetadata(928 request_id=str(group_id),929 is_prompt=True,930 seq_data={group_id: seq_data},931 sampling_params=sampling_params,932 block_tables=None,933 lora_request=dummy_lora_requests_per_seq[group_id]934 if dummy_lora_requests_per_seq else None,935 multi_modal_data=fake_multi_modal_input,936 )937 seqs.append(seq)938 939 # Run the model with the dummy inputs.940 num_layers = self.model_config.get_num_layers(self.parallel_config)941 kv_caches = [None] * num_layers942 self.execute_model(seqs, kv_caches)943 torch.cuda.synchronize()944 return945 946 def remove_all_loras(self) -> bool:947 if not self.lora_manager:948 raise RuntimeError("LoRA is not enabled.")949 return self.lora_manager.remove_all_loras()950 951 def set_active_loras(self, lora_requests: Set[LoRARequest],952 lora_mapping: LoRAMapping) -> None:953 if not self.lora_manager:954 raise RuntimeError("LoRA is not enabled.")955 self.lora_manager.set_active_loras(lora_requests, lora_mapping)956 957 def add_lora(self, lora_request: LoRARequest) -> bool:958 if not self.lora_manager:959 raise RuntimeError("LoRA is not enabled.")960 return self.lora_manager.add_lora(lora_request)961 962 def remove_lora(self, lora_id: int) -> bool:963 if not self.lora_manager:964 raise RuntimeError("LoRA is not enabled.")965 return self.lora_manager.remove_lora(lora_id)966 967 def list_loras(self) -> Set[int]:968 if not self.lora_manager:969 raise RuntimeError("LoRA is not enabled.")970 return self.lora_manager.list_loras()971 972 @torch.inference_mode()973 def capture_model(self, kv_caches: List[torch.Tensor]) -> None:974 """Cuda graph capture a model.975 976 Note that CUDA graph's performance gain is negligible if number977 of batched tokens are larger than 200. And since CUDA graph978 requires fixed sized tensors, supporting large/variable batch979 size requires high GPU memory overhead. Thus, vLLM only captures980 decoding requests. Mixed batch (chunked prefill + decoding) or981 prefill requests are not captured.982 983 Since it is used for decoding-only, it assumes there's only 1 token984 per sequence in the batch.985 """986 # NOTE(woosuk): This is a hack to ensure that the NCCL backend is never987 # deleted before the CUDA graphs.988 self.pynccl_backend = pynccl_utils.get_nccl_backend()989 990 assert not self.model_config.enforce_eager991 logger.info("Capturing the model for CUDA graphs. This may lead to "992 "unexpected consequences if the model is not static. To "993 "run the model in eager mode, set 'enforce_eager=True' or "994 "use '--enforce-eager' in the CLI.")995 logger.info("CUDA graphs can take additional 1~3 GiB memory per GPU. "996 "If you are running out of memory, consider decreasing "997 "`gpu_memory_utilization` or enforcing eager mode. "998 "You can also reduce the `max_num_seqs` as needed "999 "to decrease memory usage.")1000 start_time = time.perf_counter()1001 1002 # Prepare dummy inputs. These will be reused for all batch sizes.1003 max_batch_size = max(_BATCH_SIZES_TO_CAPTURE)1004 input_tokens = torch.zeros(max_batch_size, dtype=torch.long).cuda()1005 input_positions = torch.zeros(max_batch_size, dtype=torch.long).cuda()1006 slot_mapping = torch.empty(max_batch_size, dtype=torch.long).cuda()1007 slot_mapping.fill_(_PAD_SLOT_ID)1008 context_lens = torch.ones(max_batch_size, dtype=torch.int32).cuda()1009 block_tables = torch.from_numpy(self.graph_block_tables).cuda()1010 1011 graph_batch_size = _get_graph_batch_size(1012 self.scheduler_config.max_num_seqs)1013 batch_size_capture_list = [1014 bs for bs in _BATCH_SIZES_TO_CAPTURE if bs <= graph_batch_size1015 ]1016 1017 # NOTE(woosuk): There are 3 backends for all-reduce: custom all-reduce1018 # kernel, pynccl, and PyTorch NCCL. When using CUDA graph, we use1019 # either custom all-reduce kernel or pynccl. When not using CUDA1020 # graph, we use either custom all-reduce kernel or PyTorch NCCL.1021 # We always prioritize using custom all-reduce kernel but fall back1022 # to PyTorch or pynccl if it is disabled or not supported.1023 with custom_all_reduce.capture():1024 # NOTE: Capturing the largest batch size first may help reduce the1025 # memory usage of CUDA graph.1026 for batch_size in reversed(batch_size_capture_list):1027 # Create dummy attn_metadata.1028 decode_metadata = self.attn_backend.make_metadata(1029 is_prompt=False,1030 prompt_lens=None,1031 prompt_lens_tensor=None,1032 max_subquery_len=None,1033 max_context_len=self.max_context_len_to_capture,1034 max_prompt_len=None,1035 subquery_start_loc=None,1036 seq_start_loc=None,1037 context_lens=context_lens[:batch_size],1038 block_tables=block_tables[:batch_size],1039 use_cuda_graph=True,1040 )1041 attn_metadata = AttentionMetadata(1042 num_prefills=0,1043 num_prefill_tokens=0,1044 num_decode_tokens=batch_size,1045 slot_mapping=slot_mapping[:batch_size],1046 prefill_metadata=None,1047 decode_metadata=decode_metadata,1048 kv_cache_dtype=self.kv_cache_dtype,1049 )1050 1051 if self.lora_config:1052 lora_mapping = LoRAMapping(1053 [0] * batch_size,1054 [0] * batch_size,1055 )1056 self.set_active_loras(set(), lora_mapping)1057 1058 graph_runner = CUDAGraphRunner(self.model)1059 graph_runner.capture(1060 input_tokens[:batch_size],1061 input_positions[:batch_size],1062 kv_caches,1063 attn_metadata,1064 memory_pool=self.graph_memory_pool,1065 )1066 self.graph_memory_pool = graph_runner.graph.pool()1067 self.graph_runners[batch_size] = graph_runner1068 1069 end_time = time.perf_counter()1070 elapsed_time = end_time - start_time1071 # This usually takes < 10 seconds.1072 logger.info(f"Graph capturing finished in {elapsed_time:.0f} secs.")1073 1074 def __del__(self) -> None:1075 # Delete the CUDA graphs before deleting the pynccl communicator.1076 # NOTE(woosuk): This is necessary because otherwise deadlocks can1077 # happen.1078 # FIXME(woosuk): This is a bit hacky. Find a more robust solution.1079 # TODO(youkaichao): when we get enough user feedback that pynccl is1080 # more stable than cupy, we can remove this, e.g. in v0.4.1.1081 self.graph_runners.clear()1082 self.pynccl_backend = None1083 1084 @property1085 def vocab_size(self) -> int:1086 return self.model_config.get_vocab_size()1087 1088 1089class CUDAGraphRunner:1090 1091 def __init__(self, model: nn.Module):1092 self.model = model1093 self.input_buffers: Dict[str, torch.Tensor] = {}1094 self.output_buffers: Dict[str, torch.Tensor] = {}1095 1096 self._graph: Optional[torch.cuda.CUDAGraph] = None1097 1098 @property1099 def graph(self):1100 assert self._graph is not None1101 return self._graph1102 1103 def capture(1104 self,1105 input_ids: torch.Tensor,1106 positions: torch.Tensor,1107 kv_caches: List[torch.Tensor],1108 attn_metadata: AttentionMetadata,1109 memory_pool,1110 **kwargs,1111 ) -> None:1112 assert self._graph is None1113 # Run the model once without capturing the graph.1114 # This is to make sure that the captured graph does not include the1115 # kernel launches for initial benchmarking (e.g., Triton autotune).1116 with _maybe_pynccl():1117 self.model(1118 input_ids,1119 positions,1120 kv_caches,1121 attn_metadata,1122 **kwargs,1123 )1124 torch.cuda.synchronize()1125 1126 # Capture the graph.1127 # NOTE(woosuk): Python 3.8 does not support multi-line with statements.1128 # https://stackoverflow.com/questions/31039022/python-multi-line-with-statement1129 self._graph = torch.cuda.CUDAGraph()1130 with torch.cuda.graph(self._graph, pool=memory_pool): # noqa: SIM1171131 with _maybe_pynccl():1132 hidden_states = self.model(1133 input_ids,1134 positions,1135 kv_caches,1136 attn_metadata,1137 **kwargs,1138 )1139 torch.cuda.synchronize()1140 1141 # Save the input and output buffers.1142 self.input_buffers = {1143 "input_ids": input_ids,1144 "positions": positions,1145 "kv_caches": kv_caches,1146 "slot_mapping": attn_metadata.slot_mapping,1147 "context_lens": attn_metadata.decode_metadata.context_lens,1148 "block_tables": attn_metadata.decode_metadata.block_tables,1149 }1150 self.output_buffers = {"hidden_states": hidden_states}1151 return1152 1153 def forward(1154 self,1155 input_ids: torch.Tensor,1156 positions: torch.Tensor,1157 kv_caches: List[torch.Tensor],1158 attn_metadata: AttentionMetadata,1159 **kwargs,1160 ) -> torch.Tensor:1161 # KV caches are fixed tensors, so we don't need to copy them.1162 del kv_caches1163 1164 # Copy the input tensors to the input buffers.1165 self.input_buffers["input_ids"].copy_(input_ids, non_blocking=True)1166 self.input_buffers["positions"].copy_(positions, non_blocking=True)1167 self.input_buffers["slot_mapping"].copy_(attn_metadata.slot_mapping,1168 non_blocking=True)1169 self.input_buffers["context_lens"].copy_(1170 attn_metadata.decode_metadata.context_lens, non_blocking=True)1171 self.input_buffers["block_tables"].copy_(1172 attn_metadata.decode_metadata.block_tables, non_blocking=True)1173 # Run the graph.1174 self.graph.replay()1175 1176 # Return the output tensor.1177 return self.output_buffers["hidden_states"]1178 1179 def __call__(self, *args, **kwargs):1180 return self.forward(*args, **kwargs)1181 1182 1183@contextlib.contextmanager1184def _maybe_pynccl():1185 if pynccl_utils.is_initialized(1186 ) and not custom_all_reduce.is_initialized():1187 with with_pynccl_for_all_reduce():1188 yield1189 else:1190 yield1191 1192 1193def _get_graph_batch_size(batch_size: int) -> int:1194 """Returns the padded batch size given actual batch size.1195 1196 Batch sizes are 1, 2, 4, _BATCH_SIZE_ALIGNMENT,1197 2*_BATCH_SIZE_ALIGNMENT, 3*_BATCH_SIZE_ALIGNMENT...1198 """1199 if batch_size <= 2:1200 return batch_size