[Bugfix] Last token measurement fix

[rajveerb]

RequestMetrics' first_token_time is reported as greater than last_token_time which is false. Below is a sample output of RequestMetrics with the incorrect behavior:

RequestMetrics(arrival_time=1734697165.6934612, \
    last_token_time=1734697165.6934612, \
    first_scheduled_time=1734697165.700166, \
    first_token_time=1734697165.7171977, \
    time_in_queue=0.006704807281494141, \
    finished_time=1734697165.7414508, \
    scheduler_time=0.0014218189753592014, \
    model_forward_time=None, \
    model_execute_time=None)

FIX #9998

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[rajveerb]

Simple example of code where it fails:

from vllm import LLM, SamplingParams

prompts = [
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.0, top_p=1.0)

llm = LLM(model="facebook/opt-125m",)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
    print(output.metrics)

[rajveerb]

Cause of the issue:

When request is received by the vLLM and once it reaches line 665 in vllm/sequence.py , the last_token_time is set to the arrival time of the request. This initial value does not change.

Note: SequenceGroup.get_last_latency function is responsible for subsequent update of last_token_time variable.

SequenceGroup.get_last_latency is called at two locations inside LLMEngine._get_stats function present in vllm/engine/llm_engine.py. The first call is made at line 1704 in vllm/engine/llm_engine.py and the second call is made at line 1712 in vllm/engine/llm_engine.py. The purpose of these calls is to collect TTFT and TPOT values.

LLMEngine._get_stats function is called by LLMEngine.do_log_stats function present in vllm/engine/llm_engine.py at line 1562, if LLMEngine.log_stats is True. The issue of last_token_time > first_token_time is observed when LLMEngine.log_stats is False, due to logic on line 535 in vllm/engine/llm_engine.py. By default, LLMEngine.log_stats is set to False because EngineArgs.disable_log_stats is set to True in vllm/engine/arg_utils.py at line 125.

Note: LLMEngine.do_log_stats function is called by LLMEngine._process_model_outputs function on present in vllm/engine/llm_engine.py.

Proposal/patch for the issue:

Changes have been made to vllm/sequence.py and vllm/engine/llm_engine.py to fix the issue.

At a high level, instead of setting the value of last_token_latency in LLMEngine._get_stats function, we can set the value of last_token_latency in LLMEngine._process_model_outputs function. This will ensure that the value of last_token_latency is updated correctly. And, in LLMEngine._get_stats function, a new read-only function SequenceGroup.get_last_token_latency is called.

[rajveerb]

Below are two logs, one before fix and at the bottom post fix from running the sample code above.

Before fix:

╰─± python tester.py
INFO 12-20 07:19:11 config.py:478] This model supports multiple tasks: {'score', 'reward', 'embed', 'classify', 'generate'}. Defaulting to 'generate'.
INFO 12-20 07:19:11 llm_engine.py:249] Initializing an LLM engine (v0.6.5) with config: model='facebook/opt-125m', speculative_config=None, tokenizer='facebook/opt-125m', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=2048, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, kv_cache_dtype=auto, quantization_param_path=None, device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar'), observability_config=ObservabilityConfig(otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=facebook/opt-125m, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=False, chunked_prefill_enabled=False, use_async_output_proc=True, mm_cache_preprocessor=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={"splitting_ops":["vllm.unified_attention","vllm.unified_attention_with_output"],"candidate_compile_sizes":[],"compile_sizes":[],"capture_sizes":[256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],"max_capture_size":256}, use_cached_outputs=False, 
INFO 12-20 07:19:12 selector.py:120] Using Flash Attention backend.
INFO 12-20 07:19:13 model_runner.py:1092] Starting to load model facebook/opt-125m...
INFO 12-20 07:19:13 weight_utils.py:243] Using model weights format ['*.bin']
Loading pt checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
vllm/model_executor/model_loader/weight_utils.py:425: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  state = torch.load(bin_file, map_location="cpu")
Loading pt checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  7.05it/s]
Loading pt checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  7.05it/s]

