add the instrumentation needed to measure model forward time

benchmarks/backend_request_func.py

@@ -3,12 +3,16 @@
 import sys
 import time
 import traceback
+import uuid
 from dataclasses import dataclass, field
 from typing import List, Optional
 
 import aiohttp
 from tqdm.asyncio import tqdm
 
+from vllm import AsyncLLMEngine
+from vllm.sampling_params import SamplingParams
+
 AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
 
 
@@ -21,6 +25,7 @@ class RequestFuncInput:
     model: str
     best_of: int = 1
     use_beam_search: bool = False
+    engine: AsyncLLMEngine = None
 
 
 @dataclass
@@ -370,6 +375,59 @@ async def async_request_openai_chat_completions(
     return output
 
 
+async def async_request_vllm_engine(
+    request_func_input: RequestFuncInput,
+    pbar: Optional[tqdm] = None,
+) -> RequestFuncOutput:
+    engine = request_func_input.engine
+    sampling_params = SamplingParams(
+        temperature=0.0,
+        best_of=request_func_input.best_of,
+        max_tokens=request_func_input.output_len,
+    )
+
+    generator = engine.generate(
+        request_func_input.prompt,
+        sampling_params,
+        str(uuid.uuid4()),
+    )
+
+    assert not request_func_input.use_beam_search
+
+    output = RequestFuncOutput()
+    output.prompt_len = request_func_input.prompt_len
+
+    generated_text = ""
+    ttft = 0.0
+    st = time.perf_counter()
+    most_recent_timestamp = st
+    try:
+        async for out in generator:
+            assert len(out.outputs) == 1
+            timestamp = time.perf_counter()
+            # First token
+            if ttft == 0.0:
+                ttft = time.perf_counter() - st
+                output.ttft = ttft
+
+            # Decoding phase
+            output.itl.append(timestamp - most_recent_timestamp)
+
+            most_recent_timestamp = timestamp
+
+        output.generated_text = out.outputs[0].text
+        output.success = True
+        output.latency = time.perf_counter() - st
+    except Exception:
+        output.success = False
+        exc_info = sys.exc_info()
+        output.error = "".join(traceback.format_exception(*exc_info))
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
 # Since vllm must support Python 3.8, we can't use str.removeprefix(prefix)
 # introduced in Python 3.9
 def remove_prefix(text: str, prefix: str) -> str:
@@ -381,9 +439,10 @@ def remove_prefix(text: str, prefix: str) -> str:
 ASYNC_REQUEST_FUNCS = {
     "tgi": async_request_tgi,
     "vllm": async_request_openai_completions,
+    "vllm-engine": async_request_vllm_engine,
     "lmdeploy": async_request_openai_completions,
     "deepspeed-mii": async_request_deepspeed_mii,
     "openai": async_request_openai_completions,
     "openai-chat": async_request_openai_chat_completions,
     "tensorrt-llm": async_request_trt_llm,
-}
+}

benchmarks/benchmark_serving.py

@@ -17,6 +17,10 @@
         --dataset-path <path to dataset> \
         --request-rate <request_rate> \ # By default <request_rate> is inf
         --num-prompts <num_prompts> # By default <num_prompts> is 1000
