235 lines
9.5 KiB
Python
235 lines
9.5 KiB
Python
#
|
|
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
"""Patches for profiling-based dynamic chunk sizing.
|
|
|
|
This module patches ``EngineCore`` to:
|
|
1. Run profiling at startup (after model_executor is ready).
|
|
2. Record execution timing after each model step to refine the
|
|
history-aware chunk prediction model online.
|
|
|
|
In multiprocessing ``spawn`` mode the child process starts a fresh Python
|
|
interpreter, so class-level monkey-patches applied in the parent are lost.
|
|
To handle this we additionally wrap ``EngineCoreProc.run_engine_core``
|
|
(the subprocess entry-point): when pickle resolves the wrapper it triggers
|
|
an import of this module, which re-applies the ``EngineCore.__init__``
|
|
patches inside the child process before any ``EngineCore`` is instantiated.
|
|
"""
|
|
|
|
from vllm.logger import logger
|
|
from vllm.v1.engine.core import EngineCore, EngineCoreProc
|
|
|
|
from vllm_ascend.utils import vllm_version_is
|
|
|
|
_profiling_patches_applied = False
|
|
_original_update_from_output = None
|
|
_original_schedule = None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helper: record execution timing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _record_execution_timing(scheduler, scheduler_output, model_output):
|
|
"""Record execution timing for online model refinement.
|
|
|
|
Extracts ``execution_time_ms`` (set dynamically by the NPU model runner)
|
|
from the model output and feeds it back to the
|
|
``ProfilingChunkManager`` for incremental fitting of the history-aware
|
|
latency model.
|
|
"""
|
|
profiling_mgr = getattr(scheduler, "profiling_chunk_manager", None)
|
|
SET_TIME_COUNT = 3
|
|
if profiling_mgr is None or not profiling_mgr.is_ready:
|
|
return
|
|
|
|
# Once both the target latency and history model are calibrated,
|
|
# stop collecting timing data and disable the synchronize-and-time
|
|
# calls in the model runner to avoid unnecessary pipeline stalls.
|
|
if profiling_mgr._set_time_done and profiling_mgr.predictor.history_fitted:
|
|
try:
|
|
from vllm_ascend.ascend_config import get_ascend_config
|
|
|
|
get_ascend_config().profiling_chunk_config.need_timing = False
|
|
except RuntimeError:
|
|
pass
|
|
# Mark the scheduler so that the next scheduler_output carries
|
|
# a ``disable_profiling_timing`` flag to the worker process,
|
|
# which will set its own process-local need_timing to False.
|
|
scheduler._profiling_timing_done = True
|
|
return
|
|
|
|
elapsed_time_ms = getattr(model_output, "execution_time_ms", 0.0)
|
|
if elapsed_time_ms <= 0:
|
|
return
|
|
elapsed_time = elapsed_time_ms / 1000.0
|
|
|
|
try:
|
|
total_tokens = getattr(scheduler_output, "total_num_scheduled_tokens", 0)
|
|
if total_tokens <= 0:
|
|
return
|
|
|
|
num_scheduled_tokens = getattr(scheduler_output, "num_scheduled_tokens", {})
|
|
request_chunks = []
|
|
|
|
total_hist_tokens = 0
|
|
new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
|
|
for req in new_reqs:
|
|
req_id = getattr(req, "request_id", None) or getattr(req, "req_id", None)
|
|
if req_id and req_id in num_scheduled_tokens:
|
|
chunk_size = num_scheduled_tokens[req_id]
|
|
hist_seq_len = getattr(req, "num_computed_tokens", 0)
|
|
total_hist_tokens += hist_seq_len
|
|
if chunk_size > 0:
|
|
request_chunks.append((chunk_size, hist_seq_len))
|
|
|
|
cached_reqs = getattr(scheduler_output, "scheduled_cached_reqs", None)
|
|
if cached_reqs is not None:
|
|
req_ids = getattr(cached_reqs, "req_ids", [])
|
|
computed_tokens_list = getattr(cached_reqs, "num_computed_tokens", [])
|
|
for i, req_id in enumerate(req_ids):
|
|
if req_id in num_scheduled_tokens:
|
|
chunk_size = num_scheduled_tokens[req_id]
|
|
hist_seq_len = computed_tokens_list[i] if i < len(computed_tokens_list) else 0
|
|
total_hist_tokens += hist_seq_len
|
|
if chunk_size > 0:
|
|
request_chunks.append((chunk_size, hist_seq_len))
|
|
|
|
# is first chunk processing — collect 3 samples before marking done
|
|
if total_hist_tokens == 0 and not profiling_mgr._set_time_done:
|
|
profiling_mgr.predictor.set_target_latency(0, elapsed_time * 1000)
|
|
profiling_mgr._set_time_count += 1
|
|
if profiling_mgr._set_time_count >= SET_TIME_COUNT:
|
|
profiling_mgr._set_time_done = True
|
|
|
|
if not request_chunks:
|
|
# Cannot accurately attribute batch latency to individual
|
|
# requests — skip this sample to avoid polluting the model.
|
|
logger.debug("[ProfilingChunk] Skipping timing sample: unable to extract per-request chunk info")
|
|
return
|
|
|
|
if not profiling_mgr.predictor.history_fitted:
|
|
profiling_mgr.record_batch_execution_time(request_chunks, elapsed_time)
|
|
|
|
except (AttributeError, TypeError) as e:
|
|
logger.debug("Failed to record execution timing: %s", e)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helper: wrap scheduler.update_from_output for timing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _ensure_update_from_output_wrapped(scheduler):
|
|
"""Wrap scheduler.update_from_output to record execution timing."""
|
|
global _original_update_from_output
|
|
if _original_update_from_output is not None:
|
|
return
|
|
if not hasattr(scheduler, "profiling_chunk_manager"):
|
|
return
|
|
|
|
cls = type(scheduler)
|
|
_original_update_from_output = cls.update_from_output
|
|
|
|
def _wrapped_update_from_output(self, scheduler_output, model_output):
|
|
_record_execution_timing(self, scheduler_output, model_output)
|
|
return _original_update_from_output(self, scheduler_output, model_output)
|
|
|
|
cls.update_from_output = _wrapped_update_from_output
|
|
|
|
|
|
def _ensure_schedule_wrapped(scheduler):
|
|
"""Wrap scheduler.schedule to propagate timing-done signal via scheduler_output.
|
|
|
|
When ``_record_execution_timing`` detects that calibration is complete, it
|
|
sets ``scheduler._profiling_timing_done = True``. This wrapper copies that
|
|
flag onto every subsequent ``SchedulerOutput`` so the worker process can
|
|
read it and disable its own process-local ``need_timing``.
|
|
"""
|
|
global _original_schedule
|
|
if _original_schedule is not None:
|
|
return
|
|
if not hasattr(scheduler, "profiling_chunk_manager"):
|
|
return
|
|
|
|
cls = type(scheduler)
|
|
_original_schedule = cls.schedule
|
|
|
|
def _wrapped_schedule(self, throttle_prefills: bool = False):
|
|
if vllm_version_is("0.23.0"):
|
|
output = _original_schedule(self)
|
|
else:
|
|
output = _original_schedule(self, throttle_prefills)
|
|
if getattr(self, "_profiling_timing_done", False) and output is not None:
|
|
output.disable_profiling_timing = True
|
|
return output
|
|
|
|
cls.schedule = _wrapped_schedule
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Core: apply EngineCore.__init__ patches (idempotent)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _apply_profiling_patches():
|
|
"""Patch ``EngineCore.__init__`` to trigger profiling and timing hooks.
|
|
|
|
Safe to call multiple times; the guard ``_profiling_patches_applied``
|
|
ensures the patch is applied at most once per process.
|
|
"""
|
|
global _profiling_patches_applied
|
|
if _profiling_patches_applied:
|
|
return
|
|
_profiling_patches_applied = True
|
|
|
|
original_init = EngineCore.__init__
|
|
|
|
def _patched_engine_core_init(self, *args, **kwargs):
|
|
original_init(self, *args, **kwargs)
|
|
|
|
if hasattr(self.scheduler, "run_profiling_chunk_init"):
|
|
logger.info("[ProfilingChunk] Running profiling initialization...")
|
|
self.scheduler.run_profiling_chunk_init(self.model_executor)
|
|
|
|
_ensure_update_from_output_wrapped(self.scheduler)
|
|
_ensure_schedule_wrapped(self.scheduler)
|
|
|
|
EngineCore.__init__ = _patched_engine_core_init
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 1. Apply patches at module level for the InprocClient (in-process) path.
|
|
# ---------------------------------------------------------------------------
|
|
_apply_profiling_patches()
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 2. Wrap EngineCoreProc.run_engine_core so that spawned subprocesses
|
|
# re-apply the patches. When the child unpickles this wrapper it
|
|
# imports this module, which triggers _apply_profiling_patches() above,
|
|
# ensuring EngineCore.__init__ is patched before any instance is created.
|
|
# ---------------------------------------------------------------------------
|
|
_original_run_engine_core = EngineCoreProc.run_engine_core
|
|
|
|
|
|
def _patched_run_engine_core(*args, **kwargs):
|
|
_apply_profiling_patches()
|
|
return _original_run_engine_core(*args, **kwargs)
|
|
|
|
|
|
EngineCoreProc.run_engine_core = _patched_run_engine_core
|