209 lines
9.1 KiB
Python
209 lines
9.1 KiB
Python
import os
|
||
import queue
|
||
import signal
|
||
import sys
|
||
import threading
|
||
import time
|
||
from collections import deque
|
||
from collections.abc import Generator
|
||
from concurrent.futures import Future
|
||
from contextlib import ExitStack, contextmanager
|
||
from inspect import isclass, signature
|
||
from logging import DEBUG
|
||
from typing import Any, Callable, Optional, TypeVar, Union
|
||
|
||
import msgspec
|
||
import zmq
|
||
|
||
from vllm.config import ParallelConfig, VllmConfig
|
||
from vllm.logger import init_logger
|
||
from vllm.utils import make_zmq_socket
|
||
|
||
from vllm.v1.kv_cache_interface import KVCacheConfig
|
||
from vllm.v1.engine import (EngineCoreRequest,EngineCoreRequestType)
|
||
from vllm.v1.core.kv_cache_utils import (BlockHash,
|
||
generate_scheduler_kv_cache_config,
|
||
get_kv_cache_configs,
|
||
get_request_block_hasher,
|
||
init_none_hash)
|
||
from vllm.v1.serial_utils import MsgpackDecoder
|
||
from vllm.v1.engine.core import EngineCore
|
||
from vllm.v1.request import Request, RequestStatus
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
POLLING_TIMEOUT_S = 2.5
|
||
HANDSHAKE_TIMEOUT_MINS = 5
|
||
|
||
_R = TypeVar('_R') # Return type for collective_rpc
|
||
|
||
class EngineCoreProc(EngineCore):
|
||
"""ZMQ-wrapper for running EngineCore in background process."""
|
||
|
||
ENGINE_CORE_DEAD = b'ENGINE_CORE_DEAD'
|
||
def process_input_sockets(self, input_addresses: list[str],
|
||
coord_input_address: Optional[str],
|
||
identity: bytes, ready_event: threading.Event):
|
||
"""Input socket IO thread."""
|
||
|
||
# Msgpack serialization decoding.
|
||
add_request_decoder = MsgpackDecoder(EngineCoreRequest)
|
||
generic_decoder = MsgpackDecoder()
|
||
bulk_add_decoder = MsgpackDecoder(list[EngineCoreRequest])
|
||
|
||
with ExitStack() as stack, zmq.Context() as ctx:
|
||
input_sockets = [
|
||
stack.enter_context(
|
||
make_zmq_socket(ctx,
|
||
input_address,
|
||
zmq.DEALER,
|
||
identity=identity,
|
||
bind=False))
|
||
for input_address in input_addresses
|
||
]
|
||
if coord_input_address is None:
|
||
coord_socket = None
|
||
else:
|
||
coord_socket = stack.enter_context(
|
||
make_zmq_socket(ctx,
|
||
coord_input_address,
|
||
zmq.XSUB,
|
||
identity=identity,
|
||
bind=False))
|
||
# Send subscription message to coordinator.
|
||
coord_socket.send(b'\x01')
|
||
|
||
# Register sockets with poller.
|
||
poller = zmq.Poller()
|
||
for input_socket in input_sockets:
|
||
# Send initial message to each input socket - this is required
|
||
# before the front-end ROUTER socket can send input messages
|
||
# back to us.
|
||
input_socket.send(b'')
|
||
poller.register(input_socket, zmq.POLLIN)
|
||
|
||
if coord_socket is not None:
|
||
poller.register(coord_socket, zmq.POLLIN)
|
||
|
||
ready_event.set()
|
||
del ready_event
|
||
while True:
|
||
for input_socket, _ in poller.poll():
|
||
# (RequestType, RequestData)
|
||
type_frame, *data_frames = input_socket.recv_multipart(
|
||
copy=False)
|
||
request_type = EngineCoreRequestType(
|
||
bytes(type_frame.buffer))
|
||
|
||
if request_type == EngineCoreRequestType.ADD_BULK:
|
||
# 关键：按 list[EngineCoreRequest] 解码，然后在接收线程就地 fan-out
|
||
requests = bulk_add_decoder.decode(data_frames)
|
||
for r in requests:
|
||
r = self.preprocess_add_request(r)
|
||
self.input_queue.put_nowait((EngineCoreRequestType.ADD, r))
|
||
continue
|
||
|
||
# Deserialize the request data.
|
||
if request_type == EngineCoreRequestType.ADD:
|
||
request = add_request_decoder.decode(data_frames)
|
||
request = self.preprocess_add_request(request)
|
||
else:
|
||
request = generic_decoder.decode(data_frames)
|
||
|
||
# Push to input queue for core busy loop.
|
||
self.input_queue.put_nowait((request_type, request))
|
||
|
||
|
||
class EngineCore:
|
||
"""Inner loop of vLLM's Engine."""
|
||
|
||
def _initialize_kv_caches(
|
||
self, vllm_config: VllmConfig) -> tuple[int, int, KVCacheConfig]:
|
||
start = time.time()
|
||
|
||
# Get all kv cache needed by the model
|
||
kv_cache_specs = self.model_executor.get_kv_cache_specs()
|
||
|
||
# get_kv_cache_specs in model_runner
|
||
# for layer_name, layer in vllm_config.compilation_config.static_forward_context.items():
|
||
# print(f'layer_name = {layer_name}; layer = {layer}')
|
||
# # 只有 moe layer， 拿不到attention layer?
|
||
# # TODO for no kv cahe model
|
||
|
||
|
||
# has_kv_cache = any(kv_cache_spec for kv_cache_spec in kv_cache_specs)
|
||
|
||
# if has_kv_cache:
|
||
# if os.environ.get("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH") == "1":
|
||
# dp_group = getattr(self, "dp_group", None)
|
||
# assert dp_group is not None
|
||
# self.available_gpu_memory_for_kv_cache = \
|
||
# ParallelConfig.sync_kv_cache_memory_size(dp_group, -1)
|
||
# available_gpu_memory = [
|
||
# self.available_gpu_memory_for_kv_cache
|
||
# ] * len(kv_cache_specs)
|
||
# else:
|
||
# # Profiles the peak memory usage of the model to determine how
|
||
# # much memory can be allocated for kv cache.
|
||
# available_gpu_memory = (
|
||
# self.model_executor.determine_available_memory())
|
||
# self.available_gpu_memory_for_kv_cache = \
|
||
# available_gpu_memory[0]
|
||
# else:
|
||
# # Attention free models don't need memory for kv cache
|
||
# available_gpu_memory = [0] * len(kv_cache_specs)
|
||
|
||
memory_blocks = self.model_executor.determine_available_memory_block() # [(memory, blocks) * rank_number]
|
||
available_gpu_memory = [memory_block[0] for memory_block in memory_blocks]
|
||
num_gpu_blocks = memory_blocks[0][1]
|
||
|
||
assert len(kv_cache_specs) == len(available_gpu_memory)
|
||
|
||
kv_cache_configs = get_kv_cache_configs(vllm_config, kv_cache_specs,
|
||
available_gpu_memory)
|
||
|
||
### patch here to support long seq_length for mtp
|
||
for kv_cache_config in kv_cache_configs:
|
||
for ii in range(len(kv_cache_config.kv_cache_tensors)):
|
||
kv_cache_config.kv_cache_tensors[ii].size = kv_cache_config.kv_cache_tensors[ii].size * num_gpu_blocks // kv_cache_config.num_blocks
|
||
|
||
kv_cache_config.num_blocks = num_gpu_blocks
|
||
### patch here to support long seq_length for mtp end
|
||
scheduler_kv_cache_config = generate_scheduler_kv_cache_config(
|
||
kv_cache_configs)
|
||
num_gpu_blocks = scheduler_kv_cache_config.num_blocks
|
||
num_cpu_blocks = 0
|
||
|
||
# Initialize kv cache and warmup the execution
|
||
self.model_executor.initialize_from_config(kv_cache_configs)
|
||
|
||
elapsed = time.time() - start
|
||
logger.info(("init engine (profile, create kv cache, "
|
||
"warmup model) took %.2f seconds"), elapsed)
|
||
return num_gpu_blocks, num_cpu_blocks, scheduler_kv_cache_config
|
||
|
||
def preprocess_add_request(
|
||
self, request: EngineCoreRequest) -> tuple[Request, int]:
|
||
"""Preprocess the request.
|
||
|
||
This function could be directly used in input processing thread to allow
|
||
request initialization running in parallel with Model forward
|
||
"""
|
||
# Note on thread safety: no race condition.
|
||
# `mm_receiver_cache` is reset at the end of LLMEngine init,
|
||
# and will only be accessed in the input processing thread afterwards.
|
||
if self.mm_receiver_cache is not None and request.mm_features:
|
||
request.mm_features = (
|
||
self.mm_receiver_cache.get_and_update_features(
|
||
request.mm_features))
|
||
|
||
req = Request.from_engine_core_request(request,
|
||
self.request_block_hasher)
|
||
if req.use_structured_output:
|
||
# Note on thread safety: no race condition.
|
||
# `grammar_init` is only invoked in input processing thread. For
|
||
# `structured_output_manager`, each request is independent and
|
||
# grammar compilation is async. Scheduler always checks grammar
|
||
# compilation status before scheduling request.
|
||
self.structured_output_manager.grammar_init(req)
|
||
return req, request.current_wave |