+
+    when using tgi backend, add
+        --endpoint /generate_stream
+    to the end of the command above.
 """
 import argparse
 import asyncio
@@ -35,9 +39,49 @@
 from tqdm.asyncio import tqdm
 from transformers import PreTrainedTokenizerBase
 
+from vllm import AsyncEngineArgs, AsyncLLMEngine
+from vllm.engine.async_llm_engine import _AsyncLLMEngine
 from vllm.transformers_utils.tokenizer import get_tokenizer
 
 
+
+def print_latency(gpu_latency_set, cpu_latency_set, title=""):
+    # 10 warmup queries
+
+    gpu_latency_set = gpu_latency_set[10:]
+    cpu_latency_set = cpu_latency_set[10:]
+    count = len(gpu_latency_set)
+    if count > 0:
+        gpu_latency_set.sort()
+        cpu_latency_set.sort()
+        n50 = (count - 1) * 0.5 + 1
+        n90 = (count - 1) * 0.9 + 1
+        n95 = (count - 1) * 0.95 + 1
+        n99 = (count - 1) * 0.99 + 1
+        n999 = (count - 1) * 0.999 + 1
+
+        avg = sum(gpu_latency_set) / count
+        p50 = gpu_latency_set[int(n50) - 1]
+        p90 = gpu_latency_set[int(n90) - 1]
+        p95 = gpu_latency_set[int(n95) - 1]
+        p99 = gpu_latency_set[int(n99) - 1]
+        p999 = gpu_latency_set[int(n999) - 1]
+
+        cpu_avg = sum(cpu_latency_set) / count
+        cpu_p50 = cpu_latency_set[int(n50) - 1]
+        cpu_p90 = cpu_latency_set[int(n90) - 1]
+        cpu_p95 = cpu_latency_set[int(n95) - 1]
+        cpu_p99 = cpu_latency_set[int(n99) - 1]
+        cpu_p999 = cpu_latency_set[int(n999) - 1]
+
+        print("====== latency stats {0} ======", title)
+        print("\tAvg Latency: ({0:8.2f}: {1:8.2f}) ms".format(avg, cpu_avg))
+        print("\tP50 Latency: ({0:8.2f}: {1:8.2f}) ms".format(p50, cpu_p50))
+        print("\tP90 Latency: ({0:8.2f}: {1:8.2f}) ms".format(p90, cpu_p90))
+        print("\tP95 Latency: ({0:8.2f}: {1:8.2f}) ms".format(p95, cpu_p95))
+        print("\tP99 Latency: ({0:8.2f}: {1:8.2f}) ms".format(p99, cpu_p99))
+        print("\t999 Latency: ({0:8.2f}: {1:8.2f}) ms".format(p999, cpu_p999))
+
 @dataclass
 class BenchmarkMetrics:
     completed: int
@@ -246,6 +290,27 @@ async def benchmark(
     else:
         raise ValueError(f"Unknown backend: {backend}")
 
+    if backend == "vllm-engine":
+        engine_args = AsyncEngineArgs(
+            model=model_id,
+            tensor_parallel_size=8,
+            #pipeline_parallel_size=2,
+            #gpu_memory_utilization=0.95,
+            max_num_seqs=32,
+            max_num_batched_tokens=256,
+            enable_chunked_prefill=True,
+            disable_log_requests=True,
+            load_format="dummy",
+            #quantization="deepspeedfp",
+            #load_format="sharded_state",
+            #engine_use_ray=True,
+            enforce_eager=True,
+            #tokenizer_pool_size=16,
+        )
+        engine = AsyncLLMEngine.from_engine_args(engine_args)
+    else:
+        engine = None
+
     print(f"Traffic request rate: {request_rate}")
 
     pbar = None if disable_tqdm else tqdm(total=len(input_requests))
@@ -262,6 +327,7 @@ async def benchmark(
             output_len=output_len,
             best_of=best_of,
             use_beam_search=use_beam_search,
+            engine=engine,
         )
         tasks.append(
             asyncio.create_task(
@@ -274,6 +340,12 @@ async def benchmark(
 
     benchmark_duration = time.perf_counter() - benchmark_start_time
 
+    for key in _AsyncLLMEngine.gpu_latencies:
+        print(f'-'*50)
+        print_latency(_AsyncLLMEngine.gpu_latencies[key], _AsyncLLMEngine.cpu_latencies[key], key)
+        print(f'-'*50)
+    # print(_AsyncLLMEngine.all_latencies)
+    # exit()
     metrics, actual_output_lens = calculate_metrics(
         input_requests=input_requests,
         outputs=outputs,
@@ -593,4 +665,4 @@ def main(args: argparse.Namespace):
     )
 
     args = parser.parse_args()
-    main(args)
+    main(args)

vllm/engine/async_llm_engine.py

@@ -5,7 +5,7 @@
                     Set, Tuple, Type, Union)
 
 from transformers import PreTrainedTokenizer
-
+import torch
 import vllm.envs as envs
 from vllm.config import DecodingConfig, ModelConfig
 from vllm.core.scheduler import SchedulerOutputs
@@ -198,7 +198,18 @@ def has_new_requests(self):
 
 class _AsyncLLMEngine(LLMEngine):
     """Extension of LLMEngine to add async methods."""
-
+    gpu_latencies = {'scheduling': [], 
+                         'prepare_input': [], 
+                         'model_forward': [], 
+                         'compute_logits': [], 
+                         'sample': [],
+                         'process_outputs': []}
+    cpu_latencies = {'scheduling': [], 
+                         'prepare_input': [], 
+                         'model_forward': [], 
+                         'compute_logits': [], 
+                         'sample': [],
+                         'process_outputs': []}
     async def step_async(
             self) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
         """Performs one decoding iteration and returns newly generated results.
