diff --git a/Dockerfile b/Dockerfile index 454a085..911cbb1 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,9 +1,9 @@ -FROM nvidia/cuda:12.4.1-base-ubuntu22.04 +FROM nvidia/cuda:12.9.0-base-ubuntu22.04 RUN apt-get update -y \ && apt-get install -y python3-pip -RUN ldconfig /usr/local/cuda-12.4/compat/ +RUN ldconfig /usr/local/cuda-12.9/compat/ # Install Python dependencies COPY builder/requirements.txt /requirements.txt @@ -12,7 +12,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \ python3 -m pip install --upgrade -r /requirements.txt # Install vLLM -RUN python3 -m pip install vllm==0.11.0 +RUN python3 -m pip install vllm==0.15.0 # Setup for Option 2: Building the Image with the Model included ARG MODEL_NAME="" @@ -21,6 +21,7 @@ ARG BASE_PATH="/runpod-volume" ARG QUANTIZATION="" ARG MODEL_REVISION="" ARG TOKENIZER_REVISION="" +ARG VLLM_NIGHTLY="true" ENV MODEL_NAME=$MODEL_NAME \ MODEL_REVISION=$MODEL_REVISION \ @@ -35,6 +36,12 @@ ENV MODEL_NAME=$MODEL_NAME \ ENV PYTHONPATH="/:/vllm-workspace" +RUN if [ -n "${VLLM_NIGHTLY}" ]; then \ + pip install -U vllm --pre --index-url https://pypi.org/simple --extra-index-url https://wheels.vllm.ai/nightly && \ + apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/* && \ + pip install git+https://github.com/huggingface/transformers.git; \ +fi + COPY src /src RUN --mount=type=secret,id=HF_TOKEN,required=false \ diff --git a/builder/requirements.txt b/builder/requirements.txt index 2ab29ad..9b1a484 100644 --- a/builder/requirements.txt +++ b/builder/requirements.txt @@ -8,7 +8,7 @@ typing-extensions>=4.8.0 pydantic pydantic-settings hf-transfer -transformers>=4.57.0 +transformers>=4.57.5 bitsandbytes>=0.45.0 kernels -torch==2.6.0 +torch>=2.10.0 diff --git a/src/engine.py b/src/engine.py index 99d9132..1d9f97f 100644 --- a/src/engine.py +++ b/src/engine.py @@ -9,10 +9,13 @@ import time from vllm import AsyncLLMEngine from vllm.entrypoints.logger import RequestLogger -from vllm.entrypoints.openai.serving_chat import OpenAIServingChat -from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion -from vllm.entrypoints.openai.protocol import ChatCompletionRequest, CompletionRequest, ErrorResponse -from vllm.entrypoints.openai.serving_models import BaseModelPath, LoRAModulePath, OpenAIServingModels +from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion +from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.openai.models.protocol import BaseModelPath, LoRAModulePath from utils import DummyRequest, JobInput, BatchSize, create_error_response @@ -202,45 +205,38 @@ class OpenAIvLLMEngine(vLLMEngine): return adapters async def _initialize_engines(self): - self.model_config = await self.llm.get_model_config() self.base_model_paths = [ BaseModelPath(name=self.engine_args.model, model_path=self.engine_args.model) ] self.serving_models = OpenAIServingModels( engine_client=self.llm, - model_config=self.model_config, base_model_paths=self.base_model_paths, lora_modules=self.lora_adapters, ) await self.serving_models.init_static_loras() - + # Get chat template from vLLM tokenizer if available chat_template = None if self.tokenizer and hasattr(self.tokenizer, 'tokenizer'): chat_template = self.tokenizer.tokenizer.chat_template - + self.chat_engine = OpenAIServingChat( - engine_client=self.llm, - model_config=self.model_config, + engine_client=self.llm, models=self.serving_models, response_role=self.response_role, request_logger=None, chat_template=chat_template, chat_template_content_format="auto", - # enable_reasoning=os.getenv('ENABLE_REASONING', 'false').lower() == 'true', - reasoning_parser= os.getenv('REASONING_PARSER', "") or None, - # return_token_as_token_ids=False, + reasoning_parser=os.getenv('REASONING_PARSER', "") or None, enable_auto_tools=os.getenv('ENABLE_AUTO_TOOL_CHOICE', 'false').lower() == 'true', tool_parser=os.getenv('TOOL_CALL_PARSER', "") or None, enable_prompt_tokens_details=False ) self.completion_engine = OpenAIServingCompletion( - engine_client=self.llm, - model_config=self.model_config, + engine_client=self.llm, models=self.serving_models, request_logger=None, - # return_token_as_token_ids=False, ) async def generate(self, openai_request: JobInput): diff --git a/src/engine_args.py b/src/engine_args.py index 7d4586f..3c6b04d 100644 --- a/src/engine_args.py +++ b/src/engine_args.py @@ -38,7 +38,6 @@ DEFAULT_ARGS = { "block_size": int(os.getenv('BLOCK_SIZE', 16)), "enable_prefix_caching": os.getenv('ENABLE_PREFIX_CACHING', 'False').lower() == 'true', "disable_sliding_window": os.getenv('DISABLE_SLIDING_WINDOW', 'False').lower() == 'true', - "use_v2_block_manager": os.getenv('USE_V2_BLOCK_MANAGER', 'False').lower() == 'true', "swap_space": int(os.getenv('SWAP_SPACE', 4)), # GiB "cpu_offload_gb": int(os.getenv('CPU_OFFLOAD_GB', 0)), # GiB "max_num_batched_tokens": int(os.getenv('MAX_NUM_BATCHED_TOKENS', 0)) or None, @@ -71,29 +70,128 @@ DEFAULT_ARGS = { "device": os.getenv('DEVICE', 'auto'), "ray_workers_use_nsight": os.getenv('RAY_WORKERS_USE_NSIGHT', 'False').lower() == 'true', "num_gpu_blocks_override": int(os.getenv('NUM_GPU_BLOCKS_OVERRIDE', 0)) or None, - "num_lookahead_slots": int(os.getenv('NUM_LOOKAHEAD_SLOTS', 0)), "model_loader_extra_config": os.getenv('MODEL_LOADER_EXTRA_CONFIG', None), "ignore_patterns": os.getenv('IGNORE_PATTERNS', None), "preemption_mode": os.getenv('PREEMPTION_MODE', None), "scheduler_delay_factor": float(os.getenv('SCHEDULER_DELAY_FACTOR', 0.0)), "enable_chunked_prefill": os.getenv('ENABLE_CHUNKED_PREFILL', None), "guided_decoding_backend": os.getenv('GUIDED_DECODING_BACKEND', 'outlines'), - "speculative_model": os.getenv('SPECULATIVE_MODEL', None), - "speculative_draft_tensor_parallel_size": int(os.getenv('SPECULATIVE_DRAFT_TENSOR_PARALLEL_SIZE', 0)) or None, "enable_expert_parallel": bool(os.getenv('ENABLE_EXPERT_PARALLEL', 'False').lower() == 'true'), - "num_speculative_tokens": int(os.getenv('NUM_SPECULATIVE_TOKENS', 0)) or None, - "speculative_max_model_len": int(os.getenv('SPECULATIVE_MAX_MODEL_LEN', 0)) or None, - "speculative_disable_by_batch_size": int(os.getenv('SPECULATIVE_DISABLE_BY_BATCH_SIZE', 0)) or None, - "ngram_prompt_lookup_max": int(os.getenv('NGRAM_PROMPT_LOOKUP_MAX', 0)) or None, - "ngram_prompt_lookup_min": int(os.getenv('NGRAM_PROMPT_LOOKUP_MIN', 0)) or None, - "spec_decoding_acceptance_method": os.getenv('SPEC_DECODING_ACCEPTANCE_METHOD', 'rejection_sampler'), - "typical_acceptance_sampler_posterior_threshold": float(os.getenv('TYPICAL_ACCEPTANCE_SAMPLER_POSTERIOR_THRESHOLD', 0)) or None, - "typical_acceptance_sampler_posterior_alpha": float(os.getenv('TYPICAL_ACCEPTANCE_SAMPLER_POSTERIOR_ALPHA', 0)) or None, "qlora_adapter_name_or_path": os.getenv('QLORA_ADAPTER_NAME_OR_PATH', None), - "disable_logprobs_during_spec_decoding": os.getenv('DISABLE_LOGPROBS_DURING_SPEC_DECODING', None), "otlp_traces_endpoint": os.getenv('OTLP_TRACES_ENDPOINT', None), - "use_v2_block_manager": os.getenv('USE_V2_BLOCK_MANAGER', 'true'), } + + +def get_speculative_config(): + """ + Build speculative decoding configuration from environment variables. + + Supports two modes: + 1. Full JSON config via SPECULATIVE_CONFIG env var + 2. Individual env vars for common settings + + Speculative Methods: + - "draft_model": Use a smaller draft model for speculation + - "ngram": Use n-gram based prompt lookup (no additional model needed) + - "eagle" / "eagle3": Use EAGLE-based speculation + - "medusa": Use Medusa heads for speculation + - "mlp_speculator": Use MLP-based speculator + + Returns: + dict | None: Speculative config dictionary or None if not configured + """ + # Option 1: Full JSON configuration + spec_config_json = os.getenv('SPECULATIVE_CONFIG') + if spec_config_json: + try: + config = json.loads(spec_config_json) + logging.info(f"Using speculative config from SPECULATIVE_CONFIG: {config}") + return config + except json.JSONDecodeError as e: + logging.error(f"Failed to parse SPECULATIVE_CONFIG JSON: {e}") + return None + + # Option 2: Build config from individual environment variables + spec_method = os.getenv('SPECULATIVE_METHOD') # ngram, draft_model, eagle, eagle3, medusa, mlp_speculator + spec_model = os.getenv('SPECULATIVE_MODEL') + num_spec_tokens = os.getenv('NUM_SPECULATIVE_TOKENS') + + # N-gram specific settings + ngram_max = os.getenv('NGRAM_PROMPT_LOOKUP_MAX') + ngram_min = os.getenv('NGRAM_PROMPT_LOOKUP_MIN') + + # Check if any speculative decoding is configured + if not any([spec_method, spec_model, ngram_max]): + return None + + config = {} + + # Determine method + if spec_method: + config['method'] = spec_method + elif ngram_max and not spec_model: + config['method'] = 'ngram' + elif spec_model: + # Auto-detect method based on model name if not specified + model_lower = spec_model.lower() + if 'eagle3' in model_lower: + config['method'] = 'eagle3' + elif 'eagle' in model_lower: + config['method'] = 'eagle' + elif 'medusa' in model_lower: + config['method'] = 'medusa' + else: + config['method'] = 'draft_model' + + # Model configuration + if spec_model: + config['model'] = spec_model + + # Number of speculative tokens + if num_spec_tokens: + config['num_speculative_tokens'] = int(num_spec_tokens) + + # N-gram settings + if ngram_max: + config['prompt_lookup_max'] = int(ngram_max) + if ngram_min: + config['prompt_lookup_min'] = int(ngram_min) + + # Draft model tensor parallel size + draft_tp = os.getenv('SPECULATIVE_DRAFT_TENSOR_PARALLEL_SIZE') + if draft_tp: + config['draft_tensor_parallel_size'] = int(draft_tp) + + # Max model length for draft + spec_max_len = os.getenv('SPECULATIVE_MAX_MODEL_LEN') + if spec_max_len: + config['max_model_len'] = int(spec_max_len) + + # Disable by batch size + disable_batch = os.getenv('SPECULATIVE_DISABLE_BY_BATCH_SIZE') + if disable_batch: + config['disable_by_batch_size'] = int(disable_batch) + + # Draft model quantization + spec_quant = os.getenv('SPECULATIVE_QUANTIZATION') + if spec_quant: + config['quantization'] = spec_quant + + # Draft model revision + spec_revision = os.getenv('SPECULATIVE_MODEL_REVISION') + if spec_revision: + config['revision'] = spec_revision + + # Enforce eager mode for draft model + spec_eager = os.getenv('SPECULATIVE_ENFORCE_EAGER') + if spec_eager: + config['enforce_eager'] = spec_eager.lower() == 'true' + + if config: + logging.info(f"Built speculative config from env vars: {config}") + return config + + return None limit_mm_env = os.getenv('LIMIT_MM_PER_PROMPT') if limit_mm_env is not None: DEFAULT_ARGS["limit_mm_per_prompt"] = convert_limit_mm_per_prompt(limit_mm_env) @@ -171,9 +269,10 @@ def get_engine_args(): if os.getenv("MAX_CONTEXT_LEN_TO_CAPTURE"): args["max_seq_len_to_capture"] = int(os.getenv("MAX_CONTEXT_LEN_TO_CAPTURE")) logging.warning("Using MAX_CONTEXT_LEN_TO_CAPTURE is deprecated. Please use MAX_SEQ_LEN_TO_CAPTURE instead.") - - # if "gemma-2" in args.get("model", "").lower(): - # os.environ["VLLM_ATTENTION_BACKEND"] = "FLASHINFER" - # logging.info("Using FLASHINFER for gemma-2 model.") - + + # Add speculative decoding configuration if present + speculative_config = get_speculative_config() + if speculative_config: + args["speculative_config"] = speculative_config + return AsyncEngineArgs(**args) diff --git a/src/utils.py b/src/utils.py index ee1b927..1572f6b 100644 --- a/src/utils.py +++ b/src/utils.py @@ -3,11 +3,10 @@ import logging from http import HTTPStatus from functools import wraps from time import time -from vllm.entrypoints.openai.protocol import RequestResponseMetadata try: from vllm.utils import random_uuid - from vllm.entrypoints.openai.protocol import ErrorResponse + from vllm.entrypoints.openai.engine.protocol import ErrorResponse, RequestResponseMetadata from vllm import SamplingParams except ImportError: logging.warning("Error importing vllm, skipping related imports. This is ONLY expected when baking model into docker image from a machine without GPUs")