INFO 12-20 07:19:13 model_runner.py:1097] Loading model weights took 0.2389 GB
INFO 12-20 07:19:14 worker.py:241] Memory profiling takes 0.57 seconds
INFO 12-20 07:19:14 worker.py:241] the current vLLM instance can use total_gpu_memory (44.40GiB) x gpu_memory_utilization (0.90) = 39.96GiB
INFO 12-20 07:19:14 worker.py:241] model weights take 0.24GiB; non_torch_memory takes 0.08GiB; PyTorch activation peak memory takes 0.47GiB; the rest of the memory reserved for KV Cache is 39.17GiB.
INFO 12-20 07:19:14 gpu_executor.py:76] # GPU blocks: 71306, # CPU blocks: 7281
INFO 12-20 07:19:14 gpu_executor.py:80] Maximum concurrency for 2048 tokens per request: 557.08x
INFO 12-20 07:19:15 model_runner.py:1413] Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.
INFO 12-20 07:19:15 model_runner.py:1417] If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.
INFO 12-20 07:19:24 model_runner.py:1527] Graph capturing finished in 9 secs, took 0.38 GiB
INFO 12-20 07:19:24 llm_engine.py:446] init engine (profile, create kv cache, warmup model) took 11.21 seconds
Processed prompts: 100%|█| 1/1 [00:00<00:00, 23.50it/s, est. speed input: 141.18 toks/s, output: 376.36 toks/s
Prompt: 'The future of AI is', Generated text: ' in the hands of the people.\n\nThe future of AI is in the'
RequestMetrics(arrival_time=1734697165.6934612, \
last_token_time=1734697165.6934612, \
first_scheduled_time=1734697165.700166, \
first_token_time=1734697165.7171977, \
time_in_queue=0.006704807281494141, \
finished_time=1734697165.7414508, \
scheduler_time=0.0014218189753592014, \
model_forward_time=None, \
model_execute_time=None)
[rank0]:[W1220 07:19:26.300683543 ProcessGroupNCCL.cpp:1250] Warning: WARNING: process group has NOT been destroyed before we destruct ProcessGroupNCCL. On normal program exit, the application should call destroy_process_group to ensure that any pending NCCL operations have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. This constraint has always been present,  but this warning has only been added since PyTorch 2.4 (function operator())

Post fix:

╰─± python tester.py                                                          
INFO 12-20 09:27:59 config.py:478] This model supports multiple tasks: {'classify', 'reward', 'generate', 'score', 'embed'}. Defaulting to 'generate'.
INFO 12-20 09:27:59 llm_engine.py:249] Initializing an LLM engine (v0.6.6.dev0+g2d1b9baa.d20241220) with config: model='facebook/opt-125m', speculative_config=None, tokenizer='facebook/opt-125m', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=2048, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, kv_cache_dtype=auto, quantization_param_path=None, device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar'), observability_config=ObservabilityConfig(otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=facebook/opt-125m, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=False, chunked_prefill_enabled=False, use_async_output_proc=True, mm_cache_preprocessor=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={"splitting_ops":["vllm.unified_attention","vllm.unified_attention_with_output"],"candidate_compile_sizes":[],"compile_sizes":[],"capture_sizes":[256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],"max_capture_size":256}, use_cached_outputs=False, 
INFO 12-20 09:28:00 selector.py:120] Using Flash Attention backend.
INFO 12-20 09:28:01 model_runner.py:1092] Starting to load model facebook/opt-125m...
INFO 12-20 09:28:01 weight_utils.py:243] Using model weights format ['*.bin']
Loading pt checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
vllm/model_executor/model_loader/weight_utils.py:425: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  state = torch.load(bin_file, map_location="cpu")
Loading pt checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  6.84it/s]
Loading pt checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  6.83it/s]

INFO 12-20 09:28:01 model_runner.py:1097] Loading model weights took 0.2389 GB
INFO 12-20 09:28:02 worker.py:241] Memory profiling takes 0.59 seconds
INFO 12-20 09:28:02 worker.py:241] the current vLLM instance can use total_gpu_memory (44.40GiB) x gpu_memory_utilization (0.90) = 39.96GiB
INFO 12-20 09:28:02 worker.py:241] model weights take 0.24GiB; non_torch_memory takes 0.08GiB; PyTorch activation peak memory takes 0.47GiB; the rest of the memory reserved for KV Cache is 39.17GiB.
INFO 12-20 09:28:02 gpu_executor.py:76] # GPU blocks: 71306, # CPU blocks: 7281
INFO 12-20 09:28:02 gpu_executor.py:80] Maximum concurrency for 2048 tokens per request: 557.08x
INFO 12-20 09:28:03 model_runner.py:1413] Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.
INFO 12-20 09:28:03 model_runner.py:1417] If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.
INFO 12-20 09:28:12 model_runner.py:1527] Graph capturing finished in 9 secs, took 0.38 GiB
INFO 12-20 09:28:12 llm_engine.py:446] init engine (profile, create kv cache, warmup model) took 11.16 seconds
Processed prompts: 100%|██████████████████████████████████████████| 1/1 [00:00<00:00, 21.58it/s, est. speed input: 129.61 toks/s, output: 345.54 toks/s]
Prompt: 'The future of AI is', Generated text: ' in the hands of the people.\n\nThe future of AI is in the'
RequestMetrics(arrival_time=1734704893.9175107, \
last_token_time=1734704893.969714, \
first_scheduled_time=1734704893.9247289, \
first_token_time=1734704893.9458623, \
time_in_queue=0.007218122482299805, \
finished_time=1734704893.9697964, \
scheduler_time=0.0014244113117456436, \
model_forward_time=None, \
model_execute_time=None)
[rank0]:[W1220 09:28:14.534092501 ProcessGroupNCCL.cpp:1250] Warning: WARNING: process group has NOT been destroyed before we destruct ProcessGroupNCCL. On normal program exit, the application should call destroy_process_group to ensure that any pending NCCL operations have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. This constraint has always been present,  but this warning has only been added since PyTorch 2.4 (function operator())

[ywang96]

Thanks for finding the root cause and adding this fix! LGTM and I left two nit comments.

[DarkLight1337]

Please merge from main to fix the CI failure in entrypoints tests.

[rajveerb]

Please merge from main to fix the CI failure in entrypoints tests.
@DarkLight1337 @ywang96

Ready for merge I believe. I'm not sure of the two failing tests though.