@@ -210,8 +221,16 @@ async def step_async(
         and updates the scheduler with the model outputs. Finally, it decodes
         the sequences and returns the newly generated results.
         """
-        seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule()
 
+        start_event = torch.cuda.Event(enable_timing=True)
+        schedule_event = torch.cuda.Event(enable_timing=True)
+        new_event = torch.cuda.Event(enable_timing=True)
+        process_outputs_event = torch.cuda.Event(enable_timing=True)
+        start_event.record()
+        t0 = time.time()
+        seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule()
+        t1 = time.time()
+        schedule_event.record()
         if not scheduler_outputs.is_empty():
             # Execute the model.
             execute_model_req = ExecuteModelRequest(
@@ -224,16 +243,31 @@ async def step_async(
             )
             output = await self.model_executor.execute_model_async(
                 execute_model_req)
+            latencies = output[0][-1]
+            output = [output[0][0]]
         else:
             output = []
-
+            latencies = {}
+        new_event.record()
+        t2 = time.time()
         request_outputs = self._process_model_outputs(
             output, scheduler_outputs.scheduled_seq_groups,
             scheduler_outputs.ignored_seq_groups, seq_group_metadata_list)
-
+        t3 = time.time()
+        process_outputs_event.record()
         # Log stats.
         self.do_log_stats(scheduler_outputs, output)
-
+        torch.cuda.synchronize()
+        schedule_time = start_event.elapsed_time(schedule_event)
+        process_outputs_time = new_event.elapsed_time(process_outputs_event)
+        _AsyncLLMEngine.gpu_latencies['scheduling'].append(schedule_time)
+        _AsyncLLMEngine.cpu_latencies['scheduling'].append((t1 - t0)*1000)
+        _AsyncLLMEngine.gpu_latencies['process_outputs'].append(process_outputs_time)
+        _AsyncLLMEngine.cpu_latencies['process_outputs'].append((t3 - t2)*1000)
+        for key in latencies:
+            _AsyncLLMEngine.gpu_latencies[key].append(latencies[key][0])
+            _AsyncLLMEngine.cpu_latencies[key].append(latencies[key][1])
+        # print(f'{[_AsyncLLMEngine.all_latencies[key][-1] for key in _AsyncLLMEngine.all_latencies]}')
         return request_outputs
 
     async def encode_request_async(
@@ -326,7 +360,6 @@ def __init__(self,
         self.log_requests = log_requests
         self.max_log_len = max_log_len
         self.engine = self._init_engine(*args, **kwargs)
-
         self.background_loop: Optional[asyncio.Future] = None
         # We need to keep a reference to unshielded
         # task as well to prevent it from being garbage
@@ -481,7 +514,6 @@ async def engine_step(self) -> bool:
             request_outputs = await self.engine.step.remote()  # type: ignore
         else:
             request_outputs = await self.engine.step_async()
-
         # Put the outputs into the corresponding streams.
         for request_output in request_outputs:
             self._request_tracker.process_request_output(

vllm/worker/model_runner.py

@@ -33,7 +33,11 @@
 _BATCH_SIZES_TO_CAPTURE = [1, 2, 4] + [
     _BATCH_SIZE_ALIGNMENT * i for i in range(1, 33)
 ]
-
+start_event = None
+compute_logits_event = None
+prepare_event = None
+model_forward_event = None
+sample_event = None
 
 class ModelInput(NamedTuple):
     input_tokens: torch.Tensor
@@ -652,10 +656,21 @@ def execute_model(
         seq_group_metadata_list: List[SequenceGroupMetadata],
         kv_caches: List[torch.Tensor],
     ) -> Optional[SamplerOutput]:
+        global start_event, prepare_event, model_forward_event, compute_logits_event, sample_event
+        if start_event is None:
+            start_event = torch.cuda.Event(enable_timing=True)
+            prepare_event = torch.cuda.Event(enable_timing=True)
+            model_forward_event = torch.cuda.Event(enable_timing=True)
+            compute_logits_event = torch.cuda.Event(enable_timing=True)
+            sample_event = torch.cuda.Event(enable_timing=True)
+
+        start_event.record()
+        t0 = time.time()
         (input_tokens, input_positions, attn_metadata, sampling_metadata,
          lora_requests, lora_mapping, multi_modal_input
          ) = self.prepare_input_tensors(seq_group_metadata_list)
-
+        prepare_event.record()
+        t1 = time.time()
         if self.lora_config:
             self.set_active_loras(lora_requests, lora_mapping)
 
@@ -675,22 +690,41 @@ def execute_model(
         }
         if self.vision_language_config:
             execute_model_kwargs.update({"image_input": multi_modal_input})
+
+        # t0 = time.time()
         hidden_states = model_executable(**execute_model_kwargs)
+        # t1 = time.time()
+        # fwd_time = t1 - t0
+        model_forward_event.record()
+        t2 = time.time()
 
         # Compute the logits.
         logits = self.model.compute_logits(hidden_states, sampling_metadata)
-
+        compute_logits_event.record()
+        t3 = time.time()
         # Only perform sampling in the driver worker.
         if not self.is_driver_worker:
             return None
 
+        # t1 = time.time()
         # Sample the next token.
         output = self.model.sample(
             logits=logits,
             sampling_metadata=sampling_metadata,
         )
-
-        return output
+        sample_event.record()
+        t4 = time.time()
+        torch.cuda.synchronize()
+        # total_time = time.time() - t0
+        prepare_input_latency = start_event.elapsed_time(prepare_event)
+        model_forward_latency = prepare_event.elapsed_time(model_forward_event)
+        compute_logits_latency = model_forward_event.elapsed_time(compute_logits_event)
+        sample_latency = compute_logits_event.elapsed_time(sample_event)
+        # print(f'total_time {total_time*1000:.3f} prepare_input {prepare_input_latency:.3f} : {(t1-t0) * 1000:.3f} model_forward {model_forward_latency:.3f} : {(t2-t1) * 1000:.3f} compute_logits {compute_logits_latency:.3f} : {(t3-t2) * 1000:.3f} sample {sample_latency:.3f} : {(t4-t3) * 1000:.3f}')
+        return (output, {'prepare_input': (prepare_input_latency, (t1-t0)*1000), 
+                         'model_forward': (model_forward_latency, (t2-t1)*1000), 
+                         'compute_logits': (compute_logits_latency, (t3-t2)*1000), 
+                         'sample': (sample_latency, (t4-t3)*1000)})
 
     @torch.inference_mode()
     def profile_run(self) -> None: