mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-12 16:58:14 +00:00
[Hardware][GPU] Profiler config additional to increase it scope and annotation details (#37524)
Signed-off-by: devalshahamd <[email protected]> Signed-off-by: Deval Shah <[email protected]> Signed-off-by: Deval Shah <[email protected]> Co-authored-by: Deval Shah <[email protected]>
This commit is contained in:
devalshahamd
co-authored by
Deval Shah
Deval Shah
parent
b5433b6f50
commit
fcd2255d16
@@ -1,10 +1,15 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config import ProfilerConfig
|
||||
from vllm.config import CUDAGraphMode, ProfilerConfig
|
||||
from vllm.config.profiler import _is_uri_path
|
||||
from vllm.profiler.wrapper import WorkerProfiler
|
||||
from vllm.v1.core.sched.output import CachedRequestData
|
||||
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
|
||||
from vllm.v1.worker.gpu_worker import Worker
|
||||
|
||||
|
||||
class ConcreteWorkerProfiler(WorkerProfiler):
|
||||
@@ -236,3 +241,60 @@ class TestIsUriPath:
|
||||
def test_is_uri_path(self, path, expected):
|
||||
"""Test that _is_uri_path correctly identifies URI vs local paths."""
|
||||
assert _is_uri_path(path) == expected
|
||||
|
||||
|
||||
class TestAnnotateProfile:
|
||||
"""Tests for Worker.annotate_profile() annotation string formatting."""
|
||||
|
||||
def _annotate(self, detailed: bool) -> str:
|
||||
worker = MagicMock()
|
||||
worker.vllm_config.profiler_config.detailed_trace_annotation = detailed
|
||||
worker.profiler = MagicMock()
|
||||
|
||||
ctx_req = MagicMock(req_id="ctx1", num_computed_tokens=0)
|
||||
cached = CachedRequestData(
|
||||
req_ids=["gen1"],
|
||||
resumed_req_ids=set(),
|
||||
new_token_ids=[],
|
||||
all_token_ids={},
|
||||
new_block_ids=[],
|
||||
num_computed_tokens=[10],
|
||||
num_output_tokens=[1],
|
||||
)
|
||||
sched = MagicMock(
|
||||
scheduled_new_reqs=[ctx_req],
|
||||
scheduled_cached_reqs=cached,
|
||||
num_scheduled_tokens={"ctx1": 4, "gen1": 1},
|
||||
)
|
||||
|
||||
Worker.annotate_profile(worker, sched)
|
||||
return worker.profiler.annotate_context_manager.call_args[0][0]
|
||||
|
||||
def test_simple_format_mixed(self):
|
||||
assert self._annotate(detailed=False) == (
|
||||
"execute_context_1(4)_generation_1(1)"
|
||||
)
|
||||
|
||||
def test_detailed_format_mixed(self):
|
||||
# ctx1: sq=4, sk=4, sqsq=16, sqsk=16 | gen1: sq=1, sk=11, sqsq=1, sqsk=11 | bs=5
|
||||
assert self._annotate(detailed=True) == (
|
||||
"execute_5_context_1(sq4sk4sqsq16sqsk16)_generation_1(sq1sk11sqsq1sqsk11)"
|
||||
)
|
||||
|
||||
|
||||
def test_profiler_entered_during_capture():
|
||||
"""Profiler is used as a context manager in _warmup_and_capture,
|
||||
confirming it is active during the actual graph capture run."""
|
||||
runner = MagicMock()
|
||||
runner.compilation_config.cudagraph_num_of_warmups = 0
|
||||
mock_profiler = MagicMock()
|
||||
|
||||
GPUModelRunner._warmup_and_capture(
|
||||
runner,
|
||||
desc=MagicMock(num_tokens=4, uniform=True),
|
||||
cudagraph_runtime_mode=CUDAGraphMode.FULL,
|
||||
profiler=mock_profiler,
|
||||
)
|
||||
|
||||
mock_profiler.__enter__.assert_called_once()
|
||||
mock_profiler.__exit__.assert_called_once()
|
||||
|
||||
@@ -66,6 +66,17 @@ class ProfilerConfig:
|
||||
"""If `True`, enables memory profiling in the torch profiler.
|
||||
Disabled by default."""
|
||||
|
||||
capture_torch_profiler: bool = False
|
||||
"""If `True`, enables a torch profiler during CUDA graph capture on rank 0.
|
||||
Traces are saved to a `capture_traces` subdirectory under `torch_profiler_dir`.
|
||||
Requires `profiler` to be set to 'torch'."""
|
||||
|
||||
detailed_trace_annotation: bool = False
|
||||
"""If `True`, uses detailed annotations with roofline metrics (sk, sqsq,
|
||||
sqsk) in profiler trace events. If `False`, uses simple annotations with
|
||||
only context/generation request counts and token counts.
|
||||
Disabled by default."""
|
||||
|
||||
ignore_frontend: bool = False
|
||||
"""If `True`, disables the front-end profiling of AsyncLLM when using the
|
||||
'torch' profiler. This is needed to reduce overhead when using delay/limit options,
|
||||
@@ -144,4 +155,10 @@ class ProfilerConfig:
|
||||
if profiler_dir and not _is_uri_path(profiler_dir):
|
||||
self.torch_profiler_dir = os.path.abspath(os.path.expanduser(profiler_dir))
|
||||
|
||||
if self.capture_torch_profiler and self.profiler != "torch":
|
||||
raise ValueError(
|
||||
"capture_torch_profiler is only applicable when profiler is "
|
||||
"set to 'torch'"
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
@@ -8,7 +8,7 @@ import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable, Iterable, Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from contextlib import AbstractContextManager, contextmanager, nullcontext
|
||||
from copy import copy, deepcopy
|
||||
from dataclasses import dataclass, replace
|
||||
from functools import reduce
|
||||
@@ -6756,6 +6756,44 @@ class GPUModelRunner(
|
||||
# Capture the large shapes first so that the smaller shapes
|
||||
# can reuse the memory pool allocated for the large shapes.
|
||||
set_cudagraph_capturing_enabled(True)
|
||||
|
||||
# Setup torch profiler for graph capture traces (conditional)
|
||||
from vllm.distributed.parallel_state import get_world_group
|
||||
|
||||
local_rank = get_world_group().local_rank
|
||||
enable_profiler = (
|
||||
local_rank == 0
|
||||
) and self.vllm_config.profiler_config.capture_torch_profiler
|
||||
if enable_profiler:
|
||||
trace_dir = (
|
||||
self.vllm_config.profiler_config.torch_profiler_dir + "/capture_traces"
|
||||
)
|
||||
profiler = torch.profiler.profile(
|
||||
activities=[
|
||||
torch.profiler.ProfilerActivity.CPU,
|
||||
torch.profiler.ProfilerActivity.CUDA,
|
||||
],
|
||||
record_shapes=True,
|
||||
profile_memory=True,
|
||||
with_stack=True,
|
||||
on_trace_ready=torch.profiler.tensorboard_trace_handler(
|
||||
trace_dir,
|
||||
worker_name=f"graph_capture_rank_{local_rank}",
|
||||
use_gzip=True,
|
||||
),
|
||||
)
|
||||
logger.info_once(
|
||||
"Rank %d: Torch profiler enabled for CUDA graph capture, "
|
||||
"traces will be saved to: %s",
|
||||
local_rank,
|
||||
trace_dir,
|
||||
)
|
||||
else:
|
||||
profiler = nullcontext()
|
||||
logger.info_once(
|
||||
"Rank %d: Torch profiler disabled for CUDA graph capture", local_rank
|
||||
)
|
||||
|
||||
with self._freeze_gc(), graph_capture(device=self.device):
|
||||
torch.accelerator.synchronize()
|
||||
torch.accelerator.empty_cache()
|
||||
@@ -6768,6 +6806,7 @@ class GPUModelRunner(
|
||||
self._capture_cudagraphs(
|
||||
batch_descriptors=batch_descs,
|
||||
cudagraph_runtime_mode=runtime_mode,
|
||||
profiler=profiler,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
@@ -6811,7 +6850,10 @@ class GPUModelRunner(
|
||||
profile_seq_lens: int | None = None,
|
||||
allow_microbatching: bool = False,
|
||||
num_warmups: int | None = None,
|
||||
profiler: AbstractContextManager[Any] | None = None,
|
||||
):
|
||||
if profiler is None:
|
||||
profiler = nullcontext()
|
||||
if num_warmups is None:
|
||||
num_warmups = self.compilation_config.cudagraph_num_of_warmups
|
||||
force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL
|
||||
@@ -6827,22 +6869,29 @@ class GPUModelRunner(
|
||||
num_active_loras=desc.num_active_loras,
|
||||
profile_seq_lens=profile_seq_lens,
|
||||
)
|
||||
self._dummy_run(
|
||||
desc.num_tokens,
|
||||
cudagraph_runtime_mode=cudagraph_runtime_mode,
|
||||
uniform_decode=desc.uniform,
|
||||
allow_microbatching=allow_microbatching,
|
||||
skip_eplb=True,
|
||||
remove_lora=False,
|
||||
num_active_loras=desc.num_active_loras,
|
||||
is_graph_capturing=True,
|
||||
profile_seq_lens=profile_seq_lens,
|
||||
)
|
||||
with (
|
||||
profiler,
|
||||
torch.profiler.record_function(
|
||||
f"capture_{desc.num_tokens}_{cudagraph_runtime_mode.name}"
|
||||
),
|
||||
):
|
||||
self._dummy_run(
|
||||
desc.num_tokens,
|
||||
cudagraph_runtime_mode=cudagraph_runtime_mode,
|
||||
uniform_decode=desc.uniform,
|
||||
allow_microbatching=allow_microbatching,
|
||||
skip_eplb=True,
|
||||
remove_lora=False,
|
||||
num_active_loras=desc.num_active_loras,
|
||||
is_graph_capturing=True,
|
||||
profile_seq_lens=profile_seq_lens,
|
||||
)
|
||||
|
||||
def _capture_cudagraphs(
|
||||
self,
|
||||
batch_descriptors: list[BatchDescriptor],
|
||||
cudagraph_runtime_mode: CUDAGraphMode,
|
||||
profiler: AbstractContextManager[Any] | None = None,
|
||||
):
|
||||
assert (
|
||||
cudagraph_runtime_mode != CUDAGraphMode.NONE
|
||||
@@ -6885,6 +6934,7 @@ class GPUModelRunner(
|
||||
batch_desc,
|
||||
cudagraph_runtime_mode=cudagraph_runtime_mode,
|
||||
allow_microbatching=allow_microbatching,
|
||||
profiler=profiler,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
self.maybe_remove_all_loras(self.lora_config)
|
||||
|
||||
@@ -975,19 +975,104 @@ class Worker(WorkerBase):
|
||||
|
||||
iteration_details = compute_iteration_details(scheduler_output)
|
||||
|
||||
annotation = "".join(
|
||||
[
|
||||
"execute_context_",
|
||||
str(iteration_details.num_ctx_requests),
|
||||
"(",
|
||||
str(iteration_details.num_ctx_tokens),
|
||||
")_generation_",
|
||||
str(iteration_details.num_generation_requests),
|
||||
"(",
|
||||
str(iteration_details.num_generation_tokens),
|
||||
")",
|
||||
]
|
||||
)
|
||||
if self.vllm_config.profiler_config.detailed_trace_annotation:
|
||||
# Compute roofline-model metrics per request, split by phase
|
||||
# (context vs generation). These help estimate compute and
|
||||
# memory intensity from the trace.
|
||||
#
|
||||
# Per-request quantities:
|
||||
# query_len = number of scheduled (new) tokens for this request
|
||||
# seq_len = total sequence length (computed + scheduled tokens)
|
||||
#
|
||||
# Aggregated across requests in each phase
|
||||
# (ctx_=context, gen_=generation):
|
||||
# seq_len_sum = sum of seq_len (total KV length)
|
||||
# qq_compute = sum of query_len*query_len
|
||||
# (proxy for QK^T compute cost)
|
||||
# qk_compute = sum of query_len*seq_len
|
||||
# (proxy for QK^T compute cost for decode and
|
||||
# chunked prefill)
|
||||
# total_scheduled_tokens = scheduled tokens across all requests
|
||||
ctx_seq_len_sum = 0
|
||||
ctx_qq_compute = 0
|
||||
ctx_qk_compute = 0
|
||||
gen_seq_len_sum = 0
|
||||
gen_qq_compute = 0
|
||||
gen_qk_compute = 0
|
||||
total_scheduled_tokens = 0
|
||||
|
||||
# Build a map of req_id -> num_computed_tokens for all requests
|
||||
new_req_ids = {
|
||||
new_req.req_id for new_req in scheduler_output.scheduled_new_reqs
|
||||
}
|
||||
num_computed_tokens_ids = {
|
||||
new_req.req_id: new_req.num_computed_tokens
|
||||
for new_req in scheduler_output.scheduled_new_reqs
|
||||
}
|
||||
for req_id, num_computed_tokens in zip(
|
||||
scheduler_output.scheduled_cached_reqs.req_ids,
|
||||
scheduler_output.scheduled_cached_reqs.num_computed_tokens,
|
||||
):
|
||||
num_computed_tokens_ids[req_id] = num_computed_tokens
|
||||
|
||||
# Accumulate per-phase metrics
|
||||
for req_id, num_tokens in scheduler_output.num_scheduled_tokens.items():
|
||||
query_len = num_tokens
|
||||
total_scheduled_tokens += query_len
|
||||
seq_len = num_computed_tokens_ids.get(req_id, 0) + query_len
|
||||
if (
|
||||
scheduler_output.scheduled_cached_reqs.is_context_phase(req_id)
|
||||
or req_id in new_req_ids
|
||||
):
|
||||
ctx_seq_len_sum += seq_len
|
||||
ctx_qq_compute += query_len * query_len
|
||||
ctx_qk_compute += query_len * seq_len
|
||||
else:
|
||||
gen_seq_len_sum += seq_len
|
||||
gen_qq_compute += query_len * query_len
|
||||
gen_qk_compute += query_len * seq_len
|
||||
annotation = "".join(
|
||||
[
|
||||
"execute_",
|
||||
str(total_scheduled_tokens),
|
||||
"_context_",
|
||||
str(iteration_details.num_ctx_requests),
|
||||
"(sq",
|
||||
str(iteration_details.num_ctx_tokens),
|
||||
"sk",
|
||||
str(ctx_seq_len_sum),
|
||||
"sqsq",
|
||||
str(ctx_qq_compute),
|
||||
"sqsk",
|
||||
str(ctx_qk_compute),
|
||||
")_generation_",
|
||||
str(iteration_details.num_generation_requests),
|
||||
"(sq",
|
||||
str(iteration_details.num_generation_tokens),
|
||||
"sk",
|
||||
str(gen_seq_len_sum),
|
||||
"sqsq",
|
||||
str(gen_qq_compute),
|
||||
"sqsk",
|
||||
str(gen_qk_compute),
|
||||
")",
|
||||
]
|
||||
)
|
||||
else:
|
||||
annotation = "".join(
|
||||
[
|
||||
"execute_context_",
|
||||
str(iteration_details.num_ctx_requests),
|
||||
"(",
|
||||
str(iteration_details.num_ctx_tokens),
|
||||
")",
|
||||
"_generation_",
|
||||
str(iteration_details.num_generation_requests),
|
||||
"(",
|
||||
str(iteration_details.num_generation_tokens),
|
||||
")",
|
||||
]
|
||||
)
|
||||
return self.profiler.annotate_context_manager(annotation)
|
||||
|
||||
@torch.inference_mode()
|
||||
|
||||
Reference in New Issue
Block a user