aphrodite_engine.py 59 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468
  1. import os
  2. import time
  3. from contextlib import contextmanager
  4. from typing import TYPE_CHECKING, Any, ClassVar, Dict, Iterable, List, Optional
  5. from typing import Sequence as GenericSequence
  6. from typing import Tuple, Type, TypeVar, Union
  7. from loguru import logger
  8. from transformers import PreTrainedTokenizer
  9. from typing_extensions import assert_never
  10. from aphrodite.common.config import (CacheConfig, DecodingConfig, DeviceConfig,
  11. EngineConfig, LoadConfig, LoRAConfig,
  12. ModelConfig, ParallelConfig,
  13. PromptAdapterConfig, SchedulerConfig,
  14. SpeculativeConfig)
  15. from aphrodite.common.logger import setup_logger
  16. from aphrodite.common.outputs import (EmbeddingRequestOutput, RequestOutput,
  17. RequestOutputFactory)
  18. from aphrodite.common.pooling_params import PoolingParams
  19. from aphrodite.common.sampling_params import SamplingParams
  20. from aphrodite.common.sequence import (EmbeddingSequenceGroupOutput,
  21. ExecuteModelRequest, PoolerOutput,
  22. SamplerOutput, Sequence, SequenceGroup,
  23. SequenceGroupMetadata, SequenceStatus)
  24. from aphrodite.common.utils import Counter, Device
  25. from aphrodite.engine.args_tools import EngineArgs
  26. from aphrodite.engine.metrics_types import StatLoggerBase, Stats
  27. from aphrodite.engine.output_processor.interfaces import (
  28. SequenceGroupOutputProcessor)
  29. from aphrodite.engine.output_processor.stop_checker import StopChecker
  30. from aphrodite.engine.output_processor.util import (
  31. create_output_by_sequence_group)
  32. from aphrodite.executor.executor_base import ExecutorBase
  33. from aphrodite.executor.ray_utils import initialize_ray_cluster
  34. from aphrodite.inputs import (INPUT_REGISTRY, EncoderDecoderLLMInputs,
  35. InputRegistry, LLMInputs, PromptInputs,
  36. SingletonPromptInputs)
  37. from aphrodite.inputs.parse import is_explicit_encoder_decoder_prompt
  38. from aphrodite.lora.request import LoRARequest
  39. from aphrodite.multimodal import MultiModalDataDict
  40. from aphrodite.processing.scheduler import (ScheduledSequenceGroup, Scheduler,
  41. SchedulerOutputs)
  42. from aphrodite.prompt_adapter.request import PromptAdapterRequest
  43. from aphrodite.transformers_utils.config import try_get_generation_config
  44. from aphrodite.transformers_utils.detokenizer import Detokenizer
  45. from aphrodite.transformers_utils.tokenizer_group import (
  46. BaseTokenizerGroup, init_tokenizer_from_configs)
  47. from aphrodite.version import __version__ as APHRODITE_VERSION
  48. _LOCAL_LOGGING_INTERVAL_SEC = 5
  49. APHRODITE_USE_RAY_SPMD_WORKER = bool(
  50. os.getenv("APHRODITE_USE_RAY_SPMD_WORKER", 0))
  51. def _load_generation_config_dict(model_config: ModelConfig) -> Dict[str, Any]:
  52. config = try_get_generation_config(
  53. model_config.model,
  54. trust_remote_code=model_config.trust_remote_code,
  55. revision=model_config.revision,
  56. )
  57. if config is None:
  58. return {}
  59. return config.to_diff_dict()
  60. _O = TypeVar("_O", RequestOutput, EmbeddingRequestOutput)
  61. PromptComponents = Tuple[Optional[str], List[int],
  62. Optional[MultiModalDataDict]]
  63. DecoderPromptComponents = Tuple[Optional[str], Optional[List[int]],
  64. Optional[MultiModalDataDict]]
  65. class AphroditeEngine:
  66. """An LLM engine that receives requests and generates texts.
  67. This is the main class for the Aphrodite engine. It receives requests
  68. from clients and generates texts from the LLM. It includes a tokenizer, a
  69. language model (possibly distributed across multiple GPUs), and GPU memory
  70. space allocated for intermediate states (aka KV cache). This class utilizes
  71. iteration-level scheduling and efficient memory management to maximize the
  72. serving throughput.
  73. The `LLM` class wraps this class for offline batched inference and the
  74. `AsyncAphrodite` class wraps this class for online serving.
  75. NOTE: The config arguments are derived from the `EngineArgs` class. For the
  76. comprehensive list of arguments, see `EngineArgs`.
  77. Args:
  78. model_config: The configuration related to the LLM model.
  79. cache_config: The configuration related to the KV cache memory
  80. management.
  81. parallel_config: The configuration related to distributed execution.
  82. scheduler_config: The configuration related to the request scheduler.
  83. device_config: The configuration related to the device.
  84. lora_config (Optional): The configuration related to serving multi-LoRA.
  85. speculative_config (Optional): The configuration related to speculative
  86. decoding.
  87. executor_class: The model executor class for managing distributed
  88. execution.
  89. prompt_adapter_config (Optional): The configuration related to serving
  90. prompt adapters.
  91. log_stats: Whether to log statistics.
  92. """
  93. DO_VALIDATE_OUTPUT: ClassVar[bool] = False
  94. """A flag to toggle whether to validate the type of request output."""
  95. @classmethod
  96. @contextmanager
  97. def enable_output_validation(cls):
  98. cls.DO_VALIDATE_OUTPUT = True
  99. yield
  100. cls.DO_VALIDATE_OUTPUT = False
  101. @classmethod
  102. def validate_output(
  103. cls,
  104. output: object,
  105. output_type: Type[_O],
  106. ) -> _O:
  107. do_validate = cls.DO_VALIDATE_OUTPUT
  108. if ((TYPE_CHECKING or do_validate)
  109. and not isinstance(output, output_type)):
  110. raise TypeError(f"Expected output of type {output_type}, "
  111. f"but found type {type(output)}")
  112. return output
  113. @classmethod
  114. def validate_outputs(
  115. cls,
  116. outputs: GenericSequence[object],
  117. output_type: Type[_O],
  118. ) -> List[_O]:
  119. do_validate = cls.DO_VALIDATE_OUTPUT
  120. outputs_: List[_O]
  121. if TYPE_CHECKING or do_validate:
  122. outputs_ = []
  123. for output in outputs:
  124. if not isinstance(output, output_type):
  125. raise TypeError(f"Expected output of type {output_type}, "
  126. f"but found type {type(output)}")
  127. outputs_.append(output)
  128. else:
  129. outputs_ = outputs
  130. return outputs_
  131. tokenizer: Optional[BaseTokenizerGroup]
  132. def __init__(
  133. self,
  134. model_config: ModelConfig,
  135. cache_config: CacheConfig,
  136. parallel_config: ParallelConfig,
  137. scheduler_config: SchedulerConfig,
  138. device_config: DeviceConfig,
  139. load_config: LoadConfig,
  140. lora_config: Optional[LoRAConfig],
  141. speculative_config: Optional[SpeculativeConfig],
  142. decoding_config: Optional[DecodingConfig],
  143. prompt_adapter_config: Optional[PromptAdapterConfig],
  144. executor_class: Type[ExecutorBase],
  145. log_stats: bool,
  146. stat_loggers: Optional[Dict[str, StatLoggerBase]] = None,
  147. input_registry: InputRegistry = INPUT_REGISTRY,
  148. ) -> None:
  149. try:
  150. import aphrodite.commit_id
  151. commit_id = True
  152. except ImportError:
  153. commit_id = False
  154. config_dict = {
  155. "Model": model_config.model,
  156. "Speculative Config": speculative_config,
  157. "DataType": model_config.dtype,
  158. "Model Load Format": load_config.load_format,
  159. "Tensor Parallel Size": parallel_config.tensor_parallel_size,
  160. "Pipeline Parallel Size": parallel_config.pipeline_parallel_size,
  161. "Disable Custom All-Reduce":
  162. parallel_config.disable_custom_all_reduce,
  163. "Quantization Format": model_config.quantization,
  164. "Context Length": model_config.max_model_len,
  165. "Enforce Eager Mode": model_config.enforce_eager,
  166. "Prefix Caching": cache_config.enable_prefix_caching,
  167. "KV Cache DataType": cache_config.cache_dtype,
  168. "Device": device_config.device,
  169. "Rope Scaling": model_config.rope_scaling,
  170. "Guided Decoding Backend": decoding_config
  171. }
  172. logger.info("-" * 85)
  173. if not commit_id:
  174. logger.info(
  175. f"Initializing Aphrodite Engine (v{APHRODITE_VERSION}) "
  176. "with the following config:")
  177. else:
  178. logger.info(f"Initializing Aphrodite Engine (v{APHRODITE_VERSION} "
  179. f"commit {aphrodite.__short_commit__}) with the "
  180. "following config:")
  181. for key, value in config_dict.items():
  182. if value is not None and not ((key == "Model Load Format" or key ==\
  183. "KV Cache DataType") and value == \
  184. "auto"):
  185. logger.info(f"{key} = {value!r}")
  186. logger.info("-" * 85)
  187. # TODO: Print more configs in debug mode.
  188. from aphrodite.plugins import load_general_plugins
  189. load_general_plugins()
  190. self.model_config = model_config
  191. self.cache_config = cache_config
  192. self.lora_config = lora_config
  193. self.parallel_config = parallel_config
  194. self.scheduler_config = scheduler_config
  195. self.device_config = device_config
  196. self.speculative_config = speculative_config
  197. self.load_config = load_config
  198. self.decoding_config = decoding_config or DecodingConfig()
  199. self.prompt_adapter_config = prompt_adapter_config
  200. self.log_stats = log_stats
  201. if not self.model_config.skip_tokenizer_init:
  202. self.tokenizer = self._init_tokenizer()
  203. self.detokenizer = Detokenizer(self.tokenizer)
  204. tokenizer_group = self.get_tokenizer_group()
  205. else:
  206. self.tokenizer = None
  207. self.detokenizer = None
  208. tokenizer_group = None
  209. # Ensure that the function doesn't contain a reference to self,
  210. # to avoid engine GC issues
  211. def get_tokenizer_for_seq(sequence: Sequence) -> PreTrainedTokenizer:
  212. assert tokenizer_group, ("tokenizer_group cannot be None, "
  213. "make sure skip_tokenizer_init is False")
  214. return tokenizer_group.get_lora_tokenizer(sequence.lora_request)
  215. self.seq_counter = Counter()
  216. self.generation_config_fields = _load_generation_config_dict(
  217. model_config)
  218. self.input_registry = input_registry
  219. self.input_processor = input_registry.create_input_processor(
  220. model_config)
  221. self.model_executor = executor_class(
  222. model_config=model_config,
  223. cache_config=cache_config,
  224. parallel_config=parallel_config,
  225. scheduler_config=scheduler_config,
  226. device_config=device_config,
  227. lora_config=lora_config,
  228. speculative_config=speculative_config,
  229. load_config=load_config,
  230. prompt_adapter_config=prompt_adapter_config,
  231. )
  232. if not self.model_config.embedding_mode:
  233. self._initialize_kv_caches()
  234. if self.tokenizer:
  235. # Ping the tokenizer to ensure liveness if it runs in a
  236. # different process.
  237. self.tokenizer.ping()
  238. # Create the scheduler.
  239. # NOTE: the cache_config here have been updated with the numbers of
  240. # GPU and CPU blocks, which are profiled in the distributed executor.
  241. self.scheduler = [
  242. Scheduler(scheduler_config, cache_config, lora_config,
  243. parallel_config.pipeline_parallel_size)
  244. for _ in range(parallel_config.pipeline_parallel_size)
  245. ]
  246. # Metric Logging.
  247. if self.log_stats:
  248. if stat_loggers is not None:
  249. self.stat_loggers = stat_loggers
  250. else:
  251. # Lazy import for prometheus multiprocessing.
  252. # We need to set PROMETHEUS_MULTIPROC_DIR environment variable
  253. # before prometheus_client is imported.
  254. # See https://prometheus.github.io/client_python/multiprocess/
  255. from aphrodite.engine.metrics import (LoggingStatLogger,
  256. PrometheusStatLogger)
  257. self.stat_loggers = {
  258. "logging":
  259. LoggingStatLogger(
  260. local_interval=_LOCAL_LOGGING_INTERVAL_SEC),
  261. "prometheus":
  262. PrometheusStatLogger(
  263. local_interval=_LOCAL_LOGGING_INTERVAL_SEC,
  264. labels=dict(model_name=model_config.served_model_name),
  265. max_model_len=self.model_config.max_model_len),
  266. }
  267. self.stat_loggers["prometheus"].info("cache_config",
  268. self.cache_config)
  269. # Create sequence output processor, e.g. for beam search or
  270. # speculative decoding.
  271. self.output_processor = (
  272. SequenceGroupOutputProcessor.create_output_processor(
  273. self.scheduler_config,
  274. self.detokenizer,
  275. self.scheduler,
  276. self.seq_counter,
  277. get_tokenizer_for_seq,
  278. stop_checker=StopChecker(
  279. self.scheduler_config.max_model_len,
  280. get_tokenizer_for_seq,
  281. ),
  282. ))
  283. def _initialize_kv_caches(self) -> None:
  284. """Initialize the KV cache in the worker(s).
  285. The workers will determine the number of blocks in both the GPU cache
  286. and the swap CPU cache.
  287. """
  288. num_gpu_blocks, num_cpu_blocks = (
  289. self.model_executor.determine_num_available_blocks())
  290. if self.cache_config.num_gpu_blocks_override is not None:
  291. num_gpu_blocks_override = self.cache_config.num_gpu_blocks_override
  292. logger.info(f"Overriding {num_gpu_blocks=} with "
  293. f"{num_gpu_blocks_override=}")
  294. num_gpu_blocks = num_gpu_blocks_override
  295. self.cache_config.num_gpu_blocks = num_gpu_blocks
  296. self.cache_config.num_cpu_blocks = num_cpu_blocks
  297. self.model_executor.initialize_cache(num_gpu_blocks, num_cpu_blocks)
  298. @classmethod
  299. def _get_executor_cls(cls,
  300. engine_config: EngineConfig) -> Type[ExecutorBase]:
  301. distributed_executor_backend = (
  302. engine_config.parallel_config.distributed_executor_backend)
  303. # Initialize the cluster and specify the executor class.
  304. if isinstance(distributed_executor_backend, type):
  305. if not issubclass(distributed_executor_backend, ExecutorBase):
  306. raise TypeError(
  307. "distributed_executor_backend must be a subclass of "
  308. f"ExecutorBase. Got {distributed_executor_backend}.")
  309. if distributed_executor_backend.uses_ray: # type: ignore
  310. initialize_ray_cluster(engine_config.parallel_config)
  311. executor_class = distributed_executor_backend
  312. elif engine_config.device_config.device_type == "neuron":
  313. from aphrodite.executor.neuron_executor import NeuronExecutor
  314. executor_class = NeuronExecutor
  315. elif engine_config.device_config.device_type == "tpu":
  316. if distributed_executor_backend == "ray":
  317. initialize_ray_cluster(engine_config.parallel_config)
  318. from aphrodite.executor.ray_tpu_executor import RayTPUExecutor
  319. executor_class = RayTPUExecutor
  320. else:
  321. assert distributed_executor_backend is None
  322. from aphrodite.executor.tpu_executor import TPUExecutor
  323. executor_class = TPUExecutor
  324. elif engine_config.device_config.device_type == "cpu":
  325. from aphrodite.executor.cpu_executor import CPUExecutor
  326. executor_class = CPUExecutor
  327. elif engine_config.device_config.device_type == "openvino":
  328. from aphrodite.executor.openvino_executor import OpenVINOExecutor
  329. executor_class = OpenVINOExecutor
  330. elif engine_config.device_config.device_type == "xpu":
  331. if distributed_executor_backend == "ray":
  332. initialize_ray_cluster(engine_config.parallel_config)
  333. from aphrodite.executor.ray_xpu_executor import RayXPUExecutor
  334. executor_class = RayXPUExecutor
  335. else:
  336. from aphrodite.executor.xpu_executor import XPUExecutor
  337. executor_class = XPUExecutor
  338. elif distributed_executor_backend == "ray":
  339. initialize_ray_cluster(engine_config.parallel_config)
  340. from aphrodite.executor.ray_gpu_executor import RayGPUExecutor
  341. executor_class = RayGPUExecutor
  342. elif distributed_executor_backend == "mp":
  343. from aphrodite.executor.multiproc_gpu_executor import (
  344. MultiprocessingGPUExecutor)
  345. assert not APHRODITE_USE_RAY_SPMD_WORKER, (
  346. "multiprocessing distributed executor backend does not "
  347. "support APHRODITE_USE_RAY_SPMD_WORKER=1")
  348. executor_class = MultiprocessingGPUExecutor
  349. else:
  350. from aphrodite.executor.gpu_executor import GPUExecutor
  351. executor_class = GPUExecutor
  352. return executor_class
  353. @classmethod
  354. def from_engine_args(
  355. cls,
  356. engine_args: EngineArgs,
  357. stat_loggers: Optional[Dict[str, StatLoggerBase]] = None,
  358. ) -> "AphroditeEngine":
  359. """Creates an LLM engine from the engine arguments."""
  360. # Create the engine configs.
  361. engine_config = engine_args.create_engine_config()
  362. executor_class = cls._get_executor_cls(engine_config)
  363. # Create the LLM engine.
  364. engine = cls(
  365. **engine_config.to_dict(),
  366. executor_class=executor_class,
  367. log_stats=not engine_args.disable_log_stats,
  368. stat_loggers=stat_loggers,
  369. )
  370. return engine
  371. def __reduce__(self):
  372. # This is to ensure that the AphroditeEngine is not referenced in
  373. # the closure used to initialize Ray worker actors
  374. raise RuntimeError("AphroditeEngine should not be pickled!")
  375. def __del__(self):
  376. # Shutdown the model executor when engine is garbage collected.
  377. # Use getattr since __init__ can fail before the field is set
  378. if model_executor := getattr(self, "model_executor", None):
  379. model_executor.shutdown()
  380. MISSING_TOKENIZER_GROUP_MSG = ("Unable to get tokenizer because "
  381. "skip_tokenizer_init is True")
  382. def get_tokenizer_group(
  383. self,
  384. fail_msg: str = MISSING_TOKENIZER_GROUP_MSG) -> BaseTokenizerGroup:
  385. if self.tokenizer is None:
  386. raise ValueError(fail_msg)
  387. return self.tokenizer
  388. def get_tokenizer(
  389. self,
  390. lora_request: Optional[LoRARequest] = None
  391. ) -> "PreTrainedTokenizer":
  392. return self.get_tokenizer_group().get_lora_tokenizer(lora_request)
  393. def _init_tokenizer(self) -> BaseTokenizerGroup:
  394. return init_tokenizer_from_configs(
  395. model_config=self.model_config,
  396. scheduler_config=self.scheduler_config,
  397. parallel_config=self.parallel_config,
  398. enable_lora=bool(self.lora_config))
  399. def _verify_args(self) -> None:
  400. self.model_config.verify_with_parallel_config(self.parallel_config)
  401. self.cache_config.verify_with_parallel_config(self.parallel_config)
  402. if self.lora_config:
  403. self.lora_config.verify_with_model_config(self.model_config)
  404. self.lora_config.verify_with_scheduler_config(
  405. self.scheduler_config)
  406. if self.prompt_adapter_config:
  407. self.prompt_adapter_config.verify_with_model_config(
  408. self.model_config)
  409. def _get_bos_token_id(self,
  410. lora_request: Optional[LoRARequest] = None
  411. ) -> Optional[int]:
  412. if self.tokenizer is None:
  413. logger.warning("Using None for BOS token id because tokenizer "
  414. "is not initialized")
  415. return None
  416. return self.tokenizer.get_lora_tokenizer(lora_request).bos_token_id
  417. def _get_eos_token_id(self,
  418. lora_request: Optional[LoRARequest] = None
  419. ) -> Optional[int]:
  420. if self.tokenizer is None:
  421. logger.warning("Using None for EOS token id because tokenizer "
  422. "is not initialized")
  423. return None
  424. return self.tokenizer.get_lora_tokenizer(lora_request).eos_token_id
  425. def _get_decoder_start_token_id(self) -> Optional[int]:
  426. '''
  427. Obtain the decoder start token id employed by an encoder/decoder
  428. model. Returns None for non-encoder/decoder models or if the
  429. model config is unavailable.
  430. '''
  431. if not self.is_encoder_decoder_model():
  432. logger.warning("Using None for decoder start token id because "
  433. "this is not an encoder/decoder model.")
  434. return None
  435. if (self.model_config is None or self.model_config.hf_config is None):
  436. logger.warning("Using None for decoder start token id because "
  437. "model config is not available.")
  438. return None
  439. dec_start_token_id = getattr(self.model_config.hf_config,
  440. 'decoder_start_token_id', None)
  441. if dec_start_token_id is None:
  442. logger.warning("Falling back on <BOS> for decoder start token id "
  443. "because decoder start token id is not available.")
  444. dec_start_token_id = self._get_bos_token_id()
  445. return dec_start_token_id
  446. def _add_processed_request(
  447. self,
  448. request_id: str,
  449. processed_inputs: Union[LLMInputs, EncoderDecoderLLMInputs],
  450. params: Union[SamplingParams, PoolingParams],
  451. arrival_time: float,
  452. lora_request: Optional[LoRARequest],
  453. prompt_adapter_request: Optional[PromptAdapterRequest] = None,
  454. ) -> None:
  455. # Create the sequences.
  456. block_size = self.cache_config.block_size
  457. seq_id = next(self.seq_counter)
  458. eos_token_id = self._get_eos_token_id(lora_request)
  459. seq = Sequence(seq_id, processed_inputs, block_size, eos_token_id,
  460. lora_request, prompt_adapter_request)
  461. encoder_seq = None
  462. if 'encoder_prompt_token_ids' in processed_inputs:
  463. encoder_seq = Sequence(seq_id,
  464. processed_inputs,
  465. block_size,
  466. eos_token_id,
  467. lora_request,
  468. prompt_adapter_request,
  469. from_decoder_prompt=False)
  470. # Create a SequenceGroup based on SamplingParams or PoolingParams
  471. if isinstance(params, SamplingParams):
  472. seq_group = self._create_sequence_group_with_sampling(
  473. request_id,
  474. seq,
  475. params,
  476. arrival_time=arrival_time,
  477. lora_request=lora_request,
  478. prompt_adapter_request=prompt_adapter_request,
  479. encoder_seq=encoder_seq,
  480. )
  481. elif isinstance(params, PoolingParams):
  482. seq_group = self._create_sequence_group_with_pooling(
  483. request_id,
  484. seq,
  485. params,
  486. arrival_time=arrival_time,
  487. lora_request=lora_request,
  488. prompt_adapter_request=prompt_adapter_request,
  489. encoder_seq=encoder_seq,
  490. )
  491. else:
  492. raise ValueError(
  493. "Either SamplingParams or PoolingParams must be provided.")
  494. # Add the sequence group to the scheduler with least unfinished seqs.
  495. costs = [
  496. scheduler.get_num_unfinished_seq_groups()
  497. for scheduler in self.scheduler
  498. ]
  499. min_cost_scheduler = self.scheduler[costs.index(min(costs))]
  500. min_cost_scheduler.add_seq_group(seq_group)
  501. def stop_remote_worker_execution_loop(self) -> None:
  502. self.model_executor.stop_remote_worker_execution_loop()
  503. _LLMInputComponentsType = Tuple[str, List[int]]
  504. def _prepare_decoder_input_ids_for_generation(
  505. self,
  506. decoder_input_ids: Optional[List[int]],
  507. ) -> List[int]:
  508. """
  509. Prepares `decoder_input_ids` for generation with encoder-decoder models.
  510. Based on
  511. https://github.com/huggingface/transformers/blob/
  512. 4037a2b5b1278736e566aec12e169100275545ea/
  513. src/transformers/generation/utils.py
  514. specifically GenerationMixin._prepare_decoder_input_ids_for_generation()
  515. Arguments:
  516. * decoder_input_ids: input token ids to preprocess
  517. Returns:
  518. * Processed token list
  519. """
  520. decoder_start_token_id = self._get_decoder_start_token_id()
  521. assert decoder_start_token_id is not None
  522. if decoder_input_ids is None:
  523. # no decoder prompt input ->
  524. # use decoder_start_token_id as decoder_input_ids
  525. decoder_input_ids = self._get_default_enc_dec_decoder_prompt()
  526. if (len(decoder_input_ids) == 0
  527. or decoder_input_ids[0] != decoder_start_token_id):
  528. decoder_input_ids = [decoder_start_token_id] + decoder_input_ids
  529. return decoder_input_ids
  530. def _tokenize_prompt(
  531. self,
  532. prompt: str,
  533. request_id: str,
  534. lora_request: Optional[LoRARequest],
  535. ) -> List[int]:
  536. '''
  537. Wrapper around application of the model's tokenizer.
  538. Arguments:
  539. * prompt
  540. * request_id
  541. * lora_request
  542. Returns:
  543. * prompt token ids
  544. '''
  545. tokenizer = self.get_tokenizer_group("prompts must be None if "
  546. "skip_tokenizer_init is True")
  547. return tokenizer.encode(request_id=request_id,
  548. prompt=prompt,
  549. lora_request=lora_request)
  550. def _extract_prompt_components(
  551. self,
  552. inputs: SingletonPromptInputs,
  553. request_id: str,
  554. lora_request: Optional[LoRARequest] = None,
  555. ) -> PromptComponents:
  556. '''
  557. Extract the components of any single encoder or decoder input prompt.
  558. Arguments:
  559. * request_id
  560. * inputs: single encoder or decoder input prompt
  561. * lora_request: this is only valid for decoder prompts
  562. Returns:
  563. * prompt
  564. * prompt_token_ids
  565. * multi_modal_data
  566. '''
  567. if isinstance(inputs, str):
  568. prompt = inputs
  569. prompt_token_ids = self._tokenize_prompt(
  570. prompt,
  571. request_id=request_id,
  572. lora_request=lora_request,
  573. )
  574. multi_modal_data = None
  575. elif isinstance(inputs, dict):
  576. if "prompt_token_ids" in inputs:
  577. prompt = None
  578. prompt_token_ids = inputs["prompt_token_ids"]
  579. else:
  580. # NOTE: This extra assignment is required to pass mypy
  581. prompt = parsed_prompt = inputs["prompt"]
  582. prompt_token_ids = self._tokenize_prompt(
  583. parsed_prompt,
  584. request_id=request_id,
  585. lora_request=lora_request,
  586. )
  587. multi_modal_data = inputs.get("multi_modal_data")
  588. else:
  589. assert_never(inputs)
  590. return prompt, prompt_token_ids, multi_modal_data
  591. def _apply_prompt_adapter(
  592. self,
  593. prompt_token_ids: List[int],
  594. prompt_adapter_request: Optional[PromptAdapterRequest],
  595. ) -> List[int]:
  596. if prompt_adapter_request:
  597. prompt_token_ids = (
  598. [0] * prompt_adapter_request.prompt_adapter_num_virtual_tokens
  599. + prompt_token_ids)
  600. return prompt_token_ids
  601. def _get_default_enc_dec_decoder_prompt(self) -> List[int]:
  602. '''
  603. Specifically for encoder/decoder models:
  604. generate a default decoder prompt for when
  605. the user specifies only the encoder prompt.
  606. Encoder/decoder models utilize the decoder
  607. prompt in different ways; as new models are
  608. added, it is intended that this function
  609. will be extended to produce differing
  610. default decoder prompts, depending on the
  611. model variety.
  612. Absent a special case, the default behavior
  613. of this method is to mirror the behavior of
  614. the HuggingFace (HF) GenerationMixin for a None
  615. decoder prompt, which is to employ a logit processor
  616. setting to force the first decoded token to be <BOS>.
  617. Here, this behavior is approximated by having the
  618. "default" decoder prompt be <BOS>.
  619. However, it is possible that in the future
  620. other models may have different or more
  621. complex logic for the default decoder prompt.
  622. This motivates having a special helper method
  623. for default decoder prompts.
  624. Returns:
  625. * prompt_token_ids
  626. '''
  627. bos_token_id = self._get_bos_token_id()
  628. assert bos_token_id is not None
  629. return [bos_token_id]
  630. def _build_enc_dec_llm_inputs(
  631. self,
  632. encoder_comps: PromptComponents,
  633. decoder_comps: DecoderPromptComponents,
  634. ) -> EncoderDecoderLLMInputs:
  635. encoder_prompt, encoder_prompt_ids, encoder_mm_data = encoder_comps
  636. decoder_prompt, decoder_prompt_ids, decoder_mm_data = decoder_comps
  637. if encoder_mm_data is not None or decoder_mm_data is not None:
  638. raise ValueError("Multi-modal encoder-decoder models are "
  639. "not supported yet")
  640. decoder_prompt_ids = (
  641. self._prepare_decoder_input_ids_for_generation(decoder_prompt_ids))
  642. return EncoderDecoderLLMInputs(
  643. prompt_token_ids=decoder_prompt_ids,
  644. prompt=decoder_prompt,
  645. encoder_prompt_token_ids=encoder_prompt_ids,
  646. encoder_prompt=encoder_prompt,
  647. )
  648. def _process_encoder_decoder_prompt(
  649. self,
  650. inputs: PromptInputs,
  651. request_id: str,
  652. ) -> EncoderDecoderLLMInputs:
  653. '''
  654. For encoder/decoder models only:
  655. Process an input prompt into an
  656. :class:`EncoderDecoderLLMInputs` instance.
  657. There are two types of input prompts:
  658. singleton prompts which carry only the
  659. encoder prompt, and explicit encoder/decoder
  660. prompts which carry both the encoder and the
  661. decoder prompts as member variables.
  662. This function handles the following scenarios:
  663. * Singleton encoder prompt: extract encoder prompt
  664. token ids & infer default decoder prompt token ids
  665. * Explicit encoder/decoder prompt: extract encoder
  666. and decoder prompt token ids
  667. Note that for Explicit encoder/decoder prompts,
  668. each sub-prompt (encoder or decoder prompt) can
  669. have any possible singleton type; thus this
  670. method relies on helper functions to obtain
  671. token ids for the sub-prompts.
  672. Arguments:
  673. * inputs: an input prompt
  674. * request_id
  675. Returns:
  676. * :class:`EncoderDecoderLLMInputs` instance
  677. '''
  678. encoder_comps: PromptComponents
  679. decoder_comps: DecoderPromptComponents
  680. if is_explicit_encoder_decoder_prompt(inputs):
  681. encoder_comps = self._extract_prompt_components(
  682. inputs["encoder_prompt"],
  683. request_id=request_id,
  684. )
  685. if (decoder_input := inputs["decoder_prompt"]) is None:
  686. decoder_comps = None, None, None
  687. else:
  688. decoder_comps = self._extract_prompt_components(
  689. decoder_input,
  690. request_id=request_id,
  691. )
  692. else:
  693. encoder_comps = self._extract_prompt_components(
  694. inputs,
  695. request_id=request_id,
  696. )
  697. decoder_comps = None, None, None
  698. return self._build_enc_dec_llm_inputs(encoder_comps, decoder_comps)
  699. def _build_decoder_only_llm_inputs(
  700. self,
  701. prompt_comps: PromptComponents,
  702. prompt_adapter_request: Optional[PromptAdapterRequest],
  703. ) -> LLMInputs:
  704. prompt, prompt_token_ids, multi_modal_data = prompt_comps
  705. prompt_token_ids = self._apply_prompt_adapter(
  706. prompt_token_ids, prompt_adapter_request=prompt_adapter_request)
  707. return LLMInputs(prompt_token_ids=prompt_token_ids,
  708. prompt=prompt,
  709. multi_modal_data=multi_modal_data)
  710. def _process_decoder_only_prompt(
  711. self,
  712. inputs: SingletonPromptInputs,
  713. request_id: str,
  714. lora_request: Optional[LoRARequest] = None,
  715. prompt_adapter_request: Optional[PromptAdapterRequest] = None,
  716. ) -> LLMInputs:
  717. '''
  718. For decoder-only models:
  719. Process an input prompt into an :class:`LLMInputs` instance.
  720. Arguments:
  721. * inputs: input prompt
  722. * request_id
  723. * lora_request
  724. * prompt_adapter_request
  725. Returns:
  726. * :class:`LLMInputs` instance
  727. '''
  728. prompt_comps = self._extract_prompt_components(
  729. inputs,
  730. request_id=request_id,
  731. lora_request=lora_request,
  732. )
  733. return self._build_decoder_only_llm_inputs(
  734. prompt_comps,
  735. prompt_adapter_request=prompt_adapter_request,
  736. )
  737. def process_model_inputs(
  738. self,
  739. inputs: PromptInputs,
  740. request_id: str,
  741. lora_request: Optional[LoRARequest] = None,
  742. prompt_adapter_request: Optional[PromptAdapterRequest] = None,
  743. ) -> Union[LLMInputs, EncoderDecoderLLMInputs]:
  744. if self.is_encoder_decoder_model():
  745. # Encoder-decoder model requires special mapping of
  746. # input prompts to encoder & decoder
  747. model_inputs = self._process_encoder_decoder_prompt(
  748. inputs,
  749. request_id=request_id,
  750. )
  751. else:
  752. if is_explicit_encoder_decoder_prompt(inputs):
  753. raise ValueError("Cannot pass encoder-decoder prompt "
  754. "to decoder-only models")
  755. # Decoder-only operation
  756. model_inputs = self._process_decoder_only_prompt(
  757. inputs,
  758. request_id=request_id,
  759. lora_request=lora_request,
  760. prompt_adapter_request=prompt_adapter_request,
  761. )
  762. return self.input_processor(model_inputs)
  763. def add_request(
  764. self,
  765. request_id: str,
  766. inputs: PromptInputs,
  767. params: Union[SamplingParams, PoolingParams],
  768. arrival_time: Optional[float] = None,
  769. lora_request: Optional[LoRARequest] = None,
  770. prompt_adapter_request: Optional[PromptAdapterRequest] = None,
  771. ) -> None:
  772. """Add a request to the engine's request pool.
  773. The request is added to the request pool and will be processed by the
  774. scheduler as `engine.step()` is called. The exact scheduling policy is
  775. determined by the scheduler.
  776. Args:
  777. request_id: The unique ID of the request.
  778. prompt: The prompt string. Can be None if prompt_token_ids is
  779. provided.
  780. params: Parameters for sampling or pooling. SamplingParams
  781. for text generation. PoolingParams for pooling.
  782. prompt_token_ids: The token IDs of the prompt. If None, we
  783. use the tokenizer to convert the prompts to token IDs.
  784. arrival_time: The arrival time of the request. If None, we use
  785. the current monotonic time.
  786. multi_modal_data: Multi modal data per request.
  787. Details:
  788. - Set arrival_time to the current time if it is None.
  789. - Set prompt_token_ids to the encoded prompt if it is None.
  790. - Create `best_of` number of :class:`~aphrodite.common.sequence`
  791. objects.
  792. - Create a :class:`~aphrodite.common.sequenceGroup` object
  793. from the list of :class:`~aphrodite.common.sequence`.
  794. - Add the :class:`~aphrodite.common.sequenceGroup` object to the
  795. scheduler.
  796. Example:
  797. >>> # initialize engine
  798. >>> engine = AphroditeEngine.from_engine_args(engine_args)
  799. >>> # set request arguments
  800. >>> example_prompt = "Who is the president of the United States?"
  801. >>> sampling_params = SamplingParams(temperature=0.0)
  802. >>> request_id = 0
  803. >>>
  804. >>> # add the request to the engine
  805. >>> engine.add_request(
  806. >>> str(request_id),
  807. >>> example_prompt,
  808. >>> SamplingParams(temperature=0.0))
  809. >>> # continue the request processing
  810. >>> ...
  811. """
  812. if lora_request is not None and not self.lora_config:
  813. raise ValueError(f"Got lora_request {lora_request} but LoRA is "
  814. "not enabled!")
  815. if arrival_time is None:
  816. arrival_time = time.time()
  817. processed_inputs = self.process_model_inputs(
  818. inputs,
  819. request_id=request_id,
  820. lora_request=lora_request,
  821. prompt_adapter_request=prompt_adapter_request,
  822. )
  823. self._add_processed_request(
  824. request_id=request_id,
  825. processed_inputs=processed_inputs,
  826. params=params,
  827. arrival_time=arrival_time,
  828. lora_request=lora_request,
  829. prompt_adapter_request=prompt_adapter_request,
  830. )
  831. def _create_sequence_group_with_sampling(
  832. self,
  833. request_id: str,
  834. seq: Sequence,
  835. sampling_params: SamplingParams,
  836. arrival_time: float,
  837. lora_request: Optional[LoRARequest],
  838. prompt_adapter_request: Optional[PromptAdapterRequest] = None,
  839. encoder_seq: Optional[Sequence] = None,
  840. ) -> SequenceGroup:
  841. """Creates a SequenceGroup with SamplingParams."""
  842. max_logprobs = self.get_model_config().max_logprobs
  843. if (sampling_params.logprobs
  844. and sampling_params.logprobs > max_logprobs) or (
  845. sampling_params.prompt_logprobs
  846. and sampling_params.prompt_logprobs > max_logprobs):
  847. raise ValueError(f"Cannot request more than "
  848. f"{max_logprobs} logprobs.")
  849. # Defensive copy of SamplingParams, which are used by the sampler,
  850. # this doesn't deep-copy LogitsProcessor objects
  851. sampling_params = sampling_params.clone()
  852. sampling_params.update_from_generation_config(
  853. self.generation_config_fields, seq.eos_token_id)
  854. # Create the sequence group.
  855. seq_group = SequenceGroup(
  856. request_id=request_id,
  857. seqs=[seq],
  858. arrival_time=arrival_time,
  859. sampling_params=sampling_params,
  860. lora_request=lora_request,
  861. prompt_adapter_request=prompt_adapter_request,
  862. encoder_seq=encoder_seq)
  863. return seq_group
  864. def _create_sequence_group_with_pooling(
  865. self,
  866. request_id: str,
  867. seq: Sequence,
  868. pooling_params: PoolingParams,
  869. arrival_time: float,
  870. lora_request: Optional[LoRARequest],
  871. prompt_adapter_request: Optional[PromptAdapterRequest] = None,
  872. encoder_seq: Optional[Sequence] = None,
  873. ) -> SequenceGroup:
  874. """Creates a SequenceGroup with PoolingParams."""
  875. # Defensive copy of PoolingParams, which are used by the pooler
  876. pooling_params = pooling_params.clone()
  877. # Create the sequence group.
  878. seq_group = SequenceGroup(
  879. request_id=request_id,
  880. seqs=[seq],
  881. arrival_time=arrival_time,
  882. lora_request=lora_request,
  883. pooling_params=pooling_params,
  884. prompt_adapter_request=prompt_adapter_request,
  885. encoder_seq=encoder_seq)
  886. return seq_group
  887. def abort_request(self, request_id: Union[str, Iterable[str]]) -> None:
  888. """Aborts a request(s) with the given ID.
  889. Args:
  890. request_id: The ID(s) of the request to abort.
  891. Details:
  892. - Refer to the
  893. :meth:`~aphrodite.processing.scheduler.Scheduler.abort_seq_group`
  894. from class :class:`~aphrodite.processing.scheduler.Scheduler`.
  895. Example:
  896. >>> # initialize engine and add a request with request_id
  897. >>> request_id = str(0)
  898. >>> # abort the request
  899. >>> engine.abort_request(request_id)
  900. """
  901. for scheduler in self.scheduler:
  902. scheduler.abort_seq_group(request_id)
  903. def get_model_config(self) -> ModelConfig:
  904. """Gets the model configuration."""
  905. return self.model_config
  906. def get_parallel_config(self) -> ParallelConfig:
  907. """Gets the parallel configuration."""
  908. return self.parallel_config
  909. def get_decoding_config(self) -> DecodingConfig:
  910. """Gets the decoding configuration."""
  911. return self.decoding_config
  912. def get_scheduler_config(self) -> SchedulerConfig:
  913. """Gets the scheduler configuration."""
  914. return self.scheduler_config
  915. def get_lora_config(self) -> LoRAConfig:
  916. """Gets the LoRA configuration."""
  917. return self.lora_config
  918. def get_num_unfinished_requests(self) -> int:
  919. """Gets the number of unfinished requests."""
  920. return sum(scheduler.get_num_unfinished_seq_groups()
  921. for scheduler in self.scheduler)
  922. def has_unfinished_requests(self) -> bool:
  923. """Returns True if there are unfinished requests."""
  924. return any(scheduler.has_unfinished_seqs()
  925. for scheduler in self.scheduler)
  926. def has_unfinished_requests_for_virtual_engine(
  927. self, virtual_engine: int) -> bool:
  928. """
  929. Returns True if there are unfinished requests for the virtual engine.
  930. """
  931. return self.scheduler[virtual_engine].has_unfinished_seqs()
  932. def _process_sequence_group_outputs(
  933. self,
  934. seq_group: SequenceGroup,
  935. outputs: List[EmbeddingSequenceGroupOutput],
  936. ) -> None:
  937. seq_group.embeddings = outputs[0].embeddings
  938. for seq in seq_group.get_seqs():
  939. seq.status = SequenceStatus.FINISHED_STOPPED
  940. return
  941. def _process_model_outputs(
  942. self,
  943. output: GenericSequence[Union[SamplerOutput, PoolerOutput]],
  944. scheduled_seq_groups: List[ScheduledSequenceGroup],
  945. ignored_seq_groups: List[SequenceGroup],
  946. seq_group_metadata_list: List[SequenceGroupMetadata],
  947. ) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
  948. """Apply the model output to the sequences in the scheduled seq groups.
  949. Returns RequestOutputs that can be returned to the client.
  950. """
  951. now = time.time()
  952. # Organize outputs by [sequence group][step] instead of
  953. # [step][sequence group].
  954. output_by_sequence_group = create_output_by_sequence_group(
  955. output, num_seq_groups=len(scheduled_seq_groups))
  956. # Update the scheduled sequence groups with the model outputs.
  957. for scheduled_seq_group, outputs, seq_group_meta in zip(
  958. scheduled_seq_groups, output_by_sequence_group,
  959. seq_group_metadata_list):
  960. seq_group = scheduled_seq_group.seq_group
  961. seq_group.update_num_computed_tokens(
  962. scheduled_seq_group.token_chunk_size)
  963. if self.model_config.embedding_mode:
  964. self._process_sequence_group_outputs(seq_group, outputs)
  965. continue
  966. self.output_processor.process_prompt_logprob(seq_group, outputs)
  967. if seq_group_meta.do_sample:
  968. self.output_processor.process_outputs(seq_group, outputs)
  969. # Free the finished sequence groups.
  970. for scheduler in self.scheduler:
  971. scheduler.free_finished_seq_groups()
  972. # Create the outputs.
  973. request_outputs: List[Union[RequestOutput,
  974. EmbeddingRequestOutput]] = []
  975. for scheduled_seq_group in scheduled_seq_groups:
  976. seq_group = scheduled_seq_group.seq_group
  977. seq_group.maybe_set_first_token_time(now)
  978. request_output = RequestOutputFactory.create(seq_group)
  979. request_outputs.append(request_output)
  980. for seq_group in ignored_seq_groups:
  981. request_output = RequestOutputFactory.create(seq_group)
  982. request_outputs.append(request_output)
  983. return request_outputs
  984. def step(self) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
  985. """Performs one decoding iteration and returns newly generated results.
  986. .. figure:: https://i.imgur.com/sv2HssD.png
  987. :alt: Overview of the step function
  988. :align: center
  989. Overview of the step function.
  990. Details:
  991. - Step 1: Schedules the sequences to be executed in the next
  992. iteration and the token blocks to be swapped in/out/copy.
  993. - Depending on the scheduling policy,
  994. sequences may be `preempted/reordered`.
  995. - A Sequence Group (SG) refer to a group of sequences
  996. that are generated from the same prompt.
  997. - Step 2: Calls the distributed executor to execute the model.
  998. - Step 3: Processes the model output. This mainly includes:
  999. - Decodes the relevant outputs.
  1000. - Updates the scheduled sequence groups with model outputs
  1001. based on its `sampling parameters` (`use_beam_search` or not).
  1002. - Frees the finished sequence groups.
  1003. - Finally, it creates and returns the newly generated results.
  1004. Example:
  1005. >>> # Please see the example/ folder for more detailed examples.
  1006. >>>
  1007. >>> # initialize engine and request arguments
  1008. >>> engine = AphroditeEngine.from_engine_args(engine_args)
  1009. >>> example_inputs = [(0, "What is LLM?",
  1010. >>> SamplingParams(temperature=0.0))]
  1011. >>>
  1012. >>> # Start the engine with an event loop
  1013. >>> while True:
  1014. >>> if example_inputs:
  1015. >>> req_id, prompt, sampling_params = example_inputs.pop(0)
  1016. >>> engine.add_request(str(req_id), prompt, sampling_params)
  1017. >>>
  1018. >>> # continue the request processing
  1019. >>> request_outputs = engine.step()
  1020. >>> for request_output in request_outputs:
  1021. >>> if request_output.finished:
  1022. >>> # return or show the request output
  1023. >>>
  1024. >>> if not (engine.has_unfinished_requests() or example_inputs):
  1025. >>> break
  1026. """
  1027. if self.parallel_config.pipeline_parallel_size > 1:
  1028. raise NotImplementedError(
  1029. "Pipeline parallelism is only supported through AsyncAphrodite "
  1030. "as performance will be severely degraded otherwise.")
  1031. seq_group_metadata_list, scheduler_outputs = self.scheduler[
  1032. 0].schedule()
  1033. if not scheduler_outputs.is_empty():
  1034. finished_requests_ids = self.scheduler[
  1035. 0].get_and_reset_finished_requests_ids()
  1036. execute_model_req = ExecuteModelRequest(
  1037. seq_group_metadata_list=seq_group_metadata_list,
  1038. blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,
  1039. blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,
  1040. blocks_to_copy=scheduler_outputs.blocks_to_copy,
  1041. num_lookahead_slots=scheduler_outputs.num_lookahead_slots,
  1042. running_queue_size=scheduler_outputs.running_queue_size,
  1043. finished_requests_ids=finished_requests_ids,
  1044. )
  1045. output = self.model_executor.execute_model(
  1046. execute_model_req=execute_model_req)
  1047. else:
  1048. output = []
  1049. request_outputs = self._process_model_outputs(
  1050. output, scheduler_outputs.scheduled_seq_groups,
  1051. scheduler_outputs.ignored_seq_groups, seq_group_metadata_list)
  1052. # Log stats.
  1053. self.do_log_stats(scheduler_outputs, output)
  1054. if not self.has_unfinished_requests():
  1055. # Stop the execute model loop in parallel workers until there are
  1056. # more requests to process. This avoids waiting indefinitely in
  1057. # torch.distributed ops which may otherwise timeout, and unblocks
  1058. # the RPC thread in the workers so that they can process any other
  1059. # queued control plane messages, such as add/remove lora adapters.
  1060. self.model_executor.stop_remote_worker_execution_loop()
  1061. return request_outputs
  1062. def add_logger(self, logger_name: str, logger: StatLoggerBase) -> None:
  1063. if logger_name in self.stat_loggers:
  1064. raise KeyError(f"Logger with name {logger_name} already exists.")
  1065. self.stat_loggers[logger_name] = logger
  1066. def remove_logger(self, logger_name: str) -> None:
  1067. if logger_name not in self.stat_loggers:
  1068. raise KeyError(f"Logger with name {logger_name} does not exist.")
  1069. del self.stat_loggers[logger_name]
  1070. def do_log_stats(
  1071. self,
  1072. scheduler_outputs: Optional[SchedulerOutputs] = None,
  1073. model_output: Optional[List[SamplerOutput]] = None) -> None:
  1074. """Forced log when no requests active."""
  1075. if self.log_stats:
  1076. stats = self._get_stats(scheduler_outputs, model_output)
  1077. for loggers in self.stat_loggers.values():
  1078. loggers.log(stats)
  1079. def _get_stats(
  1080. self,
  1081. scheduler_outputs: Optional[SchedulerOutputs],
  1082. model_output: Optional[List[SamplerOutput]] = None) -> Stats:
  1083. """Get Stats to be Logged to Prometheus.
  1084. Args:
  1085. scheduler_outputs: Optional, used to populate metrics related to
  1086. the scheduled batch,
  1087. model_output: Optional, used to emit speculative decoding metrics
  1088. which are created by the workers.
  1089. """
  1090. now = time.time()
  1091. # System State
  1092. # Scheduler State
  1093. num_running_sys = sum(
  1094. len(scheduler.running) for scheduler in self.scheduler)
  1095. num_swapped_sys = sum(
  1096. len(scheduler.swapped) for scheduler in self.scheduler)
  1097. num_waiting_sys = sum(
  1098. len(scheduler.waiting) for scheduler in self.scheduler)
  1099. # KV Cache Usage in %
  1100. num_total_gpu = self.cache_config.num_gpu_blocks
  1101. gpu_cache_usage_sys = 0.
  1102. if num_total_gpu is not None:
  1103. num_free_gpu = sum(
  1104. scheduler.block_manager.get_num_free_gpu_blocks()
  1105. for scheduler in self.scheduler)
  1106. if not self.model_config.is_attention_free():
  1107. gpu_cache_usage_sys = 1.0 - (num_free_gpu / num_total_gpu)
  1108. else:
  1109. gpu_cache_usage_sys = 0.0
  1110. num_total_cpu = self.cache_config.num_cpu_blocks
  1111. cpu_cache_usage_sys = 0.
  1112. if num_total_cpu is not None and num_total_cpu > 0:
  1113. num_free_cpu = sum(
  1114. scheduler.block_manager.get_num_free_cpu_blocks()
  1115. for scheduler in self.scheduler)
  1116. if not self.model_config.is_attention_free():
  1117. cpu_cache_usage_sys = 1.0 - (num_free_cpu / num_total_cpu)
  1118. else:
  1119. cpu_cache_usage_sys = 0.0
  1120. # Prefix Cache Hit Rate. Note that we always use
  1121. # the cache hit rate of the first virtual engine.
  1122. cpu_prefix_cache_hit_rate = self.scheduler[
  1123. 0].get_prefix_cache_hit_rate(Device.CPU)
  1124. gpu_prefix_cache_hit_rate = self.scheduler[
  1125. 0].get_prefix_cache_hit_rate(Device.GPU)
  1126. # Iteration stats
  1127. num_prompt_tokens_iter = 0
  1128. num_generation_tokens_iter = 0
  1129. time_to_first_tokens_iter: List[float] = []
  1130. time_per_output_tokens_iter: List[float] = []
  1131. num_preemption_iter = (0 if scheduler_outputs is None else
  1132. scheduler_outputs.preempted)
  1133. # Request stats
  1134. # Latency
  1135. time_e2e_requests: List[float] = []
  1136. # Metadata
  1137. num_prompt_tokens_requests: List[int] = []
  1138. num_generation_tokens_requests: List[int] = []
  1139. best_of_requests: List[int] = []
  1140. n_requests: List[int] = []
  1141. finished_reason_requests: List[str] = []
  1142. # NOTE: This loop assumes prefill seq_groups are before
  1143. # decode seq_groups in scheduled_seq_groups.
  1144. if scheduler_outputs is not None:
  1145. num_generation_tokens_from_prefill_groups = 0.
  1146. # NOTE: if scheduler_outputs.num_prefill_groups > 0 and
  1147. # the len of scheduler_outputs.scheduled_seq_groups is !=
  1148. # scheduler_outputs.num_prefill_groups, this means that
  1149. # chunked prefills have been detected.
  1150. for idx, scheduled_seq_group in enumerate(
  1151. scheduler_outputs.scheduled_seq_groups):
  1152. group_was_prefill = idx < scheduler_outputs.num_prefill_groups
  1153. seq_group = scheduled_seq_group.seq_group
  1154. # NOTE: a seq_group that completed all of its prefill tokens
  1155. # in the last iteration will have seq_group.is_prefill() = False
  1156. # with group_was_prefill = True
  1157. if group_was_prefill:
  1158. # Number of prompt tokens.
  1159. num_prompt_tokens_iter += (
  1160. scheduled_seq_group.token_chunk_size)
  1161. # If the seq_group just finished the prefill state
  1162. # get TTFT.
  1163. if not seq_group.is_prefill():
  1164. latency = seq_group.get_last_latency(now)
  1165. time_to_first_tokens_iter.append(latency)
  1166. # One generation token per finished prefill.
  1167. num_generation_tokens_from_prefill_groups += (
  1168. seq_group.num_seqs())
  1169. else:
  1170. # TPOTs.
  1171. latency = seq_group.get_last_latency(now)
  1172. time_per_output_tokens_iter.append(latency)
  1173. # Because of chunked prefill, we can have a single sequence
  1174. # group that does multiple prompt_runs. To prevent logging
  1175. # the same metadata more than once per request, we standardize
  1176. # on logging request level information for finished requests,
  1177. # which can only happen once.
  1178. if seq_group.is_finished():
  1179. # Latency timings
  1180. time_e2e_requests.append(now -
  1181. seq_group.metrics.arrival_time)
  1182. # Metadata
  1183. num_prompt_tokens_requests.append(
  1184. len(seq_group.prompt_token_ids))
  1185. num_generation_tokens_requests.extend([
  1186. seq.get_output_len()
  1187. for seq in seq_group.get_finished_seqs()
  1188. ])
  1189. if seq_group.sampling_params is not None:
  1190. best_of_requests.append(
  1191. seq_group.sampling_params.best_of)
  1192. n_requests.append(seq_group.sampling_params.n)
  1193. finished_reason_requests.extend([
  1194. SequenceStatus.get_finished_reason(seq.status)
  1195. for seq in seq_group.get_finished_seqs()
  1196. ])
  1197. # Number of generation tokens.
  1198. # num_batched_tokens equals the number of prompt_tokens plus the
  1199. # number of decode_tokens in a single iteration. So,
  1200. # num_generation_tokens = num_batched_tokens - num_prompt_tokens
  1201. # + num_generation_tokens_from_prefill_groups (since we generate
  1202. # one token on prefills on iters where the prefill finishes).
  1203. num_generation_tokens_iter = (
  1204. scheduler_outputs.num_batched_tokens - num_prompt_tokens_iter +
  1205. num_generation_tokens_from_prefill_groups)
  1206. # Spec decode, if enabled, emits specialized metrics from the worker in
  1207. # sampler output.
  1208. if model_output and (model_output[0].spec_decode_worker_metrics
  1209. is not None):
  1210. spec_decode_metrics = model_output[0].spec_decode_worker_metrics
  1211. else:
  1212. spec_decode_metrics = None
  1213. return Stats(
  1214. now=now,
  1215. # System stats
  1216. # Scheduler State
  1217. num_running_sys=num_running_sys,
  1218. num_swapped_sys=num_swapped_sys,
  1219. num_waiting_sys=num_waiting_sys,
  1220. # KV Cache Usage in %
  1221. gpu_cache_usage_sys=gpu_cache_usage_sys,
  1222. cpu_cache_usage_sys=cpu_cache_usage_sys,
  1223. # Prefix Cache Hit Rate
  1224. cpu_prefix_cache_hit_rate=cpu_prefix_cache_hit_rate,
  1225. gpu_prefix_cache_hit_rate=gpu_prefix_cache_hit_rate,
  1226. # Iteration stats
  1227. num_prompt_tokens_iter=num_prompt_tokens_iter,
  1228. num_generation_tokens_iter=num_generation_tokens_iter,
  1229. time_to_first_tokens_iter=time_to_first_tokens_iter,
  1230. time_per_output_tokens_iter=time_per_output_tokens_iter,
  1231. spec_decode_metrics=spec_decode_metrics,
  1232. num_preemption_iter=num_preemption_iter,
  1233. # Request stats
  1234. # Latency
  1235. time_e2e_requests=time_e2e_requests,
  1236. # Metadata
  1237. num_prompt_tokens_requests=num_prompt_tokens_requests,
  1238. num_generation_tokens_requests=num_generation_tokens_requests,
  1239. best_of_requests=best_of_requests,
  1240. n_requests=n_requests,
  1241. finished_reason_requests=finished_reason_requests,
  1242. )
  1243. def add_lora(self, lora_request: LoRARequest) -> bool:
  1244. return self.model_executor.add_lora(lora_request)
  1245. def remove_lora(self, lora_id: int) -> bool:
  1246. return self.model_executor.remove_lora(lora_id)
  1247. def list_loras(self) -> List[int]:
  1248. return self.model_executor.list_loras()
  1249. def pin_lora(self, lora_id: int) -> bool:
  1250. return self.model_executor.pin_lora(lora_id)
  1251. def add_prompt_adapter(
  1252. self, prompt_adapter_request: PromptAdapterRequest) -> bool:
  1253. return self.model_executor.add_prompt_adapter(prompt_adapter_request)
  1254. def remove_prompt_adapter(self, prompt_adapter_id: int) -> bool:
  1255. return self.model_executor.remove_prompt_adapter(prompt_adapter_id)
  1256. def list_prompt_adapters(self) -> List[int]:
  1257. return self.model_executor.list_prompt_adapters()
  1258. def check_health(self) -> None:
  1259. if self.tokenizer:
  1260. self.tokenizer.check_health()
  1261. self.model_executor.check_health()
  1262. def is_encoder_decoder_model(self):
  1263. return self.model_config.is_encoder_decoder_model
  1264. def is_embedding_model(self):
  1265. return self.model_config.is_embedding_model
  1266. setup_logger()