diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 00000000..61c49910 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,10 @@ +# Force LF line endings for all text files +* text=auto eol=lf +*.py text eol=lf +*.sh text eol=lf +*.cu text eol=lf +*.cuh text eol=lf +*.yaml text eol=lf +*.yml text eol=lf +*.md text eol=lf +Dockerfile text eol=lf diff --git a/qwen3_6_scripts/api_server.py b/qwen3_6_scripts/api_server.py index fa5ef301..d63fc4b3 100644 --- a/qwen3_6_scripts/api_server.py +++ b/qwen3_6_scripts/api_server.py @@ -1,18 +1,18 @@ -import asyncio -import importlib -import inspect -import multiprocessing -import os -import regex as re -import signal -import socket +import asyncio +import importlib +import inspect +import multiprocessing +import os +import regex as re +import signal +import socket import sys -import tempfile +import tempfile import time -from argparse import Namespace -from contextlib import asynccontextmanager -from functools import partial -from http import HTTPStatus +from argparse import Namespace +from contextlib import asynccontextmanager +from functools import partial +from http import HTTPStatus from typing import AsyncIterator, Set @@ -365,59 +365,59 @@ def _bi100_startup_trace(message: str) -> None: _bi100_startup_trace("api_server stdlib imports complete; loading runtime dependencies") -import uvloop -from fastapi import APIRouter, FastAPI, Request -from fastapi.exceptions import RequestValidationError -from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import JSONResponse, Response, StreamingResponse -from starlette.datastructures import State -from starlette.routing import Mount -from typing_extensions import assert_never - -import vllm.envs as envs -from vllm.config import ModelConfig -from vllm.engine.arg_utils import AsyncEngineArgs -from vllm.engine.async_llm_engine import AsyncLLMEngine -from vllm.engine.multiprocessing.client import MQLLMEngineClient -from vllm.engine.multiprocessing.engine import run_mp_engine -from vllm.engine.protocol import EngineClient -from vllm.entrypoints.launcher import serve_http -from vllm.entrypoints.logger import RequestLogger -from vllm.entrypoints.openai.cli_args import (make_arg_parser, - validate_parsed_serve_args) -# yapf conflicts with isort for this block -# yapf: disable -from vllm.entrypoints.openai.protocol import (ChatCompletionRequest, - ChatCompletionResponse, - CompletionRequest, - CompletionResponse, - DetokenizeRequest, - DetokenizeResponse, - EmbeddingRequest, - EmbeddingResponse, ErrorResponse, - LoadLoraAdapterRequest, - TokenizeRequest, - TokenizeResponse, - UnloadLoraAdapterRequest) -# yapf: enable -from vllm.entrypoints.openai.serving_chat import OpenAIServingChat -from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion -from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding -from vllm.entrypoints.openai.serving_engine import BaseModelPath -from vllm.entrypoints.openai.serving_tokenization import ( - OpenAIServingTokenization) -from vllm.entrypoints.openai.tool_parsers import ToolParserManager -from vllm.reasoning import ReasoningParserManager -from vllm.logger import init_logger -from vllm.usage.usage_lib import UsageContext -from vllm.utils import FlexibleArgumentParser, get_open_zmq_ipc_path -from vllm.version import __version__ as VLLM_VERSION - -TIMEOUT_KEEP_ALIVE = 5 # seconds - -prometheus_multiproc_dir: tempfile.TemporaryDirectory - -# Cannot use __name__ (https://github.com/vllm-project/vllm/pull/4765) +import uvloop +from fastapi import APIRouter, FastAPI, Request +from fastapi.exceptions import RequestValidationError +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import JSONResponse, Response, StreamingResponse +from starlette.datastructures import State +from starlette.routing import Mount +from typing_extensions import assert_never + +import vllm.envs as envs +from vllm.config import ModelConfig +from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.engine.async_llm_engine import AsyncLLMEngine +from vllm.engine.multiprocessing.client import MQLLMEngineClient +from vllm.engine.multiprocessing.engine import run_mp_engine +from vllm.engine.protocol import EngineClient +from vllm.entrypoints.launcher import serve_http +from vllm.entrypoints.logger import RequestLogger +from vllm.entrypoints.openai.cli_args import (make_arg_parser, + validate_parsed_serve_args) +# yapf conflicts with isort for this block +# yapf: disable +from vllm.entrypoints.openai.protocol import (ChatCompletionRequest, + ChatCompletionResponse, + CompletionRequest, + CompletionResponse, + DetokenizeRequest, + DetokenizeResponse, + EmbeddingRequest, + EmbeddingResponse, ErrorResponse, + LoadLoraAdapterRequest, + TokenizeRequest, + TokenizeResponse, + UnloadLoraAdapterRequest) +# yapf: enable +from vllm.entrypoints.openai.serving_chat import OpenAIServingChat +from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion +from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding +from vllm.entrypoints.openai.serving_engine import BaseModelPath +from vllm.entrypoints.openai.serving_tokenization import ( + OpenAIServingTokenization) +from vllm.entrypoints.openai.tool_parsers import ToolParserManager +from vllm.reasoning import ReasoningParserManager +from vllm.logger import init_logger +from vllm.usage.usage_lib import UsageContext +from vllm.utils import FlexibleArgumentParser, get_open_zmq_ipc_path +from vllm.version import __version__ as VLLM_VERSION + +TIMEOUT_KEEP_ALIVE = 5 # seconds + +prometheus_multiproc_dir: tempfile.TemporaryDirectory + +# Cannot use __name__ (https://github.com/vllm-project/vllm/pull/4765) logger = init_logger('vllm.entrypoints.openai.api_server') _running_tasks: Set[asyncio.Task] = set() @@ -543,249 +543,249 @@ def _bi100_log_request_validation_4xx(raw_request, exc) -> None: @asynccontextmanager async def lifespan(app: FastAPI): - try: - if app.state.log_stats: - engine_client: EngineClient = app.state.engine_client - - async def _force_log(): - while True: - await asyncio.sleep(10.) - await engine_client.do_log_stats() - - task = asyncio.create_task(_force_log()) - _running_tasks.add(task) - task.add_done_callback(_running_tasks.remove) - else: - task = None - try: - yield - finally: - if task is not None: - task.cancel() - finally: - # Ensure app state including engine ref is gc'd - del app.state - - -@asynccontextmanager -async def build_async_engine_client( - args: Namespace) -> AsyncIterator[EngineClient]: - + try: + if app.state.log_stats: + engine_client: EngineClient = app.state.engine_client + + async def _force_log(): + while True: + await asyncio.sleep(10.) + await engine_client.do_log_stats() + + task = asyncio.create_task(_force_log()) + _running_tasks.add(task) + task.add_done_callback(_running_tasks.remove) + else: + task = None + try: + yield + finally: + if task is not None: + task.cancel() + finally: + # Ensure app state including engine ref is gc'd + del app.state + + +@asynccontextmanager +async def build_async_engine_client( + args: Namespace) -> AsyncIterator[EngineClient]: + _bi100_startup_trace("building AsyncEngineArgs") - # Context manager to handle engine_client lifecycle - # Ensures everything is shutdown and cleaned up on error/exit - engine_args = AsyncEngineArgs.from_cli_args(args) - + # Context manager to handle engine_client lifecycle + # Ensures everything is shutdown and cleaned up on error/exit + engine_args = AsyncEngineArgs.from_cli_args(args) + _bi100_startup_trace("entering engine client construction") - async with build_async_engine_client_from_engine_args( - engine_args, args.disable_frontend_multiprocessing) as engine: + async with build_async_engine_client_from_engine_args( + engine_args, args.disable_frontend_multiprocessing) as engine: _bi100_startup_trace("engine client construction completed") - yield engine - - -@asynccontextmanager -async def build_async_engine_client_from_engine_args( - engine_args: AsyncEngineArgs, - disable_frontend_multiprocessing: bool = False, -) -> AsyncIterator[EngineClient]: - """ - Create EngineClient, either: - - in-process using the AsyncLLMEngine Directly - - multiprocess using AsyncLLMEngine RPC - - Returns the Client or None if the creation failed. - """ - - # Fall back - # TODO: fill out feature matrix. - if (MQLLMEngineClient.is_unsupported_config(engine_args) - or disable_frontend_multiprocessing): - engine_config = engine_args.create_engine_config() - uses_ray = getattr(AsyncLLMEngine._get_executor_cls(engine_config), - "uses_ray", False) - - build_engine = partial(AsyncLLMEngine.from_engine_args, - engine_args=engine_args, - engine_config=engine_config, - usage_context=UsageContext.OPENAI_API_SERVER) - if uses_ray: - # Must run in main thread with ray for its signal handlers to work - engine_client = build_engine() - else: - engine_client = await asyncio.get_running_loop().run_in_executor( - None, build_engine) - - yield engine_client - return - - # Otherwise, use the multiprocessing AsyncLLMEngine. - else: - if "PROMETHEUS_MULTIPROC_DIR" not in os.environ: - # Make TemporaryDirectory for prometheus multiprocessing - # Note: global TemporaryDirectory will be automatically - # cleaned up upon exit. - global prometheus_multiproc_dir - prometheus_multiproc_dir = tempfile.TemporaryDirectory() - os.environ[ - "PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir.name - else: - logger.warning( - "Found PROMETHEUS_MULTIPROC_DIR was set by user. " - "This directory must be wiped between vLLM runs or " - "you will find inaccurate metrics. Unset the variable " - "and vLLM will properly handle cleanup.") - - # Select random path for IPC. - ipc_path = get_open_zmq_ipc_path() - logger.info("Multiprocessing frontend to use %s for IPC Path.", - ipc_path) - - # Start RPCServer in separate process (holds the LLMEngine). - # the current process might have CUDA context, - # so we need to spawn a new process - context = multiprocessing.get_context("spawn") - - engine_process = context.Process(target=run_mp_engine, - args=(engine_args, - UsageContext.OPENAI_API_SERVER, - ipc_path)) - engine_process.start() - logger.info("Started engine process with PID %d", engine_process.pid) - - # Build RPCClient, which conforms to EngineClient Protocol. - # NOTE: Actually, this is not true yet. We still need to support - # embedding models via RPC (see TODO above) - engine_config = engine_args.create_engine_config() - mp_engine_client = MQLLMEngineClient(ipc_path, engine_config) - - try: - while True: - try: - await mp_engine_client.setup() - break - except TimeoutError: - if not engine_process.is_alive(): - raise RuntimeError( - "Engine process failed to start") from None - - yield mp_engine_client # type: ignore[misc] - finally: - # Ensure rpc server process was terminated - engine_process.terminate() - - # Close all open connections to the backend - mp_engine_client.close() - - # Wait for engine process to join - engine_process.join(4) - if engine_process.exitcode is None: - # Kill if taking longer than 5 seconds to stop - engine_process.kill() - - # Lazy import for prometheus multiprocessing. - # We need to set PROMETHEUS_MULTIPROC_DIR environment variable - # before prometheus_client is imported. - # See https://prometheus.github.io/client_python/multiprocess/ - from prometheus_client import multiprocess - multiprocess.mark_process_dead(engine_process.pid) - - -router = APIRouter() - - -def mount_metrics(app: FastAPI): - # Lazy import for prometheus multiprocessing. - # We need to set PROMETHEUS_MULTIPROC_DIR environment variable - # before prometheus_client is imported. - # See https://prometheus.github.io/client_python/multiprocess/ - from prometheus_client import (CollectorRegistry, make_asgi_app, - multiprocess) - - prometheus_multiproc_dir_path = os.getenv("PROMETHEUS_MULTIPROC_DIR", None) - if prometheus_multiproc_dir_path is not None: - logger.info("vLLM to use %s as PROMETHEUS_MULTIPROC_DIR", - prometheus_multiproc_dir_path) - registry = CollectorRegistry() - multiprocess.MultiProcessCollector(registry) - - # Add prometheus asgi middleware to route /metrics requests - metrics_route = Mount("/metrics", make_asgi_app(registry=registry)) - else: - # Add prometheus asgi middleware to route /metrics requests - metrics_route = Mount("/metrics", make_asgi_app()) - - # Workaround for 307 Redirect for /metrics - metrics_route.path_regex = re.compile("^/metrics(?P.*)$") - app.routes.append(metrics_route) - - -def chat(request: Request) -> OpenAIServingChat: - return request.app.state.openai_serving_chat - - -def completion(request: Request) -> OpenAIServingCompletion: - return request.app.state.openai_serving_completion - - -def tokenization(request: Request) -> OpenAIServingTokenization: - return request.app.state.openai_serving_tokenization - - -def embedding(request: Request) -> OpenAIServingEmbedding: - return request.app.state.openai_serving_embedding - - -def engine_client(request: Request) -> EngineClient: - return request.app.state.engine_client - - -@router.get("/health") -async def health(raw_request: Request) -> Response: - """Health check.""" - await engine_client(raw_request).check_health() - return Response(status_code=200) - - -@router.post("/tokenize") -async def tokenize(request: TokenizeRequest, raw_request: Request): - generator = await tokenization(raw_request).create_tokenize(request) - if isinstance(generator, ErrorResponse): - return JSONResponse(content=generator.model_dump(), - status_code=generator.code) - elif isinstance(generator, TokenizeResponse): - return JSONResponse(content=generator.model_dump()) - - assert_never(generator) - - -@router.post("/detokenize") -async def detokenize(request: DetokenizeRequest, raw_request: Request): - generator = await tokenization(raw_request).create_detokenize(request) - if isinstance(generator, ErrorResponse): - return JSONResponse(content=generator.model_dump(), - status_code=generator.code) - elif isinstance(generator, DetokenizeResponse): - return JSONResponse(content=generator.model_dump()) - - assert_never(generator) - - -@router.get("/v1/models") -async def show_available_models(raw_request: Request): - models = await completion(raw_request).show_available_models() - return JSONResponse(content=models.model_dump()) - - -@router.get("/version") -async def show_version(): - ver = {"version": VLLM_VERSION} - return JSONResponse(content=ver) - - -@router.post("/v1/chat/completions") -async def create_chat_completion(request: ChatCompletionRequest, - raw_request: Request): - + yield engine + + +@asynccontextmanager +async def build_async_engine_client_from_engine_args( + engine_args: AsyncEngineArgs, + disable_frontend_multiprocessing: bool = False, +) -> AsyncIterator[EngineClient]: + """ + Create EngineClient, either: + - in-process using the AsyncLLMEngine Directly + - multiprocess using AsyncLLMEngine RPC + + Returns the Client or None if the creation failed. + """ + + # Fall back + # TODO: fill out feature matrix. + if (MQLLMEngineClient.is_unsupported_config(engine_args) + or disable_frontend_multiprocessing): + engine_config = engine_args.create_engine_config() + uses_ray = getattr(AsyncLLMEngine._get_executor_cls(engine_config), + "uses_ray", False) + + build_engine = partial(AsyncLLMEngine.from_engine_args, + engine_args=engine_args, + engine_config=engine_config, + usage_context=UsageContext.OPENAI_API_SERVER) + if uses_ray: + # Must run in main thread with ray for its signal handlers to work + engine_client = build_engine() + else: + engine_client = await asyncio.get_running_loop().run_in_executor( + None, build_engine) + + yield engine_client + return + + # Otherwise, use the multiprocessing AsyncLLMEngine. + else: + if "PROMETHEUS_MULTIPROC_DIR" not in os.environ: + # Make TemporaryDirectory for prometheus multiprocessing + # Note: global TemporaryDirectory will be automatically + # cleaned up upon exit. + global prometheus_multiproc_dir + prometheus_multiproc_dir = tempfile.TemporaryDirectory() + os.environ[ + "PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir.name + else: + logger.warning( + "Found PROMETHEUS_MULTIPROC_DIR was set by user. " + "This directory must be wiped between vLLM runs or " + "you will find inaccurate metrics. Unset the variable " + "and vLLM will properly handle cleanup.") + + # Select random path for IPC. + ipc_path = get_open_zmq_ipc_path() + logger.info("Multiprocessing frontend to use %s for IPC Path.", + ipc_path) + + # Start RPCServer in separate process (holds the LLMEngine). + # the current process might have CUDA context, + # so we need to spawn a new process + context = multiprocessing.get_context("spawn") + + engine_process = context.Process(target=run_mp_engine, + args=(engine_args, + UsageContext.OPENAI_API_SERVER, + ipc_path)) + engine_process.start() + logger.info("Started engine process with PID %d", engine_process.pid) + + # Build RPCClient, which conforms to EngineClient Protocol. + # NOTE: Actually, this is not true yet. We still need to support + # embedding models via RPC (see TODO above) + engine_config = engine_args.create_engine_config() + mp_engine_client = MQLLMEngineClient(ipc_path, engine_config) + + try: + while True: + try: + await mp_engine_client.setup() + break + except TimeoutError: + if not engine_process.is_alive(): + raise RuntimeError( + "Engine process failed to start") from None + + yield mp_engine_client # type: ignore[misc] + finally: + # Ensure rpc server process was terminated + engine_process.terminate() + + # Close all open connections to the backend + mp_engine_client.close() + + # Wait for engine process to join + engine_process.join(4) + if engine_process.exitcode is None: + # Kill if taking longer than 5 seconds to stop + engine_process.kill() + + # Lazy import for prometheus multiprocessing. + # We need to set PROMETHEUS_MULTIPROC_DIR environment variable + # before prometheus_client is imported. + # See https://prometheus.github.io/client_python/multiprocess/ + from prometheus_client import multiprocess + multiprocess.mark_process_dead(engine_process.pid) + + +router = APIRouter() + + +def mount_metrics(app: FastAPI): + # Lazy import for prometheus multiprocessing. + # We need to set PROMETHEUS_MULTIPROC_DIR environment variable + # before prometheus_client is imported. + # See https://prometheus.github.io/client_python/multiprocess/ + from prometheus_client import (CollectorRegistry, make_asgi_app, + multiprocess) + + prometheus_multiproc_dir_path = os.getenv("PROMETHEUS_MULTIPROC_DIR", None) + if prometheus_multiproc_dir_path is not None: + logger.info("vLLM to use %s as PROMETHEUS_MULTIPROC_DIR", + prometheus_multiproc_dir_path) + registry = CollectorRegistry() + multiprocess.MultiProcessCollector(registry) + + # Add prometheus asgi middleware to route /metrics requests + metrics_route = Mount("/metrics", make_asgi_app(registry=registry)) + else: + # Add prometheus asgi middleware to route /metrics requests + metrics_route = Mount("/metrics", make_asgi_app()) + + # Workaround for 307 Redirect for /metrics + metrics_route.path_regex = re.compile("^/metrics(?P.*)$") + app.routes.append(metrics_route) + + +def chat(request: Request) -> OpenAIServingChat: + return request.app.state.openai_serving_chat + + +def completion(request: Request) -> OpenAIServingCompletion: + return request.app.state.openai_serving_completion + + +def tokenization(request: Request) -> OpenAIServingTokenization: + return request.app.state.openai_serving_tokenization + + +def embedding(request: Request) -> OpenAIServingEmbedding: + return request.app.state.openai_serving_embedding + + +def engine_client(request: Request) -> EngineClient: + return request.app.state.engine_client + + +@router.get("/health") +async def health(raw_request: Request) -> Response: + """Health check.""" + await engine_client(raw_request).check_health() + return Response(status_code=200) + + +@router.post("/tokenize") +async def tokenize(request: TokenizeRequest, raw_request: Request): + generator = await tokenization(raw_request).create_tokenize(request) + if isinstance(generator, ErrorResponse): + return JSONResponse(content=generator.model_dump(), + status_code=generator.code) + elif isinstance(generator, TokenizeResponse): + return JSONResponse(content=generator.model_dump()) + + assert_never(generator) + + +@router.post("/detokenize") +async def detokenize(request: DetokenizeRequest, raw_request: Request): + generator = await tokenization(raw_request).create_detokenize(request) + if isinstance(generator, ErrorResponse): + return JSONResponse(content=generator.model_dump(), + status_code=generator.code) + elif isinstance(generator, DetokenizeResponse): + return JSONResponse(content=generator.model_dump()) + + assert_never(generator) + + +@router.get("/v1/models") +async def show_available_models(raw_request: Request): + models = await completion(raw_request).show_available_models() + return JSONResponse(content=models.model_dump()) + + +@router.get("/version") +async def show_version(): + ver = {"version": VLLM_VERSION} + return JSONResponse(content=ver) + + +@router.post("/v1/chat/completions") +async def create_chat_completion(request: ChatCompletionRequest, + raw_request: Request): + generator = await chat(raw_request).create_chat_completion( request, raw_request) @@ -793,288 +793,288 @@ async def create_chat_completion(request: ChatCompletionRequest, _bi100_log_chat_4xx(request, generator) return JSONResponse(content=generator.model_dump(), status_code=generator.code) - - elif isinstance(generator, ChatCompletionResponse): - return JSONResponse(content=generator.model_dump()) - - return StreamingResponse(content=generator, media_type="text/event-stream") - - -@router.post("/v1/completions") -async def create_completion(request: CompletionRequest, raw_request: Request): - generator = await completion(raw_request).create_completion( - request, raw_request) - if isinstance(generator, ErrorResponse): - return JSONResponse(content=generator.model_dump(), - status_code=generator.code) - elif isinstance(generator, CompletionResponse): - return JSONResponse(content=generator.model_dump()) - - return StreamingResponse(content=generator, media_type="text/event-stream") - - -@router.post("/v1/embeddings") -async def create_embedding(request: EmbeddingRequest, raw_request: Request): - generator = await embedding(raw_request).create_embedding( - request, raw_request) - if isinstance(generator, ErrorResponse): - return JSONResponse(content=generator.model_dump(), - status_code=generator.code) - elif isinstance(generator, EmbeddingResponse): - return JSONResponse(content=generator.model_dump()) - - assert_never(generator) - - -if envs.VLLM_TORCH_PROFILER_DIR: - logger.warning( - "Torch Profiler is enabled in the API server. This should ONLY be " - "used for local development!") - - @router.post("/start_profile") - async def start_profile(raw_request: Request): - logger.info("Starting profiler...") - await engine_client(raw_request).start_profile() - logger.info("Profiler started.") - return Response(status_code=200) - - @router.post("/stop_profile") - async def stop_profile(raw_request: Request): - logger.info("Stopping profiler...") - await engine_client(raw_request).stop_profile() - logger.info("Profiler stopped.") - return Response(status_code=200) - - -if envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING: - logger.warning( - "Lora dynamic loading & unloading is enabled in the API server. " - "This should ONLY be used for local development!") - - @router.post("/v1/load_lora_adapter") - async def load_lora_adapter(request: LoadLoraAdapterRequest, - raw_request: Request): - response = await chat(raw_request).load_lora_adapter(request) - if isinstance(response, ErrorResponse): - return JSONResponse(content=response.model_dump(), - status_code=response.code) - - response = await completion(raw_request).load_lora_adapter(request) - if isinstance(response, ErrorResponse): - return JSONResponse(content=response.model_dump(), - status_code=response.code) - - return Response(status_code=200, content=response) - - @router.post("/v1/unload_lora_adapter") - async def unload_lora_adapter(request: UnloadLoraAdapterRequest, - raw_request: Request): - response = await chat(raw_request).unload_lora_adapter(request) - if isinstance(response, ErrorResponse): - return JSONResponse(content=response.model_dump(), - status_code=response.code) - - response = await completion(raw_request).unload_lora_adapter(request) - if isinstance(response, ErrorResponse): - return JSONResponse(content=response.model_dump(), - status_code=response.code) - - return Response(status_code=200, content=response) - - -def build_app(args: Namespace) -> FastAPI: - if args.disable_fastapi_docs: - app = FastAPI(openapi_url=None, - docs_url=None, - redoc_url=None, - lifespan=lifespan) - else: - app = FastAPI(lifespan=lifespan) - app.include_router(router) - app.root_path = args.root_path - - mount_metrics(app) - - app.add_middleware( - CORSMiddleware, - allow_origins=args.allowed_origins, - allow_credentials=args.allow_credentials, - allow_methods=args.allowed_methods, - allow_headers=args.allowed_headers, - ) - + + elif isinstance(generator, ChatCompletionResponse): + return JSONResponse(content=generator.model_dump()) + + return StreamingResponse(content=generator, media_type="text/event-stream") + + +@router.post("/v1/completions") +async def create_completion(request: CompletionRequest, raw_request: Request): + generator = await completion(raw_request).create_completion( + request, raw_request) + if isinstance(generator, ErrorResponse): + return JSONResponse(content=generator.model_dump(), + status_code=generator.code) + elif isinstance(generator, CompletionResponse): + return JSONResponse(content=generator.model_dump()) + + return StreamingResponse(content=generator, media_type="text/event-stream") + + +@router.post("/v1/embeddings") +async def create_embedding(request: EmbeddingRequest, raw_request: Request): + generator = await embedding(raw_request).create_embedding( + request, raw_request) + if isinstance(generator, ErrorResponse): + return JSONResponse(content=generator.model_dump(), + status_code=generator.code) + elif isinstance(generator, EmbeddingResponse): + return JSONResponse(content=generator.model_dump()) + + assert_never(generator) + + +if envs.VLLM_TORCH_PROFILER_DIR: + logger.warning( + "Torch Profiler is enabled in the API server. This should ONLY be " + "used for local development!") + + @router.post("/start_profile") + async def start_profile(raw_request: Request): + logger.info("Starting profiler...") + await engine_client(raw_request).start_profile() + logger.info("Profiler started.") + return Response(status_code=200) + + @router.post("/stop_profile") + async def stop_profile(raw_request: Request): + logger.info("Stopping profiler...") + await engine_client(raw_request).stop_profile() + logger.info("Profiler stopped.") + return Response(status_code=200) + + +if envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING: + logger.warning( + "Lora dynamic loading & unloading is enabled in the API server. " + "This should ONLY be used for local development!") + + @router.post("/v1/load_lora_adapter") + async def load_lora_adapter(request: LoadLoraAdapterRequest, + raw_request: Request): + response = await chat(raw_request).load_lora_adapter(request) + if isinstance(response, ErrorResponse): + return JSONResponse(content=response.model_dump(), + status_code=response.code) + + response = await completion(raw_request).load_lora_adapter(request) + if isinstance(response, ErrorResponse): + return JSONResponse(content=response.model_dump(), + status_code=response.code) + + return Response(status_code=200, content=response) + + @router.post("/v1/unload_lora_adapter") + async def unload_lora_adapter(request: UnloadLoraAdapterRequest, + raw_request: Request): + response = await chat(raw_request).unload_lora_adapter(request) + if isinstance(response, ErrorResponse): + return JSONResponse(content=response.model_dump(), + status_code=response.code) + + response = await completion(raw_request).unload_lora_adapter(request) + if isinstance(response, ErrorResponse): + return JSONResponse(content=response.model_dump(), + status_code=response.code) + + return Response(status_code=200, content=response) + + +def build_app(args: Namespace) -> FastAPI: + if args.disable_fastapi_docs: + app = FastAPI(openapi_url=None, + docs_url=None, + redoc_url=None, + lifespan=lifespan) + else: + app = FastAPI(lifespan=lifespan) + app.include_router(router) + app.root_path = args.root_path + + mount_metrics(app) + + app.add_middleware( + CORSMiddleware, + allow_origins=args.allowed_origins, + allow_credentials=args.allow_credentials, + allow_methods=args.allowed_methods, + allow_headers=args.allowed_headers, + ) + @app.exception_handler(RequestValidationError) async def validation_exception_handler(raw_request, exc): _bi100_log_request_validation_4xx(raw_request, exc) chat = app.state.openai_serving_chat err = chat.create_error_response(message=str(exc)) - return JSONResponse(err.model_dump(), - status_code=HTTPStatus.BAD_REQUEST) - - if token := envs.VLLM_API_KEY or args.api_key: - - @app.middleware("http") - async def authentication(request: Request, call_next): - root_path = "" if args.root_path is None else args.root_path - if request.method == "OPTIONS": - return await call_next(request) - if not request.url.path.startswith(f"{root_path}/v1"): - return await call_next(request) - if request.headers.get("Authorization") != "Bearer " + token: - return JSONResponse(content={"error": "Unauthorized"}, - status_code=401) - return await call_next(request) - - for middleware in args.middleware: - module_path, object_name = middleware.rsplit(".", 1) - imported = getattr(importlib.import_module(module_path), object_name) - if inspect.isclass(imported): - app.add_middleware(imported) - elif inspect.iscoroutinefunction(imported): - app.middleware("http")(imported) - else: - raise ValueError(f"Invalid middleware {middleware}. " - f"Must be a function or a class.") - - return app - - -def init_app_state( - engine_client: EngineClient, - model_config: ModelConfig, - state: State, - args: Namespace, -) -> None: - if args.served_model_name is not None: - served_model_names = args.served_model_name - else: - served_model_names = [args.model] - - if args.disable_log_requests: - request_logger = None - else: - request_logger = RequestLogger(max_log_len=args.max_log_len) - - base_model_paths = [ - BaseModelPath(name=name, model_path=args.model) - for name in served_model_names - ] - - state.engine_client = engine_client - state.log_stats = not args.disable_log_stats - - state.openai_serving_chat = OpenAIServingChat( - engine_client, - model_config, - base_model_paths, - args.response_role, - lora_modules=args.lora_modules, - prompt_adapters=args.prompt_adapters, - request_logger=request_logger, - chat_template=args.chat_template, - return_tokens_as_token_ids=args.return_tokens_as_token_ids, - enable_auto_tools=args.enable_auto_tool_choice, - tool_parser=args.tool_call_parser, - reasoning_parser=getattr(args, 'reasoning_parser', None)) - state.openai_serving_completion = OpenAIServingCompletion( - engine_client, - model_config, - base_model_paths, - lora_modules=args.lora_modules, - prompt_adapters=args.prompt_adapters, - request_logger=request_logger, - return_tokens_as_token_ids=args.return_tokens_as_token_ids, - ) - state.openai_serving_embedding = OpenAIServingEmbedding( - engine_client, - model_config, - base_model_paths, - request_logger=request_logger, - ) - state.openai_serving_tokenization = OpenAIServingTokenization( - engine_client, - model_config, - base_model_paths, - lora_modules=args.lora_modules, - request_logger=request_logger, - chat_template=args.chat_template, - ) - - -async def run_server(args, **uvicorn_kwargs) -> None: + return JSONResponse(err.model_dump(), + status_code=HTTPStatus.BAD_REQUEST) + + if token := envs.VLLM_API_KEY or args.api_key: + + @app.middleware("http") + async def authentication(request: Request, call_next): + root_path = "" if args.root_path is None else args.root_path + if request.method == "OPTIONS": + return await call_next(request) + if not request.url.path.startswith(f"{root_path}/v1"): + return await call_next(request) + if request.headers.get("Authorization") != "Bearer " + token: + return JSONResponse(content={"error": "Unauthorized"}, + status_code=401) + return await call_next(request) + + for middleware in args.middleware: + module_path, object_name = middleware.rsplit(".", 1) + imported = getattr(importlib.import_module(module_path), object_name) + if inspect.isclass(imported): + app.add_middleware(imported) + elif inspect.iscoroutinefunction(imported): + app.middleware("http")(imported) + else: + raise ValueError(f"Invalid middleware {middleware}. " + f"Must be a function or a class.") + + return app + + +def init_app_state( + engine_client: EngineClient, + model_config: ModelConfig, + state: State, + args: Namespace, +) -> None: + if args.served_model_name is not None: + served_model_names = args.served_model_name + else: + served_model_names = [args.model] + + if args.disable_log_requests: + request_logger = None + else: + request_logger = RequestLogger(max_log_len=args.max_log_len) + + base_model_paths = [ + BaseModelPath(name=name, model_path=args.model) + for name in served_model_names + ] + + state.engine_client = engine_client + state.log_stats = not args.disable_log_stats + + state.openai_serving_chat = OpenAIServingChat( + engine_client, + model_config, + base_model_paths, + args.response_role, + lora_modules=args.lora_modules, + prompt_adapters=args.prompt_adapters, + request_logger=request_logger, + chat_template=args.chat_template, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + enable_auto_tools=args.enable_auto_tool_choice, + tool_parser=args.tool_call_parser, + reasoning_parser=getattr(args, 'reasoning_parser', None)) + state.openai_serving_completion = OpenAIServingCompletion( + engine_client, + model_config, + base_model_paths, + lora_modules=args.lora_modules, + prompt_adapters=args.prompt_adapters, + request_logger=request_logger, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + ) + state.openai_serving_embedding = OpenAIServingEmbedding( + engine_client, + model_config, + base_model_paths, + request_logger=request_logger, + ) + state.openai_serving_tokenization = OpenAIServingTokenization( + engine_client, + model_config, + base_model_paths, + lora_modules=args.lora_modules, + request_logger=request_logger, + chat_template=args.chat_template, + ) + + +async def run_server(args, **uvicorn_kwargs) -> None: _bi100_startup_trace("run_server entered") - logger.info("vLLM API server version %s", VLLM_VERSION) - logger.info("args: %s", args) - - if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: - ToolParserManager.import_tool_parser(args.tool_parser_plugin) - - valide_tool_parses = ToolParserManager.tool_parsers.keys() - if args.enable_auto_tool_choice \ - and args.tool_call_parser not in valide_tool_parses: - raise KeyError(f"invalid tool call parser: {args.tool_call_parser} " - f"(chose from {{ {','.join(valide_tool_parses)} }})") - - reasoning_parser = getattr(args, 'reasoning_parser', None) - if reasoning_parser: - valid_reasoning = ReasoningParserManager.list_registered() - if reasoning_parser not in valid_reasoning: - raise KeyError( - f"invalid reasoning parser: {reasoning_parser} " - f"(chose from {{ {','.join(valid_reasoning)} }})") - - # workaround to make sure that we bind the port before the engine is set up. - # This avoids race conditions with ray. - # see https://github.com/vllm-project/vllm/issues/8204 - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.bind(("", args.port)) - - def signal_handler(*_) -> None: - # Interrupt server on sigterm while initializing - raise KeyboardInterrupt("terminated") - - signal.signal(signal.SIGTERM, signal_handler) - + logger.info("vLLM API server version %s", VLLM_VERSION) + logger.info("args: %s", args) + + if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: + ToolParserManager.import_tool_parser(args.tool_parser_plugin) + + valide_tool_parses = ToolParserManager.tool_parsers.keys() + if args.enable_auto_tool_choice \ + and args.tool_call_parser not in valide_tool_parses: + raise KeyError(f"invalid tool call parser: {args.tool_call_parser} " + f"(chose from {{ {','.join(valide_tool_parses)} }})") + + reasoning_parser = getattr(args, 'reasoning_parser', None) + if reasoning_parser: + valid_reasoning = ReasoningParserManager.list_registered() + if reasoning_parser not in valid_reasoning: + raise KeyError( + f"invalid reasoning parser: {reasoning_parser} " + f"(chose from {{ {','.join(valid_reasoning)} }})") + + # workaround to make sure that we bind the port before the engine is set up. + # This avoids race conditions with ray. + # see https://github.com/vllm-project/vllm/issues/8204 + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("", args.port)) + + def signal_handler(*_) -> None: + # Interrupt server on sigterm while initializing + raise KeyboardInterrupt("terminated") + + signal.signal(signal.SIGTERM, signal_handler) + _bi100_startup_trace("starting engine client context") - async with build_async_engine_client(args) as engine_client: + async with build_async_engine_client(args) as engine_client: _bi100_startup_trace("building FastAPI application") - app = build_app(args) - + app = build_app(args) + _bi100_startup_trace("requesting model config from engine") - model_config = await engine_client.get_model_config() + model_config = await engine_client.get_model_config() _bi100_startup_trace("model config received; initializing app state") - init_app_state(engine_client, model_config, app.state, args) - + init_app_state(engine_client, model_config, app.state, args) + _bi100_startup_trace("starting HTTP server") - shutdown_task = await serve_http( - app, - host=args.host, - port=args.port, - log_level=args.uvicorn_log_level, - timeout_keep_alive=TIMEOUT_KEEP_ALIVE, - ssl_keyfile=args.ssl_keyfile, - ssl_certfile=args.ssl_certfile, - ssl_ca_certs=args.ssl_ca_certs, - ssl_cert_reqs=args.ssl_cert_reqs, - fd=sock.fileno(), - **uvicorn_kwargs, - ) - - # NB: Await server shutdown only after the backend context is exited - await shutdown_task - - -if __name__ == "__main__": + shutdown_task = await serve_http( + app, + host=args.host, + port=args.port, + log_level=args.uvicorn_log_level, + timeout_keep_alive=TIMEOUT_KEEP_ALIVE, + ssl_keyfile=args.ssl_keyfile, + ssl_certfile=args.ssl_certfile, + ssl_ca_certs=args.ssl_ca_certs, + ssl_cert_reqs=args.ssl_cert_reqs, + fd=sock.fileno(), + **uvicorn_kwargs, + ) + + # NB: Await server shutdown only after the backend context is exited + await shutdown_task + + +if __name__ == "__main__": _bi100_startup_trace("api_server __main__ entered") - # NOTE(simon): - # This section should be in sync with vllm/scripts.py for CLI entrypoints. - parser = FlexibleArgumentParser( - description="vLLM OpenAI-Compatible RESTful API server.") - parser = make_arg_parser(parser) - args = parser.parse_args() - validate_parsed_serve_args(args) + # NOTE(simon): + # This section should be in sync with vllm/scripts.py for CLI entrypoints. + parser = FlexibleArgumentParser( + description="vLLM OpenAI-Compatible RESTful API server.") + parser = make_arg_parser(parser) + args = parser.parse_args() + validate_parsed_serve_args(args) _bi100_startup_trace( f"arguments parsed model={args.model} tp={args.tensor_parallel_size} " f"max_model_len={args.max_model_len}") - - uvloop.run(run_server(args)) + + uvloop.run(run_server(args)) diff --git a/qwen3_6_scripts/chat_utils.py b/qwen3_6_scripts/chat_utils.py index 007c1e83..55e1ed3e 100644 --- a/qwen3_6_scripts/chat_utils.py +++ b/qwen3_6_scripts/chat_utils.py @@ -1,516 +1,516 @@ -import asyncio -import codecs -import json -from abc import ABC, abstractmethod -from collections import defaultdict -from functools import lru_cache, partial -from pathlib import Path -from typing import (Any, Awaitable, Dict, Generic, Iterable, List, Literal, - Mapping, Optional, Tuple, TypeVar, Union, cast) - -# yapf conflicts with isort for this block -# yapf: disable -from openai.types.chat import (ChatCompletionAssistantMessageParam, - ChatCompletionContentPartImageParam) -from openai.types.chat import ( - ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam) -from openai.types.chat import (ChatCompletionContentPartRefusalParam, - ChatCompletionContentPartTextParam) -from openai.types.chat import ( - ChatCompletionMessageParam as OpenAIChatCompletionMessageParam) -from openai.types.chat import (ChatCompletionMessageToolCallParam, - ChatCompletionToolMessageParam) -# yapf: enable -# pydantic needs the TypedDict from typing_extensions -from pydantic import ConfigDict -from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast -from typing_extensions import Required, TypeAlias, TypedDict - -from vllm.config import ModelConfig -from vllm.logger import init_logger -from vllm.multimodal import MultiModalDataDict -from vllm.multimodal.utils import (async_get_and_parse_audio, - async_get_and_parse_image, - get_and_parse_audio, get_and_parse_image) -from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer - -logger = init_logger(__name__) - - -class AudioURL(TypedDict, total=False): - url: Required[str] - """ - Either a URL of the audio or a data URL with base64 encoded audio data. - """ - - -class ChatCompletionContentPartAudioParam(TypedDict, total=False): - audio_url: Required[AudioURL] - - type: Required[Literal["audio_url"]] - """The type of the content part.""" - - -class CustomChatCompletionContentPartParam(TypedDict, total=False): - __pydantic_config__ = ConfigDict(extra="allow") # type: ignore - - type: Required[str] - """The type of the content part.""" - - -ChatCompletionContentPartParam: TypeAlias = Union[ - OpenAIChatCompletionContentPartParam, ChatCompletionContentPartAudioParam, - ChatCompletionContentPartRefusalParam, - CustomChatCompletionContentPartParam] - - -class CustomChatCompletionMessageParam(TypedDict, total=False): - """Enables custom roles in the Chat Completion API.""" - role: Required[str] - """The role of the message's author.""" - - content: Union[str, List[ChatCompletionContentPartParam]] - """The contents of the message.""" - - name: str - """An optional name for the participant. - - Provides the model information to differentiate between participants of the - same role. - """ - - tool_call_id: Optional[str] - """Tool call that this message is responding to.""" - - tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]] - """The tool calls generated by the model, such as function calls.""" - - reasoning_content: Optional[str] - """Reasoning / thinking content for assistant messages (vLLM extension). - When present in a previous assistant turn, it is rendered as - ... before the main content so the model sees its own - chain-of-thought in subsequent turns.""" - - -ChatCompletionMessageParam = Union[OpenAIChatCompletionMessageParam, - CustomChatCompletionMessageParam] - - -# TODO: Make fields ReadOnly once mypy supports it -class ConversationMessage(TypedDict, total=False): - role: Required[str] - """The role of the message's author.""" - - content: Optional[str] - """The contents of the message""" - - tool_call_id: Optional[str] - """Tool call that this message is responding to.""" - - name: Optional[str] - """The name of the function to call""" - - tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]] - """The tool calls generated by the model, such as function calls.""" - - reasoning_content: Optional[str] - """Reasoning / thinking content for assistant messages. - Passed directly to the chat template (Qwen3 reads message.reasoning_content - natively) instead of being manually wrapped in ....""" - - -ModalityStr = Literal["image", "audio", "video"] -_T = TypeVar("_T") - - -class BaseMultiModalItemTracker(ABC, Generic[_T]): - """ - Tracks multi-modal items in a given request and ensures that the number - of multi-modal items in a given request does not exceed the configured - maximum per prompt. - """ - - def __init__(self, model_config: ModelConfig, tokenizer: AnyTokenizer): - super().__init__() - - self._model_config = model_config - self._tokenizer = tokenizer - self._allowed_items = (model_config.multimodal_config.limit_per_prompt - if model_config.multimodal_config else {}) - self._consumed_items = {k: 0 for k in self._allowed_items} - - self._items: List[_T] = [] - - @staticmethod - @lru_cache(maxsize=None) - def _cached_token_str(tokenizer: AnyTokenizer, token_index: int) -> str: - return tokenizer.decode(token_index) - - def _placeholder_str(self, modality: ModalityStr, - current_count: int) -> Optional[str]: - # TODO: Let user specify how to insert image tokens into prompt - # (similar to chat template) - hf_config = self._model_config.hf_config - model_type = hf_config.model_type - - if modality == "image": - if model_type == "phi3_v": - # Workaround since this token is not defined in the tokenizer - return f"<|image_{current_count}|>" - if model_type == "minicpmv": - return "(./)" - if model_type in ("blip-2", "chatglm", "fuyu", "paligemma", - "pixtral"): - # These models do not use image tokens in the prompt - return None - if model_type == "qwen": - return f"Picture {current_count}: " - if model_type.startswith("llava"): - return self._cached_token_str(self._tokenizer, - hf_config.image_token_index) - if model_type in ("chameleon", "internvl_chat", "NVLM_D"): - return "" - if model_type == "mllama": - return "<|image|>" +import asyncio +import codecs +import json +from abc import ABC, abstractmethod +from collections import defaultdict +from functools import lru_cache, partial +from pathlib import Path +from typing import (Any, Awaitable, Dict, Generic, Iterable, List, Literal, + Mapping, Optional, Tuple, TypeVar, Union, cast) + +# yapf conflicts with isort for this block +# yapf: disable +from openai.types.chat import (ChatCompletionAssistantMessageParam, + ChatCompletionContentPartImageParam) +from openai.types.chat import ( + ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam) +from openai.types.chat import (ChatCompletionContentPartRefusalParam, + ChatCompletionContentPartTextParam) +from openai.types.chat import ( + ChatCompletionMessageParam as OpenAIChatCompletionMessageParam) +from openai.types.chat import (ChatCompletionMessageToolCallParam, + ChatCompletionToolMessageParam) +# yapf: enable +# pydantic needs the TypedDict from typing_extensions +from pydantic import ConfigDict +from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast +from typing_extensions import Required, TypeAlias, TypedDict + +from vllm.config import ModelConfig +from vllm.logger import init_logger +from vllm.multimodal import MultiModalDataDict +from vllm.multimodal.utils import (async_get_and_parse_audio, + async_get_and_parse_image, + get_and_parse_audio, get_and_parse_image) +from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer + +logger = init_logger(__name__) + + +class AudioURL(TypedDict, total=False): + url: Required[str] + """ + Either a URL of the audio or a data URL with base64 encoded audio data. + """ + + +class ChatCompletionContentPartAudioParam(TypedDict, total=False): + audio_url: Required[AudioURL] + + type: Required[Literal["audio_url"]] + """The type of the content part.""" + + +class CustomChatCompletionContentPartParam(TypedDict, total=False): + __pydantic_config__ = ConfigDict(extra="allow") # type: ignore + + type: Required[str] + """The type of the content part.""" + + +ChatCompletionContentPartParam: TypeAlias = Union[ + OpenAIChatCompletionContentPartParam, ChatCompletionContentPartAudioParam, + ChatCompletionContentPartRefusalParam, + CustomChatCompletionContentPartParam] + + +class CustomChatCompletionMessageParam(TypedDict, total=False): + """Enables custom roles in the Chat Completion API.""" + role: Required[str] + """The role of the message's author.""" + + content: Union[str, List[ChatCompletionContentPartParam]] + """The contents of the message.""" + + name: str + """An optional name for the participant. + + Provides the model information to differentiate between participants of the + same role. + """ + + tool_call_id: Optional[str] + """Tool call that this message is responding to.""" + + tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]] + """The tool calls generated by the model, such as function calls.""" + + reasoning_content: Optional[str] + """Reasoning / thinking content for assistant messages (vLLM extension). + When present in a previous assistant turn, it is rendered as + ... before the main content so the model sees its own + chain-of-thought in subsequent turns.""" + + +ChatCompletionMessageParam = Union[OpenAIChatCompletionMessageParam, + CustomChatCompletionMessageParam] + + +# TODO: Make fields ReadOnly once mypy supports it +class ConversationMessage(TypedDict, total=False): + role: Required[str] + """The role of the message's author.""" + + content: Optional[str] + """The contents of the message""" + + tool_call_id: Optional[str] + """Tool call that this message is responding to.""" + + name: Optional[str] + """The name of the function to call""" + + tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]] + """The tool calls generated by the model, such as function calls.""" + + reasoning_content: Optional[str] + """Reasoning / thinking content for assistant messages. + Passed directly to the chat template (Qwen3 reads message.reasoning_content + natively) instead of being manually wrapped in ....""" + + +ModalityStr = Literal["image", "audio", "video"] +_T = TypeVar("_T") + + +class BaseMultiModalItemTracker(ABC, Generic[_T]): + """ + Tracks multi-modal items in a given request and ensures that the number + of multi-modal items in a given request does not exceed the configured + maximum per prompt. + """ + + def __init__(self, model_config: ModelConfig, tokenizer: AnyTokenizer): + super().__init__() + + self._model_config = model_config + self._tokenizer = tokenizer + self._allowed_items = (model_config.multimodal_config.limit_per_prompt + if model_config.multimodal_config else {}) + self._consumed_items = {k: 0 for k in self._allowed_items} + + self._items: List[_T] = [] + + @staticmethod + @lru_cache(maxsize=None) + def _cached_token_str(tokenizer: AnyTokenizer, token_index: int) -> str: + return tokenizer.decode(token_index) + + def _placeholder_str(self, modality: ModalityStr, + current_count: int) -> Optional[str]: + # TODO: Let user specify how to insert image tokens into prompt + # (similar to chat template) + hf_config = self._model_config.hf_config + model_type = hf_config.model_type + + if modality == "image": + if model_type == "phi3_v": + # Workaround since this token is not defined in the tokenizer + return f"<|image_{current_count}|>" + if model_type == "minicpmv": + return "(./)" + if model_type in ("blip-2", "chatglm", "fuyu", "paligemma", + "pixtral"): + # These models do not use image tokens in the prompt + return None + if model_type == "qwen": + return f"Picture {current_count}: " + if model_type.startswith("llava"): + return self._cached_token_str(self._tokenizer, + hf_config.image_token_index) + if model_type in ("chameleon", "internvl_chat", "NVLM_D"): + return "" + if model_type == "mllama": + return "<|image|>" if model_type in ("qwen2_vl", "qwen2_5_vl", "qwen3_5", "qwen3_5_moe"): return "<|vision_start|><|image_pad|><|vision_end|>" - if model_type == "molmo": - return "" - - raise TypeError(f"Unknown model type: {model_type}") - elif modality == "audio": - if model_type == "ultravox": - return "<|reserved_special_token_0|>" - raise TypeError(f"Unknown model type: {model_type}") - elif modality == "video": - if model_type in ("qwen2_vl","qwen2_5_vl"): - return "<|vision_start|><|video_pad|><|vision_end|>" - raise TypeError(f"Unknown model type: {model_type}") - else: - raise TypeError(f"Unknown modality: {modality}") - - @staticmethod - def _combine(items: List[MultiModalDataDict]) -> MultiModalDataDict: - mm_lists: Mapping[str, List[object]] = defaultdict(list) - - # Merge all the multi-modal items - for single_mm_data in items: - for mm_key, mm_item in single_mm_data.items(): - if isinstance(mm_item, list): - mm_lists[mm_key].extend(mm_item) - else: - mm_lists[mm_key].append(mm_item) - - # Unpack any single item lists for models that don't expect multiple. - return { - mm_key: mm_list[0] if len(mm_list) == 1 else mm_list - for mm_key, mm_list in mm_lists.items() - } - - def add(self, modality: ModalityStr, item: _T) -> Optional[str]: - """ - Add a multi-modal item to the current prompt and returns the - placeholder string to use, if any. - """ - allowed_count = self._allowed_items.get(modality, 1) - current_count = self._consumed_items.get(modality, 0) + 1 - if current_count > allowed_count: - raise ValueError( - f"At most {allowed_count} {modality}(s) may be provided in " - "one request.") - - self._consumed_items[modality] = current_count - self._items.append(item) - - return self._placeholder_str(modality, current_count) - - @abstractmethod - def create_parser(self) -> "BaseMultiModalContentParser": - raise NotImplementedError - - -class MultiModalItemTracker(BaseMultiModalItemTracker[MultiModalDataDict]): - - def all_mm_data(self) -> Optional[MultiModalDataDict]: - return self._combine(self._items) if self._items else None - - def create_parser(self) -> "BaseMultiModalContentParser": - return MultiModalContentParser(self) - - -class AsyncMultiModalItemTracker( - BaseMultiModalItemTracker[Awaitable[MultiModalDataDict]]): - - async def all_mm_data(self) -> Optional[MultiModalDataDict]: - if self._items: - items = await asyncio.gather(*self._items) - return self._combine(items) - - return None - - def create_parser(self) -> "BaseMultiModalContentParser": - return AsyncMultiModalContentParser(self) - - -class BaseMultiModalContentParser(ABC): - - def __init__(self) -> None: - super().__init__() - - # multimodal placeholder_string : count - self._placeholder_counts: Dict[str, int] = defaultdict(lambda: 0) - - def _add_placeholder(self, placeholder: Optional[str]): - if placeholder: - self._placeholder_counts[placeholder] += 1 - - def mm_placeholder_counts(self) -> Dict[str, int]: - return dict(self._placeholder_counts) - - @abstractmethod - def parse_image(self, image_url: str) -> None: - raise NotImplementedError - - @abstractmethod - def parse_audio(self, audio_url: str) -> None: - raise NotImplementedError - - -class MultiModalContentParser(BaseMultiModalContentParser): - - def __init__(self, tracker: MultiModalItemTracker) -> None: - super().__init__() - - self._tracker = tracker - - def parse_image(self, image_url: str) -> None: - image = get_and_parse_image(image_url) - - placeholder = self._tracker.add("image", image) - self._add_placeholder(placeholder) - - def parse_audio(self, audio_url: str) -> None: - audio = get_and_parse_audio(audio_url) - - placeholder = self._tracker.add("audio", audio) - self._add_placeholder(placeholder) - - -class AsyncMultiModalContentParser(BaseMultiModalContentParser): - - def __init__(self, tracker: AsyncMultiModalItemTracker) -> None: - super().__init__() - - self._tracker = tracker - - def parse_image(self, image_url: str) -> None: - image_coro = async_get_and_parse_image(image_url) - - placeholder = self._tracker.add("image", image_coro) - self._add_placeholder(placeholder) - - def parse_audio(self, audio_url: str) -> None: - audio_coro = async_get_and_parse_audio(audio_url) - - placeholder = self._tracker.add("audio", audio_coro) - self._add_placeholder(placeholder) - - -def validate_chat_template(chat_template: Optional[Union[Path, str]]): - """Raises if the provided chat template appears invalid.""" - if chat_template is None: - return - - elif isinstance(chat_template, Path) and not chat_template.exists(): - raise FileNotFoundError( - "the supplied chat template path doesn't exist") - - elif isinstance(chat_template, str): - JINJA_CHARS = "{}\n" - if not any(c in chat_template - for c in JINJA_CHARS) and not Path(chat_template).exists(): - raise ValueError( - f"The supplied chat template string ({chat_template}) " - f"appears path-like, but doesn't exist!") - - else: - raise TypeError( - f"{type(chat_template)} is not a valid chat template type") - - -def load_chat_template( - chat_template: Optional[Union[Path, str]]) -> Optional[str]: - if chat_template is None: - return None - try: - with open(chat_template, "r") as f: - resolved_chat_template = f.read() - except OSError as e: - if isinstance(chat_template, Path): - raise - - JINJA_CHARS = "{}\n" - if not any(c in chat_template for c in JINJA_CHARS): - msg = (f"The supplied chat template ({chat_template}) " - f"looks like a file path, but it failed to be " - f"opened. Reason: {e}") - raise ValueError(msg) from e - - # If opening a file fails, set chat template to be args to - # ensure we decode so our escape are interpreted correctly - resolved_chat_template = codecs.decode(chat_template, "unicode_escape") - - logger.info("Using supplied chat template:\n%s", resolved_chat_template) - return resolved_chat_template - - -# TODO: Let user specify how to insert multimodal tokens into prompt -# (similar to chat template) -def _get_full_multimodal_text_prompt(placeholder_counts: Dict[str, int], - text_prompt: str) -> str: - """Combine multimodal prompts for a multimodal language model.""" - - # Look through the text prompt to check for missing placeholders - missing_placeholders: List[str] = [] - for placeholder in placeholder_counts: - - # For any existing placeholder in the text prompt, we leave it as is - placeholder_counts[placeholder] -= text_prompt.count(placeholder) - - if placeholder_counts[placeholder] < 0: - raise ValueError( - f"Found more '{placeholder}' placeholders in input prompt than " - "actual multimodal data items.") - - missing_placeholders.extend([placeholder] * - placeholder_counts[placeholder]) - - # NOTE: For now we always add missing placeholders at the front of - # the prompt. This may change to be customizable in the future. - return "\n".join(missing_placeholders + [text_prompt]) - - -# No need to validate using Pydantic again -_TextParser = partial(cast, ChatCompletionContentPartTextParam) -_ImageParser = partial(cast, ChatCompletionContentPartImageParam) -_AudioParser = partial(cast, ChatCompletionContentPartAudioParam) -_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam) -MODEL_KEEP_MULTI_MODAL_CONTENT = {'mllama'} - - -def _parse_chat_message_content_parts( - role: str, - parts: Iterable[ChatCompletionContentPartParam], - mm_tracker: BaseMultiModalItemTracker, -) -> List[ConversationMessage]: - texts: List[str] = [] - - mm_parser = mm_tracker.create_parser() - keep_multimodal_content = \ - mm_tracker._model_config.hf_config.model_type in \ - MODEL_KEEP_MULTI_MODAL_CONTENT - - has_image = False - for part in parts: - part_type = part["type"] - if part_type == "text": - text = _TextParser(part)["text"] - texts.append(text) - elif part_type == "image_url": - image_url = _ImageParser(part)["image_url"] - - if image_url.get("detail", "auto") != "auto": - logger.warning( - "'image_url.detail' is currently not supported and " - "will be ignored.") - - mm_parser.parse_image(image_url["url"]) - has_image = True - elif part_type == "audio_url": - audio_url = _AudioParser(part)["audio_url"] - - mm_parser.parse_audio(audio_url["url"]) - elif part_type == "refusal": - text = _RefusalParser(part)["refusal"] - texts.append(text) - else: - raise NotImplementedError(f"Unknown part type: {part_type}") - - text_prompt = "\n".join(texts) - if keep_multimodal_content: - text_prompt = "\n".join(texts) - role_content = [{'type': 'text', 'text': text_prompt}] - - if has_image: - role_content = [{'type': 'image'}] + role_content - return [ConversationMessage(role=role, - content=role_content)] # type: ignore - else: - mm_placeholder_counts = mm_parser.mm_placeholder_counts() - if mm_placeholder_counts: - text_prompt = _get_full_multimodal_text_prompt( - mm_placeholder_counts, text_prompt) - return [ConversationMessage(role=role, content=text_prompt)] - - -# No need to validate using Pydantic again -_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam) -_ToolParser = partial(cast, ChatCompletionToolMessageParam) - - -def _parse_chat_message_content( - message: ChatCompletionMessageParam, - mm_tracker: BaseMultiModalItemTracker, -) -> List[ConversationMessage]: - role = message["role"] - content = message.get("content") - - if content is None: - content = [] - elif isinstance(content, str): - content = [ - ChatCompletionContentPartTextParam(type="text", text=content) - ] - - result = _parse_chat_message_content_parts( - role, - content, # type: ignore - mm_tracker, - ) - - for result_msg in result: - if role == 'assistant': - parsed_msg = _AssistantParser(message) - - if "tool_calls" in parsed_msg: - result_msg["tool_calls"] = list(parsed_msg["tool_calls"]) - - # Pass reasoning content as a dedicated field so the chat template - # can render it natively (Qwen3: message.reasoning_content branch). - # Accept both "reasoning" (new vllm) and "reasoning_content" (ours). - reasoning = (message.get("reasoning") # type: ignore[arg-type] - or message.get("reasoning_content")) # type: ignore[arg-type] - if reasoning and isinstance(reasoning, str): - result_msg["reasoning_content"] = reasoning - - elif role == "tool": - parsed_msg = _ToolParser(message) - if "tool_call_id" in parsed_msg: - result_msg["tool_call_id"] = parsed_msg["tool_call_id"] - - if "name" in message and isinstance(message["name"], str): - result_msg["name"] = message["name"] - - return result - - -def _postprocess_messages(messages: List[ConversationMessage]) -> None: - # per the Transformers docs & maintainers, tool call arguments in - # assistant-role messages with tool_calls need to be dicts not JSON str - - # this is how tool-use chat templates will expect them moving forwards - # so, for messages that have tool_calls, parse the string (which we get - # from openAI format) to dict + if model_type == "molmo": + return "" + + raise TypeError(f"Unknown model type: {model_type}") + elif modality == "audio": + if model_type == "ultravox": + return "<|reserved_special_token_0|>" + raise TypeError(f"Unknown model type: {model_type}") + elif modality == "video": + if model_type in ("qwen2_vl","qwen2_5_vl"): + return "<|vision_start|><|video_pad|><|vision_end|>" + raise TypeError(f"Unknown model type: {model_type}") + else: + raise TypeError(f"Unknown modality: {modality}") + + @staticmethod + def _combine(items: List[MultiModalDataDict]) -> MultiModalDataDict: + mm_lists: Mapping[str, List[object]] = defaultdict(list) + + # Merge all the multi-modal items + for single_mm_data in items: + for mm_key, mm_item in single_mm_data.items(): + if isinstance(mm_item, list): + mm_lists[mm_key].extend(mm_item) + else: + mm_lists[mm_key].append(mm_item) + + # Unpack any single item lists for models that don't expect multiple. + return { + mm_key: mm_list[0] if len(mm_list) == 1 else mm_list + for mm_key, mm_list in mm_lists.items() + } + + def add(self, modality: ModalityStr, item: _T) -> Optional[str]: + """ + Add a multi-modal item to the current prompt and returns the + placeholder string to use, if any. + """ + allowed_count = self._allowed_items.get(modality, 1) + current_count = self._consumed_items.get(modality, 0) + 1 + if current_count > allowed_count: + raise ValueError( + f"At most {allowed_count} {modality}(s) may be provided in " + "one request.") + + self._consumed_items[modality] = current_count + self._items.append(item) + + return self._placeholder_str(modality, current_count) + + @abstractmethod + def create_parser(self) -> "BaseMultiModalContentParser": + raise NotImplementedError + + +class MultiModalItemTracker(BaseMultiModalItemTracker[MultiModalDataDict]): + + def all_mm_data(self) -> Optional[MultiModalDataDict]: + return self._combine(self._items) if self._items else None + + def create_parser(self) -> "BaseMultiModalContentParser": + return MultiModalContentParser(self) + + +class AsyncMultiModalItemTracker( + BaseMultiModalItemTracker[Awaitable[MultiModalDataDict]]): + + async def all_mm_data(self) -> Optional[MultiModalDataDict]: + if self._items: + items = await asyncio.gather(*self._items) + return self._combine(items) + + return None + + def create_parser(self) -> "BaseMultiModalContentParser": + return AsyncMultiModalContentParser(self) + + +class BaseMultiModalContentParser(ABC): + + def __init__(self) -> None: + super().__init__() + + # multimodal placeholder_string : count + self._placeholder_counts: Dict[str, int] = defaultdict(lambda: 0) + + def _add_placeholder(self, placeholder: Optional[str]): + if placeholder: + self._placeholder_counts[placeholder] += 1 + + def mm_placeholder_counts(self) -> Dict[str, int]: + return dict(self._placeholder_counts) + + @abstractmethod + def parse_image(self, image_url: str) -> None: + raise NotImplementedError + + @abstractmethod + def parse_audio(self, audio_url: str) -> None: + raise NotImplementedError + + +class MultiModalContentParser(BaseMultiModalContentParser): + + def __init__(self, tracker: MultiModalItemTracker) -> None: + super().__init__() + + self._tracker = tracker + + def parse_image(self, image_url: str) -> None: + image = get_and_parse_image(image_url) + + placeholder = self._tracker.add("image", image) + self._add_placeholder(placeholder) + + def parse_audio(self, audio_url: str) -> None: + audio = get_and_parse_audio(audio_url) + + placeholder = self._tracker.add("audio", audio) + self._add_placeholder(placeholder) + + +class AsyncMultiModalContentParser(BaseMultiModalContentParser): + + def __init__(self, tracker: AsyncMultiModalItemTracker) -> None: + super().__init__() + + self._tracker = tracker + + def parse_image(self, image_url: str) -> None: + image_coro = async_get_and_parse_image(image_url) + + placeholder = self._tracker.add("image", image_coro) + self._add_placeholder(placeholder) + + def parse_audio(self, audio_url: str) -> None: + audio_coro = async_get_and_parse_audio(audio_url) + + placeholder = self._tracker.add("audio", audio_coro) + self._add_placeholder(placeholder) + + +def validate_chat_template(chat_template: Optional[Union[Path, str]]): + """Raises if the provided chat template appears invalid.""" + if chat_template is None: + return + + elif isinstance(chat_template, Path) and not chat_template.exists(): + raise FileNotFoundError( + "the supplied chat template path doesn't exist") + + elif isinstance(chat_template, str): + JINJA_CHARS = "{}\n" + if not any(c in chat_template + for c in JINJA_CHARS) and not Path(chat_template).exists(): + raise ValueError( + f"The supplied chat template string ({chat_template}) " + f"appears path-like, but doesn't exist!") + + else: + raise TypeError( + f"{type(chat_template)} is not a valid chat template type") + + +def load_chat_template( + chat_template: Optional[Union[Path, str]]) -> Optional[str]: + if chat_template is None: + return None + try: + with open(chat_template, "r") as f: + resolved_chat_template = f.read() + except OSError as e: + if isinstance(chat_template, Path): + raise + + JINJA_CHARS = "{}\n" + if not any(c in chat_template for c in JINJA_CHARS): + msg = (f"The supplied chat template ({chat_template}) " + f"looks like a file path, but it failed to be " + f"opened. Reason: {e}") + raise ValueError(msg) from e + + # If opening a file fails, set chat template to be args to + # ensure we decode so our escape are interpreted correctly + resolved_chat_template = codecs.decode(chat_template, "unicode_escape") + + logger.info("Using supplied chat template:\n%s", resolved_chat_template) + return resolved_chat_template + + +# TODO: Let user specify how to insert multimodal tokens into prompt +# (similar to chat template) +def _get_full_multimodal_text_prompt(placeholder_counts: Dict[str, int], + text_prompt: str) -> str: + """Combine multimodal prompts for a multimodal language model.""" + + # Look through the text prompt to check for missing placeholders + missing_placeholders: List[str] = [] + for placeholder in placeholder_counts: + + # For any existing placeholder in the text prompt, we leave it as is + placeholder_counts[placeholder] -= text_prompt.count(placeholder) + + if placeholder_counts[placeholder] < 0: + raise ValueError( + f"Found more '{placeholder}' placeholders in input prompt than " + "actual multimodal data items.") + + missing_placeholders.extend([placeholder] * + placeholder_counts[placeholder]) + + # NOTE: For now we always add missing placeholders at the front of + # the prompt. This may change to be customizable in the future. + return "\n".join(missing_placeholders + [text_prompt]) + + +# No need to validate using Pydantic again +_TextParser = partial(cast, ChatCompletionContentPartTextParam) +_ImageParser = partial(cast, ChatCompletionContentPartImageParam) +_AudioParser = partial(cast, ChatCompletionContentPartAudioParam) +_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam) +MODEL_KEEP_MULTI_MODAL_CONTENT = {'mllama'} + + +def _parse_chat_message_content_parts( + role: str, + parts: Iterable[ChatCompletionContentPartParam], + mm_tracker: BaseMultiModalItemTracker, +) -> List[ConversationMessage]: + texts: List[str] = [] + + mm_parser = mm_tracker.create_parser() + keep_multimodal_content = \ + mm_tracker._model_config.hf_config.model_type in \ + MODEL_KEEP_MULTI_MODAL_CONTENT + + has_image = False + for part in parts: + part_type = part["type"] + if part_type == "text": + text = _TextParser(part)["text"] + texts.append(text) + elif part_type == "image_url": + image_url = _ImageParser(part)["image_url"] + + if image_url.get("detail", "auto") != "auto": + logger.warning( + "'image_url.detail' is currently not supported and " + "will be ignored.") + + mm_parser.parse_image(image_url["url"]) + has_image = True + elif part_type == "audio_url": + audio_url = _AudioParser(part)["audio_url"] + + mm_parser.parse_audio(audio_url["url"]) + elif part_type == "refusal": + text = _RefusalParser(part)["refusal"] + texts.append(text) + else: + raise NotImplementedError(f"Unknown part type: {part_type}") + + text_prompt = "\n".join(texts) + if keep_multimodal_content: + text_prompt = "\n".join(texts) + role_content = [{'type': 'text', 'text': text_prompt}] + + if has_image: + role_content = [{'type': 'image'}] + role_content + return [ConversationMessage(role=role, + content=role_content)] # type: ignore + else: + mm_placeholder_counts = mm_parser.mm_placeholder_counts() + if mm_placeholder_counts: + text_prompt = _get_full_multimodal_text_prompt( + mm_placeholder_counts, text_prompt) + return [ConversationMessage(role=role, content=text_prompt)] + + +# No need to validate using Pydantic again +_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam) +_ToolParser = partial(cast, ChatCompletionToolMessageParam) + + +def _parse_chat_message_content( + message: ChatCompletionMessageParam, + mm_tracker: BaseMultiModalItemTracker, +) -> List[ConversationMessage]: + role = message["role"] + content = message.get("content") + + if content is None: + content = [] + elif isinstance(content, str): + content = [ + ChatCompletionContentPartTextParam(type="text", text=content) + ] + + result = _parse_chat_message_content_parts( + role, + content, # type: ignore + mm_tracker, + ) + + for result_msg in result: + if role == 'assistant': + parsed_msg = _AssistantParser(message) + + if "tool_calls" in parsed_msg: + result_msg["tool_calls"] = list(parsed_msg["tool_calls"]) + + # Pass reasoning content as a dedicated field so the chat template + # can render it natively (Qwen3: message.reasoning_content branch). + # Accept both "reasoning" (new vllm) and "reasoning_content" (ours). + reasoning = (message.get("reasoning") # type: ignore[arg-type] + or message.get("reasoning_content")) # type: ignore[arg-type] + if reasoning and isinstance(reasoning, str): + result_msg["reasoning_content"] = reasoning + + elif role == "tool": + parsed_msg = _ToolParser(message) + if "tool_call_id" in parsed_msg: + result_msg["tool_call_id"] = parsed_msg["tool_call_id"] + + if "name" in message and isinstance(message["name"], str): + result_msg["name"] = message["name"] + + return result + + +def _postprocess_messages(messages: List[ConversationMessage]) -> None: + # per the Transformers docs & maintainers, tool call arguments in + # assistant-role messages with tool_calls need to be dicts not JSON str - + # this is how tool-use chat templates will expect them moving forwards + # so, for messages that have tool_calls, parse the string (which we get + # from openAI format) to dict for message in messages: if (message["role"] == "assistant" and "tool_calls" in message and message["tool_calls"] is not None): @@ -533,85 +533,85 @@ def _postprocess_messages(messages: List[ConversationMessage]) -> None: raise TypeError( "Tool call arguments must decode to a JSON object.") item["function"]["arguments"] = arguments - - -def parse_chat_messages( - messages: List[ChatCompletionMessageParam], - model_config: ModelConfig, - tokenizer: AnyTokenizer, -) -> Tuple[List[ConversationMessage], Optional[MultiModalDataDict]]: - conversation: List[ConversationMessage] = [] - mm_tracker = MultiModalItemTracker(model_config, tokenizer) - - for msg in messages: - sub_messages = _parse_chat_message_content(msg, mm_tracker) - - conversation.extend(sub_messages) - - _postprocess_messages(conversation) - - return conversation, mm_tracker.all_mm_data() - - -def parse_chat_messages_futures( - messages: List[ChatCompletionMessageParam], - model_config: ModelConfig, - tokenizer: AnyTokenizer, -) -> Tuple[List[ConversationMessage], Awaitable[Optional[MultiModalDataDict]]]: - conversation: List[ConversationMessage] = [] - mm_tracker = AsyncMultiModalItemTracker(model_config, tokenizer) - - for msg in messages: - sub_messages = _parse_chat_message_content(msg, mm_tracker) - - conversation.extend(sub_messages) - - _postprocess_messages(conversation) - - return conversation, mm_tracker.all_mm_data() - - -def apply_hf_chat_template( - tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], - conversation: List[ConversationMessage], - chat_template: Optional[str], - *, - tokenize: bool = False, # Different from HF's default - **kwargs: Any, -) -> str: - if chat_template is None and tokenizer.chat_template is None: - raise ValueError( - "As of transformers v4.44, default chat template is no longer " - "allowed, so you must provide a chat template if the tokenizer " - "does not define one.") - - return tokenizer.apply_chat_template( - conversation=conversation, # type: ignore[arg-type] - chat_template=chat_template, - tokenize=tokenize, - **kwargs, - ) - - -def apply_mistral_chat_template( - tokenizer: MistralTokenizer, - messages: List[ChatCompletionMessageParam], - chat_template: Optional[str] = None, - **kwargs: Any, -) -> List[int]: - if chat_template is not None: - logger.warning( - "'chat_template' cannot be overridden for mistral tokenizer.") - if "add_generation_prompt" in kwargs: - logger.warning( - "'add_generation_prompt' is not supported for mistral tokenizer, " - "so it will be ignored.") - if "continue_final_message" in kwargs: - logger.warning( - "'continue_final_message' is not supported for mistral tokenizer, " - "so it will be ignored.") - - return tokenizer.apply_chat_template( - messages=messages, - **kwargs, - ) + + +def parse_chat_messages( + messages: List[ChatCompletionMessageParam], + model_config: ModelConfig, + tokenizer: AnyTokenizer, +) -> Tuple[List[ConversationMessage], Optional[MultiModalDataDict]]: + conversation: List[ConversationMessage] = [] + mm_tracker = MultiModalItemTracker(model_config, tokenizer) + + for msg in messages: + sub_messages = _parse_chat_message_content(msg, mm_tracker) + + conversation.extend(sub_messages) + + _postprocess_messages(conversation) + + return conversation, mm_tracker.all_mm_data() + + +def parse_chat_messages_futures( + messages: List[ChatCompletionMessageParam], + model_config: ModelConfig, + tokenizer: AnyTokenizer, +) -> Tuple[List[ConversationMessage], Awaitable[Optional[MultiModalDataDict]]]: + conversation: List[ConversationMessage] = [] + mm_tracker = AsyncMultiModalItemTracker(model_config, tokenizer) + + for msg in messages: + sub_messages = _parse_chat_message_content(msg, mm_tracker) + + conversation.extend(sub_messages) + + _postprocess_messages(conversation) + + return conversation, mm_tracker.all_mm_data() + + +def apply_hf_chat_template( + tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], + conversation: List[ConversationMessage], + chat_template: Optional[str], + *, + tokenize: bool = False, # Different from HF's default + **kwargs: Any, +) -> str: + if chat_template is None and tokenizer.chat_template is None: + raise ValueError( + "As of transformers v4.44, default chat template is no longer " + "allowed, so you must provide a chat template if the tokenizer " + "does not define one.") + + return tokenizer.apply_chat_template( + conversation=conversation, # type: ignore[arg-type] + chat_template=chat_template, + tokenize=tokenize, + **kwargs, + ) + + +def apply_mistral_chat_template( + tokenizer: MistralTokenizer, + messages: List[ChatCompletionMessageParam], + chat_template: Optional[str] = None, + **kwargs: Any, +) -> List[int]: + if chat_template is not None: + logger.warning( + "'chat_template' cannot be overridden for mistral tokenizer.") + if "add_generation_prompt" in kwargs: + logger.warning( + "'add_generation_prompt' is not supported for mistral tokenizer, " + "so it will be ignored.") + if "continue_final_message" in kwargs: + logger.warning( + "'continue_final_message' is not supported for mistral tokenizer, " + "so it will be ignored.") + + return tokenizer.apply_chat_template( + messages=messages, + **kwargs, + ) diff --git a/qwen3_6_scripts/cli_args.py b/qwen3_6_scripts/cli_args.py index 292b6da5..ad0698d2 100644 --- a/qwen3_6_scripts/cli_args.py +++ b/qwen3_6_scripts/cli_args.py @@ -1,261 +1,261 @@ -""" -This file contains the command line arguments for the vLLM's -OpenAI-compatible server. It is kept in a separate file for documentation -purposes. -""" - -import argparse -import json -import ssl -from typing import List, Optional, Sequence, Union - -from vllm.engine.arg_utils import AsyncEngineArgs, nullable_str -from vllm.entrypoints.chat_utils import validate_chat_template -from vllm.entrypoints.openai.serving_engine import (LoRAModulePath, - PromptAdapterPath) -from vllm.entrypoints.openai.tool_parsers import ToolParserManager -from vllm.utils import FlexibleArgumentParser - - -class LoRAParserAction(argparse.Action): - - def __call__( - self, - parser: argparse.ArgumentParser, - namespace: argparse.Namespace, - values: Optional[Union[str, Sequence[str]]], - option_string: Optional[str] = None, - ): - if values is None: - values = [] - if isinstance(values, str): - raise TypeError("Expected values to be a list") - - lora_list: List[LoRAModulePath] = [] - for item in values: - if item in [None, '']: # Skip if item is None or empty string - continue - if '=' in item and ',' not in item: # Old format: name=path - name, path = item.split('=') - lora_list.append(LoRAModulePath(name, path)) - else: # Assume JSON format - try: - lora_dict = json.loads(item) - lora = LoRAModulePath(**lora_dict) - lora_list.append(lora) - except json.JSONDecodeError: - parser.error( - f"Invalid JSON format for --lora-modules: {item}") - except TypeError as e: - parser.error( - f"Invalid fields for --lora-modules: {item} - {str(e)}" - ) - setattr(namespace, self.dest, lora_list) - - -class PromptAdapterParserAction(argparse.Action): - - def __call__( - self, - parser: argparse.ArgumentParser, - namespace: argparse.Namespace, - values: Optional[Union[str, Sequence[str]]], - option_string: Optional[str] = None, - ): - if values is None: - values = [] - if isinstance(values, str): - raise TypeError("Expected values to be a list") - - adapter_list: List[PromptAdapterPath] = [] - for item in values: - name, path = item.split('=') - adapter_list.append(PromptAdapterPath(name, path)) - setattr(namespace, self.dest, adapter_list) - - -def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: - parser.add_argument("--host", - type=nullable_str, - default=None, - help="host name") - parser.add_argument("--port", type=int, default=8000, help="port number") - parser.add_argument( - "--uvicorn-log-level", - type=str, - default="info", - choices=['debug', 'info', 'warning', 'error', 'critical', 'trace'], - help="log level for uvicorn") - parser.add_argument("--allow-credentials", - action="store_true", - help="allow credentials") - parser.add_argument("--allowed-origins", - type=json.loads, - default=["*"], - help="allowed origins") - parser.add_argument("--allowed-methods", - type=json.loads, - default=["*"], - help="allowed methods") - parser.add_argument("--allowed-headers", - type=json.loads, - default=["*"], - help="allowed headers") - parser.add_argument("--api-key", - type=nullable_str, - default=None, - help="If provided, the server will require this key " - "to be presented in the header.") - parser.add_argument( - "--lora-modules", - type=nullable_str, - default=None, - nargs='+', - action=LoRAParserAction, - help="LoRA module configurations in either 'name=path' format" - "or JSON format. " - "Example (old format): 'name=path' " - "Example (new format): " - "'{\"name\": \"name\", \"local_path\": \"path\", " - "\"base_model_name\": \"id\"}'") - parser.add_argument( - "--prompt-adapters", - type=nullable_str, - default=None, - nargs='+', - action=PromptAdapterParserAction, - help="Prompt adapter configurations in the format name=path. " - "Multiple adapters can be specified.") - parser.add_argument("--chat-template", - type=nullable_str, - default=None, - help="The file path to the chat template, " - "or the template in single-line form " - "for the specified model") - parser.add_argument("--response-role", - type=nullable_str, - default="assistant", - help="The role name to return if " - "`request.add_generation_prompt=true`.") - parser.add_argument("--ssl-keyfile", - type=nullable_str, - default=None, - help="The file path to the SSL key file") - parser.add_argument("--ssl-certfile", - type=nullable_str, - default=None, - help="The file path to the SSL cert file") - parser.add_argument("--ssl-ca-certs", - type=nullable_str, - default=None, - help="The CA certificates file") - parser.add_argument( - "--ssl-cert-reqs", - type=int, - default=int(ssl.CERT_NONE), - help="Whether client certificate is required (see stdlib ssl module's)" - ) - parser.add_argument( - "--root-path", - type=nullable_str, - default=None, - help="FastAPI root_path when app is behind a path based routing proxy") - parser.add_argument( - "--middleware", - type=nullable_str, - action="append", - default=[], - help="Additional ASGI middleware to apply to the app. " - "We accept multiple --middleware arguments. " - "The value should be an import path. " - "If a function is provided, vLLM will add it to the server " - "using @app.middleware('http'). " - "If a class is provided, vLLM will add it to the server " - "using app.add_middleware(). ") - parser.add_argument( - "--return-tokens-as-token-ids", - action="store_true", - help="When --max-logprobs is specified, represents single tokens as " - "strings of the form 'token_id:{token_id}' so that tokens that " - "are not JSON-encodable can be identified.") - parser.add_argument( - "--disable-frontend-multiprocessing", - action="store_true", - help="If specified, will run the OpenAI frontend server in the same " - "process as the model serving engine.") - - parser.add_argument( - "--enable-auto-tool-choice", - action="store_true", - default=False, - help= - "Enable auto tool choice for supported models. Use --tool-call-parser" - "to specify which parser to use") - - valid_tool_parsers = ToolParserManager.tool_parsers.keys() - parser.add_argument( - "--tool-call-parser", - type=str, - metavar="{" + ",".join(valid_tool_parsers) + "} or name registered in " - "--tool-parser-plugin", - default=None, - help= - "Select the tool call parser depending on the model that you're using." - " This is used to parse the model-generated tool call into OpenAI API " - "format. Required for --enable-auto-tool-choice.") - - parser.add_argument( - "--tool-parser-plugin", - type=str, - default="", - help= - "Special the tool parser plugin write to parse the model-generated tool" - " into OpenAI API format, the name register in this plugin can be used " - "in --tool-call-parser.") - - parser.add_argument( - "--reasoning-parser", - type=str, - default=None, - help= - "Select the reasoning parser to split ... content into " - "reasoning_content vs content in the response. " - "Supported: qwen3") - - parser = AsyncEngineArgs.add_cli_args(parser) - - parser.add_argument('--max-log-len', - type=int, - default=None, - help='Max number of prompt characters or prompt ' - 'ID numbers being printed in log.' - '\n\nDefault: Unlimited') - - parser.add_argument( - "--disable-fastapi-docs", - action='store_true', - default=False, - help="Disable FastAPI's OpenAPI schema, Swagger UI, and ReDoc endpoint" - ) - - return parser - - -def validate_parsed_serve_args(args: argparse.Namespace): - """Quick checks for model serve args that raise prior to loading.""" - if hasattr(args, "subparser") and args.subparser != "serve": - return - - # Ensure that the chat template is valid; raises if it likely isn't - validate_chat_template(args.chat_template) - - # Enable auto tool needs a tool call parser to be valid - if args.enable_auto_tool_choice and not args.tool_call_parser: - raise TypeError("Error: --enable-auto-tool-choice requires " - "--tool-call-parser") - - -def create_parser_for_docs() -> FlexibleArgumentParser: - parser_for_docs = FlexibleArgumentParser( - prog="-m vllm.entrypoints.openai.api_server") - return make_arg_parser(parser_for_docs) +""" +This file contains the command line arguments for the vLLM's +OpenAI-compatible server. It is kept in a separate file for documentation +purposes. +""" + +import argparse +import json +import ssl +from typing import List, Optional, Sequence, Union + +from vllm.engine.arg_utils import AsyncEngineArgs, nullable_str +from vllm.entrypoints.chat_utils import validate_chat_template +from vllm.entrypoints.openai.serving_engine import (LoRAModulePath, + PromptAdapterPath) +from vllm.entrypoints.openai.tool_parsers import ToolParserManager +from vllm.utils import FlexibleArgumentParser + + +class LoRAParserAction(argparse.Action): + + def __call__( + self, + parser: argparse.ArgumentParser, + namespace: argparse.Namespace, + values: Optional[Union[str, Sequence[str]]], + option_string: Optional[str] = None, + ): + if values is None: + values = [] + if isinstance(values, str): + raise TypeError("Expected values to be a list") + + lora_list: List[LoRAModulePath] = [] + for item in values: + if item in [None, '']: # Skip if item is None or empty string + continue + if '=' in item and ',' not in item: # Old format: name=path + name, path = item.split('=') + lora_list.append(LoRAModulePath(name, path)) + else: # Assume JSON format + try: + lora_dict = json.loads(item) + lora = LoRAModulePath(**lora_dict) + lora_list.append(lora) + except json.JSONDecodeError: + parser.error( + f"Invalid JSON format for --lora-modules: {item}") + except TypeError as e: + parser.error( + f"Invalid fields for --lora-modules: {item} - {str(e)}" + ) + setattr(namespace, self.dest, lora_list) + + +class PromptAdapterParserAction(argparse.Action): + + def __call__( + self, + parser: argparse.ArgumentParser, + namespace: argparse.Namespace, + values: Optional[Union[str, Sequence[str]]], + option_string: Optional[str] = None, + ): + if values is None: + values = [] + if isinstance(values, str): + raise TypeError("Expected values to be a list") + + adapter_list: List[PromptAdapterPath] = [] + for item in values: + name, path = item.split('=') + adapter_list.append(PromptAdapterPath(name, path)) + setattr(namespace, self.dest, adapter_list) + + +def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: + parser.add_argument("--host", + type=nullable_str, + default=None, + help="host name") + parser.add_argument("--port", type=int, default=8000, help="port number") + parser.add_argument( + "--uvicorn-log-level", + type=str, + default="info", + choices=['debug', 'info', 'warning', 'error', 'critical', 'trace'], + help="log level for uvicorn") + parser.add_argument("--allow-credentials", + action="store_true", + help="allow credentials") + parser.add_argument("--allowed-origins", + type=json.loads, + default=["*"], + help="allowed origins") + parser.add_argument("--allowed-methods", + type=json.loads, + default=["*"], + help="allowed methods") + parser.add_argument("--allowed-headers", + type=json.loads, + default=["*"], + help="allowed headers") + parser.add_argument("--api-key", + type=nullable_str, + default=None, + help="If provided, the server will require this key " + "to be presented in the header.") + parser.add_argument( + "--lora-modules", + type=nullable_str, + default=None, + nargs='+', + action=LoRAParserAction, + help="LoRA module configurations in either 'name=path' format" + "or JSON format. " + "Example (old format): 'name=path' " + "Example (new format): " + "'{\"name\": \"name\", \"local_path\": \"path\", " + "\"base_model_name\": \"id\"}'") + parser.add_argument( + "--prompt-adapters", + type=nullable_str, + default=None, + nargs='+', + action=PromptAdapterParserAction, + help="Prompt adapter configurations in the format name=path. " + "Multiple adapters can be specified.") + parser.add_argument("--chat-template", + type=nullable_str, + default=None, + help="The file path to the chat template, " + "or the template in single-line form " + "for the specified model") + parser.add_argument("--response-role", + type=nullable_str, + default="assistant", + help="The role name to return if " + "`request.add_generation_prompt=true`.") + parser.add_argument("--ssl-keyfile", + type=nullable_str, + default=None, + help="The file path to the SSL key file") + parser.add_argument("--ssl-certfile", + type=nullable_str, + default=None, + help="The file path to the SSL cert file") + parser.add_argument("--ssl-ca-certs", + type=nullable_str, + default=None, + help="The CA certificates file") + parser.add_argument( + "--ssl-cert-reqs", + type=int, + default=int(ssl.CERT_NONE), + help="Whether client certificate is required (see stdlib ssl module's)" + ) + parser.add_argument( + "--root-path", + type=nullable_str, + default=None, + help="FastAPI root_path when app is behind a path based routing proxy") + parser.add_argument( + "--middleware", + type=nullable_str, + action="append", + default=[], + help="Additional ASGI middleware to apply to the app. " + "We accept multiple --middleware arguments. " + "The value should be an import path. " + "If a function is provided, vLLM will add it to the server " + "using @app.middleware('http'). " + "If a class is provided, vLLM will add it to the server " + "using app.add_middleware(). ") + parser.add_argument( + "--return-tokens-as-token-ids", + action="store_true", + help="When --max-logprobs is specified, represents single tokens as " + "strings of the form 'token_id:{token_id}' so that tokens that " + "are not JSON-encodable can be identified.") + parser.add_argument( + "--disable-frontend-multiprocessing", + action="store_true", + help="If specified, will run the OpenAI frontend server in the same " + "process as the model serving engine.") + + parser.add_argument( + "--enable-auto-tool-choice", + action="store_true", + default=False, + help= + "Enable auto tool choice for supported models. Use --tool-call-parser" + "to specify which parser to use") + + valid_tool_parsers = ToolParserManager.tool_parsers.keys() + parser.add_argument( + "--tool-call-parser", + type=str, + metavar="{" + ",".join(valid_tool_parsers) + "} or name registered in " + "--tool-parser-plugin", + default=None, + help= + "Select the tool call parser depending on the model that you're using." + " This is used to parse the model-generated tool call into OpenAI API " + "format. Required for --enable-auto-tool-choice.") + + parser.add_argument( + "--tool-parser-plugin", + type=str, + default="", + help= + "Special the tool parser plugin write to parse the model-generated tool" + " into OpenAI API format, the name register in this plugin can be used " + "in --tool-call-parser.") + + parser.add_argument( + "--reasoning-parser", + type=str, + default=None, + help= + "Select the reasoning parser to split ... content into " + "reasoning_content vs content in the response. " + "Supported: qwen3") + + parser = AsyncEngineArgs.add_cli_args(parser) + + parser.add_argument('--max-log-len', + type=int, + default=None, + help='Max number of prompt characters or prompt ' + 'ID numbers being printed in log.' + '\n\nDefault: Unlimited') + + parser.add_argument( + "--disable-fastapi-docs", + action='store_true', + default=False, + help="Disable FastAPI's OpenAPI schema, Swagger UI, and ReDoc endpoint" + ) + + return parser + + +def validate_parsed_serve_args(args: argparse.Namespace): + """Quick checks for model serve args that raise prior to loading.""" + if hasattr(args, "subparser") and args.subparser != "serve": + return + + # Ensure that the chat template is valid; raises if it likely isn't + validate_chat_template(args.chat_template) + + # Enable auto tool needs a tool call parser to be valid + if args.enable_auto_tool_choice and not args.tool_call_parser: + raise TypeError("Error: --enable-auto-tool-choice requires " + "--tool-call-parser") + + +def create_parser_for_docs() -> FlexibleArgumentParser: + parser_for_docs = FlexibleArgumentParser( + prog="-m vllm.entrypoints.openai.api_server") + return make_arg_parser(parser_for_docs) diff --git a/qwen3_6_scripts/mamba_cache.py b/qwen3_6_scripts/mamba_cache.py index a1ceea87..7537b9e9 100644 --- a/qwen3_6_scripts/mamba_cache.py +++ b/qwen3_6_scripts/mamba_cache.py @@ -1,224 +1,224 @@ -from typing import Dict, List, Optional - -import torch - -from vllm.attention.backends.abstract import AttentionMetadata - - -class MambaCacheManager: - - def __init__(self, dtype, num_mamba_layers, max_batch_size, - conv_state_shape, temporal_state_shape): - - conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) + - conv_state_shape, - dtype=dtype, - device="cuda") - temporal_state = torch.zeros(size=(num_mamba_layers, max_batch_size) + - temporal_state_shape, - dtype=dtype, - device="cuda") - - self.mamba_cache = (conv_state, temporal_state) - - # Maps between the request id and a dict that maps between the seq_id - # and its index inside the self.mamba_cache - self.mamba_cache_indices_mapping: Dict[str, Dict[int, int]] = {} - - def current_run_tensors(self, input_ids: torch.Tensor, - attn_metadata: AttentionMetadata, **kwargs): - """ - Return the tensors for the current run's conv and ssm state. - """ - if "seqlen_agnostic_capture_inputs" not in kwargs: - # We get here only on Prefill/Eager mode runs - request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"] - finished_requests_ids = kwargs["finished_requests_ids"] - - self._release_finished_requests(finished_requests_ids) - mamba_cache_tensors = self._prepare_current_run_mamba_cache( - request_ids_to_seq_ids, finished_requests_ids) - - else: - # CUDA graph capturing runs - mamba_cache_tensors = kwargs["seqlen_agnostic_capture_inputs"] - - return mamba_cache_tensors - - def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): - """ - Copy the relevant Mamba cache into the CUDA graph input buffer - that was provided during the capture runs - (JambaForCausalLM.mamba_gc_cache_buffer). - """ - assert all( - key in kwargs - for key in ["request_ids_to_seq_ids", "finished_requests_ids"]) - finished_requests_ids = kwargs["finished_requests_ids"] - request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"] - - self._release_finished_requests(finished_requests_ids) - self._prepare_current_run_mamba_cache(request_ids_to_seq_ids, - finished_requests_ids) - - def get_seqlen_agnostic_capture_inputs(self, batch_size: int): - """ - Provide the CUDA graph capture runs with a buffer in adjusted size. - The buffer is used to maintain the Mamba Cache during the CUDA graph - replay runs. - """ - return tuple(buffer[:, :batch_size] for buffer in self.mamba_cache) - - def _swap_mamba_cache(self, from_index: int, to_index: int): - assert len(self.mamba_cache) > 0 - for cache_t in self.mamba_cache: - cache_t[:, [to_index,from_index]] = \ - cache_t[:, [from_index,to_index]] - - def _copy_mamba_cache(self, from_index: int, to_index: int): - assert len(self.mamba_cache) > 0 - for cache_t in self.mamba_cache: - cache_t[:, to_index].copy_(cache_t[:, from_index], - non_blocking=True) - - def _move_out_if_already_occupied(self, index: int, - all_occupied_indices: List[int]): - if index in all_occupied_indices: - first_free_index = self._first_free_index_in_mamba_cache() - # In case occupied, move the occupied to a new empty block - self._move_cache_index_and_mappings(from_index=index, - to_index=first_free_index) - - def _assign_seq_id_to_mamba_cache_in_specific_dest(self, cur_rid: str, - seq_id: int, - destination_index: int): - """ - Assign (req_id,seq_id) pair to a `destination_index` index, if - already occupied, move the occupying index to a free index. - """ - all_occupied_indices = self._get_all_occupied_indices() - if cur_rid not in self.mamba_cache_indices_mapping: - self._move_out_if_already_occupied( - index=destination_index, - all_occupied_indices=all_occupied_indices) - for cache_t in self.mamba_cache: - cache_t[:, destination_index].zero_() - self.mamba_cache_indices_mapping[cur_rid] = { - seq_id: destination_index - } - elif seq_id not in (seq_ids2indices := - self.mamba_cache_indices_mapping[cur_rid]): - # parallel sampling , where n > 1, assume prefill have - # already happened now we only need to copy the already - # existing cache into the siblings seq_ids caches - self._move_out_if_already_occupied( - index=destination_index, - all_occupied_indices=all_occupied_indices) - index_exists = list(seq_ids2indices.values())[0] - # case of decoding n>1, copy prefill cache to decoding indices - self._copy_mamba_cache(from_index=index_exists, - to_index=destination_index) - self.mamba_cache_indices_mapping[cur_rid][ - seq_id] = destination_index - else: - # already exists - cache_index_already_exists = self.mamba_cache_indices_mapping[ - cur_rid][seq_id] - if cache_index_already_exists != destination_index: - # In case the seq id already exists but not in - # the right destination, swap it with what's occupying it - self._swap_pair_indices_and_mappings( - from_index=cache_index_already_exists, - to_index=destination_index) - - def _prepare_current_run_mamba_cache( - self, request_ids_to_seq_ids: Dict[str, list[int]], - finished_requests_ids: List[str]): - running_indices = [] - request_ids_to_seq_ids_flatten = [ - (req_id, seq_id) - for req_id, seq_ids in request_ids_to_seq_ids.items() - for seq_id in seq_ids - ] - batch_size = len(request_ids_to_seq_ids_flatten) - for dest_index, (request_id, - seq_id) in enumerate(request_ids_to_seq_ids_flatten): - if request_id in finished_requests_ids: - # Do not allocate cache index for requests that run - # and finish right after - continue - self._assign_seq_id_to_mamba_cache_in_specific_dest( - request_id, seq_id, dest_index) - running_indices.append(dest_index) - - self._clean_up_first_bs_blocks(batch_size, running_indices) - conv_state = self.mamba_cache[0][:, :batch_size] - temporal_state = self.mamba_cache[1][:, :batch_size] - - return (conv_state, temporal_state) - - def _get_all_occupied_indices(self): - return [ - cache_idx - for seq_ids2indices in self.mamba_cache_indices_mapping.values() - for cache_idx in seq_ids2indices.values() - ] - - def _clean_up_first_bs_blocks(self, batch_size: int, - indices_for_current_run: List[int]): - # move out all of the occupied but currently not running blocks - # outside of the first n blocks - destination_indices = range(batch_size) - max_possible_batch_size = self.mamba_cache[0].shape[1] - for destination_index in destination_indices: - if destination_index in self._get_all_occupied_indices() and \ - destination_index not in indices_for_current_run: - # move not running indices outside of the batch - all_other_indices = list( - range(batch_size, max_possible_batch_size)) - first_avail_index = self._first_free_index_in_mamba_cache( - all_other_indices) - self._swap_indices(from_index=destination_index, - to_index=first_avail_index) - - def _move_cache_index_and_mappings(self, from_index: int, to_index: int): - self._copy_mamba_cache(from_index=from_index, to_index=to_index) - self._update_mapping_index(from_index=from_index, to_index=to_index) - - def _swap_pair_indices_and_mappings(self, from_index: int, to_index: int): - self._swap_mamba_cache(from_index=from_index, to_index=to_index) - self._swap_mapping_index(from_index=from_index, to_index=to_index) - - def _swap_mapping_index(self, from_index: int, to_index: int): - for seq_ids2index in self.mamba_cache_indices_mapping.values(): - for seq_id, index in seq_ids2index.items(): - if from_index == index: - seq_ids2index.update({seq_id: to_index}) - elif to_index == index: - seq_ids2index.update({seq_id: from_index}) - - def _update_mapping_index(self, from_index: int, to_index: int): - for seq_ids2index in self.mamba_cache_indices_mapping.values(): - for seq_id, index in seq_ids2index.items(): - if from_index == index: - seq_ids2index.update({seq_id: to_index}) - return - - def _release_finished_requests(self, - finished_seq_groups_req_ids: List[str]): - for req_id in finished_seq_groups_req_ids: - if req_id in self.mamba_cache_indices_mapping: - self.mamba_cache_indices_mapping.pop(req_id) - - def _first_free_index_in_mamba_cache( - self, indices_range: Optional[List[int]] = None) -> int: - assert self.mamba_cache is not None - if indices_range is None: - max_possible_batch_size = self.mamba_cache[0].shape[1] - indices_range = list(range(max_possible_batch_size)) - all_occupied_indices = self._get_all_occupied_indices() - for i in indices_range: - if i not in all_occupied_indices: - return i - raise Exception("Couldn't find a free spot in the mamba cache! This" - "should never happen") +from typing import Dict, List, Optional + +import torch + +from vllm.attention.backends.abstract import AttentionMetadata + + +class MambaCacheManager: + + def __init__(self, dtype, num_mamba_layers, max_batch_size, + conv_state_shape, temporal_state_shape): + + conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) + + conv_state_shape, + dtype=dtype, + device="cuda") + temporal_state = torch.zeros(size=(num_mamba_layers, max_batch_size) + + temporal_state_shape, + dtype=dtype, + device="cuda") + + self.mamba_cache = (conv_state, temporal_state) + + # Maps between the request id and a dict that maps between the seq_id + # and its index inside the self.mamba_cache + self.mamba_cache_indices_mapping: Dict[str, Dict[int, int]] = {} + + def current_run_tensors(self, input_ids: torch.Tensor, + attn_metadata: AttentionMetadata, **kwargs): + """ + Return the tensors for the current run's conv and ssm state. + """ + if "seqlen_agnostic_capture_inputs" not in kwargs: + # We get here only on Prefill/Eager mode runs + request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"] + finished_requests_ids = kwargs["finished_requests_ids"] + + self._release_finished_requests(finished_requests_ids) + mamba_cache_tensors = self._prepare_current_run_mamba_cache( + request_ids_to_seq_ids, finished_requests_ids) + + else: + # CUDA graph capturing runs + mamba_cache_tensors = kwargs["seqlen_agnostic_capture_inputs"] + + return mamba_cache_tensors + + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): + """ + Copy the relevant Mamba cache into the CUDA graph input buffer + that was provided during the capture runs + (JambaForCausalLM.mamba_gc_cache_buffer). + """ + assert all( + key in kwargs + for key in ["request_ids_to_seq_ids", "finished_requests_ids"]) + finished_requests_ids = kwargs["finished_requests_ids"] + request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"] + + self._release_finished_requests(finished_requests_ids) + self._prepare_current_run_mamba_cache(request_ids_to_seq_ids, + finished_requests_ids) + + def get_seqlen_agnostic_capture_inputs(self, batch_size: int): + """ + Provide the CUDA graph capture runs with a buffer in adjusted size. + The buffer is used to maintain the Mamba Cache during the CUDA graph + replay runs. + """ + return tuple(buffer[:, :batch_size] for buffer in self.mamba_cache) + + def _swap_mamba_cache(self, from_index: int, to_index: int): + assert len(self.mamba_cache) > 0 + for cache_t in self.mamba_cache: + cache_t[:, [to_index,from_index]] = \ + cache_t[:, [from_index,to_index]] + + def _copy_mamba_cache(self, from_index: int, to_index: int): + assert len(self.mamba_cache) > 0 + for cache_t in self.mamba_cache: + cache_t[:, to_index].copy_(cache_t[:, from_index], + non_blocking=True) + + def _move_out_if_already_occupied(self, index: int, + all_occupied_indices: List[int]): + if index in all_occupied_indices: + first_free_index = self._first_free_index_in_mamba_cache() + # In case occupied, move the occupied to a new empty block + self._move_cache_index_and_mappings(from_index=index, + to_index=first_free_index) + + def _assign_seq_id_to_mamba_cache_in_specific_dest(self, cur_rid: str, + seq_id: int, + destination_index: int): + """ + Assign (req_id,seq_id) pair to a `destination_index` index, if + already occupied, move the occupying index to a free index. + """ + all_occupied_indices = self._get_all_occupied_indices() + if cur_rid not in self.mamba_cache_indices_mapping: + self._move_out_if_already_occupied( + index=destination_index, + all_occupied_indices=all_occupied_indices) + for cache_t in self.mamba_cache: + cache_t[:, destination_index].zero_() + self.mamba_cache_indices_mapping[cur_rid] = { + seq_id: destination_index + } + elif seq_id not in (seq_ids2indices := + self.mamba_cache_indices_mapping[cur_rid]): + # parallel sampling , where n > 1, assume prefill have + # already happened now we only need to copy the already + # existing cache into the siblings seq_ids caches + self._move_out_if_already_occupied( + index=destination_index, + all_occupied_indices=all_occupied_indices) + index_exists = list(seq_ids2indices.values())[0] + # case of decoding n>1, copy prefill cache to decoding indices + self._copy_mamba_cache(from_index=index_exists, + to_index=destination_index) + self.mamba_cache_indices_mapping[cur_rid][ + seq_id] = destination_index + else: + # already exists + cache_index_already_exists = self.mamba_cache_indices_mapping[ + cur_rid][seq_id] + if cache_index_already_exists != destination_index: + # In case the seq id already exists but not in + # the right destination, swap it with what's occupying it + self._swap_pair_indices_and_mappings( + from_index=cache_index_already_exists, + to_index=destination_index) + + def _prepare_current_run_mamba_cache( + self, request_ids_to_seq_ids: Dict[str, list[int]], + finished_requests_ids: List[str]): + running_indices = [] + request_ids_to_seq_ids_flatten = [ + (req_id, seq_id) + for req_id, seq_ids in request_ids_to_seq_ids.items() + for seq_id in seq_ids + ] + batch_size = len(request_ids_to_seq_ids_flatten) + for dest_index, (request_id, + seq_id) in enumerate(request_ids_to_seq_ids_flatten): + if request_id in finished_requests_ids: + # Do not allocate cache index for requests that run + # and finish right after + continue + self._assign_seq_id_to_mamba_cache_in_specific_dest( + request_id, seq_id, dest_index) + running_indices.append(dest_index) + + self._clean_up_first_bs_blocks(batch_size, running_indices) + conv_state = self.mamba_cache[0][:, :batch_size] + temporal_state = self.mamba_cache[1][:, :batch_size] + + return (conv_state, temporal_state) + + def _get_all_occupied_indices(self): + return [ + cache_idx + for seq_ids2indices in self.mamba_cache_indices_mapping.values() + for cache_idx in seq_ids2indices.values() + ] + + def _clean_up_first_bs_blocks(self, batch_size: int, + indices_for_current_run: List[int]): + # move out all of the occupied but currently not running blocks + # outside of the first n blocks + destination_indices = range(batch_size) + max_possible_batch_size = self.mamba_cache[0].shape[1] + for destination_index in destination_indices: + if destination_index in self._get_all_occupied_indices() and \ + destination_index not in indices_for_current_run: + # move not running indices outside of the batch + all_other_indices = list( + range(batch_size, max_possible_batch_size)) + first_avail_index = self._first_free_index_in_mamba_cache( + all_other_indices) + self._swap_indices(from_index=destination_index, + to_index=first_avail_index) + + def _move_cache_index_and_mappings(self, from_index: int, to_index: int): + self._copy_mamba_cache(from_index=from_index, to_index=to_index) + self._update_mapping_index(from_index=from_index, to_index=to_index) + + def _swap_pair_indices_and_mappings(self, from_index: int, to_index: int): + self._swap_mamba_cache(from_index=from_index, to_index=to_index) + self._swap_mapping_index(from_index=from_index, to_index=to_index) + + def _swap_mapping_index(self, from_index: int, to_index: int): + for seq_ids2index in self.mamba_cache_indices_mapping.values(): + for seq_id, index in seq_ids2index.items(): + if from_index == index: + seq_ids2index.update({seq_id: to_index}) + elif to_index == index: + seq_ids2index.update({seq_id: from_index}) + + def _update_mapping_index(self, from_index: int, to_index: int): + for seq_ids2index in self.mamba_cache_indices_mapping.values(): + for seq_id, index in seq_ids2index.items(): + if from_index == index: + seq_ids2index.update({seq_id: to_index}) + return + + def _release_finished_requests(self, + finished_seq_groups_req_ids: List[str]): + for req_id in finished_seq_groups_req_ids: + if req_id in self.mamba_cache_indices_mapping: + self.mamba_cache_indices_mapping.pop(req_id) + + def _first_free_index_in_mamba_cache( + self, indices_range: Optional[List[int]] = None) -> int: + assert self.mamba_cache is not None + if indices_range is None: + max_possible_batch_size = self.mamba_cache[0].shape[1] + indices_range = list(range(max_possible_batch_size)) + all_occupied_indices = self._get_all_occupied_indices() + for i in indices_range: + if i not in all_occupied_indices: + return i + raise Exception("Couldn't find a free spot in the mamba cache! This" + "should never happen") diff --git a/qwen3_6_scripts/patch_vllm_qwen3_5.py b/qwen3_6_scripts/patch_vllm_qwen3_5.py index 55313fc9..d9731594 100644 --- a/qwen3_6_scripts/patch_vllm_qwen3_5.py +++ b/qwen3_6_scripts/patch_vllm_qwen3_5.py @@ -1,11 +1,11 @@ """ Patches the vLLM model registry and deploys the Qwen3_5 model file. - + Deploy steps on the remote machine: 1. patch_ops.sh locates vLLM with importlib.util.find_spec. 2. cp modified_scripts/qwen3_5.py into the detected vllm model directory. 2. python3 modified_scripts/patch_vllm_qwen3_5.py - + The registry patch installs Qwen3.6 aliases so /model/config.json does not need to be edited by hand. """ @@ -26,8 +26,8 @@ EXPECTED_REGISTRY_ENTRIES = ( '"Qwen3_6ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM")', '"Qwen3_6MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM")', ) - - + + def main(): print(f"=== Patching {REGISTRY} ===") replace_once( @@ -42,7 +42,7 @@ def main(): ' "Qwen3_6MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),', required=True, already_contains='"Qwen3_6MoeForCausalLM"') - + print("\n=== Static verification ===") model_source = MODEL.read_text(encoding="utf-8") tree = ast.parse(model_source, filename=str(MODEL)) @@ -67,7 +67,7 @@ def main(): print(f" registry aliases verified: {len(EXPECTED_REGISTRY_ENTRIES)}") print("\nDone. Registry aliases installed; do not edit /model/config.json.") - - -if __name__ == "__main__": - main() + + +if __name__ == "__main__": + main() diff --git a/qwen3_6_scripts/patch_vllm_tool_parser.py b/qwen3_6_scripts/patch_vllm_tool_parser.py index 18463a53..e0bdd40c 100644 --- a/qwen3_6_scripts/patch_vllm_tool_parser.py +++ b/qwen3_6_scripts/patch_vllm_tool_parser.py @@ -1,23 +1,23 @@ -""" -Patches vLLM 0.6.3 to register Qwen3CoderToolParser under the name "qwen3_coder". - +""" +Patches vLLM 0.6.3 to register Qwen3CoderToolParser under the name "qwen3_coder". + Deploy steps on the remote machine (already called by patch_ops.sh): 1. patch_ops.sh locates vLLM with importlib.util.find_spec. 2. cp qwen3coder_tool_parser.py into the detected vllm tool_parsers. 2. python3 patch_vllm_tool_parser.py - -Usage after patching: - --tool-call-parser qwen3_coder --enable-auto-tool-choice -""" - + +Usage after patching: + --tool-call-parser qwen3_coder --enable-auto-tool-choice +""" + from patch_utils import ensure_dir, package_root, replace_once VLLM_ROOT = package_root("vllm") TOOL_PARSERS_DIR = VLLM_ROOT / "entrypoints" / "openai" / "tool_parsers" INIT_FILE = TOOL_PARSERS_DIR / "__init__.py" - - -def main(): + + +def main(): ensure_dir(TOOL_PARSERS_DIR) print(f"=== Patching {INIT_FILE} ===") @@ -35,23 +35,23 @@ def main(): ' "Qwen3CoderToolParser"\n]', required=True, already_contains='"Qwen3CoderToolParser"') - - print("\n=== Verification ===") - try: - import importlib.util - spec = importlib.util.spec_from_file_location( - "qwen3coder_tool_parser", + + print("\n=== Verification ===") + try: + import importlib.util + spec = importlib.util.spec_from_file_location( + "qwen3coder_tool_parser", str(TOOL_PARSERS_DIR / "qwen3coder_tool_parser.py"), - ) - mod = importlib.util.module_from_spec(spec) - print(f" Module spec loaded: {spec.name}") - print(" (full import requires torch/vllm runtime — skipping exec)") - except Exception as e: + ) + mod = importlib.util.module_from_spec(spec) + print(f" Module spec loaded: {spec.name}") + print(" (full import requires torch/vllm runtime — skipping exec)") + except Exception as e: print(f" [optional] spec check failed: {e}") - - print("\nDone. Start vLLM server with:") - print(" --tool-call-parser qwen3_coder --enable-auto-tool-choice") - - -if __name__ == "__main__": - main() + + print("\nDone. Start vLLM server with:") + print(" --tool-call-parser qwen3_coder --enable-auto-tool-choice") + + +if __name__ == "__main__": + main() diff --git a/qwen3_6_scripts/patch_xformers_sdpa_batch.py b/qwen3_6_scripts/patch_xformers_sdpa_batch.py index 72315b58..59d494b3 100644 --- a/qwen3_6_scripts/patch_xformers_sdpa_batch.py +++ b/qwen3_6_scripts/patch_xformers_sdpa_batch.py @@ -1,29 +1,29 @@ """ -策略：批量（block-diagonal）fallback — 纯 PyTorch 数学实现 -============================================================= -构建块对角 causal mask，对整批序列一次 matmul + softmax， -完全绕开所有硬件 flash attention kernel。 - -背景： - ixformer flshattF: head_dim > 128 报错拒绝 - cudnnFlashAttnForward: 接受 head_dim=256，但数值结果错误（输出全"!"） - 两者大概率是同一硬件单元，ixformer 提前拦截了硬件不支持的配置。 - 纯 matmul 路径完全绕开硬件 flash attention，数值正确。 - -优点： - 数值正确。 - 并发请求 prefill attention 在 GPU 上真正并行（一次大 matmul）。 - -缺点： - 峰值显存 = total_tokens² × H × dtype_size - total_tokens 受 --max-num-batched-tokens 控制，max-model-len 控制不住。 - -内存参考（fp16，H_local=6，--max-num-batched-tokens=T）： - T=2048 → 峰值 ~50 MB - T=4096 → 峰值 ~200 MB - T=8192 → 峰值 ~800 MB - T=16384 → 峰值 ~3.2 GB - +策略：批量（block-diagonal）fallback — 纯 PyTorch 数学实现 +============================================================= +构建块对角 causal mask，对整批序列一次 matmul + softmax， +完全绕开所有硬件 flash attention kernel。 + +背景： + ixformer flshattF: head_dim > 128 报错拒绝 + cudnnFlashAttnForward: 接受 head_dim=256，但数值结果错误（输出全"!"） + 两者大概率是同一硬件单元，ixformer 提前拦截了硬件不支持的配置。 + 纯 matmul 路径完全绕开硬件 flash attention，数值正确。 + +优点： + 数值正确。 + 并发请求 prefill attention 在 GPU 上真正并行（一次大 matmul）。 + +缺点： + 峰值显存 = total_tokens² × H × dtype_size + total_tokens 受 --max-num-batched-tokens 控制，max-model-len 控制不住。 + +内存参考（fp16，H_local=6，--max-num-batched-tokens=T）： + T=2048 → 峰值 ~50 MB + T=4096 → 峰值 ~200 MB + T=8192 → 峰值 ~800 MB + T=16384 → 峰值 ~3.2 GB + Deploy: python3 modified_scripts/patch_xformers_sdpa_batch.py """ @@ -31,126 +31,126 @@ Deploy: from patch_utils import package_root, replace_once XFORMERS_PATH = package_root("vllm") / "attention" / "backends" / "xformers.py" - -FALLBACK_METHOD = ''' - def _run_sdpa_fallback( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: "XFormersMetadata", - ) -> torch.Tensor: - """批量纯数学 attention fallback。 - - 构建块对角 causal mask（等价于 ixformer BlockDiagonalCausalMask）， - 对整批序列一次 matmul + softmax，GPU 并行处理所有序列。 - - 块对角 mask 结构（seq1 len=3，seq2 len=2）： - s1,0 s1,1 s1,2 s2,0 s2,1 - s1,0 [ 0 -inf -inf -inf -inf ] - s1,1 [ 0 0 -inf -inf -inf ] - s1,2 [ 0 0 0 -inf -inf ] - s2,0 [-inf -inf -inf 0 -inf ] - s2,1 [-inf -inf -inf 0 0 ] - - softmax 在 float32 下计算防止 float16 溢出，结果转回原始 dtype。 - - Args: - query : [1, total_prefill_tokens, num_heads, head_dim] - key : [1, total_prefill_tokens, num_kv_heads, head_dim] - value : [1, total_prefill_tokens, num_kv_heads, head_dim] - Returns: - [1, total_prefill_tokens, num_heads, head_dim] - """ - assert attn_metadata.seq_lens is not None - orig_dtype = query.dtype - total_tokens = query.shape[1] - - # ── 构建块对角 causal mask [T, T] ──────────────────────────────── - # 全部初始化为 -inf，再对每条序列的对角块填入下三角 0 - mask = torch.full( - (total_tokens, total_tokens), - float("-inf"), - dtype=torch.float32, - device=query.device, - ) - start = 0 - for seq_len in attn_metadata.seq_lens: - end = start + seq_len - mask[start:end, start:end] = torch.tril( - torch.zeros(seq_len, seq_len, - dtype=torch.float32, device=query.device) - ) - start = end - - # ── [1, H, T, D]，.contiguous() ────────────────────────────────── - q_all = query.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - k_all = key.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - v_all = value.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - - # ── GQA：展开 KV heads ──────────────────────────────────────────── - if k_all.shape[1] != q_all.shape[1]: - n = q_all.shape[1] // k_all.shape[1] - k_all = k_all.repeat_interleave(n, dim=1).contiguous() - v_all = v_all.repeat_interleave(n, dim=1).contiguous() - - # ── 纯数学 attention（float32 防溢出）──────────────────────────── - # [1, H, T, T] - attn_w = torch.matmul(q_all.float(), k_all.float().transpose(-2, -1)) - attn_w = attn_w * self.scale - attn_w = attn_w + mask # 加法广播：mask [T,T] → [1, H, T, T] - attn_w = torch.softmax(attn_w, dim=-1) - - out = torch.matmul(attn_w, v_all.float()).to(orig_dtype) - # [1, H, T, D] → [1, T, H, D] - return out.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - -''' - -OLD_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op = self.attn_op - ) - return out.view_as(original_query)\ -""" - -NEW_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - if self.head_size > 128: - out = self._run_sdpa_fallback(query, key, value, attn_metadata) - else: - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op=self.attn_op, - ) - return out.view_as(original_query)\ -""" - -INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward(" - - + +FALLBACK_METHOD = ''' + def _run_sdpa_fallback( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: "XFormersMetadata", + ) -> torch.Tensor: + """批量纯数学 attention fallback。 + + 构建块对角 causal mask（等价于 ixformer BlockDiagonalCausalMask）， + 对整批序列一次 matmul + softmax，GPU 并行处理所有序列。 + + 块对角 mask 结构（seq1 len=3，seq2 len=2）： + s1,0 s1,1 s1,2 s2,0 s2,1 + s1,0 [ 0 -inf -inf -inf -inf ] + s1,1 [ 0 0 -inf -inf -inf ] + s1,2 [ 0 0 0 -inf -inf ] + s2,0 [-inf -inf -inf 0 -inf ] + s2,1 [-inf -inf -inf 0 0 ] + + softmax 在 float32 下计算防止 float16 溢出，结果转回原始 dtype。 + + Args: + query : [1, total_prefill_tokens, num_heads, head_dim] + key : [1, total_prefill_tokens, num_kv_heads, head_dim] + value : [1, total_prefill_tokens, num_kv_heads, head_dim] + Returns: + [1, total_prefill_tokens, num_heads, head_dim] + """ + assert attn_metadata.seq_lens is not None + orig_dtype = query.dtype + total_tokens = query.shape[1] + + # ── 构建块对角 causal mask [T, T] ──────────────────────────────── + # 全部初始化为 -inf，再对每条序列的对角块填入下三角 0 + mask = torch.full( + (total_tokens, total_tokens), + float("-inf"), + dtype=torch.float32, + device=query.device, + ) + start = 0 + for seq_len in attn_metadata.seq_lens: + end = start + seq_len + mask[start:end, start:end] = torch.tril( + torch.zeros(seq_len, seq_len, + dtype=torch.float32, device=query.device) + ) + start = end + + # ── [1, H, T, D]，.contiguous() ────────────────────────────────── + q_all = query.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) + k_all = key.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) + v_all = value.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) + + # ── GQA：展开 KV heads ──────────────────────────────────────────── + if k_all.shape[1] != q_all.shape[1]: + n = q_all.shape[1] // k_all.shape[1] + k_all = k_all.repeat_interleave(n, dim=1).contiguous() + v_all = v_all.repeat_interleave(n, dim=1).contiguous() + + # ── 纯数学 attention（float32 防溢出）──────────────────────────── + # [1, H, T, T] + attn_w = torch.matmul(q_all.float(), k_all.float().transpose(-2, -1)) + attn_w = attn_w * self.scale + attn_w = attn_w + mask # 加法广播：mask [T,T] → [1, H, T, T] + attn_w = torch.softmax(attn_w, dim=-1) + + out = torch.matmul(attn_w, v_all.float()).to(orig_dtype) + # [1, H, T, D] → [1, T, H, D] + return out.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) + +''' + +OLD_XFORMER_BLOCK = """\ + self.attn_op = xops.fmha.flash.FwOp() + if self.alibi_slopes is None: + # Add the batch dimension. + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + out = xops.memory_efficient_attention_forward( + query, + key, + value, + attn_bias=attn_bias[0], + p=0.0, + scale=self.scale, + op = self.attn_op + ) + return out.view_as(original_query)\ +""" + +NEW_XFORMER_BLOCK = """\ + self.attn_op = xops.fmha.flash.FwOp() + if self.alibi_slopes is None: + # Add the batch dimension. + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + if self.head_size > 128: + out = self._run_sdpa_fallback(query, key, value, attn_metadata) + else: + out = xops.memory_efficient_attention_forward( + query, + key, + value, + attn_bias=attn_bias[0], + p=0.0, + scale=self.scale, + op=self.attn_op, + ) + return out.view_as(original_query)\ +""" + +INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward(" + + def patch_file(path): replace_once( path, @@ -164,14 +164,14 @@ def patch_file(path): NEW_XFORMER_BLOCK, required=True, already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)") - - -def main(): - print("=== patch_xformers_sdpa_batch (batch, pure-math) ===") - print(f"Target: {XFORMERS_PATH}") - patch_file(XFORMERS_PATH) - print("\nDone.") - - -if __name__ == "__main__": - main() + + +def main(): + print("=== patch_xformers_sdpa_batch (batch, pure-math) ===") + print(f"Target: {XFORMERS_PATH}") + patch_file(XFORMERS_PATH) + print("\nDone.") + + +if __name__ == "__main__": + main() diff --git a/qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py b/qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py index f5212087..78a33836 100644 --- a/qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py +++ b/qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py @@ -1,26 +1,26 @@ -""" -策略：批量（block-diagonal）— F.scaled_dot_product_attention，可走硬件 kernel -============================================================================= -构建块对角 causal mask，对整批序列一次 F.scaled_dot_product_attention。 -与 patch_xformers_sdpa_batch.py（纯 matmul）的区别： - SDPA 会根据 PyTorch/驱动能力分发到最优 kernel（Flash Attention / - mem-efficient attention / math fallback），而不是固定走 cublas matmul。 - -历史说明： - 该方案最早因输出全"!"而被弃用，后续排查确认"!"由 mamba_cache.py bug - 引起，与 attention 实现无关。当前恢复此方案用于性能对比测试。 - -已知硬件限制（BI-V100）： - cudnnFlashAttnForward 不支持 is_causal=True（报错）。 - 本实现使用 is_causal=False + 显式块对角 additive mask 规避此限制。 - 若 SDPA 仍分发到有问题的 kernel，回退到 patch_xformers_sdpa_batch.py。 - -优点（vs 纯 matmul）： - SDPA 可分发到 Flash Attention kernel → O(L) 显存、更快的 CUDA kernel。 - -缺点： - 依赖硬件 kernel 行为，若 kernel 有 bug 则数值错误（需与 matmul 版对比验证）。 - +""" +策略：批量（block-diagonal）— F.scaled_dot_product_attention，可走硬件 kernel +============================================================================= +构建块对角 causal mask，对整批序列一次 F.scaled_dot_product_attention。 +与 patch_xformers_sdpa_batch.py（纯 matmul）的区别： + SDPA 会根据 PyTorch/驱动能力分发到最优 kernel（Flash Attention / + mem-efficient attention / math fallback），而不是固定走 cublas matmul。 + +历史说明： + 该方案最早因输出全"!"而被弃用，后续排查确认"!"由 mamba_cache.py bug + 引起，与 attention 实现无关。当前恢复此方案用于性能对比测试。 + +已知硬件限制（BI-V100）： + cudnnFlashAttnForward 不支持 is_causal=True（报错）。 + 本实现使用 is_causal=False + 显式块对角 additive mask 规避此限制。 + 若 SDPA 仍分发到有问题的 kernel，回退到 patch_xformers_sdpa_batch.py。 + +优点（vs 纯 matmul）： + SDPA 可分发到 Flash Attention kernel → O(L) 显存、更快的 CUDA kernel。 + +缺点： + 依赖硬件 kernel 行为，若 kernel 有 bug 则数值错误（需与 matmul 版对比验证）。 + Deploy: python3 modified_scripts/patch_xformers_sdpa_batch_kernel.py """ @@ -28,128 +28,128 @@ Deploy: from patch_utils import package_root, replace_once XFORMERS_PATH = package_root("vllm") / "attention" / "backends" / "xformers.py" - -FALLBACK_METHOD = ''' - def _run_sdpa_fallback( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: "XFormersMetadata", - ) -> torch.Tensor: - """批量 F.scaled_dot_product_attention fallback（可走硬件 kernel）。 - - 构建块对角 causal mask，对整批序列一次 SDPA 调用。 - SDPA 可分发到 Flash Attention / mem-efficient attention kernel。 - is_causal=False + 显式 additive mask，规避 cudnnFlashAttnForward - 不支持 is_causal=True 的限制。 - - 块对角 mask（seq1 len=3，seq2 len=2）： - s1,0 s1,1 s1,2 s2,0 s2,1 - s1,0 [ 0 -inf -inf -inf -inf ] - s1,1 [ 0 0 -inf -inf -inf ] - s1,2 [ 0 0 0 -inf -inf ] - s2,0 [-inf -inf -inf 0 -inf ] - s2,1 [-inf -inf -inf 0 0 ] - - Args: - query : [1, total_prefill_tokens, num_heads, head_dim] - key : [1, total_prefill_tokens, num_kv_heads, head_dim] - value : [1, total_prefill_tokens, num_kv_heads, head_dim] - Returns: - [1, total_prefill_tokens, num_heads, head_dim] - """ - import torch.nn.functional as F - - assert attn_metadata.seq_lens is not None - orig_dtype = query.dtype - total_tokens = query.shape[1] - - # ── 块对角 causal mask [T, T] ───────────────────────────────────── - mask = torch.full( - (total_tokens, total_tokens), - float("-inf"), - dtype=orig_dtype, - device=query.device, - ) - start = 0 - for seq_len in attn_metadata.seq_lens: - end = start + seq_len - mask[start:end, start:end] = torch.tril( - torch.zeros(seq_len, seq_len, dtype=orig_dtype, device=query.device) - ) - start = end - - # ── [1, H, T, D] ────────────────────────────────────────────────── - q_all = query.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - k_all = key.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - v_all = value.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - - # ── GQA：展开 KV heads ──────────────────────────────────────────── - if k_all.shape[1] != q_all.shape[1]: - n = q_all.shape[1] // k_all.shape[1] - k_all = k_all.repeat_interleave(n, dim=1).contiguous() - v_all = v_all.repeat_interleave(n, dim=1).contiguous() - - # ── F.scaled_dot_product_attention（可走硬件 kernel）───────────── - # is_causal=False：避免 cudnnFlashAttnForward "not support causal mode" - # attn_mask 传 additive float mask（非 bool），SDPA 选择 math/kernel 路径 - out = F.scaled_dot_product_attention( - q_all, k_all, v_all, - attn_mask=mask, - dropout_p=0.0, - is_causal=False, - scale=self.scale, - ) - # [1, H, T, D] → [1, T, H, D] - return out.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - -''' - -OLD_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op = self.attn_op - ) - return out.view_as(original_query)\ -""" - -NEW_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - if self.head_size > 128: - out = self._run_sdpa_fallback(query, key, value, attn_metadata) - else: - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op=self.attn_op, - ) - return out.view_as(original_query)\ -""" - -INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward(" - - + +FALLBACK_METHOD = ''' + def _run_sdpa_fallback( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: "XFormersMetadata", + ) -> torch.Tensor: + """批量 F.scaled_dot_product_attention fallback（可走硬件 kernel）。 + + 构建块对角 causal mask，对整批序列一次 SDPA 调用。 + SDPA 可分发到 Flash Attention / mem-efficient attention kernel。 + is_causal=False + 显式 additive mask，规避 cudnnFlashAttnForward + 不支持 is_causal=True 的限制。 + + 块对角 mask（seq1 len=3，seq2 len=2）： + s1,0 s1,1 s1,2 s2,0 s2,1 + s1,0 [ 0 -inf -inf -inf -inf ] + s1,1 [ 0 0 -inf -inf -inf ] + s1,2 [ 0 0 0 -inf -inf ] + s2,0 [-inf -inf -inf 0 -inf ] + s2,1 [-inf -inf -inf 0 0 ] + + Args: + query : [1, total_prefill_tokens, num_heads, head_dim] + key : [1, total_prefill_tokens, num_kv_heads, head_dim] + value : [1, total_prefill_tokens, num_kv_heads, head_dim] + Returns: + [1, total_prefill_tokens, num_heads, head_dim] + """ + import torch.nn.functional as F + + assert attn_metadata.seq_lens is not None + orig_dtype = query.dtype + total_tokens = query.shape[1] + + # ── 块对角 causal mask [T, T] ───────────────────────────────────── + mask = torch.full( + (total_tokens, total_tokens), + float("-inf"), + dtype=orig_dtype, + device=query.device, + ) + start = 0 + for seq_len in attn_metadata.seq_lens: + end = start + seq_len + mask[start:end, start:end] = torch.tril( + torch.zeros(seq_len, seq_len, dtype=orig_dtype, device=query.device) + ) + start = end + + # ── [1, H, T, D] ────────────────────────────────────────────────── + q_all = query.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) + k_all = key.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) + v_all = value.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) + + # ── GQA：展开 KV heads ──────────────────────────────────────────── + if k_all.shape[1] != q_all.shape[1]: + n = q_all.shape[1] // k_all.shape[1] + k_all = k_all.repeat_interleave(n, dim=1).contiguous() + v_all = v_all.repeat_interleave(n, dim=1).contiguous() + + # ── F.scaled_dot_product_attention（可走硬件 kernel）───────────── + # is_causal=False：避免 cudnnFlashAttnForward "not support causal mode" + # attn_mask 传 additive float mask（非 bool），SDPA 选择 math/kernel 路径 + out = F.scaled_dot_product_attention( + q_all, k_all, v_all, + attn_mask=mask, + dropout_p=0.0, + is_causal=False, + scale=self.scale, + ) + # [1, H, T, D] → [1, T, H, D] + return out.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) + +''' + +OLD_XFORMER_BLOCK = """\ + self.attn_op = xops.fmha.flash.FwOp() + if self.alibi_slopes is None: + # Add the batch dimension. + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + out = xops.memory_efficient_attention_forward( + query, + key, + value, + attn_bias=attn_bias[0], + p=0.0, + scale=self.scale, + op = self.attn_op + ) + return out.view_as(original_query)\ +""" + +NEW_XFORMER_BLOCK = """\ + self.attn_op = xops.fmha.flash.FwOp() + if self.alibi_slopes is None: + # Add the batch dimension. + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + if self.head_size > 128: + out = self._run_sdpa_fallback(query, key, value, attn_metadata) + else: + out = xops.memory_efficient_attention_forward( + query, + key, + value, + attn_bias=attn_bias[0], + p=0.0, + scale=self.scale, + op=self.attn_op, + ) + return out.view_as(original_query)\ +""" + +INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward(" + + def patch_file(path): replace_once( path, @@ -163,14 +163,14 @@ def patch_file(path): NEW_XFORMER_BLOCK, required=True, already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)") - - -def main(): - print("=== patch_xformers_sdpa_batch_kernel (batch, F.sdpa + kernel dispatch) ===") - print(f"Target: {XFORMERS_PATH}") - patch_file(XFORMERS_PATH) - print("\nDone.") - - -if __name__ == "__main__": - main() + + +def main(): + print("=== patch_xformers_sdpa_batch_kernel (batch, F.sdpa + kernel dispatch) ===") + print(f"Target: {XFORMERS_PATH}") + patch_file(XFORMERS_PATH) + print("\nDone.") + + +if __name__ == "__main__": + main() diff --git a/qwen3_6_scripts/patch_xformers_sdpa_seq.py b/qwen3_6_scripts/patch_xformers_sdpa_seq.py index 26c72589..9f40bd46 100644 --- a/qwen3_6_scripts/patch_xformers_sdpa_seq.py +++ b/qwen3_6_scripts/patch_xformers_sdpa_seq.py @@ -1,39 +1,39 @@ -""" -策略：顺序（per-sequence）fallback — 纯 PyTorch 数学实现 -========================================================== -逐条序列用 matmul + softmax 手写 attention，完全绕开所有硬件 -flash attention kernel（ixformer / cudnnFlashAttnForward）。 - -背景： - Iluvatar cudnnFlashAttnForward 存在两个已知问题： - 1. 不支持 is_causal=True（报错） - 2. 使用 attn_mask 路径时数值结果不正确（静默错误，输出全为"!"） - 与华为昇腾 910B4 上 llama.cpp --flash-attn off 修复同类问题的原理相同。 - 纯数学路径（matmul + softmax）在任何 PyTorch 后端上结果都正确。 - -优点： - 数值正确，不依赖任何硬件特定 attention kernel。 - 峰值显存 = max(seq_len)² × H × dtype_size，由 --max-model-len 控制。 - -缺点： - 并发请求的 prefill attention 串行执行。 - O(L²) 显存（无 flash attention 的 O(L) 优化）。 - -内存参考（fp16，H_local=6）： - max-model-len=4096 → 峰值 ~200 MB - max-model-len=8192 → 峰值 ~800 MB - max-model-len=16384 → 峰值 ~3.2 GB - -额外 patch（arg_utils.py）： - vllm 0.6.3 在 max_model_len > 32K 时会自动开启 chunked prefill（无命令行 - 关闭选项），原意是防止 profiling OOM。但 _run_sdpa_fallback 已通过 Q-tiling - 解决了该问题，chunked prefill 反而会把推理路径从 _run_sdpa_fallback 切换到 - _forward_prefix_pytorch，属于不必要的行为变更，因此一并禁用该自动逻辑。 - -Deploy: - python3 modified_scripts/patch_xformers_sdpa_seq.py -""" - +""" +策略：顺序（per-sequence）fallback — 纯 PyTorch 数学实现 +========================================================== +逐条序列用 matmul + softmax 手写 attention，完全绕开所有硬件 +flash attention kernel（ixformer / cudnnFlashAttnForward）。 + +背景： + Iluvatar cudnnFlashAttnForward 存在两个已知问题： + 1. 不支持 is_causal=True（报错） + 2. 使用 attn_mask 路径时数值结果不正确（静默错误，输出全为"!"） + 与华为昇腾 910B4 上 llama.cpp --flash-attn off 修复同类问题的原理相同。 + 纯数学路径（matmul + softmax）在任何 PyTorch 后端上结果都正确。 + +优点： + 数值正确，不依赖任何硬件特定 attention kernel。 + 峰值显存 = max(seq_len)² × H × dtype_size，由 --max-model-len 控制。 + +缺点： + 并发请求的 prefill attention 串行执行。 + O(L²) 显存（无 flash attention 的 O(L) 优化）。 + +内存参考（fp16，H_local=6）： + max-model-len=4096 → 峰值 ~200 MB + max-model-len=8192 → 峰值 ~800 MB + max-model-len=16384 → 峰值 ~3.2 GB + +额外 patch（arg_utils.py）： + vllm 0.6.3 在 max_model_len > 32K 时会自动开启 chunked prefill（无命令行 + 关闭选项），原意是防止 profiling OOM。但 _run_sdpa_fallback 已通过 Q-tiling + 解决了该问题，chunked prefill 反而会把推理路径从 _run_sdpa_fallback 切换到 + _forward_prefix_pytorch，属于不必要的行为变更，因此一并禁用该自动逻辑。 + +Deploy: + python3 modified_scripts/patch_xformers_sdpa_seq.py +""" + from patch_utils import package_root, replace_one_of, replace_once VLLM_ROOT = package_root("vllm") @@ -44,24 +44,24 @@ LOGITS_PROC_PATH = ( OUTLINES_DECODING_PATH = ( VLLM_ROOT / "model_executor" / "guided_decoding" / "outlines_decoding.py") - -# _apply_logits_processors crashes when seq_groups is None (intermediate -# chunked-prefill chunks on the driver rank). Add an early-return guard. -_LP_OLD_BLOCK = """\ -def _apply_logits_processors( - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, -) -> torch.Tensor: - found_logits_processors = False\ -""" - + +# _apply_logits_processors crashes when seq_groups is None (intermediate +# chunked-prefill chunks on the driver rank). Add an early-return guard. +_LP_OLD_BLOCK = """\ +def _apply_logits_processors( + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, +) -> torch.Tensor: + found_logits_processors = False\ +""" + _LP_NEW_BLOCK = """\ -def _apply_logits_processors( - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, -) -> torch.Tensor: - if sampling_metadata.seq_groups is None: # intermediate chunked-prefill chunk - return logits +def _apply_logits_processors( + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, +) -> torch.Tensor: + if sampling_metadata.seq_groups is None: # intermediate chunked-prefill chunk + return logits found_logits_processors = False\ """ @@ -116,26 +116,26 @@ _ws : JSON_WS? JSON_STRING: /"(\\["\\\/bfnrt]|\\u[0-9a-fA-F]{4}|[^"\\\x00-\x1f])*"/ JSON_WS: /[ \t\r\n]{1,4}/ %import common.SIGNED_NUMBER''' - -# vllm 0.6.3 自动开启 chunked prefill 的原始块 -_ARG_OLD_BLOCK = """\ - if (is_gpu and not use_sliding_window and not use_spec_decode - and not self.enable_lora - and not self.enable_prompt_adapter): - self.enable_chunked_prefill = True - logger.warning( - "Chunked prefill is enabled by default for models with " - "max_model_len > 32K. Currently, chunked prefill might " - "not work with some features or models. If you " - "encounter any issues, please disable chunked prefill " - "by setting --enable-chunked-prefill=False.")\ -""" - + +# vllm 0.6.3 自动开启 chunked prefill 的原始块 +_ARG_OLD_BLOCK = """\ + if (is_gpu and not use_sliding_window and not use_spec_decode + and not self.enable_lora + and not self.enable_prompt_adapter): + self.enable_chunked_prefill = True + logger.warning( + "Chunked prefill is enabled by default for models with " + "max_model_len > 32K. Currently, chunked prefill might " + "not work with some features or models. If you " + "encounter any issues, please disable chunked prefill " + "by setting --enable-chunked-prefill=False.")\ +""" + _ARG_NEW_BLOCK = """\ - if (is_gpu and not use_sliding_window and not use_spec_decode - and not self.enable_lora - and not self.enable_prompt_adapter): - pass # skip auto-enable: Q-tiling in _run_sdpa_fallback + if (is_gpu and not use_sliding_window and not use_spec_decode + and not self.enable_lora + and not self.enable_prompt_adapter): + pass # skip auto-enable: Q-tiling in _run_sdpa_fallback # handles long-context memory without chunked prefill\ """ @@ -163,146 +163,146 @@ _MM_PREFIX_NEW_BLOCK = """\ "supported for multimodal models and has been disabled.") self.enable_prefix_caching = False\ """ - -FALLBACK_METHOD = ''' - def _run_sdpa_fallback( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: "XFormersMetadata", - ) -> torch.Tensor: - """Use ixformer flash_attn_varlen_func for head_dim > 128. - - Verified on real BI-V100: flash_attn_func handles head_dim=256 - correctly (diff < 0.004, no NaN). For seq >= 1024, faster than - PyTorch matmul. For profiling, sequences can be 20K+ tokens — this - is dramatically faster than the previous Python Q-tiling fallback. - - Falls back to pure-math if flash_attn is unavailable. - """ - import ixformer as _ixf - - assert attn_metadata.seq_lens is not None - orig_dtype = query.dtype - num_seqs = len(attn_metadata.seq_lens) - - q_flat = query.squeeze(0) # [T, H, D] - k_flat = key.squeeze(0) # [T, Hkv, D] - v_flat = value.squeeze(0) - - # Build cu_seqlens from seq_lens - seq_lens_list = list(attn_metadata.seq_lens) - cu_seqlens = torch.zeros(num_seqs + 1, dtype=torch.int32, - device=query.device) - for i, sl in enumerate(seq_lens_list): - cu_seqlens[i + 1] = cu_seqlens[i] + sl - max_seqlen = max(seq_lens_list) - - try: - # Skip flash_attn during profiling — OOMs on large dummy batch - import os - if os.environ.get("BI100_IN_STARTUP_PROFILE") == "1": - raise RuntimeError("skip flash_attn during profiling") - out = _ixf.flash_attn_varlen_func( - q_flat.to(torch.float16), - k_flat.to(torch.float16), - v_flat.to(torch.float16), - cu_seqlens, cu_seqlens, - max_seqlen, max_seqlen, - causal=True, - ) - return out.to(orig_dtype).unsqueeze(0) - except Exception: - pass - - # Fallback: pure-math Q-tiling (original implementation) - _Q_CHUNK = 256 - - # During profiling, skip expensive attention — return zeros. - # Profiling only measures memory footprint, not output correctness. - if os.environ.get("BI100_IN_STARTUP_PROFILE") == "1": - return torch.zeros_like(query) - - if (attn_metadata.query_start_loc is not None - and len(attn_metadata.query_start_loc) == num_seqs + 1): - q_lens = [ - int(attn_metadata.query_start_loc[i + 1].item()) - - int(attn_metadata.query_start_loc[i].item()) - for i in range(num_seqs) - ] - else: - q_lens = seq_lens_list - - output = torch.empty_like(q_flat) - seq_start = 0 - for q_len in q_lens: - seq_end = seq_start + q_len - k_s = k_flat[seq_start:seq_end].permute(1, 0, 2).float() - v_s = v_flat[seq_start:seq_end].permute(1, 0, 2).float() - if k_s.shape[0] != self.num_heads: - n = self.num_heads // k_s.shape[0] - k_s = k_s.repeat_interleave(n, dim=0).contiguous() - v_s = v_s.repeat_interleave(n, dim=0).contiguous() - k_pos = torch.arange(q_len, device=query.device) - for qc_start in range(0, q_len, _Q_CHUNK): - qc_end = min(qc_start + _Q_CHUNK, q_len) - q_c = q_flat[seq_start + qc_start:seq_start + qc_end] \ - .permute(1, 0, 2).float() - attn_w = torch.matmul(q_c, k_s.transpose(-2, -1)) * self.scale - qc_q_pos = torch.arange(qc_start, qc_end, device=query.device) - mask = k_pos.unsqueeze(0) > qc_q_pos.unsqueeze(1) - attn_w = attn_w.masked_fill(mask.unsqueeze(0), float("-inf")) - attn_w = torch.softmax(attn_w, dim=-1) - out_c = torch.matmul(attn_w, v_s).to(orig_dtype) - output[seq_start + qc_start:seq_start + qc_end] = ( - out_c.permute(1, 0, 2)) - seq_start = seq_end - return output.unsqueeze(0) - -''' - -OLD_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op = self.attn_op - ) - return out.view_as(original_query)\ -""" - -NEW_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - if self.head_size > 128: - out = self._run_sdpa_fallback(query, key, value, attn_metadata) - else: - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op=self.attn_op, - ) - return out.view_as(original_query)\ -""" - + +FALLBACK_METHOD = ''' + def _run_sdpa_fallback( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: "XFormersMetadata", + ) -> torch.Tensor: + """Use ixformer flash_attn_varlen_func for head_dim > 128. + + Verified on real BI-V100: flash_attn_func handles head_dim=256 + correctly (diff < 0.004, no NaN). For seq >= 1024, faster than + PyTorch matmul. For profiling, sequences can be 20K+ tokens — this + is dramatically faster than the previous Python Q-tiling fallback. + + Falls back to pure-math if flash_attn is unavailable. + """ + import ixformer as _ixf + + assert attn_metadata.seq_lens is not None + orig_dtype = query.dtype + num_seqs = len(attn_metadata.seq_lens) + + q_flat = query.squeeze(0) # [T, H, D] + k_flat = key.squeeze(0) # [T, Hkv, D] + v_flat = value.squeeze(0) + + # Build cu_seqlens from seq_lens + seq_lens_list = list(attn_metadata.seq_lens) + cu_seqlens = torch.zeros(num_seqs + 1, dtype=torch.int32, + device=query.device) + for i, sl in enumerate(seq_lens_list): + cu_seqlens[i + 1] = cu_seqlens[i] + sl + max_seqlen = max(seq_lens_list) + + try: + # Skip flash_attn during profiling — OOMs on large dummy batch + import os + if os.environ.get("BI100_IN_STARTUP_PROFILE") == "1": + raise RuntimeError("skip flash_attn during profiling") + out = _ixf.flash_attn_varlen_func( + q_flat.to(torch.float16), + k_flat.to(torch.float16), + v_flat.to(torch.float16), + cu_seqlens, cu_seqlens, + max_seqlen, max_seqlen, + causal=True, + ) + return out.to(orig_dtype).unsqueeze(0) + except Exception: + pass + + # Fallback: pure-math Q-tiling (original implementation) + _Q_CHUNK = 256 + + # During profiling, skip expensive attention — return zeros. + # Profiling only measures memory footprint, not output correctness. + if os.environ.get("BI100_IN_STARTUP_PROFILE") == "1": + return torch.zeros_like(query) + + if (attn_metadata.query_start_loc is not None + and len(attn_metadata.query_start_loc) == num_seqs + 1): + q_lens = [ + int(attn_metadata.query_start_loc[i + 1].item()) - + int(attn_metadata.query_start_loc[i].item()) + for i in range(num_seqs) + ] + else: + q_lens = seq_lens_list + + output = torch.empty_like(q_flat) + seq_start = 0 + for q_len in q_lens: + seq_end = seq_start + q_len + k_s = k_flat[seq_start:seq_end].permute(1, 0, 2).float() + v_s = v_flat[seq_start:seq_end].permute(1, 0, 2).float() + if k_s.shape[0] != self.num_heads: + n = self.num_heads // k_s.shape[0] + k_s = k_s.repeat_interleave(n, dim=0).contiguous() + v_s = v_s.repeat_interleave(n, dim=0).contiguous() + k_pos = torch.arange(q_len, device=query.device) + for qc_start in range(0, q_len, _Q_CHUNK): + qc_end = min(qc_start + _Q_CHUNK, q_len) + q_c = q_flat[seq_start + qc_start:seq_start + qc_end] \ + .permute(1, 0, 2).float() + attn_w = torch.matmul(q_c, k_s.transpose(-2, -1)) * self.scale + qc_q_pos = torch.arange(qc_start, qc_end, device=query.device) + mask = k_pos.unsqueeze(0) > qc_q_pos.unsqueeze(1) + attn_w = attn_w.masked_fill(mask.unsqueeze(0), float("-inf")) + attn_w = torch.softmax(attn_w, dim=-1) + out_c = torch.matmul(attn_w, v_s).to(orig_dtype) + output[seq_start + qc_start:seq_start + qc_end] = ( + out_c.permute(1, 0, 2)) + seq_start = seq_end + return output.unsqueeze(0) + +''' + +OLD_XFORMER_BLOCK = """\ + self.attn_op = xops.fmha.flash.FwOp() + if self.alibi_slopes is None: + # Add the batch dimension. + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + out = xops.memory_efficient_attention_forward( + query, + key, + value, + attn_bias=attn_bias[0], + p=0.0, + scale=self.scale, + op = self.attn_op + ) + return out.view_as(original_query)\ +""" + +NEW_XFORMER_BLOCK = """\ + self.attn_op = xops.fmha.flash.FwOp() + if self.alibi_slopes is None: + # Add the batch dimension. + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + if self.head_size > 128: + out = self._run_sdpa_fallback(query, key, value, attn_metadata) + else: + out = xops.memory_efficient_attention_forward( + query, + key, + value, + attn_bias=attn_bias[0], + p=0.0, + scale=self.scale, + op=self.attn_op, + ) + return out.view_as(original_query)\ +""" + INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward(" _PREFIX_CALL_OLD_BLOCK = """\ @@ -367,8 +367,8 @@ def patch_file(path): required=True, already_contains=( "is_causal_decoder=(attn_type == AttentionType.DECODER)")) - - + + def patch_arg_utils(path): replace_once( path, @@ -382,8 +382,8 @@ def patch_arg_utils(path): _MM_PREFIX_NEW_BLOCK, required=True, already_contains="Keeping prefix caching enabled for the Qwen3.6") - - + + def patch_logits_processor(path): replace_once( path, @@ -402,27 +402,27 @@ def patch_outlines_json_grammar(path): ], required=True, already_contains="JSON_WS:") - - -def main(): - print("=== patch_xformers_sdpa_seq (sequential, pure-math) ===") - print(f"Target: {XFORMERS_PATH}") - patch_file(XFORMERS_PATH) - - print("\n=== patch_arg_utils (disable chunked-prefill auto-enable) ===") - print(f"Target: {ARG_UTILS_PATH}") - patch_arg_utils(ARG_UTILS_PATH) - - print("\n=== patch_logits_processor (seq_groups=None guard for chunked prefill) ===") + + +def main(): + print("=== patch_xformers_sdpa_seq (sequential, pure-math) ===") + print(f"Target: {XFORMERS_PATH}") + patch_file(XFORMERS_PATH) + + print("\n=== patch_arg_utils (disable chunked-prefill auto-enable) ===") + print(f"Target: {ARG_UTILS_PATH}") + patch_arg_utils(ARG_UTILS_PATH) + + print("\n=== patch_logits_processor (seq_groups=None guard for chunked prefill) ===") print(f"Target: {LOGITS_PROC_PATH}") patch_logits_processor(LOGITS_PROC_PATH) print("\n=== patch_outlines_json_grammar (reject raw control chars) ===") print(f"Target: {OUTLINES_DECODING_PATH}") patch_outlines_json_grammar(OUTLINES_DECODING_PATH) - - print("\nDone.") - - -if __name__ == "__main__": - main() + + print("\nDone.") + + +if __name__ == "__main__": + main() diff --git a/qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py b/qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py index 39633d6d..2ef68660 100644 --- a/qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py +++ b/qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py @@ -1,22 +1,22 @@ -""" -策略：顺序（per-sequence）— F.scaled_dot_product_attention，可走硬件 kernel -============================================================================= -逐条序列调用 F.scaled_dot_product_attention，is_causal=False + 显式因果 mask。 -与 patch_xformers_sdpa_seq.py（纯 matmul）的区别： - SDPA 可分发到 Flash Attention / mem-efficient attention kernel， - 而纯 matmul 固定走 cublas。 - -硬件限制（BI-V100）： - cudnnFlashAttnForward 不支持 is_causal=True（直接报错）。 - 必须使用 is_causal=False + 显式 additive causal mask。 - 每条序列单独构造上三角 -inf mask，peak 显存 = max(seq_len)² × dtype， - 比 batch 版的 total_tokens² 小得多。 - -与 batch_kernel 的对比： - seq_kernel: 显存小，peak = max_single_seq²；并发 prefill 串行排队 - batch_kernel: 显存大，peak = total_tokens²；并发 prefill 一次并行处理， - 通过 --max-num-batched-tokens 控制 total_tokens 上限 - +""" +策略：顺序（per-sequence）— F.scaled_dot_product_attention，可走硬件 kernel +============================================================================= +逐条序列调用 F.scaled_dot_product_attention，is_causal=False + 显式因果 mask。 +与 patch_xformers_sdpa_seq.py（纯 matmul）的区别： + SDPA 可分发到 Flash Attention / mem-efficient attention kernel， + 而纯 matmul 固定走 cublas。 + +硬件限制（BI-V100）： + cudnnFlashAttnForward 不支持 is_causal=True（直接报错）。 + 必须使用 is_causal=False + 显式 additive causal mask。 + 每条序列单独构造上三角 -inf mask，peak 显存 = max(seq_len)² × dtype， + 比 batch 版的 total_tokens² 小得多。 + +与 batch_kernel 的对比： + seq_kernel: 显存小，peak = max_single_seq²；并发 prefill 串行排队 + batch_kernel: 显存大，peak = total_tokens²；并发 prefill 一次并行处理， + 通过 --max-num-batched-tokens 控制 total_tokens 上限 + Deploy: python3 modified_scripts/patch_xformers_sdpa_seq_kernel.py """ @@ -24,122 +24,122 @@ Deploy: from patch_utils import package_root, replace_once XFORMERS_PATH = package_root("vllm") / "attention" / "backends" / "xformers.py" - -FALLBACK_METHOD = ''' - def _run_sdpa_fallback( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: "XFormersMetadata", - ) -> torch.Tensor: - """顺序 F.scaled_dot_product_attention fallback（可走硬件 kernel）。 - - 逐条序列调用 SDPA，is_causal=False + 显式上三角 additive mask。 - cudnnFlashAttnForward 不支持 is_causal=True，必须用显式 mask。 - 逐序列构造 mask，peak 显存 = max(seq_len)² × dtype（远小于 batch 版）。 - - Args: - query : [1, total_prefill_tokens, num_heads, head_dim] - key : [1, total_prefill_tokens, num_kv_heads, head_dim] - value : [1, total_prefill_tokens, num_kv_heads, head_dim] - Returns: - [1, total_prefill_tokens, num_heads, head_dim] - """ - import torch.nn.functional as F - - assert attn_metadata.seq_lens is not None - orig_dtype = query.dtype - - q_flat = query.squeeze(0) # [T, H, D] - k_flat = key.squeeze(0) # [T, Hkv, D] - v_flat = value.squeeze(0) - - output = torch.empty_like(q_flat) - start = 0 - for seq_len in attn_metadata.seq_lens: - end = start + seq_len - # [1, H, L, D] - q_s = q_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0) - k_s = k_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0) - v_s = v_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0) - - # GQA：展开 KV heads - if k_s.shape[1] != q_s.shape[1]: - n = q_s.shape[1] // k_s.shape[1] - k_s = k_s.repeat_interleave(n, dim=1).contiguous() - v_s = v_s.repeat_interleave(n, dim=1).contiguous() - - # 逐序列因果 mask [L, L]，上三角 -inf - causal_mask = torch.tril( - torch.zeros(seq_len, seq_len, dtype=orig_dtype, device=q_s.device) - ) - causal_mask = causal_mask.masked_fill( - torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, - device=q_s.device), diagonal=1), - float("-inf"), - ) - - # is_causal=False + 显式 mask，规避 cudnnFlashAttnForward 不支持 is_causal=True - out_s = F.scaled_dot_product_attention( - q_s, k_s, v_s, - attn_mask=causal_mask, - dropout_p=0.0, - is_causal=False, - scale=self.scale, - ) - # [1, H, L, D] → [L, H, D] - output[start:end] = out_s.squeeze(0).permute(1, 0, 2).to(orig_dtype) - start = end - - return output.unsqueeze(0) # [1, T, H, D] - -''' - -OLD_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op = self.attn_op - ) - return out.view_as(original_query)\ -""" - -NEW_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - if self.head_size > 128: - out = self._run_sdpa_fallback(query, key, value, attn_metadata) - else: - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op=self.attn_op, - ) - return out.view_as(original_query)\ -""" - -INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward(" - - + +FALLBACK_METHOD = ''' + def _run_sdpa_fallback( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: "XFormersMetadata", + ) -> torch.Tensor: + """顺序 F.scaled_dot_product_attention fallback（可走硬件 kernel）。 + + 逐条序列调用 SDPA，is_causal=False + 显式上三角 additive mask。 + cudnnFlashAttnForward 不支持 is_causal=True，必须用显式 mask。 + 逐序列构造 mask，peak 显存 = max(seq_len)² × dtype（远小于 batch 版）。 + + Args: + query : [1, total_prefill_tokens, num_heads, head_dim] + key : [1, total_prefill_tokens, num_kv_heads, head_dim] + value : [1, total_prefill_tokens, num_kv_heads, head_dim] + Returns: + [1, total_prefill_tokens, num_heads, head_dim] + """ + import torch.nn.functional as F + + assert attn_metadata.seq_lens is not None + orig_dtype = query.dtype + + q_flat = query.squeeze(0) # [T, H, D] + k_flat = key.squeeze(0) # [T, Hkv, D] + v_flat = value.squeeze(0) + + output = torch.empty_like(q_flat) + start = 0 + for seq_len in attn_metadata.seq_lens: + end = start + seq_len + # [1, H, L, D] + q_s = q_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0) + k_s = k_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0) + v_s = v_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0) + + # GQA：展开 KV heads + if k_s.shape[1] != q_s.shape[1]: + n = q_s.shape[1] // k_s.shape[1] + k_s = k_s.repeat_interleave(n, dim=1).contiguous() + v_s = v_s.repeat_interleave(n, dim=1).contiguous() + + # 逐序列因果 mask [L, L]，上三角 -inf + causal_mask = torch.tril( + torch.zeros(seq_len, seq_len, dtype=orig_dtype, device=q_s.device) + ) + causal_mask = causal_mask.masked_fill( + torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, + device=q_s.device), diagonal=1), + float("-inf"), + ) + + # is_causal=False + 显式 mask，规避 cudnnFlashAttnForward 不支持 is_causal=True + out_s = F.scaled_dot_product_attention( + q_s, k_s, v_s, + attn_mask=causal_mask, + dropout_p=0.0, + is_causal=False, + scale=self.scale, + ) + # [1, H, L, D] → [L, H, D] + output[start:end] = out_s.squeeze(0).permute(1, 0, 2).to(orig_dtype) + start = end + + return output.unsqueeze(0) # [1, T, H, D] + +''' + +OLD_XFORMER_BLOCK = """\ + self.attn_op = xops.fmha.flash.FwOp() + if self.alibi_slopes is None: + # Add the batch dimension. + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + out = xops.memory_efficient_attention_forward( + query, + key, + value, + attn_bias=attn_bias[0], + p=0.0, + scale=self.scale, + op = self.attn_op + ) + return out.view_as(original_query)\ +""" + +NEW_XFORMER_BLOCK = """\ + self.attn_op = xops.fmha.flash.FwOp() + if self.alibi_slopes is None: + # Add the batch dimension. + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + if self.head_size > 128: + out = self._run_sdpa_fallback(query, key, value, attn_metadata) + else: + out = xops.memory_efficient_attention_forward( + query, + key, + value, + attn_bias=attn_bias[0], + p=0.0, + scale=self.scale, + op=self.attn_op, + ) + return out.view_as(original_query)\ +""" + +INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward(" + + def patch_file(path): replace_once( path, @@ -153,14 +153,14 @@ def patch_file(path): NEW_XFORMER_BLOCK, required=True, already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)") - - -def main(): - print("=== patch_xformers_sdpa_seq_kernel (seq, F.sdpa + kernel dispatch) ===") - print(f"Target: {XFORMERS_PATH}") - patch_file(XFORMERS_PATH) - print("\nDone.") - - -if __name__ == "__main__": - main() + + +def main(): + print("=== patch_xformers_sdpa_seq_kernel (seq, F.sdpa + kernel dispatch) ===") + print(f"Target: {XFORMERS_PATH}") + patch_file(XFORMERS_PATH) + print("\nDone.") + + +if __name__ == "__main__": + main() diff --git a/qwen3_6_scripts/protocol.py b/qwen3_6_scripts/protocol.py index 638925a6..646be6ee 100644 --- a/qwen3_6_scripts/protocol.py +++ b/qwen3_6_scripts/protocol.py @@ -2,141 +2,141 @@ # https://github.com/lm-sys/FastChat/blob/168ccc29d3f7edc50823016105c024fe2282732a/fastchat/protocol/openai_api_protocol.py import json import time -from argparse import Namespace -from typing import Any, Dict, List, Literal, Optional, Union - -import torch -from openai.types.chat import ChatCompletionContentPartParam -from pydantic import BaseModel, ConfigDict, Field, model_validator -from typing_extensions import Annotated, Required, TypedDict - -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam -from vllm.pooling_params import PoolingParams -from vllm.sampling_params import (BeamSearchParams, GuidedDecodingParams, - RequestOutputKind, SamplingParams) -from vllm.sequence import Logprob -from vllm.utils import random_uuid - -# torch is mocked during docs generation, -# so we have to provide the values as literals -_MOCK_LONG_INFO = Namespace(min=-9223372036854775808, max=9223372036854775807) -_LONG_INFO: Union["torch.iinfo", Namespace] - -try: - from sphinx.ext.autodoc.mock import _MockModule - - if isinstance(torch, _MockModule): - _LONG_INFO = _MOCK_LONG_INFO - else: - _LONG_INFO = torch.iinfo(torch.long) -except ModuleNotFoundError: - _LONG_INFO = torch.iinfo(torch.long) - -assert _LONG_INFO.min == _MOCK_LONG_INFO.min -assert _LONG_INFO.max == _MOCK_LONG_INFO.max - - -class CustomChatCompletionMessageParam(TypedDict, total=False): - """Enables custom roles in the Chat Completion API.""" - role: Required[str] - """The role of the message's author.""" - - content: Union[str, List[ChatCompletionContentPartParam]] - """The contents of the message.""" - - name: str - """An optional name for the participant. - - Provides the model information to differentiate between participants of the - same role. - """ - - tool_call_id: Optional[str] - - tool_calls: Optional[List[dict]] - - -class OpenAIBaseModel(BaseModel): - # OpenAI API does not allow extra fields - model_config = ConfigDict(extra="allow") - - -class ErrorResponse(OpenAIBaseModel): - object: str = "error" - message: str - type: str - param: Optional[str] = None - code: int - - -class ModelPermission(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"modelperm-{random_uuid()}") - object: str = "model_permission" - created: int = Field(default_factory=lambda: int(time.time())) - allow_create_engine: bool = False - allow_sampling: bool = True - allow_logprobs: bool = True - allow_search_indices: bool = False - allow_view: bool = True - allow_fine_tuning: bool = False - organization: str = "*" - group: Optional[str] = None - is_blocking: bool = False - - -class ModelCard(OpenAIBaseModel): - id: str - object: str = "model" - created: int = Field(default_factory=lambda: int(time.time())) - owned_by: str = "vllm" - root: Optional[str] = None - parent: Optional[str] = None - max_model_len: Optional[int] = None - permission: List[ModelPermission] = Field(default_factory=list) - - -class ModelList(OpenAIBaseModel): - object: str = "list" - data: List[ModelCard] = Field(default_factory=list) - - -class PromptTokensDetails(OpenAIBaseModel): - cached_tokens: int = 0 - - -class UsageInfo(OpenAIBaseModel): - prompt_tokens: int = 0 - total_tokens: int = 0 - completion_tokens: Optional[int] = 0 - reasoning_tokens: Optional[int] = None - prompt_tokens_details: Optional[PromptTokensDetails] = None - - -class RequestResponseMetadata(BaseModel): - request_id: str - final_usage_info: Optional[UsageInfo] = None - - -class JsonSchemaResponseFormat(OpenAIBaseModel): - name: str - description: Optional[str] = None - # schema is the field in openai but that causes conflicts with pydantic so - # instead use json_schema with an alias - json_schema: Optional[Dict[str, Any]] = Field(default=None, alias='schema') - strict: Optional[bool] = None - - -class ResponseFormat(OpenAIBaseModel): - # type must be "json_schema", "json_object" or "text" - type: Literal["text", "json_object", "json_schema"] - json_schema: Optional[JsonSchemaResponseFormat] = None - - -class StreamOptions(OpenAIBaseModel): - include_usage: Optional[bool] = True - continuous_usage_stats: Optional[bool] = True - - +from argparse import Namespace +from typing import Any, Dict, List, Literal, Optional, Union + +import torch +from openai.types.chat import ChatCompletionContentPartParam +from pydantic import BaseModel, ConfigDict, Field, model_validator +from typing_extensions import Annotated, Required, TypedDict + +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam +from vllm.pooling_params import PoolingParams +from vllm.sampling_params import (BeamSearchParams, GuidedDecodingParams, + RequestOutputKind, SamplingParams) +from vllm.sequence import Logprob +from vllm.utils import random_uuid + +# torch is mocked during docs generation, +# so we have to provide the values as literals +_MOCK_LONG_INFO = Namespace(min=-9223372036854775808, max=9223372036854775807) +_LONG_INFO: Union["torch.iinfo", Namespace] + +try: + from sphinx.ext.autodoc.mock import _MockModule + + if isinstance(torch, _MockModule): + _LONG_INFO = _MOCK_LONG_INFO + else: + _LONG_INFO = torch.iinfo(torch.long) +except ModuleNotFoundError: + _LONG_INFO = torch.iinfo(torch.long) + +assert _LONG_INFO.min == _MOCK_LONG_INFO.min +assert _LONG_INFO.max == _MOCK_LONG_INFO.max + + +class CustomChatCompletionMessageParam(TypedDict, total=False): + """Enables custom roles in the Chat Completion API.""" + role: Required[str] + """The role of the message's author.""" + + content: Union[str, List[ChatCompletionContentPartParam]] + """The contents of the message.""" + + name: str + """An optional name for the participant. + + Provides the model information to differentiate between participants of the + same role. + """ + + tool_call_id: Optional[str] + + tool_calls: Optional[List[dict]] + + +class OpenAIBaseModel(BaseModel): + # OpenAI API does not allow extra fields + model_config = ConfigDict(extra="allow") + + +class ErrorResponse(OpenAIBaseModel): + object: str = "error" + message: str + type: str + param: Optional[str] = None + code: int + + +class ModelPermission(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"modelperm-{random_uuid()}") + object: str = "model_permission" + created: int = Field(default_factory=lambda: int(time.time())) + allow_create_engine: bool = False + allow_sampling: bool = True + allow_logprobs: bool = True + allow_search_indices: bool = False + allow_view: bool = True + allow_fine_tuning: bool = False + organization: str = "*" + group: Optional[str] = None + is_blocking: bool = False + + +class ModelCard(OpenAIBaseModel): + id: str + object: str = "model" + created: int = Field(default_factory=lambda: int(time.time())) + owned_by: str = "vllm" + root: Optional[str] = None + parent: Optional[str] = None + max_model_len: Optional[int] = None + permission: List[ModelPermission] = Field(default_factory=list) + + +class ModelList(OpenAIBaseModel): + object: str = "list" + data: List[ModelCard] = Field(default_factory=list) + + +class PromptTokensDetails(OpenAIBaseModel): + cached_tokens: int = 0 + + +class UsageInfo(OpenAIBaseModel): + prompt_tokens: int = 0 + total_tokens: int = 0 + completion_tokens: Optional[int] = 0 + reasoning_tokens: Optional[int] = None + prompt_tokens_details: Optional[PromptTokensDetails] = None + + +class RequestResponseMetadata(BaseModel): + request_id: str + final_usage_info: Optional[UsageInfo] = None + + +class JsonSchemaResponseFormat(OpenAIBaseModel): + name: str + description: Optional[str] = None + # schema is the field in openai but that causes conflicts with pydantic so + # instead use json_schema with an alias + json_schema: Optional[Dict[str, Any]] = Field(default=None, alias='schema') + strict: Optional[bool] = None + + +class ResponseFormat(OpenAIBaseModel): + # type must be "json_schema", "json_object" or "text" + type: Literal["text", "json_object", "json_schema"] + json_schema: Optional[JsonSchemaResponseFormat] = None + + +class StreamOptions(OpenAIBaseModel): + include_usage: Optional[bool] = True + continuous_usage_stats: Optional[bool] = True + + class FunctionDefinition(OpenAIBaseModel): name: str description: Optional[str] = None @@ -154,43 +154,43 @@ class FunctionDefinition(OpenAIBaseModel): "Function tools with strict=true are not supported by this " "runtime.") return self - - -class ChatCompletionToolsParam(OpenAIBaseModel): - type: Literal["function"] = "function" - function: FunctionDefinition - - -class ChatCompletionNamedFunction(OpenAIBaseModel): - name: str - - -class ChatCompletionNamedToolChoiceParam(OpenAIBaseModel): - function: ChatCompletionNamedFunction - type: Literal["function"] = "function" - - -class ChatCompletionRequest(OpenAIBaseModel): - # Ordered by official OpenAI API documentation - # https://platform.openai.com/docs/api-reference/chat/create - messages: List[ChatCompletionMessageParam] - model: str - frequency_penalty: Optional[float] = 0.0 - logit_bias: Optional[Dict[str, float]] = None - logprobs: Optional[bool] = False - top_logprobs: Optional[int] = 0 - max_tokens: Optional[int] = None - # OpenAI newer API field — treat as alias for max_tokens - max_completion_tokens: Optional[int] = None - n: Optional[int] = 1 - presence_penalty: Optional[float] = 0.0 - response_format: Optional[ResponseFormat] = None - seed: Optional[int] = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max) - stop: Optional[Union[str, List[str]]] = Field(default_factory=list) - stream: Optional[bool] = False - stream_options: Optional[StreamOptions] = None - temperature: Optional[float] = 0.7 - top_p: Optional[float] = 1.0 + + +class ChatCompletionToolsParam(OpenAIBaseModel): + type: Literal["function"] = "function" + function: FunctionDefinition + + +class ChatCompletionNamedFunction(OpenAIBaseModel): + name: str + + +class ChatCompletionNamedToolChoiceParam(OpenAIBaseModel): + function: ChatCompletionNamedFunction + type: Literal["function"] = "function" + + +class ChatCompletionRequest(OpenAIBaseModel): + # Ordered by official OpenAI API documentation + # https://platform.openai.com/docs/api-reference/chat/create + messages: List[ChatCompletionMessageParam] + model: str + frequency_penalty: Optional[float] = 0.0 + logit_bias: Optional[Dict[str, float]] = None + logprobs: Optional[bool] = False + top_logprobs: Optional[int] = 0 + max_tokens: Optional[int] = None + # OpenAI newer API field — treat as alias for max_tokens + max_completion_tokens: Optional[int] = None + n: Optional[int] = 1 + presence_penalty: Optional[float] = 0.0 + response_format: Optional[ResponseFormat] = None + seed: Optional[int] = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max) + stop: Optional[Union[str, List[str]]] = Field(default_factory=list) + stream: Optional[bool] = False + stream_options: Optional[StreamOptions] = None + temperature: Optional[float] = 0.7 + top_p: Optional[float] = 1.0 tools: Optional[List[ChatCompletionToolsParam]] = None tool_choice: Optional[Union[Literal["none"], Literal["auto"], ChatCompletionNamedToolChoiceParam]] = "none" @@ -199,146 +199,146 @@ class ChatCompletionRequest(OpenAIBaseModel): # NOTE this will be ignored by VLLM -- the model determines the behavior parallel_tool_calls: Optional[bool] = False - user: Optional[str] = None - - # doc: begin-chat-completion-sampling-params - best_of: Optional[int] = None - use_beam_search: bool = False - top_k: int = -1 - min_p: float = 0.0 - repetition_penalty: float = 1.0 - length_penalty: float = 1.0 - stop_token_ids: Optional[List[int]] = Field(default_factory=list) - include_stop_str_in_output: bool = False - ignore_eos: bool = False - min_tokens: int = 0 - skip_special_tokens: bool = True - spaces_between_special_tokens: bool = True + user: Optional[str] = None + + # doc: begin-chat-completion-sampling-params + best_of: Optional[int] = None + use_beam_search: bool = False + top_k: int = -1 + min_p: float = 0.0 + repetition_penalty: float = 1.0 + length_penalty: float = 1.0 + stop_token_ids: Optional[List[int]] = Field(default_factory=list) + include_stop_str_in_output: bool = False + ignore_eos: bool = False + min_tokens: int = 0 + skip_special_tokens: bool = True + spaces_between_special_tokens: bool = True truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None prompt_logprobs: Optional[int] = None bi100_prompt_logprobs_sample_positions: Optional[List[int]] = None - # doc: end-chat-completion-sampling-params - - # doc: begin-chat-completion-extra-params - echo: bool = Field( - default=False, - description=( - "If true, the new message will be prepended with the last message " - "if they belong to the same role."), - ) - add_generation_prompt: bool = Field( - default=True, - description= - ("If true, the generation prompt will be added to the chat template. " - "This is a parameter used by chat template in tokenizer config of the " - "model."), - ) - continue_final_message: bool = Field( - default=False, - description= - ("If this is set, the chat will be formatted so that the final " - "message in the chat is open-ended, without any EOS tokens. The " - "model will continue this message rather than starting a new one. " - "This allows you to \"prefill\" part of the model's response for it. " - "Cannot be used at the same time as `add_generation_prompt`."), - ) - add_special_tokens: bool = Field( - default=False, - description=( - "If true, special tokens (e.g. BOS) will be added to the prompt " - "on top of what is added by the chat template. " - "For most models, the chat template takes care of adding the " - "special tokens so this should be set to false (as is the " - "default)."), - ) - documents: Optional[List[Dict[str, str]]] = Field( - default=None, - description= - ("A list of dicts representing documents that will be accessible to " - "the model if it is performing RAG (retrieval-augmented generation)." - " If the template does not support RAG, this argument will have no " - "effect. We recommend that each document should be a dict containing " - "\"title\" and \"text\" keys."), - ) - chat_template: Optional[str] = Field( - default=None, - description=( - "A Jinja template to use for this conversion. " - "As of transformers v4.44, default chat template is no longer " - "allowed, so you must provide a chat template if the tokenizer " - "does not define one."), - ) - chat_template_kwargs: Optional[Dict[str, Any]] = Field( - default=None, - description=("Additional kwargs to pass to the template renderer. " - "Will be accessible by the chat template."), - ) - guided_json: Optional[Union[str, dict, BaseModel]] = Field( - default=None, - description=("If specified, the output will follow the JSON schema."), - ) - guided_regex: Optional[str] = Field( - default=None, - description=( - "If specified, the output will follow the regex pattern."), - ) - guided_choice: Optional[List[str]] = Field( - default=None, - description=( - "If specified, the output will be exactly one of the choices."), - ) - guided_grammar: Optional[str] = Field( - default=None, - description=( - "If specified, the output will follow the context free grammar."), - ) - guided_decoding_backend: Optional[str] = Field( - default=None, - description=( - "If specified, will override the default guided decoding backend " - "of the server for this specific request. If set, must be either " - "'outlines' / 'lm-format-enforcer'")) - guided_whitespace_pattern: Optional[str] = Field( - default=None, - description=( - "If specified, will override the default whitespace pattern " - "for guided json decoding.")) - priority: int = Field( - default=0, - description=( - "The priority of the request (lower means earlier handling; " - "default: 0). Any priority other than 0 will raise an error " - "if the served model does not use priority scheduling.")) - - # doc: end-chat-completion-extra-params - - def to_beam_search_params(self, - default_max_tokens: int) -> BeamSearchParams: - max_tokens = self.max_tokens - if max_tokens is None: - max_tokens = default_max_tokens - - n = self.n if self.n is not None else 1 - temperature = self.temperature if self.temperature is not None else 0.0 - - return BeamSearchParams( - beam_width=n, - max_tokens=max_tokens, - ignore_eos=self.ignore_eos, - temperature=temperature, - length_penalty=self.length_penalty, - ) - - def to_sampling_params(self, default_max_tokens: int) -> SamplingParams: - max_tokens = self.max_tokens - if max_tokens is None: - max_tokens = default_max_tokens - - prompt_logprobs = self.prompt_logprobs - if prompt_logprobs is None and self.echo: - prompt_logprobs = self.top_logprobs - - guided_json_object = None + # doc: end-chat-completion-sampling-params + + # doc: begin-chat-completion-extra-params + echo: bool = Field( + default=False, + description=( + "If true, the new message will be prepended with the last message " + "if they belong to the same role."), + ) + add_generation_prompt: bool = Field( + default=True, + description= + ("If true, the generation prompt will be added to the chat template. " + "This is a parameter used by chat template in tokenizer config of the " + "model."), + ) + continue_final_message: bool = Field( + default=False, + description= + ("If this is set, the chat will be formatted so that the final " + "message in the chat is open-ended, without any EOS tokens. The " + "model will continue this message rather than starting a new one. " + "This allows you to \"prefill\" part of the model's response for it. " + "Cannot be used at the same time as `add_generation_prompt`."), + ) + add_special_tokens: bool = Field( + default=False, + description=( + "If true, special tokens (e.g. BOS) will be added to the prompt " + "on top of what is added by the chat template. " + "For most models, the chat template takes care of adding the " + "special tokens so this should be set to false (as is the " + "default)."), + ) + documents: Optional[List[Dict[str, str]]] = Field( + default=None, + description= + ("A list of dicts representing documents that will be accessible to " + "the model if it is performing RAG (retrieval-augmented generation)." + " If the template does not support RAG, this argument will have no " + "effect. We recommend that each document should be a dict containing " + "\"title\" and \"text\" keys."), + ) + chat_template: Optional[str] = Field( + default=None, + description=( + "A Jinja template to use for this conversion. " + "As of transformers v4.44, default chat template is no longer " + "allowed, so you must provide a chat template if the tokenizer " + "does not define one."), + ) + chat_template_kwargs: Optional[Dict[str, Any]] = Field( + default=None, + description=("Additional kwargs to pass to the template renderer. " + "Will be accessible by the chat template."), + ) + guided_json: Optional[Union[str, dict, BaseModel]] = Field( + default=None, + description=("If specified, the output will follow the JSON schema."), + ) + guided_regex: Optional[str] = Field( + default=None, + description=( + "If specified, the output will follow the regex pattern."), + ) + guided_choice: Optional[List[str]] = Field( + default=None, + description=( + "If specified, the output will be exactly one of the choices."), + ) + guided_grammar: Optional[str] = Field( + default=None, + description=( + "If specified, the output will follow the context free grammar."), + ) + guided_decoding_backend: Optional[str] = Field( + default=None, + description=( + "If specified, will override the default guided decoding backend " + "of the server for this specific request. If set, must be either " + "'outlines' / 'lm-format-enforcer'")) + guided_whitespace_pattern: Optional[str] = Field( + default=None, + description=( + "If specified, will override the default whitespace pattern " + "for guided json decoding.")) + priority: int = Field( + default=0, + description=( + "The priority of the request (lower means earlier handling; " + "default: 0). Any priority other than 0 will raise an error " + "if the served model does not use priority scheduling.")) + + # doc: end-chat-completion-extra-params + + def to_beam_search_params(self, + default_max_tokens: int) -> BeamSearchParams: + max_tokens = self.max_tokens + if max_tokens is None: + max_tokens = default_max_tokens + + n = self.n if self.n is not None else 1 + temperature = self.temperature if self.temperature is not None else 0.0 + + return BeamSearchParams( + beam_width=n, + max_tokens=max_tokens, + ignore_eos=self.ignore_eos, + temperature=temperature, + length_penalty=self.length_penalty, + ) + + def to_sampling_params(self, default_max_tokens: int) -> SamplingParams: + max_tokens = self.max_tokens + if max_tokens is None: + max_tokens = default_max_tokens + + prompt_logprobs = self.prompt_logprobs + if prompt_logprobs is None and self.echo: + prompt_logprobs = self.top_logprobs + + guided_json_object = None guided_json_from_schema = None if self.response_format is not None: if self.response_format.type == "json_object": @@ -346,94 +346,94 @@ class ChatCompletionRequest(OpenAIBaseModel): # this vLLM/Outlines build. A generic object schema has the # same API semantics and uses the stable regex backend. guided_json_from_schema = {"type": "object"} - elif (self.response_format.type == "json_schema" - and self.response_format.json_schema is not None - and self.response_format.json_schema.json_schema is not None): - guided_json_from_schema = \ - self.response_format.json_schema.json_schema - - guided_decoding = GuidedDecodingParams.from_optional( - json=(self._get_guided_json_from_tool() - or self.guided_json - or guided_json_from_schema), - regex=self.guided_regex, - choice=self.guided_choice, - grammar=self.guided_grammar, - json_object=guided_json_object, - backend=self.guided_decoding_backend, - whitespace_pattern=self.guided_whitespace_pattern) - - return SamplingParams.from_optional( - n=self.n, - best_of=self.best_of, - presence_penalty=self.presence_penalty, - frequency_penalty=self.frequency_penalty, - repetition_penalty=self.repetition_penalty, - temperature=self.temperature, - top_p=self.top_p, - top_k=self.top_k, - min_p=self.min_p, - seed=self.seed, - stop=self.stop, - stop_token_ids=self.stop_token_ids, + elif (self.response_format.type == "json_schema" + and self.response_format.json_schema is not None + and self.response_format.json_schema.json_schema is not None): + guided_json_from_schema = \ + self.response_format.json_schema.json_schema + + guided_decoding = GuidedDecodingParams.from_optional( + json=(self._get_guided_json_from_tool() + or self.guided_json + or guided_json_from_schema), + regex=self.guided_regex, + choice=self.guided_choice, + grammar=self.guided_grammar, + json_object=guided_json_object, + backend=self.guided_decoding_backend, + whitespace_pattern=self.guided_whitespace_pattern) + + return SamplingParams.from_optional( + n=self.n, + best_of=self.best_of, + presence_penalty=self.presence_penalty, + frequency_penalty=self.frequency_penalty, + repetition_penalty=self.repetition_penalty, + temperature=self.temperature, + top_p=self.top_p, + top_k=self.top_k, + min_p=self.min_p, + seed=self.seed, + stop=self.stop, + stop_token_ids=self.stop_token_ids, logprobs=self.top_logprobs if self.logprobs else None, prompt_logprobs=prompt_logprobs, prompt_logprob_positions=( self.bi100_prompt_logprobs_sample_positions), - ignore_eos=self.ignore_eos, - max_tokens=max_tokens, - min_tokens=self.min_tokens, - skip_special_tokens=self.skip_special_tokens, - spaces_between_special_tokens=self.spaces_between_special_tokens, - include_stop_str_in_output=self.include_stop_str_in_output, - truncate_prompt_tokens=self.truncate_prompt_tokens, - output_kind=RequestOutputKind.DELTA if self.stream \ - else RequestOutputKind.FINAL_ONLY, - guided_decoding=guided_decoding, - logit_bias=self.logit_bias) - - def _get_guided_json_from_tool( - self) -> Optional[Union[str, dict, BaseModel]]: - # user has chosen to not use any tool - if self.tool_choice == "none" or self.tools is None: - return None - - # user has chosen to use a named tool - if type(self.tool_choice) is ChatCompletionNamedToolChoiceParam: - tool_name = self.tool_choice.function.name - tools = {tool.function.name: tool.function for tool in self.tools} - if tool_name not in tools: - raise ValueError( - f"Tool '{tool_name}' has not been passed in `tools`.") - tool = tools[tool_name] - return tool.parameters - - return None - - @model_validator(mode="before") - @classmethod - def fold_max_completion_tokens(cls, data): - """OpenAI newer API: max_completion_tokens → max_tokens alias.""" - if isinstance(data, dict): - mct = data.pop("max_completion_tokens", None) - if mct is not None and data.get("max_tokens") is None: - data["max_tokens"] = mct - return data - - @model_validator(mode="before") - @classmethod - def normalize_messages(cls, data): - """Normalize incoming messages before pydantic union validation. - - Real-world clients (e.g. from other providers) send assistant tool_call - messages with content=null, which fails the strict Union type check. - Replace null content with "" so validation passes. - reasoning_content is intentionally kept — chat_utils.py wraps it as - ... for multi-turn reasoning history. - """ - messages = data.get("messages") - if not isinstance(messages, list): - return data + ignore_eos=self.ignore_eos, + max_tokens=max_tokens, + min_tokens=self.min_tokens, + skip_special_tokens=self.skip_special_tokens, + spaces_between_special_tokens=self.spaces_between_special_tokens, + include_stop_str_in_output=self.include_stop_str_in_output, + truncate_prompt_tokens=self.truncate_prompt_tokens, + output_kind=RequestOutputKind.DELTA if self.stream \ + else RequestOutputKind.FINAL_ONLY, + guided_decoding=guided_decoding, + logit_bias=self.logit_bias) + + def _get_guided_json_from_tool( + self) -> Optional[Union[str, dict, BaseModel]]: + # user has chosen to not use any tool + if self.tool_choice == "none" or self.tools is None: + return None + + # user has chosen to use a named tool + if type(self.tool_choice) is ChatCompletionNamedToolChoiceParam: + tool_name = self.tool_choice.function.name + tools = {tool.function.name: tool.function for tool in self.tools} + if tool_name not in tools: + raise ValueError( + f"Tool '{tool_name}' has not been passed in `tools`.") + tool = tools[tool_name] + return tool.parameters + + return None + + @model_validator(mode="before") + @classmethod + def fold_max_completion_tokens(cls, data): + """OpenAI newer API: max_completion_tokens → max_tokens alias.""" + if isinstance(data, dict): + mct = data.pop("max_completion_tokens", None) + if mct is not None and data.get("max_tokens") is None: + data["max_tokens"] = mct + return data + + @model_validator(mode="before") + @classmethod + def normalize_messages(cls, data): + """Normalize incoming messages before pydantic union validation. + + Real-world clients (e.g. from other providers) send assistant tool_call + messages with content=null, which fails the strict Union type check. + Replace null content with "" so validation passes. + reasoning_content is intentionally kept — chat_utils.py wraps it as + ... for multi-turn reasoning history. + """ + messages = data.get("messages") + if not isinstance(messages, list): + return data normalized = [] for msg in messages: if not isinstance(msg, dict): @@ -567,31 +567,31 @@ class ChatCompletionRequest(OpenAIBaseModel): @model_validator(mode="before") @classmethod def validate_stream_options(cls, data): - if data.get("stream_options") and not data.get("stream"): - raise ValueError( - "Stream options can only be defined when `stream=True`.") - - return data - - @model_validator(mode="before") - @classmethod + if data.get("stream_options") and not data.get("stream"): + raise ValueError( + "Stream options can only be defined when `stream=True`.") + + return data + + @model_validator(mode="before") + @classmethod def check_logprobs(cls, data): if (prompt_logprobs := data.get("prompt_logprobs")) is not None: - if data.get("stream") and prompt_logprobs > 0: - raise ValueError( - "`prompt_logprobs` are not available when `stream=True`.") - - if prompt_logprobs < 0: - raise ValueError("`prompt_logprobs` must be a positive value.") - - if (top_logprobs := data.get("top_logprobs")) is not None: - if top_logprobs < 0: - raise ValueError("`top_logprobs` must be a positive value.") - - if not data.get("logprobs"): - raise ValueError( - "when using `top_logprobs`, `logprobs` must be set to true." - ) + if data.get("stream") and prompt_logprobs > 0: + raise ValueError( + "`prompt_logprobs` are not available when `stream=True`.") + + if prompt_logprobs < 0: + raise ValueError("`prompt_logprobs` must be a positive value.") + + if (top_logprobs := data.get("top_logprobs")) is not None: + if top_logprobs < 0: + raise ValueError("`top_logprobs` must be a positive value.") + + if not data.get("logprobs"): + raise ValueError( + "when using `top_logprobs`, `logprobs` must be set to true." + ) return data @@ -626,39 +626,39 @@ class ChatCompletionRequest(OpenAIBaseModel): "BI100 sampled prompt logprobs require positive " "`prompt_logprobs`.") return data - - @model_validator(mode="before") - @classmethod - def check_guided_decoding_count(cls, data): - if isinstance(data, ValueError): - raise data - - guide_count = sum([ - "guided_json" in data and data["guided_json"] is not None, - "guided_regex" in data and data["guided_regex"] is not None, - "guided_choice" in data and data["guided_choice"] is not None - ]) - # you can only use one kind of guided decoding - if guide_count > 1: - raise ValueError( - "You can only use one kind of guided decoding " - "('guided_json', 'guided_regex' or 'guided_choice').") + + @model_validator(mode="before") + @classmethod + def check_guided_decoding_count(cls, data): + if isinstance(data, ValueError): + raise data + + guide_count = sum([ + "guided_json" in data and data["guided_json"] is not None, + "guided_regex" in data and data["guided_regex"] is not None, + "guided_choice" in data and data["guided_choice"] is not None + ]) + # you can only use one kind of guided decoding + if guide_count > 1: + raise ValueError( + "You can only use one kind of guided decoding " + "('guided_json', 'guided_regex' or 'guided_choice').") # you can only either use guided decoding or a forced tool, not both if guide_count > 0 and data.get("tool_choice", "none") not in ("none", "auto"): - raise ValueError( - "You can only either use guided decoding or tools, not both.") - return data - - @model_validator(mode="before") - @classmethod - def check_tool_usage(cls, data): - - # if "tool_choice" is not specified but tools are provided, - # default to "auto" tool_choice - if "tool_choice" not in data and data.get("tools"): - data["tool_choice"] = "auto" - + raise ValueError( + "You can only either use guided decoding or tools, not both.") + return data + + @model_validator(mode="before") + @classmethod + def check_tool_usage(cls, data): + + # if "tool_choice" is not specified but tools are provided, + # default to "auto" tool_choice + if "tool_choice" not in data and data.get("tools"): + data["tool_choice"] = "auto" + # if "tool_choice" is specified -- validation if "tool_choice" in data: if data["tool_choice"] == "none": @@ -667,8 +667,8 @@ class ChatCompletionRequest(OpenAIBaseModel): # ensure that if "tool choice" is specified, tools are present if "tools" not in data or data["tools"] is None: raise ValueError( - "When using `tool_choice`, `tools` must be set.") - + "When using `tool_choice`, `tools` must be set.") + # make sure that tool choice is either a named tool # OR that it's set to "auto"/"none" if data["tool_choice"] != "auto" and not isinstance( @@ -676,511 +676,511 @@ class ChatCompletionRequest(OpenAIBaseModel): raise ValueError( "`tool_choice` must be a named tool, \"auto\", or " "\"none\".") - - # ensure that if "tool_choice" is specified as an object, - # it matches a valid tool - if isinstance(data["tool_choice"], dict): - valid_tool = False - specified_function = data["tool_choice"]["function"] - if not specified_function: - raise ValueError( - "Incorrectly formatted `tool_choice`. Should be like " - "`{\"type\": \"function\"," - " \"function\": {\"name\": \"my_function\"}}`") - specified_function_name = specified_function["name"] - if not specified_function_name: - raise ValueError( - "Incorrectly formatted `tool_choice`. Should be like " - "`{\"type\": \"function\", " - "\"function\": {\"name\": \"my_function\"}}`") - for tool in data["tools"]: - if tool["function"]["name"] == specified_function_name: - valid_tool = True - break - if not valid_tool: - raise ValueError( - "The tool specified in `tool_choice` does not match any" - " of the specified `tools`") - return data - - @model_validator(mode="before") - @classmethod - def check_generation_prompt(cls, data): - if data.get("continue_final_message") and data.get( - "add_generation_prompt"): - raise ValueError("Cannot set both `continue_final_message` and " - "`add_generation_prompt` to True.") - return data - - -class CompletionRequest(OpenAIBaseModel): - # Ordered by official OpenAI API documentation - # https://platform.openai.com/docs/api-reference/completions/create - model: str - prompt: Union[List[int], List[List[int]], str, List[str]] - best_of: Optional[int] = None - echo: Optional[bool] = False - frequency_penalty: Optional[float] = 0.0 - logit_bias: Optional[Dict[str, float]] = None - logprobs: Optional[int] = None - max_tokens: Optional[int] = 16 - n: int = 1 - presence_penalty: Optional[float] = 0.0 - seed: Optional[int] = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max) - stop: Optional[Union[str, List[str]]] = Field(default_factory=list) - stream: Optional[bool] = False - stream_options: Optional[StreamOptions] = None - suffix: Optional[str] = None - temperature: Optional[float] = 1.0 - top_p: Optional[float] = 1.0 - user: Optional[str] = None - - # doc: begin-completion-sampling-params - use_beam_search: bool = False - top_k: int = -1 - min_p: float = 0.0 - repetition_penalty: float = 1.0 - length_penalty: float = 1.0 - stop_token_ids: Optional[List[int]] = Field(default_factory=list) - include_stop_str_in_output: bool = False - ignore_eos: bool = False - min_tokens: int = 0 - skip_special_tokens: bool = True - spaces_between_special_tokens: bool = True - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None - allowed_token_ids: Optional[List[int]] = None - prompt_logprobs: Optional[int] = None - # doc: end-completion-sampling-params - - # doc: begin-completion-extra-params - add_special_tokens: bool = Field( - default=True, - description=( - "If true (the default), special tokens (e.g. BOS) will be added to " - "the prompt."), - ) - response_format: Optional[ResponseFormat] = Field( - default=None, - description= - ("Similar to chat completion, this parameter specifies the format of " - "output. Only {'type': 'json_object'} or {'type': 'text' } is " - "supported."), - ) - guided_json: Optional[Union[str, dict, BaseModel]] = Field( - default=None, - description="If specified, the output will follow the JSON schema.", - ) - guided_regex: Optional[str] = Field( - default=None, - description=( - "If specified, the output will follow the regex pattern."), - ) - guided_choice: Optional[List[str]] = Field( - default=None, - description=( - "If specified, the output will be exactly one of the choices."), - ) - guided_grammar: Optional[str] = Field( - default=None, - description=( - "If specified, the output will follow the context free grammar."), - ) - guided_decoding_backend: Optional[str] = Field( - default=None, - description=( - "If specified, will override the default guided decoding backend " - "of the server for this specific request. If set, must be one of " - "'outlines' / 'lm-format-enforcer'")) - guided_whitespace_pattern: Optional[str] = Field( - default=None, - description=( - "If specified, will override the default whitespace pattern " - "for guided json decoding.")) - priority: int = Field( - default=0, - description=( - "The priority of the request (lower means earlier handling; " - "default: 0). Any priority other than 0 will raise an error " - "if the served model does not use priority scheduling.")) - - # doc: end-completion-extra-params - - def to_beam_search_params(self, - default_max_tokens: int) -> BeamSearchParams: - max_tokens = self.max_tokens - if max_tokens is None: - max_tokens = default_max_tokens - - n = self.n if self.n is not None else 1 - temperature = self.temperature if self.temperature is not None else 0.0 - - return BeamSearchParams( - beam_width=n, - max_tokens=max_tokens, - ignore_eos=self.ignore_eos, - temperature=temperature, - length_penalty=self.length_penalty, - ) - - def to_sampling_params(self, default_max_tokens: int) -> SamplingParams: - max_tokens = self.max_tokens - if max_tokens is None: - max_tokens = default_max_tokens - - prompt_logprobs = self.prompt_logprobs - if prompt_logprobs is None and self.echo: - prompt_logprobs = self.logprobs - - echo_without_generation = self.echo and self.max_tokens == 0 - - guided_json_object = None + + # ensure that if "tool_choice" is specified as an object, + # it matches a valid tool + if isinstance(data["tool_choice"], dict): + valid_tool = False + specified_function = data["tool_choice"]["function"] + if not specified_function: + raise ValueError( + "Incorrectly formatted `tool_choice`. Should be like " + "`{\"type\": \"function\"," + " \"function\": {\"name\": \"my_function\"}}`") + specified_function_name = specified_function["name"] + if not specified_function_name: + raise ValueError( + "Incorrectly formatted `tool_choice`. Should be like " + "`{\"type\": \"function\", " + "\"function\": {\"name\": \"my_function\"}}`") + for tool in data["tools"]: + if tool["function"]["name"] == specified_function_name: + valid_tool = True + break + if not valid_tool: + raise ValueError( + "The tool specified in `tool_choice` does not match any" + " of the specified `tools`") + return data + + @model_validator(mode="before") + @classmethod + def check_generation_prompt(cls, data): + if data.get("continue_final_message") and data.get( + "add_generation_prompt"): + raise ValueError("Cannot set both `continue_final_message` and " + "`add_generation_prompt` to True.") + return data + + +class CompletionRequest(OpenAIBaseModel): + # Ordered by official OpenAI API documentation + # https://platform.openai.com/docs/api-reference/completions/create + model: str + prompt: Union[List[int], List[List[int]], str, List[str]] + best_of: Optional[int] = None + echo: Optional[bool] = False + frequency_penalty: Optional[float] = 0.0 + logit_bias: Optional[Dict[str, float]] = None + logprobs: Optional[int] = None + max_tokens: Optional[int] = 16 + n: int = 1 + presence_penalty: Optional[float] = 0.0 + seed: Optional[int] = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max) + stop: Optional[Union[str, List[str]]] = Field(default_factory=list) + stream: Optional[bool] = False + stream_options: Optional[StreamOptions] = None + suffix: Optional[str] = None + temperature: Optional[float] = 1.0 + top_p: Optional[float] = 1.0 + user: Optional[str] = None + + # doc: begin-completion-sampling-params + use_beam_search: bool = False + top_k: int = -1 + min_p: float = 0.0 + repetition_penalty: float = 1.0 + length_penalty: float = 1.0 + stop_token_ids: Optional[List[int]] = Field(default_factory=list) + include_stop_str_in_output: bool = False + ignore_eos: bool = False + min_tokens: int = 0 + skip_special_tokens: bool = True + spaces_between_special_tokens: bool = True + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None + allowed_token_ids: Optional[List[int]] = None + prompt_logprobs: Optional[int] = None + # doc: end-completion-sampling-params + + # doc: begin-completion-extra-params + add_special_tokens: bool = Field( + default=True, + description=( + "If true (the default), special tokens (e.g. BOS) will be added to " + "the prompt."), + ) + response_format: Optional[ResponseFormat] = Field( + default=None, + description= + ("Similar to chat completion, this parameter specifies the format of " + "output. Only {'type': 'json_object'} or {'type': 'text' } is " + "supported."), + ) + guided_json: Optional[Union[str, dict, BaseModel]] = Field( + default=None, + description="If specified, the output will follow the JSON schema.", + ) + guided_regex: Optional[str] = Field( + default=None, + description=( + "If specified, the output will follow the regex pattern."), + ) + guided_choice: Optional[List[str]] = Field( + default=None, + description=( + "If specified, the output will be exactly one of the choices."), + ) + guided_grammar: Optional[str] = Field( + default=None, + description=( + "If specified, the output will follow the context free grammar."), + ) + guided_decoding_backend: Optional[str] = Field( + default=None, + description=( + "If specified, will override the default guided decoding backend " + "of the server for this specific request. If set, must be one of " + "'outlines' / 'lm-format-enforcer'")) + guided_whitespace_pattern: Optional[str] = Field( + default=None, + description=( + "If specified, will override the default whitespace pattern " + "for guided json decoding.")) + priority: int = Field( + default=0, + description=( + "The priority of the request (lower means earlier handling; " + "default: 0). Any priority other than 0 will raise an error " + "if the served model does not use priority scheduling.")) + + # doc: end-completion-extra-params + + def to_beam_search_params(self, + default_max_tokens: int) -> BeamSearchParams: + max_tokens = self.max_tokens + if max_tokens is None: + max_tokens = default_max_tokens + + n = self.n if self.n is not None else 1 + temperature = self.temperature if self.temperature is not None else 0.0 + + return BeamSearchParams( + beam_width=n, + max_tokens=max_tokens, + ignore_eos=self.ignore_eos, + temperature=temperature, + length_penalty=self.length_penalty, + ) + + def to_sampling_params(self, default_max_tokens: int) -> SamplingParams: + max_tokens = self.max_tokens + if max_tokens is None: + max_tokens = default_max_tokens + + prompt_logprobs = self.prompt_logprobs + if prompt_logprobs is None and self.echo: + prompt_logprobs = self.logprobs + + echo_without_generation = self.echo and self.max_tokens == 0 + + guided_json_object = None guided_json_from_schema = None if self.response_format is not None: if self.response_format.type == "json_object": # Keep CompletionRequest aligned with ChatCompletionRequest. guided_json_from_schema = {"type": "object"} - elif (self.response_format.type == "json_schema" - and self.response_format.json_schema is not None - and self.response_format.json_schema.json_schema is not None): - guided_json_from_schema = \ - self.response_format.json_schema.json_schema - - guided_decoding = GuidedDecodingParams.from_optional( - json=self.guided_json or guided_json_from_schema, - regex=self.guided_regex, - choice=self.guided_choice, - grammar=self.guided_grammar, - json_object=guided_json_object, - backend=self.guided_decoding_backend, - whitespace_pattern=self.guided_whitespace_pattern) - - return SamplingParams.from_optional( - n=self.n, - best_of=self.best_of, - presence_penalty=self.presence_penalty, - frequency_penalty=self.frequency_penalty, - repetition_penalty=self.repetition_penalty, - temperature=self.temperature, - top_p=self.top_p, - top_k=self.top_k, - min_p=self.min_p, - seed=self.seed, - stop=self.stop, - stop_token_ids=self.stop_token_ids, - logprobs=self.logprobs, - ignore_eos=self.ignore_eos, - max_tokens=max_tokens if not echo_without_generation else 1, - min_tokens=self.min_tokens, - prompt_logprobs=prompt_logprobs, - skip_special_tokens=self.skip_special_tokens, - spaces_between_special_tokens=self.spaces_between_special_tokens, - include_stop_str_in_output=self.include_stop_str_in_output, - truncate_prompt_tokens=self.truncate_prompt_tokens, - output_kind=RequestOutputKind.DELTA if self.stream \ - else RequestOutputKind.FINAL_ONLY, - guided_decoding=guided_decoding, - logit_bias=self.logit_bias, - allowed_token_ids=self.allowed_token_ids) - - @model_validator(mode="before") - @classmethod - def check_guided_decoding_count(cls, data): - guide_count = sum([ - "guided_json" in data and data["guided_json"] is not None, - "guided_regex" in data and data["guided_regex"] is not None, - "guided_choice" in data and data["guided_choice"] is not None - ]) - if guide_count > 1: - raise ValueError( - "You can only use one kind of guided decoding " - "('guided_json', 'guided_regex' or 'guided_choice').") - return data - - @model_validator(mode="before") - @classmethod - def check_logprobs(cls, data): - if (prompt_logprobs := data.get("prompt_logprobs")) is not None: - if data.get("stream") and prompt_logprobs > 0: - raise ValueError( - "`prompt_logprobs` are not available when `stream=True`.") - - if prompt_logprobs < 0: - raise ValueError("`prompt_logprobs` must be a positive value.") - - if (logprobs := data.get("logprobs")) is not None and logprobs < 0: - raise ValueError("`logprobs` must be a positive value.") - - return data - - @model_validator(mode="before") - @classmethod - def validate_stream_options(cls, data): - if data.get("stream_options") and not data.get("stream"): - raise ValueError( - "Stream options can only be defined when `stream=True`.") - - return data - - -class EmbeddingRequest(OpenAIBaseModel): - # Ordered by official OpenAI API documentation - # https://platform.openai.com/docs/api-reference/embeddings - model: str - input: Union[List[int], List[List[int]], str, List[str]] - encoding_format: Literal["float", "base64"] = "float" - dimensions: Optional[int] = None - user: Optional[str] = None - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None - - # doc: begin-embedding-pooling-params - additional_data: Optional[Any] = None - - # doc: end-embedding-pooling-params - - # doc: begin-embedding-extra-params - priority: int = Field( - default=0, - description=( - "The priority of the request (lower means earlier handling; " - "default: 0). Any priority other than 0 will raise an error " - "if the served model does not use priority scheduling.")) - - # doc: end-embedding-extra-params - - def to_pooling_params(self): - return PoolingParams(additional_data=self.additional_data) - - -class CompletionLogProbs(OpenAIBaseModel): - text_offset: List[int] = Field(default_factory=list) - token_logprobs: List[Optional[float]] = Field(default_factory=list) - tokens: List[str] = Field(default_factory=list) - top_logprobs: List[Optional[Dict[str, - float]]] = Field(default_factory=list) - - -class CompletionResponseChoice(OpenAIBaseModel): - index: int - text: str - logprobs: Optional[CompletionLogProbs] = None - finish_reason: Optional[str] = None - stop_reason: Optional[Union[int, str]] = Field( - default=None, - description=( - "The stop string or token id that caused the completion " - "to stop, None if the completion finished for some other reason " - "including encountering the EOS token"), - ) - prompt_logprobs: Optional[List[Optional[Dict[int, Logprob]]]] = None - - -class CompletionResponse(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"cmpl-{random_uuid()}") - object: str = "text_completion" - created: int = Field(default_factory=lambda: int(time.time())) - model: str - choices: List[CompletionResponseChoice] - usage: UsageInfo - - -class CompletionResponseStreamChoice(OpenAIBaseModel): - index: int - text: str - logprobs: Optional[CompletionLogProbs] = None - finish_reason: Optional[str] = None - stop_reason: Optional[Union[int, str]] = Field( - default=None, - description=( - "The stop string or token id that caused the completion " - "to stop, None if the completion finished for some other reason " - "including encountering the EOS token"), - ) - - -class CompletionStreamResponse(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"cmpl-{random_uuid()}") - object: str = "text_completion" - created: int = Field(default_factory=lambda: int(time.time())) - model: str - choices: List[CompletionResponseStreamChoice] - usage: Optional[UsageInfo] = Field(default=None) - - -class EmbeddingResponseData(OpenAIBaseModel): - index: int - object: str = "embedding" - embedding: Union[List[float], str] - - -class EmbeddingResponse(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"cmpl-{random_uuid()}") - object: str = "list" - created: int = Field(default_factory=lambda: int(time.time())) - model: str - data: List[EmbeddingResponseData] - usage: UsageInfo - - -class FunctionCall(OpenAIBaseModel): - name: str - arguments: str - - -class ToolCall(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"chatcmpl-tool-{random_uuid()}") - type: Literal["function"] = "function" - function: FunctionCall - - -class DeltaFunctionCall(BaseModel): - name: Optional[str] = None - arguments: Optional[str] = None - - -# a tool call delta where everything is optional -class DeltaToolCall(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"chatcmpl-tool-{random_uuid()}") - type: Literal["function"] = "function" - index: int - function: Optional[DeltaFunctionCall] = None - - -class ExtractedToolCallInformation(BaseModel): - # indicate if tools were called - tools_called: bool - - # extracted tool calls - tool_calls: List[ToolCall] - - # content - per OpenAI spec, content AND tool calls can be returned rarely - # But some models will do this intentionally - content: Optional[str] = None - - -class ChatMessage(OpenAIBaseModel): - role: str - reasoning_content: Optional[str] = None - content: Optional[str] = None - tool_calls: List[ToolCall] = Field(default_factory=list) - - -class ChatCompletionLogProb(OpenAIBaseModel): - token: str - logprob: float = -9999.0 - bytes: Optional[List[int]] = None - - -class ChatCompletionLogProbsContent(ChatCompletionLogProb): - top_logprobs: List[ChatCompletionLogProb] = Field(default_factory=list) - - -class ChatCompletionLogProbs(OpenAIBaseModel): - content: Optional[List[ChatCompletionLogProbsContent]] = None - - -class ChatCompletionResponseChoice(OpenAIBaseModel): - index: int - message: ChatMessage - logprobs: Optional[ChatCompletionLogProbs] = None - # per OpenAI spec this is the default - finish_reason: Optional[str] = "stop" - # not part of the OpenAI spec but included in vLLM for legacy reasons - stop_reason: Optional[Union[int, str]] = None - - -class ChatCompletionResponse(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"chatcmpl-{random_uuid()}") - object: Literal["chat.completion"] = "chat.completion" - created: int = Field(default_factory=lambda: int(time.time())) - model: str - choices: List[ChatCompletionResponseChoice] - usage: UsageInfo - prompt_logprobs: Optional[List[Optional[Dict[int, Logprob]]]] = None - - -class DeltaMessage(OpenAIBaseModel): - role: Optional[str] = None - reasoning_content: Optional[str] = None - content: Optional[str] = None - tool_calls: List[DeltaToolCall] = Field(default_factory=list) - - -class ChatCompletionResponseStreamChoice(OpenAIBaseModel): - index: int - delta: DeltaMessage - logprobs: Optional[ChatCompletionLogProbs] = None - finish_reason: Optional[str] = None - stop_reason: Optional[Union[int, str]] = None - - -class ChatCompletionStreamResponse(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"chatcmpl-{random_uuid()}") - object: Literal["chat.completion.chunk"] = "chat.completion.chunk" - created: int = Field(default_factory=lambda: int(time.time())) - model: str - choices: List[ChatCompletionResponseStreamChoice] - usage: Optional[UsageInfo] = Field(default=None) - - -class BatchRequestInput(OpenAIBaseModel): - """ - The per-line object of the batch input file. - - NOTE: Currently only the `/v1/chat/completions` endpoint is supported. - """ - - # A developer-provided per-request id that will be used to match outputs to - # inputs. Must be unique for each request in a batch. - custom_id: str - - # The HTTP method to be used for the request. Currently only POST is - # supported. - method: str - - # The OpenAI API relative URL to be used for the request. Currently - # /v1/chat/completions is supported. - url: str - - # The parameters of the request. - body: Union[ChatCompletionRequest, EmbeddingRequest] - - -class BatchResponseData(OpenAIBaseModel): - # HTTP status code of the response. - status_code: int = 200 - - # An unique identifier for the API request. - request_id: str - - # The body of the response. - body: Optional[Union[ChatCompletionResponse, EmbeddingResponse]] = None - - -class BatchRequestOutput(OpenAIBaseModel): - """ - The per-line object of the batch output and error files - """ - - id: str - - # A developer-provided per-request id that will be used to match outputs to - # inputs. - custom_id: str - - response: Optional[BatchResponseData] - - # For requests that failed with a non-HTTP error, this will contain more - # information on the cause of the failure. - error: Optional[Any] - - -class TokenizeCompletionRequest(OpenAIBaseModel): - model: str - prompt: str - - add_special_tokens: bool = Field(default=True) - - + elif (self.response_format.type == "json_schema" + and self.response_format.json_schema is not None + and self.response_format.json_schema.json_schema is not None): + guided_json_from_schema = \ + self.response_format.json_schema.json_schema + + guided_decoding = GuidedDecodingParams.from_optional( + json=self.guided_json or guided_json_from_schema, + regex=self.guided_regex, + choice=self.guided_choice, + grammar=self.guided_grammar, + json_object=guided_json_object, + backend=self.guided_decoding_backend, + whitespace_pattern=self.guided_whitespace_pattern) + + return SamplingParams.from_optional( + n=self.n, + best_of=self.best_of, + presence_penalty=self.presence_penalty, + frequency_penalty=self.frequency_penalty, + repetition_penalty=self.repetition_penalty, + temperature=self.temperature, + top_p=self.top_p, + top_k=self.top_k, + min_p=self.min_p, + seed=self.seed, + stop=self.stop, + stop_token_ids=self.stop_token_ids, + logprobs=self.logprobs, + ignore_eos=self.ignore_eos, + max_tokens=max_tokens if not echo_without_generation else 1, + min_tokens=self.min_tokens, + prompt_logprobs=prompt_logprobs, + skip_special_tokens=self.skip_special_tokens, + spaces_between_special_tokens=self.spaces_between_special_tokens, + include_stop_str_in_output=self.include_stop_str_in_output, + truncate_prompt_tokens=self.truncate_prompt_tokens, + output_kind=RequestOutputKind.DELTA if self.stream \ + else RequestOutputKind.FINAL_ONLY, + guided_decoding=guided_decoding, + logit_bias=self.logit_bias, + allowed_token_ids=self.allowed_token_ids) + + @model_validator(mode="before") + @classmethod + def check_guided_decoding_count(cls, data): + guide_count = sum([ + "guided_json" in data and data["guided_json"] is not None, + "guided_regex" in data and data["guided_regex"] is not None, + "guided_choice" in data and data["guided_choice"] is not None + ]) + if guide_count > 1: + raise ValueError( + "You can only use one kind of guided decoding " + "('guided_json', 'guided_regex' or 'guided_choice').") + return data + + @model_validator(mode="before") + @classmethod + def check_logprobs(cls, data): + if (prompt_logprobs := data.get("prompt_logprobs")) is not None: + if data.get("stream") and prompt_logprobs > 0: + raise ValueError( + "`prompt_logprobs` are not available when `stream=True`.") + + if prompt_logprobs < 0: + raise ValueError("`prompt_logprobs` must be a positive value.") + + if (logprobs := data.get("logprobs")) is not None and logprobs < 0: + raise ValueError("`logprobs` must be a positive value.") + + return data + + @model_validator(mode="before") + @classmethod + def validate_stream_options(cls, data): + if data.get("stream_options") and not data.get("stream"): + raise ValueError( + "Stream options can only be defined when `stream=True`.") + + return data + + +class EmbeddingRequest(OpenAIBaseModel): + # Ordered by official OpenAI API documentation + # https://platform.openai.com/docs/api-reference/embeddings + model: str + input: Union[List[int], List[List[int]], str, List[str]] + encoding_format: Literal["float", "base64"] = "float" + dimensions: Optional[int] = None + user: Optional[str] = None + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None + + # doc: begin-embedding-pooling-params + additional_data: Optional[Any] = None + + # doc: end-embedding-pooling-params + + # doc: begin-embedding-extra-params + priority: int = Field( + default=0, + description=( + "The priority of the request (lower means earlier handling; " + "default: 0). Any priority other than 0 will raise an error " + "if the served model does not use priority scheduling.")) + + # doc: end-embedding-extra-params + + def to_pooling_params(self): + return PoolingParams(additional_data=self.additional_data) + + +class CompletionLogProbs(OpenAIBaseModel): + text_offset: List[int] = Field(default_factory=list) + token_logprobs: List[Optional[float]] = Field(default_factory=list) + tokens: List[str] = Field(default_factory=list) + top_logprobs: List[Optional[Dict[str, + float]]] = Field(default_factory=list) + + +class CompletionResponseChoice(OpenAIBaseModel): + index: int + text: str + logprobs: Optional[CompletionLogProbs] = None + finish_reason: Optional[str] = None + stop_reason: Optional[Union[int, str]] = Field( + default=None, + description=( + "The stop string or token id that caused the completion " + "to stop, None if the completion finished for some other reason " + "including encountering the EOS token"), + ) + prompt_logprobs: Optional[List[Optional[Dict[int, Logprob]]]] = None + + +class CompletionResponse(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"cmpl-{random_uuid()}") + object: str = "text_completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[CompletionResponseChoice] + usage: UsageInfo + + +class CompletionResponseStreamChoice(OpenAIBaseModel): + index: int + text: str + logprobs: Optional[CompletionLogProbs] = None + finish_reason: Optional[str] = None + stop_reason: Optional[Union[int, str]] = Field( + default=None, + description=( + "The stop string or token id that caused the completion " + "to stop, None if the completion finished for some other reason " + "including encountering the EOS token"), + ) + + +class CompletionStreamResponse(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"cmpl-{random_uuid()}") + object: str = "text_completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[CompletionResponseStreamChoice] + usage: Optional[UsageInfo] = Field(default=None) + + +class EmbeddingResponseData(OpenAIBaseModel): + index: int + object: str = "embedding" + embedding: Union[List[float], str] + + +class EmbeddingResponse(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"cmpl-{random_uuid()}") + object: str = "list" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + data: List[EmbeddingResponseData] + usage: UsageInfo + + +class FunctionCall(OpenAIBaseModel): + name: str + arguments: str + + +class ToolCall(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"chatcmpl-tool-{random_uuid()}") + type: Literal["function"] = "function" + function: FunctionCall + + +class DeltaFunctionCall(BaseModel): + name: Optional[str] = None + arguments: Optional[str] = None + + +# a tool call delta where everything is optional +class DeltaToolCall(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"chatcmpl-tool-{random_uuid()}") + type: Literal["function"] = "function" + index: int + function: Optional[DeltaFunctionCall] = None + + +class ExtractedToolCallInformation(BaseModel): + # indicate if tools were called + tools_called: bool + + # extracted tool calls + tool_calls: List[ToolCall] + + # content - per OpenAI spec, content AND tool calls can be returned rarely + # But some models will do this intentionally + content: Optional[str] = None + + +class ChatMessage(OpenAIBaseModel): + role: str + reasoning_content: Optional[str] = None + content: Optional[str] = None + tool_calls: List[ToolCall] = Field(default_factory=list) + + +class ChatCompletionLogProb(OpenAIBaseModel): + token: str + logprob: float = -9999.0 + bytes: Optional[List[int]] = None + + +class ChatCompletionLogProbsContent(ChatCompletionLogProb): + top_logprobs: List[ChatCompletionLogProb] = Field(default_factory=list) + + +class ChatCompletionLogProbs(OpenAIBaseModel): + content: Optional[List[ChatCompletionLogProbsContent]] = None + + +class ChatCompletionResponseChoice(OpenAIBaseModel): + index: int + message: ChatMessage + logprobs: Optional[ChatCompletionLogProbs] = None + # per OpenAI spec this is the default + finish_reason: Optional[str] = "stop" + # not part of the OpenAI spec but included in vLLM for legacy reasons + stop_reason: Optional[Union[int, str]] = None + + +class ChatCompletionResponse(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"chatcmpl-{random_uuid()}") + object: Literal["chat.completion"] = "chat.completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[ChatCompletionResponseChoice] + usage: UsageInfo + prompt_logprobs: Optional[List[Optional[Dict[int, Logprob]]]] = None + + +class DeltaMessage(OpenAIBaseModel): + role: Optional[str] = None + reasoning_content: Optional[str] = None + content: Optional[str] = None + tool_calls: List[DeltaToolCall] = Field(default_factory=list) + + +class ChatCompletionResponseStreamChoice(OpenAIBaseModel): + index: int + delta: DeltaMessage + logprobs: Optional[ChatCompletionLogProbs] = None + finish_reason: Optional[str] = None + stop_reason: Optional[Union[int, str]] = None + + +class ChatCompletionStreamResponse(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"chatcmpl-{random_uuid()}") + object: Literal["chat.completion.chunk"] = "chat.completion.chunk" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[ChatCompletionResponseStreamChoice] + usage: Optional[UsageInfo] = Field(default=None) + + +class BatchRequestInput(OpenAIBaseModel): + """ + The per-line object of the batch input file. + + NOTE: Currently only the `/v1/chat/completions` endpoint is supported. + """ + + # A developer-provided per-request id that will be used to match outputs to + # inputs. Must be unique for each request in a batch. + custom_id: str + + # The HTTP method to be used for the request. Currently only POST is + # supported. + method: str + + # The OpenAI API relative URL to be used for the request. Currently + # /v1/chat/completions is supported. + url: str + + # The parameters of the request. + body: Union[ChatCompletionRequest, EmbeddingRequest] + + +class BatchResponseData(OpenAIBaseModel): + # HTTP status code of the response. + status_code: int = 200 + + # An unique identifier for the API request. + request_id: str + + # The body of the response. + body: Optional[Union[ChatCompletionResponse, EmbeddingResponse]] = None + + +class BatchRequestOutput(OpenAIBaseModel): + """ + The per-line object of the batch output and error files + """ + + id: str + + # A developer-provided per-request id that will be used to match outputs to + # inputs. + custom_id: str + + response: Optional[BatchResponseData] + + # For requests that failed with a non-HTTP error, this will contain more + # information on the cause of the failure. + error: Optional[Any] + + +class TokenizeCompletionRequest(OpenAIBaseModel): + model: str + prompt: str + + add_special_tokens: bool = Field(default=True) + + class TokenizeChatRequest(OpenAIBaseModel): model: str messages: List[ChatCompletionMessageParam] @@ -1189,40 +1189,40 @@ class TokenizeChatRequest(OpenAIBaseModel): continue_final_message: bool = Field(default=False) add_special_tokens: bool = Field(default=False) chat_template_kwargs: Optional[Dict[str, Any]] = Field(default=None) - - @model_validator(mode="before") - @classmethod - def check_generation_prompt(cls, data): - if data.get("continue_final_message") and data.get( - "add_generation_prompt"): - raise ValueError("Cannot set both `continue_final_message` and " - "`add_generation_prompt` to True.") - return data - - -TokenizeRequest = Union[TokenizeCompletionRequest, TokenizeChatRequest] - - -class TokenizeResponse(OpenAIBaseModel): - count: int - max_model_len: int - tokens: List[int] - - -class DetokenizeRequest(OpenAIBaseModel): - model: str - tokens: List[int] - - -class DetokenizeResponse(OpenAIBaseModel): - prompt: str - - -class LoadLoraAdapterRequest(BaseModel): - lora_name: str - lora_path: str - - -class UnloadLoraAdapterRequest(BaseModel): - lora_name: str - lora_int_id: Optional[int] = Field(default=None) + + @model_validator(mode="before") + @classmethod + def check_generation_prompt(cls, data): + if data.get("continue_final_message") and data.get( + "add_generation_prompt"): + raise ValueError("Cannot set both `continue_final_message` and " + "`add_generation_prompt` to True.") + return data + + +TokenizeRequest = Union[TokenizeCompletionRequest, TokenizeChatRequest] + + +class TokenizeResponse(OpenAIBaseModel): + count: int + max_model_len: int + tokens: List[int] + + +class DetokenizeRequest(OpenAIBaseModel): + model: str + tokens: List[int] + + +class DetokenizeResponse(OpenAIBaseModel): + prompt: str + + +class LoadLoraAdapterRequest(BaseModel): + lora_name: str + lora_path: str + + +class UnloadLoraAdapterRequest(BaseModel): + lora_name: str + lora_int_id: Optional[int] = Field(default=None) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 93f634c6..8103948e 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -1,5 +1,5 @@ # Inference-only Qwen3.6-35B-A3B (Qwen3_5 MoE architecture) for Iluvatar BI-V100. -# Pure-PyTorch DeltaNet (no fla / causal_conv1d dependency). +# Pure-PyTorch DeltaNet (no fla / causal_conv1d dependency). # Includes the native Qwen3.6 vision tower; MTP remains unsupported. from functools import lru_cache, partial @@ -9,7 +9,7 @@ import sys import time from typing import (Any, Dict, Iterable, List, Literal, Mapping, Optional, Tuple, TypedDict, Union) - + def _bi100_model_trace(message: str) -> None: if os.getenv("BI100_EXECUTOR_STARTUP_DEBUG") == "1": stamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) @@ -55,25 +55,25 @@ if not hasattr(_qwen2_vl_image_processing, "make_batched_videos"): from vllm.attention import Attention, AttentionMetadata from vllm.config import (CacheConfig, LoRAConfig, MultiModalConfig, SchedulerConfig) -from vllm.distributed import (get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce) -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.layernorm import GemmaRMSNorm -from vllm.model_executor.layers.linear import (ColumnParallelLinear, - MergedColumnParallelLinear, - ReplicatedLinear, - RowParallelLinear) -from vllm.model_executor.layers.fused_moe import FusedMoE -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.distributed import (get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce) +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.layernorm import GemmaRMSNorm +from vllm.model_executor.layers.linear import (ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear) +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import ( MRotaryEmbedding, _apply_rotary_emb) -from vllm.model_executor.layers.sampler import Sampler, SamplerOutput -from vllm.model_executor.layers.vocab_parallel_embedding import ( - ParallelLMHead, VocabParallelEmbedding) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, sharded_weight_loader) +from vllm.model_executor.layers.sampler import Sampler, SamplerOutput +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, VocabParallelEmbedding) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, sharded_weight_loader) from vllm.model_executor.models.mamba_cache import MambaCacheManager from vllm.model_executor.models.qwen2_vl import (Qwen2VisionAttention, Qwen2VisionRotaryEmbedding) @@ -85,8 +85,8 @@ from vllm.multimodal import (MULTIMODAL_REGISTRY, MultiModalDataDict, from vllm.multimodal.base import MultiModalData from vllm.sequence import IntermediateTensors, SequenceData from vllm.transformers_utils.tokenizer import get_tokenizer -from vllm.worker.model_runner import (_BATCH_SIZES_TO_CAPTURE, - _get_graph_batch_size) +from vllm.worker.model_runner import (_BATCH_SIZES_TO_CAPTURE, + _get_graph_batch_size) from vllm.logger import init_logger from vllm.bi100_env import env_bool, env_int from vllm.bi100_profile import (bi100_profile_event_enabled, @@ -635,7 +635,7 @@ def input_processor_for_qwen36(ctx: InputContext, # --------------------------------------------------------------------------- # Pure-PyTorch DeltaNet kernels (fallbacks from transformers 5.2.0) # --------------------------------------------------------------------------- - + def _l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor: return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) @@ -673,172 +673,172 @@ def _validate_gdn_prefix_key(key: Any) -> Tuple[int, bytes]: or not isinstance(key[1], bytes) or len(key[1]) != 32): raise RuntimeError(f"invalid GDN prefix key: {key!r}") return key - - -def _torch_causal_conv1d_update( - hidden_states: torch.Tensor, # (batch, channels, seq=1) - conv_state: torch.Tensor, # (batch, channels, state_len) modified in-place - weight: torch.Tensor, # (channels, kernel_size) - bias: Optional[torch.Tensor] = None, - activation: Optional[str] = None, -) -> torch.Tensor: - _, channels, seq_len = hidden_states.shape - state_len = conv_state.shape[-1] - cat = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) - conv_state.copy_(cat[:, :, -state_len:]) - out = F.conv1d(cat, weight.unsqueeze(1), bias, padding=0, groups=channels) - out = out[:, :, -seq_len:] - if activation is not None: - out = F.silu(out) - return out.to(hidden_states.dtype) - - -def _torch_chunk_gated_delta_rule( - query: torch.Tensor, # (batch, seq, num_heads, head_k_dim) - key: torch.Tensor, - value: torch.Tensor, # (batch, seq, num_heads, head_v_dim) - g: torch.Tensor, # (batch, seq, num_heads) - beta: torch.Tensor, # (batch, seq, num_heads) - chunk_size: int = 64, - initial_state: Optional[torch.Tensor] = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, -) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + + +def _torch_causal_conv1d_update( + hidden_states: torch.Tensor, # (batch, channels, seq=1) + conv_state: torch.Tensor, # (batch, channels, state_len) modified in-place + weight: torch.Tensor, # (channels, kernel_size) + bias: Optional[torch.Tensor] = None, + activation: Optional[str] = None, +) -> torch.Tensor: + _, channels, seq_len = hidden_states.shape + state_len = conv_state.shape[-1] + cat = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) + conv_state.copy_(cat[:, :, -state_len:]) + out = F.conv1d(cat, weight.unsqueeze(1), bias, padding=0, groups=channels) + out = out[:, :, -seq_len:] + if activation is not None: + out = F.silu(out) + return out.to(hidden_states.dtype) + + +def _torch_chunk_gated_delta_rule( + query: torch.Tensor, # (batch, seq, num_heads, head_k_dim) + key: torch.Tensor, + value: torch.Tensor, # (batch, seq, num_heads, head_v_dim) + g: torch.Tensor, # (batch, seq, num_heads) + beta: torch.Tensor, # (batch, seq, num_heads) + chunk_size: int = 64, + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: if use_qk_l2norm_in_kernel: query = _l2norm(query) key = _l2norm(key) - # Transpose to (batch, num_heads, seq, dim) - query, key, value, beta, g = [ - x.transpose(1, 2).contiguous().to(torch.float32) - for x in (query, key, value, beta, g) - ] - batch, num_heads, seq_len, k_dim = key.shape - v_dim = value.shape[-1] - pad = (chunk_size - seq_len % chunk_size) % chunk_size - query = F.pad(query, (0, 0, 0, pad)) - key = F.pad(key, (0, 0, 0, pad)) - value = F.pad(value, (0, 0, 0, pad)) - beta = F.pad(beta, (0, pad)) - g = F.pad(g, (0, pad)) - total_len = seq_len + pad - scale = 1.0 / (query.shape[-1] ** 0.5) - query = query * scale - - v_beta = value * beta.unsqueeze(-1) - k_beta = key * beta.unsqueeze(-1) - query, key, value, k_beta, v_beta = [ - x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) - for x in (query, key, value, k_beta, v_beta) - ] - g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) - mask_upper = torch.triu( - torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), - diagonal=0) - - g = g.cumsum(dim=-1) - decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() - attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask_upper, 0) - for i in range(1, chunk_size): - row = attn[..., i, :i].clone() - sub = attn[..., :i, :i].clone() - attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) - attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) - value = attn @ v_beta - k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) - - last_state = ( - torch.zeros(batch, num_heads, k_dim, v_dim, dtype=value.dtype, device=value.device) - if initial_state is None - else initial_state.to(value) - ) - core_out = torch.zeros_like(value) - mask_upper2 = torch.triu( - torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), - diagonal=1) - - for i in range(total_len // chunk_size): - q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] - attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0) - v_prime = k_cumdecay[:, :, i] @ last_state - v_new = v_i - v_prime - attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state - core_out[:, :, i] = attn_inter + attn_i @ v_new - last_state = ( - last_state * g[:, :, i, -1, None, None].exp() - + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]) - .transpose(-1, -2) @ v_new - ) - + # Transpose to (batch, num_heads, seq, dim) + query, key, value, beta, g = [ + x.transpose(1, 2).contiguous().to(torch.float32) + for x in (query, key, value, beta, g) + ] + batch, num_heads, seq_len, k_dim = key.shape + v_dim = value.shape[-1] + pad = (chunk_size - seq_len % chunk_size) % chunk_size + query = F.pad(query, (0, 0, 0, pad)) + key = F.pad(key, (0, 0, 0, pad)) + value = F.pad(value, (0, 0, 0, pad)) + beta = F.pad(beta, (0, pad)) + g = F.pad(g, (0, pad)) + total_len = seq_len + pad + scale = 1.0 / (query.shape[-1] ** 0.5) + query = query * scale + + v_beta = value * beta.unsqueeze(-1) + k_beta = key * beta.unsqueeze(-1) + query, key, value, k_beta, v_beta = [ + x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) + for x in (query, key, value, k_beta, v_beta) + ] + g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) + mask_upper = torch.triu( + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), + diagonal=0) + + g = g.cumsum(dim=-1) + decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() + attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask_upper, 0) + for i in range(1, chunk_size): + row = attn[..., i, :i].clone() + sub = attn[..., :i, :i].clone() + attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) + value = attn @ v_beta + k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) + + last_state = ( + torch.zeros(batch, num_heads, k_dim, v_dim, dtype=value.dtype, device=value.device) + if initial_state is None + else initial_state.to(value) + ) + core_out = torch.zeros_like(value) + mask_upper2 = torch.triu( + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), + diagonal=1) + + for i in range(total_len // chunk_size): + q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] + attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0) + v_prime = k_cumdecay[:, :, i] @ last_state + v_new = v_i - v_prime + attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state + core_out[:, :, i] = attn_inter + attn_i @ v_new + last_state = ( + last_state * g[:, :, i, -1, None, None].exp() + + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]) + .transpose(-1, -2) @ v_new + ) + if not output_final_state: last_state = None core_out = core_out.reshape(batch, num_heads, -1, v_dim)[:, :, :seq_len] core_out = core_out.transpose(1, 2).contiguous() return core_out, last_state - -def _torch_recurrent_gated_delta_rule( - query: torch.Tensor, # (batch, 1, num_heads, head_k_dim) - key: torch.Tensor, - value: torch.Tensor, - g: torch.Tensor, # (batch, 1, num_heads) - beta: torch.Tensor, - initial_state: Optional[torch.Tensor] = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, -) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + +def _torch_recurrent_gated_delta_rule( + query: torch.Tensor, # (batch, 1, num_heads, head_k_dim) + key: torch.Tensor, + value: torch.Tensor, + g: torch.Tensor, # (batch, 1, num_heads) + beta: torch.Tensor, + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: if use_qk_l2norm_in_kernel: query = _l2norm(query) key = _l2norm(key) - query, key, value, beta, g = [ - x.transpose(1, 2).contiguous().to(torch.float32) - for x in (query, key, value, beta, g) - ] - batch, num_heads, seq_len, k_dim = key.shape - v_dim = value.shape[-1] - scale = 1.0 / (query.shape[-1] ** 0.5) - query = query * scale - - core_out = torch.zeros(batch, num_heads, seq_len, v_dim, - dtype=value.dtype, device=value.device) - last_state = ( - torch.zeros(batch, num_heads, k_dim, v_dim, - dtype=value.dtype, device=value.device) - if initial_state is None - else initial_state.to(value) - ) - for t in range(seq_len): - q_t = query[:, :, t] - k_t = key[:, :, t] - v_t = value[:, :, t] - g_t = g[:, :, t].exp().unsqueeze(-1).unsqueeze(-1) - beta_t = beta[:, :, t].unsqueeze(-1) - last_state = last_state * g_t - kv_mem = (last_state * k_t.unsqueeze(-1)).sum(dim=-2) - delta = (v_t - kv_mem) * beta_t - last_state = last_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2) - core_out[:, :, t] = (last_state * q_t.unsqueeze(-1)).sum(dim=-2) - + query, key, value, beta, g = [ + x.transpose(1, 2).contiguous().to(torch.float32) + for x in (query, key, value, beta, g) + ] + batch, num_heads, seq_len, k_dim = key.shape + v_dim = value.shape[-1] + scale = 1.0 / (query.shape[-1] ** 0.5) + query = query * scale + + core_out = torch.zeros(batch, num_heads, seq_len, v_dim, + dtype=value.dtype, device=value.device) + last_state = ( + torch.zeros(batch, num_heads, k_dim, v_dim, + dtype=value.dtype, device=value.device) + if initial_state is None + else initial_state.to(value) + ) + for t in range(seq_len): + q_t = query[:, :, t] + k_t = key[:, :, t] + v_t = value[:, :, t] + g_t = g[:, :, t].exp().unsqueeze(-1).unsqueeze(-1) + beta_t = beta[:, :, t].unsqueeze(-1) + last_state = last_state * g_t + kv_mem = (last_state * k_t.unsqueeze(-1)).sum(dim=-2) + delta = (v_t - kv_mem) * beta_t + last_state = last_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2) + core_out[:, :, t] = (last_state * q_t.unsqueeze(-1)).sum(dim=-2) + if not output_final_state: last_state = None core_out = core_out.transpose(1, 2).contiguous() return core_out, last_state - - -# --------------------------------------------------------------------------- -# Gated RMSNorm (for DeltaNet output normalisation) -# --------------------------------------------------------------------------- - + + +# --------------------------------------------------------------------------- +# Gated RMSNorm (for DeltaNet output normalisation) +# --------------------------------------------------------------------------- + class Qwen3_5RMSNormGated(nn.Module): - def __init__(self, hidden_size: int, eps: float = 1e-6): - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + def forward(self, hidden_states: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - hs = hidden_states.to(torch.float32) - variance = hs.pow(2).mean(-1, keepdim=True) - hs = hs * torch.rsqrt(variance + self.variance_epsilon) + input_dtype = hidden_states.dtype + hs = hidden_states.to(torch.float32) + variance = hs.pow(2).mean(-1, keepdim=True) + hs = hs * torch.rsqrt(variance + self.variance_epsilon) hs = self.weight * hs.to(input_dtype) return (hs * F.silu(gate.to(torch.float32))).to(input_dtype) @@ -941,29 +941,29 @@ def _load_full_attention_qgkv_weight(params_dict, name: str, # --------------------------------------------------------------------------- # Gated DeltaNet (linear_attention layers) # --------------------------------------------------------------------------- - -class GatedDeltaNet(nn.Module): - def __init__( - self, - text_cfg, - layer_idx: int, - quant_config: Optional[QuantizationConfig] = None, - ) -> None: - super().__init__() - self.layer_idx = layer_idx - self.hidden_size = text_cfg.hidden_size + +class GatedDeltaNet(nn.Module): + def __init__( + self, + text_cfg, + layer_idx: int, + quant_config: Optional[QuantizationConfig] = None, + ) -> None: + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = text_cfg.hidden_size self.num_v_heads = text_cfg.linear_num_value_heads # checkpoint: 32 self.num_k_heads = text_cfg.linear_num_key_heads # checkpoint: 16 - self.head_k_dim = text_cfg.linear_key_head_dim # 128 - self.head_v_dim = text_cfg.linear_value_head_dim # 128 - self.key_dim = self.num_k_heads * self.head_k_dim # 2048 + self.head_k_dim = text_cfg.linear_key_head_dim # 128 + self.head_v_dim = text_cfg.linear_value_head_dim # 128 + self.key_dim = self.num_k_heads * self.head_k_dim # 2048 self.value_dim = self.num_v_heads * self.head_v_dim # checkpoint: 4096 self.conv_dim = self.key_dim * 2 + self.value_dim # checkpoint: 8192 - self.conv_kernel_size = text_cfg.linear_conv_kernel_dim # 4 + self.conv_kernel_size = text_cfg.linear_conv_kernel_dim # 4 self.head_expand_ratio = self.num_v_heads // self.num_k_heads # checkpoint: 2 - - tp_size = get_tensor_model_parallel_world_size() - + + tp_size = get_tensor_model_parallel_world_size() + # Keep each logical projection independently TP-sharded while executing # one GEMM. Per-rank output order is [q, k, v, z, beta, decay]. self.in_proj_qkvzba = MergedColumnParallelLinear( @@ -971,103 +971,103 @@ class GatedDeltaNet(nn.Module): [self.key_dim, self.key_dim, self.value_dim, self.value_dim, self.num_v_heads, self.num_v_heads], bias=False, quant_config=quant_config) - self.out_proj = RowParallelLinear( - self.value_dim, self.hidden_size, - bias=False, quant_config=quant_config) - - # Depthwise conv weight — sharded along channel dim (dim 0) - local_conv_dim = self.conv_dim // tp_size - self.conv1d_weight = nn.Parameter( - torch.empty(local_conv_dim, 1, self.conv_kernel_size)) - set_weight_attrs(self.conv1d_weight, { - "weight_loader": self._conv1d_weight_loader}) - - # Per-head scalar parameters — sharded along dim 0 - local_num_v = self.num_v_heads // tp_size - self.A_log = nn.Parameter(torch.zeros(local_num_v)) - self.dt_bias = nn.Parameter(torch.zeros(local_num_v)) - set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)}) - set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) - - # Gated RMSNorm on head_v_dim — replicated (head_v_dim=128 is small) + self.out_proj = RowParallelLinear( + self.value_dim, self.hidden_size, + bias=False, quant_config=quant_config) + + # Depthwise conv weight — sharded along channel dim (dim 0) + local_conv_dim = self.conv_dim // tp_size + self.conv1d_weight = nn.Parameter( + torch.empty(local_conv_dim, 1, self.conv_kernel_size)) + set_weight_attrs(self.conv1d_weight, { + "weight_loader": self._conv1d_weight_loader}) + + # Per-head scalar parameters — sharded along dim 0 + local_num_v = self.num_v_heads // tp_size + self.A_log = nn.Parameter(torch.zeros(local_num_v)) + self.dt_bias = nn.Parameter(torch.zeros(local_num_v)) + set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)}) + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) + + # Gated RMSNorm on head_v_dim — replicated (head_v_dim=128 is small) self.norm = Qwen3_5RMSNormGated(self.head_v_dim, eps=text_cfg.rms_norm_eps) self.captured_conv_states: Dict[int, torch.Tensor] = {} self.captured_temporal_states: Dict[int, torch.Tensor] = {} - - def _conv1d_weight_loader(self, param: torch.Tensor, - loaded_weight: torch.Tensor) -> None: + + def _conv1d_weight_loader(self, param: torch.Tensor, + loaded_weight: torch.Tensor) -> None: # loaded_weight is ordered as [q, k, v] along its channel dimension. - # Must gather channels in the same non-contiguous pattern that - # MergedColumnParallelLinear uses for in_proj_qkv, so that each rank's - # conv1d_weight[i] applies to the correct in_proj_qkv output channel. - tp_rank = get_tensor_model_parallel_rank() - tp_size = get_tensor_model_parallel_world_size() - key_local = self.key_dim // tp_size # 512 with TP=4 + # Must gather channels in the same non-contiguous pattern that + # MergedColumnParallelLinear uses for in_proj_qkv, so that each rank's + # conv1d_weight[i] applies to the correct in_proj_qkv output channel. + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + key_local = self.key_dim // tp_size # 512 with TP=4 val_local = self.value_dim // tp_size # 1024 with TP=4 - q_s = loaded_weight[tp_rank * key_local : (tp_rank + 1) * key_local] - k_s = loaded_weight[self.key_dim + tp_rank * key_local : - self.key_dim + (tp_rank + 1) * key_local] - v_s = loaded_weight[2 * self.key_dim + tp_rank * val_local : - 2 * self.key_dim + (tp_rank + 1) * val_local] - param.data.copy_(torch.cat([q_s, k_s, v_s], dim=0)) - - def forward( - self, - hidden_states: torch.Tensor, # (total_tokens, hidden_size) - attn_metadata: AttentionMetadata, + q_s = loaded_weight[tp_rank * key_local : (tp_rank + 1) * key_local] + k_s = loaded_weight[self.key_dim + tp_rank * key_local : + self.key_dim + (tp_rank + 1) * key_local] + v_s = loaded_weight[2 * self.key_dim + tp_rank * val_local : + 2 * self.key_dim + (tp_rank + 1) * val_local] + param.data.copy_(torch.cat([q_s, k_s, v_s], dim=0)) + + def forward( + self, + hidden_states: torch.Tensor, # (total_tokens, hidden_size) + attn_metadata: AttentionMetadata, conv_state: torch.Tensor, # (batch, local_conv_dim, kernel-1) in-place temporal_state: torch.Tensor, # (batch, local_v_heads, k_dim, v_dim) in-place capture_offsets: Optional[Iterable[int]] = None, segment_offsets: Optional[Iterable[int]] = None, ) -> torch.Tensor: - tp_size = get_tensor_model_parallel_world_size() - local_key_dim = self.key_dim // tp_size - local_val_dim = self.value_dim // tp_size - local_num_v = self.num_v_heads // tp_size - local_num_k = self.num_k_heads // tp_size + tp_size = get_tensor_model_parallel_world_size() + local_key_dim = self.key_dim // tp_size + local_val_dim = self.value_dim // tp_size + local_num_v = self.num_v_heads // tp_size + local_num_k = self.num_k_heads // tp_size local_conv_dim = self.conv_dim // tp_size self.captured_conv_states = {} self.captured_temporal_states = {} - - is_prefill = attn_metadata.num_prefill_tokens > 0 - + + is_prefill = attn_metadata.num_prefill_tokens > 0 + projected, _ = self.in_proj_qkvzba(hidden_states) mixed_qkv_all, z_all, b_all, a_all = torch.split( projected, [local_conv_dim, local_val_dim, local_num_v, local_num_v], dim=-1, ) - - if is_prefill: - seq_starts = attn_metadata.query_start_loc.tolist() - outputs = [] - state_len = self.conv_kernel_size - 1 - weight_2d = self.conv1d_weight.squeeze(1) # (local_conv_dim, kernel) - + + if is_prefill: + seq_starts = attn_metadata.query_start_loc.tolist() + outputs = [] + state_len = self.conv_kernel_size - 1 + weight_2d = self.conv1d_weight.squeeze(1) # (local_conv_dim, kernel) + for si in range(len(seq_starts) - 1): - s, e = int(seq_starts[si]), int(seq_starts[si + 1]) - seq_len = e - s - - # Shape: (1, local_conv_dim, seq_len) - mixed_qkv = (mixed_qkv_all[s:e] - .transpose(0, 1).unsqueeze(0) - .to(weight_2d.dtype)) - - # Load prev conv state BEFORE overwriting (needed for causal conv padding). - # For first prefill of a request: mamba_cache is zeros → correct. - # For chunked prefill chunk 2+: carries last state_len tokens from prev chunk. - prev_conv = conv_state[si:si + 1].clone().to(weight_2d.dtype) # [1, local_conv_dim, state_len] - - # Save conv state (last state_len positions) - if seq_len >= state_len: - conv_state[si].copy_(mixed_qkv[0, :, -state_len:]) - else: - conv_state[si, :, state_len - seq_len:].copy_( - mixed_qkv[0]) - conv_state[si, :, :state_len - seq_len] = 0 - - # Causal conv: left-pad with previous conv state (not zeros). + s, e = int(seq_starts[si]), int(seq_starts[si + 1]) + seq_len = e - s + + # Shape: (1, local_conv_dim, seq_len) + mixed_qkv = (mixed_qkv_all[s:e] + .transpose(0, 1).unsqueeze(0) + .to(weight_2d.dtype)) + + # Load prev conv state BEFORE overwriting (needed for causal conv padding). + # For first prefill of a request: mamba_cache is zeros → correct. + # For chunked prefill chunk 2+: carries last state_len tokens from prev chunk. + prev_conv = conv_state[si:si + 1].clone().to(weight_2d.dtype) # [1, local_conv_dim, state_len] + + # Save conv state (last state_len positions) + if seq_len >= state_len: + conv_state[si].copy_(mixed_qkv[0, :, -state_len:]) + else: + conv_state[si, :, state_len - seq_len:].copy_( + mixed_qkv[0]) + conv_state[si, :, :state_len - seq_len] = 0 + + # Causal conv: left-pad with previous conv state (not zeros). padded = torch.cat([prev_conv, mixed_qkv], dim=2) seq_capture_offsets = (set(capture_offsets or ()) if si == 0 else set()) @@ -1079,33 +1079,33 @@ class GatedDeltaNet(nn.Module): 0, :, capture_offset: capture_offset + state_len].clone() mixed_qkv_conv = F.conv1d( - padded, self.conv1d_weight, - bias=None, padding=0, groups=local_conv_dim) - mixed_qkv_conv = F.silu(mixed_qkv_conv) - # (1, seq_len, local_conv_dim) - mixed_qkv_conv = mixed_qkv_conv.squeeze(0).transpose(0, 1).unsqueeze(0) - - q, k, v = torch.split( - mixed_qkv_conv, - [local_key_dim, local_key_dim, local_val_dim], dim=-1) - q = q.reshape(1, seq_len, local_num_k, self.head_k_dim) - k = k.reshape(1, seq_len, local_num_k, self.head_k_dim) - v = v.reshape(1, seq_len, local_num_v, self.head_v_dim) - - beta = b_all[s:e].sigmoid().unsqueeze(0) # (1, seq_len, local_num_v) - g = (-self.A_log.float().exp() - * F.softplus(a_all[s:e].float() + self.dt_bias) - ).unsqueeze(0) # (1, seq_len, local_num_v) - - # Expand k/q to match num_v_heads - q = q.repeat_interleave(self.head_expand_ratio, dim=2) - k = k.repeat_interleave(self.head_expand_ratio, dim=2) - - # Sub-sequence chunking: call _torch_chunk_gated_delta_rule - # on _DNN_CHUNK tokens at a time to cap peak memory. - # Full 18K: tensors [1,6,282,64,64]=220 MB each → ~990 MB/call. - # With _DNN_CHUNK=4096: [1,6,64,64,64]=6 MB each → ~137 MB/call. - # State is chained via initial_state / output_final_state. + padded, self.conv1d_weight, + bias=None, padding=0, groups=local_conv_dim) + mixed_qkv_conv = F.silu(mixed_qkv_conv) + # (1, seq_len, local_conv_dim) + mixed_qkv_conv = mixed_qkv_conv.squeeze(0).transpose(0, 1).unsqueeze(0) + + q, k, v = torch.split( + mixed_qkv_conv, + [local_key_dim, local_key_dim, local_val_dim], dim=-1) + q = q.reshape(1, seq_len, local_num_k, self.head_k_dim) + k = k.reshape(1, seq_len, local_num_k, self.head_k_dim) + v = v.reshape(1, seq_len, local_num_v, self.head_v_dim) + + beta = b_all[s:e].sigmoid().unsqueeze(0) # (1, seq_len, local_num_v) + g = (-self.A_log.float().exp() + * F.softplus(a_all[s:e].float() + self.dt_bias) + ).unsqueeze(0) # (1, seq_len, local_num_v) + + # Expand k/q to match num_v_heads + q = q.repeat_interleave(self.head_expand_ratio, dim=2) + k = k.repeat_interleave(self.head_expand_ratio, dim=2) + + # Sub-sequence chunking: call _torch_chunk_gated_delta_rule + # on _DNN_CHUNK tokens at a time to cap peak memory. + # Full 18K: tensors [1,6,282,64,64]=220 MB each → ~990 MB/call. + # With _DNN_CHUNK=4096: [1,6,64,64,64]=6 MB each → ~137 MB/call. + # State is chained via initial_state / output_final_state. cur_state = temporal_state[si:si + 1].clone() core_out_parts = [] segment_ends = _gdn_segment_ends( @@ -1134,14 +1134,14 @@ class GatedDeltaNet(nn.Module): self.captured_temporal_states[sc_end] = ( cur_state[0].clone()) sc_start = sc_end - if cur_state is not None: - temporal_state[si].copy_(cur_state[0]) - # [1, seq_len, num_v_heads, head_v_dim] - core_out = torch.cat(core_out_parts, dim=1) - - # Gate + norm + output proj - z = z_all[s:e].reshape(seq_len, local_num_v, self.head_v_dim) - core_out = core_out.reshape(seq_len, local_num_v, self.head_v_dim) + if cur_state is not None: + temporal_state[si].copy_(cur_state[0]) + # [1, seq_len, num_v_heads, head_v_dim] + core_out = torch.cat(core_out_parts, dim=1) + + # Gate + norm + output proj + z = z_all[s:e].reshape(seq_len, local_num_v, self.head_v_dim) + core_out = core_out.reshape(seq_len, local_num_v, self.head_v_dim) normed = self.norm( core_out.reshape(-1, self.head_v_dim), z.reshape(-1, self.head_v_dim)) @@ -1155,17 +1155,17 @@ class GatedDeltaNet(nn.Module): result = torch.cat(outputs, dim=0) return _check_gdn_finite( result, layer_idx=self.layer_idx, stage="prefill-output") - - else: - # Decode: one token per sequence - num_seqs = hidden_states.shape[0] - weight_2d = self.conv1d_weight.squeeze(1) - - # (num_seqs, local_conv_dim, 1) - mixed_qkv = (mixed_qkv_all - .to(weight_2d.dtype) - .unsqueeze(-1)) - + + else: + # Decode: one token per sequence + num_seqs = hidden_states.shape[0] + weight_2d = self.conv1d_weight.squeeze(1) + + # (num_seqs, local_conv_dim, 1) + mixed_qkv = (mixed_qkv_all + .to(weight_2d.dtype) + .unsqueeze(-1)) + if _USE_COREX_GDN_CAUSAL_CONV: mixed_qkv_conv = _corex_gdn_causal_conv.causal_conv_update( conv_state, mixed_qkv, weight_2d) @@ -1173,9 +1173,9 @@ class GatedDeltaNet(nn.Module): mixed_qkv_conv = _torch_causal_conv1d_update( mixed_qkv, conv_state, weight_2d, bias=None, activation='silu') - # (num_seqs, local_conv_dim, 1) → (num_seqs, 1, local_conv_dim) - mixed_qkv_conv = mixed_qkv_conv.squeeze(-1).unsqueeze(1) - + # (num_seqs, local_conv_dim, 1) → (num_seqs, 1, local_conv_dim) + mixed_qkv_conv = mixed_qkv_conv.squeeze(-1).unsqueeze(1) + packed_mixed_qkv = mixed_qkv_conv.squeeze(1) use_corex_packed_decode = ( _USE_COREX_GDN_PACKED_DECODE @@ -1309,12 +1309,12 @@ class GatedDeltaNet(nn.Module): out, _ = self.out_proj(normed) return _check_gdn_finite( out, layer_idx=self.layer_idx, stage="decode-output") - - -# --------------------------------------------------------------------------- -# Full Attention (with gated q — unique to Qwen3.5) -# --------------------------------------------------------------------------- - + + +# --------------------------------------------------------------------------- +# Full Attention (with gated q — unique to Qwen3.5) +# --------------------------------------------------------------------------- + class Qwen3_5AttentionHeadRMSNorm(GemmaRMSNorm): def forward_cuda( self, @@ -1342,88 +1342,88 @@ class Qwen3_5AttentionHeadRMSNorm(GemmaRMSNorm): class Qwen3_5FullAttention(nn.Module): - def __init__( - self, - text_cfg, - layer_idx: int, - cache_config: Optional[CacheConfig] = None, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - ) -> None: - super().__init__() - self.layer_idx = layer_idx - self.hidden_size = text_cfg.hidden_size # 5120 - self.num_heads = text_cfg.num_attention_heads # 24 - self.num_kv_heads = text_cfg.num_key_value_heads # 4 - self.head_dim = text_cfg.head_dim # 256 - self.rms_norm_eps = text_cfg.rms_norm_eps - + def __init__( + self, + text_cfg, + layer_idx: int, + cache_config: Optional[CacheConfig] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = text_cfg.hidden_size # 5120 + self.num_heads = text_cfg.num_attention_heads # 24 + self.num_kv_heads = text_cfg.num_key_value_heads # 4 + self.head_dim = text_cfg.head_dim # 256 + self.rms_norm_eps = text_cfg.rms_norm_eps + tp_size = get_tensor_model_parallel_world_size() self.local_num_heads = self.num_heads // tp_size self.scaling = self.head_dim ** -0.5 self.use_packed_local_qgkv = tp_size > self.num_kv_heads - - # When num_kv_heads < tp_size we cannot shard KV further (would give - # fractional heads per rank). Use ReplicatedLinear so every rank holds - # all KV heads; local_num_kv_heads equals the full count. - # When num_kv_heads >= tp_size standard ColumnParallel sharding applies. - if tp_size > self.num_kv_heads: - # GQA-aware TP sharding: ixformer kernel only supports num_kv_heads=1 - # per rank. With num_kv_heads=2 < tp_size=4 we cannot shard KV - # evenly, but we CAN assign each rank the ONE KV head that serves - # its Q heads: - # q_per_kv = num_heads // num_kv_heads (e.g. 16//2 = 8) - # Rank r uses KV head r * local_num_heads // q_per_kv - # e.g. ranks 0,1 → KV head 0; ranks 2,3 → KV head 1. - # We replicate all KV heads to every rank and select in forward(). - self.proj_kv_heads = self.num_kv_heads # heads available from projection - self.local_num_kv_heads = 1 # heads after rank-local selection - self.q_per_kv_global = self.num_heads // self.num_kv_heads + + # When num_kv_heads < tp_size we cannot shard KV further (would give + # fractional heads per rank). Use ReplicatedLinear so every rank holds + # all KV heads; local_num_kv_heads equals the full count. + # When num_kv_heads >= tp_size standard ColumnParallel sharding applies. + if tp_size > self.num_kv_heads: + # GQA-aware TP sharding: ixformer kernel only supports num_kv_heads=1 + # per rank. With num_kv_heads=2 < tp_size=4 we cannot shard KV + # evenly, but we CAN assign each rank the ONE KV head that serves + # its Q heads: + # q_per_kv = num_heads // num_kv_heads (e.g. 16//2 = 8) + # Rank r uses KV head r * local_num_heads // q_per_kv + # e.g. ranks 0,1 → KV head 0; ranks 2,3 → KV head 1. + # We replicate all KV heads to every rank and select in forward(). + self.proj_kv_heads = self.num_kv_heads # heads available from projection + self.local_num_kv_heads = 1 # heads after rank-local selection + self.q_per_kv_global = self.num_heads // self.num_kv_heads local_qg_dim = self.local_num_heads * self.head_dim * 2 replicated_kv_dim = self.num_kv_heads * self.head_dim self.qgkv_proj = ReplicatedLinear( self.hidden_size, local_qg_dim + 2 * replicated_kv_dim, bias=False, quant_config=quant_config, prefix=f"{prefix}.qgkv_proj") - else: - # Standard sharding: each rank gets num_kv_heads // tp_size heads. - self.local_num_kv_heads = self.num_kv_heads // tp_size - self.proj_kv_heads = self.local_num_kv_heads # already sharded - self.q_per_kv_global = None - self.k_proj = ColumnParallelLinear( - self.hidden_size, self.num_kv_heads * self.head_dim, - bias=False, quant_config=quant_config, - prefix=f"{prefix}.k_proj") - self.v_proj = ColumnParallelLinear( - self.hidden_size, self.num_kv_heads * self.head_dim, - bias=False, quant_config=quant_config, - prefix=f"{prefix}.v_proj") - - self.local_q_dim = self.local_num_heads * self.head_dim - self.local_kv_dim = self.local_num_kv_heads * self.head_dim - + else: + # Standard sharding: each rank gets num_kv_heads // tp_size heads. + self.local_num_kv_heads = self.num_kv_heads // tp_size + self.proj_kv_heads = self.local_num_kv_heads # already sharded + self.q_per_kv_global = None + self.k_proj = ColumnParallelLinear( + self.hidden_size, self.num_kv_heads * self.head_dim, + bias=False, quant_config=quant_config, + prefix=f"{prefix}.k_proj") + self.v_proj = ColumnParallelLinear( + self.hidden_size, self.num_kv_heads * self.head_dim, + bias=False, quant_config=quant_config, + prefix=f"{prefix}.v_proj") + + self.local_q_dim = self.local_num_heads * self.head_dim + self.local_kv_dim = self.local_num_kv_heads * self.head_dim + if not self.use_packed_local_qgkv: # q_proj includes gate: output = num_heads * head_dim * 2 self.q_proj = ColumnParallelLinear( self.hidden_size, self.num_heads * self.head_dim * 2, bias=False, quant_config=quant_config, prefix=f"{prefix}.q_proj") - self.o_proj = RowParallelLinear( - self.num_heads * self.head_dim, self.hidden_size, - bias=False, quant_config=quant_config, - prefix=f"{prefix}.o_proj") - + self.o_proj = RowParallelLinear( + self.num_heads * self.head_dim, self.hidden_size, + bias=False, quant_config=quant_config, + prefix=f"{prefix}.o_proj") + self.q_norm = Qwen3_5AttentionHeadRMSNorm( self.head_dim, eps=self.rms_norm_eps) self.k_norm = Qwen3_5AttentionHeadRMSNorm( self.head_dim, eps=self.rms_norm_eps) - - # Partial RoPE: rotary_dim = head_dim * partial_rotary_factor = 256 * 0.25 = 64 - rope_params = getattr(text_cfg, "rope_parameters", {}) or {} - rope_theta = rope_params.get("rope_theta", 10_000_000) - partial_factor = rope_params.get("partial_rotary_factor", 0.25) - rotary_dim = int(self.head_dim * partial_factor) - + + # Partial RoPE: rotary_dim = head_dim * partial_rotary_factor = 256 * 0.25 = 64 + rope_params = getattr(text_cfg, "rope_parameters", {}) or {} + rope_theta = rope_params.get("rope_theta", 10_000_000) + partial_factor = rope_params.get("partial_rotary_factor", 0.25) + rotary_dim = int(self.head_dim * partial_factor) + self.rotary_emb = Qwen3_5InterleavedMRotaryEmbedding( head_size=self.head_dim, rotary_dim=rotary_dim, @@ -1433,26 +1433,26 @@ class Qwen3_5FullAttention(nn.Module): dtype=torch.get_default_dtype(), mrope_section=rope_params.get("mrope_section", [11, 11, 10]), ) - - self.attn = Attention( - self.local_num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.local_num_kv_heads, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.attn", - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - kv_cache: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> torch.Tensor: - total_tokens = hidden_states.shape[0] - + + self.attn = Attention( + self.local_num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.local_num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + total_tokens = hidden_states.shape[0] + with bi100_timer("full_attn.project_qgkv"): if self.use_packed_local_qgkv: projected, _ = self.qgkv_proj(hidden_states) @@ -1503,65 +1503,65 @@ class Qwen3_5FullAttention(nn.Module): with bi100_timer("full_attn.output_proj"): output, _ = self.o_proj(attn_out) return output - - -# --------------------------------------------------------------------------- -# MLP (SwiGLU, same as Qwen2/Qwen3) -# --------------------------------------------------------------------------- - -class Qwen3_5MLP(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - hidden_act: str, - quant_config: Optional[QuantizationConfig] = None, - ) -> None: - super().__init__() - self.gate_up_proj = MergedColumnParallelLinear( - hidden_size, [intermediate_size] * 2, - bias=False, quant_config=quant_config) - self.down_proj = RowParallelLinear( - intermediate_size, hidden_size, - bias=False, quant_config=quant_config) - if hidden_act != "silu": - raise ValueError(f"Unsupported activation: {hidden_act}") - self.act_fn = SiluAndMul() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - gate_up, _ = self.gate_up_proj(x) - x = self.act_fn(gate_up) - x, _ = self.down_proj(x) - return x - - -# --------------------------------------------------------------------------- -# MoE sparse block (Qwen3.5-MoE / Qwen3.6-35B-A3B) -# --------------------------------------------------------------------------- - -class Qwen3_5MoeSparseBlock(nn.Module): - """Replaces Qwen3_5MLP for qwen3_5_moe_text layers. - - FusedMoE is used ONLY for weight storage and loading (create_weights / - weight_loader are pure PyTorch). Its forward kernel is bypassed because - ixformer on BI-V100 lacks vllm_moe_topk_softmax / vllm_invoke_fused_moe_kernel. - Routing and expert computation use a pure-PyTorch loop instead. - - Shared expert uses RowParallelLinear(reduce_results=False) so both paths - produce partial (pre-all-reduce) outputs that are combined before a single - all-reduce. - """ - - def __init__( - self, - text_cfg, - quant_config: Optional[QuantizationConfig] = None, - ) -> None: - super().__init__() - hidden_size = text_cfg.hidden_size - self.num_experts = text_cfg.num_experts - self.top_k = text_cfg.num_experts_per_tok - + + +# --------------------------------------------------------------------------- +# MLP (SwiGLU, same as Qwen2/Qwen3) +# --------------------------------------------------------------------------- + +class Qwen3_5MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, [intermediate_size] * 2, + bias=False, quant_config=quant_config) + self.down_proj = RowParallelLinear( + intermediate_size, hidden_size, + bias=False, quant_config=quant_config) + if hidden_act != "silu": + raise ValueError(f"Unsupported activation: {hidden_act}") + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +# --------------------------------------------------------------------------- +# MoE sparse block (Qwen3.5-MoE / Qwen3.6-35B-A3B) +# --------------------------------------------------------------------------- + +class Qwen3_5MoeSparseBlock(nn.Module): + """Replaces Qwen3_5MLP for qwen3_5_moe_text layers. + + FusedMoE is used ONLY for weight storage and loading (create_weights / + weight_loader are pure PyTorch). Its forward kernel is bypassed because + ixformer on BI-V100 lacks vllm_moe_topk_softmax / vllm_invoke_fused_moe_kernel. + Routing and expert computation use a pure-PyTorch loop instead. + + Shared expert uses RowParallelLinear(reduce_results=False) so both paths + produce partial (pre-all-reduce) outputs that are combined before a single + all-reduce. + """ + + def __init__( + self, + text_cfg, + quant_config: Optional[QuantizationConfig] = None, + ) -> None: + super().__init__() + hidden_size = text_cfg.hidden_size + self.num_experts = text_cfg.num_experts + self.top_k = text_cfg.num_experts_per_tok + # Router and scalar shared-expert gate read the same hidden state. Keep # their checkpoint shards in one replicated weight so forward needs a # single GEMM for 256 + 1 outputs. @@ -1570,24 +1570,24 @@ class Qwen3_5MoeSparseBlock(nn.Module): bias=False, quant_config=quant_config) self.router_shared_gate.weight.weight_loader = \ self._router_shared_gate_weight_loader - - # FusedMoE: only used for weight storage + weight_loader. - # Forward is bypassed — see _pure_pytorch_experts(). - self.experts = FusedMoE( - num_experts=text_cfg.num_experts, - top_k=text_cfg.num_experts_per_tok, - hidden_size=hidden_size, - intermediate_size=text_cfg.moe_intermediate_size, - reduce_results=False, # we do the all-reduce ourselves below - renormalize=True, - quant_config=quant_config, - ) - - # Shared expert: defer all-reduce to combine with routed output first - shared_size = text_cfg.shared_expert_intermediate_size - self.shared_expert_gate_up = MergedColumnParallelLinear( - hidden_size, [shared_size] * 2, bias=False, - quant_config=quant_config) + + # FusedMoE: only used for weight storage + weight_loader. + # Forward is bypassed — see _pure_pytorch_experts(). + self.experts = FusedMoE( + num_experts=text_cfg.num_experts, + top_k=text_cfg.num_experts_per_tok, + hidden_size=hidden_size, + intermediate_size=text_cfg.moe_intermediate_size, + reduce_results=False, # we do the all-reduce ourselves below + renormalize=True, + quant_config=quant_config, + ) + + # Shared expert: defer all-reduce to combine with routed output first + shared_size = text_cfg.shared_expert_intermediate_size + self.shared_expert_gate_up = MergedColumnParallelLinear( + hidden_size, [shared_size] * 2, bias=False, + quant_config=quant_config) self.shared_expert_down = RowParallelLinear( shared_size, hidden_size, bias=False, reduce_results=False, quant_config=quant_config) @@ -1614,19 +1614,19 @@ class Qwen3_5MoeSparseBlock(nn.Module): "unexpected router/shared gate weight shape: " f"expected {expected}, got {tuple(loaded_weight.shape)}") param.data.narrow(0, offset, rows).copy_(loaded_weight) - - def _pure_pytorch_experts( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor: - """Pure-PyTorch MoE (ixformer has no MoE kernels on BI-V100). - - w13_weight: (num_experts, 2*inter_per_partition, hidden) [TP-sharded] - w2_weight: (num_experts, hidden, inter_per_partition) [TP-sharded] - Output is partial (pre-all-reduce), same contract as FusedMoE - with reduce_results=False. - """ + + def _pure_pytorch_experts( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + """Pure-PyTorch MoE (ixformer has no MoE kernels on BI-V100). + + w13_weight: (num_experts, 2*inter_per_partition, hidden) [TP-sharded] + w2_weight: (num_experts, hidden, inter_per_partition) [TP-sharded] + Output is partial (pre-all-reduce), same contract as FusedMoE + with reduce_results=False. + """ # Fused topk+softmax: single CUB kernel vs 2 PyTorch ops. # Source: xllm/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh if _USE_COREX_MOE_TOPK_SOFTMAX: @@ -1639,16 +1639,16 @@ class Qwen3_5MoeSparseBlock(nn.Module): router_logits.float(), self.top_k, dim=-1) # (T, top_k) topk_weights = torch.softmax(topk_logits, dim=-1) topk_weights = topk_weights.to(hidden_states.dtype) - - w13 = self.experts.w13_weight # (E, 2*I, H) - w2 = self.experts.w2_weight # (E, H, I) - - T = hidden_states.shape[0] - if T == 1: - # Fast path: single token (decode). - # Batched GEMM: replace top_k separate F.linear calls with 2 fused ops. - # gate_up: 1 large GEMM (1,H) × (K*2*I,H)^T → (1, K*2*I) - # down: 1 bmm (K,H,I) @ (K,I,1) → (K,H) + + w13 = self.experts.w13_weight # (E, 2*I, H) + w2 = self.experts.w2_weight # (E, H, I) + + T = hidden_states.shape[0] + if T == 1: + # Fast path: single token (decode). + # Batched GEMM: replace top_k separate F.linear calls with 2 fused ops. + # gate_up: 1 large GEMM (1,H) × (K*2*I,H)^T → (1, K*2*I) + # down: 1 bmm (K,H,I) @ (K,I,1) → (K,H) # Total: 3 kernel launches vs previous 16 (top_k*2). eids = topk_ids[0] # (K,) ws = topk_weights[0].to(hidden_states.dtype) # (K,) @@ -1695,23 +1695,23 @@ class Qwen3_5MoeSparseBlock(nn.Module): else: w13_sel = w13[eids] # (K, 2*I, H) w2_sel = w2[eids] # (K, H, I) - - H = hidden_states.shape[-1] - - gate_up = F.linear( - hidden_states, - w13_sel.reshape(-1, H), # (K*2*I, H) — contiguous after indexing - ) # (1, K*2*I) - gate_up = gate_up.view(self.top_k, -1) # (K, 2*I) + + H = hidden_states.shape[-1] + + gate_up = F.linear( + hidden_states, + w13_sel.reshape(-1, H), # (K*2*I, H) — contiguous after indexing + ) # (1, K*2*I) + gate_up = gate_up.view(self.top_k, -1) # (K, 2*I) if _USE_FUSED_MOE_ACTIVATION: act = self.act_fn(gate_up) # (K, I) else: gate, up = gate_up.chunk(2, dim=-1) act = F.silu(gate) * up - - # bmm: (K,H,I) @ (K,I,1) → (K,H,1) → (K,H) - expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H) - + + # bmm: (K,H,I) @ (K,I,1) → (K,H,1) → (K,H) + expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H) + if (_USE_COREX_MOE_EXACT_REDUCE and expert_out.dtype == torch.float16 and ws.dtype == torch.float16 @@ -1759,9 +1759,9 @@ class Qwen3_5MoeSparseBlock(nn.Module): weights = sorted_weights[start:end].unsqueeze(-1) out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype)) start = end - - return out # partial, all-reduce done in forward() - + + return out # partial, all-reduce done in forward() + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: with bi100_timer("moe.router"): router_and_shared_gate, _ = self.router_shared_gate(hidden_states) @@ -1782,22 +1782,22 @@ class Qwen3_5MoeSparseBlock(nn.Module): with bi100_timer("moe.all_reduce"): out = tensor_model_parallel_all_reduce(out) return out - - -# --------------------------------------------------------------------------- -# Decoder layer (dispatches to GatedDeltaNet or Qwen3_5FullAttention) -# --------------------------------------------------------------------------- - - -class Qwen3_5DecoderLayer(nn.Module): - def __init__( - self, - text_cfg, - layer_idx: int, - layer_type: str, - cache_config: Optional[CacheConfig] = None, - quant_config: Optional[QuantizationConfig] = None, - ) -> None: + + +# --------------------------------------------------------------------------- +# Decoder layer (dispatches to GatedDeltaNet or Qwen3_5FullAttention) +# --------------------------------------------------------------------------- + + +class Qwen3_5DecoderLayer(nn.Module): + def __init__( + self, + text_cfg, + layer_idx: int, + layer_type: str, + cache_config: Optional[CacheConfig] = None, + quant_config: Optional[QuantizationConfig] = None, + ) -> None: super().__init__() self.layer_idx = layer_idx self.layer_type = layer_type @@ -1805,39 +1805,39 @@ class Qwen3_5DecoderLayer(nn.Module): os.getenv("BI100_DIAGNOSTIC_LAYER_TRACE") == "1") self.input_layernorm = GemmaRMSNorm(text_cfg.hidden_size, eps=text_cfg.rms_norm_eps) - self.post_attention_layernorm = GemmaRMSNorm(text_cfg.hidden_size, - eps=text_cfg.rms_norm_eps) - - if layer_type == "linear_attention": - self.linear_attn = GatedDeltaNet(text_cfg, layer_idx, - quant_config=quant_config) - else: - self.self_attn = Qwen3_5FullAttention( - text_cfg, layer_idx, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"layers.{layer_idx}.self_attn", - ) - - if getattr(text_cfg, 'model_type', '') == 'qwen3_5_moe_text': - self.mlp = Qwen3_5MoeSparseBlock(text_cfg, quant_config=quant_config) - else: - self.mlp = Qwen3_5MLP( - hidden_size=text_cfg.hidden_size, - intermediate_size=text_cfg.intermediate_size, - hidden_act=text_cfg.hidden_act, - quant_config=quant_config, - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - kv_cache: Optional[torch.Tensor], - attn_metadata: AttentionMetadata, - residual: Optional[torch.Tensor], - # Only for linear_attention layers: - conv_state: Optional[torch.Tensor] = None, + self.post_attention_layernorm = GemmaRMSNorm(text_cfg.hidden_size, + eps=text_cfg.rms_norm_eps) + + if layer_type == "linear_attention": + self.linear_attn = GatedDeltaNet(text_cfg, layer_idx, + quant_config=quant_config) + else: + self.self_attn = Qwen3_5FullAttention( + text_cfg, layer_idx, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"layers.{layer_idx}.self_attn", + ) + + if getattr(text_cfg, 'model_type', '') == 'qwen3_5_moe_text': + self.mlp = Qwen3_5MoeSparseBlock(text_cfg, quant_config=quant_config) + else: + self.mlp = Qwen3_5MLP( + hidden_size=text_cfg.hidden_size, + intermediate_size=text_cfg.intermediate_size, + hidden_act=text_cfg.hidden_act, + quant_config=quant_config, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_cache: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + residual: Optional[torch.Tensor], + # Only for linear_attention layers: + conv_state: Optional[torch.Tensor] = None, temporal_state: Optional[torch.Tensor] = None, gdn_capture_offsets: Optional[Iterable[int]] = None, gdn_segment_offsets: Optional[Iterable[int]] = None, @@ -1881,9 +1881,9 @@ class Qwen3_5DecoderLayer(nn.Module): ) return hidden_states, residual - - -# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- # Full transformer model # --------------------------------------------------------------------------- @@ -1918,14 +1918,14 @@ class Qwen3_5Model(nn.Module): self.kv_cache_count = kv_cache_count self.embed_tokens = VocabParallelEmbedding( text_cfg.vocab_size, text_cfg.hidden_size) - self.layers = nn.ModuleList([ - Qwen3_5DecoderLayer( - text_cfg, i, text_cfg.layer_types[i], - cache_config=cache_config, quant_config=quant_config) - for i in range(text_cfg.num_hidden_layers) - ]) - self.norm = GemmaRMSNorm(text_cfg.hidden_size, eps=text_cfg.rms_norm_eps) - + self.layers = nn.ModuleList([ + Qwen3_5DecoderLayer( + text_cfg, i, text_cfg.layer_types[i], + cache_config=cache_config, quant_config=quant_config) + for i in range(text_cfg.num_hidden_layers) + ]) + self.norm = GemmaRMSNorm(text_cfg.hidden_size, eps=text_cfg.rms_norm_eps) + def forward( self, input_ids: torch.Tensor, @@ -1942,9 +1942,9 @@ class Qwen3_5Model(nn.Module): with bi100_timer("model.embed"): hidden_states = (self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds) - residual = None - - attn_idx = 0 + residual = None + + attn_idx = 0 linear_idx = 0 capture_offsets = tuple(gdn_capture_offsets or ()) captured_conv_states: Dict[int, List[torch.Tensor]] = { @@ -1953,13 +1953,13 @@ class Qwen3_5Model(nn.Module): captured_temporal_states: Dict[int, List[torch.Tensor]] = { offset: [] for offset in capture_offsets } - for layer in self.layers: - if layer.layer_type == "linear_attention": - hidden_states, residual = layer( - positions, hidden_states, - kv_cache=None, - attn_metadata=attn_metadata, - residual=residual, + for layer in self.layers: + if layer.layer_type == "linear_attention": + hidden_states, residual = layer( + positions, hidden_states, + kv_cache=None, + attn_metadata=attn_metadata, + residual=residual, conv_state=conv_states[linear_idx], temporal_state=temporal_states[linear_idx], gdn_capture_offsets=capture_offsets, @@ -1970,14 +1970,14 @@ class Qwen3_5Model(nn.Module): layer.linear_attn.captured_conv_states[offset]) captured_temporal_states[offset].append( layer.linear_attn.captured_temporal_states[offset]) - linear_idx += 1 - else: - kv_cache = kv_caches[attn_idx] - hidden_states, residual = layer( - positions, hidden_states, - kv_cache=kv_cache, - attn_metadata=attn_metadata, - residual=residual, + linear_idx += 1 + else: + kv_cache = kv_caches[attn_idx] + hidden_states, residual = layer( + positions, hidden_states, + kv_cache=kv_cache, + attn_metadata=attn_metadata, + residual=residual, ) attn_idx += 1 @@ -1992,42 +1992,42 @@ class Qwen3_5Model(nn.Module): for offset, states in captured_temporal_states.items() } return hidden_states - - -# --------------------------------------------------------------------------- -# Top-level CausalLM wrapper with MambaCacheManager -# --------------------------------------------------------------------------- - + + +# --------------------------------------------------------------------------- +# Top-level CausalLM wrapper with MambaCacheManager +# --------------------------------------------------------------------------- + class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsMultiModal): - - has_inner_state = True - supports_lora = True - - packed_modules_mapping = { - "gate_up_proj": ["gate_proj", "up_proj"], - } - - supported_lora_modules = [ - "gate_up_proj", - "down_proj", - "o_proj", - ] - embedding_modules = {} - embedding_padding_modules = [] - - def __init__( - self, - config, # Qwen3_5Config (top-level) - cache_config: Optional[CacheConfig] = None, - quant_config: Optional[QuantizationConfig] = None, + + has_inner_state = True + supports_lora = True + + packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"], + } + + supported_lora_modules = [ + "gate_up_proj", + "down_proj", + "o_proj", + ] + embedding_modules = {} + embedding_padding_modules = [] + + def __init__( + self, + config, # Qwen3_5Config (top-level) + cache_config: Optional[CacheConfig] = None, + quant_config: Optional[QuantizationConfig] = None, lora_config: Optional[LoRAConfig] = None, scheduler_config: Optional[SchedulerConfig] = None, multimodal_config: Optional[MultiModalConfig] = None, prefix: str = "", - ) -> None: + ) -> None: _bi100_model_trace("Qwen3_5ForCausalLM initialization begin") - super().__init__() + super().__init__() self.config = config self.scheduler_config = scheduler_config self.multimodal_config = multimodal_config @@ -2042,10 +2042,10 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, "type": "mrope", "mrope_section": mrope_sections, } - - # Pre-compute counts - self.num_linear_layers = sum( - 1 for lt in text_cfg.layer_types if lt == "linear_attention") + + # Pre-compute counts + self.num_linear_layers = sum( + 1 for lt in text_cfg.layer_types if lt == "linear_attention") self.num_attn_layers = sum( 1 for lt in text_cfg.layer_types if lt == "full_attention") layers_block_type = getattr( @@ -2079,14 +2079,14 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, ) # DeltaNet state dimensions (per layer, per sequence, TP-sharded) - tp_size = get_tensor_model_parallel_world_size() - self.conv_dim = (text_cfg.linear_num_key_heads * text_cfg.linear_key_head_dim * 2 - + text_cfg.linear_num_value_heads * text_cfg.linear_value_head_dim) - self.num_v_heads = text_cfg.linear_num_value_heads - self.head_k_dim = text_cfg.linear_key_head_dim - self.head_v_dim = text_cfg.linear_value_head_dim - self.conv_kernel_size = text_cfg.linear_conv_kernel_dim - + tp_size = get_tensor_model_parallel_world_size() + self.conv_dim = (text_cfg.linear_num_key_heads * text_cfg.linear_key_head_dim * 2 + + text_cfg.linear_num_value_heads * text_cfg.linear_value_head_dim) + self.num_v_heads = text_cfg.linear_num_value_heads + self.head_k_dim = text_cfg.linear_key_head_dim + self.head_v_dim = text_cfg.linear_value_head_dim + self.conv_kernel_size = text_cfg.linear_conv_kernel_dim + self.model = Qwen3_5Model( text_cfg, cache_config=cache_config, @@ -2098,33 +2098,33 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, config.vision_config, quant_config=None, ) - - self.lm_head = ParallelLMHead( - text_cfg.vocab_size, text_cfg.hidden_size, - quant_config=quant_config, - ) - - self.logits_processor = LogitsProcessor(text_cfg.vocab_size) - self.sampler = Sampler() - - # Lazy initialised in first forward call - self.mamba_cache: Optional[MambaCacheManager] = None - + + self.lm_head = ParallelLMHead( + text_cfg.vocab_size, text_cfg.hidden_size, + quant_config=quant_config, + ) + + self.logits_processor = LogitsProcessor(text_cfg.vocab_size) + self.sampler = Sampler() + + # Lazy initialised in first forward call + self.mamba_cache: Optional[MambaCacheManager] = None + # Scheduler-owned recurrent prefix states. Keys are stable chained # content hashes, never recyclable physical KV block ids. self._gdn_prefix_cache: Dict[ Tuple[int, bytes], Tuple[torch.Tensor, torch.Tensor]] = {} - self._block_size: int = (cache_config.block_size - if cache_config is not None else 16) + self._block_size: int = (cache_config.block_size + if cache_config is not None else 16) self._startup_forward_traced = False _bi100_model_trace("Qwen3_5ForCausalLM initialization complete") - + def _get_mamba_cache_shape(self): - tp_size = get_tensor_model_parallel_world_size() - # Each sequence's state is stored in float32 - conv_state_shape = (self.conv_dim // tp_size, self.conv_kernel_size - 1) - temporal_state_shape = ( - self.num_v_heads // tp_size, self.head_k_dim, self.head_v_dim) + tp_size = get_tensor_model_parallel_world_size() + # Each sequence's state is stored in float32 + conv_state_shape = (self.conv_dim // tp_size, self.conv_kernel_size - 1) + temporal_state_shape = ( + self.num_v_heads // tp_size, self.head_k_dim, self.head_v_dim) return conv_state_shape, temporal_state_shape @staticmethod @@ -2180,33 +2180,33 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, image_input["data"], grid_thw=image_input["image_grid_thw"], ) - + @bi100_profile_transaction def forward( self, input_ids: torch.Tensor, - positions: torch.Tensor, - kv_caches: List[torch.Tensor], - attn_metadata: AttentionMetadata, - intermediate_tensors: Optional[IntermediateTensors] = None, - **kwargs, - ) -> torch.Tensor: + positions: torch.Tensor, + kv_caches: List[torch.Tensor], + attn_metadata: AttentionMetadata, + intermediate_tensors: Optional[IntermediateTensors] = None, + **kwargs, + ) -> torch.Tensor: if not self._startup_forward_traced: self._startup_forward_traced = True _bi100_model_trace("first model forward entered") - if self.mamba_cache is None: - if self.scheduler_config is not None: - max_batch_size = _get_graph_batch_size( - self.scheduler_config.max_num_seqs) - else: - max_batch_size = max(_BATCH_SIZES_TO_CAPTURE) + 2 - self.mamba_cache = MambaCacheManager( - torch.float32, - self.num_linear_layers, - max_batch_size, - *self._get_mamba_cache_shape(), - ) - + if self.mamba_cache is None: + if self.scheduler_config is not None: + max_batch_size = _get_graph_batch_size( + self.scheduler_config.max_num_seqs) + else: + max_batch_size = max(_BATCH_SIZES_TO_CAPTURE) + 2 + self.mamba_cache = MambaCacheManager( + torch.float32, + self.num_linear_layers, + max_batch_size, + *self._get_mamba_cache_shape(), + ) + gdn_restore_key = kwargs.pop("gdn_restore_key", None) gdn_capture_points = kwargs.pop("gdn_capture_points", None) or [] gdn_evict_keys = kwargs.pop("gdn_evict_keys", None) or [] @@ -2214,10 +2214,10 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, mamba_tensors = self.mamba_cache.current_run_tensors( input_ids, attn_metadata, **kwargs) - # conv_states: (num_linear_layers, batch, local_conv_dim, kernel-1) - # temporal_states: (num_linear_layers, batch, local_num_v, k_dim, v_dim) - conv_states, temporal_states = mamba_tensors - + # conv_states: (num_linear_layers, batch, local_conv_dim, kernel-1) + # temporal_states: (num_linear_layers, batch, local_num_v, k_dim, v_dim) + conv_states, temporal_states = mamba_tensors + _is_single_seq_prefill = ( attn_metadata is not None and attn_metadata.num_prefill_tokens > 0 @@ -2324,7 +2324,7 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, captured_conv.detach().cpu().clone(), captured_temporal.detach().cpu().clone(), ) - + if bi100_profile_event_enabled(): profile_prefill_tokens = int( getattr(attn_metadata, "num_prefill_tokens", 0) or 0) @@ -2355,57 +2355,57 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, gdn_evict_keys=len(gdn_evict_keys), ) return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - sampling_metadata: SamplingMetadata, - ) -> Optional[torch.Tensor]: - # All TP ranks must call logits_processor to participate in the NCCL - # gather inside lm_head. Non-driver ranks return None after the gather. - # With chunked prefill, intermediate chunks have seq_groups=None on all - # ranks; _apply_logits_processors is guarded against this in - # logits_processor.py (patched by patch_xformers_sdpa_seq.py). - logits = self.logits_processor(self.lm_head, hidden_states, - sampling_metadata) - return logits - - def sample( - self, - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, - ) -> Optional[SamplerOutput]: - return self.sampler(logits, sampling_metadata) - - def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): - return self.mamba_cache.copy_inputs_before_cuda_graphs( - input_buffers, **kwargs) - - def get_seqlen_agnostic_capture_inputs(self, batch_size: int): - return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size) - - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + + def compute_logits( + self, + hidden_states: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[torch.Tensor]: + # All TP ranks must call logits_processor to participate in the NCCL + # gather inside lm_head. Non-driver ranks return None after the gather. + # With chunked prefill, intermediate chunks have seq_groups=None on all + # ranks; _apply_logits_processors is guarded against this in + # logits_processor.py (patched by patch_xformers_sdpa_seq.py). + logits = self.logits_processor(self.lm_head, hidden_states, + sampling_metadata) + return logits + + def sample( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[SamplerOutput]: + return self.sampler(logits, sampling_metadata) + + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): + return self.mamba_cache.copy_inputs_before_cuda_graphs( + input_buffers, **kwargs) + + def get_seqlen_agnostic_capture_inputs(self, batch_size: int): + return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): _bi100_model_trace("dense load_weights begin") loaded_count = 0 - stacked_params_mapping = [ - # (param_name, weight_name, shard_id) - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - params_dict = dict(self.named_parameters()) - - for name, loaded_weight in weights: + stacked_params_mapping = [ + # (param_name, weight_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + + for name, loaded_weight in weights: loaded_count += 1 - # Skip vision and MTP branches - if (name.startswith("model.visual") - or name.startswith("mtp.") - or name.startswith("model.mtp")): - continue - - # Prefix remapping: checkpoint may wrap under language_model - if name.startswith("model.language_model."): - name = "model." + name[len("model.language_model."):] - + # Skip vision and MTP branches + if (name.startswith("model.visual") + or name.startswith("mtp.") + or name.startswith("model.mtp")): + continue + + # Prefix remapping: checkpoint may wrap under language_model + if name.startswith("model.language_model."): + name = "model." + name[len("model.language_model."):] + # Skip positional embedding caches if "rotary_emb.inv_freq" in name: continue @@ -2419,80 +2419,80 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, continue # Remap conv1d.weight → conv1d_weight - # The conv has depth (1) dim in the checkpoint that we handle separately - if ".linear_attn.conv1d.weight" in name: - name = name.replace(".linear_attn.conv1d.weight", - ".linear_attn.conv1d_weight") - - # Stacked param loading (gate_up_proj) - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - if name.endswith(".bias") and name not in params_dict: - break - if name not in params_dict: - break - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - if name.endswith(".bias") and name not in params_dict: - continue - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", - default_weight_loader) - weight_loader(param, loaded_weight) + # The conv has depth (1) dim in the checkpoint that we handle separately + if ".linear_attn.conv1d.weight" in name: + name = name.replace(".linear_attn.conv1d.weight", + ".linear_attn.conv1d_weight") + + # Stacked param loading (gate_up_proj) + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + if name.endswith(".bias") and name not in params_dict: + break + if name not in params_dict: + break + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, loaded_weight) _bi100_model_trace(f"dense load_weights complete items={loaded_count}") - - -# --------------------------------------------------------------------------- -# Qwen3.6-35B-A3B (Qwen3_5-MoE architecture) -# --------------------------------------------------------------------------- - + + +# --------------------------------------------------------------------------- +# Qwen3.6-35B-A3B (Qwen3_5-MoE architecture) +# --------------------------------------------------------------------------- + @MULTIMODAL_REGISTRY.register_image_input_mapper(qwen36_image_input_mapper) @MULTIMODAL_REGISTRY.register_max_image_tokens(get_max_qwen36_image_tokens) @INPUT_REGISTRY.register_dummy_data(dummy_data_for_qwen36) @INPUT_REGISTRY.register_input_processor(input_processor_for_qwen36) class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM): - """Qwen3.6-35B-A3B: same hybrid-attention backbone as 27B, dense MLP - replaced by Qwen3_5MoeSparseBlock (256 routed experts + shared expert). - Only load_weights differs from the dense variant. - """ - + """Qwen3.6-35B-A3B: same hybrid-attention backbone as 27B, dense MLP + replaced by Qwen3_5MoeSparseBlock (256 routed experts + shared expert). + Only load_weights differs from the dense variant. + """ + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): _bi100_model_trace("MoE load_weights begin") loaded_count = 0 vision_loaded_count = 0 - # Checkpoint key format for this model (transformers Qwen3_5MoeExperts): - # mlp.experts.gate_up_proj shape (num_experts, 2*intermediate, hidden) - # mlp.experts.down_proj shape (num_experts, hidden, intermediate) + # Checkpoint key format for this model (transformers Qwen3_5MoeExperts): + # mlp.experts.gate_up_proj shape (num_experts, 2*intermediate, hidden) + # mlp.experts.down_proj shape (num_experts, hidden, intermediate) # mlp.gate.weight shape (num_experts, hidden) [router] # mlp.shared_expert_gate.weight shape (1, hidden) - # mlp.shared_expert.{gate,up,down}_proj.weight [shared MLP] - # Our FusedMoE stores: - # mlp.experts.w13_weight shape (num_experts, 2*intermediate//tp, hidden) - # mlp.experts.w2_weight shape (num_experts, hidden, intermediate//tp) + # mlp.shared_expert.{gate,up,down}_proj.weight [shared MLP] + # Our FusedMoE stores: + # mlp.experts.w13_weight shape (num_experts, 2*intermediate//tp, hidden) + # mlp.experts.w2_weight shape (num_experts, hidden, intermediate//tp) # Our router/shared gate stores both tensors in one (num_experts+1, H) # replicated weight. Our shared expert stores: - # mlp.shared_expert_gate_up.weight (merged gate+up) - # mlp.shared_expert_down.weight - - stacked_params_mapping = [ - # (param_name, weight_name, shard_id) - # shared expert - ("shared_expert_gate_up", "shared_expert.gate_proj", 0), - ("shared_expert_gate_up", "shared_expert.up_proj", 1), - # linear_attention dense proj (same as 27B) - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - - params_dict = dict(self.named_parameters()) - + # mlp.shared_expert_gate_up.weight (merged gate+up) + # mlp.shared_expert_down.weight + + stacked_params_mapping = [ + # (param_name, weight_name, shard_id) + # shared expert + ("shared_expert_gate_up", "shared_expert.gate_proj", 0), + ("shared_expert_gate_up", "shared_expert.up_proj", 1), + # linear_attention dense proj (same as 27B) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: loaded_count += 1 if name.startswith("model.visual."): @@ -2525,14 +2525,14 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM): if (name.startswith("mtp.") or name.startswith("model.mtp")): continue - - # Prefix remapping for VL checkpoint (Qwen3_5MoeForConditionalGeneration): - # model.language_model.model.{layers,embed_tokens,norm} -> model.{...} - # model.language_model.lm_head -> lm_head - # Prefix remapping: checkpoint may wrap under language_model - if name.startswith("model.language_model."): - name = "model." + name[len("model.language_model."):] - + + # Prefix remapping for VL checkpoint (Qwen3_5MoeForConditionalGeneration): + # model.language_model.model.{layers,embed_tokens,norm} -> model.{...} + # model.language_model.lm_head -> lm_head + # Prefix remapping: checkpoint may wrap under language_model + if name.startswith("model.language_model."): + name = "model." + name[len("model.language_model."):] + if "rotary_emb.inv_freq" in name: continue @@ -2565,95 +2565,95 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM): continue if ".linear_attn.conv1d.weight" in name: - name = name.replace(".linear_attn.conv1d.weight", - ".linear_attn.conv1d_weight") - - # --- Fused routed-expert weights (all experts in one tensor) --- - - if "mlp.experts.gate_up_proj" in name: - # loaded_weight: (num_experts, 2*intermediate, hidden) - w13_name = name.replace("mlp.experts.gate_up_proj", - "mlp.experts.w13_weight") - if w13_name not in params_dict: - continue - param = params_dict[w13_name] - n_exp = loaded_weight.shape[0] - inter = loaded_weight.shape[1] // 2 - gate_w = loaded_weight[:, :inter, :].contiguous() - up_w = loaded_weight[:, inter:, :].contiguous() - for eid in range(n_exp): - param.weight_loader(param, gate_w[eid], "w1_weight", "w1", eid) - param.weight_loader(param, up_w[eid], "w3_weight", "w3", eid) - continue - - if "mlp.experts.down_proj" in name: - # loaded_weight: (num_experts, hidden, intermediate) - w2_name = name.replace("mlp.experts.down_proj", - "mlp.experts.w2_weight") - if w2_name not in params_dict: - continue - param = params_dict[w2_name] - n_exp = loaded_weight.shape[0] - for eid in range(n_exp): - param.weight_loader(param, loaded_weight[eid], "w2_weight", "w2", eid) - continue - - # --- Shared expert down_proj rename --- - if "mlp.shared_expert.down_proj" in name: - name = name.replace("mlp.shared_expert.down_proj", - "mlp.shared_expert_down") - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - continue - - # --- Individual expert weights (FT checkpoint: experts.{i}.{proj}.weight) --- - # Standard transformers fine-tuning saves each expert separately instead of - # the pre-merged (num_experts, ...) tensors in the original checkpoint. - if ".mlp.experts." in name: - parts = name.split(".mlp.experts.", 1) - expert_rest = parts[1] # e.g. "0.gate_proj.weight" - dot_pos = expert_rest.find(".") - if dot_pos > 0 and expert_rest[:dot_pos].isdigit(): - eid = int(expert_rest[:dot_pos]) - proj_raw = expert_rest[dot_pos + 1:] - proj = proj_raw[:-7] if proj_raw.endswith(".weight") else proj_raw - prefix = parts[0] # e.g. "model.layers.0" - if proj == "gate_proj": - w13_name = f"{prefix}.mlp.experts.w13_weight" - if w13_name in params_dict: - param = params_dict[w13_name] - param.weight_loader(param, loaded_weight, "w1_weight", "w1", eid) - elif proj == "up_proj": - w13_name = f"{prefix}.mlp.experts.w13_weight" - if w13_name in params_dict: - param = params_dict[w13_name] - param.weight_loader(param, loaded_weight, "w3_weight", "w3", eid) - elif proj == "down_proj": - w2_name = f"{prefix}.mlp.experts.w2_weight" - if w2_name in params_dict: - param = params_dict[w2_name] - param.weight_loader(param, loaded_weight, "w2_weight", "w2", eid) - continue - - # --- Stacked / standard weights --- - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - if name not in params_dict: - break - param = params_dict[name] - param.weight_loader(param, loaded_weight, shard_id) - break - else: - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) + name = name.replace(".linear_attn.conv1d.weight", + ".linear_attn.conv1d_weight") + + # --- Fused routed-expert weights (all experts in one tensor) --- + + if "mlp.experts.gate_up_proj" in name: + # loaded_weight: (num_experts, 2*intermediate, hidden) + w13_name = name.replace("mlp.experts.gate_up_proj", + "mlp.experts.w13_weight") + if w13_name not in params_dict: + continue + param = params_dict[w13_name] + n_exp = loaded_weight.shape[0] + inter = loaded_weight.shape[1] // 2 + gate_w = loaded_weight[:, :inter, :].contiguous() + up_w = loaded_weight[:, inter:, :].contiguous() + for eid in range(n_exp): + param.weight_loader(param, gate_w[eid], "w1_weight", "w1", eid) + param.weight_loader(param, up_w[eid], "w3_weight", "w3", eid) + continue + + if "mlp.experts.down_proj" in name: + # loaded_weight: (num_experts, hidden, intermediate) + w2_name = name.replace("mlp.experts.down_proj", + "mlp.experts.w2_weight") + if w2_name not in params_dict: + continue + param = params_dict[w2_name] + n_exp = loaded_weight.shape[0] + for eid in range(n_exp): + param.weight_loader(param, loaded_weight[eid], "w2_weight", "w2", eid) + continue + + # --- Shared expert down_proj rename --- + if "mlp.shared_expert.down_proj" in name: + name = name.replace("mlp.shared_expert.down_proj", + "mlp.shared_expert_down") + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + continue + + # --- Individual expert weights (FT checkpoint: experts.{i}.{proj}.weight) --- + # Standard transformers fine-tuning saves each expert separately instead of + # the pre-merged (num_experts, ...) tensors in the original checkpoint. + if ".mlp.experts." in name: + parts = name.split(".mlp.experts.", 1) + expert_rest = parts[1] # e.g. "0.gate_proj.weight" + dot_pos = expert_rest.find(".") + if dot_pos > 0 and expert_rest[:dot_pos].isdigit(): + eid = int(expert_rest[:dot_pos]) + proj_raw = expert_rest[dot_pos + 1:] + proj = proj_raw[:-7] if proj_raw.endswith(".weight") else proj_raw + prefix = parts[0] # e.g. "model.layers.0" + if proj == "gate_proj": + w13_name = f"{prefix}.mlp.experts.w13_weight" + if w13_name in params_dict: + param = params_dict[w13_name] + param.weight_loader(param, loaded_weight, "w1_weight", "w1", eid) + elif proj == "up_proj": + w13_name = f"{prefix}.mlp.experts.w13_weight" + if w13_name in params_dict: + param = params_dict[w13_name] + param.weight_loader(param, loaded_weight, "w3_weight", "w3", eid) + elif proj == "down_proj": + w2_name = f"{prefix}.mlp.experts.w2_weight" + if w2_name in params_dict: + param = params_dict[w2_name] + param.weight_loader(param, loaded_weight, "w2_weight", "w2", eid) + continue + + # --- Stacked / standard weights --- + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + if name not in params_dict: + break + param = params_dict[name] + param.weight_loader(param, loaded_weight, shard_id) + break + else: + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) _bi100_model_trace( f"MoE load_weights complete items={loaded_count} " f"vision_items={vision_loaded_count}") diff --git a/qwen3_6_scripts/qwen3_5/__init__.py b/qwen3_6_scripts/qwen3_5/__init__.py index 6988a594..168e15ff 100644 --- a/qwen3_6_scripts/qwen3_5/__init__.py +++ b/qwen3_6_scripts/qwen3_5/__init__.py @@ -1,3 +1,3 @@ -from .configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig, Qwen3_5VisionConfig - -__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5VisionConfig"] +from .configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig, Qwen3_5VisionConfig + +__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5VisionConfig"] diff --git a/qwen3_6_scripts/qwen3_5/configuration_qwen3_5.py b/qwen3_6_scripts/qwen3_5/configuration_qwen3_5.py index a12d2464..2820db3c 100644 --- a/qwen3_6_scripts/qwen3_5/configuration_qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5/configuration_qwen3_5.py @@ -1,20 +1,20 @@ -# Adapted from transformers 5.2.0 for compatibility with transformers 4.55.3 + torch 2.1.0 -# Stubs layer_type_validation and RopeParameters which do not exist in 4.55.3 - +# Adapted from transformers 5.2.0 for compatibility with transformers 4.55.3 + torch 2.1.0 +# Stubs layer_type_validation and RopeParameters which do not exist in 4.55.3 + import os from typing import Optional, List - -from ...configuration_utils import PretrainedConfig as PreTrainedConfig - -# --- Local stubs for APIs not present in transformers 4.55.3 --- -# Always use these definitions; do NOT import from the older transformers -# as same-named functions there have incompatible signatures. - + +from ...configuration_utils import PretrainedConfig as PreTrainedConfig + +# --- Local stubs for APIs not present in transformers 4.55.3 --- +# Always use these definitions; do NOT import from the older transformers +# as same-named functions there have incompatible signatures. + def layer_type_validation(layer_types, num_hidden_layers=None, attention=True): - allowed = {"full_attention", "linear_attention"} - if not all(lt in allowed for lt in layer_types): - raise ValueError(f"layer_types entries must be in {allowed}, got {layer_types}") - if num_hidden_layers is not None and num_hidden_layers != len(layer_types): + allowed = {"full_attention", "linear_attention"} + if not all(lt in allowed for lt in layer_types): + raise ValueError(f"layer_types entries must be in {allowed}, got {layer_types}") + if num_hidden_layers is not None and num_hidden_layers != len(layer_types): raise ValueError( f"num_hidden_layers ({num_hidden_layers}) != len(layer_types) ({len(layer_types)})" ) @@ -57,7 +57,7 @@ def _vllm_layers_block_type( "attention" if layer_type == "full_attention" else layer_type for layer_type in layer_types ] - + try: from typing import TypedDict except ImportError: @@ -68,161 +68,161 @@ else: rope_type: str partial_rotary_factor: float factor: float - -# --- End stubs --- - - -class Qwen3_5TextConfig(PreTrainedConfig): - r""" + +# --- End stubs --- + + +class Qwen3_5TextConfig(PreTrainedConfig): + r""" Configuration for the text backbone of Qwen3.5 / Qwen3.6-35B-A3B models. - model_type is "qwen3_5_text" (used internally by the nested config). - """ - - model_type = "qwen3_5_text" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=248320, - hidden_size=4096, - intermediate_size=12288, - num_hidden_layers=32, - num_attention_heads=16, - num_key_value_heads=4, - hidden_act="silu", - max_position_embeddings=32768, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - tie_word_embeddings=False, - rope_parameters=None, - attention_bias=False, - attention_dropout=0.0, - head_dim=256, - linear_conv_kernel_dim=4, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=32, - layer_types=None, - pad_token_id=None, - bos_token_id=None, - eos_token_id=None, - **kwargs, - ): - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id - self.tie_word_embeddings = tie_word_embeddings - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.head_dim = head_dim - self.rope_parameters = rope_parameters - kwargs.setdefault("partial_rotary_factor", 0.25) - - self.layer_types = layer_types - if self.layer_types is None: - interval_pattern = kwargs.get("full_attention_interval", 4) - self.layer_types = [ - "linear_attention" if bool((i + 1) % interval_pattern) else "full_attention" - for i in range(self.num_hidden_layers) - ] - layer_type_validation(self.layer_types, self.num_hidden_layers) - - self.linear_conv_kernel_dim = linear_conv_kernel_dim - self.linear_key_head_dim = linear_key_head_dim - self.linear_value_head_dim = linear_value_head_dim - self.linear_num_key_heads = linear_num_key_heads - self.linear_num_value_heads = linear_num_value_heads - super().__init__(**kwargs) - - -class Qwen3_5VisionConfig(PreTrainedConfig): - model_type = "qwen3_5_vision" - - def __init__( - self, - depth=27, - hidden_size=1152, - hidden_act="gelu_pytorch_tanh", - intermediate_size=4304, - num_heads=16, - in_channels=3, - patch_size=16, - spatial_merge_size=2, - temporal_patch_size=2, - out_hidden_size=3584, - num_position_embeddings=2304, - initializer_range=0.02, - **kwargs, - ): - super().__init__(**kwargs) - self.depth = depth - self.hidden_size = hidden_size - self.hidden_act = hidden_act - self.intermediate_size = intermediate_size - self.num_heads = num_heads - self.in_channels = in_channels - self.patch_size = patch_size - self.spatial_merge_size = spatial_merge_size - self.temporal_patch_size = temporal_patch_size - self.out_hidden_size = out_hidden_size - self.num_position_embeddings = num_position_embeddings - self.initializer_range = initializer_range - - -class Qwen3_5Config(PreTrainedConfig): - r""" + model_type is "qwen3_5_text" (used internally by the nested config). + """ + + model_type = "qwen3_5_text" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=248320, + hidden_size=4096, + intermediate_size=12288, + num_hidden_layers=32, + num_attention_heads=16, + num_key_value_heads=4, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_parameters=None, + attention_bias=False, + attention_dropout=0.0, + head_dim=256, + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + layer_types=None, + pad_token_id=None, + bos_token_id=None, + eos_token_id=None, + **kwargs, + ): + self.pad_token_id = pad_token_id + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.tie_word_embeddings = tie_word_embeddings + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.head_dim = head_dim + self.rope_parameters = rope_parameters + kwargs.setdefault("partial_rotary_factor", 0.25) + + self.layer_types = layer_types + if self.layer_types is None: + interval_pattern = kwargs.get("full_attention_interval", 4) + self.layer_types = [ + "linear_attention" if bool((i + 1) % interval_pattern) else "full_attention" + for i in range(self.num_hidden_layers) + ] + layer_type_validation(self.layer_types, self.num_hidden_layers) + + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_key_head_dim = linear_key_head_dim + self.linear_value_head_dim = linear_value_head_dim + self.linear_num_key_heads = linear_num_key_heads + self.linear_num_value_heads = linear_num_value_heads + super().__init__(**kwargs) + + +class Qwen3_5VisionConfig(PreTrainedConfig): + model_type = "qwen3_5_vision" + + def __init__( + self, + depth=27, + hidden_size=1152, + hidden_act="gelu_pytorch_tanh", + intermediate_size=4304, + num_heads=16, + in_channels=3, + patch_size=16, + spatial_merge_size=2, + temporal_patch_size=2, + out_hidden_size=3584, + num_position_embeddings=2304, + initializer_range=0.02, + **kwargs, + ): + super().__init__(**kwargs) + self.depth = depth + self.hidden_size = hidden_size + self.hidden_act = hidden_act + self.intermediate_size = intermediate_size + self.num_heads = num_heads + self.in_channels = in_channels + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + self.temporal_patch_size = temporal_patch_size + self.out_hidden_size = out_hidden_size + self.num_position_embeddings = num_position_embeddings + self.initializer_range = initializer_range + + +class Qwen3_5Config(PreTrainedConfig): + r""" Top-level configuration for Qwen3.5 / Qwen3.6-35B-A3B. - model_type = "qwen3_5" matches the model card / config.json. - Wraps Qwen3_5TextConfig (and optionally Qwen3_5VisionConfig for multimodal use). - For vLLM text-only inference only text_config is consumed. - """ - - model_type = "qwen3_5" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - text_config=None, - vision_config=None, - image_token_id=248056, - video_token_id=248057, - vision_start_token_id=248053, - vision_end_token_id=248054, + model_type = "qwen3_5" matches the model card / config.json. + Wraps Qwen3_5TextConfig (and optionally Qwen3_5VisionConfig for multimodal use). + For vLLM text-only inference only text_config is consumed. + """ + + model_type = "qwen3_5" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config=None, + vision_config=None, + image_token_id=248056, + video_token_id=248057, + vision_start_token_id=248053, + vision_end_token_id=248054, tie_word_embeddings=False, **kwargs, ): serialized_mode = kwargs.pop(HYBRID_KV_ACCOUNTING_CONFIG, None) serialized_layers = kwargs.pop("layers_block_type", None) if isinstance(text_config, dict): - self.text_config = Qwen3_5TextConfig(**text_config) - elif text_config is None: - self.text_config = Qwen3_5TextConfig() - else: - self.text_config = text_config - - if isinstance(vision_config, dict): - self.vision_config = Qwen3_5VisionConfig(**vision_config) - elif vision_config is None: - self.vision_config = Qwen3_5VisionConfig() - else: - self.vision_config = vision_config - - self.image_token_id = image_token_id - self.video_token_id = video_token_id - self.vision_start_token_id = vision_start_token_id + self.text_config = Qwen3_5TextConfig(**text_config) + elif text_config is None: + self.text_config = Qwen3_5TextConfig() + else: + self.text_config = text_config + + if isinstance(vision_config, dict): + self.vision_config = Qwen3_5VisionConfig(**vision_config) + elif vision_config is None: + self.vision_config = Qwen3_5VisionConfig() + else: + self.vision_config = vision_config + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id self.vision_end_token_id = vision_end_token_id self.tie_word_embeddings = tie_word_embeddings super().__init__(**kwargs) @@ -237,6 +237,6 @@ class Qwen3_5Config(PreTrainedConfig): f"{HYBRID_KV_ACCOUNTING_CONFIG}={mode!r}") setattr(self, HYBRID_KV_ACCOUNTING_CONFIG, mode) self.layers_block_type = layers_block_type - - -__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5VisionConfig"] + + +__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5VisionConfig"] diff --git a/qwen3_6_scripts/qwen3_5_moe/__init__.py b/qwen3_6_scripts/qwen3_5_moe/__init__.py index 1c27df9f..6376ee8b 100644 --- a/qwen3_6_scripts/qwen3_5_moe/__init__.py +++ b/qwen3_6_scripts/qwen3_5_moe/__init__.py @@ -1,3 +1,3 @@ -from .configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig - -__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"] +from .configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig + +__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"] diff --git a/qwen3_6_scripts/qwen3_5_moe/configuration_qwen3_5_moe.py b/qwen3_6_scripts/qwen3_5_moe/configuration_qwen3_5_moe.py index 34667d86..0489e3ee 100644 --- a/qwen3_6_scripts/qwen3_5_moe/configuration_qwen3_5_moe.py +++ b/qwen3_6_scripts/qwen3_5_moe/configuration_qwen3_5_moe.py @@ -1,20 +1,20 @@ -# Adapted from transformers 5.2.0 for compatibility with transformers 4.55.3 + torch 2.1.0 -# Source: transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py -# Stubs layer_type_validation and RopeParameters which do not exist in 4.55.3 -# Removes ignore_keys_at_rope_validation / base_model_tp_plan / base_model_pp_plan -# which are 5.x-only and irrelevant for vLLM inference. - +# Adapted from transformers 5.2.0 for compatibility with transformers 4.55.3 + torch 2.1.0 +# Source: transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py +# Stubs layer_type_validation and RopeParameters which do not exist in 4.55.3 +# Removes ignore_keys_at_rope_validation / base_model_tp_plan / base_model_pp_plan +# which are 5.x-only and irrelevant for vLLM inference. + import os from typing import Optional - -from ...configuration_utils import PretrainedConfig as PreTrainedConfig - -# --- Local stubs for APIs not present in transformers 4.55.3 --- + +from ...configuration_utils import PretrainedConfig as PreTrainedConfig + +# --- Local stubs for APIs not present in transformers 4.55.3 --- def layer_type_validation(layer_types, num_hidden_layers=None, attention=True): - allowed = {"full_attention", "linear_attention"} - if not all(lt in allowed for lt in layer_types): - raise ValueError(f"layer_types entries must be in {allowed}, got {layer_types}") - if num_hidden_layers is not None and num_hidden_layers != len(layer_types): + allowed = {"full_attention", "linear_attention"} + if not all(lt in allowed for lt in layer_types): + raise ValueError(f"layer_types entries must be in {allowed}, got {layer_types}") + if num_hidden_layers is not None and num_hidden_layers != len(layer_types): raise ValueError( f"num_hidden_layers ({num_hidden_layers}) != len(layer_types) ({len(layer_types)})" ) @@ -57,7 +57,7 @@ def _vllm_layers_block_type( "attention" if layer_type == "full_attention" else layer_type for layer_type in layer_types ] - + try: from typing import TypedDict except ImportError: @@ -68,171 +68,171 @@ else: rope_type: str partial_rotary_factor: float factor: float - -# --- End stubs --- - - -class Qwen3_5MoeTextConfig(PreTrainedConfig): - r""" - Configuration for the text backbone of Qwen3.5-MoE / Qwen3.6-35B-A3B models. - model_type is "qwen3_5_moe_text" (used internally by the nested config). - """ - - model_type = "qwen3_5_moe_text" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=248320, - hidden_size=2048, - num_hidden_layers=40, - num_attention_heads=16, - num_key_value_heads=2, - hidden_act="silu", - max_position_embeddings=32768, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - tie_word_embeddings=False, - rope_parameters=None, - attention_bias=False, - attention_dropout=0.0, - head_dim=256, - linear_conv_kernel_dim=4, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=32, - moe_intermediate_size=512, - shared_expert_intermediate_size=512, - num_experts_per_tok=8, - num_experts=256, - output_router_logits=False, - router_aux_loss_coef=0.001, - layer_types=None, - pad_token_id=None, - bos_token_id=None, - eos_token_id=None, - **kwargs, - ): - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id - self.tie_word_embeddings = tie_word_embeddings - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.head_dim = head_dim - self.rope_parameters = rope_parameters - kwargs.setdefault("partial_rotary_factor", 0.25) - - self.layer_types = layer_types - if self.layer_types is None: - interval_pattern = kwargs.get("full_attention_interval", 4) - self.layer_types = [ - "linear_attention" if bool((i + 1) % interval_pattern) else "full_attention" - for i in range(self.num_hidden_layers) - ] - layer_type_validation(self.layer_types, self.num_hidden_layers) - - self.linear_conv_kernel_dim = linear_conv_kernel_dim - self.linear_key_head_dim = linear_key_head_dim - self.linear_value_head_dim = linear_value_head_dim - self.linear_num_key_heads = linear_num_key_heads - self.linear_num_value_heads = linear_num_value_heads - self.moe_intermediate_size = moe_intermediate_size - self.shared_expert_intermediate_size = shared_expert_intermediate_size - self.num_experts_per_tok = num_experts_per_tok - self.num_experts = num_experts - self.output_router_logits = output_router_logits - self.router_aux_loss_coef = router_aux_loss_coef - super().__init__(**kwargs) - - -class Qwen3_5MoeVisionConfig(PreTrainedConfig): - model_type = "qwen3_5_moe" - - def __init__( - self, - depth=27, - hidden_size=1152, - hidden_act="gelu_pytorch_tanh", - intermediate_size=4304, - num_heads=16, - in_channels=3, - patch_size=16, - spatial_merge_size=2, - temporal_patch_size=2, - out_hidden_size=3584, - num_position_embeddings=2304, - initializer_range=0.02, - **kwargs, - ): - super().__init__(**kwargs) - self.depth = depth - self.hidden_size = hidden_size - self.hidden_act = hidden_act - self.intermediate_size = intermediate_size - self.num_heads = num_heads - self.in_channels = in_channels - self.patch_size = patch_size - self.spatial_merge_size = spatial_merge_size - self.temporal_patch_size = temporal_patch_size - self.out_hidden_size = out_hidden_size - self.num_position_embeddings = num_position_embeddings - self.initializer_range = initializer_range - - -class Qwen3_5MoeConfig(PreTrainedConfig): - r""" - Top-level configuration for Qwen3.5-MoE / Qwen3.6-35B-A3B. - model_type = "qwen3_5_moe" matches the model card / config.json. - Wraps Qwen3_5MoeTextConfig (and optionally Qwen3_5MoeVisionConfig). - For vLLM text-only inference only text_config is consumed. - """ - - model_type = "qwen3_5_moe" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - text_config=None, - vision_config=None, - image_token_id=248056, - video_token_id=248057, - vision_start_token_id=248053, - vision_end_token_id=248054, + +# --- End stubs --- + + +class Qwen3_5MoeTextConfig(PreTrainedConfig): + r""" + Configuration for the text backbone of Qwen3.5-MoE / Qwen3.6-35B-A3B models. + model_type is "qwen3_5_moe_text" (used internally by the nested config). + """ + + model_type = "qwen3_5_moe_text" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=248320, + hidden_size=2048, + num_hidden_layers=40, + num_attention_heads=16, + num_key_value_heads=2, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_parameters=None, + attention_bias=False, + attention_dropout=0.0, + head_dim=256, + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + moe_intermediate_size=512, + shared_expert_intermediate_size=512, + num_experts_per_tok=8, + num_experts=256, + output_router_logits=False, + router_aux_loss_coef=0.001, + layer_types=None, + pad_token_id=None, + bos_token_id=None, + eos_token_id=None, + **kwargs, + ): + self.pad_token_id = pad_token_id + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.tie_word_embeddings = tie_word_embeddings + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.head_dim = head_dim + self.rope_parameters = rope_parameters + kwargs.setdefault("partial_rotary_factor", 0.25) + + self.layer_types = layer_types + if self.layer_types is None: + interval_pattern = kwargs.get("full_attention_interval", 4) + self.layer_types = [ + "linear_attention" if bool((i + 1) % interval_pattern) else "full_attention" + for i in range(self.num_hidden_layers) + ] + layer_type_validation(self.layer_types, self.num_hidden_layers) + + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_key_head_dim = linear_key_head_dim + self.linear_value_head_dim = linear_value_head_dim + self.linear_num_key_heads = linear_num_key_heads + self.linear_num_value_heads = linear_num_value_heads + self.moe_intermediate_size = moe_intermediate_size + self.shared_expert_intermediate_size = shared_expert_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.num_experts = num_experts + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + super().__init__(**kwargs) + + +class Qwen3_5MoeVisionConfig(PreTrainedConfig): + model_type = "qwen3_5_moe" + + def __init__( + self, + depth=27, + hidden_size=1152, + hidden_act="gelu_pytorch_tanh", + intermediate_size=4304, + num_heads=16, + in_channels=3, + patch_size=16, + spatial_merge_size=2, + temporal_patch_size=2, + out_hidden_size=3584, + num_position_embeddings=2304, + initializer_range=0.02, + **kwargs, + ): + super().__init__(**kwargs) + self.depth = depth + self.hidden_size = hidden_size + self.hidden_act = hidden_act + self.intermediate_size = intermediate_size + self.num_heads = num_heads + self.in_channels = in_channels + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + self.temporal_patch_size = temporal_patch_size + self.out_hidden_size = out_hidden_size + self.num_position_embeddings = num_position_embeddings + self.initializer_range = initializer_range + + +class Qwen3_5MoeConfig(PreTrainedConfig): + r""" + Top-level configuration for Qwen3.5-MoE / Qwen3.6-35B-A3B. + model_type = "qwen3_5_moe" matches the model card / config.json. + Wraps Qwen3_5MoeTextConfig (and optionally Qwen3_5MoeVisionConfig). + For vLLM text-only inference only text_config is consumed. + """ + + model_type = "qwen3_5_moe" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config=None, + vision_config=None, + image_token_id=248056, + video_token_id=248057, + vision_start_token_id=248053, + vision_end_token_id=248054, tie_word_embeddings=False, **kwargs, ): serialized_mode = kwargs.pop(HYBRID_KV_ACCOUNTING_CONFIG, None) serialized_layers = kwargs.pop("layers_block_type", None) if isinstance(text_config, dict): - self.text_config = Qwen3_5MoeTextConfig(**text_config) - elif text_config is None: - self.text_config = Qwen3_5MoeTextConfig() - else: - self.text_config = text_config - - if isinstance(vision_config, dict): - self.vision_config = Qwen3_5MoeVisionConfig(**vision_config) - elif vision_config is None: - self.vision_config = Qwen3_5MoeVisionConfig() - else: - self.vision_config = vision_config - - self.image_token_id = image_token_id - self.video_token_id = video_token_id - self.vision_start_token_id = vision_start_token_id + self.text_config = Qwen3_5MoeTextConfig(**text_config) + elif text_config is None: + self.text_config = Qwen3_5MoeTextConfig() + else: + self.text_config = text_config + + if isinstance(vision_config, dict): + self.vision_config = Qwen3_5MoeVisionConfig(**vision_config) + elif vision_config is None: + self.vision_config = Qwen3_5MoeVisionConfig() + else: + self.vision_config = vision_config + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id self.vision_end_token_id = vision_end_token_id self.tie_word_embeddings = tie_word_embeddings super().__init__(**kwargs) @@ -247,6 +247,6 @@ class Qwen3_5MoeConfig(PreTrainedConfig): f"{HYBRID_KV_ACCOUNTING_CONFIG}={mode!r}") setattr(self, HYBRID_KV_ACCOUNTING_CONFIG, mode) self.layers_block_type = layers_block_type - - -__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"] + + +__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"] diff --git a/qwen3_6_scripts/qwen3coder_tool_parser.py b/qwen3_6_scripts/qwen3coder_tool_parser.py index 04cbc6f6..66be4567 100644 --- a/qwen3_6_scripts/qwen3coder_tool_parser.py +++ b/qwen3_6_scripts/qwen3coder_tool_parser.py @@ -1,172 +1,172 @@ -import ast -import json -import uuid -from typing import Any, Dict, List, Optional, Sequence, Union - -import regex as re - -from vllm.entrypoints.openai.protocol import (ChatCompletionRequest, - ChatCompletionToolsParam, - DeltaFunctionCall, DeltaMessage, - DeltaToolCall, - ExtractedToolCallInformation, - FunctionCall, ToolCall) -from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import ( - ToolParser, ToolParserManager) -from vllm.logger import init_logger -from vllm.transformers_utils.tokenizer import AnyTokenizer - -logger = init_logger(__name__) - - -@ToolParserManager.register_module("qwen3_coder") -class Qwen3CoderToolParser(ToolParser): - """ - Tool parser for Qwen3 models using XML-style tool call format: - - value - - - Port of vllm-original qwen3coder_tool_parser.py to vllm 0.6.3 API. - """ - - def __init__(self, tokenizer: AnyTokenizer): - super().__init__(tokenizer) - - self.current_tool_name_sent: bool = False - self.prev_tool_call_arr: List[Dict] = [] - # Base class uses int; we override with string IDs - self.current_tool_id: Optional[str] = None # type: ignore[assignment] - self.streamed_args_for_tool: List[str] = [] - - self.tool_call_start_token: str = "" - self.tool_call_end_token: str = "" - self.tool_call_prefix: str = "(.*?)", re.DOTALL) - self.tool_call_regex = re.compile( - r"(.*?)|(.*?)$", re.DOTALL) - self.tool_call_function_regex = re.compile( - r"||(?=)|$)", - re.DOTALL) - - if not self.model_tokenizer: - raise ValueError( - "The model tokenizer must be passed to the ToolParser " - "constructor during construction.") - - self.tool_call_start_token_id = self.vocab.get( - self.tool_call_start_token) - self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token) - - if (self.tool_call_start_token_id is None - or self.tool_call_end_token_id is None): - raise RuntimeError( - "Qwen3 XML Tool parser could not locate tool call start/end " - "tokens in the tokenizer!") - - logger.debug("vLLM Successfully imported tool parser %s !", - self.__class__.__name__) - - - def _generate_tool_call_id(self) -> str: - return f"call_{uuid.uuid4().hex[:24]}" - - def _reset_streaming_state(self) -> None: - self.current_tool_index = 0 - self.is_tool_call_started = False - self.header_sent = False - self.current_tool_id = None - self.current_function_name: Optional[str] = None - self.current_param_name: Optional[str] = None - self.current_param_value: str = "" - self.param_count = 0 - self.in_param = False - self.in_function = False - self.accumulated_text: str = "" - self.json_started = False - self.json_closed = False - self.accumulated_params: Dict[str, Any] = {} - self.streaming_request: Optional[ChatCompletionRequest] = None - - def _get_arguments_config( - self, func_name: str, - tools: Optional[List[ChatCompletionToolsParam]]) -> Dict: - if tools is None: - return {} - for config in tools: - if not hasattr(config, "type") or not ( - hasattr(config, "function") - and hasattr(config.function, "name")): - continue - if config.type == "function" and config.function.name == func_name: - if not hasattr(config.function, "parameters"): - return {} - params = config.function.parameters - if isinstance(params, dict) and "properties" in params: - return params["properties"] - elif isinstance(params, dict): - return params - else: - return {} - logger.debug("Tool '%s' is not defined in the tools list.", func_name) - return {} - - def _convert_param_value(self, param_value: str, param_name: str, - param_config: Dict, func_name: str) -> Any: - if param_value.lower() == "null": - return None - - if param_name not in param_config: - if param_config != {}: - logger.debug( - "Parsed parameter '%s' is not defined in tool '%s', " - "returning string value.", param_name, func_name) - return param_value - - if (isinstance(param_config[param_name], dict) - and "type" in param_config[param_name]): - param_type = str( - param_config[param_name]["type"]).strip().lower() - else: - param_type = "string" - - if param_type in ["string", "str", "text", "varchar", "char", "enum"]: - return param_value - elif (param_type.startswith("int") or param_type.startswith("uint") - or param_type.startswith("long") - or param_type.startswith("short") - or param_type.startswith("unsigned")): - try: - return int(param_value) - except (ValueError, TypeError): - return param_value - elif param_type.startswith("num") or param_type.startswith("float"): - try: - v = float(param_value) - return int(v) if v - int(v) == 0 else v - except (ValueError, TypeError): - return param_value - elif param_type in ["boolean", "bool", "binary"]: - lower = param_value.lower() - if lower not in ["true", "false"]: - logger.debug( - "Parameter '%s' value '%s' is not boolean in tool '%s'.", - param_name, param_value, func_name) - return lower == "true" - else: - if (param_type in ["object", "array", "arr"] - or param_type.startswith("dict") - or param_type.startswith("list")): +import ast +import json +import uuid +from typing import Any, Dict, List, Optional, Sequence, Union + +import regex as re + +from vllm.entrypoints.openai.protocol import (ChatCompletionRequest, + ChatCompletionToolsParam, + DeltaFunctionCall, DeltaMessage, + DeltaToolCall, + ExtractedToolCallInformation, + FunctionCall, ToolCall) +from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import ( + ToolParser, ToolParserManager) +from vllm.logger import init_logger +from vllm.transformers_utils.tokenizer import AnyTokenizer + +logger = init_logger(__name__) + + +@ToolParserManager.register_module("qwen3_coder") +class Qwen3CoderToolParser(ToolParser): + """ + Tool parser for Qwen3 models using XML-style tool call format: + + value + + + Port of vllm-original qwen3coder_tool_parser.py to vllm 0.6.3 API. + """ + + def __init__(self, tokenizer: AnyTokenizer): + super().__init__(tokenizer) + + self.current_tool_name_sent: bool = False + self.prev_tool_call_arr: List[Dict] = [] + # Base class uses int; we override with string IDs + self.current_tool_id: Optional[str] = None # type: ignore[assignment] + self.streamed_args_for_tool: List[str] = [] + + self.tool_call_start_token: str = "" + self.tool_call_end_token: str = "" + self.tool_call_prefix: str = "(.*?)", re.DOTALL) + self.tool_call_regex = re.compile( + r"(.*?)|(.*?)$", re.DOTALL) + self.tool_call_function_regex = re.compile( + r"||(?=)|$)", + re.DOTALL) + + if not self.model_tokenizer: + raise ValueError( + "The model tokenizer must be passed to the ToolParser " + "constructor during construction.") + + self.tool_call_start_token_id = self.vocab.get( + self.tool_call_start_token) + self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token) + + if (self.tool_call_start_token_id is None + or self.tool_call_end_token_id is None): + raise RuntimeError( + "Qwen3 XML Tool parser could not locate tool call start/end " + "tokens in the tokenizer!") + + logger.debug("vLLM Successfully imported tool parser %s !", + self.__class__.__name__) + + + def _generate_tool_call_id(self) -> str: + return f"call_{uuid.uuid4().hex[:24]}" + + def _reset_streaming_state(self) -> None: + self.current_tool_index = 0 + self.is_tool_call_started = False + self.header_sent = False + self.current_tool_id = None + self.current_function_name: Optional[str] = None + self.current_param_name: Optional[str] = None + self.current_param_value: str = "" + self.param_count = 0 + self.in_param = False + self.in_function = False + self.accumulated_text: str = "" + self.json_started = False + self.json_closed = False + self.accumulated_params: Dict[str, Any] = {} + self.streaming_request: Optional[ChatCompletionRequest] = None + + def _get_arguments_config( + self, func_name: str, + tools: Optional[List[ChatCompletionToolsParam]]) -> Dict: + if tools is None: + return {} + for config in tools: + if not hasattr(config, "type") or not ( + hasattr(config, "function") + and hasattr(config.function, "name")): + continue + if config.type == "function" and config.function.name == func_name: + if not hasattr(config.function, "parameters"): + return {} + params = config.function.parameters + if isinstance(params, dict) and "properties" in params: + return params["properties"] + elif isinstance(params, dict): + return params + else: + return {} + logger.debug("Tool '%s' is not defined in the tools list.", func_name) + return {} + + def _convert_param_value(self, param_value: str, param_name: str, + param_config: Dict, func_name: str) -> Any: + if param_value.lower() == "null": + return None + + if param_name not in param_config: + if param_config != {}: + logger.debug( + "Parsed parameter '%s' is not defined in tool '%s', " + "returning string value.", param_name, func_name) + return param_value + + if (isinstance(param_config[param_name], dict) + and "type" in param_config[param_name]): + param_type = str( + param_config[param_name]["type"]).strip().lower() + else: + param_type = "string" + + if param_type in ["string", "str", "text", "varchar", "char", "enum"]: + return param_value + elif (param_type.startswith("int") or param_type.startswith("uint") + or param_type.startswith("long") + or param_type.startswith("short") + or param_type.startswith("unsigned")): + try: + return int(param_value) + except (ValueError, TypeError): + return param_value + elif param_type.startswith("num") or param_type.startswith("float"): + try: + v = float(param_value) + return int(v) if v - int(v) == 0 else v + except (ValueError, TypeError): + return param_value + elif param_type in ["boolean", "bool", "binary"]: + lower = param_value.lower() + if lower not in ["true", "false"]: + logger.debug( + "Parameter '%s' value '%s' is not boolean in tool '%s'.", + param_name, param_value, func_name) + return lower == "true" + else: + if (param_type in ["object", "array", "arr"] + or param_type.startswith("dict") + or param_type.startswith("list")): try: return json.loads(param_value) except (json.JSONDecodeError, TypeError, ValueError): @@ -186,334 +186,334 @@ class Qwen3CoderToolParser(ToolParser): func_name, exc_info=True) return param_value - - def _parse_xml_function_call( - self, function_call_str: str, - tools: Optional[List[ChatCompletionToolsParam]]) -> ToolCall: - end_index = function_call_str.index(">") - function_name = function_call_str[:end_index] - param_config = self._get_arguments_config(function_name, tools) - parameters = function_call_str[end_index + 1:] - param_dict: Dict[str, Any] = {} - for match_text in self.tool_call_parameter_regex.findall(parameters): - idx = match_text.index(">") - param_name = match_text[:idx] - param_value = str(match_text[idx + 1:]) - if param_value.startswith("\n"): - param_value = param_value[1:] - if param_value.endswith("\n"): - param_value = param_value[:-1] - param_dict[param_name] = self._convert_param_value( - param_value, param_name, param_config, function_name) - return ToolCall( - type="function", - function=FunctionCall( - name=function_name, - arguments=json.dumps(param_dict, ensure_ascii=False))) - - def _get_function_calls(self, model_output: str) -> List[str]: - matched_ranges = self.tool_call_regex.findall(model_output) - raw_tool_calls = [ - match[0] if match[0] else match[1] for match in matched_ranges - ] - if not raw_tool_calls: - raw_tool_calls = [model_output] - raw_function_calls: List[tuple] = [] - for tool_call in raw_tool_calls: - raw_function_calls.extend( - self.tool_call_function_regex.findall(tool_call)) - return [match[0] if match[0] else match[1] - for match in raw_function_calls] - - def extract_tool_calls( - self, model_output: str, - request: ChatCompletionRequest) -> ExtractedToolCallInformation: - if self.tool_call_prefix not in model_output: - return ExtractedToolCallInformation(tools_called=False, - tool_calls=[], - content=model_output) - try: - function_calls = self._get_function_calls(model_output) - if not function_calls: - return ExtractedToolCallInformation(tools_called=False, - tool_calls=[], - content=model_output) - - tool_calls = [ - self._parse_xml_function_call(fc, request.tools) - for fc in function_calls - ] - - self.prev_tool_call_arr.clear() - for tc in tool_calls: - self.prev_tool_call_arr.append({ - "name": tc.function.name, - "arguments": tc.function.arguments, - }) - - content_index = model_output.find(self.tool_call_start_token) - idx = model_output.find(self.tool_call_prefix) - content_index = content_index if content_index >= 0 else idx - content = model_output[:content_index] - - return ExtractedToolCallInformation( - tools_called=bool(tool_calls), - tool_calls=tool_calls, - content=content if content else None, - ) - except Exception: - logger.exception("Error extracting tool call from response.") - return ExtractedToolCallInformation(tools_called=False, - tool_calls=[], - content=model_output) - - def extract_tool_calls_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - request: ChatCompletionRequest, - ) -> Union[DeltaMessage, None]: - if not previous_text: - self._reset_streaming_state() - self.streaming_request = request - - if not delta_text: - if delta_token_ids and self.tool_call_end_token_id not in delta_token_ids: - complete_calls = len( - self.tool_call_complete_regex.findall(current_text)) - if complete_calls > 0 and self.prev_tool_call_arr: - open_calls = ( - current_text.count(self.tool_call_start_token) - - current_text.count(self.tool_call_end_token)) - if open_calls == 0: - return DeltaMessage(content="") - elif not self.is_tool_call_started and current_text: - return DeltaMessage(content="") - return None - - self.accumulated_text = current_text - - if self.json_closed and not self.in_function: - tool_ends = current_text.count(self.tool_call_end_token) - if tool_ends > self.current_tool_index: - self.current_tool_index += 1 - self.header_sent = False - self.param_count = 0 - self.json_started = False - self.json_closed = False - self.accumulated_params = {} - tool_starts = current_text.count(self.tool_call_start_token) - if self.current_tool_index >= tool_starts: - self.is_tool_call_started = False - return None - - if not self.is_tool_call_started: - if (self.tool_call_start_token_id in delta_token_ids - or self.tool_call_start_token in delta_text): - self.is_tool_call_started = True - if self.tool_call_start_token in delta_text: - content_before = delta_text[:delta_text.index( - self.tool_call_start_token)] - if content_before: - return DeltaMessage(content=content_before) - return None - else: - if (current_text.rstrip().endswith(self.tool_call_end_token) - and delta_text.strip() == ""): - return None - return DeltaMessage(content=delta_text) - - tool_starts_count = current_text.count(self.tool_call_start_token) - if self.current_tool_index >= tool_starts_count: - return None - - # Locate the current tool call's text slice - tool_start_positions: List[int] = [] - search = 0 - while True: - search = current_text.find(self.tool_call_start_token, search) - if search == -1: - break - tool_start_positions.append(search) - search += len(self.tool_call_start_token) - - if self.current_tool_index >= len(tool_start_positions): - return None - - tool_start_idx = tool_start_positions[self.current_tool_index] - tool_end_idx = current_text.find(self.tool_call_end_token, - tool_start_idx) - if tool_end_idx == -1: - tool_text = current_text[tool_start_idx:] - else: - tool_text = current_text[tool_start_idx:tool_end_idx + - len(self.tool_call_end_token)] - - if not self.header_sent: - if self.tool_call_prefix in tool_text: - func_start = (tool_text.find(self.tool_call_prefix) + - len(self.tool_call_prefix)) - func_end = tool_text.find(">", func_start) - if func_end != -1: - self.current_function_name = tool_text[func_start:func_end] - self.current_tool_id = self._generate_tool_call_id() - self.header_sent = True - self.in_function = True - self.prev_tool_call_arr.append({ - "name": self.current_function_name, - "arguments": "{}", - }) - self.streamed_args_for_tool.append("") - return DeltaMessage(tool_calls=[ - DeltaToolCall( - index=self.current_tool_index, - id=self.current_tool_id, - function=DeltaFunctionCall( - name=self.current_function_name, - arguments=""), - type="function", - ) - ]) - return None - - if self.in_function: - if not self.json_started: - self.json_started = True - self.streamed_args_for_tool[self.current_tool_index] += "{" - return DeltaMessage(tool_calls=[ - DeltaToolCall( - index=self.current_tool_index, - function=DeltaFunctionCall(arguments="{"), - ) - ]) - - # Collect all complete parameters in one pass (speculative-decode safe) - param_starts: List[int] = [] - search = 0 - while True: - search = tool_text.find(self.parameter_prefix, search) - if search == -1: - break - param_starts.append(search) - search += len(self.parameter_prefix) - - json_fragments: List[str] = [] - while not self.in_param and self.param_count < len(param_starts): - param_idx = param_starts[self.param_count] - param_start = param_idx + len(self.parameter_prefix) - remaining = tool_text[param_start:] - - if ">" not in remaining: - break - - name_end = remaining.find(">") - current_param_name = remaining[:name_end] - value_start = param_start + name_end + 1 - value_text = tool_text[value_start:] - if value_text.startswith("\n"): - value_text = value_text[1:] - - param_end_idx = value_text.find(self.parameter_end_token) - if param_end_idx == -1: - next_param = value_text.find(self.parameter_prefix) - func_end = value_text.find(self.function_end_token) - if next_param != -1 and (func_end == -1 - or next_param < func_end): - param_end_idx = next_param - elif func_end != -1: - param_end_idx = func_end - else: - tool_end_in_value = value_text.find( - self.tool_call_end_token) - if tool_end_in_value != -1: - param_end_idx = tool_end_in_value - else: - break - - if param_end_idx == -1: - break - - param_value = value_text[:param_end_idx] - if param_value.endswith("\n"): - param_value = param_value[:-1] - - self.accumulated_params[current_param_name] = param_value - param_config = self._get_arguments_config( - self.current_function_name or "", - self.streaming_request.tools - if self.streaming_request else None) - converted = self._convert_param_value( - param_value, current_param_name, param_config, - self.current_function_name or "") - serialized = json.dumps(converted, ensure_ascii=False) - + + def _parse_xml_function_call( + self, function_call_str: str, + tools: Optional[List[ChatCompletionToolsParam]]) -> ToolCall: + end_index = function_call_str.index(">") + function_name = function_call_str[:end_index] + param_config = self._get_arguments_config(function_name, tools) + parameters = function_call_str[end_index + 1:] + param_dict: Dict[str, Any] = {} + for match_text in self.tool_call_parameter_regex.findall(parameters): + idx = match_text.index(">") + param_name = match_text[:idx] + param_value = str(match_text[idx + 1:]) + if param_value.startswith("\n"): + param_value = param_value[1:] + if param_value.endswith("\n"): + param_value = param_value[:-1] + param_dict[param_name] = self._convert_param_value( + param_value, param_name, param_config, function_name) + return ToolCall( + type="function", + function=FunctionCall( + name=function_name, + arguments=json.dumps(param_dict, ensure_ascii=False))) + + def _get_function_calls(self, model_output: str) -> List[str]: + matched_ranges = self.tool_call_regex.findall(model_output) + raw_tool_calls = [ + match[0] if match[0] else match[1] for match in matched_ranges + ] + if not raw_tool_calls: + raw_tool_calls = [model_output] + raw_function_calls: List[tuple] = [] + for tool_call in raw_tool_calls: + raw_function_calls.extend( + self.tool_call_function_regex.findall(tool_call)) + return [match[0] if match[0] else match[1] + for match in raw_function_calls] + + def extract_tool_calls( + self, model_output: str, + request: ChatCompletionRequest) -> ExtractedToolCallInformation: + if self.tool_call_prefix not in model_output: + return ExtractedToolCallInformation(tools_called=False, + tool_calls=[], + content=model_output) + try: + function_calls = self._get_function_calls(model_output) + if not function_calls: + return ExtractedToolCallInformation(tools_called=False, + tool_calls=[], + content=model_output) + + tool_calls = [ + self._parse_xml_function_call(fc, request.tools) + for fc in function_calls + ] + + self.prev_tool_call_arr.clear() + for tc in tool_calls: + self.prev_tool_call_arr.append({ + "name": tc.function.name, + "arguments": tc.function.arguments, + }) + + content_index = model_output.find(self.tool_call_start_token) + idx = model_output.find(self.tool_call_prefix) + content_index = content_index if content_index >= 0 else idx + content = model_output[:content_index] + + return ExtractedToolCallInformation( + tools_called=bool(tool_calls), + tool_calls=tool_calls, + content=content if content else None, + ) + except Exception: + logger.exception("Error extracting tool call from response.") + return ExtractedToolCallInformation(tools_called=False, + tool_calls=[], + content=model_output) + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request: ChatCompletionRequest, + ) -> Union[DeltaMessage, None]: + if not previous_text: + self._reset_streaming_state() + self.streaming_request = request + + if not delta_text: + if delta_token_ids and self.tool_call_end_token_id not in delta_token_ids: + complete_calls = len( + self.tool_call_complete_regex.findall(current_text)) + if complete_calls > 0 and self.prev_tool_call_arr: + open_calls = ( + current_text.count(self.tool_call_start_token) - + current_text.count(self.tool_call_end_token)) + if open_calls == 0: + return DeltaMessage(content="") + elif not self.is_tool_call_started and current_text: + return DeltaMessage(content="") + return None + + self.accumulated_text = current_text + + if self.json_closed and not self.in_function: + tool_ends = current_text.count(self.tool_call_end_token) + if tool_ends > self.current_tool_index: + self.current_tool_index += 1 + self.header_sent = False + self.param_count = 0 + self.json_started = False + self.json_closed = False + self.accumulated_params = {} + tool_starts = current_text.count(self.tool_call_start_token) + if self.current_tool_index >= tool_starts: + self.is_tool_call_started = False + return None + + if not self.is_tool_call_started: + if (self.tool_call_start_token_id in delta_token_ids + or self.tool_call_start_token in delta_text): + self.is_tool_call_started = True + if self.tool_call_start_token in delta_text: + content_before = delta_text[:delta_text.index( + self.tool_call_start_token)] + if content_before: + return DeltaMessage(content=content_before) + return None + else: + if (current_text.rstrip().endswith(self.tool_call_end_token) + and delta_text.strip() == ""): + return None + return DeltaMessage(content=delta_text) + + tool_starts_count = current_text.count(self.tool_call_start_token) + if self.current_tool_index >= tool_starts_count: + return None + + # Locate the current tool call's text slice + tool_start_positions: List[int] = [] + search = 0 + while True: + search = current_text.find(self.tool_call_start_token, search) + if search == -1: + break + tool_start_positions.append(search) + search += len(self.tool_call_start_token) + + if self.current_tool_index >= len(tool_start_positions): + return None + + tool_start_idx = tool_start_positions[self.current_tool_index] + tool_end_idx = current_text.find(self.tool_call_end_token, + tool_start_idx) + if tool_end_idx == -1: + tool_text = current_text[tool_start_idx:] + else: + tool_text = current_text[tool_start_idx:tool_end_idx + + len(self.tool_call_end_token)] + + if not self.header_sent: + if self.tool_call_prefix in tool_text: + func_start = (tool_text.find(self.tool_call_prefix) + + len(self.tool_call_prefix)) + func_end = tool_text.find(">", func_start) + if func_end != -1: + self.current_function_name = tool_text[func_start:func_end] + self.current_tool_id = self._generate_tool_call_id() + self.header_sent = True + self.in_function = True + self.prev_tool_call_arr.append({ + "name": self.current_function_name, + "arguments": "{}", + }) + self.streamed_args_for_tool.append("") + return DeltaMessage(tool_calls=[ + DeltaToolCall( + index=self.current_tool_index, + id=self.current_tool_id, + function=DeltaFunctionCall( + name=self.current_function_name, + arguments=""), + type="function", + ) + ]) + return None + + if self.in_function: + if not self.json_started: + self.json_started = True + self.streamed_args_for_tool[self.current_tool_index] += "{" + return DeltaMessage(tool_calls=[ + DeltaToolCall( + index=self.current_tool_index, + function=DeltaFunctionCall(arguments="{"), + ) + ]) + + # Collect all complete parameters in one pass (speculative-decode safe) + param_starts: List[int] = [] + search = 0 + while True: + search = tool_text.find(self.parameter_prefix, search) + if search == -1: + break + param_starts.append(search) + search += len(self.parameter_prefix) + + json_fragments: List[str] = [] + while not self.in_param and self.param_count < len(param_starts): + param_idx = param_starts[self.param_count] + param_start = param_idx + len(self.parameter_prefix) + remaining = tool_text[param_start:] + + if ">" not in remaining: + break + + name_end = remaining.find(">") + current_param_name = remaining[:name_end] + value_start = param_start + name_end + 1 + value_text = tool_text[value_start:] + if value_text.startswith("\n"): + value_text = value_text[1:] + + param_end_idx = value_text.find(self.parameter_end_token) + if param_end_idx == -1: + next_param = value_text.find(self.parameter_prefix) + func_end = value_text.find(self.function_end_token) + if next_param != -1 and (func_end == -1 + or next_param < func_end): + param_end_idx = next_param + elif func_end != -1: + param_end_idx = func_end + else: + tool_end_in_value = value_text.find( + self.tool_call_end_token) + if tool_end_in_value != -1: + param_end_idx = tool_end_in_value + else: + break + + if param_end_idx == -1: + break + + param_value = value_text[:param_end_idx] + if param_value.endswith("\n"): + param_value = param_value[:-1] + + self.accumulated_params[current_param_name] = param_value + param_config = self._get_arguments_config( + self.current_function_name or "", + self.streaming_request.tools + if self.streaming_request else None) + converted = self._convert_param_value( + param_value, current_param_name, param_config, + self.current_function_name or "") + serialized = json.dumps(converted, ensure_ascii=False) + sep = "" if self.param_count == 0 else ", " key = json.dumps(current_param_name, ensure_ascii=False) json_fragments.append(f"{sep}{key}: {serialized}") - self.param_count += 1 - - if json_fragments: - combined = "".join(json_fragments) - if self.current_tool_index < len(self.streamed_args_for_tool): - self.streamed_args_for_tool[ - self.current_tool_index] += combined - else: - logger.warning( - "streamed_args_for_tool out of sync: index=%d len=%d", - self.current_tool_index, - len(self.streamed_args_for_tool)) - return DeltaMessage(tool_calls=[ - DeltaToolCall( - index=self.current_tool_index, - function=DeltaFunctionCall(arguments=combined), - ) - ]) - - # Emit closing brace when is seen (after params are done) - if not self.json_closed and self.function_end_token in tool_text: - self.json_closed = True - func_start = (tool_text.find(self.tool_call_prefix) + - len(self.tool_call_prefix)) - func_content_end = tool_text.find(self.function_end_token, - func_start) - if func_content_end != -1: - try: - parsed_tool = self._parse_xml_function_call( - tool_text[func_start:func_content_end], - self.streaming_request.tools - if self.streaming_request else None) - if self.current_tool_index < len( - self.prev_tool_call_arr): - self.prev_tool_call_arr[ - self.current_tool_index]["arguments"] = ( - parsed_tool.function.arguments) - except Exception: - logger.debug("Failed to parse tool call during " - "streaming: %s", - tool_text, - exc_info=True) - - if self.current_tool_index < len(self.streamed_args_for_tool): - self.streamed_args_for_tool[ - self.current_tool_index] += "}" - else: - logger.warning( - "streamed_args_for_tool out of sync: index=%d len=%d", - self.current_tool_index, - len(self.streamed_args_for_tool)) - - result = DeltaMessage(tool_calls=[ - DeltaToolCall( - index=self.current_tool_index, - function=DeltaFunctionCall(arguments="}"), - ) - ]) - self.in_function = False - self.accumulated_params = {} - return result - - return None + self.param_count += 1 + + if json_fragments: + combined = "".join(json_fragments) + if self.current_tool_index < len(self.streamed_args_for_tool): + self.streamed_args_for_tool[ + self.current_tool_index] += combined + else: + logger.warning( + "streamed_args_for_tool out of sync: index=%d len=%d", + self.current_tool_index, + len(self.streamed_args_for_tool)) + return DeltaMessage(tool_calls=[ + DeltaToolCall( + index=self.current_tool_index, + function=DeltaFunctionCall(arguments=combined), + ) + ]) + + # Emit closing brace when is seen (after params are done) + if not self.json_closed and self.function_end_token in tool_text: + self.json_closed = True + func_start = (tool_text.find(self.tool_call_prefix) + + len(self.tool_call_prefix)) + func_content_end = tool_text.find(self.function_end_token, + func_start) + if func_content_end != -1: + try: + parsed_tool = self._parse_xml_function_call( + tool_text[func_start:func_content_end], + self.streaming_request.tools + if self.streaming_request else None) + if self.current_tool_index < len( + self.prev_tool_call_arr): + self.prev_tool_call_arr[ + self.current_tool_index]["arguments"] = ( + parsed_tool.function.arguments) + except Exception: + logger.debug("Failed to parse tool call during " + "streaming: %s", + tool_text, + exc_info=True) + + if self.current_tool_index < len(self.streamed_args_for_tool): + self.streamed_args_for_tool[ + self.current_tool_index] += "}" + else: + logger.warning( + "streamed_args_for_tool out of sync: index=%d len=%d", + self.current_tool_index, + len(self.streamed_args_for_tool)) + + result = DeltaMessage(tool_calls=[ + DeltaToolCall( + index=self.current_tool_index, + function=DeltaFunctionCall(arguments="}"), + ) + ]) + self.in_function = False + self.accumulated_params = {} + return result + + return None diff --git a/qwen3_6_scripts/reasoning/__init__.py b/qwen3_6_scripts/reasoning/__init__.py index b32c3f51..9be24806 100644 --- a/qwen3_6_scripts/reasoning/__init__.py +++ b/qwen3_6_scripts/reasoning/__init__.py @@ -1,16 +1,16 @@ -""" +""" Reasoning parser module for vLLM 0.6.3 (BI-V100 / Qwen3.6-35B-A3B adaptation). - -Usage: --reasoning-parser qwen3 -""" - -from vllm.reasoning.abs_reasoning_parsers import ReasoningParser, ReasoningParserManager - -__all__ = ["ReasoningParser", "ReasoningParserManager"] - -# Lazy-register Qwen3 parser; imported on first get_reasoning_parser("qwen3"). -ReasoningParserManager.register_lazy( - "qwen3", - "vllm.reasoning.qwen3_reasoning_parser", - "Qwen3ReasoningParser", -) + +Usage: --reasoning-parser qwen3 +""" + +from vllm.reasoning.abs_reasoning_parsers import ReasoningParser, ReasoningParserManager + +__all__ = ["ReasoningParser", "ReasoningParserManager"] + +# Lazy-register Qwen3 parser; imported on first get_reasoning_parser("qwen3"). +ReasoningParserManager.register_lazy( + "qwen3", + "vllm.reasoning.qwen3_reasoning_parser", + "Qwen3ReasoningParser", +) diff --git a/qwen3_6_scripts/reasoning/abs_reasoning_parsers.py b/qwen3_6_scripts/reasoning/abs_reasoning_parsers.py index 8d91f7fe..c6141074 100644 --- a/qwen3_6_scripts/reasoning/abs_reasoning_parsers.py +++ b/qwen3_6_scripts/reasoning/abs_reasoning_parsers.py @@ -1,243 +1,243 @@ -""" -Abstract reasoning parser base classes for vLLM 0.6.3. -Adapted from vllm-original/vllm/reasoning/abs_reasoning_parsers.py: - - Removed vllm.entrypoints.mcp, vllm.utils.collection_utils, import_utils - - DeltaMessage from vllm 0.6.3 protocol path - - TokenizerLike -> AnyTokenizer - - ReasoningParserManager: simplified eager + lazy registration -""" - -import importlib -from abc import abstractmethod -from collections.abc import Iterable, Sequence -from functools import cached_property -from typing import Any, Optional, TYPE_CHECKING - -if TYPE_CHECKING: - from vllm.entrypoints.openai.protocol import DeltaMessage - from vllm.transformers_utils.tokenizer import AnyTokenizer -else: - DeltaMessage = Any - AnyTokenizer = Any - - -class ReasoningParser: - """Abstract base for all reasoning parsers.""" - - def __init__(self, tokenizer: "AnyTokenizer", *args, **kwargs): - self.model_tokenizer = tokenizer - - @cached_property - def vocab(self) -> dict: - return self.model_tokenizer.get_vocab() - - @abstractmethod - def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: - """Return True once the reasoning block has closed in input_ids.""" - - def is_reasoning_end_streaming( - self, input_ids: Sequence[int], delta_ids: Iterable[int] - ) -> bool: - return self.is_reasoning_end(input_ids) - - @abstractmethod - def extract_content_ids(self, input_ids: list) -> list: - """Return token ids that belong to the content (post-reasoning) part.""" - - def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: - return 0 - - @abstractmethod - def extract_reasoning( - self, model_output: str, request: Any - ) -> "tuple[Optional[str], Optional[str]]": - """ - Split a complete model output into (reasoning_text, content_text). - Either part may be None. - """ - - @abstractmethod - def extract_reasoning_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - ) -> Optional["DeltaMessage"]: - """ - Extract reasoning from a streaming delta. - Returns a DeltaMessage with reasoning_content and/or content set, - or None if this delta should be suppressed (control token). - """ - - -class BaseThinkingReasoningParser(ReasoningParser): - """ - Base for parsers that use ... delimiters. - Subclasses define start_token / end_token properties. - """ - - @property - @abstractmethod - def start_token(self) -> str: - raise NotImplementedError - - @property - @abstractmethod - def end_token(self) -> str: - raise NotImplementedError - - def __init__(self, tokenizer: "AnyTokenizer", *args, **kwargs): - super().__init__(tokenizer, *args, **kwargs) - - if not self.model_tokenizer: - raise ValueError("Tokenizer must be passed to ReasoningParser.") - if not self.start_token or not self.end_token: - raise ValueError("start_token and end_token must be defined.") - - self.start_token_id: Optional[int] = self.vocab.get(self.start_token) - self.end_token_id: Optional[int] = self.vocab.get(self.end_token) - if self.start_token_id is None or self.end_token_id is None: - raise RuntimeError( - f"{self.__class__.__name__}: could not find think tokens " - f"'{self.start_token}'/'{self.end_token}' in tokenizer vocab." - ) - - def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: - for token_id in reversed(input_ids): - if token_id == self.start_token_id: - return False - if token_id == self.end_token_id: - return True - return False - - def is_reasoning_end_streaming( - self, input_ids: Sequence[int], delta_ids: Iterable[int] - ) -> bool: - return self.end_token_id in delta_ids - - def extract_content_ids(self, input_ids: list) -> list: - if self.end_token_id not in input_ids[:-1]: - return [] - return input_ids[input_ids.index(self.end_token_id) + 1:] - - def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: - count = 0 - depth = 0 - for tid in token_ids: - if tid == self.start_token_id: - depth += 1 - elif tid == self.end_token_id: - if depth > 0: - depth -= 1 - elif depth > 0: - count += 1 - return count - - def extract_reasoning( - self, model_output: str, request: Any - ) -> "tuple[Optional[str], Optional[str]]": - # Strip if the model generated it (old-style template). - parts = model_output.partition(self.start_token) - model_output = parts[2] if parts[1] else parts[0] - - if self.end_token not in model_output: - return model_output, None - reasoning, _, content = model_output.partition(self.end_token) - return reasoning, content or None - - def extract_reasoning_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - ) -> Optional["DeltaMessage"]: - from vllm.entrypoints.openai.protocol import DeltaMessage as _DeltaMessage - - # Suppress lone control tokens. - if len(delta_token_ids) == 1 and delta_token_ids[0] in ( - self.start_token_id, self.end_token_id - ): - return None - - start_in_prev = self.start_token_id in previous_token_ids - start_in_delta = self.start_token_id in delta_token_ids - end_in_prev = self.end_token_id in previous_token_ids - end_in_delta = self.end_token_id in delta_token_ids - - if start_in_prev: - if end_in_delta: - end_idx = delta_text.find(self.end_token) - reasoning = delta_text[:end_idx] if end_idx >= 0 else "" - content = delta_text[end_idx + len(self.end_token):] if end_idx >= 0 else None - return _DeltaMessage( - reasoning_content=reasoning or None, - content=content or None, - ) - elif end_in_prev: - return _DeltaMessage(content=delta_text) - else: - return _DeltaMessage(reasoning_content=delta_text) - - elif start_in_delta: - if end_in_delta: - start_idx = delta_text.find(self.start_token) - end_idx = delta_text.find(self.end_token) - reasoning = delta_text[start_idx + len(self.start_token):end_idx] - content = delta_text[end_idx + len(self.end_token):] - return _DeltaMessage( - reasoning_content=reasoning or None, - content=content or None, - ) - else: - return _DeltaMessage(reasoning_content=delta_text) - - else: - return _DeltaMessage(content=delta_text) - - -class ReasoningParserManager: - """ - Registry for ReasoningParser implementations. - Supports eager and lazy registration. - """ - - _parsers: dict = {} # name -> class (eager) - _lazy: dict = {} # name -> (module_path, class_name) - - @classmethod - def register_module(cls, name: str, parser_cls: type) -> None: - """Eagerly register a ReasoningParser class.""" - if not issubclass(parser_cls, ReasoningParser): - raise TypeError(f"{parser_cls} is not a ReasoningParser subclass.") - cls._parsers[name] = parser_cls - - @classmethod - def register_lazy(cls, name: str, module_path: str, class_name: str) -> None: - """Register a parser for deferred import.""" - cls._lazy[name] = (module_path, class_name) - - @classmethod - def get_reasoning_parser(cls, name: str) -> type: - if name in cls._parsers: - return cls._parsers[name] - if name in cls._lazy: - module_path, class_name = cls._lazy[name] - mod = importlib.import_module(module_path) - parser_cls = getattr(mod, class_name) - cls._parsers[name] = parser_cls - return parser_cls - registered = sorted(set(cls._parsers) | set(cls._lazy)) - raise KeyError( - f"Reasoning parser '{name}' not found. " - f"Available: {registered}" - ) - - @classmethod - def list_registered(cls) -> list: - return sorted(set(cls._parsers) | set(cls._lazy)) +""" +Abstract reasoning parser base classes for vLLM 0.6.3. +Adapted from vllm-original/vllm/reasoning/abs_reasoning_parsers.py: + - Removed vllm.entrypoints.mcp, vllm.utils.collection_utils, import_utils + - DeltaMessage from vllm 0.6.3 protocol path + - TokenizerLike -> AnyTokenizer + - ReasoningParserManager: simplified eager + lazy registration +""" + +import importlib +from abc import abstractmethod +from collections.abc import Iterable, Sequence +from functools import cached_property +from typing import Any, Optional, TYPE_CHECKING + +if TYPE_CHECKING: + from vllm.entrypoints.openai.protocol import DeltaMessage + from vllm.transformers_utils.tokenizer import AnyTokenizer +else: + DeltaMessage = Any + AnyTokenizer = Any + + +class ReasoningParser: + """Abstract base for all reasoning parsers.""" + + def __init__(self, tokenizer: "AnyTokenizer", *args, **kwargs): + self.model_tokenizer = tokenizer + + @cached_property + def vocab(self) -> dict: + return self.model_tokenizer.get_vocab() + + @abstractmethod + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: + """Return True once the reasoning block has closed in input_ids.""" + + def is_reasoning_end_streaming( + self, input_ids: Sequence[int], delta_ids: Iterable[int] + ) -> bool: + return self.is_reasoning_end(input_ids) + + @abstractmethod + def extract_content_ids(self, input_ids: list) -> list: + """Return token ids that belong to the content (post-reasoning) part.""" + + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + return 0 + + @abstractmethod + def extract_reasoning( + self, model_output: str, request: Any + ) -> "tuple[Optional[str], Optional[str]]": + """ + Split a complete model output into (reasoning_text, content_text). + Either part may be None. + """ + + @abstractmethod + def extract_reasoning_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + ) -> Optional["DeltaMessage"]: + """ + Extract reasoning from a streaming delta. + Returns a DeltaMessage with reasoning_content and/or content set, + or None if this delta should be suppressed (control token). + """ + + +class BaseThinkingReasoningParser(ReasoningParser): + """ + Base for parsers that use ... delimiters. + Subclasses define start_token / end_token properties. + """ + + @property + @abstractmethod + def start_token(self) -> str: + raise NotImplementedError + + @property + @abstractmethod + def end_token(self) -> str: + raise NotImplementedError + + def __init__(self, tokenizer: "AnyTokenizer", *args, **kwargs): + super().__init__(tokenizer, *args, **kwargs) + + if not self.model_tokenizer: + raise ValueError("Tokenizer must be passed to ReasoningParser.") + if not self.start_token or not self.end_token: + raise ValueError("start_token and end_token must be defined.") + + self.start_token_id: Optional[int] = self.vocab.get(self.start_token) + self.end_token_id: Optional[int] = self.vocab.get(self.end_token) + if self.start_token_id is None or self.end_token_id is None: + raise RuntimeError( + f"{self.__class__.__name__}: could not find think tokens " + f"'{self.start_token}'/'{self.end_token}' in tokenizer vocab." + ) + + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: + for token_id in reversed(input_ids): + if token_id == self.start_token_id: + return False + if token_id == self.end_token_id: + return True + return False + + def is_reasoning_end_streaming( + self, input_ids: Sequence[int], delta_ids: Iterable[int] + ) -> bool: + return self.end_token_id in delta_ids + + def extract_content_ids(self, input_ids: list) -> list: + if self.end_token_id not in input_ids[:-1]: + return [] + return input_ids[input_ids.index(self.end_token_id) + 1:] + + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + count = 0 + depth = 0 + for tid in token_ids: + if tid == self.start_token_id: + depth += 1 + elif tid == self.end_token_id: + if depth > 0: + depth -= 1 + elif depth > 0: + count += 1 + return count + + def extract_reasoning( + self, model_output: str, request: Any + ) -> "tuple[Optional[str], Optional[str]]": + # Strip if the model generated it (old-style template). + parts = model_output.partition(self.start_token) + model_output = parts[2] if parts[1] else parts[0] + + if self.end_token not in model_output: + return model_output, None + reasoning, _, content = model_output.partition(self.end_token) + return reasoning, content or None + + def extract_reasoning_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + ) -> Optional["DeltaMessage"]: + from vllm.entrypoints.openai.protocol import DeltaMessage as _DeltaMessage + + # Suppress lone control tokens. + if len(delta_token_ids) == 1 and delta_token_ids[0] in ( + self.start_token_id, self.end_token_id + ): + return None + + start_in_prev = self.start_token_id in previous_token_ids + start_in_delta = self.start_token_id in delta_token_ids + end_in_prev = self.end_token_id in previous_token_ids + end_in_delta = self.end_token_id in delta_token_ids + + if start_in_prev: + if end_in_delta: + end_idx = delta_text.find(self.end_token) + reasoning = delta_text[:end_idx] if end_idx >= 0 else "" + content = delta_text[end_idx + len(self.end_token):] if end_idx >= 0 else None + return _DeltaMessage( + reasoning_content=reasoning or None, + content=content or None, + ) + elif end_in_prev: + return _DeltaMessage(content=delta_text) + else: + return _DeltaMessage(reasoning_content=delta_text) + + elif start_in_delta: + if end_in_delta: + start_idx = delta_text.find(self.start_token) + end_idx = delta_text.find(self.end_token) + reasoning = delta_text[start_idx + len(self.start_token):end_idx] + content = delta_text[end_idx + len(self.end_token):] + return _DeltaMessage( + reasoning_content=reasoning or None, + content=content or None, + ) + else: + return _DeltaMessage(reasoning_content=delta_text) + + else: + return _DeltaMessage(content=delta_text) + + +class ReasoningParserManager: + """ + Registry for ReasoningParser implementations. + Supports eager and lazy registration. + """ + + _parsers: dict = {} # name -> class (eager) + _lazy: dict = {} # name -> (module_path, class_name) + + @classmethod + def register_module(cls, name: str, parser_cls: type) -> None: + """Eagerly register a ReasoningParser class.""" + if not issubclass(parser_cls, ReasoningParser): + raise TypeError(f"{parser_cls} is not a ReasoningParser subclass.") + cls._parsers[name] = parser_cls + + @classmethod + def register_lazy(cls, name: str, module_path: str, class_name: str) -> None: + """Register a parser for deferred import.""" + cls._lazy[name] = (module_path, class_name) + + @classmethod + def get_reasoning_parser(cls, name: str) -> type: + if name in cls._parsers: + return cls._parsers[name] + if name in cls._lazy: + module_path, class_name = cls._lazy[name] + mod = importlib.import_module(module_path) + parser_cls = getattr(mod, class_name) + cls._parsers[name] = parser_cls + return parser_cls + registered = sorted(set(cls._parsers) | set(cls._lazy)) + raise KeyError( + f"Reasoning parser '{name}' not found. " + f"Available: {registered}" + ) + + @classmethod + def list_registered(cls) -> list: + return sorted(set(cls._parsers) | set(cls._lazy)) diff --git a/qwen3_6_scripts/reasoning/qwen3_reasoning_parser.py b/qwen3_6_scripts/reasoning/qwen3_reasoning_parser.py index c6ad6cad..f7fddfec 100644 --- a/qwen3_6_scripts/reasoning/qwen3_reasoning_parser.py +++ b/qwen3_6_scripts/reasoning/qwen3_reasoning_parser.py @@ -1,40 +1,40 @@ -""" -Reasoning parser for Qwen3 / Qwen3.5 / Qwen3.6 model family. -Adapted from vllm-original/vllm/reasoning/qwen3_reasoning_parser.py. - -The model uses ... to wrap chain-of-thought output. -For Qwen3.5+ the chat template injects into the prompt, so only - appears in the generated tokens; older templates generate -themselves. Both styles are handled. -""" - -from typing import Optional, Sequence, Any - -from vllm.reasoning.abs_reasoning_parsers import ( - BaseThinkingReasoningParser, - ReasoningParserManager, -) - - -class Qwen3ReasoningParser(BaseThinkingReasoningParser): - - def __init__(self, tokenizer: Any, *args, **kwargs): - super().__init__(tokenizer, *args, **kwargs) - chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} - self.thinking_enabled = chat_kwargs.get("enable_thinking", True) - - @property - def start_token(self) -> str: - return "" - - @property - def end_token(self) -> str: - return "" - - def extract_reasoning( - self, model_output: str, request: Any - ) -> "tuple[Optional[str], Optional[str]]": - # Strip if the model generated it (old template / edge case). +""" +Reasoning parser for Qwen3 / Qwen3.5 / Qwen3.6 model family. +Adapted from vllm-original/vllm/reasoning/qwen3_reasoning_parser.py. + +The model uses ... to wrap chain-of-thought output. +For Qwen3.5+ the chat template injects into the prompt, so only + appears in the generated tokens; older templates generate +themselves. Both styles are handled. +""" + +from typing import Optional, Sequence, Any + +from vllm.reasoning.abs_reasoning_parsers import ( + BaseThinkingReasoningParser, + ReasoningParserManager, +) + + +class Qwen3ReasoningParser(BaseThinkingReasoningParser): + + def __init__(self, tokenizer: Any, *args, **kwargs): + super().__init__(tokenizer, *args, **kwargs) + chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} + self.thinking_enabled = chat_kwargs.get("enable_thinking", True) + + @property + def start_token(self) -> str: + return "" + + @property + def end_token(self) -> str: + return "" + + def extract_reasoning( + self, model_output: str, request: Any + ) -> "tuple[Optional[str], Optional[str]]": + # Strip if the model generated it (old template / edge case). parts = model_output.partition(self.start_token) model_output = parts[2] if parts[1] else parts[0] @@ -47,66 +47,66 @@ class Qwen3ReasoningParser(BaseThinkingReasoningParser): if self.end_token not in model_output: # Thinking enabled but output truncated before . return model_output, None - - reasoning, _, content = model_output.partition(self.end_token) - return reasoning, content or None - - def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: - token_ids = list(token_ids) - if self.start_token_id in token_ids: - # Old-style template: model generates itself. - # Use depth-counting from the base class. - return super().count_reasoning_tokens(token_ids) - elif self.end_token_id in token_ids: - # New-style template (Qwen3.5+): is injected into the - # prompt, so output starts already inside the thinking block. - # Every token before is a reasoning token. - return token_ids.index(self.end_token_id) - else: - # No in output: either truncated (all reasoning) - # or thinking disabled (none). - return len(token_ids) if self.thinking_enabled else 0 - - def extract_reasoning_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - ): - from vllm.entrypoints.openai.protocol import DeltaMessage - - if not self.thinking_enabled: - return DeltaMessage(content=delta_text) if delta_text else None - - # Strip from delta if the model generates it itself. - if self.start_token_id in delta_token_ids: - start_idx = delta_text.find(self.start_token) - if start_idx >= 0: - delta_text = delta_text[start_idx + len(self.start_token):] - - if self.end_token_id in delta_token_ids: - end_idx = delta_text.find(self.end_token) - if end_idx >= 0: - reasoning = delta_text[:end_idx] - content = delta_text[end_idx + len(self.end_token):] - if not reasoning and not content: - return None - return DeltaMessage( - reasoning_content=reasoning or None, - content=content or None, - ) - return None - - if not delta_text: - return None - elif self.end_token_id in previous_token_ids: - return DeltaMessage(content=delta_text) - else: - return DeltaMessage(reasoning_content=delta_text) - - -# Register immediately when this module is imported. -ReasoningParserManager.register_module("qwen3", Qwen3ReasoningParser) + + reasoning, _, content = model_output.partition(self.end_token) + return reasoning, content or None + + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + token_ids = list(token_ids) + if self.start_token_id in token_ids: + # Old-style template: model generates itself. + # Use depth-counting from the base class. + return super().count_reasoning_tokens(token_ids) + elif self.end_token_id in token_ids: + # New-style template (Qwen3.5+): is injected into the + # prompt, so output starts already inside the thinking block. + # Every token before is a reasoning token. + return token_ids.index(self.end_token_id) + else: + # No in output: either truncated (all reasoning) + # or thinking disabled (none). + return len(token_ids) if self.thinking_enabled else 0 + + def extract_reasoning_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + ): + from vllm.entrypoints.openai.protocol import DeltaMessage + + if not self.thinking_enabled: + return DeltaMessage(content=delta_text) if delta_text else None + + # Strip from delta if the model generates it itself. + if self.start_token_id in delta_token_ids: + start_idx = delta_text.find(self.start_token) + if start_idx >= 0: + delta_text = delta_text[start_idx + len(self.start_token):] + + if self.end_token_id in delta_token_ids: + end_idx = delta_text.find(self.end_token) + if end_idx >= 0: + reasoning = delta_text[:end_idx] + content = delta_text[end_idx + len(self.end_token):] + if not reasoning and not content: + return None + return DeltaMessage( + reasoning_content=reasoning or None, + content=content or None, + ) + return None + + if not delta_text: + return None + elif self.end_token_id in previous_token_ids: + return DeltaMessage(content=delta_text) + else: + return DeltaMessage(reasoning_content=delta_text) + + +# Register immediately when this module is imported. +ReasoningParserManager.register_module("qwen3", Qwen3ReasoningParser) diff --git a/qwen3_6_scripts/scheduler.py b/qwen3_6_scripts/scheduler.py index ac5155aa..2f27246f 100644 --- a/qwen3_6_scripts/scheduler.py +++ b/qwen3_6_scripts/scheduler.py @@ -1,20 +1,20 @@ -import enum -import os -import random -import time +import enum +import os +import random +import time from collections import deque -from dataclasses import dataclass, field -from typing import (Callable, Deque, Dict, Iterable, List, Optional, Set, - Tuple, Union) - -from vllm.config import CacheConfig, LoRAConfig, SchedulerConfig -from vllm.core.interfaces import AllocStatus, BlockSpaceManager -from vllm.logger import init_logger -from vllm.lora.request import LoRARequest -from vllm.prompt_adapter.request import PromptAdapterRequest -from vllm.sequence import (Sequence, SequenceData, SequenceGroup, - SequenceGroupMetadata, SequenceGroupMetadataDelta, - SequenceStatus) +from dataclasses import dataclass, field +from typing import (Callable, Deque, Dict, Iterable, List, Optional, Set, + Tuple, Union) + +from vllm.config import CacheConfig, LoRAConfig, SchedulerConfig +from vllm.core.interfaces import AllocStatus, BlockSpaceManager +from vllm.logger import init_logger +from vllm.lora.request import LoRARequest +from vllm.prompt_adapter.request import PromptAdapterRequest +from vllm.sequence import (Sequence, SequenceData, SequenceGroup, + SequenceGroupMetadata, SequenceGroupMetadataDelta, + SequenceStatus) from vllm.utils import Device, PyObjectCache try: @@ -38,14 +38,14 @@ except ImportError: # Local source-tree tests. gdn_restore_alignment, gdn_restore_mode_from_env, keys_from_block_hashes, restore_key_is_eligible, strict_prefix_block_count) - -logger = init_logger(__name__) - -# Test-only. If configured, decode is preempted with -# ARTIFICIAL_PREEMPTION_PROB% probability. -ENABLE_ARTIFICIAL_PREEMPT = bool( - os.getenv("VLLM_TEST_ENABLE_ARTIFICIAL_PREEMPT", False)) # noqa -ARTIFICIAL_PREEMPTION_PROB = 0.5 + +logger = init_logger(__name__) + +# Test-only. If configured, decode is preempted with +# ARTIFICIAL_PREEMPTION_PROB% probability. +ENABLE_ARTIFICIAL_PREEMPT = bool( + os.getenv("VLLM_TEST_ENABLE_ARTIFICIAL_PREEMPT", False)) # noqa +ARTIFICIAL_PREEMPTION_PROB = 0.5 ARTIFICIAL_PREEMPTION_MAX_CNT = 500 @@ -88,34 +88,34 @@ def _plan_gdn_prefix_fast_forward( or physical_query_tokens >= nominal_chunk_size)): return fallback return logical_chunk_size, physical_query_tokens - - -class PreemptionMode(enum.Enum): - """Preemption modes. - - 1. Swapping: Swap out the blocks of the preempted sequences to CPU memory - and swap them back in when the sequences are resumed. - 2. Recomputation: Discard the blocks of the preempted sequences and - recompute them when the sequences are resumed, treating the sequences as - new prompts. - """ - SWAP = enum.auto() - RECOMPUTE = enum.auto() - - -@dataclass -class SchedulingBudget: - """The available slots for scheduling. - - TODO(sang): Right now, the budget is request_id-aware meaning it can ignore - budget update from the same request_id. It is because in normal scheduling - path, we update RUNNING num_seqs ahead of time, meaning it could be - updated more than once when scheduling RUNNING requests. Since this won't - happen if we only have chunked prefill scheduling, we can remove this - feature from the API when chunked prefill is enabled by default. - """ - token_budget: int - max_num_seqs: int + + +class PreemptionMode(enum.Enum): + """Preemption modes. + + 1. Swapping: Swap out the blocks of the preempted sequences to CPU memory + and swap them back in when the sequences are resumed. + 2. Recomputation: Discard the blocks of the preempted sequences and + recompute them when the sequences are resumed, treating the sequences as + new prompts. + """ + SWAP = enum.auto() + RECOMPUTE = enum.auto() + + +@dataclass +class SchedulingBudget: + """The available slots for scheduling. + + TODO(sang): Right now, the budget is request_id-aware meaning it can ignore + budget update from the same request_id. It is because in normal scheduling + path, we update RUNNING num_seqs ahead of time, meaning it could be + updated more than once when scheduling RUNNING requests. Since this won't + happen if we only have chunked prefill scheduling, we can remove this + feature from the API when chunked prefill is enabled by default. + """ + token_budget: int + max_num_seqs: int _request_ids_num_batched_tokens: Set[str] = field(default_factory=set) _request_ids_num_curr_seqs: Set[str] = field(default_factory=set) _num_batched_tokens: int = 0 @@ -123,16 +123,16 @@ class SchedulingBudget: _request_num_scheduled_tokens: Dict[str, int] = field( default_factory=dict) _num_curr_seqs: int = 0 - - def can_schedule(self, *, num_new_tokens: int, num_new_seqs: int): - assert num_new_tokens != 0 - assert num_new_seqs != 0 - return (self.num_batched_tokens + num_new_tokens <= self.token_budget - and self.num_curr_seqs + num_new_seqs <= self.max_num_seqs) - - def remaining_token_budget(self): - return self.token_budget - self.num_batched_tokens - + + def can_schedule(self, *, num_new_tokens: int, num_new_seqs: int): + assert num_new_tokens != 0 + assert num_new_seqs != 0 + return (self.num_batched_tokens + num_new_tokens <= self.token_budget + and self.num_curr_seqs + num_new_seqs <= self.max_num_seqs) + + def remaining_token_budget(self): + return self.token_budget - self.num_batched_tokens + def add_num_batched_tokens( self, req_id: str, num_batched_tokens: int, num_scheduled_tokens: Optional[int] = None): @@ -145,290 +145,290 @@ class SchedulingBudget: self._num_batched_tokens += num_batched_tokens self._num_scheduled_tokens += num_scheduled_tokens self._request_num_scheduled_tokens[req_id] = num_scheduled_tokens - - def subtract_num_batched_tokens(self, req_id: str, - num_batched_tokens: int): + + def subtract_num_batched_tokens(self, req_id: str, + num_batched_tokens: int): if req_id in self._request_ids_num_batched_tokens: self._request_ids_num_batched_tokens.remove(req_id) self._num_batched_tokens -= num_batched_tokens self._num_scheduled_tokens -= ( self._request_num_scheduled_tokens.pop(req_id)) - - def add_num_seqs(self, req_id: str, num_curr_seqs: int): - if req_id in self._request_ids_num_curr_seqs: - return - - self._request_ids_num_curr_seqs.add(req_id) - self._num_curr_seqs += num_curr_seqs - - def subtract_num_seqs(self, req_id: str, num_curr_seqs: int): - if req_id in self._request_ids_num_curr_seqs: - self._request_ids_num_curr_seqs.remove(req_id) - self._num_curr_seqs -= num_curr_seqs - - @property + + def add_num_seqs(self, req_id: str, num_curr_seqs: int): + if req_id in self._request_ids_num_curr_seqs: + return + + self._request_ids_num_curr_seqs.add(req_id) + self._num_curr_seqs += num_curr_seqs + + def subtract_num_seqs(self, req_id: str, num_curr_seqs: int): + if req_id in self._request_ids_num_curr_seqs: + self._request_ids_num_curr_seqs.remove(req_id) + self._num_curr_seqs -= num_curr_seqs + + @property def num_batched_tokens(self): return self._num_batched_tokens @property def num_scheduled_tokens(self): return self._num_scheduled_tokens - - @property - def num_curr_seqs(self): - return self._num_curr_seqs - - -@dataclass -class ScheduledSequenceGroup: - # A sequence group that's scheduled. - seq_group: SequenceGroup - # The total chunk size (number of tokens) to process for next iteration. - # 1 for decoding. Same as prompt tokens for prefill, but if prefill is - # chunked, it can be smaller than that. - token_chunk_size: int - - -@dataclass -class SchedulerOutputs: - """The scheduling decision made from a scheduler.""" - # Scheduled sequence groups. - scheduled_seq_groups: Iterable[ScheduledSequenceGroup] - # Number of prefill groups scheduled. - num_prefill_groups: int - # Total number of batched tokens. - num_batched_tokens: int - # Blocks to swap in. List of CPU -> GPU block number. - blocks_to_swap_in: List[Tuple[int, int]] - # Blocks to swap out. List of GPU -> CPU block number. - blocks_to_swap_out: List[Tuple[int, int]] - # Blocks to copy. Source to dest block. - blocks_to_copy: List[Tuple[int, int]] - # Sequence groups that are going to be ignored. - ignored_seq_groups: List[SequenceGroup] - # The number of slots for lookahead decoding. - num_lookahead_slots: int - # The number of requests in the running queue - running_queue_size: int + + @property + def num_curr_seqs(self): + return self._num_curr_seqs + + +@dataclass +class ScheduledSequenceGroup: + # A sequence group that's scheduled. + seq_group: SequenceGroup + # The total chunk size (number of tokens) to process for next iteration. + # 1 for decoding. Same as prompt tokens for prefill, but if prefill is + # chunked, it can be smaller than that. + token_chunk_size: int + + +@dataclass +class SchedulerOutputs: + """The scheduling decision made from a scheduler.""" + # Scheduled sequence groups. + scheduled_seq_groups: Iterable[ScheduledSequenceGroup] + # Number of prefill groups scheduled. + num_prefill_groups: int + # Total number of batched tokens. + num_batched_tokens: int + # Blocks to swap in. List of CPU -> GPU block number. + blocks_to_swap_in: List[Tuple[int, int]] + # Blocks to swap out. List of GPU -> CPU block number. + blocks_to_swap_out: List[Tuple[int, int]] + # Blocks to copy. Source to dest block. + blocks_to_copy: List[Tuple[int, int]] + # Sequence groups that are going to be ignored. + ignored_seq_groups: List[SequenceGroup] + # The number of slots for lookahead decoding. + num_lookahead_slots: int + # The number of requests in the running queue + running_queue_size: int preempted: int def __post_init__(self): # Request-level preemption cannot swap both ways in one step. The # content-addressed CPU tier appends its ordered maps after creation. assert not (self.blocks_to_swap_in and self.blocks_to_swap_out) - - self.num_loras: int = len(self.lora_requests) - if self.num_loras > 0: - self._sort_by_lora_ids() - - self.num_prompt_adapters: int = len(self.prompt_adapter_requests) - - def is_empty(self) -> bool: - # NOTE: We do not consider the ignored sequence groups. - return (not self.scheduled_seq_groups and not self.blocks_to_swap_in - and not self.blocks_to_swap_out and not self.blocks_to_copy) - - def _sort_by_lora_ids(self): - self.scheduled_seq_groups = sorted( - self.scheduled_seq_groups, - key=lambda g: (g.seq_group.lora_int_id, g.seq_group.request_id)) - - @property - def lora_requests(self) -> Set[LoRARequest]: - return { - g.seq_group.lora_request - for g in self.scheduled_seq_groups - if g.seq_group.lora_request is not None - } - - @property - def prompt_adapter_requests(self) -> Set[PromptAdapterRequest]: - return { - g.seq_group.prompt_adapter_request - for g in self.scheduled_seq_groups - if g.seq_group.prompt_adapter_request is not None - } - - -@dataclass -class SchedulerRunningOutputs: - """The requests that are scheduled from a running queue. - - Could contain prefill (prefill that's chunked) or decodes. If there's not - enough memory, it can be preempted (for recompute) or swapped out. - """ - # Selected sequences that are running and in a decoding phase. - decode_seq_groups: List[ScheduledSequenceGroup] - # Selected sequences that are running and in a prefill phase. - # I.e., it means the prefill has been chunked. - prefill_seq_groups: List[ScheduledSequenceGroup] - # The preempted sequences. - preempted: List[SequenceGroup] - # Sequences that are swapped out. - swapped_out: List[SequenceGroup] - # The blocks to swap out. - blocks_to_swap_out: List[Tuple[int, int]] - # The blocks to copy. - blocks_to_copy: List[Tuple[int, int]] - # The number of slots for lookahead decoding. - num_lookahead_slots: int - - # Optimization for fast-access to seq_group lists - decode_seq_groups_list: List[SequenceGroup] - prefill_seq_groups_list: List[SequenceGroup] - - @classmethod - def create_empty(cls) -> "SchedulerRunningOutputs": - return SchedulerRunningOutputs( - decode_seq_groups=[], - prefill_seq_groups=[], - preempted=[], - swapped_out=[], - blocks_to_swap_out=[], - blocks_to_copy=[], - num_lookahead_slots=0, - decode_seq_groups_list=[], - prefill_seq_groups_list=[], - ) - - -@dataclass -class SchedulerSwappedInOutputs: - """The requests that are scheduled from a swap queue. - - Could contain prefill (prefill that's chunked) or decodes. - """ - # Selected sequences that are going to be swapped in and is in a - # decoding phase. - decode_seq_groups: List[ScheduledSequenceGroup] - # Selected sequences that are going to be swapped in and in a prefill - # phase. I.e., it means the prefill has been chunked. - prefill_seq_groups: List[ScheduledSequenceGroup] - # The blocks to swap in. - blocks_to_swap_in: List[Tuple[int, int]] - # The blocks to copy. - blocks_to_copy: List[Tuple[int, int]] - # The number of slots for lookahead decoding. - num_lookahead_slots: int - # Infeasible sequence groups. - infeasible_seq_groups: List[SequenceGroup] - - @classmethod - def create_empty(cls) -> "SchedulerSwappedInOutputs": - return SchedulerSwappedInOutputs( - decode_seq_groups=[], - prefill_seq_groups=[], - blocks_to_swap_in=[], - blocks_to_copy=[], - num_lookahead_slots=0, - infeasible_seq_groups=[], - ) - - -@dataclass -class SchedulerPrefillOutputs: - """The requests that are scheduled from a waiting queue. - - Could contain a fresh prefill requests or preempted requests that need - to be recomputed from scratch. - """ - # Selected sequences for prefill. - seq_groups: List[ScheduledSequenceGroup] - # Ignored sequence groups. - ignored_seq_groups: List[SequenceGroup] - num_lookahead_slots: int - - @classmethod - def create_empty(cls) -> "SchedulerPrefillOutputs": - return SchedulerPrefillOutputs( - seq_groups=[], - ignored_seq_groups=[], - num_lookahead_slots=0, - ) - - -def seq_group_metadata_builder(): - return SequenceGroupMetadata(request_id="", - is_prompt=False, - seq_data={}, - sampling_params=None, - block_tables={}) - - -def scheduler_running_outputs_builder(): - return SchedulerRunningOutputs(decode_seq_groups=[], - prefill_seq_groups=[], - preempted=[], - swapped_out=[], - blocks_to_swap_out=[], - blocks_to_copy=[], - num_lookahead_slots=0, - prefill_seq_groups_list=[], - decode_seq_groups_list=[]) - - -def scheduled_seq_group_builder(): - return ScheduledSequenceGroup(SequenceGroup("", [], -1), - token_chunk_size=0) - # return ScheduledSequenceGroup(seq_group=None, token_chunk_size=0) - - -class Scheduler: - - def __init__( - self, - scheduler_config: SchedulerConfig, - cache_config: CacheConfig, - lora_config: Optional[LoRAConfig], - pipeline_parallel_size: int = 1, - output_proc_callback: Optional[Callable] = None, - ) -> None: - self.scheduler_config = scheduler_config - self.cache_config = cache_config - # Note for LoRA scheduling: the current policy is extremely - # simple and NOT fair. It can lead to starvation of some - # LoRAs. This should be improved in the future. - self.lora_config = lora_config - - version = "v1" - if self.scheduler_config.use_v2_block_manager: - version = "v2" - if (self.scheduler_config.embedding_mode - or self.cache_config.is_attention_free): - version = "placeholder" - - BlockSpaceManagerImpl = BlockSpaceManager.get_block_space_manager_class( - version) - - num_gpu_blocks = cache_config.num_gpu_blocks - if num_gpu_blocks: - num_gpu_blocks //= pipeline_parallel_size - - num_cpu_blocks = cache_config.num_cpu_blocks - if num_cpu_blocks: - num_cpu_blocks //= pipeline_parallel_size - - # Create the block space manager. - self.block_manager = BlockSpaceManagerImpl( - block_size=self.cache_config.block_size, - num_gpu_blocks=num_gpu_blocks, - num_cpu_blocks=num_cpu_blocks, - sliding_window=self.cache_config.sliding_window, - enable_caching=self.cache_config.enable_prefix_caching) - - # Sequence groups in the WAITING state. - # Contain new prefill or preempted requests. - self.waiting: Deque[SequenceGroup] = deque() - # Sequence groups in the RUNNING state. - # Contain decode requests. - self.running: Deque[SequenceGroup] = deque() - # Sequence groups in the SWAPPED state. - # Contain decode requests that are swapped out. - self.swapped: Deque[SequenceGroup] = deque() - # Sequence groups finished requests ids since last step iteration. - # It lets the model know that any state associated with these requests - # can and must be released after the current step. - # This is used to evict the finished requests from the Mamba cache. + + self.num_loras: int = len(self.lora_requests) + if self.num_loras > 0: + self._sort_by_lora_ids() + + self.num_prompt_adapters: int = len(self.prompt_adapter_requests) + + def is_empty(self) -> bool: + # NOTE: We do not consider the ignored sequence groups. + return (not self.scheduled_seq_groups and not self.blocks_to_swap_in + and not self.blocks_to_swap_out and not self.blocks_to_copy) + + def _sort_by_lora_ids(self): + self.scheduled_seq_groups = sorted( + self.scheduled_seq_groups, + key=lambda g: (g.seq_group.lora_int_id, g.seq_group.request_id)) + + @property + def lora_requests(self) -> Set[LoRARequest]: + return { + g.seq_group.lora_request + for g in self.scheduled_seq_groups + if g.seq_group.lora_request is not None + } + + @property + def prompt_adapter_requests(self) -> Set[PromptAdapterRequest]: + return { + g.seq_group.prompt_adapter_request + for g in self.scheduled_seq_groups + if g.seq_group.prompt_adapter_request is not None + } + + +@dataclass +class SchedulerRunningOutputs: + """The requests that are scheduled from a running queue. + + Could contain prefill (prefill that's chunked) or decodes. If there's not + enough memory, it can be preempted (for recompute) or swapped out. + """ + # Selected sequences that are running and in a decoding phase. + decode_seq_groups: List[ScheduledSequenceGroup] + # Selected sequences that are running and in a prefill phase. + # I.e., it means the prefill has been chunked. + prefill_seq_groups: List[ScheduledSequenceGroup] + # The preempted sequences. + preempted: List[SequenceGroup] + # Sequences that are swapped out. + swapped_out: List[SequenceGroup] + # The blocks to swap out. + blocks_to_swap_out: List[Tuple[int, int]] + # The blocks to copy. + blocks_to_copy: List[Tuple[int, int]] + # The number of slots for lookahead decoding. + num_lookahead_slots: int + + # Optimization for fast-access to seq_group lists + decode_seq_groups_list: List[SequenceGroup] + prefill_seq_groups_list: List[SequenceGroup] + + @classmethod + def create_empty(cls) -> "SchedulerRunningOutputs": + return SchedulerRunningOutputs( + decode_seq_groups=[], + prefill_seq_groups=[], + preempted=[], + swapped_out=[], + blocks_to_swap_out=[], + blocks_to_copy=[], + num_lookahead_slots=0, + decode_seq_groups_list=[], + prefill_seq_groups_list=[], + ) + + +@dataclass +class SchedulerSwappedInOutputs: + """The requests that are scheduled from a swap queue. + + Could contain prefill (prefill that's chunked) or decodes. + """ + # Selected sequences that are going to be swapped in and is in a + # decoding phase. + decode_seq_groups: List[ScheduledSequenceGroup] + # Selected sequences that are going to be swapped in and in a prefill + # phase. I.e., it means the prefill has been chunked. + prefill_seq_groups: List[ScheduledSequenceGroup] + # The blocks to swap in. + blocks_to_swap_in: List[Tuple[int, int]] + # The blocks to copy. + blocks_to_copy: List[Tuple[int, int]] + # The number of slots for lookahead decoding. + num_lookahead_slots: int + # Infeasible sequence groups. + infeasible_seq_groups: List[SequenceGroup] + + @classmethod + def create_empty(cls) -> "SchedulerSwappedInOutputs": + return SchedulerSwappedInOutputs( + decode_seq_groups=[], + prefill_seq_groups=[], + blocks_to_swap_in=[], + blocks_to_copy=[], + num_lookahead_slots=0, + infeasible_seq_groups=[], + ) + + +@dataclass +class SchedulerPrefillOutputs: + """The requests that are scheduled from a waiting queue. + + Could contain a fresh prefill requests or preempted requests that need + to be recomputed from scratch. + """ + # Selected sequences for prefill. + seq_groups: List[ScheduledSequenceGroup] + # Ignored sequence groups. + ignored_seq_groups: List[SequenceGroup] + num_lookahead_slots: int + + @classmethod + def create_empty(cls) -> "SchedulerPrefillOutputs": + return SchedulerPrefillOutputs( + seq_groups=[], + ignored_seq_groups=[], + num_lookahead_slots=0, + ) + + +def seq_group_metadata_builder(): + return SequenceGroupMetadata(request_id="", + is_prompt=False, + seq_data={}, + sampling_params=None, + block_tables={}) + + +def scheduler_running_outputs_builder(): + return SchedulerRunningOutputs(decode_seq_groups=[], + prefill_seq_groups=[], + preempted=[], + swapped_out=[], + blocks_to_swap_out=[], + blocks_to_copy=[], + num_lookahead_slots=0, + prefill_seq_groups_list=[], + decode_seq_groups_list=[]) + + +def scheduled_seq_group_builder(): + return ScheduledSequenceGroup(SequenceGroup("", [], -1), + token_chunk_size=0) + # return ScheduledSequenceGroup(seq_group=None, token_chunk_size=0) + + +class Scheduler: + + def __init__( + self, + scheduler_config: SchedulerConfig, + cache_config: CacheConfig, + lora_config: Optional[LoRAConfig], + pipeline_parallel_size: int = 1, + output_proc_callback: Optional[Callable] = None, + ) -> None: + self.scheduler_config = scheduler_config + self.cache_config = cache_config + # Note for LoRA scheduling: the current policy is extremely + # simple and NOT fair. It can lead to starvation of some + # LoRAs. This should be improved in the future. + self.lora_config = lora_config + + version = "v1" + if self.scheduler_config.use_v2_block_manager: + version = "v2" + if (self.scheduler_config.embedding_mode + or self.cache_config.is_attention_free): + version = "placeholder" + + BlockSpaceManagerImpl = BlockSpaceManager.get_block_space_manager_class( + version) + + num_gpu_blocks = cache_config.num_gpu_blocks + if num_gpu_blocks: + num_gpu_blocks //= pipeline_parallel_size + + num_cpu_blocks = cache_config.num_cpu_blocks + if num_cpu_blocks: + num_cpu_blocks //= pipeline_parallel_size + + # Create the block space manager. + self.block_manager = BlockSpaceManagerImpl( + block_size=self.cache_config.block_size, + num_gpu_blocks=num_gpu_blocks, + num_cpu_blocks=num_cpu_blocks, + sliding_window=self.cache_config.sliding_window, + enable_caching=self.cache_config.enable_prefix_caching) + + # Sequence groups in the WAITING state. + # Contain new prefill or preempted requests. + self.waiting: Deque[SequenceGroup] = deque() + # Sequence groups in the RUNNING state. + # Contain decode requests. + self.running: Deque[SequenceGroup] = deque() + # Sequence groups in the SWAPPED state. + # Contain decode requests that are swapped out. + self.swapped: Deque[SequenceGroup] = deque() + # Sequence groups finished requests ids since last step iteration. + # It lets the model know that any state associated with these requests + # can and must be released after the current step. + # This is used to evict the finished requests from the Mamba cache. self._finished_requests_ids: List[str] = list() self._gdn_prefix_policy = GdnPrefixStatePolicy( gdn_cache_policy_from_env()) @@ -443,73 +443,73 @@ class Scheduler: str, Optional[GdnPrefixKey]] = {} self._gdn_request_capture_targets: Dict[ str, Tuple[GdnPrefixKey, ...]] = {} - # Time at previous scheduling step - self.prev_time = 0.0 - # Did we schedule a prompt at previous step? - self.prev_prompt = False - # Latency of the last prompt step - self.last_prompt_latency = 0.0 - # preemption mode, RECOMPUTE or SWAP - self.user_specified_preemption_mode = scheduler_config.preemption_mode - - # The following field is test-only. It is used to inject artificial - # preemption. - self.enable_artificial_preemption = ENABLE_ARTIFICIAL_PREEMPT - self.artificial_preempt_cnt = (ARTIFICIAL_PREEMPTION_MAX_CNT - if self.enable_artificial_preemption - else 0) - self.num_cumulative_preemption: int = 0 - - # Used to cache python objects - self._seq_group_metadata_cache: List[PyObjectCache] = [] - self._scheduler_running_outputs_cache: List[PyObjectCache] = [] - self._scheduled_seq_group_cache: List[PyObjectCache] = [] - - # For async output processing, we need to swap cache buffers between - # iterations. I.e. since the output processing is lagged one step, - # we cannot reuse the cached objects immediately when the schedule() - # is called again, but only when schedule() is called the second time. - self.output_proc_callback = output_proc_callback - self.use_async_output_proc = self.output_proc_callback is not None - self.num_cache_iters = 2 if self.use_async_output_proc else 1 - - self.cache_id = 0 - for i in range(self.num_cache_iters): - self._seq_group_metadata_cache.append( - PyObjectCache(seq_group_metadata_builder)) - self._scheduler_running_outputs_cache.append( - PyObjectCache(scheduler_running_outputs_builder)) - self._scheduled_seq_group_cache.append( - PyObjectCache(scheduled_seq_group_builder)) - - # For async postprocessor, the extra decode run cannot be done - # when the request reaches max_model_len. In this case, the request - # will be stopped during schedule() call and added to this stop list - # for processing and deallocation by the free_finished_seq_groups() - self._async_stopped: List[SequenceGroup] = [] - - @property - def next_cache_id(self): - return (self.cache_id + 1) % self.num_cache_iters - - @property - def lora_enabled(self) -> bool: - return bool(self.lora_config) - - @property - def num_decoding_tokens_per_seq(self) -> int: - """The number of new tokens.""" - return 1 - - def add_seq_group(self, seq_group: SequenceGroup) -> None: - # Add sequence groups to the waiting queue. - self.waiting.append(seq_group) - - def _add_seq_group_to_running(self, seq_group: SequenceGroup) -> None: - # Add sequence groups to the running queue. - # Only for testing purposes. - self.running.append(seq_group) - + # Time at previous scheduling step + self.prev_time = 0.0 + # Did we schedule a prompt at previous step? + self.prev_prompt = False + # Latency of the last prompt step + self.last_prompt_latency = 0.0 + # preemption mode, RECOMPUTE or SWAP + self.user_specified_preemption_mode = scheduler_config.preemption_mode + + # The following field is test-only. It is used to inject artificial + # preemption. + self.enable_artificial_preemption = ENABLE_ARTIFICIAL_PREEMPT + self.artificial_preempt_cnt = (ARTIFICIAL_PREEMPTION_MAX_CNT + if self.enable_artificial_preemption + else 0) + self.num_cumulative_preemption: int = 0 + + # Used to cache python objects + self._seq_group_metadata_cache: List[PyObjectCache] = [] + self._scheduler_running_outputs_cache: List[PyObjectCache] = [] + self._scheduled_seq_group_cache: List[PyObjectCache] = [] + + # For async output processing, we need to swap cache buffers between + # iterations. I.e. since the output processing is lagged one step, + # we cannot reuse the cached objects immediately when the schedule() + # is called again, but only when schedule() is called the second time. + self.output_proc_callback = output_proc_callback + self.use_async_output_proc = self.output_proc_callback is not None + self.num_cache_iters = 2 if self.use_async_output_proc else 1 + + self.cache_id = 0 + for i in range(self.num_cache_iters): + self._seq_group_metadata_cache.append( + PyObjectCache(seq_group_metadata_builder)) + self._scheduler_running_outputs_cache.append( + PyObjectCache(scheduler_running_outputs_builder)) + self._scheduled_seq_group_cache.append( + PyObjectCache(scheduled_seq_group_builder)) + + # For async postprocessor, the extra decode run cannot be done + # when the request reaches max_model_len. In this case, the request + # will be stopped during schedule() call and added to this stop list + # for processing and deallocation by the free_finished_seq_groups() + self._async_stopped: List[SequenceGroup] = [] + + @property + def next_cache_id(self): + return (self.cache_id + 1) % self.num_cache_iters + + @property + def lora_enabled(self) -> bool: + return bool(self.lora_config) + + @property + def num_decoding_tokens_per_seq(self) -> int: + """The number of new tokens.""" + return 1 + + def add_seq_group(self, seq_group: SequenceGroup) -> None: + # Add sequence groups to the waiting queue. + self.waiting.append(seq_group) + + def _add_seq_group_to_running(self, seq_group: SequenceGroup) -> None: + # Add sequence groups to the running queue. + # Only for testing purposes. + self.running.append(seq_group) + def _add_seq_group_to_swapped(self, seq_group: SequenceGroup) -> None: # Add sequence groups to the swapped queue. # Only for testing purposes. @@ -556,45 +556,45 @@ class Scheduler: return capped_chunk_size, capped_query_tokens def abort_seq_group(self, request_id: Union[str, Iterable[str]]) -> None: - """Aborts a sequence group with the given ID. - - Check if the sequence group with the given ID - is present in any of the state queue. - If present, remove the sequence group from the state queue. - Also, if any of the sequences in the sequence group is not finished, - free the sequence with status `FINISHED_ABORTED`. - Otherwise, do nothing. - - Args: - request_id: The ID(s) of the sequence group to abort. - """ - if isinstance(request_id, str): - request_id = (request_id, ) - request_ids = set(request_id) - for state_queue in [self.waiting, self.running, self.swapped]: - aborted_groups: List[SequenceGroup] = [] - for seq_group in state_queue: - if not request_ids: - # Using 'break' here may add two extra iterations, - # but is acceptable to reduce complexity. - break - if seq_group.request_id in request_ids: - # Appending aborted group into pending list. - aborted_groups.append(seq_group) - request_ids.remove(seq_group.request_id) - for aborted_group in aborted_groups: - # Remove the sequence group from the state queue. - state_queue.remove(aborted_group) - # Remove the aborted request from the Mamba cache. - self._finished_requests_ids.append(aborted_group.request_id) - for seq in aborted_group.get_seqs(): - if seq.is_finished(): - continue - seq.status = SequenceStatus.FINISHED_ABORTED - self.free_seq(seq) - - self._free_seq_group_cross_attn_blocks(aborted_group) - + """Aborts a sequence group with the given ID. + + Check if the sequence group with the given ID + is present in any of the state queue. + If present, remove the sequence group from the state queue. + Also, if any of the sequences in the sequence group is not finished, + free the sequence with status `FINISHED_ABORTED`. + Otherwise, do nothing. + + Args: + request_id: The ID(s) of the sequence group to abort. + """ + if isinstance(request_id, str): + request_id = (request_id, ) + request_ids = set(request_id) + for state_queue in [self.waiting, self.running, self.swapped]: + aborted_groups: List[SequenceGroup] = [] + for seq_group in state_queue: + if not request_ids: + # Using 'break' here may add two extra iterations, + # but is acceptable to reduce complexity. + break + if seq_group.request_id in request_ids: + # Appending aborted group into pending list. + aborted_groups.append(seq_group) + request_ids.remove(seq_group.request_id) + for aborted_group in aborted_groups: + # Remove the sequence group from the state queue. + state_queue.remove(aborted_group) + # Remove the aborted request from the Mamba cache. + self._finished_requests_ids.append(aborted_group.request_id) + for seq in aborted_group.get_seqs(): + if seq.is_finished(): + continue + seq.status = SequenceStatus.FINISHED_ABORTED + self.free_seq(seq) + + self._free_seq_group_cross_attn_blocks(aborted_group) + def _free_seq_group_cross_attn_blocks( self, seq_group: SequenceGroup, @@ -611,78 +611,78 @@ class Scheduler: self.block_manager, "release_request_cache_namespace", None) if release_namespace is not None: release_namespace(seq_group.request_id) - - def has_unfinished_seqs(self) -> bool: - return len(self.waiting) != 0 or len(self.running) != 0 or len( - self.swapped) != 0 - - def get_prefix_cache_hit_rate(self, device: Device) -> float: - return self.block_manager.get_prefix_cache_hit_rate(device) - - def get_num_unfinished_seq_groups(self) -> int: - return len(self.waiting) + len(self.running) + len(self.swapped) - - def get_and_reset_finished_requests_ids(self) -> List[str]: - """Flushes the list of request ids of previously finished seq_groups.""" - finished_requests_ids = self._finished_requests_ids - self._finished_requests_ids = list() - return finished_requests_ids - - def _schedule_running( - self, - budget: SchedulingBudget, - curr_loras: Optional[Set[int]], - enable_chunking: bool = False, - ) -> SchedulerRunningOutputs: - """Schedule sequence groups that are running. - - Running queue should include decode and chunked prefill requests. - - Args: - budget: The scheduling budget. The argument is in-place updated - when any decodes are preempted. - curr_loras: Currently batched lora request ids. The argument is - in-place updated when any decodes are preempted. - enable_chunking: If True, seq group can be chunked and only a - chunked number of tokens are scheduled if - `budget.num_batched_tokens` has not enough capacity to schedule - all tokens. - - Returns: - SchedulerRunningOutputs. - """ - ret: SchedulerRunningOutputs = \ - self._scheduler_running_outputs_cache[self.cache_id].get_object() - ret.blocks_to_swap_out.clear() - ret.blocks_to_copy.clear() - ret.decode_seq_groups.clear() - ret.prefill_seq_groups.clear() - ret.preempted.clear() - ret.swapped_out.clear() - - ret.num_lookahead_slots = self._get_num_lookahead_slots( - is_prefill=False, enable_chunking=enable_chunking) - - ret.decode_seq_groups_list.clear() - ret.prefill_seq_groups_list.clear() - - # Blocks that need to be swapped or copied before model execution. - blocks_to_swap_out: List[Tuple[int, int]] = ret.blocks_to_swap_out - blocks_to_copy: List[Tuple[int, int]] = ret.blocks_to_copy - - decode_seq_groups: List[ScheduledSequenceGroup] = ret.decode_seq_groups - prefill_seq_groups: List[ - ScheduledSequenceGroup] = ret.prefill_seq_groups - preempted: List[SequenceGroup] = ret.preempted - swapped_out: List[SequenceGroup] = ret.swapped_out - - running_queue = self.running - assert len(self._async_stopped) == 0 - while running_queue: - seq_group = running_queue[0] - num_running_tokens = self._get_num_new_tokens( - seq_group, SequenceStatus.RUNNING, enable_chunking, budget) - + + def has_unfinished_seqs(self) -> bool: + return len(self.waiting) != 0 or len(self.running) != 0 or len( + self.swapped) != 0 + + def get_prefix_cache_hit_rate(self, device: Device) -> float: + return self.block_manager.get_prefix_cache_hit_rate(device) + + def get_num_unfinished_seq_groups(self) -> int: + return len(self.waiting) + len(self.running) + len(self.swapped) + + def get_and_reset_finished_requests_ids(self) -> List[str]: + """Flushes the list of request ids of previously finished seq_groups.""" + finished_requests_ids = self._finished_requests_ids + self._finished_requests_ids = list() + return finished_requests_ids + + def _schedule_running( + self, + budget: SchedulingBudget, + curr_loras: Optional[Set[int]], + enable_chunking: bool = False, + ) -> SchedulerRunningOutputs: + """Schedule sequence groups that are running. + + Running queue should include decode and chunked prefill requests. + + Args: + budget: The scheduling budget. The argument is in-place updated + when any decodes are preempted. + curr_loras: Currently batched lora request ids. The argument is + in-place updated when any decodes are preempted. + enable_chunking: If True, seq group can be chunked and only a + chunked number of tokens are scheduled if + `budget.num_batched_tokens` has not enough capacity to schedule + all tokens. + + Returns: + SchedulerRunningOutputs. + """ + ret: SchedulerRunningOutputs = \ + self._scheduler_running_outputs_cache[self.cache_id].get_object() + ret.blocks_to_swap_out.clear() + ret.blocks_to_copy.clear() + ret.decode_seq_groups.clear() + ret.prefill_seq_groups.clear() + ret.preempted.clear() + ret.swapped_out.clear() + + ret.num_lookahead_slots = self._get_num_lookahead_slots( + is_prefill=False, enable_chunking=enable_chunking) + + ret.decode_seq_groups_list.clear() + ret.prefill_seq_groups_list.clear() + + # Blocks that need to be swapped or copied before model execution. + blocks_to_swap_out: List[Tuple[int, int]] = ret.blocks_to_swap_out + blocks_to_copy: List[Tuple[int, int]] = ret.blocks_to_copy + + decode_seq_groups: List[ScheduledSequenceGroup] = ret.decode_seq_groups + prefill_seq_groups: List[ + ScheduledSequenceGroup] = ret.prefill_seq_groups + preempted: List[SequenceGroup] = ret.preempted + swapped_out: List[SequenceGroup] = ret.swapped_out + + running_queue = self.running + assert len(self._async_stopped) == 0 + while running_queue: + seq_group = running_queue[0] + num_running_tokens = self._get_num_new_tokens( + seq_group, SequenceStatus.RUNNING, enable_chunking, budget) + if num_running_tokens == 0: # No budget => Stop break @@ -691,402 +691,402 @@ class Scheduler: seq_group, num_running_tokens, num_running_tokens) running_queue.popleft() - - # With async postprocessor, an extra decode run is done - # to process the final tokens. The check below avoids this extra - # decode run when the model max len is reached, in order to avoid - # a memory overflow. - if self.use_async_output_proc and seq_group.seqs[0].get_len( - ) > self.scheduler_config.max_model_len: - self._async_stopped.append(seq_group) - continue - - # NOTE(woosuk): Preemption happens only when there is no available - # slot to keep all the sequence groups in the RUNNING state. - while not self._can_append_slots(seq_group, enable_chunking): - budget.subtract_num_batched_tokens(seq_group.request_id, - num_running_tokens) - num_running_seqs = seq_group.get_max_num_running_seqs() - budget.subtract_num_seqs(seq_group.request_id, - num_running_seqs) - - if (curr_loras is not None and seq_group.lora_int_id > 0 - and seq_group.lora_int_id in curr_loras): - curr_loras.remove(seq_group.lora_int_id) - - # Determine victim sequence - cont_loop = True - if running_queue: - # Preempt the lowest-priority sequence group. - victim_seq_group = running_queue.pop() - else: - # No other sequence group can be preempted. - # Preempt the current sequence group. - # Note: This is also where we stop this loop - # (since there is nothing else to preempt) - victim_seq_group = seq_group - cont_loop = False - - # With async postprocessor, before preempting a sequence - # we need to ensure it has no pending async postprocessor - do_preempt = True - if self.use_async_output_proc: - assert self.output_proc_callback is not None - self.output_proc_callback( - request_id=victim_seq_group.request_id) - - # It may be that the async pending "victim_seq_group" - # becomes finished, in which case we simply free it. - if victim_seq_group.is_finished(): - self._free_finished_seq_group(victim_seq_group) - do_preempt = False - - # Do preemption - if do_preempt: - preempted_mode = self._preempt(victim_seq_group, - blocks_to_swap_out) - if preempted_mode == PreemptionMode.RECOMPUTE: - preempted.append(victim_seq_group) - else: - swapped_out.append(victim_seq_group) - - if not cont_loop: - break - else: - self._append_slots(seq_group, blocks_to_copy, enable_chunking) - is_prefill = seq_group.is_prefill() - - scheduled_seq_group: ScheduledSequenceGroup = \ - self._scheduled_seq_group_cache[self.cache_id].get_object() - scheduled_seq_group.seq_group = seq_group - if is_prefill: - scheduled_seq_group.token_chunk_size = num_running_tokens - prefill_seq_groups.append(scheduled_seq_group) - ret.prefill_seq_groups_list.append(seq_group) - else: - scheduled_seq_group.token_chunk_size = 1 - decode_seq_groups.append(scheduled_seq_group) - ret.decode_seq_groups_list.append(seq_group) - - budget.add_num_batched_tokens(seq_group.request_id, - num_running_tokens) - # OPTIMIZATION: Note that get_max_num_running_seqs is - # expensive. For the default scheduling chase where - # enable_chunking is False, num_seqs are updated before running - # this method, so we don't have to update it again here. - if enable_chunking: - num_running_seqs = seq_group.get_max_num_running_seqs() - budget.add_num_seqs(seq_group.request_id, num_running_seqs) - if curr_loras is not None and seq_group.lora_int_id > 0: - curr_loras.add(seq_group.lora_int_id) - - self._scheduler_running_outputs_cache[self.next_cache_id].reset() - self._scheduled_seq_group_cache[self.next_cache_id].reset() - - return ret - - def _schedule_swapped( - self, - budget: SchedulingBudget, - curr_loras: Optional[Set[int]], - enable_chunking: bool = False, - ) -> SchedulerSwappedInOutputs: - """Schedule sequence groups that are swapped out. - - It schedules swapped requests as long as it fits `budget` and - curr_loras <= max_lora from the scheduling config. The input arguments - `budget` and `curr_loras` are updated based on scheduled seq_groups. - - Args: - budget: The scheduling budget. The argument is in-place updated - when any requests are swapped in. - curr_loras: Currently batched lora request ids. The argument is - in-place updated when any requests are swapped in. - enable_chunking: If True, seq group can be chunked and only a - chunked number of tokens are scheduled if - `budget.num_batched_tokens` has not enough capacity to schedule - all tokens. - - Returns: - SchedulerSwappedInOutputs. - """ - # Blocks that need to be swapped or copied before model execution. - blocks_to_swap_in: List[Tuple[int, int]] = [] - blocks_to_copy: List[Tuple[int, int]] = [] - decode_seq_groups: List[ScheduledSequenceGroup] = [] - prefill_seq_groups: List[ScheduledSequenceGroup] = [] - infeasible_seq_groups: List[SequenceGroup] = [] - - swapped_queue = self.swapped - - leftover_swapped: Deque[SequenceGroup] = deque() - while swapped_queue: - seq_group = swapped_queue[0] - - # If the sequence group cannot be swapped in, stop. - is_prefill = seq_group.is_prefill() - alloc_status = self.block_manager.can_swap_in( - seq_group, - self._get_num_lookahead_slots(is_prefill, enable_chunking)) - if alloc_status == AllocStatus.LATER: - break - elif alloc_status == AllocStatus.NEVER: - logger.warning( - "Failing the request %s because there's not enough kv " - "cache blocks to run the entire sequence.", - seq_group.request_id) - for seq in seq_group.get_seqs(): - seq.status = SequenceStatus.FINISHED_IGNORED - infeasible_seq_groups.append(seq_group) - swapped_queue.popleft() - continue - - lora_int_id = 0 - if self.lora_enabled: - lora_int_id = seq_group.lora_int_id - assert curr_loras is not None - assert self.lora_config is not None - if (lora_int_id > 0 and (lora_int_id not in curr_loras) - and len(curr_loras) >= self.lora_config.max_loras): - # We don't have a space for another LoRA, so - # we ignore this request for now. - leftover_swapped.appendleft(seq_group) - swapped_queue.popleft() - continue - - # The total number of sequences in the RUNNING state should not - # exceed the maximum number of sequences. - num_new_seqs = seq_group.get_max_num_running_seqs() - num_new_tokens = self._get_num_new_tokens(seq_group, - SequenceStatus.SWAPPED, - enable_chunking, budget) - - if (num_new_tokens == 0 - or not budget.can_schedule(num_new_tokens=num_new_tokens, - num_new_seqs=num_new_seqs)): - break - - if lora_int_id > 0 and curr_loras is not None: - curr_loras.add(lora_int_id) - swapped_queue.popleft() - self._swap_in(seq_group, blocks_to_swap_in) - self._append_slots(seq_group, blocks_to_copy, enable_chunking) - is_prefill = seq_group.is_prefill() - if is_prefill: - prefill_seq_groups.append( - ScheduledSequenceGroup(seq_group, - token_chunk_size=num_new_tokens)) - else: - decode_seq_groups.append( - ScheduledSequenceGroup(seq_group, token_chunk_size=1)) - budget.add_num_batched_tokens(seq_group.request_id, num_new_tokens) - budget.add_num_seqs(seq_group.request_id, num_new_seqs) - - swapped_queue.extendleft(leftover_swapped) - - return SchedulerSwappedInOutputs( - decode_seq_groups=decode_seq_groups, - prefill_seq_groups=prefill_seq_groups, - blocks_to_swap_in=blocks_to_swap_in, - blocks_to_copy=blocks_to_copy, - num_lookahead_slots=self._get_num_lookahead_slots( - is_prefill=False, enable_chunking=enable_chunking), - infeasible_seq_groups=infeasible_seq_groups, - ) - - def _get_prompt_limit(self, seq_group: SequenceGroup) -> int: - if self.scheduler_config.chunked_prefill_enabled and \ - not self.scheduler_config.is_multi_step: - prompt_limit = self.scheduler_config.max_model_len - else: - prompt_limit = min(self.scheduler_config.max_model_len, - self.scheduler_config.max_num_batched_tokens) - - # Model is fine tuned with long context. Return the fine tuned max_len. - if (seq_group.lora_request - and seq_group.lora_request.long_lora_max_len): - assert prompt_limit <= seq_group.lora_request.long_lora_max_len - return seq_group.lora_request.long_lora_max_len - else: - return prompt_limit - - def _get_priority(self, - seq_group: SequenceGroup) -> Tuple[Optional[int], float]: - """ Get the priority of the sequence group. - Highest preference to user-defined priority, followed by arrival time. - Args: - seq_group: The sequence group input. - Returns: - The priority of the sequence group. - """ - return seq_group.priority, seq_group.arrival_time - - def _schedule_priority_preemption( - self, - budget: SchedulingBudget, - ) -> int: - """Sorts waiting and running queue. Also, force preempt requests - from the running queue if their priority is lower. - Priority-based preemption is used with the priority policy. - Args: - budget: The scheduling budget. The argument is in-place updated - when any requests are scheduled. - Returns: - A count of priority-based preemptions. - """ - - waiting_queue = self.waiting - - running_queue = deque(sorted(self.running, key=self._get_priority)) - - blocks_to_swap_out: List[Tuple[int, int]] = [] - force_preemption_count = 0 - - if waiting_queue: - seq_group = waiting_queue.popleft() - num_new_seqs = seq_group.get_max_num_running_seqs() - num_new_tokens = self._get_num_new_tokens(seq_group, - SequenceStatus.WAITING, - False, budget) - - #Only preempt if priority inversion exists - while running_queue and self._get_priority( - running_queue[-1]) > self._get_priority(seq_group): - #Only preempt if waiting sequence cannot be allocated - can_allocate = self.block_manager.can_allocate(seq_group) - if (num_new_tokens and can_allocate == AllocStatus.OK - and budget.can_schedule(num_new_tokens=num_new_tokens, - num_new_seqs=num_new_seqs)): - break - - #Adjust budget to remove the victim sequence group - vseq_group = running_queue.pop() - num_running_tokens = self._get_num_new_tokens( - vseq_group, SequenceStatus.RUNNING, False, budget) - budget.subtract_num_batched_tokens(vseq_group.request_id, - num_running_tokens) - num_running_seqs = vseq_group.get_max_num_running_seqs() - budget.subtract_num_seqs(vseq_group.request_id, - num_running_seqs) - - #Preempt out the victim sequence group - self._preempt(vseq_group, blocks_to_swap_out, - PreemptionMode.RECOMPUTE) - waiting_queue.appendleft(vseq_group) - force_preemption_count += 1 - #Put the sequence back into the waiting queue - waiting_queue.appendleft(seq_group) - - waiting_queue = deque(sorted(waiting_queue, key=self._get_priority)) - - self.waiting = waiting_queue - self.running = running_queue - return force_preemption_count - - def _schedule_prefills( - self, - budget: SchedulingBudget, - curr_loras: Optional[Set[int]], - enable_chunking: bool = False, - ) -> SchedulerPrefillOutputs: - """Schedule sequence groups that are in prefill stage. - - Note that the current scheduler treats PREEMPTED_FOR_RECOMPUTE - as a new prefill (that starts from beginning -> most recently generated - tokens). - - It schedules waiting requests as long as it fits `budget` and - curr_loras <= max_lora from the scheduling config. The input arguments - `budget` and `curr_loras` are updated based on scheduled seq_groups. - - Args: - budget: The scheduling budget. The argument is in-place updated - when any requests are scheduled. - curr_loras: Currently batched lora request ids. The argument is - in-place updated when any requests are scheduled. - enable_chunking: If True, seq group can be chunked and only a - chunked number of tokens are scheduled if - `budget.num_batched_tokens` has not enough capacity to schedule - all tokens. - - Returns: - SchedulerPrefillOutputs. - """ - ignored_seq_groups: List[SequenceGroup] = [] - seq_groups: List[ScheduledSequenceGroup] = [] - - waiting_queue = self.waiting - - leftover_waiting_sequences: Deque[SequenceGroup] = deque() - while self._passed_delay(time.time()) and waiting_queue: - seq_group = waiting_queue[0] - - waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING) - assert len(waiting_seqs) == 1, ( - "Waiting sequence group should have only one prompt " - "sequence.") - num_new_tokens = self._get_num_new_tokens(seq_group, - SequenceStatus.WAITING, - enable_chunking, budget) - if not enable_chunking: - num_prompt_tokens = waiting_seqs[0].get_len() - assert num_new_tokens == num_prompt_tokens - - prompt_limit = self._get_prompt_limit(seq_group) - if num_new_tokens > prompt_limit: - logger.warning( - "Input prompt (%d tokens) is too long" - " and exceeds limit of %d", num_new_tokens, prompt_limit) - for seq in waiting_seqs: - seq.status = SequenceStatus.FINISHED_IGNORED - ignored_seq_groups.append(seq_group) - waiting_queue.popleft() - continue - - num_lookahead_slots: int = 0 - if self.scheduler_config.is_multi_step and enable_chunking: - num_lookahead_slots = self._get_num_lookahead_slots( - True, enable_chunking) - - # If the sequence group cannot be allocated, stop. - can_allocate = self.block_manager.can_allocate( - seq_group, num_lookahead_slots=num_lookahead_slots) - if can_allocate == AllocStatus.LATER: - break - elif can_allocate == AllocStatus.NEVER: - logger.warning( - "Input prompt (%d tokens) + lookahead slots (%d) is " - "too long and exceeds the capacity of block_manager", - num_new_tokens, num_lookahead_slots) - for seq in waiting_seqs: - seq.status = SequenceStatus.FINISHED_IGNORED - ignored_seq_groups.append(seq_group) - waiting_queue.popleft() - continue - - lora_int_id = 0 - if self.lora_enabled: - lora_int_id = seq_group.lora_int_id - assert curr_loras is not None - assert self.lora_config is not None - if (self.lora_enabled and lora_int_id > 0 - and lora_int_id not in curr_loras - and len(curr_loras) >= self.lora_config.max_loras): - # We don't have a space for another LoRA, so - # we ignore this request for now. - leftover_waiting_sequences.appendleft(seq_group) - waiting_queue.popleft() - continue - - num_new_seqs = seq_group.get_max_num_running_seqs() - if (num_new_tokens == 0 - or not budget.can_schedule(num_new_tokens=num_new_tokens, - num_new_seqs=num_new_seqs)): - break - - # Can schedule this request. - if curr_loras is not None and lora_int_id > 0: - curr_loras.add(lora_int_id) + + # With async postprocessor, an extra decode run is done + # to process the final tokens. The check below avoids this extra + # decode run when the model max len is reached, in order to avoid + # a memory overflow. + if self.use_async_output_proc and seq_group.seqs[0].get_len( + ) > self.scheduler_config.max_model_len: + self._async_stopped.append(seq_group) + continue + + # NOTE(woosuk): Preemption happens only when there is no available + # slot to keep all the sequence groups in the RUNNING state. + while not self._can_append_slots(seq_group, enable_chunking): + budget.subtract_num_batched_tokens(seq_group.request_id, + num_running_tokens) + num_running_seqs = seq_group.get_max_num_running_seqs() + budget.subtract_num_seqs(seq_group.request_id, + num_running_seqs) + + if (curr_loras is not None and seq_group.lora_int_id > 0 + and seq_group.lora_int_id in curr_loras): + curr_loras.remove(seq_group.lora_int_id) + + # Determine victim sequence + cont_loop = True + if running_queue: + # Preempt the lowest-priority sequence group. + victim_seq_group = running_queue.pop() + else: + # No other sequence group can be preempted. + # Preempt the current sequence group. + # Note: This is also where we stop this loop + # (since there is nothing else to preempt) + victim_seq_group = seq_group + cont_loop = False + + # With async postprocessor, before preempting a sequence + # we need to ensure it has no pending async postprocessor + do_preempt = True + if self.use_async_output_proc: + assert self.output_proc_callback is not None + self.output_proc_callback( + request_id=victim_seq_group.request_id) + + # It may be that the async pending "victim_seq_group" + # becomes finished, in which case we simply free it. + if victim_seq_group.is_finished(): + self._free_finished_seq_group(victim_seq_group) + do_preempt = False + + # Do preemption + if do_preempt: + preempted_mode = self._preempt(victim_seq_group, + blocks_to_swap_out) + if preempted_mode == PreemptionMode.RECOMPUTE: + preempted.append(victim_seq_group) + else: + swapped_out.append(victim_seq_group) + + if not cont_loop: + break + else: + self._append_slots(seq_group, blocks_to_copy, enable_chunking) + is_prefill = seq_group.is_prefill() + + scheduled_seq_group: ScheduledSequenceGroup = \ + self._scheduled_seq_group_cache[self.cache_id].get_object() + scheduled_seq_group.seq_group = seq_group + if is_prefill: + scheduled_seq_group.token_chunk_size = num_running_tokens + prefill_seq_groups.append(scheduled_seq_group) + ret.prefill_seq_groups_list.append(seq_group) + else: + scheduled_seq_group.token_chunk_size = 1 + decode_seq_groups.append(scheduled_seq_group) + ret.decode_seq_groups_list.append(seq_group) + + budget.add_num_batched_tokens(seq_group.request_id, + num_running_tokens) + # OPTIMIZATION: Note that get_max_num_running_seqs is + # expensive. For the default scheduling chase where + # enable_chunking is False, num_seqs are updated before running + # this method, so we don't have to update it again here. + if enable_chunking: + num_running_seqs = seq_group.get_max_num_running_seqs() + budget.add_num_seqs(seq_group.request_id, num_running_seqs) + if curr_loras is not None and seq_group.lora_int_id > 0: + curr_loras.add(seq_group.lora_int_id) + + self._scheduler_running_outputs_cache[self.next_cache_id].reset() + self._scheduled_seq_group_cache[self.next_cache_id].reset() + + return ret + + def _schedule_swapped( + self, + budget: SchedulingBudget, + curr_loras: Optional[Set[int]], + enable_chunking: bool = False, + ) -> SchedulerSwappedInOutputs: + """Schedule sequence groups that are swapped out. + + It schedules swapped requests as long as it fits `budget` and + curr_loras <= max_lora from the scheduling config. The input arguments + `budget` and `curr_loras` are updated based on scheduled seq_groups. + + Args: + budget: The scheduling budget. The argument is in-place updated + when any requests are swapped in. + curr_loras: Currently batched lora request ids. The argument is + in-place updated when any requests are swapped in. + enable_chunking: If True, seq group can be chunked and only a + chunked number of tokens are scheduled if + `budget.num_batched_tokens` has not enough capacity to schedule + all tokens. + + Returns: + SchedulerSwappedInOutputs. + """ + # Blocks that need to be swapped or copied before model execution. + blocks_to_swap_in: List[Tuple[int, int]] = [] + blocks_to_copy: List[Tuple[int, int]] = [] + decode_seq_groups: List[ScheduledSequenceGroup] = [] + prefill_seq_groups: List[ScheduledSequenceGroup] = [] + infeasible_seq_groups: List[SequenceGroup] = [] + + swapped_queue = self.swapped + + leftover_swapped: Deque[SequenceGroup] = deque() + while swapped_queue: + seq_group = swapped_queue[0] + + # If the sequence group cannot be swapped in, stop. + is_prefill = seq_group.is_prefill() + alloc_status = self.block_manager.can_swap_in( + seq_group, + self._get_num_lookahead_slots(is_prefill, enable_chunking)) + if alloc_status == AllocStatus.LATER: + break + elif alloc_status == AllocStatus.NEVER: + logger.warning( + "Failing the request %s because there's not enough kv " + "cache blocks to run the entire sequence.", + seq_group.request_id) + for seq in seq_group.get_seqs(): + seq.status = SequenceStatus.FINISHED_IGNORED + infeasible_seq_groups.append(seq_group) + swapped_queue.popleft() + continue + + lora_int_id = 0 + if self.lora_enabled: + lora_int_id = seq_group.lora_int_id + assert curr_loras is not None + assert self.lora_config is not None + if (lora_int_id > 0 and (lora_int_id not in curr_loras) + and len(curr_loras) >= self.lora_config.max_loras): + # We don't have a space for another LoRA, so + # we ignore this request for now. + leftover_swapped.appendleft(seq_group) + swapped_queue.popleft() + continue + + # The total number of sequences in the RUNNING state should not + # exceed the maximum number of sequences. + num_new_seqs = seq_group.get_max_num_running_seqs() + num_new_tokens = self._get_num_new_tokens(seq_group, + SequenceStatus.SWAPPED, + enable_chunking, budget) + + if (num_new_tokens == 0 + or not budget.can_schedule(num_new_tokens=num_new_tokens, + num_new_seqs=num_new_seqs)): + break + + if lora_int_id > 0 and curr_loras is not None: + curr_loras.add(lora_int_id) + swapped_queue.popleft() + self._swap_in(seq_group, blocks_to_swap_in) + self._append_slots(seq_group, blocks_to_copy, enable_chunking) + is_prefill = seq_group.is_prefill() + if is_prefill: + prefill_seq_groups.append( + ScheduledSequenceGroup(seq_group, + token_chunk_size=num_new_tokens)) + else: + decode_seq_groups.append( + ScheduledSequenceGroup(seq_group, token_chunk_size=1)) + budget.add_num_batched_tokens(seq_group.request_id, num_new_tokens) + budget.add_num_seqs(seq_group.request_id, num_new_seqs) + + swapped_queue.extendleft(leftover_swapped) + + return SchedulerSwappedInOutputs( + decode_seq_groups=decode_seq_groups, + prefill_seq_groups=prefill_seq_groups, + blocks_to_swap_in=blocks_to_swap_in, + blocks_to_copy=blocks_to_copy, + num_lookahead_slots=self._get_num_lookahead_slots( + is_prefill=False, enable_chunking=enable_chunking), + infeasible_seq_groups=infeasible_seq_groups, + ) + + def _get_prompt_limit(self, seq_group: SequenceGroup) -> int: + if self.scheduler_config.chunked_prefill_enabled and \ + not self.scheduler_config.is_multi_step: + prompt_limit = self.scheduler_config.max_model_len + else: + prompt_limit = min(self.scheduler_config.max_model_len, + self.scheduler_config.max_num_batched_tokens) + + # Model is fine tuned with long context. Return the fine tuned max_len. + if (seq_group.lora_request + and seq_group.lora_request.long_lora_max_len): + assert prompt_limit <= seq_group.lora_request.long_lora_max_len + return seq_group.lora_request.long_lora_max_len + else: + return prompt_limit + + def _get_priority(self, + seq_group: SequenceGroup) -> Tuple[Optional[int], float]: + """ Get the priority of the sequence group. + Highest preference to user-defined priority, followed by arrival time. + Args: + seq_group: The sequence group input. + Returns: + The priority of the sequence group. + """ + return seq_group.priority, seq_group.arrival_time + + def _schedule_priority_preemption( + self, + budget: SchedulingBudget, + ) -> int: + """Sorts waiting and running queue. Also, force preempt requests + from the running queue if their priority is lower. + Priority-based preemption is used with the priority policy. + Args: + budget: The scheduling budget. The argument is in-place updated + when any requests are scheduled. + Returns: + A count of priority-based preemptions. + """ + + waiting_queue = self.waiting + + running_queue = deque(sorted(self.running, key=self._get_priority)) + + blocks_to_swap_out: List[Tuple[int, int]] = [] + force_preemption_count = 0 + + if waiting_queue: + seq_group = waiting_queue.popleft() + num_new_seqs = seq_group.get_max_num_running_seqs() + num_new_tokens = self._get_num_new_tokens(seq_group, + SequenceStatus.WAITING, + False, budget) + + #Only preempt if priority inversion exists + while running_queue and self._get_priority( + running_queue[-1]) > self._get_priority(seq_group): + #Only preempt if waiting sequence cannot be allocated + can_allocate = self.block_manager.can_allocate(seq_group) + if (num_new_tokens and can_allocate == AllocStatus.OK + and budget.can_schedule(num_new_tokens=num_new_tokens, + num_new_seqs=num_new_seqs)): + break + + #Adjust budget to remove the victim sequence group + vseq_group = running_queue.pop() + num_running_tokens = self._get_num_new_tokens( + vseq_group, SequenceStatus.RUNNING, False, budget) + budget.subtract_num_batched_tokens(vseq_group.request_id, + num_running_tokens) + num_running_seqs = vseq_group.get_max_num_running_seqs() + budget.subtract_num_seqs(vseq_group.request_id, + num_running_seqs) + + #Preempt out the victim sequence group + self._preempt(vseq_group, blocks_to_swap_out, + PreemptionMode.RECOMPUTE) + waiting_queue.appendleft(vseq_group) + force_preemption_count += 1 + #Put the sequence back into the waiting queue + waiting_queue.appendleft(seq_group) + + waiting_queue = deque(sorted(waiting_queue, key=self._get_priority)) + + self.waiting = waiting_queue + self.running = running_queue + return force_preemption_count + + def _schedule_prefills( + self, + budget: SchedulingBudget, + curr_loras: Optional[Set[int]], + enable_chunking: bool = False, + ) -> SchedulerPrefillOutputs: + """Schedule sequence groups that are in prefill stage. + + Note that the current scheduler treats PREEMPTED_FOR_RECOMPUTE + as a new prefill (that starts from beginning -> most recently generated + tokens). + + It schedules waiting requests as long as it fits `budget` and + curr_loras <= max_lora from the scheduling config. The input arguments + `budget` and `curr_loras` are updated based on scheduled seq_groups. + + Args: + budget: The scheduling budget. The argument is in-place updated + when any requests are scheduled. + curr_loras: Currently batched lora request ids. The argument is + in-place updated when any requests are scheduled. + enable_chunking: If True, seq group can be chunked and only a + chunked number of tokens are scheduled if + `budget.num_batched_tokens` has not enough capacity to schedule + all tokens. + + Returns: + SchedulerPrefillOutputs. + """ + ignored_seq_groups: List[SequenceGroup] = [] + seq_groups: List[ScheduledSequenceGroup] = [] + + waiting_queue = self.waiting + + leftover_waiting_sequences: Deque[SequenceGroup] = deque() + while self._passed_delay(time.time()) and waiting_queue: + seq_group = waiting_queue[0] + + waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING) + assert len(waiting_seqs) == 1, ( + "Waiting sequence group should have only one prompt " + "sequence.") + num_new_tokens = self._get_num_new_tokens(seq_group, + SequenceStatus.WAITING, + enable_chunking, budget) + if not enable_chunking: + num_prompt_tokens = waiting_seqs[0].get_len() + assert num_new_tokens == num_prompt_tokens + + prompt_limit = self._get_prompt_limit(seq_group) + if num_new_tokens > prompt_limit: + logger.warning( + "Input prompt (%d tokens) is too long" + " and exceeds limit of %d", num_new_tokens, prompt_limit) + for seq in waiting_seqs: + seq.status = SequenceStatus.FINISHED_IGNORED + ignored_seq_groups.append(seq_group) + waiting_queue.popleft() + continue + + num_lookahead_slots: int = 0 + if self.scheduler_config.is_multi_step and enable_chunking: + num_lookahead_slots = self._get_num_lookahead_slots( + True, enable_chunking) + + # If the sequence group cannot be allocated, stop. + can_allocate = self.block_manager.can_allocate( + seq_group, num_lookahead_slots=num_lookahead_slots) + if can_allocate == AllocStatus.LATER: + break + elif can_allocate == AllocStatus.NEVER: + logger.warning( + "Input prompt (%d tokens) + lookahead slots (%d) is " + "too long and exceeds the capacity of block_manager", + num_new_tokens, num_lookahead_slots) + for seq in waiting_seqs: + seq.status = SequenceStatus.FINISHED_IGNORED + ignored_seq_groups.append(seq_group) + waiting_queue.popleft() + continue + + lora_int_id = 0 + if self.lora_enabled: + lora_int_id = seq_group.lora_int_id + assert curr_loras is not None + assert self.lora_config is not None + if (self.lora_enabled and lora_int_id > 0 + and lora_int_id not in curr_loras + and len(curr_loras) >= self.lora_config.max_loras): + # We don't have a space for another LoRA, so + # we ignore this request for now. + leftover_waiting_sequences.appendleft(seq_group) + waiting_queue.popleft() + continue + + num_new_seqs = seq_group.get_max_num_running_seqs() + if (num_new_tokens == 0 + or not budget.can_schedule(num_new_tokens=num_new_tokens, + num_new_seqs=num_new_seqs)): + break + + # Can schedule this request. + if curr_loras is not None and lora_int_id > 0: + curr_loras.add(lora_int_id) waiting_queue.popleft() self._allocate_and_set_running(seq_group) @@ -1168,22 +1168,22 @@ class Scheduler: seq_group, num_new_tokens, budget_token_count)) if enable_chunking and self.scheduler_config.is_multi_step: - blocks_to_copy: List[Tuple[int, int]] = [] - # init_multi_step_from_lookahead_slots happens in append_slots - self._append_slots(seq_group, blocks_to_copy, enable_chunking) - # This assert will trip when a copy-on-write happens. This is - # not a concern as the very first sequence-group block - # allocation happens above. Still, we have the assert to - # catch any edge-cases. - assert not blocks_to_copy - else: - seq_group.init_multi_step_from_lookahead_slots( - num_lookahead_slots, - num_scheduler_steps=self.scheduler_config. - num_scheduler_steps, - is_multi_step=self.scheduler_config.is_multi_step, - enable_chunking=enable_chunking) - + blocks_to_copy: List[Tuple[int, int]] = [] + # init_multi_step_from_lookahead_slots happens in append_slots + self._append_slots(seq_group, blocks_to_copy, enable_chunking) + # This assert will trip when a copy-on-write happens. This is + # not a concern as the very first sequence-group block + # allocation happens above. Still, we have the assert to + # catch any edge-cases. + assert not blocks_to_copy + else: + seq_group.init_multi_step_from_lookahead_slots( + num_lookahead_slots, + num_scheduler_steps=self.scheduler_config. + num_scheduler_steps, + is_multi_step=self.scheduler_config.is_multi_step, + enable_chunking=enable_chunking) + seq_groups.append( ScheduledSequenceGroup(seq_group=seq_group, token_chunk_size=num_new_tokens)) @@ -1191,250 +1191,250 @@ class Scheduler: seq_group.request_id, budget_token_count, num_scheduled_tokens=num_new_tokens) - budget.add_num_seqs(seq_group.request_id, num_new_seqs) - - # Queue requests that couldn't be scheduled. - waiting_queue.extendleft(leftover_waiting_sequences) - if len(seq_groups) > 0: - self.prev_prompt = True - - return SchedulerPrefillOutputs( - seq_groups=seq_groups, - ignored_seq_groups=ignored_seq_groups, - num_lookahead_slots=self._get_num_lookahead_slots( - is_prefill=True, enable_chunking=enable_chunking)) - - def _schedule_default(self) -> SchedulerOutputs: - """Schedule queued requests. - - The current policy is designed to optimize the throughput. First, - it batches as many prefill requests as possible. And it schedules - decodes. If there's a pressure on GPU memory, decode requests can - be swapped or preempted. - """ - # Include running requests to the budget. - budget = SchedulingBudget( - token_budget=self.scheduler_config.max_num_batched_tokens, - max_num_seqs=self.scheduler_config.max_num_seqs, - ) - # Make sure we include num running seqs before scheduling prefill, - # so that we don't schedule beyond max_num_seqs for prefill. - for seq_group in self.running: - budget.add_num_seqs(seq_group.request_id, - seq_group.get_max_num_running_seqs()) - curr_loras = set( - seq_group.lora_int_id for seq_group in self.running - if seq_group.lora_int_id > 0) if self.lora_enabled else None - - prefills = SchedulerPrefillOutputs.create_empty() - running_scheduled = SchedulerRunningOutputs.create_empty() - swapped_in = SchedulerSwappedInOutputs.create_empty() - - # If any requests are swapped, prioritized swapped requests. - if not self.swapped: - prefills = self._schedule_prefills(budget, - curr_loras, - enable_chunking=False) - - if len(prefills.seq_groups - ) == 0 and self.scheduler_config.policy == "priority": - self._schedule_priority_preemption(budget) - - # Don't schedule decodes if prefills are scheduled. - # NOTE: If `_schedule_prefills` doesn't enable chunking, self.running - # only contains decode requests, not chunked prefills. - if len(prefills.seq_groups) == 0: - running_scheduled = self._schedule_running(budget, - curr_loras, - enable_chunking=False) - - # If any sequence group is preempted, do not swap in any sequence - # group. because it means there's no slot for new running requests. - if len(running_scheduled.preempted) + len( - running_scheduled.swapped_out) == 0: - swapped_in = self._schedule_swapped(budget, curr_loras) - - assert (budget.num_batched_tokens <= - self.scheduler_config.max_num_batched_tokens) - assert budget.num_curr_seqs <= self.scheduler_config.max_num_seqs - - # Update waiting requests. - self.waiting.extendleft(running_scheduled.preempted) - # Update new running requests. - if len(prefills.seq_groups) > 0: - self.running.extend([s.seq_group for s in prefills.seq_groups]) - - self.running.extend(running_scheduled.decode_seq_groups_list) - - if len(swapped_in.decode_seq_groups) > 0: - self.running.extend( - [s.seq_group for s in swapped_in.decode_seq_groups]) - - # Update swapped requests. - self.swapped.extend(running_scheduled.swapped_out) - preempted = (len(running_scheduled.preempted) + - len(running_scheduled.swapped_out)) - - # There should be no prefill from running queue because this policy - # doesn't allow chunked prefills. - assert len(running_scheduled.prefill_seq_groups) == 0 - assert len(swapped_in.prefill_seq_groups) == 0 - - # Merge lists - num_prefill_groups = len(prefills.seq_groups) - if num_prefill_groups > 0: - scheduled_seq_groups = prefills.seq_groups - scheduled_seq_groups.extend(running_scheduled.decode_seq_groups) - else: - scheduled_seq_groups = running_scheduled.decode_seq_groups - scheduled_seq_groups.extend(swapped_in.decode_seq_groups) - - blocks_to_copy = running_scheduled.blocks_to_copy - blocks_to_copy.extend(swapped_in.blocks_to_copy) - - ignored_seq_groups = prefills.ignored_seq_groups - ignored_seq_groups.extend(swapped_in.infeasible_seq_groups) - - return SchedulerOutputs( - scheduled_seq_groups=scheduled_seq_groups, - num_prefill_groups=num_prefill_groups, + budget.add_num_seqs(seq_group.request_id, num_new_seqs) + + # Queue requests that couldn't be scheduled. + waiting_queue.extendleft(leftover_waiting_sequences) + if len(seq_groups) > 0: + self.prev_prompt = True + + return SchedulerPrefillOutputs( + seq_groups=seq_groups, + ignored_seq_groups=ignored_seq_groups, + num_lookahead_slots=self._get_num_lookahead_slots( + is_prefill=True, enable_chunking=enable_chunking)) + + def _schedule_default(self) -> SchedulerOutputs: + """Schedule queued requests. + + The current policy is designed to optimize the throughput. First, + it batches as many prefill requests as possible. And it schedules + decodes. If there's a pressure on GPU memory, decode requests can + be swapped or preempted. + """ + # Include running requests to the budget. + budget = SchedulingBudget( + token_budget=self.scheduler_config.max_num_batched_tokens, + max_num_seqs=self.scheduler_config.max_num_seqs, + ) + # Make sure we include num running seqs before scheduling prefill, + # so that we don't schedule beyond max_num_seqs for prefill. + for seq_group in self.running: + budget.add_num_seqs(seq_group.request_id, + seq_group.get_max_num_running_seqs()) + curr_loras = set( + seq_group.lora_int_id for seq_group in self.running + if seq_group.lora_int_id > 0) if self.lora_enabled else None + + prefills = SchedulerPrefillOutputs.create_empty() + running_scheduled = SchedulerRunningOutputs.create_empty() + swapped_in = SchedulerSwappedInOutputs.create_empty() + + # If any requests are swapped, prioritized swapped requests. + if not self.swapped: + prefills = self._schedule_prefills(budget, + curr_loras, + enable_chunking=False) + + if len(prefills.seq_groups + ) == 0 and self.scheduler_config.policy == "priority": + self._schedule_priority_preemption(budget) + + # Don't schedule decodes if prefills are scheduled. + # NOTE: If `_schedule_prefills` doesn't enable chunking, self.running + # only contains decode requests, not chunked prefills. + if len(prefills.seq_groups) == 0: + running_scheduled = self._schedule_running(budget, + curr_loras, + enable_chunking=False) + + # If any sequence group is preempted, do not swap in any sequence + # group. because it means there's no slot for new running requests. + if len(running_scheduled.preempted) + len( + running_scheduled.swapped_out) == 0: + swapped_in = self._schedule_swapped(budget, curr_loras) + + assert (budget.num_batched_tokens <= + self.scheduler_config.max_num_batched_tokens) + assert budget.num_curr_seqs <= self.scheduler_config.max_num_seqs + + # Update waiting requests. + self.waiting.extendleft(running_scheduled.preempted) + # Update new running requests. + if len(prefills.seq_groups) > 0: + self.running.extend([s.seq_group for s in prefills.seq_groups]) + + self.running.extend(running_scheduled.decode_seq_groups_list) + + if len(swapped_in.decode_seq_groups) > 0: + self.running.extend( + [s.seq_group for s in swapped_in.decode_seq_groups]) + + # Update swapped requests. + self.swapped.extend(running_scheduled.swapped_out) + preempted = (len(running_scheduled.preempted) + + len(running_scheduled.swapped_out)) + + # There should be no prefill from running queue because this policy + # doesn't allow chunked prefills. + assert len(running_scheduled.prefill_seq_groups) == 0 + assert len(swapped_in.prefill_seq_groups) == 0 + + # Merge lists + num_prefill_groups = len(prefills.seq_groups) + if num_prefill_groups > 0: + scheduled_seq_groups = prefills.seq_groups + scheduled_seq_groups.extend(running_scheduled.decode_seq_groups) + else: + scheduled_seq_groups = running_scheduled.decode_seq_groups + scheduled_seq_groups.extend(swapped_in.decode_seq_groups) + + blocks_to_copy = running_scheduled.blocks_to_copy + blocks_to_copy.extend(swapped_in.blocks_to_copy) + + ignored_seq_groups = prefills.ignored_seq_groups + ignored_seq_groups.extend(swapped_in.infeasible_seq_groups) + + return SchedulerOutputs( + scheduled_seq_groups=scheduled_seq_groups, + num_prefill_groups=num_prefill_groups, num_batched_tokens=budget.num_scheduled_tokens, - blocks_to_swap_in=swapped_in.blocks_to_swap_in, - blocks_to_swap_out=running_scheduled.blocks_to_swap_out, - blocks_to_copy=blocks_to_copy, - ignored_seq_groups=ignored_seq_groups, - num_lookahead_slots=running_scheduled.num_lookahead_slots, - running_queue_size=len(self.running), - preempted=preempted, - ) - - def _schedule_chunked_prefill(self) -> SchedulerOutputs: - """Schedule queued requests. - - Chunked prefill allows to chunk prefill requests, batch them together - with decode requests. This policy 1. schedule as many decoding requests - as possible. 2. schedule chunked prefill requests that are not - finished. 3. schedule swapped request. 4. schedule new prefill - requests. - - The policy can sustain the high GPU utilization because it can put - prefill and decodes requests to the same batch, while it improves - inter token latency because decodes requests don't need to be blocked - by prefill requests. - """ - budget = SchedulingBudget( - token_budget=self.scheduler_config.max_num_batched_tokens, - max_num_seqs=self.scheduler_config.max_num_seqs, - ) - curr_loras: Set[int] = set() - - prefills = SchedulerPrefillOutputs.create_empty() - swapped_in = SchedulerSwappedInOutputs.create_empty() - - # Decoding should be always scheduled first by fcfs. - running_scheduled = self._schedule_running(budget, - curr_loras, - enable_chunking=True) - - # Schedule swapped out requests. - # If preemption happens, it means we don't have space for swap-in. - if len(running_scheduled.preempted) + len( - running_scheduled.swapped_out) == 0: - swapped_in = self._schedule_swapped(budget, curr_loras) - - # Schedule new prefills. - prefills = self._schedule_prefills(budget, - curr_loras, - enable_chunking=True) - - assert (budget.num_batched_tokens <= - self.scheduler_config.max_num_batched_tokens) - assert budget.num_curr_seqs <= self.scheduler_config.max_num_seqs - - # Update waiting requests. - self.waiting.extendleft(running_scheduled.preempted) - - # Update new running requests. - # By default, vLLM scheduler prioritizes prefills. - # Once chunked prefill is enabled, - # the policy is changed to prioritize decode requests. - self.running.extend( - [s.seq_group for s in swapped_in.decode_seq_groups]) - self.running.extend( - [s.seq_group for s in swapped_in.prefill_seq_groups]) - self.running.extend( - [s.seq_group for s in running_scheduled.decode_seq_groups]) - self.running.extend( - [s.seq_group for s in running_scheduled.prefill_seq_groups]) - self.running.extend([s.seq_group for s in prefills.seq_groups]) - - # Update swapped requests. - self.swapped.extend(running_scheduled.swapped_out) - return SchedulerOutputs( - scheduled_seq_groups=(prefills.seq_groups + - running_scheduled.prefill_seq_groups + - swapped_in.prefill_seq_groups + - running_scheduled.decode_seq_groups + - swapped_in.decode_seq_groups), - num_prefill_groups=(len(prefills.seq_groups) + - len(swapped_in.prefill_seq_groups) + - len(running_scheduled.prefill_seq_groups)), + blocks_to_swap_in=swapped_in.blocks_to_swap_in, + blocks_to_swap_out=running_scheduled.blocks_to_swap_out, + blocks_to_copy=blocks_to_copy, + ignored_seq_groups=ignored_seq_groups, + num_lookahead_slots=running_scheduled.num_lookahead_slots, + running_queue_size=len(self.running), + preempted=preempted, + ) + + def _schedule_chunked_prefill(self) -> SchedulerOutputs: + """Schedule queued requests. + + Chunked prefill allows to chunk prefill requests, batch them together + with decode requests. This policy 1. schedule as many decoding requests + as possible. 2. schedule chunked prefill requests that are not + finished. 3. schedule swapped request. 4. schedule new prefill + requests. + + The policy can sustain the high GPU utilization because it can put + prefill and decodes requests to the same batch, while it improves + inter token latency because decodes requests don't need to be blocked + by prefill requests. + """ + budget = SchedulingBudget( + token_budget=self.scheduler_config.max_num_batched_tokens, + max_num_seqs=self.scheduler_config.max_num_seqs, + ) + curr_loras: Set[int] = set() + + prefills = SchedulerPrefillOutputs.create_empty() + swapped_in = SchedulerSwappedInOutputs.create_empty() + + # Decoding should be always scheduled first by fcfs. + running_scheduled = self._schedule_running(budget, + curr_loras, + enable_chunking=True) + + # Schedule swapped out requests. + # If preemption happens, it means we don't have space for swap-in. + if len(running_scheduled.preempted) + len( + running_scheduled.swapped_out) == 0: + swapped_in = self._schedule_swapped(budget, curr_loras) + + # Schedule new prefills. + prefills = self._schedule_prefills(budget, + curr_loras, + enable_chunking=True) + + assert (budget.num_batched_tokens <= + self.scheduler_config.max_num_batched_tokens) + assert budget.num_curr_seqs <= self.scheduler_config.max_num_seqs + + # Update waiting requests. + self.waiting.extendleft(running_scheduled.preempted) + + # Update new running requests. + # By default, vLLM scheduler prioritizes prefills. + # Once chunked prefill is enabled, + # the policy is changed to prioritize decode requests. + self.running.extend( + [s.seq_group for s in swapped_in.decode_seq_groups]) + self.running.extend( + [s.seq_group for s in swapped_in.prefill_seq_groups]) + self.running.extend( + [s.seq_group for s in running_scheduled.decode_seq_groups]) + self.running.extend( + [s.seq_group for s in running_scheduled.prefill_seq_groups]) + self.running.extend([s.seq_group for s in prefills.seq_groups]) + + # Update swapped requests. + self.swapped.extend(running_scheduled.swapped_out) + return SchedulerOutputs( + scheduled_seq_groups=(prefills.seq_groups + + running_scheduled.prefill_seq_groups + + swapped_in.prefill_seq_groups + + running_scheduled.decode_seq_groups + + swapped_in.decode_seq_groups), + num_prefill_groups=(len(prefills.seq_groups) + + len(swapped_in.prefill_seq_groups) + + len(running_scheduled.prefill_seq_groups)), num_batched_tokens=budget.num_scheduled_tokens, - blocks_to_swap_in=swapped_in.blocks_to_swap_in, - blocks_to_swap_out=running_scheduled.blocks_to_swap_out, - blocks_to_copy=running_scheduled.blocks_to_copy + - swapped_in.blocks_to_copy, - ignored_seq_groups=prefills.ignored_seq_groups + - swapped_in.infeasible_seq_groups, - num_lookahead_slots=running_scheduled.num_lookahead_slots, - running_queue_size=len(self.running), - preempted=(len(running_scheduled.preempted) + - len(running_scheduled.swapped_out)), - ) - - def _schedule(self) -> SchedulerOutputs: - """Schedule queued requests.""" - if self.scheduler_config.chunked_prefill_enabled: - return self._schedule_chunked_prefill() - else: - return self._schedule_default() - - def _can_append_slots(self, seq_group: SequenceGroup, - enable_chunking: bool) -> bool: - """Determine whether or not we have enough space in the KV cache to - continue generation of the sequence group. - """ - # It is True only for testing case to trigger artificial preemption. - if (self.enable_artificial_preemption - and random.uniform(0, 1) < ARTIFICIAL_PREEMPTION_PROB - and self.artificial_preempt_cnt > 0): - self.artificial_preempt_cnt -= 1 - return False - - is_prefill = seq_group.is_prefill() - num_lookahead_slots = self._get_num_lookahead_slots( - is_prefill, enable_chunking) - - if is_prefill and num_lookahead_slots > 0: - # Appending prefill slots only happens multi-step and - # chunked-prefill are enabled together. - assert self.scheduler_config.is_multi_step and enable_chunking - - return self.block_manager.can_append_slots( - seq_group=seq_group, num_lookahead_slots=num_lookahead_slots) - - def _allow_async_output_proc(self, seq_group: SequenceGroup) -> bool: - # async_output_proc is allowed only when we have a single sequence - # in the sequence group - no_single_seq = seq_group.sampling_params is None or ( - seq_group.sampling_params.n == 1) - return no_single_seq - - def schedule( - self - ) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs, bool]: - # Schedule sequence groups. - # This function call changes the internal states of the scheduler + blocks_to_swap_in=swapped_in.blocks_to_swap_in, + blocks_to_swap_out=running_scheduled.blocks_to_swap_out, + blocks_to_copy=running_scheduled.blocks_to_copy + + swapped_in.blocks_to_copy, + ignored_seq_groups=prefills.ignored_seq_groups + + swapped_in.infeasible_seq_groups, + num_lookahead_slots=running_scheduled.num_lookahead_slots, + running_queue_size=len(self.running), + preempted=(len(running_scheduled.preempted) + + len(running_scheduled.swapped_out)), + ) + + def _schedule(self) -> SchedulerOutputs: + """Schedule queued requests.""" + if self.scheduler_config.chunked_prefill_enabled: + return self._schedule_chunked_prefill() + else: + return self._schedule_default() + + def _can_append_slots(self, seq_group: SequenceGroup, + enable_chunking: bool) -> bool: + """Determine whether or not we have enough space in the KV cache to + continue generation of the sequence group. + """ + # It is True only for testing case to trigger artificial preemption. + if (self.enable_artificial_preemption + and random.uniform(0, 1) < ARTIFICIAL_PREEMPTION_PROB + and self.artificial_preempt_cnt > 0): + self.artificial_preempt_cnt -= 1 + return False + + is_prefill = seq_group.is_prefill() + num_lookahead_slots = self._get_num_lookahead_slots( + is_prefill, enable_chunking) + + if is_prefill and num_lookahead_slots > 0: + # Appending prefill slots only happens multi-step and + # chunked-prefill are enabled together. + assert self.scheduler_config.is_multi_step and enable_chunking + + return self.block_manager.can_append_slots( + seq_group=seq_group, num_lookahead_slots=num_lookahead_slots) + + def _allow_async_output_proc(self, seq_group: SequenceGroup) -> bool: + # async_output_proc is allowed only when we have a single sequence + # in the sequence group + no_single_seq = seq_group.sampling_params is None or ( + seq_group.sampling_params.n == 1) + return no_single_seq + + def schedule( + self + ) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs, bool]: + # Schedule sequence groups. + # This function call changes the internal states of the scheduler # such as self.running, self.swapped, and self.waiting. scheduler_start_time = time.perf_counter() @@ -1458,49 +1458,49 @@ class Scheduler: scheduler_outputs.blocks_to_swap_in.extend(prefix_swap_in) scheduler_outputs.blocks_to_swap_out.extend(prefix_swap_out) now = time.time() - - if not self.cache_config.enable_prefix_caching: - common_computed_block_nums = [] - - allow_async_output_proc: bool = self.use_async_output_proc - - # Create input data structures. - seq_group_metadata_list: List[SequenceGroupMetadata] = [] - for i, scheduled_seq_group in enumerate( - scheduler_outputs.scheduled_seq_groups): - seq_group = scheduled_seq_group.seq_group - token_chunk_size = scheduled_seq_group.token_chunk_size - seq_group.maybe_set_first_scheduled_time(now) - - seq_group_metadata = self._seq_group_metadata_cache[ - self.cache_id].get_object() - seq_group_metadata.seq_data.clear() - seq_group_metadata.block_tables.clear() - - # seq_id -> SequenceData - seq_data: Dict[int, SequenceData] = {} - # seq_id -> physical block numbers - block_tables: Dict[int, List[int]] = {} - - if seq_group.is_encoder_decoder(): - # Encoder associated with SequenceGroup - encoder_seq = seq_group.get_encoder_seq() - assert encoder_seq is not None - encoder_seq_data = encoder_seq.data - # Block table for cross-attention - # Also managed at SequenceGroup level - cross_block_table = self.block_manager.get_cross_block_table( - seq_group) - else: - encoder_seq_data = None - cross_block_table = None - - for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): - seq_id = seq.seq_id - seq_data[seq_id] = seq.data - block_tables[seq_id] = self.block_manager.get_block_table(seq) - self.block_manager.access_all_blocks_in_seq(seq, now) - + + if not self.cache_config.enable_prefix_caching: + common_computed_block_nums = [] + + allow_async_output_proc: bool = self.use_async_output_proc + + # Create input data structures. + seq_group_metadata_list: List[SequenceGroupMetadata] = [] + for i, scheduled_seq_group in enumerate( + scheduler_outputs.scheduled_seq_groups): + seq_group = scheduled_seq_group.seq_group + token_chunk_size = scheduled_seq_group.token_chunk_size + seq_group.maybe_set_first_scheduled_time(now) + + seq_group_metadata = self._seq_group_metadata_cache[ + self.cache_id].get_object() + seq_group_metadata.seq_data.clear() + seq_group_metadata.block_tables.clear() + + # seq_id -> SequenceData + seq_data: Dict[int, SequenceData] = {} + # seq_id -> physical block numbers + block_tables: Dict[int, List[int]] = {} + + if seq_group.is_encoder_decoder(): + # Encoder associated with SequenceGroup + encoder_seq = seq_group.get_encoder_seq() + assert encoder_seq is not None + encoder_seq_data = encoder_seq.data + # Block table for cross-attention + # Also managed at SequenceGroup level + cross_block_table = self.block_manager.get_cross_block_table( + seq_group) + else: + encoder_seq_data = None + cross_block_table = None + + for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): + seq_id = seq.seq_id + seq_data[seq_id] = seq.data + block_tables[seq_id] = self.block_manager.get_block_table(seq) + self.block_manager.access_all_blocks_in_seq(seq, now) + common_computed_block_nums = [] if self.cache_config.enable_prefix_caching: raw_computed_block_nums = list( @@ -1508,11 +1508,11 @@ class Scheduler: seq_group.get_seqs(status=SequenceStatus.RUNNING))) if not seq_group.is_prefill(): common_computed_block_nums = raw_computed_block_nums - - do_sample = True - is_prompt = seq_group.is_prefill() - # We should send the metadata to workers when the first prefill - # is sent. Subsequent requests could be chunked prefill or decode. + + do_sample = True + is_prompt = seq_group.is_prefill() + # We should send the metadata to workers when the first prefill + # is sent. Subsequent requests could be chunked prefill or decode. is_first_prefill = False gdn_restore_key = None gdn_capture_points = None @@ -1523,8 +1523,8 @@ class Scheduler: gdn_evict_keys = [] gdn_segment_offsets = [] seqs = seq_group.get_seqs() - # Prefill has only 1 sequence. - assert len(seqs) == 1 + # Prefill has only 1 sequence. + assert len(seqs) == 1 num_computed_tokens = seqs[0].data.get_num_computed_tokens() is_first_prefill = num_computed_tokens == 0 logical_end_tokens = min( @@ -1606,11 +1606,11 @@ class Scheduler: seqs[0], len(raw_computed_block_nums), gdn_restore_key, capture_actions, gdn_evict_keys, self._gdn_prefix_policy.policy) - # In the next iteration, all prompt tokens are not computed. - # It means the prefill is chunked, and we don't need sampling. - # NOTE: We use get_len instead of get_prompt_len because when - # a sequence is preempted, prefill includes previous generated - # output tokens. + # In the next iteration, all prompt tokens are not computed. + # It means the prefill is chunked, and we don't need sampling. + # NOTE: We use get_len instead of get_prompt_len because when + # a sequence is preempted, prefill includes previous generated + # output tokens. if (token_chunk_size + num_computed_tokens < seqs[0].data.get_len()): do_sample = False @@ -1620,49 +1620,49 @@ class Scheduler: None) self._gdn_request_capture_targets.pop(seq_group.request_id, None) - - # It assumes the scheduled_seq_groups is ordered by - # prefill < decoding. - if is_first_prefill or not self.scheduler_config.send_delta_data: - seq_group_metadata = SequenceGroupMetadata( - request_id=seq_group.request_id, - is_prompt=is_prompt, - seq_data=seq_data, - sampling_params=seq_group.sampling_params, - block_tables=block_tables, - do_sample=do_sample, - pooling_params=seq_group.pooling_params, - token_chunk_size=token_chunk_size, - lora_request=seq_group.lora_request, - computed_block_nums=common_computed_block_nums, - encoder_seq_data=encoder_seq_data, - cross_block_table=cross_block_table, - state=seq_group.state, - # `multi_modal_data` will only be present for the 1st comm - # between engine and worker. - # the subsequent comms can still use delta, but - # `multi_modal_data` will be None. - multi_modal_data=seq_group.multi_modal_data - if scheduler_outputs.num_prefill_groups > 0 else None, - mm_processor_kwargs=seq_group.mm_processor_kwargs, + + # It assumes the scheduled_seq_groups is ordered by + # prefill < decoding. + if is_first_prefill or not self.scheduler_config.send_delta_data: + seq_group_metadata = SequenceGroupMetadata( + request_id=seq_group.request_id, + is_prompt=is_prompt, + seq_data=seq_data, + sampling_params=seq_group.sampling_params, + block_tables=block_tables, + do_sample=do_sample, + pooling_params=seq_group.pooling_params, + token_chunk_size=token_chunk_size, + lora_request=seq_group.lora_request, + computed_block_nums=common_computed_block_nums, + encoder_seq_data=encoder_seq_data, + cross_block_table=cross_block_table, + state=seq_group.state, + # `multi_modal_data` will only be present for the 1st comm + # between engine and worker. + # the subsequent comms can still use delta, but + # `multi_modal_data` will be None. + multi_modal_data=seq_group.multi_modal_data + if scheduler_outputs.num_prefill_groups > 0 else None, + mm_processor_kwargs=seq_group.mm_processor_kwargs, prompt_adapter_request=seq_group.prompt_adapter_request, gdn_restore_key=gdn_restore_key, gdn_capture_points=gdn_capture_points, gdn_evict_keys=gdn_evict_keys, gdn_segment_offsets=gdn_segment_offsets, ) - else: - # When SPMD mode is enabled, we only send delta data except for - # the first request to reduce serialization cost. - seq_data_delta = {} - for id, data in seq_data.items(): - seq_data_delta[id] = data.get_delta_and_reset() - seq_group_metadata = SequenceGroupMetadataDelta( - seq_data_delta, - seq_group.request_id, - block_tables, - is_prompt, - do_sample=do_sample, + else: + # When SPMD mode is enabled, we only send delta data except for + # the first request to reduce serialization cost. + seq_data_delta = {} + for id, data in seq_data.items(): + seq_data_delta[id] = data.get_delta_and_reset() + seq_group_metadata = SequenceGroupMetadataDelta( + seq_data_delta, + seq_group.request_id, + block_tables, + is_prompt, + do_sample=do_sample, token_chunk_size=token_chunk_size, computed_block_nums=common_computed_block_nums, gdn_restore_key=gdn_restore_key, @@ -1670,55 +1670,55 @@ class Scheduler: gdn_evict_keys=gdn_evict_keys, gdn_segment_offsets=gdn_segment_offsets, ) - seq_group_metadata_list.append(seq_group_metadata) - - if allow_async_output_proc: - allow_async_output_proc = self._allow_async_output_proc( - seq_group) - - # Now that the batch has been created, we can assume all blocks in the - # batch will have been computed before the next scheduling invocation. - # This is because the engine assumes that a failure in model execution - # will crash the vLLM instance / will not retry. - for scheduled_seq_group in scheduler_outputs.scheduled_seq_groups: - self.block_manager.mark_blocks_as_computed( - scheduled_seq_group.seq_group, - scheduled_seq_group.token_chunk_size) - - self._seq_group_metadata_cache[self.next_cache_id].reset() - - scheduler_time = time.perf_counter() - scheduler_start_time - # Add this to scheduler time to all the sequences that are currently - # running. This will help estimate if the scheduler is a significant - # component in the e2e latency. - for seq_group in self.running: - if seq_group is not None and seq_group.metrics is not None: - if seq_group.metrics.scheduler_time is not None: - seq_group.metrics.scheduler_time += scheduler_time - else: - seq_group.metrics.scheduler_time = scheduler_time - - # Move to next cache (if exists) - self.cache_id = self.next_cache_id - - # Return results - return (seq_group_metadata_list, scheduler_outputs, - allow_async_output_proc) - - def fork_seq(self, parent_seq: Sequence, child_seq: Sequence) -> None: - self.block_manager.fork(parent_seq, child_seq) - - def free_seq(self, seq: Sequence) -> None: - """Free a sequence from a block table.""" - self.block_manager.free(seq) - - def _free_finished_seqs(self, seq_group: SequenceGroup) -> None: - """Free finished seqs in a sequence group.""" - for seq in seq_group.get_seqs(): - if seq.is_finished(): - self.free_seq(seq) - - def _free_finished_seq_group(self, seq_group: SequenceGroup) -> None: + seq_group_metadata_list.append(seq_group_metadata) + + if allow_async_output_proc: + allow_async_output_proc = self._allow_async_output_proc( + seq_group) + + # Now that the batch has been created, we can assume all blocks in the + # batch will have been computed before the next scheduling invocation. + # This is because the engine assumes that a failure in model execution + # will crash the vLLM instance / will not retry. + for scheduled_seq_group in scheduler_outputs.scheduled_seq_groups: + self.block_manager.mark_blocks_as_computed( + scheduled_seq_group.seq_group, + scheduled_seq_group.token_chunk_size) + + self._seq_group_metadata_cache[self.next_cache_id].reset() + + scheduler_time = time.perf_counter() - scheduler_start_time + # Add this to scheduler time to all the sequences that are currently + # running. This will help estimate if the scheduler is a significant + # component in the e2e latency. + for seq_group in self.running: + if seq_group is not None and seq_group.metrics is not None: + if seq_group.metrics.scheduler_time is not None: + seq_group.metrics.scheduler_time += scheduler_time + else: + seq_group.metrics.scheduler_time = scheduler_time + + # Move to next cache (if exists) + self.cache_id = self.next_cache_id + + # Return results + return (seq_group_metadata_list, scheduler_outputs, + allow_async_output_proc) + + def fork_seq(self, parent_seq: Sequence, child_seq: Sequence) -> None: + self.block_manager.fork(parent_seq, child_seq) + + def free_seq(self, seq: Sequence) -> None: + """Free a sequence from a block table.""" + self.block_manager.free(seq) + + def _free_finished_seqs(self, seq_group: SequenceGroup) -> None: + """Free finished seqs in a sequence group.""" + for seq in seq_group.get_seqs(): + if seq.is_finished(): + self.free_seq(seq) + + def _free_finished_seq_group(self, seq_group: SequenceGroup) -> None: if seq_group.is_finished(): # Free cross-attention block table, if it exists self._free_seq_group_cross_attn_blocks(seq_group) @@ -1727,269 +1727,269 @@ class Scheduler: self._gdn_request_capture_targets.pop(seq_group.request_id, None) # Add the finished requests to the finished requests list. - # This list will be used to update the Mamba cache in the - # next step. - self._finished_requests_ids.append(seq_group.request_id) - - # Free finished seqs - self._free_finished_seqs(seq_group) - - def free_finished_seq_groups(self) -> None: - remaining: Deque[SequenceGroup] = deque() - for seq_group in self.running: - self._free_finished_seq_group(seq_group) - if not seq_group.is_finished(): - remaining.append(seq_group) - - self.running = remaining - - # Handle async stopped sequence groups - # (ones that reached max model len) - if self._async_stopped: - for seq_group in self._async_stopped: - self._free_seq_group_cross_attn_blocks(seq_group) - self._finished_requests_ids.append(seq_group.request_id) - - # Free finished seqs - self._free_finished_seqs(seq_group) - - self._async_stopped.clear() - - def _allocate_and_set_running(self, seq_group: SequenceGroup) -> None: - self.block_manager.allocate(seq_group) - for seq in seq_group.get_seqs(status=SequenceStatus.WAITING): - seq.status = SequenceStatus.RUNNING - - def _append_slots(self, - seq_group: SequenceGroup, - blocks_to_copy: List[Tuple[int, int]], - enable_chunking: bool = False) -> None: - """Appends new slots to the sequences in the given sequence group. - - Args: - seq_group (SequenceGroup): The sequence group containing the - sequences to append slots to. - blocks_to_copy (List[Tuple[int, int]]): A list of tuple of two - ints, the first int is the source block index, and the second - int is the destination block index. This list is updated with - the new source and destination block indices for the appended - slots. - enable_chunking (bool): True if chunked prefill is enabled. - """ - is_prefill: bool = seq_group.is_prefill() - num_lookahead_slots: int = self._get_num_lookahead_slots( - is_prefill, enable_chunking) - - seq_group.init_multi_step_from_lookahead_slots( - num_lookahead_slots, - num_scheduler_steps=self.scheduler_config.num_scheduler_steps, - is_multi_step=self.scheduler_config.is_multi_step, - enable_chunking=enable_chunking) - - seq_status: Optional[SequenceStatus] = SequenceStatus.RUNNING - if self.scheduler_config.is_multi_step and enable_chunking: - # In multi-step chunked-prefill any sequence type can have - # slots appended. - seq_status = None - - for seq in seq_group.get_seqs(status=seq_status): - cows = self.block_manager.append_slots(seq, num_lookahead_slots) - if len(cows) > 0: - blocks_to_copy.extend(cows) - - def _preempt( - self, - seq_group: SequenceGroup, - blocks_to_swap_out: List[Tuple[int, int]], - preemption_mode: Optional[PreemptionMode] = None, - ) -> PreemptionMode: - # If preemption mode is not specified, we determine the mode as follows: - # We use recomputation by default since it incurs lower overhead than - # swapping. However, when the sequence group has multiple sequences - # (e.g., beam search), recomputation is not currently supported. In - # such a case, we use swapping instead. - # FIXME(woosuk): This makes our scheduling policy a bit bizarre. - # As swapped sequences are prioritized over waiting sequences, - # sequence groups with multiple sequences are implicitly prioritized - # over sequence groups with a single sequence. - # TODO(woosuk): Support recomputation for sequence groups with multiple - # sequences. This may require a more sophisticated CUDA kernel. - if self.user_specified_preemption_mode is None: - if seq_group.get_max_num_running_seqs() == 1: - preemption_mode = PreemptionMode.RECOMPUTE - else: - preemption_mode = PreemptionMode.SWAP - - elif self.user_specified_preemption_mode == "swap": - preemption_mode = PreemptionMode.SWAP - else: - preemption_mode = PreemptionMode.RECOMPUTE - - if self.num_cumulative_preemption % 50 == 0: - logger.warning( - "Sequence group %s is preempted by %s mode because there is " - "not enough KV cache space. This can affect the end-to-end " - "performance. Increase gpu_memory_utilization or " - "tensor_parallel_size to provide more KV cache memory. " - "total_num_cumulative_preemption=%d", seq_group.request_id, - preemption_mode, self.num_cumulative_preemption + 1) - self.num_cumulative_preemption += 1 - - if preemption_mode == PreemptionMode.RECOMPUTE: - self._preempt_by_recompute(seq_group) - elif preemption_mode == PreemptionMode.SWAP: - self._preempt_by_swap(seq_group, blocks_to_swap_out) - else: - raise AssertionError("Invalid preemption mode.") - return preemption_mode - - def _preempt_by_recompute( - self, - seq_group: SequenceGroup, - ) -> None: - seqs = seq_group.get_seqs(status=SequenceStatus.RUNNING) - assert len(seqs) == 1 - for seq in seqs: - seq.status = SequenceStatus.WAITING - self.free_seq(seq) - seq.reset_state_for_recompute() - - def _preempt_by_swap( - self, - seq_group: SequenceGroup, - blocks_to_swap_out: List[Tuple[int, int]], - ) -> None: - self._swap_out(seq_group, blocks_to_swap_out) - - def _swap_in( - self, - seq_group: SequenceGroup, - blocks_to_swap_in: List[Tuple[int, int]], - ) -> None: - mapping = self.block_manager.swap_in(seq_group) - blocks_to_swap_in.extend(mapping) - for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED): - seq.status = SequenceStatus.RUNNING - - def _swap_out( - self, - seq_group: SequenceGroup, - blocks_to_swap_out: List[Tuple[int, int]], - ) -> None: - if not self.block_manager.can_swap_out(seq_group): - # FIXME(woosuk): Abort the sequence group instead of aborting the - # entire engine. - raise RuntimeError( - "Aborted due to the lack of CPU swap space. Please increase " - "the swap space to avoid this error.") - mapping = self.block_manager.swap_out(seq_group) - blocks_to_swap_out.extend(mapping) - for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): - seq.status = SequenceStatus.SWAPPED - - def _passed_delay(self, now: float) -> bool: - if self.prev_prompt: - self.last_prompt_latency = now - self.prev_time - self.prev_time, self.prev_prompt = now, False - # Delay scheduling prompts to let waiting queue fill up - if self.scheduler_config.delay_factor > 0 and self.waiting: - earliest_arrival_time = min( - [e.metrics.arrival_time for e in self.waiting]) - passed_delay = ( - (now - earliest_arrival_time) > - (self.scheduler_config.delay_factor * self.last_prompt_latency) - or not self.running) - else: - passed_delay = True - return passed_delay - - def _get_num_lookahead_slots(self, is_prefill: bool, - enable_chunking: bool) -> int: - """The number of slots to allocate per sequence per step, beyond known - token ids. Speculative decoding uses these slots to store KV activations - of tokens which may or may not be accepted. - - Speculative decoding does not yet support prefill, so we do not perform - lookahead allocation for prefill. - - When chunking is enabled with multi-step, we allocate lookahead slots - for the prefills for when the prefills turn into decodes in the first - step. - """ - if is_prefill: - if self.scheduler_config.is_multi_step and enable_chunking: - # num_lookahead_slots was introduced in the context of decodes, - # in Speculative Decoding. - # When the num_scheduler_steps is 8, say, then the - # num_lookahead_slots is 7. Meaning, we are doing a 1-step of - # decode anyways and we wish to do 7 more. - # - # "lookaheads" for prefills, is introduced in support for - # Chunked-Prefill in Multi-Step. - return self.scheduler_config.num_lookahead_slots + 1 - else: - return 0 - - return self.scheduler_config.num_lookahead_slots - - def _get_num_new_tokens(self, seq_group: SequenceGroup, - status: SequenceStatus, enable_chunking: bool, - budget: SchedulingBudget) -> int: - """Get the next new tokens to compute for a given sequence group - that's in a given `status`. - - The API could chunk the number of tokens to compute based on `budget` - if `enable_chunking` is True. If a sequence group has multiple - sequences (e.g., running beam search), it means it is in decoding - phase, so chunking doesn't happen. - - Returns 0 if the new token cannot be computed due to token budget. - """ - num_new_tokens = 0 - seqs = seq_group.get_seqs(status=status) - for seq in seqs: - num_new_tokens += seq.get_num_new_tokens() - assert num_new_tokens > 0 - # Chunk if a running request cannot fit in the given budget. - # If number of seq > 1, it means it is doing beam search - # in a decode phase. Do not chunk. - if enable_chunking and len(seqs) == 1: - remaining_token_budget = budget.remaining_token_budget() - if self.scheduler_config.is_multi_step: - # The current multi-step + chunked prefill capability does - # not actually support chunking prompts. - # - # Therefore, `num_new_tokens` is computed in the same fashion - # for both multi-step+chunked-prefill & - # multi-step+chunked-prefill+APC - # - # Prompts with more tokens than the current remaining budget - # are postponed to future scheduler steps - if num_new_tokens > self._get_prompt_limit(seq_group): - # If the seq_group is in prompt-stage, pass the - # num_new_tokens as-is so the caller can ignore - # the sequence. - pass - else: - num_new_tokens = 0 \ - if num_new_tokens > remaining_token_budget \ - else num_new_tokens - elif self.cache_config.enable_prefix_caching: - # When prefix caching is enabled, we always allocate - # the number of new tokens that is dividable by the block - # size to avoid partial block matching. - block_size = self.cache_config.block_size - remainder = budget.token_budget % block_size - if remainder != 0: - raise ValueError("When enabling chunked prefill and " - "prefix caching, max_num_batched_tokens " - "(chunk size) must be dividable by " - "block size, but got chunk_size " - f"({budget.token_budget}) % block_size " - f"({block_size}) = {remainder}") - if remaining_token_budget < num_new_tokens: - num_new_tokens = (remaining_token_budget // - block_size) * block_size - else: - num_new_tokens = min(num_new_tokens, remaining_token_budget) - return num_new_tokens + # This list will be used to update the Mamba cache in the + # next step. + self._finished_requests_ids.append(seq_group.request_id) + + # Free finished seqs + self._free_finished_seqs(seq_group) + + def free_finished_seq_groups(self) -> None: + remaining: Deque[SequenceGroup] = deque() + for seq_group in self.running: + self._free_finished_seq_group(seq_group) + if not seq_group.is_finished(): + remaining.append(seq_group) + + self.running = remaining + + # Handle async stopped sequence groups + # (ones that reached max model len) + if self._async_stopped: + for seq_group in self._async_stopped: + self._free_seq_group_cross_attn_blocks(seq_group) + self._finished_requests_ids.append(seq_group.request_id) + + # Free finished seqs + self._free_finished_seqs(seq_group) + + self._async_stopped.clear() + + def _allocate_and_set_running(self, seq_group: SequenceGroup) -> None: + self.block_manager.allocate(seq_group) + for seq in seq_group.get_seqs(status=SequenceStatus.WAITING): + seq.status = SequenceStatus.RUNNING + + def _append_slots(self, + seq_group: SequenceGroup, + blocks_to_copy: List[Tuple[int, int]], + enable_chunking: bool = False) -> None: + """Appends new slots to the sequences in the given sequence group. + + Args: + seq_group (SequenceGroup): The sequence group containing the + sequences to append slots to. + blocks_to_copy (List[Tuple[int, int]]): A list of tuple of two + ints, the first int is the source block index, and the second + int is the destination block index. This list is updated with + the new source and destination block indices for the appended + slots. + enable_chunking (bool): True if chunked prefill is enabled. + """ + is_prefill: bool = seq_group.is_prefill() + num_lookahead_slots: int = self._get_num_lookahead_slots( + is_prefill, enable_chunking) + + seq_group.init_multi_step_from_lookahead_slots( + num_lookahead_slots, + num_scheduler_steps=self.scheduler_config.num_scheduler_steps, + is_multi_step=self.scheduler_config.is_multi_step, + enable_chunking=enable_chunking) + + seq_status: Optional[SequenceStatus] = SequenceStatus.RUNNING + if self.scheduler_config.is_multi_step and enable_chunking: + # In multi-step chunked-prefill any sequence type can have + # slots appended. + seq_status = None + + for seq in seq_group.get_seqs(status=seq_status): + cows = self.block_manager.append_slots(seq, num_lookahead_slots) + if len(cows) > 0: + blocks_to_copy.extend(cows) + + def _preempt( + self, + seq_group: SequenceGroup, + blocks_to_swap_out: List[Tuple[int, int]], + preemption_mode: Optional[PreemptionMode] = None, + ) -> PreemptionMode: + # If preemption mode is not specified, we determine the mode as follows: + # We use recomputation by default since it incurs lower overhead than + # swapping. However, when the sequence group has multiple sequences + # (e.g., beam search), recomputation is not currently supported. In + # such a case, we use swapping instead. + # FIXME(woosuk): This makes our scheduling policy a bit bizarre. + # As swapped sequences are prioritized over waiting sequences, + # sequence groups with multiple sequences are implicitly prioritized + # over sequence groups with a single sequence. + # TODO(woosuk): Support recomputation for sequence groups with multiple + # sequences. This may require a more sophisticated CUDA kernel. + if self.user_specified_preemption_mode is None: + if seq_group.get_max_num_running_seqs() == 1: + preemption_mode = PreemptionMode.RECOMPUTE + else: + preemption_mode = PreemptionMode.SWAP + + elif self.user_specified_preemption_mode == "swap": + preemption_mode = PreemptionMode.SWAP + else: + preemption_mode = PreemptionMode.RECOMPUTE + + if self.num_cumulative_preemption % 50 == 0: + logger.warning( + "Sequence group %s is preempted by %s mode because there is " + "not enough KV cache space. This can affect the end-to-end " + "performance. Increase gpu_memory_utilization or " + "tensor_parallel_size to provide more KV cache memory. " + "total_num_cumulative_preemption=%d", seq_group.request_id, + preemption_mode, self.num_cumulative_preemption + 1) + self.num_cumulative_preemption += 1 + + if preemption_mode == PreemptionMode.RECOMPUTE: + self._preempt_by_recompute(seq_group) + elif preemption_mode == PreemptionMode.SWAP: + self._preempt_by_swap(seq_group, blocks_to_swap_out) + else: + raise AssertionError("Invalid preemption mode.") + return preemption_mode + + def _preempt_by_recompute( + self, + seq_group: SequenceGroup, + ) -> None: + seqs = seq_group.get_seqs(status=SequenceStatus.RUNNING) + assert len(seqs) == 1 + for seq in seqs: + seq.status = SequenceStatus.WAITING + self.free_seq(seq) + seq.reset_state_for_recompute() + + def _preempt_by_swap( + self, + seq_group: SequenceGroup, + blocks_to_swap_out: List[Tuple[int, int]], + ) -> None: + self._swap_out(seq_group, blocks_to_swap_out) + + def _swap_in( + self, + seq_group: SequenceGroup, + blocks_to_swap_in: List[Tuple[int, int]], + ) -> None: + mapping = self.block_manager.swap_in(seq_group) + blocks_to_swap_in.extend(mapping) + for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED): + seq.status = SequenceStatus.RUNNING + + def _swap_out( + self, + seq_group: SequenceGroup, + blocks_to_swap_out: List[Tuple[int, int]], + ) -> None: + if not self.block_manager.can_swap_out(seq_group): + # FIXME(woosuk): Abort the sequence group instead of aborting the + # entire engine. + raise RuntimeError( + "Aborted due to the lack of CPU swap space. Please increase " + "the swap space to avoid this error.") + mapping = self.block_manager.swap_out(seq_group) + blocks_to_swap_out.extend(mapping) + for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): + seq.status = SequenceStatus.SWAPPED + + def _passed_delay(self, now: float) -> bool: + if self.prev_prompt: + self.last_prompt_latency = now - self.prev_time + self.prev_time, self.prev_prompt = now, False + # Delay scheduling prompts to let waiting queue fill up + if self.scheduler_config.delay_factor > 0 and self.waiting: + earliest_arrival_time = min( + [e.metrics.arrival_time for e in self.waiting]) + passed_delay = ( + (now - earliest_arrival_time) > + (self.scheduler_config.delay_factor * self.last_prompt_latency) + or not self.running) + else: + passed_delay = True + return passed_delay + + def _get_num_lookahead_slots(self, is_prefill: bool, + enable_chunking: bool) -> int: + """The number of slots to allocate per sequence per step, beyond known + token ids. Speculative decoding uses these slots to store KV activations + of tokens which may or may not be accepted. + + Speculative decoding does not yet support prefill, so we do not perform + lookahead allocation for prefill. + + When chunking is enabled with multi-step, we allocate lookahead slots + for the prefills for when the prefills turn into decodes in the first + step. + """ + if is_prefill: + if self.scheduler_config.is_multi_step and enable_chunking: + # num_lookahead_slots was introduced in the context of decodes, + # in Speculative Decoding. + # When the num_scheduler_steps is 8, say, then the + # num_lookahead_slots is 7. Meaning, we are doing a 1-step of + # decode anyways and we wish to do 7 more. + # + # "lookaheads" for prefills, is introduced in support for + # Chunked-Prefill in Multi-Step. + return self.scheduler_config.num_lookahead_slots + 1 + else: + return 0 + + return self.scheduler_config.num_lookahead_slots + + def _get_num_new_tokens(self, seq_group: SequenceGroup, + status: SequenceStatus, enable_chunking: bool, + budget: SchedulingBudget) -> int: + """Get the next new tokens to compute for a given sequence group + that's in a given `status`. + + The API could chunk the number of tokens to compute based on `budget` + if `enable_chunking` is True. If a sequence group has multiple + sequences (e.g., running beam search), it means it is in decoding + phase, so chunking doesn't happen. + + Returns 0 if the new token cannot be computed due to token budget. + """ + num_new_tokens = 0 + seqs = seq_group.get_seqs(status=status) + for seq in seqs: + num_new_tokens += seq.get_num_new_tokens() + assert num_new_tokens > 0 + # Chunk if a running request cannot fit in the given budget. + # If number of seq > 1, it means it is doing beam search + # in a decode phase. Do not chunk. + if enable_chunking and len(seqs) == 1: + remaining_token_budget = budget.remaining_token_budget() + if self.scheduler_config.is_multi_step: + # The current multi-step + chunked prefill capability does + # not actually support chunking prompts. + # + # Therefore, `num_new_tokens` is computed in the same fashion + # for both multi-step+chunked-prefill & + # multi-step+chunked-prefill+APC + # + # Prompts with more tokens than the current remaining budget + # are postponed to future scheduler steps + if num_new_tokens > self._get_prompt_limit(seq_group): + # If the seq_group is in prompt-stage, pass the + # num_new_tokens as-is so the caller can ignore + # the sequence. + pass + else: + num_new_tokens = 0 \ + if num_new_tokens > remaining_token_budget \ + else num_new_tokens + elif self.cache_config.enable_prefix_caching: + # When prefix caching is enabled, we always allocate + # the number of new tokens that is dividable by the block + # size to avoid partial block matching. + block_size = self.cache_config.block_size + remainder = budget.token_budget % block_size + if remainder != 0: + raise ValueError("When enabling chunked prefill and " + "prefix caching, max_num_batched_tokens " + "(chunk size) must be dividable by " + "block size, but got chunk_size " + f"({budget.token_budget}) % block_size " + f"({block_size}) = {remainder}") + if remaining_token_budget < num_new_tokens: + num_new_tokens = (remaining_token_budget // + block_size) * block_size + else: + num_new_tokens = min(num_new_tokens, remaining_token_budget) + return num_new_tokens diff --git a/qwen3_6_scripts/sequence.py b/qwen3_6_scripts/sequence.py index b0bd8ffc..f9b41db2 100644 --- a/qwen3_6_scripts/sequence.py +++ b/qwen3_6_scripts/sequence.py @@ -1,943 +1,943 @@ -"""Sequence and its related classes.""" -import copy -import enum -from abc import ABC, abstractmethod -from array import array -from collections import defaultdict -from dataclasses import dataclass -from functools import cached_property, reduce -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional -from typing import Sequence as GenericSequence -from typing import Set, Tuple, Union, cast - -import msgspec -import torch - -from vllm.inputs import EncoderDecoderLLMInputs, LLMInputs -from vllm.inputs.parse import is_valid_encoder_decoder_llm_inputs -from vllm.lora.request import LoRARequest -from vllm.pooling_params import PoolingParams -from vllm.prompt_adapter.request import PromptAdapterRequest -from vllm.sampling_params import SamplingParams -from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics - -if TYPE_CHECKING: - from vllm.multimodal.base import MultiModalDataDict - -VLLM_TOKEN_ID_ARRAY_TYPE = "l" - -VLLM_INVALID_TOKEN_ID = -1 - - -# We use dataclass for now because it is used for -# openai server output, and msgspec is not serializable. -# TODO(sang): Fix it. -@dataclass -class Logprob: - """Infos for supporting OpenAI compatible logprobs and token ranks. - - Attributes: - logprob: The logprob of chosen token - rank: The vocab rank of chosen token (>=1) - decoded_token: The decoded chosen token index - """ - logprob: float - rank: Optional[int] = None - decoded_token: Optional[str] = None - - -# {token_id -> logprob} per each sequence group. None if the corresponding -# sequence group doesn't require prompt logprob. -PromptLogprobs = List[Optional[Dict[int, Logprob]]] -# {token_id -> logprob} for each sequence group. -SampleLogprobs = List[Dict[int, Logprob]] - - -class SequenceStatus(enum.IntEnum): - """Status of a sequence.""" - WAITING = 0 - RUNNING = 1 - SWAPPED = 2 - # Note: anything after SWAPPED (2) will be considered - # as a finished status. - FINISHED_STOPPED = 3 - FINISHED_LENGTH_CAPPED = 4 - FINISHED_ABORTED = 5 - FINISHED_IGNORED = 6 - - @staticmethod - def is_finished(status: "SequenceStatus") -> bool: - return status > SequenceStatus.SWAPPED - - @staticmethod - def get_finished_reason(status: "SequenceStatus") -> Union[str, None]: - if status == SequenceStatus.FINISHED_STOPPED: - finish_reason = "stop" - elif status == SequenceStatus.FINISHED_LENGTH_CAPPED: - finish_reason = "length" - elif status == SequenceStatus.FINISHED_ABORTED: - finish_reason = "abort" - elif status == SequenceStatus.FINISHED_IGNORED: - # The ignored sequences are the sequences whose prompt lengths - # are longer than the model's length cap. Therefore, the stop - # reason should also be "length" as in OpenAI API. - finish_reason = "length" - else: - finish_reason = None - return finish_reason - - -class SequenceStage(enum.Enum): - PREFILL = enum.auto() - DECODE = enum.auto() - - -@dataclass -class RequestMetrics: - """Metrics associated with a request. - - Attributes: - arrival_time: The time when the request arrived. - first_scheduled_time: The time when the request was first scheduled. - first_token_time: The time when the first token was generated. - time_in_queue: The time the request spent in the queue. - finished_time: The time when the request was finished. - scheduler_time: The time spent in the scheduler when this request was - being considered by the scheduler. - model_forward_time: The time spent in the model forward pass when this - request was in the batch. - model_execute_time: The time spent in the model execute function. This - will include model forward, block/sync across - workers, cpu-gpu sync time and sampling time. - """ - arrival_time: float - last_token_time: float - first_scheduled_time: Optional[float] - first_token_time: Optional[float] - time_in_queue: Optional[float] - finished_time: Optional[float] = None - scheduler_time: Optional[float] = None - model_forward_time: Optional[float] = None - model_execute_time: Optional[float] = None - num_cached_tokens: Optional[int] = None - - -class SequenceDataDelta( - msgspec.Struct, - array_like=True, # type: ignore[call-arg] - omit_defaults=True): # type: ignore[call-arg] - """Delta SequenceData to send to workers per step.""" - # A new token to be appended to existing SequenceData. - new_output_token_ids: List[int] - # Overwriting existing `cumulative_logprob` - new_cumulative_logprob: float - # Overwriting existing `num_computed_tokens`. - new_num_computed_tokens: int - # Overwriting existing `stage`. - new_stage: SequenceStage - - -class SequenceData(msgspec.Struct, - omit_defaults=True): # type: ignore[call-arg] - """Data associated with a sequence. - - Args: - prompt_token_ids: The token IDs of the prompt. - output_token_ids: The token IDs of the output. Set to an empty list if - None. - - Attributes: - prompt_token_ids: The token IDs of the prompt. - output_token_ids: The token IDs of the output. - cumulative_logprob: The cumulative log probability of the output. - """ - # NOTE: we cannot use Union[List, array] because msgspec cannot support - # union of 2 list types. - _prompt_token_ids: array - _output_token_ids: array = msgspec.field( - default_factory=lambda: array(VLLM_TOKEN_ID_ARRAY_TYPE, [])) - - ### The below fields should not be passed as an argument ### - _cumulative_logprob: float = 0.0 - _prompt_token_ids_tuple: Tuple[int, - ...] = msgspec.field(default_factory=tuple) - # The number of tokens that are computed (that run against the model). - _num_computed_tokens: int = 0 - _stage: SequenceStage = SequenceStage.PREFILL - _cached_all_token_ids: List[int] = msgspec.field(default_factory=list) - - # It is used to get delta input. It is reset when `get_delta_and_reset` - # is called. - _new_appended_tokens: List[int] = msgspec.field(default_factory=list) - - # It is used to compute mrope_position_ids. - _mrope_position_delta: Optional[int] = None - - @staticmethod - def from_token_counts(*token_counts: Tuple[int, int]) -> "SequenceData": - if len(token_counts) == 0: - return SequenceData.from_seqs([]) - - arrs = [ - array(VLLM_TOKEN_ID_ARRAY_TYPE, [token_id]) * count - for token_id, count in token_counts - ] - - return SequenceData(reduce(array.__add__, arrs)) - - @staticmethod - def from_seqs( - prompt_token_ids: GenericSequence[int], - output_token_ids: Optional[GenericSequence[int]] = None, - ) -> "SequenceData": - prompt_token_ids_arr = array(VLLM_TOKEN_ID_ARRAY_TYPE, - prompt_token_ids) - - if output_token_ids is None: - return SequenceData(prompt_token_ids_arr) - - output_token_ids_arr = array(VLLM_TOKEN_ID_ARRAY_TYPE, - output_token_ids) - - return SequenceData(prompt_token_ids_arr, - _output_token_ids=output_token_ids_arr) - - def __post_init__(self) -> None: - assert self._prompt_token_ids.typecode == "l" - assert self._output_token_ids.typecode == "l" - self._prompt_token_ids_tuple: Tuple[int, ...] = tuple( - self._prompt_token_ids) - self._update_cached_all_tokens() - - def _update_cached_all_tokens(self): - assert isinstance(self._prompt_token_ids, array) - assert isinstance(self._output_token_ids, array) - self._cached_all_token_ids: List[int] = list(self._prompt_token_ids + - self._output_token_ids) - - @property - def cumulative_logprob(self) -> float: - return self._cumulative_logprob - - @property - def prompt_token_ids(self) -> Tuple[int, ...]: - return self._prompt_token_ids_tuple - - @prompt_token_ids.setter - def prompt_token_ids(self, new_prompt_token_ids) -> None: - raise NotImplementedError - - @property - def prompt_token_ids_array(self) -> array: - """Return the prompt token ids in array type. - - Note that the array is in "I" type, and it is not compatible - with torch.long (2 bytes vs 4 bytes). So beware of the usage. - """ - return self._prompt_token_ids - - @property - def output_token_ids(self) -> Tuple[int, ...]: - return tuple(self._output_token_ids) - - @output_token_ids.setter - def output_token_ids(self, new_output_token_ids: List[int]) -> None: - self._output_token_ids = array(VLLM_TOKEN_ID_ARRAY_TYPE, - new_output_token_ids) - self._update_cached_all_tokens() - - @property - def output_token_ids_array(self) -> array: - """Return the prompt token ids in array type. - - Note that the array is in "I" type, and it is not compatible - with torch.long (2 bytes vs 4 bytes). So beware of the usage. - """ - assert isinstance(self._output_token_ids, array) - return self._output_token_ids - - @property - def mrope_position_delta(self) -> Optional[int]: - return self._mrope_position_delta - - @mrope_position_delta.setter - def mrope_position_delta(self, new_mrope_position_delta): - self._mrope_position_delta = new_mrope_position_delta - - def append_token_id(self, token_id: int, logprob: float) -> None: - self._output_token_ids.append(token_id) - self._new_appended_tokens.append(token_id) - self._cached_all_token_ids.append(token_id) - self._cumulative_logprob += logprob - - def get_len(self) -> int: - return len(self._output_token_ids) + len(self._prompt_token_ids) - - def get_prompt_len(self) -> int: - return len(self._prompt_token_ids) - - def get_output_len(self) -> int: - return len(self._output_token_ids) - - def get_token_ids(self) -> List[int]: - return self._cached_all_token_ids - - def get_prefix_token_ids( - self, num_tokens: int - ) -> Tuple[Tuple[int, ...], Optional[Tuple[int, ...]]]: - """Get prefix tokens, and make the return value hashable""" - prompt_length = self.get_prompt_len() - if num_tokens > prompt_length: - return (self._prompt_token_ids_tuple, - tuple(self._output_token_ids[:num_tokens - prompt_length])) - else: - return (self._prompt_token_ids_tuple[:num_tokens], None) - - def get_num_computed_tokens(self) -> int: - """Return the number of prefill tokens that are already computed.""" - return self._num_computed_tokens - - def update_num_computed_tokens(self, num_new_computed_tokens: int): - """Update number of tokens computed so far.""" - self._num_computed_tokens += num_new_computed_tokens - assert self._num_computed_tokens <= self.get_len(), ( - self._num_computed_tokens, self.get_len()) - # If all tokens are computed, it means it is in decoding phase. - if self.get_num_uncomputed_tokens() == 0: - self._stage = SequenceStage.DECODE - - def reset_state_for_recompute(self) -> None: - """Reset the number of computed tokens from this sequence. It is - supposed to be called when a sequence needs to be started from - the beginning again (e.g., sequence is preempted). - """ - self._num_computed_tokens = 0 - self._stage = SequenceStage.PREFILL - self._new_appended_tokens = [] - - def get_num_uncomputed_tokens(self) -> int: - """Return the number of prefill tokens that are not computed.""" - # we use `get_len()` which includes prompt_len + output_len instead - # of prompt_len here. This is because during recompute we need to - # prefill for both prompt and output. - return self.get_len() - self.get_num_computed_tokens() - - def get_last_token_id(self) -> int: - if not self._output_token_ids: - return self._prompt_token_ids[-1] - return self._output_token_ids[-1] - - def get_prompt_token_ids(self) -> Tuple[int, ...]: - return self.prompt_token_ids - - def get_output_token_ids(self) -> Tuple[int, ...]: - return self.output_token_ids - - def get_delta_and_reset(self) -> SequenceDataDelta: - delta = SequenceDataDelta(self._new_appended_tokens, - self._cumulative_logprob, - self.get_num_computed_tokens(), self.stage) - # Reset delta state. - self._new_appended_tokens = [] - return delta - - def apply_delta(self, delta: SequenceDataDelta): - self._num_computed_tokens = delta.new_num_computed_tokens - self._cumulative_logprob = delta.new_cumulative_logprob - self._stage = delta.new_stage - self._output_token_ids.extend(delta.new_output_token_ids) - self._cached_all_token_ids.extend(delta.new_output_token_ids) - - @property - def stage(self) -> SequenceStage: - return self._stage - - def __repr__(self) -> str: - return (f"SequenceData(" - f"prompt_token_ids={self._prompt_token_ids}, " - f"output_token_ids={self.output_token_ids}, " - f"cumulative_logprob={self.cumulative_logprob}, " - f"get_num_computed_tokens={self.get_num_computed_tokens()}") - - -class Sequence: - """Stores the data, status, and block information of a sequence. - - The sequence is constructed from the LLMInputs instance passed - in through the `inputs` constructor argument. - - For encoder/decoder models, LLMInputs encapsulates both a - decoder and encoder prompt, creating an ambiguity about which - prompt to construct the sequence from. The `from_decoder_prompt` - constructor argument signals whether to construct the Sequence - from the LLMInputs decoder prompt, or encoder prompt. - - Args: - seq_id: The ID of the sequence. - inputs: The inputs of the sequence. - block_size: The block size of the sequence. Should be the same as the - block size used by the block manager and cache engine. - eos_token_id: The end-of-sequence (EOS) token id recognized by this LLM. - lora_request: LoRA request. - prompt_adapter_request: Prompt Adapter request. - from_decoder_prompt: Construct Sequence from LLMInputs decoder prompt - (True) or encoder prompt (False.) Must be True - for decoder-only model. - - """ - - def __init__( - self, - seq_id: int, - inputs: "LLMInputs", - block_size: int, - eos_token_id: Optional[int] = None, - lora_request: Optional[LoRARequest] = None, - prompt_adapter_request: Optional[PromptAdapterRequest] = None, - from_decoder_prompt: bool = True, - ) -> None: - self.seq_id = seq_id - self.inputs = inputs - self.block_size = block_size - self.eos_token_id = eos_token_id - self.lora_request = lora_request - self.prompt_adapter_request = prompt_adapter_request - self.from_decoder_prompt = from_decoder_prompt - - # For decoder-only models, a Sequence is constructed - # from an LLMInputs instance (the `inputs` arg.) - # - # For encoder/decoder models the same `inputs` - # instance could be utilized to construct either an - # encoder sequence or a decoder sequence, because - # `LLMInputs` has both decoder- and encoder-oriented - # member variables (i.e. it encapsulates both an encoder - # and a decoder prompt.) The decision of which type of sequence - # to generate is determined by the `from_decoder_prompt` argument. - # - # When constructing a encoder sequence - # (`from_decoder_prompt` False) it matters that - # the `LLMInputs` instance stored in `inputs` is valid - # in the sense that its encoder-related member variables are - # populated; below, an exception is raised if this is - # not the case. - # - # When constructing a decoder sequence (`from_decoder_prompt` True) - # it does not matter whether `inputs` has its encoder-related - # member variables populated. - if not (from_decoder_prompt - or is_valid_encoder_decoder_llm_inputs(inputs)): - raise ValueError("Cannot extract encoder input prompt from " - f"invalid input {inputs}; did you forget the " - "encoder input prompt fields?") - - self.data = SequenceData.from_seqs(self.prompt_token_ids) - self.output_logprobs: SampleLogprobs = [] - self.output_text = "" - - self.status = SequenceStatus.WAITING - self.stop_reason: Union[int, str, None] = None - - # These are used to keep track of delta outputs - self._last_output_token_ids_offset: int = 0 - self._last_output_text_offset: int = 0 - - # Used for incremental detokenization - self.prefix_offset = 0 - self.read_offset = 0 - # Input + output tokens - self.tokens: Optional[List[str]] = None - - @property - def n_blocks(self) -> int: - return (self.get_len() + self.block_size - 1) // self.block_size - - @cached_property - def prompt(self) -> Optional[str]: - # Select decoder or encoder input prompt str, as appropriate - prompt_key: str = ("prompt" - if self.from_decoder_prompt else "encoder_prompt") - - return cast(Optional[str], self.inputs.get(prompt_key)) - - @cached_property - def prompt_token_ids(self) -> List[int]: - # Select decoder or encoder input prompt token ids, as appropriate - prompt_token_ids_key: str = ("prompt_token_ids" - if self.from_decoder_prompt else - "encoder_prompt_token_ids") - - # Cache computed prompt token ids - return cast(List[int], self.inputs.get(prompt_token_ids_key)) - - @property - def multi_modal_data(self) -> "MultiModalDataDict": - if self.inputs.get("multi_modal_data") and self.inputs.get( - "encoder_multi_modal_data"): - raise ValueError( - "Multi-modal data in both encoder and decoder is not supported." - ) - inputs = self.inputs - return self.inputs.get("multi_modal_data") or (cast( - EncoderDecoderLLMInputs, - inputs).get("encoder_multi_modal_data")) or {} - - @property - def mm_processor_kwargs(self) -> Dict[str, Any]: - return self.inputs.get("mm_processor_kwargs") or {} - - @property - def lora_int_id(self) -> int: - return self.lora_request.lora_int_id if self.lora_request else 0 - - @property - def prompt_adapter_id(self) -> int: - return self.prompt_adapter_request.prompt_adapter_id \ - if self.prompt_adapter_request else 0 - - def get_output_text_to_return(self, buffer_length: int, - delta: bool) -> str: - """If delta is True, only new text since the last call to - this method is returned""" - - # We return the full output text if the sequence is finished. - truncate = buffer_length and not self.is_finished() - if not delta: - return self.output_text[:-buffer_length] if truncate else ( - self.output_text) - length = len(self.output_text) - if truncate: - length -= buffer_length - last_offset = self._last_output_text_offset - if last_offset < length: - self._last_output_text_offset = length - return self.output_text[last_offset:length] - return "" - - def get_output_token_ids_to_return( - self, delta: bool) -> Union[GenericSequence[int], int]: - """If delta is True, only new tokens since the last call to - this method are returned""" - if not delta: - return self.get_output_token_ids() - - output_len = self.get_output_len() - - # Get the number of new tokens - num_new_tokens = output_len - self._last_output_token_ids_offset - self._last_output_token_ids_offset = output_len - - # Return new tokens - if num_new_tokens == 0: - # During chunked prefill steps with no output yet, num_new_tokens=0. - # Python's [-0:] == [0:] returns the ENTIRE list — guard against this. - return [] - - if num_new_tokens == 1: - # Optimization for single decode token case - # (which is what we have most of the time) - return self.data._cached_all_token_ids[-1] - - return self.data._cached_all_token_ids[-num_new_tokens:] - - def hash_of_block(self, logical_idx: int) -> int: - # TODO This can produce incorrect hash when block size > prompt size - - # Compute the number of tokens in the sequence - # TODO: The current hashing function is O(L^2). We should optimize - # this in the future. - num_tokens = self.num_hashed_tokens_of_block(logical_idx) - hashed_tokens = self.data.get_prefix_token_ids(num_tokens) - return hash((hashed_tokens, self.lora_int_id)) - - def num_hashed_tokens_of_block(self, logical_idx: int): - return logical_idx * self.block_size + self.block_size - - def reset_state_for_recompute(self): - """Reset the sequence states for recomputation.""" - self.data.reset_state_for_recompute() - - def append_token_id(self, token_id: int, logprobs: Dict[int, - Logprob]) -> None: - assert token_id in logprobs - self.output_logprobs.append(logprobs) - self.data.append_token_id(token_id, logprobs[token_id].logprob) - - def get_len(self) -> int: - return self.data.get_len() - - def get_prompt_len(self) -> int: - return self.data.get_prompt_len() - - def get_output_len(self) -> int: - return self.data.get_output_len() - - def get_token_ids(self) -> List[int]: - return self.data.get_token_ids() - - def get_prompt_token_ids(self) -> Tuple[int, ...]: - return self.data.get_prompt_token_ids() - - def get_last_token_id(self) -> int: - return self.data.get_last_token_id() - - def get_output_token_ids(self) -> Tuple[int, ...]: - return self.data.get_output_token_ids() - - def get_cumulative_logprob(self) -> float: - return self.data.cumulative_logprob - - def is_finished(self) -> bool: - return SequenceStatus.is_finished(self.status) - - def fork(self, new_seq_id: int) -> "Sequence": - new_seq = copy.deepcopy(self) - new_seq.seq_id = new_seq_id - return new_seq - - def get_num_new_tokens(self) -> int: - """Get the number of new tokens to be computed. - - Returns: - The new number of tokens to be computed. I.e., 1 for decode, or - the remaining prompt size for prefill. - """ - if self.data.stage == SequenceStage.DECODE: - return 1 - return self.data.get_num_uncomputed_tokens() - - def is_prefill(self) -> bool: - return self.data.stage == SequenceStage.PREFILL - - def __repr__(self) -> str: - return (f"Sequence(seq_id={self.seq_id}, " - f"status={self.status.name}, " - f"num_blocks={self.n_blocks}, ") - - -class SequenceGroupState(msgspec.Struct, - omit_defaults=True): # type: ignore[call-arg] - """Mutable state tied to a specific sequence group""" - - # for multi-step decoding - num_steps: int = 1 - current_step: int = 0 - - @property - def remaining_steps(self) -> int: - return self.num_steps - self.current_step - - -class SequenceGroup: - """A group of sequences that are generated from the same prompt. - - Args: - request_id: The ID of the request. - seqs: The list of sequences. - sampling_params: The sampling parameters used to generate the outputs. - arrival_time: The arrival time of the request. - lora_request: LoRA request. - embeddings: The embeddings vectors of the prompt of the sequence group - for an embedding model. - pooling_params: The pooling parameters used to generate the pooling - for an embedding model. - encoder_seq: Optional, the single encoder sequence. Should be None - unless you are working with an encoder/decoder model. - trace_headers: OpenTelemetry trace headers. - prompt_adapter_request: Prompt Adapter request. - priority: User-defined priority of the request. - """ - - def __init__( - self, - request_id: str, - seqs: List[Sequence], - arrival_time: float, - sampling_params: Optional[SamplingParams] = None, - lora_request: Optional[LoRARequest] = None, - embeddings: Optional[List[float]] = None, - pooling_params: Optional[PoolingParams] = None, - encoder_seq: Optional[Sequence] = None, - trace_headers: Optional[Mapping[str, str]] = None, - prompt_adapter_request: Optional[PromptAdapterRequest] = None, - priority: int = 0, - ) -> None: - self.request_id = request_id - self.seqs = seqs - self.arrival_time = arrival_time - self.is_single_seq = len(seqs) == 1 - self.seqs_dict = {seq.seq_id: seq for seq in seqs} - - self.sampling_params = sampling_params - self.metrics = RequestMetrics(arrival_time=arrival_time, - last_token_time=arrival_time, - first_scheduled_time=None, - first_token_time=None, - time_in_queue=None) - self.lora_request = lora_request - self.prompt_logprobs: Optional[PromptLogprobs] = None - self.state = SequenceGroupState() - self.embeddings = embeddings - self.pooling_params = pooling_params - self.prompt_adapter_request = prompt_adapter_request - self.encoder_seq = encoder_seq - self.trace_headers = trace_headers - self.priority = priority - - self.cached_request_output = None - - @property - def prompt(self) -> Optional[str]: - # All sequences in the group should have the same prompt. - # We use the prompt of an arbitrary sequence. - return self.seqs[0].prompt - - @property - def prompt_token_ids(self) -> List[int]: - # All sequences in the group should have the same prompt. - # We use the prompt of an arbitrary sequence. - return self.seqs[0].prompt_token_ids - - @property - def encoder_prompt(self) -> Optional[str]: - # There are either 0 or 1 encoder sequences - # If one is present, its prompt is distinct - # from the decoder's. - return (self.encoder_seq.prompt - if self.encoder_seq is not None else None) - - @property - def encoder_prompt_token_ids(self) -> Optional[List[int]]: - # There are either 0 or 1 encoder sequences - # If one is present, its prompt token ids are - # distinct from the decoder's. - return (self.encoder_seq.prompt_token_ids - if self.encoder_seq is not None else None) - - @property - def multi_modal_data(self) -> "MultiModalDataDict": - # All sequences in the group should have the same multi-modal data. - # We use the multi-modal data of an arbitrary sequence. - return self.seqs[0].multi_modal_data - - @property - def mm_processor_kwargs(self) -> Dict[str, Any]: - # As with multi-modal data, all sequences in the group should have the - # same processor kwargs (i.e., mm_processor_kwargs are optionally - # provided per request; note that are independent of whether the model - # decoder-only or an encoder-decoder). - return self.seqs[0].mm_processor_kwargs - - @property - def lora_int_id(self) -> int: - return self.lora_request.lora_int_id if self.lora_request else 0 - - @property - def prompt_adapter_id(self) -> int: - return self.prompt_adapter_request.prompt_adapter_id \ - if self.prompt_adapter_request else 0 - - @property - def prompt_adapter_num_virtual_tokens(self) -> int: - return self.prompt_adapter_request.prompt_adapter_num_virtual_tokens\ - if self.prompt_adapter_request else 0 - - def init_multi_step(self, num_steps: int) -> None: - self.state.num_steps = num_steps - self.state.current_step = 0 - - def init_multi_step_from_lookahead_slots(self, num_lookahead_slots: int, - num_scheduler_steps: int, - is_multi_step: bool, - enable_chunking: bool) -> None: - - if not is_multi_step: - self.init_multi_step(num_steps=num_scheduler_steps) - return - - # Multi-Step case - is_prefill = self.is_prefill() - - # The asserts below reflect the expectations of the current system. - if is_prefill and enable_chunking: - assert num_lookahead_slots == num_scheduler_steps - self.init_multi_step(num_steps=num_lookahead_slots) - else: - is_decode: bool = not is_prefill - # If it is a prefill, num_lookahead_slots must be 0 - assert num_lookahead_slots == 0 or is_decode - # If it is a decode, num_lookahead_slots + 1 must match - # the scheduler steps. - assert num_lookahead_slots + 1 == num_scheduler_steps or is_prefill - self.init_multi_step(num_steps=num_lookahead_slots + 1) - - def get_last_latency(self, now: float) -> Optional[float]: - """Sets the last token time for Request level timings.""" - # If still in prefill phase, raise Error. - if self.is_prefill(): - raise ValueError( - "seq_group.get_last_latency() should not be called " - "if the seq_group is in prefill phase.") - - # Otherwise return token latency. - latency = now - self.metrics.last_token_time - self.metrics.last_token_time = now - return latency - - def maybe_set_first_token_time(self, time: float) -> None: - """Sets the first token time for Request level timings.""" - # Note: in a case where a sequence_group is swapped and - # recomputed, the time between iterations is counted - # in TPOT, rather than recalculating TTFT (since from the ) - # POV of the user, there is simply a long generation delay. - if (self.metrics.first_token_time is None - and self.seqs[0].get_output_len() == 1): - self.metrics.first_token_time = time - - def maybe_set_first_scheduled_time(self, time: float) -> None: - """Sets the first scheduled time and time in queue for Request - level timings.""" - if self.metrics.first_scheduled_time is None: - self.metrics.first_scheduled_time = time - self.metrics.time_in_queue = time - self.metrics.arrival_time - - def set_finished_time(self, time: Optional[float]) -> None: - """Sets the finished time for Request level timings.""" - self.metrics.finished_time = time - - def get_max_num_running_seqs(self) -> int: - """The maximum number of sequences running in parallel in the remaining - lifetime of the request.""" - if self.sampling_params: - n = self.sampling_params.n - assert isinstance(n, int) - if n > self.num_seqs(): - # At prompt stage, the sequence group is not yet filled up - # and only have one sequence running. However, in the - # generation stage, we will have `n` sequences - # running. - return n - # At sampling stages, return the number of actual sequences - # that are not finished yet. - return self.num_unfinished_seqs() - - def get_seqs( - self, - status: Optional[SequenceStatus] = None, - ) -> List[Sequence]: - if status is None: - return self.seqs - - if self.is_single_seq: - return self.seqs if self.seqs[0].status == status else [] - - return [seq for seq in self.seqs if seq.status == status] - - def is_encoder_decoder(self) -> bool: - return self.encoder_seq is not None - - def get_encoder_seq(self) -> Optional[Sequence]: - return self.encoder_seq - - def get_unfinished_seqs(self) -> List[Sequence]: - if self.is_single_seq: - return self.seqs if not self.seqs[0].is_finished() else [] - - return [seq for seq in self.seqs if not seq.is_finished()] - - def get_finished_seqs(self) -> List[Sequence]: - if self.is_single_seq: - return self.seqs if self.seqs[0].is_finished() else [] - - return [seq for seq in self.seqs if seq.is_finished()] - - def update_num_computed_tokens(self, num_new_computed_tokens: int): - """Update number of tokens computed so far.""" - for seq in self.seqs: - if not seq.is_finished(): - seq.data.update_num_computed_tokens(num_new_computed_tokens) - - def get_num_uncomputed_tokens(self) -> int: - num_uncomputed_tokens = 0 - for seq in self.seqs: - if not seq.is_finished(): - num_uncomputed_tokens += seq.data.get_num_uncomputed_tokens() - return num_uncomputed_tokens - - def num_seqs(self, status: Optional[SequenceStatus] = None) -> int: - # Optimization. We don't need to call get_seqs if we don't need to - # filter by states. - if status is None: - return len(self.seqs) - - if self.is_single_seq: - return 1 if self.seqs[0].status == status else 0 - - return len(self.get_seqs(status)) - - def num_unfinished_seqs(self) -> int: - if self.is_single_seq: - return 1 if not self.seqs[0].is_finished() else 0 - - return len(self.get_unfinished_seqs()) - - def num_finished_seqs(self) -> int: - if self.is_single_seq: - return 1 if self.seqs[0].is_finished() else 0 - - return len(self.get_finished_seqs()) - - def find(self, seq_id: int) -> Sequence: - if seq_id not in self.seqs_dict: - raise ValueError(f"Sequence {seq_id} not found.") - return self.seqs_dict[seq_id] - - def add(self, seq: Sequence) -> None: - if seq.seq_id in self.seqs_dict: - raise ValueError(f"Sequence {seq.seq_id} already exists.") - self.seqs_dict[seq.seq_id] = seq - self.seqs.append(seq) - self.is_single_seq = len(self.seqs) == 1 - - def remove(self, seq_id: int) -> None: - seq = self.seqs_dict.pop(seq_id, None) - if seq is None: - raise ValueError(f"Sequence {seq_id} not found.") - self.seqs.remove(seq) - self.is_single_seq = len(self.seqs) == 1 - - def is_finished(self) -> bool: - if self.is_single_seq: - return self.seqs[0].is_finished() - - return all(seq.is_finished() for seq in self.seqs) - - def is_prefill(self) -> bool: - # Every sequence should be in the same stage. - return self.seqs[0].is_prefill() - - def __repr__(self) -> str: - return (f"SequenceGroup(request_id={self.request_id}, " - f"sampling_params={self.sampling_params}, " - f"num_seqs={len(self.seqs)})") - - +"""Sequence and its related classes.""" +import copy +import enum +from abc import ABC, abstractmethod +from array import array +from collections import defaultdict +from dataclasses import dataclass +from functools import cached_property, reduce +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional +from typing import Sequence as GenericSequence +from typing import Set, Tuple, Union, cast + +import msgspec +import torch + +from vllm.inputs import EncoderDecoderLLMInputs, LLMInputs +from vllm.inputs.parse import is_valid_encoder_decoder_llm_inputs +from vllm.lora.request import LoRARequest +from vllm.pooling_params import PoolingParams +from vllm.prompt_adapter.request import PromptAdapterRequest +from vllm.sampling_params import SamplingParams +from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics + +if TYPE_CHECKING: + from vllm.multimodal.base import MultiModalDataDict + +VLLM_TOKEN_ID_ARRAY_TYPE = "l" + +VLLM_INVALID_TOKEN_ID = -1 + + +# We use dataclass for now because it is used for +# openai server output, and msgspec is not serializable. +# TODO(sang): Fix it. +@dataclass +class Logprob: + """Infos for supporting OpenAI compatible logprobs and token ranks. + + Attributes: + logprob: The logprob of chosen token + rank: The vocab rank of chosen token (>=1) + decoded_token: The decoded chosen token index + """ + logprob: float + rank: Optional[int] = None + decoded_token: Optional[str] = None + + +# {token_id -> logprob} per each sequence group. None if the corresponding +# sequence group doesn't require prompt logprob. +PromptLogprobs = List[Optional[Dict[int, Logprob]]] +# {token_id -> logprob} for each sequence group. +SampleLogprobs = List[Dict[int, Logprob]] + + +class SequenceStatus(enum.IntEnum): + """Status of a sequence.""" + WAITING = 0 + RUNNING = 1 + SWAPPED = 2 + # Note: anything after SWAPPED (2) will be considered + # as a finished status. + FINISHED_STOPPED = 3 + FINISHED_LENGTH_CAPPED = 4 + FINISHED_ABORTED = 5 + FINISHED_IGNORED = 6 + + @staticmethod + def is_finished(status: "SequenceStatus") -> bool: + return status > SequenceStatus.SWAPPED + + @staticmethod + def get_finished_reason(status: "SequenceStatus") -> Union[str, None]: + if status == SequenceStatus.FINISHED_STOPPED: + finish_reason = "stop" + elif status == SequenceStatus.FINISHED_LENGTH_CAPPED: + finish_reason = "length" + elif status == SequenceStatus.FINISHED_ABORTED: + finish_reason = "abort" + elif status == SequenceStatus.FINISHED_IGNORED: + # The ignored sequences are the sequences whose prompt lengths + # are longer than the model's length cap. Therefore, the stop + # reason should also be "length" as in OpenAI API. + finish_reason = "length" + else: + finish_reason = None + return finish_reason + + +class SequenceStage(enum.Enum): + PREFILL = enum.auto() + DECODE = enum.auto() + + +@dataclass +class RequestMetrics: + """Metrics associated with a request. + + Attributes: + arrival_time: The time when the request arrived. + first_scheduled_time: The time when the request was first scheduled. + first_token_time: The time when the first token was generated. + time_in_queue: The time the request spent in the queue. + finished_time: The time when the request was finished. + scheduler_time: The time spent in the scheduler when this request was + being considered by the scheduler. + model_forward_time: The time spent in the model forward pass when this + request was in the batch. + model_execute_time: The time spent in the model execute function. This + will include model forward, block/sync across + workers, cpu-gpu sync time and sampling time. + """ + arrival_time: float + last_token_time: float + first_scheduled_time: Optional[float] + first_token_time: Optional[float] + time_in_queue: Optional[float] + finished_time: Optional[float] = None + scheduler_time: Optional[float] = None + model_forward_time: Optional[float] = None + model_execute_time: Optional[float] = None + num_cached_tokens: Optional[int] = None + + +class SequenceDataDelta( + msgspec.Struct, + array_like=True, # type: ignore[call-arg] + omit_defaults=True): # type: ignore[call-arg] + """Delta SequenceData to send to workers per step.""" + # A new token to be appended to existing SequenceData. + new_output_token_ids: List[int] + # Overwriting existing `cumulative_logprob` + new_cumulative_logprob: float + # Overwriting existing `num_computed_tokens`. + new_num_computed_tokens: int + # Overwriting existing `stage`. + new_stage: SequenceStage + + +class SequenceData(msgspec.Struct, + omit_defaults=True): # type: ignore[call-arg] + """Data associated with a sequence. + + Args: + prompt_token_ids: The token IDs of the prompt. + output_token_ids: The token IDs of the output. Set to an empty list if + None. + + Attributes: + prompt_token_ids: The token IDs of the prompt. + output_token_ids: The token IDs of the output. + cumulative_logprob: The cumulative log probability of the output. + """ + # NOTE: we cannot use Union[List, array] because msgspec cannot support + # union of 2 list types. + _prompt_token_ids: array + _output_token_ids: array = msgspec.field( + default_factory=lambda: array(VLLM_TOKEN_ID_ARRAY_TYPE, [])) + + ### The below fields should not be passed as an argument ### + _cumulative_logprob: float = 0.0 + _prompt_token_ids_tuple: Tuple[int, + ...] = msgspec.field(default_factory=tuple) + # The number of tokens that are computed (that run against the model). + _num_computed_tokens: int = 0 + _stage: SequenceStage = SequenceStage.PREFILL + _cached_all_token_ids: List[int] = msgspec.field(default_factory=list) + + # It is used to get delta input. It is reset when `get_delta_and_reset` + # is called. + _new_appended_tokens: List[int] = msgspec.field(default_factory=list) + + # It is used to compute mrope_position_ids. + _mrope_position_delta: Optional[int] = None + + @staticmethod + def from_token_counts(*token_counts: Tuple[int, int]) -> "SequenceData": + if len(token_counts) == 0: + return SequenceData.from_seqs([]) + + arrs = [ + array(VLLM_TOKEN_ID_ARRAY_TYPE, [token_id]) * count + for token_id, count in token_counts + ] + + return SequenceData(reduce(array.__add__, arrs)) + + @staticmethod + def from_seqs( + prompt_token_ids: GenericSequence[int], + output_token_ids: Optional[GenericSequence[int]] = None, + ) -> "SequenceData": + prompt_token_ids_arr = array(VLLM_TOKEN_ID_ARRAY_TYPE, + prompt_token_ids) + + if output_token_ids is None: + return SequenceData(prompt_token_ids_arr) + + output_token_ids_arr = array(VLLM_TOKEN_ID_ARRAY_TYPE, + output_token_ids) + + return SequenceData(prompt_token_ids_arr, + _output_token_ids=output_token_ids_arr) + + def __post_init__(self) -> None: + assert self._prompt_token_ids.typecode == "l" + assert self._output_token_ids.typecode == "l" + self._prompt_token_ids_tuple: Tuple[int, ...] = tuple( + self._prompt_token_ids) + self._update_cached_all_tokens() + + def _update_cached_all_tokens(self): + assert isinstance(self._prompt_token_ids, array) + assert isinstance(self._output_token_ids, array) + self._cached_all_token_ids: List[int] = list(self._prompt_token_ids + + self._output_token_ids) + + @property + def cumulative_logprob(self) -> float: + return self._cumulative_logprob + + @property + def prompt_token_ids(self) -> Tuple[int, ...]: + return self._prompt_token_ids_tuple + + @prompt_token_ids.setter + def prompt_token_ids(self, new_prompt_token_ids) -> None: + raise NotImplementedError + + @property + def prompt_token_ids_array(self) -> array: + """Return the prompt token ids in array type. + + Note that the array is in "I" type, and it is not compatible + with torch.long (2 bytes vs 4 bytes). So beware of the usage. + """ + return self._prompt_token_ids + + @property + def output_token_ids(self) -> Tuple[int, ...]: + return tuple(self._output_token_ids) + + @output_token_ids.setter + def output_token_ids(self, new_output_token_ids: List[int]) -> None: + self._output_token_ids = array(VLLM_TOKEN_ID_ARRAY_TYPE, + new_output_token_ids) + self._update_cached_all_tokens() + + @property + def output_token_ids_array(self) -> array: + """Return the prompt token ids in array type. + + Note that the array is in "I" type, and it is not compatible + with torch.long (2 bytes vs 4 bytes). So beware of the usage. + """ + assert isinstance(self._output_token_ids, array) + return self._output_token_ids + + @property + def mrope_position_delta(self) -> Optional[int]: + return self._mrope_position_delta + + @mrope_position_delta.setter + def mrope_position_delta(self, new_mrope_position_delta): + self._mrope_position_delta = new_mrope_position_delta + + def append_token_id(self, token_id: int, logprob: float) -> None: + self._output_token_ids.append(token_id) + self._new_appended_tokens.append(token_id) + self._cached_all_token_ids.append(token_id) + self._cumulative_logprob += logprob + + def get_len(self) -> int: + return len(self._output_token_ids) + len(self._prompt_token_ids) + + def get_prompt_len(self) -> int: + return len(self._prompt_token_ids) + + def get_output_len(self) -> int: + return len(self._output_token_ids) + + def get_token_ids(self) -> List[int]: + return self._cached_all_token_ids + + def get_prefix_token_ids( + self, num_tokens: int + ) -> Tuple[Tuple[int, ...], Optional[Tuple[int, ...]]]: + """Get prefix tokens, and make the return value hashable""" + prompt_length = self.get_prompt_len() + if num_tokens > prompt_length: + return (self._prompt_token_ids_tuple, + tuple(self._output_token_ids[:num_tokens - prompt_length])) + else: + return (self._prompt_token_ids_tuple[:num_tokens], None) + + def get_num_computed_tokens(self) -> int: + """Return the number of prefill tokens that are already computed.""" + return self._num_computed_tokens + + def update_num_computed_tokens(self, num_new_computed_tokens: int): + """Update number of tokens computed so far.""" + self._num_computed_tokens += num_new_computed_tokens + assert self._num_computed_tokens <= self.get_len(), ( + self._num_computed_tokens, self.get_len()) + # If all tokens are computed, it means it is in decoding phase. + if self.get_num_uncomputed_tokens() == 0: + self._stage = SequenceStage.DECODE + + def reset_state_for_recompute(self) -> None: + """Reset the number of computed tokens from this sequence. It is + supposed to be called when a sequence needs to be started from + the beginning again (e.g., sequence is preempted). + """ + self._num_computed_tokens = 0 + self._stage = SequenceStage.PREFILL + self._new_appended_tokens = [] + + def get_num_uncomputed_tokens(self) -> int: + """Return the number of prefill tokens that are not computed.""" + # we use `get_len()` which includes prompt_len + output_len instead + # of prompt_len here. This is because during recompute we need to + # prefill for both prompt and output. + return self.get_len() - self.get_num_computed_tokens() + + def get_last_token_id(self) -> int: + if not self._output_token_ids: + return self._prompt_token_ids[-1] + return self._output_token_ids[-1] + + def get_prompt_token_ids(self) -> Tuple[int, ...]: + return self.prompt_token_ids + + def get_output_token_ids(self) -> Tuple[int, ...]: + return self.output_token_ids + + def get_delta_and_reset(self) -> SequenceDataDelta: + delta = SequenceDataDelta(self._new_appended_tokens, + self._cumulative_logprob, + self.get_num_computed_tokens(), self.stage) + # Reset delta state. + self._new_appended_tokens = [] + return delta + + def apply_delta(self, delta: SequenceDataDelta): + self._num_computed_tokens = delta.new_num_computed_tokens + self._cumulative_logprob = delta.new_cumulative_logprob + self._stage = delta.new_stage + self._output_token_ids.extend(delta.new_output_token_ids) + self._cached_all_token_ids.extend(delta.new_output_token_ids) + + @property + def stage(self) -> SequenceStage: + return self._stage + + def __repr__(self) -> str: + return (f"SequenceData(" + f"prompt_token_ids={self._prompt_token_ids}, " + f"output_token_ids={self.output_token_ids}, " + f"cumulative_logprob={self.cumulative_logprob}, " + f"get_num_computed_tokens={self.get_num_computed_tokens()}") + + +class Sequence: + """Stores the data, status, and block information of a sequence. + + The sequence is constructed from the LLMInputs instance passed + in through the `inputs` constructor argument. + + For encoder/decoder models, LLMInputs encapsulates both a + decoder and encoder prompt, creating an ambiguity about which + prompt to construct the sequence from. The `from_decoder_prompt` + constructor argument signals whether to construct the Sequence + from the LLMInputs decoder prompt, or encoder prompt. + + Args: + seq_id: The ID of the sequence. + inputs: The inputs of the sequence. + block_size: The block size of the sequence. Should be the same as the + block size used by the block manager and cache engine. + eos_token_id: The end-of-sequence (EOS) token id recognized by this LLM. + lora_request: LoRA request. + prompt_adapter_request: Prompt Adapter request. + from_decoder_prompt: Construct Sequence from LLMInputs decoder prompt + (True) or encoder prompt (False.) Must be True + for decoder-only model. + + """ + + def __init__( + self, + seq_id: int, + inputs: "LLMInputs", + block_size: int, + eos_token_id: Optional[int] = None, + lora_request: Optional[LoRARequest] = None, + prompt_adapter_request: Optional[PromptAdapterRequest] = None, + from_decoder_prompt: bool = True, + ) -> None: + self.seq_id = seq_id + self.inputs = inputs + self.block_size = block_size + self.eos_token_id = eos_token_id + self.lora_request = lora_request + self.prompt_adapter_request = prompt_adapter_request + self.from_decoder_prompt = from_decoder_prompt + + # For decoder-only models, a Sequence is constructed + # from an LLMInputs instance (the `inputs` arg.) + # + # For encoder/decoder models the same `inputs` + # instance could be utilized to construct either an + # encoder sequence or a decoder sequence, because + # `LLMInputs` has both decoder- and encoder-oriented + # member variables (i.e. it encapsulates both an encoder + # and a decoder prompt.) The decision of which type of sequence + # to generate is determined by the `from_decoder_prompt` argument. + # + # When constructing a encoder sequence + # (`from_decoder_prompt` False) it matters that + # the `LLMInputs` instance stored in `inputs` is valid + # in the sense that its encoder-related member variables are + # populated; below, an exception is raised if this is + # not the case. + # + # When constructing a decoder sequence (`from_decoder_prompt` True) + # it does not matter whether `inputs` has its encoder-related + # member variables populated. + if not (from_decoder_prompt + or is_valid_encoder_decoder_llm_inputs(inputs)): + raise ValueError("Cannot extract encoder input prompt from " + f"invalid input {inputs}; did you forget the " + "encoder input prompt fields?") + + self.data = SequenceData.from_seqs(self.prompt_token_ids) + self.output_logprobs: SampleLogprobs = [] + self.output_text = "" + + self.status = SequenceStatus.WAITING + self.stop_reason: Union[int, str, None] = None + + # These are used to keep track of delta outputs + self._last_output_token_ids_offset: int = 0 + self._last_output_text_offset: int = 0 + + # Used for incremental detokenization + self.prefix_offset = 0 + self.read_offset = 0 + # Input + output tokens + self.tokens: Optional[List[str]] = None + + @property + def n_blocks(self) -> int: + return (self.get_len() + self.block_size - 1) // self.block_size + + @cached_property + def prompt(self) -> Optional[str]: + # Select decoder or encoder input prompt str, as appropriate + prompt_key: str = ("prompt" + if self.from_decoder_prompt else "encoder_prompt") + + return cast(Optional[str], self.inputs.get(prompt_key)) + + @cached_property + def prompt_token_ids(self) -> List[int]: + # Select decoder or encoder input prompt token ids, as appropriate + prompt_token_ids_key: str = ("prompt_token_ids" + if self.from_decoder_prompt else + "encoder_prompt_token_ids") + + # Cache computed prompt token ids + return cast(List[int], self.inputs.get(prompt_token_ids_key)) + + @property + def multi_modal_data(self) -> "MultiModalDataDict": + if self.inputs.get("multi_modal_data") and self.inputs.get( + "encoder_multi_modal_data"): + raise ValueError( + "Multi-modal data in both encoder and decoder is not supported." + ) + inputs = self.inputs + return self.inputs.get("multi_modal_data") or (cast( + EncoderDecoderLLMInputs, + inputs).get("encoder_multi_modal_data")) or {} + + @property + def mm_processor_kwargs(self) -> Dict[str, Any]: + return self.inputs.get("mm_processor_kwargs") or {} + + @property + def lora_int_id(self) -> int: + return self.lora_request.lora_int_id if self.lora_request else 0 + + @property + def prompt_adapter_id(self) -> int: + return self.prompt_adapter_request.prompt_adapter_id \ + if self.prompt_adapter_request else 0 + + def get_output_text_to_return(self, buffer_length: int, + delta: bool) -> str: + """If delta is True, only new text since the last call to + this method is returned""" + + # We return the full output text if the sequence is finished. + truncate = buffer_length and not self.is_finished() + if not delta: + return self.output_text[:-buffer_length] if truncate else ( + self.output_text) + length = len(self.output_text) + if truncate: + length -= buffer_length + last_offset = self._last_output_text_offset + if last_offset < length: + self._last_output_text_offset = length + return self.output_text[last_offset:length] + return "" + + def get_output_token_ids_to_return( + self, delta: bool) -> Union[GenericSequence[int], int]: + """If delta is True, only new tokens since the last call to + this method are returned""" + if not delta: + return self.get_output_token_ids() + + output_len = self.get_output_len() + + # Get the number of new tokens + num_new_tokens = output_len - self._last_output_token_ids_offset + self._last_output_token_ids_offset = output_len + + # Return new tokens + if num_new_tokens == 0: + # During chunked prefill steps with no output yet, num_new_tokens=0. + # Python's [-0:] == [0:] returns the ENTIRE list — guard against this. + return [] + + if num_new_tokens == 1: + # Optimization for single decode token case + # (which is what we have most of the time) + return self.data._cached_all_token_ids[-1] + + return self.data._cached_all_token_ids[-num_new_tokens:] + + def hash_of_block(self, logical_idx: int) -> int: + # TODO This can produce incorrect hash when block size > prompt size + + # Compute the number of tokens in the sequence + # TODO: The current hashing function is O(L^2). We should optimize + # this in the future. + num_tokens = self.num_hashed_tokens_of_block(logical_idx) + hashed_tokens = self.data.get_prefix_token_ids(num_tokens) + return hash((hashed_tokens, self.lora_int_id)) + + def num_hashed_tokens_of_block(self, logical_idx: int): + return logical_idx * self.block_size + self.block_size + + def reset_state_for_recompute(self): + """Reset the sequence states for recomputation.""" + self.data.reset_state_for_recompute() + + def append_token_id(self, token_id: int, logprobs: Dict[int, + Logprob]) -> None: + assert token_id in logprobs + self.output_logprobs.append(logprobs) + self.data.append_token_id(token_id, logprobs[token_id].logprob) + + def get_len(self) -> int: + return self.data.get_len() + + def get_prompt_len(self) -> int: + return self.data.get_prompt_len() + + def get_output_len(self) -> int: + return self.data.get_output_len() + + def get_token_ids(self) -> List[int]: + return self.data.get_token_ids() + + def get_prompt_token_ids(self) -> Tuple[int, ...]: + return self.data.get_prompt_token_ids() + + def get_last_token_id(self) -> int: + return self.data.get_last_token_id() + + def get_output_token_ids(self) -> Tuple[int, ...]: + return self.data.get_output_token_ids() + + def get_cumulative_logprob(self) -> float: + return self.data.cumulative_logprob + + def is_finished(self) -> bool: + return SequenceStatus.is_finished(self.status) + + def fork(self, new_seq_id: int) -> "Sequence": + new_seq = copy.deepcopy(self) + new_seq.seq_id = new_seq_id + return new_seq + + def get_num_new_tokens(self) -> int: + """Get the number of new tokens to be computed. + + Returns: + The new number of tokens to be computed. I.e., 1 for decode, or + the remaining prompt size for prefill. + """ + if self.data.stage == SequenceStage.DECODE: + return 1 + return self.data.get_num_uncomputed_tokens() + + def is_prefill(self) -> bool: + return self.data.stage == SequenceStage.PREFILL + + def __repr__(self) -> str: + return (f"Sequence(seq_id={self.seq_id}, " + f"status={self.status.name}, " + f"num_blocks={self.n_blocks}, ") + + +class SequenceGroupState(msgspec.Struct, + omit_defaults=True): # type: ignore[call-arg] + """Mutable state tied to a specific sequence group""" + + # for multi-step decoding + num_steps: int = 1 + current_step: int = 0 + + @property + def remaining_steps(self) -> int: + return self.num_steps - self.current_step + + +class SequenceGroup: + """A group of sequences that are generated from the same prompt. + + Args: + request_id: The ID of the request. + seqs: The list of sequences. + sampling_params: The sampling parameters used to generate the outputs. + arrival_time: The arrival time of the request. + lora_request: LoRA request. + embeddings: The embeddings vectors of the prompt of the sequence group + for an embedding model. + pooling_params: The pooling parameters used to generate the pooling + for an embedding model. + encoder_seq: Optional, the single encoder sequence. Should be None + unless you are working with an encoder/decoder model. + trace_headers: OpenTelemetry trace headers. + prompt_adapter_request: Prompt Adapter request. + priority: User-defined priority of the request. + """ + + def __init__( + self, + request_id: str, + seqs: List[Sequence], + arrival_time: float, + sampling_params: Optional[SamplingParams] = None, + lora_request: Optional[LoRARequest] = None, + embeddings: Optional[List[float]] = None, + pooling_params: Optional[PoolingParams] = None, + encoder_seq: Optional[Sequence] = None, + trace_headers: Optional[Mapping[str, str]] = None, + prompt_adapter_request: Optional[PromptAdapterRequest] = None, + priority: int = 0, + ) -> None: + self.request_id = request_id + self.seqs = seqs + self.arrival_time = arrival_time + self.is_single_seq = len(seqs) == 1 + self.seqs_dict = {seq.seq_id: seq for seq in seqs} + + self.sampling_params = sampling_params + self.metrics = RequestMetrics(arrival_time=arrival_time, + last_token_time=arrival_time, + first_scheduled_time=None, + first_token_time=None, + time_in_queue=None) + self.lora_request = lora_request + self.prompt_logprobs: Optional[PromptLogprobs] = None + self.state = SequenceGroupState() + self.embeddings = embeddings + self.pooling_params = pooling_params + self.prompt_adapter_request = prompt_adapter_request + self.encoder_seq = encoder_seq + self.trace_headers = trace_headers + self.priority = priority + + self.cached_request_output = None + + @property + def prompt(self) -> Optional[str]: + # All sequences in the group should have the same prompt. + # We use the prompt of an arbitrary sequence. + return self.seqs[0].prompt + + @property + def prompt_token_ids(self) -> List[int]: + # All sequences in the group should have the same prompt. + # We use the prompt of an arbitrary sequence. + return self.seqs[0].prompt_token_ids + + @property + def encoder_prompt(self) -> Optional[str]: + # There are either 0 or 1 encoder sequences + # If one is present, its prompt is distinct + # from the decoder's. + return (self.encoder_seq.prompt + if self.encoder_seq is not None else None) + + @property + def encoder_prompt_token_ids(self) -> Optional[List[int]]: + # There are either 0 or 1 encoder sequences + # If one is present, its prompt token ids are + # distinct from the decoder's. + return (self.encoder_seq.prompt_token_ids + if self.encoder_seq is not None else None) + + @property + def multi_modal_data(self) -> "MultiModalDataDict": + # All sequences in the group should have the same multi-modal data. + # We use the multi-modal data of an arbitrary sequence. + return self.seqs[0].multi_modal_data + + @property + def mm_processor_kwargs(self) -> Dict[str, Any]: + # As with multi-modal data, all sequences in the group should have the + # same processor kwargs (i.e., mm_processor_kwargs are optionally + # provided per request; note that are independent of whether the model + # decoder-only or an encoder-decoder). + return self.seqs[0].mm_processor_kwargs + + @property + def lora_int_id(self) -> int: + return self.lora_request.lora_int_id if self.lora_request else 0 + + @property + def prompt_adapter_id(self) -> int: + return self.prompt_adapter_request.prompt_adapter_id \ + if self.prompt_adapter_request else 0 + + @property + def prompt_adapter_num_virtual_tokens(self) -> int: + return self.prompt_adapter_request.prompt_adapter_num_virtual_tokens\ + if self.prompt_adapter_request else 0 + + def init_multi_step(self, num_steps: int) -> None: + self.state.num_steps = num_steps + self.state.current_step = 0 + + def init_multi_step_from_lookahead_slots(self, num_lookahead_slots: int, + num_scheduler_steps: int, + is_multi_step: bool, + enable_chunking: bool) -> None: + + if not is_multi_step: + self.init_multi_step(num_steps=num_scheduler_steps) + return + + # Multi-Step case + is_prefill = self.is_prefill() + + # The asserts below reflect the expectations of the current system. + if is_prefill and enable_chunking: + assert num_lookahead_slots == num_scheduler_steps + self.init_multi_step(num_steps=num_lookahead_slots) + else: + is_decode: bool = not is_prefill + # If it is a prefill, num_lookahead_slots must be 0 + assert num_lookahead_slots == 0 or is_decode + # If it is a decode, num_lookahead_slots + 1 must match + # the scheduler steps. + assert num_lookahead_slots + 1 == num_scheduler_steps or is_prefill + self.init_multi_step(num_steps=num_lookahead_slots + 1) + + def get_last_latency(self, now: float) -> Optional[float]: + """Sets the last token time for Request level timings.""" + # If still in prefill phase, raise Error. + if self.is_prefill(): + raise ValueError( + "seq_group.get_last_latency() should not be called " + "if the seq_group is in prefill phase.") + + # Otherwise return token latency. + latency = now - self.metrics.last_token_time + self.metrics.last_token_time = now + return latency + + def maybe_set_first_token_time(self, time: float) -> None: + """Sets the first token time for Request level timings.""" + # Note: in a case where a sequence_group is swapped and + # recomputed, the time between iterations is counted + # in TPOT, rather than recalculating TTFT (since from the ) + # POV of the user, there is simply a long generation delay. + if (self.metrics.first_token_time is None + and self.seqs[0].get_output_len() == 1): + self.metrics.first_token_time = time + + def maybe_set_first_scheduled_time(self, time: float) -> None: + """Sets the first scheduled time and time in queue for Request + level timings.""" + if self.metrics.first_scheduled_time is None: + self.metrics.first_scheduled_time = time + self.metrics.time_in_queue = time - self.metrics.arrival_time + + def set_finished_time(self, time: Optional[float]) -> None: + """Sets the finished time for Request level timings.""" + self.metrics.finished_time = time + + def get_max_num_running_seqs(self) -> int: + """The maximum number of sequences running in parallel in the remaining + lifetime of the request.""" + if self.sampling_params: + n = self.sampling_params.n + assert isinstance(n, int) + if n > self.num_seqs(): + # At prompt stage, the sequence group is not yet filled up + # and only have one sequence running. However, in the + # generation stage, we will have `n` sequences + # running. + return n + # At sampling stages, return the number of actual sequences + # that are not finished yet. + return self.num_unfinished_seqs() + + def get_seqs( + self, + status: Optional[SequenceStatus] = None, + ) -> List[Sequence]: + if status is None: + return self.seqs + + if self.is_single_seq: + return self.seqs if self.seqs[0].status == status else [] + + return [seq for seq in self.seqs if seq.status == status] + + def is_encoder_decoder(self) -> bool: + return self.encoder_seq is not None + + def get_encoder_seq(self) -> Optional[Sequence]: + return self.encoder_seq + + def get_unfinished_seqs(self) -> List[Sequence]: + if self.is_single_seq: + return self.seqs if not self.seqs[0].is_finished() else [] + + return [seq for seq in self.seqs if not seq.is_finished()] + + def get_finished_seqs(self) -> List[Sequence]: + if self.is_single_seq: + return self.seqs if self.seqs[0].is_finished() else [] + + return [seq for seq in self.seqs if seq.is_finished()] + + def update_num_computed_tokens(self, num_new_computed_tokens: int): + """Update number of tokens computed so far.""" + for seq in self.seqs: + if not seq.is_finished(): + seq.data.update_num_computed_tokens(num_new_computed_tokens) + + def get_num_uncomputed_tokens(self) -> int: + num_uncomputed_tokens = 0 + for seq in self.seqs: + if not seq.is_finished(): + num_uncomputed_tokens += seq.data.get_num_uncomputed_tokens() + return num_uncomputed_tokens + + def num_seqs(self, status: Optional[SequenceStatus] = None) -> int: + # Optimization. We don't need to call get_seqs if we don't need to + # filter by states. + if status is None: + return len(self.seqs) + + if self.is_single_seq: + return 1 if self.seqs[0].status == status else 0 + + return len(self.get_seqs(status)) + + def num_unfinished_seqs(self) -> int: + if self.is_single_seq: + return 1 if not self.seqs[0].is_finished() else 0 + + return len(self.get_unfinished_seqs()) + + def num_finished_seqs(self) -> int: + if self.is_single_seq: + return 1 if self.seqs[0].is_finished() else 0 + + return len(self.get_finished_seqs()) + + def find(self, seq_id: int) -> Sequence: + if seq_id not in self.seqs_dict: + raise ValueError(f"Sequence {seq_id} not found.") + return self.seqs_dict[seq_id] + + def add(self, seq: Sequence) -> None: + if seq.seq_id in self.seqs_dict: + raise ValueError(f"Sequence {seq.seq_id} already exists.") + self.seqs_dict[seq.seq_id] = seq + self.seqs.append(seq) + self.is_single_seq = len(self.seqs) == 1 + + def remove(self, seq_id: int) -> None: + seq = self.seqs_dict.pop(seq_id, None) + if seq is None: + raise ValueError(f"Sequence {seq_id} not found.") + self.seqs.remove(seq) + self.is_single_seq = len(self.seqs) == 1 + + def is_finished(self) -> bool: + if self.is_single_seq: + return self.seqs[0].is_finished() + + return all(seq.is_finished() for seq in self.seqs) + + def is_prefill(self) -> bool: + # Every sequence should be in the same stage. + return self.seqs[0].is_prefill() + + def __repr__(self) -> str: + return (f"SequenceGroup(request_id={self.request_id}, " + f"sampling_params={self.sampling_params}, " + f"num_seqs={len(self.seqs)})") + + class SequenceGroupMetadataDelta( - msgspec.Struct, - tag=True, # type: ignore[call-arg] - array_like=True, # type: ignore[call-arg] - omit_defaults=True): # type: ignore[call-arg] - """Delta of SequenceGroupMetadata. - - After sending the first SequenceGroupMetadata, vLLM scheduler - only sends delta to reduce the data payload size. - """ - seq_data_delta: Dict[int, SequenceDataDelta] - request_id: str - block_tables: Dict[int, List[int]] - is_prompt: bool - do_sample: bool = True - token_chunk_size: Optional[int] = None + msgspec.Struct, + tag=True, # type: ignore[call-arg] + array_like=True, # type: ignore[call-arg] + omit_defaults=True): # type: ignore[call-arg] + """Delta of SequenceGroupMetadata. + + After sending the first SequenceGroupMetadata, vLLM scheduler + only sends delta to reduce the data payload size. + """ + seq_data_delta: Dict[int, SequenceDataDelta] + request_id: str + block_tables: Dict[int, List[int]] + is_prompt: bool + do_sample: bool = True + token_chunk_size: Optional[int] = None computed_block_nums: Optional[List[int]] = None state: Optional[SequenceGroupState] = msgspec.field( default_factory=lambda: SequenceGroupState()) @@ -947,70 +947,70 @@ class SequenceGroupMetadataDelta( gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None gdn_segment_offsets: Optional[List[int]] = None - - + + class SequenceGroupMetadata( - msgspec.Struct, - tag=True, # type: ignore[call-arg] - array_like=True, # type: ignore[call-arg] - omit_defaults=True): # type: ignore[call-arg] - """Metadata for a sequence group. Used to create `AttentionMetadata`. - - Args: - request_id: The ID of the request. - is_prompt: Whether the request is at prompt stage. - seq_data: The sequence data. (Seq id -> sequence data) - sampling_params: The sampling parameters used to generate the outputs. - block_tables: The block tables. (Seq id -> list of physical block - numbers) - do_sample: True if sampling is required. Sampling is not required when - e.g., prefill is chunked, and the current iteration only computes - query tokens for prefill, we don't need sampling. - token_chunk_size: The number of tokens to be processed (per sequence). - None if chunking is not required. - lora_request: LoRA request. - computed_block_nums: The block numbers that are already computed, - used in prefix caching. - state: Internal state tied to this sequence group. - multi_modal_data: Multi modal data. - mm_processor_kwargs: Multimodal input processor / mapper overrides. - encoder_seq_data: Optional sequence data for encoder prompt - (SequenceGroup.encoder_seq). Should be None - unless you are working with an encoder/decoder - model. - cross_block_table: Optional cross-attention block table associated - with the encoder prompt - (SequenceGroup.encoder_seq). Should be None - unless you are working with an encoder/decoder - model. - prompt_adapter_request: Prompt Adapter request. - """ - - request_id: str - is_prompt: bool - seq_data: Dict[int, SequenceData] - sampling_params: Optional[SamplingParams] - block_tables: Dict[int, List[int]] - do_sample: bool = True - pooling_params: Optional[PoolingParams] = None - lora_request: Optional[LoRARequest] = None - computed_block_nums: Optional[List[int]] = None - state: Optional[SequenceGroupState] = msgspec.field( - default_factory=lambda: SequenceGroupState()) - # "MultiModalDataDict" types. We have to use Any due to msgspec - # doesn't allow to have union of 2 different dicts. - multi_modal_data: Optional[Any] = None - mm_processor_kwargs: Optional[Dict[str, Any]] = None - encoder_seq_data: Optional[SequenceData] = None - cross_block_table: Optional[List[int]] = None - prompt_adapter_request: Optional[PromptAdapterRequest] = None - token_chunk_size: Optional[int] = None - - ### Stateful fields that are lazily defined. ### - # The number of speculative tokens adopted in this request. - # None means specuative decoding is not used. - # Zero means speculative decoding is disabled for some reasons. - # TODO: We should maintain this states out of the sequence group. + msgspec.Struct, + tag=True, # type: ignore[call-arg] + array_like=True, # type: ignore[call-arg] + omit_defaults=True): # type: ignore[call-arg] + """Metadata for a sequence group. Used to create `AttentionMetadata`. + + Args: + request_id: The ID of the request. + is_prompt: Whether the request is at prompt stage. + seq_data: The sequence data. (Seq id -> sequence data) + sampling_params: The sampling parameters used to generate the outputs. + block_tables: The block tables. (Seq id -> list of physical block + numbers) + do_sample: True if sampling is required. Sampling is not required when + e.g., prefill is chunked, and the current iteration only computes + query tokens for prefill, we don't need sampling. + token_chunk_size: The number of tokens to be processed (per sequence). + None if chunking is not required. + lora_request: LoRA request. + computed_block_nums: The block numbers that are already computed, + used in prefix caching. + state: Internal state tied to this sequence group. + multi_modal_data: Multi modal data. + mm_processor_kwargs: Multimodal input processor / mapper overrides. + encoder_seq_data: Optional sequence data for encoder prompt + (SequenceGroup.encoder_seq). Should be None + unless you are working with an encoder/decoder + model. + cross_block_table: Optional cross-attention block table associated + with the encoder prompt + (SequenceGroup.encoder_seq). Should be None + unless you are working with an encoder/decoder + model. + prompt_adapter_request: Prompt Adapter request. + """ + + request_id: str + is_prompt: bool + seq_data: Dict[int, SequenceData] + sampling_params: Optional[SamplingParams] + block_tables: Dict[int, List[int]] + do_sample: bool = True + pooling_params: Optional[PoolingParams] = None + lora_request: Optional[LoRARequest] = None + computed_block_nums: Optional[List[int]] = None + state: Optional[SequenceGroupState] = msgspec.field( + default_factory=lambda: SequenceGroupState()) + # "MultiModalDataDict" types. We have to use Any due to msgspec + # doesn't allow to have union of 2 different dicts. + multi_modal_data: Optional[Any] = None + mm_processor_kwargs: Optional[Dict[str, Any]] = None + encoder_seq_data: Optional[SequenceData] = None + cross_block_table: Optional[List[int]] = None + prompt_adapter_request: Optional[PromptAdapterRequest] = None + token_chunk_size: Optional[int] = None + + ### Stateful fields that are lazily defined. ### + # The number of speculative tokens adopted in this request. + # None means specuative decoding is not used. + # Zero means speculative decoding is disabled for some reasons. + # TODO: We should maintain this states out of the sequence group. num_speculative_tokens: Optional[int] = None # BI100 hybrid prefix-cache actions. These are internal scheduler-to-worker # metadata and never surface through the OpenAI API. @@ -1018,49 +1018,49 @@ class SequenceGroupMetadata( gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None gdn_segment_offsets: Optional[List[int]] = None - - def __post_init__(self): - if self.seq_data is not None and self.token_chunk_size is None: - if self.is_prompt: - self.token_chunk_size = next(iter( - self.seq_data.values())).get_len() - else: - self.token_chunk_size = 1 - - @property - def lora_int_id(self) -> int: - return self.lora_request.lora_int_id if self.lora_request else 0 - - @property - def prompt_adapter_id(self) -> int: - return self.prompt_adapter_request.prompt_adapter_id \ - if self.prompt_adapter_request else 0 - - @property - def prompt_adapter_num_virtual_tokens(self) -> int: - return self.prompt_adapter_request.prompt_adapter_num_virtual_tokens \ - if self.prompt_adapter_request else 0 - - # Multi-Step Chunked-Prefill property - @property - def is_single_step_prompt(self) -> bool: - # do_sample is true, only when the token_chunk_size matches the - # num_uncomputed_tokens of the sequence. This indicates that - # the prompt will finish processing in a single `execute_model` - # step. - return self.is_prompt and self.do_sample - - def get_first_seq_id(self) -> int: - # This is an efficient way of fetching the seq_id when - # we know this SequenceGroup has only one sequence. - return next(iter(self.seq_data)) - - def apply_delta(self, - sequence_group_metadata_delta: SequenceGroupMetadataDelta): - for id, delta in sequence_group_metadata_delta.seq_data_delta.items(): - self.seq_data[id].apply_delta(delta) - assert self.request_id == sequence_group_metadata_delta.request_id - self.block_tables = sequence_group_metadata_delta.block_tables + + def __post_init__(self): + if self.seq_data is not None and self.token_chunk_size is None: + if self.is_prompt: + self.token_chunk_size = next(iter( + self.seq_data.values())).get_len() + else: + self.token_chunk_size = 1 + + @property + def lora_int_id(self) -> int: + return self.lora_request.lora_int_id if self.lora_request else 0 + + @property + def prompt_adapter_id(self) -> int: + return self.prompt_adapter_request.prompt_adapter_id \ + if self.prompt_adapter_request else 0 + + @property + def prompt_adapter_num_virtual_tokens(self) -> int: + return self.prompt_adapter_request.prompt_adapter_num_virtual_tokens \ + if self.prompt_adapter_request else 0 + + # Multi-Step Chunked-Prefill property + @property + def is_single_step_prompt(self) -> bool: + # do_sample is true, only when the token_chunk_size matches the + # num_uncomputed_tokens of the sequence. This indicates that + # the prompt will finish processing in a single `execute_model` + # step. + return self.is_prompt and self.do_sample + + def get_first_seq_id(self) -> int: + # This is an efficient way of fetching the seq_id when + # we know this SequenceGroup has only one sequence. + return next(iter(self.seq_data)) + + def apply_delta(self, + sequence_group_metadata_delta: SequenceGroupMetadataDelta): + for id, delta in sequence_group_metadata_delta.seq_data_delta.items(): + self.seq_data[id].apply_delta(delta) + assert self.request_id == sequence_group_metadata_delta.request_id + self.block_tables = sequence_group_metadata_delta.block_tables self.token_chunk_size = sequence_group_metadata_delta.token_chunk_size self.do_sample = sequence_group_metadata_delta.do_sample self.is_prompt = sequence_group_metadata_delta.is_prompt @@ -1072,335 +1072,335 @@ class SequenceGroupMetadata( self.gdn_evict_keys = sequence_group_metadata_delta.gdn_evict_keys self.gdn_segment_offsets = ( sequence_group_metadata_delta.gdn_segment_offsets) - - def finish_step(self) -> None: - assert self.state is not None - assert self.state.current_step < self.state.num_steps, \ - f"current step {self.state.current_step}, num_steps {self.state.num_steps}" # noqa - self.state.current_step += 1 - - -class SequenceOutput( - msgspec.Struct, - omit_defaults=True, # type: ignore[call-arg] - array_like=True): # type: ignore[call-arg] - """The model output associated with a sequence. - - Args: - parent_seq_id: The ID of the parent sequence (for forking in beam - search). - output_token: The output token ID. - logprobs: The logprobs of the output token. - (Token id -> logP(x_i+1 | x_0, ..., x_i)) - """ - parent_seq_id: int - output_token: int - logprobs: Dict[int, Logprob] - - def __repr__(self) -> str: - return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, " - f"output_token={self.output_token}, " - f"logprobs={self.logprobs})") - - def __eq__(self, other: object) -> bool: - if not isinstance(other, SequenceOutput): - raise NotImplementedError() - equal = (self.parent_seq_id == other.parent_seq_id - and self.output_token == other.output_token) - log_probs_equal = other.logprobs == self.logprobs - return equal and log_probs_equal - - -class SequenceGroupOutput(ABC): - """The base class for model outputs associated with a sequence group.""" - - @abstractmethod - def __repr__(self) -> str: - pass - - @abstractmethod - def __eq__(self, other: object) -> bool: - pass - - -class CompletionSequenceGroupOutput( - msgspec.Struct, - omit_defaults=True, # type: ignore[call-arg] - array_like=True): # type: ignore[call-arg] - __metaclass__ = SequenceGroupOutput - """The model output associated with a completion sequence group.""" - samples: List[SequenceOutput] - # Prompt logprob for each prompt query token. - prompt_logprobs: Optional[PromptLogprobs] - - def __repr__(self) -> str: - return (f"CompletionSequenceGroupOutput(samples={self.samples}, " - f"prompt_logprobs={self.prompt_logprobs})") - - def __eq__(self, other: object) -> bool: - if not isinstance(other, CompletionSequenceGroupOutput): - raise NotImplementedError() - return (self.samples == other.samples - and self.prompt_logprobs == other.prompt_logprobs) - - -class EmbeddingSequenceGroupOutput( - msgspec.Struct, - omit_defaults=True, # type: ignore[call-arg] - array_like=True, # type: ignore[call-arg] -): - """The model output associated with an embedding sequence group.""" - __metaclass__ = SequenceGroupOutput - embeddings: List[int] - - def __repr__(self) -> str: - return (f"EmbeddingSequenceGroupOutput(" - f"embeddings_shape={len(self.embeddings)})") - - def __eq__(self, other: object) -> bool: - if not isinstance(other, EmbeddingSequenceGroupOutput): - raise NotImplementedError() - return self.embeddings == other.embeddings - - -# cannot use msgspec.Struct here because Dynamo does not support it -@dataclass -class IntermediateTensors: - """For all pipeline stages except the last, we need to return the hidden - states and residuals to be sent to the next stage. This data structure - contains the hidden states and residuals for a request. - """ - - tensors: Dict[str, torch.Tensor] - - def __getitem__(self, key: Union[str, slice]): - if isinstance(key, str): - return self.tensors[key] - elif isinstance(key, slice): - return self.__class__({k: v[key] for k, v in self.tensors.items()}) - - def __setitem__(self, key: str, value): - self.tensors[key] = value - - def __len__(self): - return len(self.tensors) - - def __eq__(self, other: object): - return isinstance(other, self.__class__) and self - - def __repr__(self) -> str: - return f"IntermediateTensors(tensors={self.tensors})" - - -class PoolerOutput( - msgspec.Struct, - omit_defaults=True, # type: ignore[call-arg] - array_like=True): # type: ignore[call-arg] - """The output from a pooling operation in the embedding model.""" - outputs: List[EmbeddingSequenceGroupOutput] - - spec_decode_worker_metrics: Optional[SpecDecodeWorkerMetrics] = None - - def __getitem__(self, idx: int): - return self.outputs[idx] - - def __setitem__(self, idx: int, value): - self.outputs[idx] = value - - def __len__(self): - return len(self.outputs) - - def __eq__(self, other: object): - return isinstance(other, - self.__class__) and self.outputs == other.outputs - - -def get_all_seq_ids( - seq_group_metadata_list: List[SequenceGroupMetadata]) -> List[int]: - """Given a list of SequenceGroupMetadata, create a list of all - sequence ids. - """ - return [seq_id for sg in seq_group_metadata_list for seq_id in sg.seq_data] - - -def get_all_seq_ids_and_request_ids( - seq_group_metadata_list: List[SequenceGroupMetadata] -) -> Tuple[List[int], Dict[str, Set[int]]]: - """Given a list of SequenceGroupMetadata, create a list of all - sequence ids. - """ - seq_ids: List[int] = [] - request_id_seq_ids_mapping: Dict[str, Set[int]] = defaultdict(set) - for sg in seq_group_metadata_list: - for seq_id in sg.seq_data: - seq_ids.append(seq_id) - request_id_seq_ids_mapping[sg.request_id].add(seq_id) - return seq_ids, request_id_seq_ids_mapping - - -class HiddenStates(msgspec.Struct, array_like=True, - omit_defaults=True): # type: ignore[call-arg] - """Hidden states corresponding to in-progress sequences. - Used in speculative decoding to pass hidden states from - the target model to the proposer model. - - seq_ids are the sequence ids of each entry of the batch - dimension of the hidden_states tensor""" - # Scorer hidden states. For prefill step, it is used for hidden states of - # all tokens, whereas for decode step, it use used for last accepted tokens. - hidden_states: torch.Tensor - # The sequence group metadata list. Only needed for decode step. - seq_group_metadata_list: Optional[List[SequenceGroupMetadata]] = None - # Scorer hidden states of the 2nd last token proposed by the proposer ( - # irrespective of whether it was accepted or not). Only used for cases when - # last proposed token is accepted (i.e., in case of bonus tokens). For the - # case of no bonus tokens, these are ignored. - second_last_token_hidden_states: Optional[torch.Tensor] = None - - _seq_ids: List[int] = msgspec.field(default_factory=list) - - def __post_init__(self): - if self.seq_group_metadata_list is not None: - assert len(self.seq_group_metadata_list) == len(self.hidden_states) - self._seq_ids = get_all_seq_ids(self.seq_group_metadata_list) - - @property - def seq_ids(self) -> List[int]: - return self._seq_ids - - def update(self, - hidden_states: torch.Tensor, - seq_group_metadata_list: List[SequenceGroupMetadata], - second_last_token_hidden_states: Optional[torch.Tensor] = None): - """Update hidden states from target model invocation. Only used for - decode steps""" - assert len(seq_group_metadata_list) == len(hidden_states) - self._seq_ids.extend(get_all_seq_ids(seq_group_metadata_list)) - self.hidden_states = torch.cat([self.hidden_states, hidden_states]) - - if self.second_last_token_hidden_states is not None: - # Adding dummy hidden_states to this to maintain same shape - self.second_last_token_hidden_states = torch.cat([ - self.second_last_token_hidden_states, - torch.zeros_like(hidden_states) - if second_last_token_hidden_states is None else - second_last_token_hidden_states - ]) - - def prune(self, - seq_group_metadata_list: List[SequenceGroupMetadata]) -> None: - """Prune to provided list of sequence ids. Only used for decode steps. - """ - # Currently this prunes all seq_ids not present in - # seq_group_metadata_list which might cause problems where a sequence - # may be "paused" then "resumed" later. This should only prune sequences - # which are confirmed to be aborted. - seq_ids = get_all_seq_ids(seq_group_metadata_list) - if seq_ids != self._seq_ids: - # Batch contents changed - prune removed sequences. - index = [self._seq_ids.index(seq_id) for seq_id in seq_ids] - self.hidden_states = self.hidden_states[index] - if self.second_last_token_hidden_states is not None: - self.second_last_token_hidden_states = self\ - .second_last_token_hidden_states[index] - self._seq_ids = seq_ids - - def expand_with_bonus_tokens( - self, seq_with_bonus_token_in_last_step: set) -> None: - """Expand hidden states for sequences with bonus tokens. This is in - alignment with `MultiStepWorker._expand_execute_model_request`.""" - if self.second_last_token_hidden_states is None \ - or not seq_with_bonus_token_in_last_step: - return - - index = [] - for seq_id in self._seq_ids: - i = self._seq_ids.index(seq_id) - if seq_id in seq_with_bonus_token_in_last_step: - index.append(i + len(self._seq_ids)) - index.append(i) - - self.hidden_states = torch.cat( - [self.hidden_states, self.second_last_token_hidden_states])[index] - - -class ExecuteModelRequest( - msgspec.Struct, - array_like=True, # type: ignore[call-arg] - omit_defaults=True): # type: ignore[call-arg] - """The model execution request, containing CPU metadata only. The LLM - engine should create an instance of this class for each request batch.""" - # The sequence group metadata list. - seq_group_metadata_list: List[Union[SequenceGroupMetadata, - SequenceGroupMetadataDelta]] - # Blocks to swap in. List of CPU -> GPU block number. - blocks_to_swap_in: List[Tuple[int, - int]] = msgspec.field(default_factory=list) - # Blocks to swap out. List of GPU -> CPU block number. - blocks_to_swap_out: List[Tuple[int, - int]] = msgspec.field(default_factory=list) - # Blocks to copy. Source to dest block. - blocks_to_copy: List[Tuple[int, int]] = msgspec.field(default_factory=list) - # Virtual engine ID for pipeline parallel. - virtual_engine: int = 0 - # The number of slots for lookahead decoding. - num_lookahead_slots: int = 0 - # The number of requests in the running queue. - running_queue_size: int = 0 - # Optional hidden states from prior step. - previous_hidden_states: Optional[HiddenStates] = None - # The number of forward steps to run. - num_steps: int = 1 - # Finished request ids since last step. - finished_requests_ids: List[str] = msgspec.field(default_factory=list) - # The last sampled token ids for multi step decoding. - last_sampled_token_ids: Optional[torch.Tensor] = None - # Async callback - async_callback: Optional[Callable] = None - - @property - def is_first_multi_step(self) -> bool: - # TODO(will) make this be able to handle batches with variable number of - # steps - assert len(self.seq_group_metadata_list) > 0 - first_seq_group = self.seq_group_metadata_list[0] - assert first_seq_group.state is not None - return first_seq_group.state.current_step == 0 - - @property - def is_last_step(self) -> bool: - # TODO(will) make this be able to handle batches with variable number of - # steps - assert len(self.seq_group_metadata_list) > 0 - first_seq_group = self.seq_group_metadata_list[0] - assert first_seq_group.state is not None - return first_seq_group.state.remaining_steps == 1 - - @property - def current_step(self) -> int: - # TODO(will) make this be able to handle batches with variable number of - # steps - assert len(self.seq_group_metadata_list) > 0 - state = self.seq_group_metadata_list[0].state - assert state is not None - return state.current_step - - def clone( - self, seq_group_metadata_list: List[Union[SequenceGroupMetadata, - SequenceGroupMetadataDelta]] - ) -> "ExecuteModelRequest": - """Clone the request with a new sequence group metadata list.""" - return ExecuteModelRequest( - seq_group_metadata_list=seq_group_metadata_list, - blocks_to_swap_in=self.blocks_to_swap_in.copy(), - blocks_to_swap_out=self.blocks_to_swap_out.copy(), - blocks_to_copy=self.blocks_to_copy.copy(), - virtual_engine=self.virtual_engine, - num_lookahead_slots=self.num_lookahead_slots, - running_queue_size=self.running_queue_size, - previous_hidden_states=self.previous_hidden_states, - num_steps=self.num_steps, - finished_requests_ids=self.finished_requests_ids, - last_sampled_token_ids=self.last_sampled_token_ids.clone() - if self.last_sampled_token_ids is not None else None, - async_callback=self.async_callback) + + def finish_step(self) -> None: + assert self.state is not None + assert self.state.current_step < self.state.num_steps, \ + f"current step {self.state.current_step}, num_steps {self.state.num_steps}" # noqa + self.state.current_step += 1 + + +class SequenceOutput( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + array_like=True): # type: ignore[call-arg] + """The model output associated with a sequence. + + Args: + parent_seq_id: The ID of the parent sequence (for forking in beam + search). + output_token: The output token ID. + logprobs: The logprobs of the output token. + (Token id -> logP(x_i+1 | x_0, ..., x_i)) + """ + parent_seq_id: int + output_token: int + logprobs: Dict[int, Logprob] + + def __repr__(self) -> str: + return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, " + f"output_token={self.output_token}, " + f"logprobs={self.logprobs})") + + def __eq__(self, other: object) -> bool: + if not isinstance(other, SequenceOutput): + raise NotImplementedError() + equal = (self.parent_seq_id == other.parent_seq_id + and self.output_token == other.output_token) + log_probs_equal = other.logprobs == self.logprobs + return equal and log_probs_equal + + +class SequenceGroupOutput(ABC): + """The base class for model outputs associated with a sequence group.""" + + @abstractmethod + def __repr__(self) -> str: + pass + + @abstractmethod + def __eq__(self, other: object) -> bool: + pass + + +class CompletionSequenceGroupOutput( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + array_like=True): # type: ignore[call-arg] + __metaclass__ = SequenceGroupOutput + """The model output associated with a completion sequence group.""" + samples: List[SequenceOutput] + # Prompt logprob for each prompt query token. + prompt_logprobs: Optional[PromptLogprobs] + + def __repr__(self) -> str: + return (f"CompletionSequenceGroupOutput(samples={self.samples}, " + f"prompt_logprobs={self.prompt_logprobs})") + + def __eq__(self, other: object) -> bool: + if not isinstance(other, CompletionSequenceGroupOutput): + raise NotImplementedError() + return (self.samples == other.samples + and self.prompt_logprobs == other.prompt_logprobs) + + +class EmbeddingSequenceGroupOutput( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + array_like=True, # type: ignore[call-arg] +): + """The model output associated with an embedding sequence group.""" + __metaclass__ = SequenceGroupOutput + embeddings: List[int] + + def __repr__(self) -> str: + return (f"EmbeddingSequenceGroupOutput(" + f"embeddings_shape={len(self.embeddings)})") + + def __eq__(self, other: object) -> bool: + if not isinstance(other, EmbeddingSequenceGroupOutput): + raise NotImplementedError() + return self.embeddings == other.embeddings + + +# cannot use msgspec.Struct here because Dynamo does not support it +@dataclass +class IntermediateTensors: + """For all pipeline stages except the last, we need to return the hidden + states and residuals to be sent to the next stage. This data structure + contains the hidden states and residuals for a request. + """ + + tensors: Dict[str, torch.Tensor] + + def __getitem__(self, key: Union[str, slice]): + if isinstance(key, str): + return self.tensors[key] + elif isinstance(key, slice): + return self.__class__({k: v[key] for k, v in self.tensors.items()}) + + def __setitem__(self, key: str, value): + self.tensors[key] = value + + def __len__(self): + return len(self.tensors) + + def __eq__(self, other: object): + return isinstance(other, self.__class__) and self + + def __repr__(self) -> str: + return f"IntermediateTensors(tensors={self.tensors})" + + +class PoolerOutput( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + array_like=True): # type: ignore[call-arg] + """The output from a pooling operation in the embedding model.""" + outputs: List[EmbeddingSequenceGroupOutput] + + spec_decode_worker_metrics: Optional[SpecDecodeWorkerMetrics] = None + + def __getitem__(self, idx: int): + return self.outputs[idx] + + def __setitem__(self, idx: int, value): + self.outputs[idx] = value + + def __len__(self): + return len(self.outputs) + + def __eq__(self, other: object): + return isinstance(other, + self.__class__) and self.outputs == other.outputs + + +def get_all_seq_ids( + seq_group_metadata_list: List[SequenceGroupMetadata]) -> List[int]: + """Given a list of SequenceGroupMetadata, create a list of all + sequence ids. + """ + return [seq_id for sg in seq_group_metadata_list for seq_id in sg.seq_data] + + +def get_all_seq_ids_and_request_ids( + seq_group_metadata_list: List[SequenceGroupMetadata] +) -> Tuple[List[int], Dict[str, Set[int]]]: + """Given a list of SequenceGroupMetadata, create a list of all + sequence ids. + """ + seq_ids: List[int] = [] + request_id_seq_ids_mapping: Dict[str, Set[int]] = defaultdict(set) + for sg in seq_group_metadata_list: + for seq_id in sg.seq_data: + seq_ids.append(seq_id) + request_id_seq_ids_mapping[sg.request_id].add(seq_id) + return seq_ids, request_id_seq_ids_mapping + + +class HiddenStates(msgspec.Struct, array_like=True, + omit_defaults=True): # type: ignore[call-arg] + """Hidden states corresponding to in-progress sequences. + Used in speculative decoding to pass hidden states from + the target model to the proposer model. + + seq_ids are the sequence ids of each entry of the batch + dimension of the hidden_states tensor""" + # Scorer hidden states. For prefill step, it is used for hidden states of + # all tokens, whereas for decode step, it use used for last accepted tokens. + hidden_states: torch.Tensor + # The sequence group metadata list. Only needed for decode step. + seq_group_metadata_list: Optional[List[SequenceGroupMetadata]] = None + # Scorer hidden states of the 2nd last token proposed by the proposer ( + # irrespective of whether it was accepted or not). Only used for cases when + # last proposed token is accepted (i.e., in case of bonus tokens). For the + # case of no bonus tokens, these are ignored. + second_last_token_hidden_states: Optional[torch.Tensor] = None + + _seq_ids: List[int] = msgspec.field(default_factory=list) + + def __post_init__(self): + if self.seq_group_metadata_list is not None: + assert len(self.seq_group_metadata_list) == len(self.hidden_states) + self._seq_ids = get_all_seq_ids(self.seq_group_metadata_list) + + @property + def seq_ids(self) -> List[int]: + return self._seq_ids + + def update(self, + hidden_states: torch.Tensor, + seq_group_metadata_list: List[SequenceGroupMetadata], + second_last_token_hidden_states: Optional[torch.Tensor] = None): + """Update hidden states from target model invocation. Only used for + decode steps""" + assert len(seq_group_metadata_list) == len(hidden_states) + self._seq_ids.extend(get_all_seq_ids(seq_group_metadata_list)) + self.hidden_states = torch.cat([self.hidden_states, hidden_states]) + + if self.second_last_token_hidden_states is not None: + # Adding dummy hidden_states to this to maintain same shape + self.second_last_token_hidden_states = torch.cat([ + self.second_last_token_hidden_states, + torch.zeros_like(hidden_states) + if second_last_token_hidden_states is None else + second_last_token_hidden_states + ]) + + def prune(self, + seq_group_metadata_list: List[SequenceGroupMetadata]) -> None: + """Prune to provided list of sequence ids. Only used for decode steps. + """ + # Currently this prunes all seq_ids not present in + # seq_group_metadata_list which might cause problems where a sequence + # may be "paused" then "resumed" later. This should only prune sequences + # which are confirmed to be aborted. + seq_ids = get_all_seq_ids(seq_group_metadata_list) + if seq_ids != self._seq_ids: + # Batch contents changed - prune removed sequences. + index = [self._seq_ids.index(seq_id) for seq_id in seq_ids] + self.hidden_states = self.hidden_states[index] + if self.second_last_token_hidden_states is not None: + self.second_last_token_hidden_states = self\ + .second_last_token_hidden_states[index] + self._seq_ids = seq_ids + + def expand_with_bonus_tokens( + self, seq_with_bonus_token_in_last_step: set) -> None: + """Expand hidden states for sequences with bonus tokens. This is in + alignment with `MultiStepWorker._expand_execute_model_request`.""" + if self.second_last_token_hidden_states is None \ + or not seq_with_bonus_token_in_last_step: + return + + index = [] + for seq_id in self._seq_ids: + i = self._seq_ids.index(seq_id) + if seq_id in seq_with_bonus_token_in_last_step: + index.append(i + len(self._seq_ids)) + index.append(i) + + self.hidden_states = torch.cat( + [self.hidden_states, self.second_last_token_hidden_states])[index] + + +class ExecuteModelRequest( + msgspec.Struct, + array_like=True, # type: ignore[call-arg] + omit_defaults=True): # type: ignore[call-arg] + """The model execution request, containing CPU metadata only. The LLM + engine should create an instance of this class for each request batch.""" + # The sequence group metadata list. + seq_group_metadata_list: List[Union[SequenceGroupMetadata, + SequenceGroupMetadataDelta]] + # Blocks to swap in. List of CPU -> GPU block number. + blocks_to_swap_in: List[Tuple[int, + int]] = msgspec.field(default_factory=list) + # Blocks to swap out. List of GPU -> CPU block number. + blocks_to_swap_out: List[Tuple[int, + int]] = msgspec.field(default_factory=list) + # Blocks to copy. Source to dest block. + blocks_to_copy: List[Tuple[int, int]] = msgspec.field(default_factory=list) + # Virtual engine ID for pipeline parallel. + virtual_engine: int = 0 + # The number of slots for lookahead decoding. + num_lookahead_slots: int = 0 + # The number of requests in the running queue. + running_queue_size: int = 0 + # Optional hidden states from prior step. + previous_hidden_states: Optional[HiddenStates] = None + # The number of forward steps to run. + num_steps: int = 1 + # Finished request ids since last step. + finished_requests_ids: List[str] = msgspec.field(default_factory=list) + # The last sampled token ids for multi step decoding. + last_sampled_token_ids: Optional[torch.Tensor] = None + # Async callback + async_callback: Optional[Callable] = None + + @property + def is_first_multi_step(self) -> bool: + # TODO(will) make this be able to handle batches with variable number of + # steps + assert len(self.seq_group_metadata_list) > 0 + first_seq_group = self.seq_group_metadata_list[0] + assert first_seq_group.state is not None + return first_seq_group.state.current_step == 0 + + @property + def is_last_step(self) -> bool: + # TODO(will) make this be able to handle batches with variable number of + # steps + assert len(self.seq_group_metadata_list) > 0 + first_seq_group = self.seq_group_metadata_list[0] + assert first_seq_group.state is not None + return first_seq_group.state.remaining_steps == 1 + + @property + def current_step(self) -> int: + # TODO(will) make this be able to handle batches with variable number of + # steps + assert len(self.seq_group_metadata_list) > 0 + state = self.seq_group_metadata_list[0].state + assert state is not None + return state.current_step + + def clone( + self, seq_group_metadata_list: List[Union[SequenceGroupMetadata, + SequenceGroupMetadataDelta]] + ) -> "ExecuteModelRequest": + """Clone the request with a new sequence group metadata list.""" + return ExecuteModelRequest( + seq_group_metadata_list=seq_group_metadata_list, + blocks_to_swap_in=self.blocks_to_swap_in.copy(), + blocks_to_swap_out=self.blocks_to_swap_out.copy(), + blocks_to_copy=self.blocks_to_copy.copy(), + virtual_engine=self.virtual_engine, + num_lookahead_slots=self.num_lookahead_slots, + running_queue_size=self.running_queue_size, + previous_hidden_states=self.previous_hidden_states, + num_steps=self.num_steps, + finished_requests_ids=self.finished_requests_ids, + last_sampled_token_ids=self.last_sampled_token_ids.clone() + if self.last_sampled_token_ids is not None else None, + async_callback=self.async_callback) diff --git a/qwen3_6_scripts/serving_chat.py b/qwen3_6_scripts/serving_chat.py index 366253b1..14b3456f 100644 --- a/qwen3_6_scripts/serving_chat.py +++ b/qwen3_6_scripts/serving_chat.py @@ -1,47 +1,47 @@ -import asyncio -import json -import time -from typing import (AsyncGenerator, AsyncIterator, Callable, Dict, Final, List, - Optional) -from typing import Sequence as GenericSequence -from typing import Union - -from fastapi import Request - -from vllm.config import ModelConfig -from vllm.engine.async_llm_engine import AsyncLLMEngine -from vllm.engine.multiprocessing.client import MQLLMEngineClient -from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import (ConversationMessage, - apply_hf_chat_template, - apply_mistral_chat_template, - load_chat_template, - parse_chat_messages_futures) -from vllm.entrypoints.logger import RequestLogger -from vllm.entrypoints.openai.protocol import ( - ChatCompletionLogProb, ChatCompletionLogProbs, - ChatCompletionLogProbsContent, ChatCompletionNamedToolChoiceParam, - ChatCompletionRequest, ChatCompletionResponse, - ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice, - ChatCompletionStreamResponse, ChatMessage, DeltaFunctionCall, DeltaMessage, - DeltaToolCall, ErrorResponse, FunctionCall, RequestResponseMetadata, - PromptTokensDetails, ToolCall, UsageInfo) -from vllm.entrypoints.openai.serving_engine import (BaseModelPath, - LoRAModulePath, - OpenAIServing, - PromptAdapterPath, - TextTokensPrompt) -from vllm.entrypoints.openai.tool_parsers import ToolParser, ToolParserManager -from vllm.inputs import TokensPrompt -from vllm.logger import init_logger -from vllm.outputs import CompletionOutput, RequestOutput -from vllm.sampling_params import BeamSearchParams, SamplingParams -from vllm.sequence import Logprob -from vllm.tracing import (contains_trace_headers, extract_trace_headers, - log_tracing_disabled_warning) -from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer -from vllm.utils import iterate_with_cancellation, random_uuid - +import asyncio +import json +import time +from typing import (AsyncGenerator, AsyncIterator, Callable, Dict, Final, List, + Optional) +from typing import Sequence as GenericSequence +from typing import Union + +from fastapi import Request + +from vllm.config import ModelConfig +from vllm.engine.async_llm_engine import AsyncLLMEngine +from vllm.engine.multiprocessing.client import MQLLMEngineClient +from vllm.engine.protocol import EngineClient +from vllm.entrypoints.chat_utils import (ConversationMessage, + apply_hf_chat_template, + apply_mistral_chat_template, + load_chat_template, + parse_chat_messages_futures) +from vllm.entrypoints.logger import RequestLogger +from vllm.entrypoints.openai.protocol import ( + ChatCompletionLogProb, ChatCompletionLogProbs, + ChatCompletionLogProbsContent, ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, ChatCompletionResponse, + ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, ChatMessage, DeltaFunctionCall, DeltaMessage, + DeltaToolCall, ErrorResponse, FunctionCall, RequestResponseMetadata, + PromptTokensDetails, ToolCall, UsageInfo) +from vllm.entrypoints.openai.serving_engine import (BaseModelPath, + LoRAModulePath, + OpenAIServing, + PromptAdapterPath, + TextTokensPrompt) +from vllm.entrypoints.openai.tool_parsers import ToolParser, ToolParserManager +from vllm.inputs import TokensPrompt +from vllm.logger import init_logger +from vllm.outputs import CompletionOutput, RequestOutput +from vllm.sampling_params import BeamSearchParams, SamplingParams +from vllm.sequence import Logprob +from vllm.tracing import (contains_trace_headers, extract_trace_headers, + log_tracing_disabled_warning) +from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer +from vllm.utils import iterate_with_cancellation, random_uuid + logger = init_logger(__name__) @@ -197,90 +197,90 @@ def _merge_sequential_chat_responses( class OpenAIServingChat(OpenAIServing): - - def __init__(self, - engine_client: EngineClient, - model_config: ModelConfig, - base_model_paths: List[BaseModelPath], - response_role: str, - *, - lora_modules: Optional[List[LoRAModulePath]], - prompt_adapters: Optional[List[PromptAdapterPath]], - request_logger: Optional[RequestLogger], - chat_template: Optional[str], - return_tokens_as_token_ids: bool = False, - enable_auto_tools: bool = False, - tool_parser: Optional[str] = None, - reasoning_parser: Optional[str] = None): - super().__init__(engine_client=engine_client, - model_config=model_config, - base_model_paths=base_model_paths, - lora_modules=lora_modules, - prompt_adapters=prompt_adapters, - request_logger=request_logger, - return_tokens_as_token_ids=return_tokens_as_token_ids) - - self.response_role = response_role - self.use_tool_use_model_template = False - self.chat_template = load_chat_template(chat_template) - - # set up tool use - self.enable_auto_tools: bool = enable_auto_tools - if self.enable_auto_tools: - logger.info( - "\"auto\" tool choice has been enabled please note that while" - " the parallel_tool_calls client option is preset for " - "compatibility reasons, it will be ignored.") - - self.tool_parser: Optional[Callable[[AnyTokenizer], ToolParser]] = None - if self.enable_auto_tools: - try: - self.tool_parser = ToolParserManager.get_tool_parser( - tool_parser) - except Exception as e: - raise TypeError("Error: --enable-auto-tool-choice requires " - f"tool_parser:'{tool_parser}' which has not " - "been registered") from e - - # set up reasoning parser - self.reasoning_parser_cls = None - if reasoning_parser: - try: - from vllm.reasoning import ReasoningParserManager - self.reasoning_parser_cls = \ - ReasoningParserManager.get_reasoning_parser(reasoning_parser) - logger.info("Reasoning parser '%s' enabled.", reasoning_parser) - except Exception as e: - raise TypeError( - f"Error: --reasoning-parser '{reasoning_parser}' could not " - "be loaded. Make sure vllm/reasoning/ is installed." - ) from e - - async def create_chat_completion( - self, - request: ChatCompletionRequest, - raw_request: Optional[Request] = None, - ) -> Union[AsyncGenerator[str, None], ChatCompletionResponse, - ErrorResponse]: - """Completion API similar to OpenAI's API. - - See https://platform.openai.com/docs/api-reference/chat/create - for the API specification. This API mimics the OpenAI - ChatCompletion API. - - """ + + def __init__(self, + engine_client: EngineClient, + model_config: ModelConfig, + base_model_paths: List[BaseModelPath], + response_role: str, + *, + lora_modules: Optional[List[LoRAModulePath]], + prompt_adapters: Optional[List[PromptAdapterPath]], + request_logger: Optional[RequestLogger], + chat_template: Optional[str], + return_tokens_as_token_ids: bool = False, + enable_auto_tools: bool = False, + tool_parser: Optional[str] = None, + reasoning_parser: Optional[str] = None): + super().__init__(engine_client=engine_client, + model_config=model_config, + base_model_paths=base_model_paths, + lora_modules=lora_modules, + prompt_adapters=prompt_adapters, + request_logger=request_logger, + return_tokens_as_token_ids=return_tokens_as_token_ids) + + self.response_role = response_role + self.use_tool_use_model_template = False + self.chat_template = load_chat_template(chat_template) + + # set up tool use + self.enable_auto_tools: bool = enable_auto_tools + if self.enable_auto_tools: + logger.info( + "\"auto\" tool choice has been enabled please note that while" + " the parallel_tool_calls client option is preset for " + "compatibility reasons, it will be ignored.") + + self.tool_parser: Optional[Callable[[AnyTokenizer], ToolParser]] = None + if self.enable_auto_tools: + try: + self.tool_parser = ToolParserManager.get_tool_parser( + tool_parser) + except Exception as e: + raise TypeError("Error: --enable-auto-tool-choice requires " + f"tool_parser:'{tool_parser}' which has not " + "been registered") from e + + # set up reasoning parser + self.reasoning_parser_cls = None + if reasoning_parser: + try: + from vllm.reasoning import ReasoningParserManager + self.reasoning_parser_cls = \ + ReasoningParserManager.get_reasoning_parser(reasoning_parser) + logger.info("Reasoning parser '%s' enabled.", reasoning_parser) + except Exception as e: + raise TypeError( + f"Error: --reasoning-parser '{reasoning_parser}' could not " + "be loaded. Make sure vllm/reasoning/ is installed." + ) from e + + async def create_chat_completion( + self, + request: ChatCompletionRequest, + raw_request: Optional[Request] = None, + ) -> Union[AsyncGenerator[str, None], ChatCompletionResponse, + ErrorResponse]: + """Completion API similar to OpenAI's API. + + See https://platform.openai.com/docs/api-reference/chat/create + for the API specification. This API mimics the OpenAI + ChatCompletion API. + + """ if not request.messages: return self.create_error_response( "messages must contain at least one message") - error_check_ret = await self._check_model(request) - if error_check_ret is not None: - logger.error("Error with model %s", error_check_ret) - return error_check_ret - - # If the engine is dead, raise the engine's DEAD_ERROR. - # This is required for the streaming case, where we return a - # success status before we actually start generating text :). + error_check_ret = await self._check_model(request) + if error_check_ret is not None: + logger.error("Error with model %s", error_check_ret) + return error_check_ret + + # If the engine is dead, raise the engine's DEAD_ERROR. + # This is required for the streaming case, where we return a + # success status before we actually start generating text :). if self.engine_client.errored: raise self.engine_client.dead_error @@ -302,169 +302,169 @@ class OpenAIServingChat(OpenAIServing): f"Use n<={max_num_seqs} or omit n.") try: - ( - lora_request, - prompt_adapter_request, - ) = self._maybe_get_adapters(request) - - model_config = self.model_config - tokenizer = await self.engine_client.get_tokenizer(lora_request) - - conversation, mm_data_future = parse_chat_messages_futures( - request.messages, model_config, tokenizer) - - tool_dicts = None if request.tools is None else [ - tool.model_dump() for tool in request.tools - ] - - prompt: Union[str, List[int]] - is_mistral_tokenizer = isinstance(tokenizer, MistralTokenizer) - if is_mistral_tokenizer: - prompt = apply_mistral_chat_template( - tokenizer, - messages=request.messages, - chat_template=request.chat_template or self.chat_template, - add_generation_prompt=request.add_generation_prompt, - continue_final_message=request.continue_final_message, - tools=tool_dicts, - documents=request.documents, - **(request.chat_template_kwargs or {}), - ) - else: - prompt = apply_hf_chat_template( - tokenizer, - conversation=conversation, - chat_template=request.chat_template or self.chat_template, - add_generation_prompt=request.add_generation_prompt, - continue_final_message=request.continue_final_message, - tools=tool_dicts, - documents=request.documents, - **(request.chat_template_kwargs or {}), - ) - except Exception as e: - logger.exception("Error in applying chat template from request") - return self.create_error_response(str(e)) - - try: - mm_data = await mm_data_future - except Exception as e: - logger.exception("Error in loading multi-modal data") - return self.create_error_response(str(e)) - + ( + lora_request, + prompt_adapter_request, + ) = self._maybe_get_adapters(request) + + model_config = self.model_config + tokenizer = await self.engine_client.get_tokenizer(lora_request) + + conversation, mm_data_future = parse_chat_messages_futures( + request.messages, model_config, tokenizer) + + tool_dicts = None if request.tools is None else [ + tool.model_dump() for tool in request.tools + ] + + prompt: Union[str, List[int]] + is_mistral_tokenizer = isinstance(tokenizer, MistralTokenizer) + if is_mistral_tokenizer: + prompt = apply_mistral_chat_template( + tokenizer, + messages=request.messages, + chat_template=request.chat_template or self.chat_template, + add_generation_prompt=request.add_generation_prompt, + continue_final_message=request.continue_final_message, + tools=tool_dicts, + documents=request.documents, + **(request.chat_template_kwargs or {}), + ) + else: + prompt = apply_hf_chat_template( + tokenizer, + conversation=conversation, + chat_template=request.chat_template or self.chat_template, + add_generation_prompt=request.add_generation_prompt, + continue_final_message=request.continue_final_message, + tools=tool_dicts, + documents=request.documents, + **(request.chat_template_kwargs or {}), + ) + except Exception as e: + logger.exception("Error in applying chat template from request") + return self.create_error_response(str(e)) + + try: + mm_data = await mm_data_future + except Exception as e: + logger.exception("Error in loading multi-modal data") + return self.create_error_response(str(e)) + # validation for OpenAI tools - # tool_choice = "required" is not supported - if request.tool_choice == "required": - return self.create_error_response( - "tool_choice = \"required\" is not supported!") - - if not is_mistral_tokenizer and request.tool_choice == "auto" and not ( - self.enable_auto_tools and self.tool_parser is not None): - # for hf tokenizers, "auto" tools requires - # --enable-auto-tool-choice and --tool-call-parser - return self.create_error_response( - "\"auto\" tool choice requires " - "--enable-auto-tool-choice and --tool-call-parser to be set") - - request_id = f"chat-{random_uuid()}" - - request_metadata = RequestResponseMetadata(request_id=request_id) - if raw_request: - raw_request.state.request_metadata = request_metadata - - try: - if self.enable_auto_tools and self.tool_parser: - request = self.tool_parser(tokenizer).adjust_request( - request=request) - - if isinstance(prompt, str): - prompt_inputs = self._tokenize_prompt_input( - request, - tokenizer, - prompt, - truncate_prompt_tokens=request.truncate_prompt_tokens, - add_special_tokens=request.add_special_tokens, - ) - else: - assert isinstance(prompt, list) and isinstance( - prompt[0], int - ), "Prompt has to be either a string or a list of token ids" - prompt_inputs = TextTokensPrompt( - prompt=tokenizer.decode(prompt), prompt_token_ids=prompt) - - assert prompt_inputs is not None - - sampling_params: Union[SamplingParams, BeamSearchParams] - # OpenAI API: max_completion_tokens takes precedence over max_tokens - if request.max_completion_tokens is not None and request.max_tokens is None: - request.max_tokens = request.max_completion_tokens - default_max_tokens = self.max_model_len - len( - prompt_inputs["prompt_token_ids"]) - if request.use_beam_search: - sampling_params = request.to_beam_search_params( - default_max_tokens) - else: - sampling_params = request.to_sampling_params( - default_max_tokens) - - self._log_inputs(request_id, - prompt_inputs, - params=sampling_params, - lora_request=lora_request, - prompt_adapter_request=prompt_adapter_request) - - engine_inputs = TokensPrompt( - prompt_token_ids=prompt_inputs["prompt_token_ids"]) - if mm_data is not None: - engine_inputs["multi_modal_data"] = mm_data - - is_tracing_enabled = (await - self.engine_client.is_tracing_enabled()) - trace_headers = None - if is_tracing_enabled and raw_request: - trace_headers = extract_trace_headers(raw_request.headers) - if (not is_tracing_enabled and raw_request - and contains_trace_headers(raw_request.headers)): - log_tracing_disabled_warning() - - if isinstance(sampling_params, BeamSearchParams): - assert isinstance(self.engine_client, - (AsyncLLMEngine, - MQLLMEngineClient)), \ - "Beam search is only supported with" \ - "AsyncLLMEngine and MQLLMEngineClient." - result_generator = self.engine_client.beam_search( - engine_inputs['prompt_token_ids'], - request_id, - sampling_params, - ) - else: - result_generator = self.engine_client.generate( - engine_inputs, - sampling_params, - request_id, - lora_request=lora_request, - trace_headers=trace_headers, - prompt_adapter_request=prompt_adapter_request, - priority=request.priority, - ) - except ValueError as e: - # TODO: Use a vllm-specific Validation Error - return self.create_error_response(str(e)) - - if raw_request: - result_generator = iterate_with_cancellation( - result_generator, raw_request.is_disconnected) - - # Streaming response - if request.stream: - return self.chat_completion_stream_generator( - request, result_generator, request_id, conversation, tokenizer, - request_metadata, raw_request=raw_request) - - try: - return await self.chat_completion_full_generator( - request, result_generator, request_id, conversation, tokenizer, - request_metadata, raw_request=raw_request) + # tool_choice = "required" is not supported + if request.tool_choice == "required": + return self.create_error_response( + "tool_choice = \"required\" is not supported!") + + if not is_mistral_tokenizer and request.tool_choice == "auto" and not ( + self.enable_auto_tools and self.tool_parser is not None): + # for hf tokenizers, "auto" tools requires + # --enable-auto-tool-choice and --tool-call-parser + return self.create_error_response( + "\"auto\" tool choice requires " + "--enable-auto-tool-choice and --tool-call-parser to be set") + + request_id = f"chat-{random_uuid()}" + + request_metadata = RequestResponseMetadata(request_id=request_id) + if raw_request: + raw_request.state.request_metadata = request_metadata + + try: + if self.enable_auto_tools and self.tool_parser: + request = self.tool_parser(tokenizer).adjust_request( + request=request) + + if isinstance(prompt, str): + prompt_inputs = self._tokenize_prompt_input( + request, + tokenizer, + prompt, + truncate_prompt_tokens=request.truncate_prompt_tokens, + add_special_tokens=request.add_special_tokens, + ) + else: + assert isinstance(prompt, list) and isinstance( + prompt[0], int + ), "Prompt has to be either a string or a list of token ids" + prompt_inputs = TextTokensPrompt( + prompt=tokenizer.decode(prompt), prompt_token_ids=prompt) + + assert prompt_inputs is not None + + sampling_params: Union[SamplingParams, BeamSearchParams] + # OpenAI API: max_completion_tokens takes precedence over max_tokens + if request.max_completion_tokens is not None and request.max_tokens is None: + request.max_tokens = request.max_completion_tokens + default_max_tokens = self.max_model_len - len( + prompt_inputs["prompt_token_ids"]) + if request.use_beam_search: + sampling_params = request.to_beam_search_params( + default_max_tokens) + else: + sampling_params = request.to_sampling_params( + default_max_tokens) + + self._log_inputs(request_id, + prompt_inputs, + params=sampling_params, + lora_request=lora_request, + prompt_adapter_request=prompt_adapter_request) + + engine_inputs = TokensPrompt( + prompt_token_ids=prompt_inputs["prompt_token_ids"]) + if mm_data is not None: + engine_inputs["multi_modal_data"] = mm_data + + is_tracing_enabled = (await + self.engine_client.is_tracing_enabled()) + trace_headers = None + if is_tracing_enabled and raw_request: + trace_headers = extract_trace_headers(raw_request.headers) + if (not is_tracing_enabled and raw_request + and contains_trace_headers(raw_request.headers)): + log_tracing_disabled_warning() + + if isinstance(sampling_params, BeamSearchParams): + assert isinstance(self.engine_client, + (AsyncLLMEngine, + MQLLMEngineClient)), \ + "Beam search is only supported with" \ + "AsyncLLMEngine and MQLLMEngineClient." + result_generator = self.engine_client.beam_search( + engine_inputs['prompt_token_ids'], + request_id, + sampling_params, + ) + else: + result_generator = self.engine_client.generate( + engine_inputs, + sampling_params, + request_id, + lora_request=lora_request, + trace_headers=trace_headers, + prompt_adapter_request=prompt_adapter_request, + priority=request.priority, + ) + except ValueError as e: + # TODO: Use a vllm-specific Validation Error + return self.create_error_response(str(e)) + + if raw_request: + result_generator = iterate_with_cancellation( + result_generator, raw_request.is_disconnected) + + # Streaming response + if request.stream: + return self.chat_completion_stream_generator( + request, result_generator, request_id, conversation, tokenizer, + request_metadata, raw_request=raw_request) + + try: + return await self.chat_completion_full_generator( + request, result_generator, request_id, conversation, tokenizer, + request_metadata, raw_request=raw_request) except ValueError as e: # TODO: Use a vllm-specific Validation Error return self.create_error_response(str(e)) @@ -522,32 +522,32 @@ class OpenAIServingChat(OpenAIServing): return response def get_chat_request_role(self, request: ChatCompletionRequest) -> str: - if request.add_generation_prompt: - return self.response_role - return request.messages[-1]["role"] - - async def chat_completion_stream_generator( - self, - request: ChatCompletionRequest, - result_generator: AsyncIterator[RequestOutput], - request_id: str, - conversation: List[ConversationMessage], - tokenizer: AnyTokenizer, - request_metadata: RequestResponseMetadata, - raw_request: Optional[Request] = None, - ) -> AsyncGenerator[str, None]: - model_name = self.base_model_paths[0].name - created_time = int(time.time()) - chunk_object_type: Final = "chat.completion.chunk" - first_iteration = True - - # Send response for each token for each request.n (index) - num_choices = 1 if request.n is None else request.n - previous_num_tokens = [0] * num_choices - finish_reason_sent = [False] * num_choices - num_prompt_tokens = 0 - num_cached_tokens: Optional[int] = None - + if request.add_generation_prompt: + return self.response_role + return request.messages[-1]["role"] + + async def chat_completion_stream_generator( + self, + request: ChatCompletionRequest, + result_generator: AsyncIterator[RequestOutput], + request_id: str, + conversation: List[ConversationMessage], + tokenizer: AnyTokenizer, + request_metadata: RequestResponseMetadata, + raw_request: Optional[Request] = None, + ) -> AsyncGenerator[str, None]: + model_name = self.base_model_paths[0].name + created_time = int(time.time()) + chunk_object_type: Final = "chat.completion.chunk" + first_iteration = True + + # Send response for each token for each request.n (index) + num_choices = 1 if request.n is None else request.n + previous_num_tokens = [0] * num_choices + finish_reason_sent = [False] * num_choices + num_prompt_tokens = 0 + num_cached_tokens: Optional[int] = None + if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam): tool_choice_function_name = request.tool_choice.function.name else: @@ -556,219 +556,219 @@ class OpenAIServingChat(OpenAIServing): [f"chatcmpl-tool-{random_uuid()}" for _ in range(num_choices)] if tool_choice_function_name else []) named_tool_header_sent = [False] * num_choices - - # Determine whether tools are in use with "auto" tool choice - tool_choice_auto = ( - not tool_choice_function_name - and self._should_stream_with_auto_tool_parsing(request)) - - use_reasoning = self.reasoning_parser_cls is not None - - all_previous_token_ids: Optional[List[List[int]]] - # previous_texts / all_previous_token_ids are needed for both tool - # parsing and reasoning parsing (both require full-history context). - if tool_choice_auto or use_reasoning: - previous_texts = [""] * num_choices - all_previous_token_ids = [[]] * num_choices - else: - previous_texts, all_previous_token_ids = None, None - - # Prepare the tool parser if it's needed - try: - if tool_choice_auto and self.tool_parser: - tool_parsers: List[Optional[ToolParser]] = [ - self.tool_parser(tokenizer) - ] * num_choices - else: - tool_parsers = [None] * num_choices - except RuntimeError as e: - logger.error("Error in tool parser creation: %s", e) - data = self.create_streaming_error_response(str(e)) - yield f"data: {data}\n\n" - yield "data: [DONE]\n\n" - return - - # Prepare reasoning parsers (one instance per choice for state isolation) - reasoning_parsers: List[Optional[object]] = [None] * num_choices - reasoning_end_arr: List[bool] = [False] * num_choices - reasoning_token_counts: List[int] = [0] * num_choices - if use_reasoning: - try: - reasoning_parsers = [ - self.reasoning_parser_cls( - tokenizer, - chat_template_kwargs=request.chat_template_kwargs) - for _ in range(num_choices) - ] - # If thinking is disabled per-request, mark reasoning as - # already ended so the tool-auto branch is reachable. - for idx, rp in enumerate(reasoning_parsers): - if hasattr(rp, 'thinking_enabled') and not rp.thinking_enabled: - reasoning_end_arr[idx] = True - except RuntimeError as e: - logger.error("Error in reasoning parser creation: %s", e) - data = self.create_streaming_error_response(str(e)) - yield f"data: {data}\n\n" - yield "data: [DONE]\n\n" - return - - # Background task: poll is_disconnected() every 300 ms and abort the - # engine request as soon as the client goes away. This catches the - # case where the HTTP layer (Starlette/uvicorn) does not actively read - # the receive channel during streaming, so is_disconnected() in - # iterate_with_cancellation never fires during fast decode. - _disconnect_watcher: Optional[asyncio.Task] = None - if raw_request is not None: - async def _watch_disconnect() -> None: - try: - while True: - if await raw_request.is_disconnected(): - logger.info( - "Client disconnected (decode watcher), " - "aborting request %s", request_id) - await self.engine_client.abort(request_id) - return - await asyncio.sleep(0.3) - except asyncio.CancelledError: - pass - _disconnect_watcher = asyncio.ensure_future(_watch_disconnect()) - - try: - async for res in result_generator: - if res.prompt_token_ids is not None: - num_prompt_tokens = len(res.prompt_token_ids) - if res.encoder_prompt_token_ids is not None: - num_prompt_tokens += len(res.encoder_prompt_token_ids) - if (num_cached_tokens is None - and res.metrics is not None - and res.metrics.num_cached_tokens is not None): - num_cached_tokens = res.metrics.num_cached_tokens - - # We need to do it here, because if there are exceptions in - # the result_generator, it needs to be sent as the FIRST - # response (by the try...catch). - if first_iteration: - # Send first response for each request.n (index) with - # the role - role = self.get_chat_request_role(request) - - # NOTE num_choices defaults to 1 so this usually executes - # once per request - for i in range(num_choices): - tool_parser = tool_parsers[i] - choice_data = ChatCompletionResponseStreamChoice( - index=i, - delta=DeltaMessage( - role=role, - content="", - ), - logprobs=None, - finish_reason=None) - chunk = ChatCompletionStreamResponse( - id=request_id, - object=chunk_object_type, - created=created_time, - choices=[choice_data], - model=model_name) - - # if usage should be included - if (request.stream_options - and request.stream_options.include_usage): - # if continuous usage stats are requested, add it - if request.stream_options.continuous_usage_stats: - usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - completion_tokens=0, - total_tokens=num_prompt_tokens) - chunk.usage = usage - # otherwise don't - else: - chunk.usage = None - - data = chunk.model_dump_json(exclude_unset=True) - yield f"data: {data}\n\n" - - # Send response to echo the input portion of the - # last message - if request.echo or request.continue_final_message: - last_msg_content: str = "" - if conversation and "content" in conversation[ - -1] and conversation[-1].get("role") == role: - last_msg_content = conversation[-1]["content"] or "" - - if last_msg_content: - for i in range(num_choices): - choice_data = ( - ChatCompletionResponseStreamChoice( - index=i, - delta=DeltaMessage( - content=last_msg_content), - logprobs=None, - finish_reason=None)) - chunk = ChatCompletionStreamResponse( - id=request_id, - object=chunk_object_type, - created=created_time, - choices=[choice_data], - model=model_name) - if (request.stream_options and - request.stream_options.include_usage): - if (request.stream_options. - continuous_usage_stats): - usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - completion_tokens=0, - total_tokens=num_prompt_tokens) - chunk.usage = usage - else: - chunk.usage = None - - data = chunk.model_dump_json( - exclude_unset=True) - yield f"data: {data}\n\n" - first_iteration = False - - for output in res.outputs: - i = output.index - tool_parser = tool_parsers[i] - - if finish_reason_sent[i]: - continue - - if request.logprobs and request.top_logprobs is not None: - assert output.logprobs is not None, ( - "Did not output logprobs") - logprobs = self._create_chat_logprobs( - token_ids=output.token_ids, - top_logprobs=output.logprobs, - tokenizer=tokenizer, - num_output_top_logprobs=request.top_logprobs, - ) - else: - logprobs = None - - delta_text = output.text - delta_message: Optional[DeltaMessage] - - # Maintain text/token history when either reasoning or - # auto-tool parsing is active. - assert previous_texts is not None or not ( - tool_choice_auto or use_reasoning) - if previous_texts is not None: - assert all_previous_token_ids is not None - previous_text = previous_texts[i] - previous_token_ids = all_previous_token_ids[i] - current_text = previous_text + delta_text - current_token_ids = previous_token_ids + list( - output.token_ids) - previous_texts[i] = current_text - all_previous_token_ids[i] = current_token_ids - else: - previous_text = "" - previous_token_ids = [] - current_text = delta_text - current_token_ids = list(output.token_ids) - + + # Determine whether tools are in use with "auto" tool choice + tool_choice_auto = ( + not tool_choice_function_name + and self._should_stream_with_auto_tool_parsing(request)) + + use_reasoning = self.reasoning_parser_cls is not None + + all_previous_token_ids: Optional[List[List[int]]] + # previous_texts / all_previous_token_ids are needed for both tool + # parsing and reasoning parsing (both require full-history context). + if tool_choice_auto or use_reasoning: + previous_texts = [""] * num_choices + all_previous_token_ids = [[]] * num_choices + else: + previous_texts, all_previous_token_ids = None, None + + # Prepare the tool parser if it's needed + try: + if tool_choice_auto and self.tool_parser: + tool_parsers: List[Optional[ToolParser]] = [ + self.tool_parser(tokenizer) + ] * num_choices + else: + tool_parsers = [None] * num_choices + except RuntimeError as e: + logger.error("Error in tool parser creation: %s", e) + data = self.create_streaming_error_response(str(e)) + yield f"data: {data}\n\n" + yield "data: [DONE]\n\n" + return + + # Prepare reasoning parsers (one instance per choice for state isolation) + reasoning_parsers: List[Optional[object]] = [None] * num_choices + reasoning_end_arr: List[bool] = [False] * num_choices + reasoning_token_counts: List[int] = [0] * num_choices + if use_reasoning: + try: + reasoning_parsers = [ + self.reasoning_parser_cls( + tokenizer, + chat_template_kwargs=request.chat_template_kwargs) + for _ in range(num_choices) + ] + # If thinking is disabled per-request, mark reasoning as + # already ended so the tool-auto branch is reachable. + for idx, rp in enumerate(reasoning_parsers): + if hasattr(rp, 'thinking_enabled') and not rp.thinking_enabled: + reasoning_end_arr[idx] = True + except RuntimeError as e: + logger.error("Error in reasoning parser creation: %s", e) + data = self.create_streaming_error_response(str(e)) + yield f"data: {data}\n\n" + yield "data: [DONE]\n\n" + return + + # Background task: poll is_disconnected() every 300 ms and abort the + # engine request as soon as the client goes away. This catches the + # case where the HTTP layer (Starlette/uvicorn) does not actively read + # the receive channel during streaming, so is_disconnected() in + # iterate_with_cancellation never fires during fast decode. + _disconnect_watcher: Optional[asyncio.Task] = None + if raw_request is not None: + async def _watch_disconnect() -> None: + try: + while True: + if await raw_request.is_disconnected(): + logger.info( + "Client disconnected (decode watcher), " + "aborting request %s", request_id) + await self.engine_client.abort(request_id) + return + await asyncio.sleep(0.3) + except asyncio.CancelledError: + pass + _disconnect_watcher = asyncio.ensure_future(_watch_disconnect()) + + try: + async for res in result_generator: + if res.prompt_token_ids is not None: + num_prompt_tokens = len(res.prompt_token_ids) + if res.encoder_prompt_token_ids is not None: + num_prompt_tokens += len(res.encoder_prompt_token_ids) + if (num_cached_tokens is None + and res.metrics is not None + and res.metrics.num_cached_tokens is not None): + num_cached_tokens = res.metrics.num_cached_tokens + + # We need to do it here, because if there are exceptions in + # the result_generator, it needs to be sent as the FIRST + # response (by the try...catch). + if first_iteration: + # Send first response for each request.n (index) with + # the role + role = self.get_chat_request_role(request) + + # NOTE num_choices defaults to 1 so this usually executes + # once per request + for i in range(num_choices): + tool_parser = tool_parsers[i] + choice_data = ChatCompletionResponseStreamChoice( + index=i, + delta=DeltaMessage( + role=role, + content="", + ), + logprobs=None, + finish_reason=None) + chunk = ChatCompletionStreamResponse( + id=request_id, + object=chunk_object_type, + created=created_time, + choices=[choice_data], + model=model_name) + + # if usage should be included + if (request.stream_options + and request.stream_options.include_usage): + # if continuous usage stats are requested, add it + if request.stream_options.continuous_usage_stats: + usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=0, + total_tokens=num_prompt_tokens) + chunk.usage = usage + # otherwise don't + else: + chunk.usage = None + + data = chunk.model_dump_json(exclude_unset=True) + yield f"data: {data}\n\n" + + # Send response to echo the input portion of the + # last message + if request.echo or request.continue_final_message: + last_msg_content: str = "" + if conversation and "content" in conversation[ + -1] and conversation[-1].get("role") == role: + last_msg_content = conversation[-1]["content"] or "" + + if last_msg_content: + for i in range(num_choices): + choice_data = ( + ChatCompletionResponseStreamChoice( + index=i, + delta=DeltaMessage( + content=last_msg_content), + logprobs=None, + finish_reason=None)) + chunk = ChatCompletionStreamResponse( + id=request_id, + object=chunk_object_type, + created=created_time, + choices=[choice_data], + model=model_name) + if (request.stream_options and + request.stream_options.include_usage): + if (request.stream_options. + continuous_usage_stats): + usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=0, + total_tokens=num_prompt_tokens) + chunk.usage = usage + else: + chunk.usage = None + + data = chunk.model_dump_json( + exclude_unset=True) + yield f"data: {data}\n\n" + first_iteration = False + + for output in res.outputs: + i = output.index + tool_parser = tool_parsers[i] + + if finish_reason_sent[i]: + continue + + if request.logprobs and request.top_logprobs is not None: + assert output.logprobs is not None, ( + "Did not output logprobs") + logprobs = self._create_chat_logprobs( + token_ids=output.token_ids, + top_logprobs=output.logprobs, + tokenizer=tokenizer, + num_output_top_logprobs=request.top_logprobs, + ) + else: + logprobs = None + + delta_text = output.text + delta_message: Optional[DeltaMessage] + + # Maintain text/token history when either reasoning or + # auto-tool parsing is active. + assert previous_texts is not None or not ( + tool_choice_auto or use_reasoning) + if previous_texts is not None: + assert all_previous_token_ids is not None + previous_text = previous_texts[i] + previous_token_ids = all_previous_token_ids[i] + current_text = previous_text + delta_text + current_token_ids = previous_token_ids + list( + output.token_ids) + previous_texts[i] = current_text + all_previous_token_ids[i] = current_token_ids + else: + previous_text = "" + previous_token_ids = [] + current_text = delta_text + current_token_ids = list(output.token_ids) + # handle streaming deltas for tools with named tool_choice if tool_choice_function_name: first_named_delta = _consume_named_tool_header_slot( @@ -782,104 +782,104 @@ class OpenAIServingChat(OpenAIServing): first_named_delta, )) ]) - - # handle reasoning: route through reasoning parser while - # has not yet been seen. - elif use_reasoning and not reasoning_end_arr[i]: - r_parser = reasoning_parsers[i] - delta_message = r_parser.extract_reasoning_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=delta_text, - previous_token_ids=previous_token_ids, - current_token_ids=current_token_ids, - delta_token_ids=output.token_ids, - ) - # Mark reasoning as ended when end token appears. - if r_parser.end_token_id in current_token_ids: - reasoning_end_arr[i] = True - - # handle streaming deltas for tools with "auto" tool choice - # (only reached after reasoning block, if any, has ended) - elif tool_choice_auto: - assert tool_parser is not None - delta_message = ( - tool_parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=delta_text, - previous_token_ids=previous_token_ids, - current_token_ids=current_token_ids, - delta_token_ids=output.token_ids, - request=request)) - - # handle streaming just a content delta - else: - delta_message = DeltaMessage(content=delta_text) - - # set the previous values for the next iteration - previous_num_tokens[i] += len(output.token_ids) - - # if the message delta is None (e.g. because it was a - # "control token" for tool calls or the parser otherwise - # wasn't ready to send a token, then - # get the next token without streaming a chunk. - # However, if this is the finish token we must NOT skip — - # the finish block updates reasoning_token_counts, sets - # finish_reason_sent, and flushes the final usage chunk. - if delta_message is None: - if output.finish_reason is None: - continue - delta_message = DeltaMessage() - - if output.finish_reason is None: - # Send token-by-token response for each request.n - - choice_data = ChatCompletionResponseStreamChoice( - index=i, - delta=delta_message, - logprobs=logprobs, - finish_reason=None) - chunk = ChatCompletionStreamResponse( - id=request_id, - object=chunk_object_type, - created=created_time, - choices=[choice_data], - model=model_name) - - # handle usage stats if requested & if continuous - if (request.stream_options - and request.stream_options.include_usage): - if request.stream_options.continuous_usage_stats: - completion_tokens = len(output.token_ids) - usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=num_prompt_tokens + - completion_tokens, - ) - chunk.usage = usage - else: - chunk.usage = None - - data = chunk.model_dump_json(exclude_unset=True) - yield f"data: {data}\n\n" - - # if the model is finished generating - else: - # check to make sure we haven't "forgotten" to stream - # any tokens that were generated but previously - # matched by partial json parsing - # only happens if we are NOT using guided decoding - auto_tools_called = False - if tool_parser: - auto_tools_called = len( - tool_parser.prev_tool_call_arr) > 0 - index = len(tool_parser.prev_tool_call_arr - ) - 1 if auto_tools_called else 0 - else: - index = 0 - + + # handle reasoning: route through reasoning parser while + # has not yet been seen. + elif use_reasoning and not reasoning_end_arr[i]: + r_parser = reasoning_parsers[i] + delta_message = r_parser.extract_reasoning_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=previous_token_ids, + current_token_ids=current_token_ids, + delta_token_ids=output.token_ids, + ) + # Mark reasoning as ended when end token appears. + if r_parser.end_token_id in current_token_ids: + reasoning_end_arr[i] = True + + # handle streaming deltas for tools with "auto" tool choice + # (only reached after reasoning block, if any, has ended) + elif tool_choice_auto: + assert tool_parser is not None + delta_message = ( + tool_parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=previous_token_ids, + current_token_ids=current_token_ids, + delta_token_ids=output.token_ids, + request=request)) + + # handle streaming just a content delta + else: + delta_message = DeltaMessage(content=delta_text) + + # set the previous values for the next iteration + previous_num_tokens[i] += len(output.token_ids) + + # if the message delta is None (e.g. because it was a + # "control token" for tool calls or the parser otherwise + # wasn't ready to send a token, then + # get the next token without streaming a chunk. + # However, if this is the finish token we must NOT skip — + # the finish block updates reasoning_token_counts, sets + # finish_reason_sent, and flushes the final usage chunk. + if delta_message is None: + if output.finish_reason is None: + continue + delta_message = DeltaMessage() + + if output.finish_reason is None: + # Send token-by-token response for each request.n + + choice_data = ChatCompletionResponseStreamChoice( + index=i, + delta=delta_message, + logprobs=logprobs, + finish_reason=None) + chunk = ChatCompletionStreamResponse( + id=request_id, + object=chunk_object_type, + created=created_time, + choices=[choice_data], + model=model_name) + + # handle usage stats if requested & if continuous + if (request.stream_options + and request.stream_options.include_usage): + if request.stream_options.continuous_usage_stats: + completion_tokens = len(output.token_ids) + usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=num_prompt_tokens + + completion_tokens, + ) + chunk.usage = usage + else: + chunk.usage = None + + data = chunk.model_dump_json(exclude_unset=True) + yield f"data: {data}\n\n" + + # if the model is finished generating + else: + # check to make sure we haven't "forgotten" to stream + # any tokens that were generated but previously + # matched by partial json parsing + # only happens if we are NOT using guided decoding + auto_tools_called = False + if tool_parser: + auto_tools_called = len( + tool_parser.prev_tool_call_arr) > 0 + index = len(tool_parser.prev_tool_call_arr + ) - 1 if auto_tools_called else 0 + else: + index = 0 + if self._should_check_for_unstreamed_tool_arg_tokens( delta_message, output) and tool_parser: # get the expected call based on partial JSON @@ -887,206 +887,206 @@ class OpenAIServingChat(OpenAIServing): expected_call = _serialize_tool_arguments( tool_parser.prev_tool_call_arr[index].get( "arguments", {})) - - # get what we've streamed so far for arguments - # for the current tool - actual_call = tool_parser.streamed_args_for_tool[ - index] - - # check to see if there's anything left to stream - remaining_call = expected_call.replace( - actual_call, "", 1) - - # set that as a delta message - delta_message = DeltaMessage(tool_calls=[ - DeltaToolCall(index=index, - function=DeltaFunctionCall( - arguments=remaining_call). - model_dump(exclude_none=True)) - ]) - - # Count reasoning tokens for this choice at finish time. - if use_reasoning and all_previous_token_ids is not None: - r_parser = reasoning_parsers[i] - reasoning_token_counts[i] = \ - r_parser.count_reasoning_tokens( - all_previous_token_ids[i]) - - # Send the finish response for each request.n only once - choice_data = ChatCompletionResponseStreamChoice( - index=i, - delta=delta_message, - logprobs=logprobs, + + # get what we've streamed so far for arguments + # for the current tool + actual_call = tool_parser.streamed_args_for_tool[ + index] + + # check to see if there's anything left to stream + remaining_call = expected_call.replace( + actual_call, "", 1) + + # set that as a delta message + delta_message = DeltaMessage(tool_calls=[ + DeltaToolCall(index=index, + function=DeltaFunctionCall( + arguments=remaining_call). + model_dump(exclude_none=True)) + ]) + + # Count reasoning tokens for this choice at finish time. + if use_reasoning and all_previous_token_ids is not None: + r_parser = reasoning_parsers[i] + reasoning_token_counts[i] = \ + r_parser.count_reasoning_tokens( + all_previous_token_ids[i]) + + # Send the finish response for each request.n only once + choice_data = ChatCompletionResponseStreamChoice( + index=i, + delta=delta_message, + logprobs=logprobs, finish_reason=("tool_calls" if ( auto_tools_called or tool_choice_function_name) else output.finish_reason), - stop_reason=output.stop_reason) - chunk = ChatCompletionStreamResponse( - id=request_id, - object=chunk_object_type, - created=created_time, - choices=[choice_data], - model=model_name) - if (request.stream_options - and request.stream_options.include_usage): - if request.stream_options.continuous_usage_stats: - completion_tokens = len(output.token_ids) - usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=num_prompt_tokens + - completion_tokens, - ) - chunk.usage = usage - else: - chunk.usage = None - data = chunk.model_dump_json(exclude_unset=True) - yield f"data: {data}\n\n" - finish_reason_sent[i] = True - - # once the final token is handled, if stream_options.include_usage - # is sent, send the usage - if (request.stream_options - and request.stream_options.include_usage): + stop_reason=output.stop_reason) + chunk = ChatCompletionStreamResponse( + id=request_id, + object=chunk_object_type, + created=created_time, + choices=[choice_data], + model=model_name) + if (request.stream_options + and request.stream_options.include_usage): + if request.stream_options.continuous_usage_stats: + completion_tokens = len(output.token_ids) + usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=num_prompt_tokens + + completion_tokens, + ) + chunk.usage = usage + else: + chunk.usage = None + data = chunk.model_dump_json(exclude_unset=True) + yield f"data: {data}\n\n" + finish_reason_sent[i] = True + + # once the final token is handled, if stream_options.include_usage + # is sent, send the usage + if (request.stream_options + and request.stream_options.include_usage): completion_tokens = sum(previous_num_tokens) - total_reasoning = sum(reasoning_token_counts) if use_reasoning else None - final_usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=num_prompt_tokens + completion_tokens, - reasoning_tokens=total_reasoning, - prompt_tokens_details=( - PromptTokensDetails(cached_tokens=num_cached_tokens) - if num_cached_tokens is not None else None), - ) - - final_usage_chunk = ChatCompletionStreamResponse( - id=request_id, - object=chunk_object_type, - created=created_time, - choices=[], - model=model_name, - usage=final_usage) - final_usage_data = (final_usage_chunk.model_dump_json( - exclude_unset=True, exclude_none=True)) - yield f"data: {final_usage_data}\n\n" - - # report to FastAPI middleware aggregate usage across all choices - num_completion_tokens = sum(previous_num_tokens) - total_reasoning = sum(reasoning_token_counts) if use_reasoning else None - request_metadata.final_usage_info = UsageInfo( - prompt_tokens=num_prompt_tokens, - completion_tokens=num_completion_tokens, - total_tokens=num_prompt_tokens + num_completion_tokens, - reasoning_tokens=total_reasoning) - - except asyncio.CancelledError: - # Client disconnected via CancelledError path; abort engine request. - await self.engine_client.abort(request_id) - return - except ValueError as e: - # TODO: Use a vllm-specific Validation Error - logger.error("error in chat completion stream generator: %s", e) - data = self.create_streaming_error_response(str(e)) - yield f"data: {data}\n\n" - finally: - # Stop the disconnect watcher (it may already be done if it fired). - if _disconnect_watcher is not None and not _disconnect_watcher.done(): - _disconnect_watcher.cancel() - try: - await _disconnect_watcher - except asyncio.CancelledError: - pass - # Covers GeneratorExit when Starlette calls aclose() on disconnect - # during decode (tokens arrive fast so CancelledError path is not - # always triggered). abort() is a no-op for already-finished requests. - await self.engine_client.abort(request_id) - # Send the final done message after all response.n are finished - yield "data: [DONE]\n\n" - - async def chat_completion_full_generator( - self, - request: ChatCompletionRequest, - result_generator: AsyncIterator[RequestOutput], - request_id: str, - conversation: List[ConversationMessage], - tokenizer: AnyTokenizer, - request_metadata: RequestResponseMetadata, - raw_request: Optional[Request] = None, - ) -> Union[ErrorResponse, ChatCompletionResponse]: - - model_name = self.base_model_paths[0].name - created_time = int(time.time()) - final_res: Optional[RequestOutput] = None - - # Background watcher: same logic as the streaming path — polls - # is_disconnected() every 300 ms so that a client disconnect during - # non-streaming decode is caught even when uvicorn isn't actively - # reading the receive channel. - _disconnect_watcher: Optional[asyncio.Task] = None - if raw_request is not None: - async def _watch_disconnect() -> None: - try: - while True: - if await raw_request.is_disconnected(): - logger.info( - "Client disconnected (non-stream watcher), " - "aborting request %s", request_id) - await self.engine_client.abort(request_id) - return - await asyncio.sleep(0.3) - except asyncio.CancelledError: - pass - _disconnect_watcher = asyncio.ensure_future(_watch_disconnect()) - - try: - async for res in result_generator: - final_res = res - except asyncio.CancelledError: - await self.engine_client.abort(request_id) - return self.create_error_response("Client disconnected") - finally: - if _disconnect_watcher is not None and not _disconnect_watcher.done(): - _disconnect_watcher.cancel() - try: - await _disconnect_watcher - except asyncio.CancelledError: - pass - await self.engine_client.abort(request_id) - - assert final_res is not None - - choices: List[ChatCompletionResponseChoice] = [] - - role = self.get_chat_request_role(request) - for output in final_res.outputs: - token_ids = output.token_ids - out_logprobs = output.logprobs - - if request.logprobs and request.top_logprobs is not None: - assert out_logprobs is not None, "Did not output logprobs" - logprobs = self._create_chat_logprobs( - token_ids=token_ids, - top_logprobs=out_logprobs, - num_output_top_logprobs=request.top_logprobs, - tokenizer=tokenizer, - ) - else: - logprobs = None - - # In the OpenAI API the finish_reason is "tools_called" - # if the tool choice is auto and the model produced a tool - # call. The same is not true for named function calls - auto_tools_called = False - - # Extract reasoning content if parser is configured. - # output_text is what remains after stripping .... - reasoning_text: Optional[str] = None - output_text: str = output.text + total_reasoning = sum(reasoning_token_counts) if use_reasoning else None + final_usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=num_prompt_tokens + completion_tokens, + reasoning_tokens=total_reasoning, + prompt_tokens_details=( + PromptTokensDetails(cached_tokens=num_cached_tokens) + if num_cached_tokens is not None else None), + ) + + final_usage_chunk = ChatCompletionStreamResponse( + id=request_id, + object=chunk_object_type, + created=created_time, + choices=[], + model=model_name, + usage=final_usage) + final_usage_data = (final_usage_chunk.model_dump_json( + exclude_unset=True, exclude_none=True)) + yield f"data: {final_usage_data}\n\n" + + # report to FastAPI middleware aggregate usage across all choices + num_completion_tokens = sum(previous_num_tokens) + total_reasoning = sum(reasoning_token_counts) if use_reasoning else None + request_metadata.final_usage_info = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=num_completion_tokens, + total_tokens=num_prompt_tokens + num_completion_tokens, + reasoning_tokens=total_reasoning) + + except asyncio.CancelledError: + # Client disconnected via CancelledError path; abort engine request. + await self.engine_client.abort(request_id) + return + except ValueError as e: + # TODO: Use a vllm-specific Validation Error + logger.error("error in chat completion stream generator: %s", e) + data = self.create_streaming_error_response(str(e)) + yield f"data: {data}\n\n" + finally: + # Stop the disconnect watcher (it may already be done if it fired). + if _disconnect_watcher is not None and not _disconnect_watcher.done(): + _disconnect_watcher.cancel() + try: + await _disconnect_watcher + except asyncio.CancelledError: + pass + # Covers GeneratorExit when Starlette calls aclose() on disconnect + # during decode (tokens arrive fast so CancelledError path is not + # always triggered). abort() is a no-op for already-finished requests. + await self.engine_client.abort(request_id) + # Send the final done message after all response.n are finished + yield "data: [DONE]\n\n" + + async def chat_completion_full_generator( + self, + request: ChatCompletionRequest, + result_generator: AsyncIterator[RequestOutput], + request_id: str, + conversation: List[ConversationMessage], + tokenizer: AnyTokenizer, + request_metadata: RequestResponseMetadata, + raw_request: Optional[Request] = None, + ) -> Union[ErrorResponse, ChatCompletionResponse]: + + model_name = self.base_model_paths[0].name + created_time = int(time.time()) + final_res: Optional[RequestOutput] = None + + # Background watcher: same logic as the streaming path — polls + # is_disconnected() every 300 ms so that a client disconnect during + # non-streaming decode is caught even when uvicorn isn't actively + # reading the receive channel. + _disconnect_watcher: Optional[asyncio.Task] = None + if raw_request is not None: + async def _watch_disconnect() -> None: + try: + while True: + if await raw_request.is_disconnected(): + logger.info( + "Client disconnected (non-stream watcher), " + "aborting request %s", request_id) + await self.engine_client.abort(request_id) + return + await asyncio.sleep(0.3) + except asyncio.CancelledError: + pass + _disconnect_watcher = asyncio.ensure_future(_watch_disconnect()) + + try: + async for res in result_generator: + final_res = res + except asyncio.CancelledError: + await self.engine_client.abort(request_id) + return self.create_error_response("Client disconnected") + finally: + if _disconnect_watcher is not None and not _disconnect_watcher.done(): + _disconnect_watcher.cancel() + try: + await _disconnect_watcher + except asyncio.CancelledError: + pass + await self.engine_client.abort(request_id) + + assert final_res is not None + + choices: List[ChatCompletionResponseChoice] = [] + + role = self.get_chat_request_role(request) + for output in final_res.outputs: + token_ids = output.token_ids + out_logprobs = output.logprobs + + if request.logprobs and request.top_logprobs is not None: + assert out_logprobs is not None, "Did not output logprobs" + logprobs = self._create_chat_logprobs( + token_ids=token_ids, + top_logprobs=out_logprobs, + num_output_top_logprobs=request.top_logprobs, + tokenizer=tokenizer, + ) + else: + logprobs = None + + # In the OpenAI API the finish_reason is "tools_called" + # if the tool choice is auto and the model produced a tool + # call. The same is not true for named function calls + auto_tools_called = False + + # Extract reasoning content if parser is configured. + # output_text is what remains after stripping .... + reasoning_text: Optional[str] = None + output_text: str = output.text if self.reasoning_parser_cls: - r_parser = self.reasoning_parser_cls( - tokenizer, - chat_template_kwargs=request.chat_template_kwargs) + r_parser = self.reasoning_parser_cls( + tokenizer, + chat_template_kwargs=request.chat_template_kwargs) reasoning_text, extracted = r_parser.extract_reasoning( output.text, request) output_text = extracted or "" @@ -1095,20 +1095,20 @@ class OpenAIServingChat(OpenAIServing): reasoning_text, output_text = \ _reclassify_named_guided_json( reasoning_text, output_text) - + named_tool_called = False # if auto tools are not enabled, and a named tool choice using - # outlines is not being used - if (not self.enable_auto_tools - or not self.tool_parser) and not isinstance( - request.tool_choice, - ChatCompletionNamedToolChoiceParam): - message = ChatMessage(role=role, - reasoning_content=reasoning_text, - content=output_text) - - # if the request uses tools and specified a tool choice + # outlines is not being used + if (not self.enable_auto_tools + or not self.tool_parser) and not isinstance( + request.tool_choice, + ChatCompletionNamedToolChoiceParam): + message = ChatMessage(role=role, + reasoning_content=reasoning_text, + content=output_text) + + # if the request uses tools and specified a tool choice elif request.tool_choice and type( request.tool_choice) is ChatCompletionNamedToolChoiceParam: @@ -1132,109 +1132,109 @@ class OpenAIServingChat(OpenAIServing): parsed_named_tool_calls, ) message = ChatMessage( - role=role, - reasoning_content=reasoning_text, - content="", - tool_calls=[ - ToolCall(function=FunctionCall( - name=request.tool_choice.function.name, + role=role, + reasoning_content=reasoning_text, + content="", + tool_calls=[ + ToolCall(function=FunctionCall( + name=request.tool_choice.function.name, arguments=named_arguments)) - ]) - - # if the request doesn't use tool choice - # OR specifies to not use a tool - elif not request.tool_choice or request.tool_choice == "none": - - message = ChatMessage(role=role, - reasoning_content=reasoning_text, - content=output_text) - - # handle when there are tools and tool choice is auto - elif request.tools and ( - request.tool_choice == "auto" - or request.tool_choice is None) and self.enable_auto_tools \ - and self.tool_parser: - - try: - tool_parser = self.tool_parser(tokenizer) - except RuntimeError as e: - logger.error("Error in tool parser creation: %s", e) - return self.create_error_response(str(e)) - - # Parse tool calls from the post-reasoning content. - tool_call_info = tool_parser.extract_tool_calls( - output_text, request=request) - auto_tools_called = tool_call_info.tools_called - if tool_call_info.tools_called: - message = ChatMessage( - role=role, - reasoning_content=reasoning_text, - content=tool_call_info.content, - tool_calls=tool_call_info.tool_calls) - else: - message = ChatMessage(role=role, - reasoning_content=reasoning_text, - content=output_text) - - # undetermined case that is still important to handle - else: - logger.error( - "Error in chat_completion_full_generator - cannot determine" - " if tools should be extracted. Returning a standard chat " - "completion.") - message = ChatMessage(role=role, - reasoning_content=reasoning_text, - content=output_text) - - choice_data = ChatCompletionResponseChoice( - index=output.index, - message=message, - logprobs=logprobs, + ]) + + # if the request doesn't use tool choice + # OR specifies to not use a tool + elif not request.tool_choice or request.tool_choice == "none": + + message = ChatMessage(role=role, + reasoning_content=reasoning_text, + content=output_text) + + # handle when there are tools and tool choice is auto + elif request.tools and ( + request.tool_choice == "auto" + or request.tool_choice is None) and self.enable_auto_tools \ + and self.tool_parser: + + try: + tool_parser = self.tool_parser(tokenizer) + except RuntimeError as e: + logger.error("Error in tool parser creation: %s", e) + return self.create_error_response(str(e)) + + # Parse tool calls from the post-reasoning content. + tool_call_info = tool_parser.extract_tool_calls( + output_text, request=request) + auto_tools_called = tool_call_info.tools_called + if tool_call_info.tools_called: + message = ChatMessage( + role=role, + reasoning_content=reasoning_text, + content=tool_call_info.content, + tool_calls=tool_call_info.tool_calls) + else: + message = ChatMessage(role=role, + reasoning_content=reasoning_text, + content=output_text) + + # undetermined case that is still important to handle + else: + logger.error( + "Error in chat_completion_full_generator - cannot determine" + " if tools should be extracted. Returning a standard chat " + "completion.") + message = ChatMessage(role=role, + reasoning_content=reasoning_text, + content=output_text) + + choice_data = ChatCompletionResponseChoice( + index=output.index, + message=message, + logprobs=logprobs, finish_reason="tool_calls" if ( auto_tools_called or named_tool_called) else output.finish_reason if output.finish_reason else "stop", - stop_reason=output.stop_reason) - choices.append(choice_data) - - if request.echo or request.continue_final_message: - last_msg_content = "" - if conversation and "content" in conversation[-1] and conversation[ - -1].get("role") == role: - last_msg_content = conversation[-1]["content"] or "" - - for choice in choices: - full_message = last_msg_content + (choice.message.content - or "") - choice.message.content = full_message - - assert final_res.prompt_token_ids is not None - num_prompt_tokens = len(final_res.prompt_token_ids) - if final_res.encoder_prompt_token_ids is not None: - num_prompt_tokens += len(final_res.encoder_prompt_token_ids) - num_generated_tokens = sum( - len(output.token_ids) for output in final_res.outputs) - total_reasoning_tokens: Optional[int] = None - if self.reasoning_parser_cls: - rp = self.reasoning_parser_cls( - tokenizer, - chat_template_kwargs=request.chat_template_kwargs) - total_reasoning_tokens = sum( - rp.count_reasoning_tokens(list(output.token_ids)) - for output in final_res.outputs) - num_cached_tokens = (final_res.metrics.num_cached_tokens - if final_res.metrics is not None else None) - usage = UsageInfo( - prompt_tokens=num_prompt_tokens, - completion_tokens=num_generated_tokens, - total_tokens=num_prompt_tokens + num_generated_tokens, - reasoning_tokens=total_reasoning_tokens, - prompt_tokens_details=( - PromptTokensDetails(cached_tokens=num_cached_tokens) - if num_cached_tokens is not None else None), - ) - - request_metadata.final_usage_info = usage - + stop_reason=output.stop_reason) + choices.append(choice_data) + + if request.echo or request.continue_final_message: + last_msg_content = "" + if conversation and "content" in conversation[-1] and conversation[ + -1].get("role") == role: + last_msg_content = conversation[-1]["content"] or "" + + for choice in choices: + full_message = last_msg_content + (choice.message.content + or "") + choice.message.content = full_message + + assert final_res.prompt_token_ids is not None + num_prompt_tokens = len(final_res.prompt_token_ids) + if final_res.encoder_prompt_token_ids is not None: + num_prompt_tokens += len(final_res.encoder_prompt_token_ids) + num_generated_tokens = sum( + len(output.token_ids) for output in final_res.outputs) + total_reasoning_tokens: Optional[int] = None + if self.reasoning_parser_cls: + rp = self.reasoning_parser_cls( + tokenizer, + chat_template_kwargs=request.chat_template_kwargs) + total_reasoning_tokens = sum( + rp.count_reasoning_tokens(list(output.token_ids)) + for output in final_res.outputs) + num_cached_tokens = (final_res.metrics.num_cached_tokens + if final_res.metrics is not None else None) + usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=num_generated_tokens, + total_tokens=num_prompt_tokens + num_generated_tokens, + reasoning_tokens=total_reasoning_tokens, + prompt_tokens_details=( + PromptTokensDetails(cached_tokens=num_cached_tokens) + if num_cached_tokens is not None else None), + ) + + request_metadata.final_usage_info = usage + prompt_logprobs = final_res.prompt_logprobs sample_positions = request.bi100_prompt_logprobs_sample_positions if sample_positions is not None: @@ -1253,109 +1253,109 @@ class OpenAIServingChat(OpenAIServing): ] response = ChatCompletionResponse( - id=request_id, - created=created_time, - model=model_name, - choices=choices, - usage=usage, + id=request_id, + created=created_time, + model=model_name, + choices=choices, + usage=usage, prompt_logprobs=prompt_logprobs, - ) - - return response - - def _get_top_logprobs( - self, logprobs: Dict[int, Logprob], top_logprobs: Optional[int], - tokenizer: AnyTokenizer) -> List[ChatCompletionLogProb]: - return [ - ChatCompletionLogProb(token=(token := self._get_decoded_token( - p[1], - p[0], - tokenizer, - return_as_token_id=self.return_tokens_as_token_ids)), - logprob=max(p[1].logprob, -9999.0), - bytes=list( - token.encode("utf-8", errors="replace"))) - for i, p in enumerate(logprobs.items()) - if top_logprobs and i < top_logprobs - ] - - def _create_chat_logprobs( - self, - token_ids: GenericSequence[int], - top_logprobs: GenericSequence[Optional[Dict[int, Logprob]]], - tokenizer: AnyTokenizer, - num_output_top_logprobs: Optional[int] = None, - ) -> ChatCompletionLogProbs: - """Create OpenAI-style logprobs.""" - logprobs_content: List[ChatCompletionLogProbsContent] = [] - - for i, token_id in enumerate(token_ids): - step_top_logprobs = top_logprobs[i] - if step_top_logprobs is None: - token = tokenizer.decode(token_id) - if self.return_tokens_as_token_ids: - token = f"token_id:{token_id}" - - logprobs_content.append( - ChatCompletionLogProbsContent( - token=token, - bytes=list(token.encode("utf-8", errors="replace")), - )) - else: - step_token = step_top_logprobs[token_id] - step_decoded = step_token.decoded_token - - logprobs_content.append( - ChatCompletionLogProbsContent( - token=self._get_decoded_token( - step_token, - token_id, - tokenizer, - self.return_tokens_as_token_ids, - ), - logprob=max(step_token.logprob, -9999.0), - bytes=None if step_decoded is None else list( - step_decoded.encode("utf-8", errors="replace")), - top_logprobs=self._get_top_logprobs( - step_top_logprobs, - num_output_top_logprobs, - tokenizer, - ), - )) - - return ChatCompletionLogProbs(content=logprobs_content) - - def _should_stream_with_auto_tool_parsing(self, - request: ChatCompletionRequest): - """ - Utility function to check if streamed tokens should go through the tool - call parser that was configured. - - We only want to do this IF user-provided tools are set, a tool parser - is configured, "auto" tool choice is enabled, and the request's tool - choice field indicates that "auto" tool choice should be used. - """ - return (request.tools and self.tool_parser and self.enable_auto_tools - and request.tool_choice in ['auto', None]) - - def _should_check_for_unstreamed_tool_arg_tokens( - self, - delta_message: Optional[DeltaMessage], - output: CompletionOutput, - ) -> bool: - """ - Check to see if we should check for unstreamed tool arguments tokens. - This is only applicable when auto tool parsing is enabled, the delta - is a tool call with arguments. - """ - - # yapf: disable - return bool( - # if there is a delta message that includes tool calls which - # include a function that has arguments - output.finish_reason is not None - and self.enable_auto_tools and self.tool_parser and delta_message - and delta_message.tool_calls and delta_message.tool_calls[0] - and delta_message.tool_calls[0].function - and delta_message.tool_calls[0].function.arguments is not None - ) + ) + + return response + + def _get_top_logprobs( + self, logprobs: Dict[int, Logprob], top_logprobs: Optional[int], + tokenizer: AnyTokenizer) -> List[ChatCompletionLogProb]: + return [ + ChatCompletionLogProb(token=(token := self._get_decoded_token( + p[1], + p[0], + tokenizer, + return_as_token_id=self.return_tokens_as_token_ids)), + logprob=max(p[1].logprob, -9999.0), + bytes=list( + token.encode("utf-8", errors="replace"))) + for i, p in enumerate(logprobs.items()) + if top_logprobs and i < top_logprobs + ] + + def _create_chat_logprobs( + self, + token_ids: GenericSequence[int], + top_logprobs: GenericSequence[Optional[Dict[int, Logprob]]], + tokenizer: AnyTokenizer, + num_output_top_logprobs: Optional[int] = None, + ) -> ChatCompletionLogProbs: + """Create OpenAI-style logprobs.""" + logprobs_content: List[ChatCompletionLogProbsContent] = [] + + for i, token_id in enumerate(token_ids): + step_top_logprobs = top_logprobs[i] + if step_top_logprobs is None: + token = tokenizer.decode(token_id) + if self.return_tokens_as_token_ids: + token = f"token_id:{token_id}" + + logprobs_content.append( + ChatCompletionLogProbsContent( + token=token, + bytes=list(token.encode("utf-8", errors="replace")), + )) + else: + step_token = step_top_logprobs[token_id] + step_decoded = step_token.decoded_token + + logprobs_content.append( + ChatCompletionLogProbsContent( + token=self._get_decoded_token( + step_token, + token_id, + tokenizer, + self.return_tokens_as_token_ids, + ), + logprob=max(step_token.logprob, -9999.0), + bytes=None if step_decoded is None else list( + step_decoded.encode("utf-8", errors="replace")), + top_logprobs=self._get_top_logprobs( + step_top_logprobs, + num_output_top_logprobs, + tokenizer, + ), + )) + + return ChatCompletionLogProbs(content=logprobs_content) + + def _should_stream_with_auto_tool_parsing(self, + request: ChatCompletionRequest): + """ + Utility function to check if streamed tokens should go through the tool + call parser that was configured. + + We only want to do this IF user-provided tools are set, a tool parser + is configured, "auto" tool choice is enabled, and the request's tool + choice field indicates that "auto" tool choice should be used. + """ + return (request.tools and self.tool_parser and self.enable_auto_tools + and request.tool_choice in ['auto', None]) + + def _should_check_for_unstreamed_tool_arg_tokens( + self, + delta_message: Optional[DeltaMessage], + output: CompletionOutput, + ) -> bool: + """ + Check to see if we should check for unstreamed tool arguments tokens. + This is only applicable when auto tool parsing is enabled, the delta + is a tool call with arguments. + """ + + # yapf: disable + return bool( + # if there is a delta message that includes tool calls which + # include a function that has arguments + output.finish_reason is not None + and self.enable_auto_tools and self.tool_parser and delta_message + and delta_message.tool_calls and delta_message.tool_calls[0] + and delta_message.tool_calls[0].function + and delta_message.tool_calls[0].function.arguments is not None + ) diff --git a/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py b/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py index 40bfd231..b553dc74 100644 --- a/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py +++ b/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py @@ -1,43 +1,43 @@ -import math -from typing import List, Optional - -from vllm.core.block.common import BlockList +import math +from typing import List, Optional + +from vllm.core.block.common import BlockList from vllm.core.block.interfaces import Block, DeviceAwareBlockAllocator -from vllm.utils import Device, cdiv, chunk_list - - -class BlockTable: - """A class to manage blocks for a specific sequence. - - The BlockTable maps a sequence of tokens to a list of blocks, where each - block represents a contiguous memory allocation for a portion of the - sequence. The blocks are managed by a DeviceAwareBlockAllocator, which is - responsible for allocating and freeing memory for the blocks. - - Args: - block_size (int): The maximum number of tokens that can be stored in a - single block. - block_allocator (DeviceAwareBlockAllocator): The block allocator used to - manage memory for the blocks. - _blocks (Optional[List[Block]], optional): An optional list of existing - blocks to initialize the BlockTable with. If not provided, an empty - BlockTable is created. - max_block_sliding_window (Optional[int], optional): The number of - blocks to keep around for each sequance. If None, all blocks - are kept (eg., when sliding window is not used). - It should at least fit the sliding window size of the model. - - Attributes: - _block_size (int): The maximum number of tokens that can be stored in a - single block. - _allocator (DeviceAwareBlockAllocator): The block allocator used to - manage memory for the blocks. - _blocks (Optional[List[Block]]): The list of blocks managed by this - BlockTable. - _num_full_slots (int): The number of tokens currently stored in the - blocks. - """ - +from vllm.utils import Device, cdiv, chunk_list + + +class BlockTable: + """A class to manage blocks for a specific sequence. + + The BlockTable maps a sequence of tokens to a list of blocks, where each + block represents a contiguous memory allocation for a portion of the + sequence. The blocks are managed by a DeviceAwareBlockAllocator, which is + responsible for allocating and freeing memory for the blocks. + + Args: + block_size (int): The maximum number of tokens that can be stored in a + single block. + block_allocator (DeviceAwareBlockAllocator): The block allocator used to + manage memory for the blocks. + _blocks (Optional[List[Block]], optional): An optional list of existing + blocks to initialize the BlockTable with. If not provided, an empty + BlockTable is created. + max_block_sliding_window (Optional[int], optional): The number of + blocks to keep around for each sequance. If None, all blocks + are kept (eg., when sliding window is not used). + It should at least fit the sliding window size of the model. + + Attributes: + _block_size (int): The maximum number of tokens that can be stored in a + single block. + _allocator (DeviceAwareBlockAllocator): The block allocator used to + manage memory for the blocks. + _blocks (Optional[List[Block]]): The list of blocks managed by this + BlockTable. + _num_full_slots (int): The number of tokens currently stored in the + blocks. + """ + def __init__( self, block_size: int, @@ -52,55 +52,55 @@ class BlockTable: if _blocks is None: _blocks = [] self._blocks: BlockList = BlockList(_blocks) - - self._max_block_sliding_window = max_block_sliding_window - self._num_full_slots = self._get_num_token_ids() - - @staticmethod - def get_num_required_blocks(token_ids: List[int], - block_size: int, - num_lookahead_slots: int = 0) -> int: - """Calculates the minimum number of blocks required to store a given - sequence of token IDs along with any look-ahead slots that may be - required (like in multi-step + chunked-prefill). - - This assumes worst-case scenario, where every block requires a new - allocation (e.g. ignoring prefix caching). - - Args: - token_ids (List[int]): The sequence of token IDs to be stored. - block_size (int): The maximum number of tokens that can be stored in - a single block. - num_lookahead_slots (int): look-ahead slots that the sequence may - require. - - Returns: - int: The minimum number of blocks required to store the given - sequence of token IDs along with any required look-ahead slots. - """ - return cdiv(len(token_ids) + num_lookahead_slots, block_size) - + + self._max_block_sliding_window = max_block_sliding_window + self._num_full_slots = self._get_num_token_ids() + + @staticmethod + def get_num_required_blocks(token_ids: List[int], + block_size: int, + num_lookahead_slots: int = 0) -> int: + """Calculates the minimum number of blocks required to store a given + sequence of token IDs along with any look-ahead slots that may be + required (like in multi-step + chunked-prefill). + + This assumes worst-case scenario, where every block requires a new + allocation (e.g. ignoring prefix caching). + + Args: + token_ids (List[int]): The sequence of token IDs to be stored. + block_size (int): The maximum number of tokens that can be stored in + a single block. + num_lookahead_slots (int): look-ahead slots that the sequence may + require. + + Returns: + int: The minimum number of blocks required to store the given + sequence of token IDs along with any required look-ahead slots. + """ + return cdiv(len(token_ids) + num_lookahead_slots, block_size) + def allocate(self, token_ids: List[int], device: Device = Device.GPU) -> None: - """Allocates memory blocks for storing the given sequence of token IDs. - - This method allocates the required number of blocks to store the given - sequence of token IDs. - - Args: - token_ids (List[int]): The sequence of token IDs to be stored. - device (Device, optional): The device on which the blocks should be - allocated. Defaults to Device.GPU. - """ + """Allocates memory blocks for storing the given sequence of token IDs. + + This method allocates the required number of blocks to store the given + sequence of token IDs. + + Args: + token_ids (List[int]): The sequence of token IDs to be stored. + device (Device, optional): The device on which the blocks should be + allocated. Defaults to Device.GPU. + """ assert not self._is_allocated assert token_ids blocks = self._allocate_blocks_for_token_ids(prev_block=None, token_ids=token_ids, device=device) - self.update(blocks) - self._num_full_slots = len(token_ids) - + self.update(blocks) + self._num_full_slots = len(token_ids) + def update(self, blocks: List[Block]) -> None: """Resets the table to the newly provided blocks (with their corresponding block ids) @@ -115,106 +115,106 @@ class BlockTable: if block_hash is not None: content_hashes.append(block_hash) return content_hashes - - def append_token_ids(self, - token_ids: List[int], - num_lookahead_slots: int = 0, - num_computed_slots: Optional[int] = None) -> None: - """Appends a sequence of token IDs to the existing blocks in the - BlockTable. - - This method appends the given sequence of token IDs to the existing - blocks in the BlockTable. If there is not enough space in the existing - blocks, new blocks are allocated using the `ensure_num_empty_slots` - method to accommodate the additional tokens. - - The token IDs are divided into chunks of size `block_size` (except for - the first chunk, which may be smaller), and each chunk is appended to a - separate block. - - Args: - token_ids (List[int]): The sequence of token IDs to be appended. - num_computed_slots (Optional[int]): The number of KV cache slots - that are already filled (computed). - When sliding window is enabled, this is used to compute how many - blocks to drop at the front of the sequence. - Without sliding window, None can be passed. - Without chunked prefill, it should be the same as - _num_full_slots. - """ - assert self._is_allocated, "no blocks have been allocated" - assert len(self._blocks) > 0 - - # Drop blocks that are no longer needed due to sliding window - if self._max_block_sliding_window is not None: - null_block = self._allocator.allocate_or_get_null_block() - assert num_computed_slots is not None - end_block_idx = (num_computed_slots // - self._block_size) - self._max_block_sliding_window - for idx in range(0, end_block_idx): - b = self._blocks[idx] - if b is not null_block: - self._allocator.free(b) - self._blocks[idx] = null_block - - # Ensure there are enough empty slots for the new tokens plus - # lookahead slots - self.ensure_num_empty_slots(num_empty_slots=len(token_ids) + - num_lookahead_slots) - - # Update the blocks with the new tokens - first_block_idx = self._num_full_slots // self._block_size - token_blocks = self._chunk_token_blocks_for_append(token_ids) - - for i, token_block in enumerate(token_blocks): - self._blocks.append_token_ids(first_block_idx + i, token_block) - - self._num_full_slots += len(token_ids) - - def ensure_num_empty_slots(self, num_empty_slots: int) -> None: - """Ensures that the BlockTable has at least the specified number of - empty slots available. - - This method checks if the BlockTable has enough empty slots (i.e., - available space) to accommodate the requested number of tokens. If not, - it allocates additional blocks on the GPU to ensure that the required - number of empty slots is available. - - Args: - num_empty_slots (int): The minimum number of empty slots required. - """ - # Currently the block table only supports - # appending tokens to GPU blocks. - device = Device.GPU - assert self._is_allocated - - if self._num_empty_slots >= num_empty_slots: - return - - slots_to_allocate = num_empty_slots - self._num_empty_slots - blocks_to_allocate = cdiv(slots_to_allocate, self._block_size) - - for _ in range(blocks_to_allocate): - assert len(self._blocks) > 0 - self._blocks.append( - self._allocator.allocate_mutable_block( - prev_block=self._blocks[-1], device=device)) - - def fork(self) -> "BlockTable": - """Creates a new BlockTable instance with a copy of the blocks from the - current instance. - - This method creates a new BlockTable instance with the same block size, - block allocator, and a copy of the blocks from the current instance. The - new BlockTable has its own independent set of blocks, but shares the - same underlying memory allocation with the original BlockTable. - - Returns: - BlockTable: A new BlockTable instance with a copy of the blocks from - the current instance. - """ - assert self._is_allocated - assert len(self._blocks) > 0 + + def append_token_ids(self, + token_ids: List[int], + num_lookahead_slots: int = 0, + num_computed_slots: Optional[int] = None) -> None: + """Appends a sequence of token IDs to the existing blocks in the + BlockTable. + + This method appends the given sequence of token IDs to the existing + blocks in the BlockTable. If there is not enough space in the existing + blocks, new blocks are allocated using the `ensure_num_empty_slots` + method to accommodate the additional tokens. + + The token IDs are divided into chunks of size `block_size` (except for + the first chunk, which may be smaller), and each chunk is appended to a + separate block. + + Args: + token_ids (List[int]): The sequence of token IDs to be appended. + num_computed_slots (Optional[int]): The number of KV cache slots + that are already filled (computed). + When sliding window is enabled, this is used to compute how many + blocks to drop at the front of the sequence. + Without sliding window, None can be passed. + Without chunked prefill, it should be the same as + _num_full_slots. + """ + assert self._is_allocated, "no blocks have been allocated" + assert len(self._blocks) > 0 + + # Drop blocks that are no longer needed due to sliding window + if self._max_block_sliding_window is not None: + null_block = self._allocator.allocate_or_get_null_block() + assert num_computed_slots is not None + end_block_idx = (num_computed_slots // + self._block_size) - self._max_block_sliding_window + for idx in range(0, end_block_idx): + b = self._blocks[idx] + if b is not null_block: + self._allocator.free(b) + self._blocks[idx] = null_block + + # Ensure there are enough empty slots for the new tokens plus + # lookahead slots + self.ensure_num_empty_slots(num_empty_slots=len(token_ids) + + num_lookahead_slots) + + # Update the blocks with the new tokens + first_block_idx = self._num_full_slots // self._block_size + token_blocks = self._chunk_token_blocks_for_append(token_ids) + + for i, token_block in enumerate(token_blocks): + self._blocks.append_token_ids(first_block_idx + i, token_block) + + self._num_full_slots += len(token_ids) + + def ensure_num_empty_slots(self, num_empty_slots: int) -> None: + """Ensures that the BlockTable has at least the specified number of + empty slots available. + + This method checks if the BlockTable has enough empty slots (i.e., + available space) to accommodate the requested number of tokens. If not, + it allocates additional blocks on the GPU to ensure that the required + number of empty slots is available. + + Args: + num_empty_slots (int): The minimum number of empty slots required. + """ + # Currently the block table only supports + # appending tokens to GPU blocks. + device = Device.GPU + assert self._is_allocated + + if self._num_empty_slots >= num_empty_slots: + return + + slots_to_allocate = num_empty_slots - self._num_empty_slots + blocks_to_allocate = cdiv(slots_to_allocate, self._block_size) + + for _ in range(blocks_to_allocate): + assert len(self._blocks) > 0 + self._blocks.append( + self._allocator.allocate_mutable_block( + prev_block=self._blocks[-1], device=device)) + + def fork(self) -> "BlockTable": + """Creates a new BlockTable instance with a copy of the blocks from the + current instance. + + This method creates a new BlockTable instance with the same block size, + block allocator, and a copy of the blocks from the current instance. The + new BlockTable has its own independent set of blocks, but shares the + same underlying memory allocation with the original BlockTable. + + Returns: + BlockTable: A new BlockTable instance with a copy of the blocks from + the current instance. + """ + assert self._is_allocated + assert len(self._blocks) > 0 forked_blocks = self._allocator.fork(self._blocks[-1]) return BlockTable( block_size=self._block_size, @@ -223,84 +223,84 @@ class BlockTable: max_block_sliding_window=self._max_block_sliding_window, cache_namespace=self._cache_namespace, ) - - def free(self) -> None: - """Frees the memory occupied by the blocks in the BlockTable. - - This method iterates over all the blocks in the `_blocks` list and calls - the `free` method of the `_allocator` object to release the memory - occupied by each block. After freeing all the blocks, the `_blocks` list - is set to `None`. - """ - for block in self.blocks: - self._allocator.free(block) - self._blocks.reset() - - @property - def physical_block_ids(self) -> List[int]: - """Returns a list of physical block indices for the blocks in the - BlockTable. - - This property returns a list of integers, where each integer represents - the physical block index of a corresponding block in the `_blocks` list. - The physical block index is a unique identifier for the memory location - occupied by the block. - - Returns: - List[int]: A list of physical block indices for the blocks in the - BlockTable. - """ - return self._blocks.ids() - - def get_unseen_token_ids(self, sequence_token_ids: List[int]) -> List[int]: - """Get the number of "unseen" tokens in the sequence. - - Unseen tokens are tokens in the sequence corresponding to this block - table, but are not yet appended to this block table. - - Args: - sequence_token_ids (List[int]): The list of token ids in the - sequence. - - Returns: - List[int]: The postfix of sequence_token_ids that has not yet been - appended to the block table. - """ - - # Since the block table is append-only, the unseen token ids are the - # ones after the appended ones. - return sequence_token_ids[self.num_full_slots:] - + + def free(self) -> None: + """Frees the memory occupied by the blocks in the BlockTable. + + This method iterates over all the blocks in the `_blocks` list and calls + the `free` method of the `_allocator` object to release the memory + occupied by each block. After freeing all the blocks, the `_blocks` list + is set to `None`. + """ + for block in self.blocks: + self._allocator.free(block) + self._blocks.reset() + + @property + def physical_block_ids(self) -> List[int]: + """Returns a list of physical block indices for the blocks in the + BlockTable. + + This property returns a list of integers, where each integer represents + the physical block index of a corresponding block in the `_blocks` list. + The physical block index is a unique identifier for the memory location + occupied by the block. + + Returns: + List[int]: A list of physical block indices for the blocks in the + BlockTable. + """ + return self._blocks.ids() + + def get_unseen_token_ids(self, sequence_token_ids: List[int]) -> List[int]: + """Get the number of "unseen" tokens in the sequence. + + Unseen tokens are tokens in the sequence corresponding to this block + table, but are not yet appended to this block table. + + Args: + sequence_token_ids (List[int]): The list of token ids in the + sequence. + + Returns: + List[int]: The postfix of sequence_token_ids that has not yet been + appended to the block table. + """ + + # Since the block table is append-only, the unseen token ids are the + # ones after the appended ones. + return sequence_token_ids[self.num_full_slots:] + def _allocate_blocks_for_token_ids(self, prev_block: Optional[Block], token_ids: List[int], device: Device) -> List[Block]: blocks: List[Block] = [] block_token_ids = [] - tail_token_ids = [] - for cur_token_ids in chunk_list(token_ids, self._block_size): - if len(cur_token_ids) == self._block_size: - block_token_ids.append(cur_token_ids) - else: - tail_token_ids.append(cur_token_ids) - + tail_token_ids = [] + for cur_token_ids in chunk_list(token_ids, self._block_size): + if len(cur_token_ids) == self._block_size: + block_token_ids.append(cur_token_ids) + else: + tail_token_ids.append(cur_token_ids) + if block_token_ids: blocks.extend(self._allocate_immutable_blocks( prev_block=prev_block, block_token_ids=block_token_ids, device=device)) prev_block = blocks[-1] - - if tail_token_ids: - assert len(tail_token_ids) == 1 - cur_token_ids = tail_token_ids[0] - + + if tail_token_ids: + assert len(tail_token_ids) == 1 + cur_token_ids = tail_token_ids[0] + block = self._allocate_mutable_block(prev_block=prev_block, device=device) - block.append_token_ids(cur_token_ids) - - blocks.append(block) - + block.append_token_ids(cur_token_ids) + + blocks.append(block) + return blocks def _allocate_mutable_block(self, prev_block: Optional[Block], @@ -372,85 +372,85 @@ class BlockTable: prev_block, block_token_ids=block_token_ids, device=device) - - def _get_all_token_ids(self) -> List[int]: - # NOTE: This function is O(seq_len); use sparingly. - token_ids: List[int] = [] - - if not self._is_allocated: - return token_ids - - for block in self.blocks: - token_ids.extend(block.token_ids) - - return token_ids - - def _get_num_token_ids(self) -> int: - res = 0 - for block in self.blocks: - res += len(block.token_ids) - - return res - - @property - def _is_allocated(self) -> bool: - return len(self._blocks) > 0 - - @property - def blocks(self) -> List[Block]: - return self._blocks.list() - - @property - def _num_empty_slots(self) -> int: - assert self._is_allocated - return len(self._blocks) * self._block_size - self._num_full_slots - - @property - def num_full_slots(self) -> int: - """Returns the total number of tokens currently stored in the - BlockTable. - - Returns: - int: The total number of tokens currently stored in the BlockTable. - """ - return self._num_full_slots - - def get_num_blocks_touched_by_append_slots( - self, token_ids: List[int], num_lookahead_slots: int) -> int: - """Determine how many blocks will be "touched" by appending the token - ids. - - This is required for the scheduler to determine whether a sequence can - continue generation, or if it must be preempted. - """ - # Math below is equivalent to: - # all_token_ids = token_ids + [-1] * num_lookahead_slots - # token_blocks = self._chunk_token_blocks_for_append(all_token_ids) - # return len(token_blocks) - - num_token_ids = len(token_ids) + num_lookahead_slots - first_chunk_size = self._block_size - (self._num_full_slots % - self._block_size) - num_token_blocks = (1 + math.ceil( - (num_token_ids - first_chunk_size) / self._block_size)) - return num_token_blocks - - def _chunk_token_blocks_for_append( - self, token_ids: List[int]) -> List[List[int]]: - """Split the token ids into block-sized chunks so they can be easily - appended to blocks. The first such "token block" may have less token ids - than the block size, since the last allocated block may be partially - full. - - If no token ids are provided, then no chunks are returned. - """ - - if not token_ids: - return [] - - first_chunk_size = self._block_size - (self._num_full_slots % - self._block_size) - token_blocks = [token_ids[:first_chunk_size]] - token_blocks.extend( - chunk_list(token_ids[first_chunk_size:], self._block_size)) - return token_blocks + + def _get_all_token_ids(self) -> List[int]: + # NOTE: This function is O(seq_len); use sparingly. + token_ids: List[int] = [] + + if not self._is_allocated: + return token_ids + + for block in self.blocks: + token_ids.extend(block.token_ids) + + return token_ids + + def _get_num_token_ids(self) -> int: + res = 0 + for block in self.blocks: + res += len(block.token_ids) + + return res + + @property + def _is_allocated(self) -> bool: + return len(self._blocks) > 0 + + @property + def blocks(self) -> List[Block]: + return self._blocks.list() + + @property + def _num_empty_slots(self) -> int: + assert self._is_allocated + return len(self._blocks) * self._block_size - self._num_full_slots + + @property + def num_full_slots(self) -> int: + """Returns the total number of tokens currently stored in the + BlockTable. + + Returns: + int: The total number of tokens currently stored in the BlockTable. + """ + return self._num_full_slots + + def get_num_blocks_touched_by_append_slots( + self, token_ids: List[int], num_lookahead_slots: int) -> int: + """Determine how many blocks will be "touched" by appending the token + ids. + + This is required for the scheduler to determine whether a sequence can + continue generation, or if it must be preempted. + """ + # Math below is equivalent to: + # all_token_ids = token_ids + [-1] * num_lookahead_slots + # token_blocks = self._chunk_token_blocks_for_append(all_token_ids) + # return len(token_blocks) + + num_token_ids = len(token_ids) + num_lookahead_slots + first_chunk_size = self._block_size - (self._num_full_slots % + self._block_size) + num_token_blocks = (1 + math.ceil( + (num_token_ids - first_chunk_size) / self._block_size)) + return num_token_blocks + + def _chunk_token_blocks_for_append( + self, token_ids: List[int]) -> List[List[int]]: + """Split the token ids into block-sized chunks so they can be easily + appended to blocks. The first such "token block" may have less token ids + than the block size, since the last allocated block may be partially + full. + + If no token ids are provided, then no chunks are returned. + """ + + if not token_ids: + return [] + + first_chunk_size = self._block_size - (self._num_full_slots % + self._block_size) + token_blocks = [token_ids[:first_chunk_size]] + token_blocks.extend( + chunk_list(token_ids[first_chunk_size:], self._block_size)) + return token_blocks diff --git a/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py b/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py index 9628289c..09a4bdd2 100644 --- a/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py +++ b/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py @@ -4,56 +4,56 @@ from vllm.core.block.cpu_kv_content_cache import (CpuKvContentCache, cpu_kv_offload_enabled) from vllm.core.block.interfaces import (Block, BlockAllocator, BlockId, DeviceAwareBlockAllocator) -from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator -from vllm.core.block.prefix_caching_block import PrefixCachingBlockAllocator -from vllm.utils import Device - - -class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): - """A block allocator that can allocate blocks on both CPU and GPU memory. - - This class implements the `DeviceAwareBlockAllocator` interface and provides - functionality for allocating and managing blocks of memory on both CPU and - GPU devices. - - The `CpuGpuBlockAllocator` maintains separate memory pools for CPU and GPU - blocks, and allows for allocation, deallocation, forking, and swapping of - blocks across these memory pools. - """ - - @staticmethod - def create( - allocator_type: str, - num_gpu_blocks: int, - num_cpu_blocks: int, - block_size: int, - ) -> DeviceAwareBlockAllocator: - """Creates a CpuGpuBlockAllocator instance with the specified - configuration. - - This static method creates and returns a CpuGpuBlockAllocator instance - based on the provided parameters. It initializes the CPU and GPU block - allocators with the specified number of blocks, block size, and - allocator type. - - Args: - allocator_type (str): The type of block allocator to use for CPU - and GPU blocks. Currently supported values are "naive" and - "prefix_caching". - num_gpu_blocks (int): The number of blocks to allocate for GPU - memory. - num_cpu_blocks (int): The number of blocks to allocate for CPU - memory. - block_size (int): The size of each block in number of tokens. - - Returns: - DeviceAwareBlockAllocator: A CpuGpuBlockAllocator instance with the - specified configuration. - - Notes: - - The block IDs are assigned contiguously, with GPU block IDs coming - before CPU block IDs. - """ +from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator +from vllm.core.block.prefix_caching_block import PrefixCachingBlockAllocator +from vllm.utils import Device + + +class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): + """A block allocator that can allocate blocks on both CPU and GPU memory. + + This class implements the `DeviceAwareBlockAllocator` interface and provides + functionality for allocating and managing blocks of memory on both CPU and + GPU devices. + + The `CpuGpuBlockAllocator` maintains separate memory pools for CPU and GPU + blocks, and allows for allocation, deallocation, forking, and swapping of + blocks across these memory pools. + """ + + @staticmethod + def create( + allocator_type: str, + num_gpu_blocks: int, + num_cpu_blocks: int, + block_size: int, + ) -> DeviceAwareBlockAllocator: + """Creates a CpuGpuBlockAllocator instance with the specified + configuration. + + This static method creates and returns a CpuGpuBlockAllocator instance + based on the provided parameters. It initializes the CPU and GPU block + allocators with the specified number of blocks, block size, and + allocator type. + + Args: + allocator_type (str): The type of block allocator to use for CPU + and GPU blocks. Currently supported values are "naive" and + "prefix_caching". + num_gpu_blocks (int): The number of blocks to allocate for GPU + memory. + num_cpu_blocks (int): The number of blocks to allocate for CPU + memory. + block_size (int): The size of each block in number of tokens. + + Returns: + DeviceAwareBlockAllocator: A CpuGpuBlockAllocator instance with the + specified configuration. + + Notes: + - The block IDs are assigned contiguously, with GPU block IDs coming + before CPU block IDs. + """ content_offload = cpu_kv_offload_enabled() if content_offload and allocator_type != "prefix_caching": raise RuntimeError( @@ -63,38 +63,38 @@ class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): "BI100_CPU_KV_OFFLOAD=1 requires at least one CPU KV block") block_ids = list(range(num_gpu_blocks + num_cpu_blocks)) - gpu_block_ids = block_ids[:num_gpu_blocks] - cpu_block_ids = block_ids[num_gpu_blocks:] - - if allocator_type == "naive": - gpu_allocator: BlockAllocator = NaiveBlockAllocator( - create_block=NaiveBlock, # type: ignore - num_blocks=num_gpu_blocks, - block_size=block_size, - block_ids=gpu_block_ids, - ) - - cpu_allocator: BlockAllocator = NaiveBlockAllocator( - create_block=NaiveBlock, # type: ignore - num_blocks=num_cpu_blocks, - block_size=block_size, - block_ids=cpu_block_ids, - ) - elif allocator_type == "prefix_caching": - gpu_allocator = PrefixCachingBlockAllocator( - num_blocks=num_gpu_blocks, - block_size=block_size, - block_ids=gpu_block_ids, - ) - - cpu_allocator = PrefixCachingBlockAllocator( - num_blocks=num_cpu_blocks, - block_size=block_size, - block_ids=cpu_block_ids, - ) - else: - raise ValueError(f"Unknown allocator type {allocator_type=}") - + gpu_block_ids = block_ids[:num_gpu_blocks] + cpu_block_ids = block_ids[num_gpu_blocks:] + + if allocator_type == "naive": + gpu_allocator: BlockAllocator = NaiveBlockAllocator( + create_block=NaiveBlock, # type: ignore + num_blocks=num_gpu_blocks, + block_size=block_size, + block_ids=gpu_block_ids, + ) + + cpu_allocator: BlockAllocator = NaiveBlockAllocator( + create_block=NaiveBlock, # type: ignore + num_blocks=num_cpu_blocks, + block_size=block_size, + block_ids=cpu_block_ids, + ) + elif allocator_type == "prefix_caching": + gpu_allocator = PrefixCachingBlockAllocator( + num_blocks=num_gpu_blocks, + block_size=block_size, + block_ids=gpu_block_ids, + ) + + cpu_allocator = PrefixCachingBlockAllocator( + num_blocks=num_cpu_blocks, + block_size=block_size, + block_ids=cpu_block_ids, + ) + else: + raise ValueError(f"Unknown allocator type {allocator_type=}") + return CpuGpuBlockAllocator( cpu_block_allocator=cpu_allocator, gpu_block_allocator=gpu_allocator, @@ -105,21 +105,21 @@ class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): def __init__(self, cpu_block_allocator: BlockAllocator, gpu_block_allocator: BlockAllocator, cpu_content_cache: Optional[CpuKvContentCache] = None): - assert not ( - cpu_block_allocator.all_block_ids - & gpu_block_allocator.all_block_ids - ), "cpu and gpu block allocators can't have intersection of block ids" - - self._allocators = { - Device.CPU: cpu_block_allocator, - Device.GPU: gpu_block_allocator, - } - + assert not ( + cpu_block_allocator.all_block_ids + & gpu_block_allocator.all_block_ids + ), "cpu and gpu block allocators can't have intersection of block ids" + + self._allocators = { + Device.CPU: cpu_block_allocator, + Device.GPU: gpu_block_allocator, + } + self._swap_mapping: Dict[int, int] = {} self._null_block: Optional[Block] = None self._cpu_content_cache = cpu_content_cache - - self._block_ids_to_allocator: Dict[int, BlockAllocator] = {} + + self._block_ids_to_allocator: Dict[int, BlockAllocator] = {} for _, allocator in self._allocators.items(): for block_id in allocator.all_block_ids: self._block_ids_to_allocator[block_id] = allocator @@ -164,236 +164,236 @@ class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): assert self._cpu_content_cache is not None gpu_slot = self.get_physical_block_id(Device.GPU, gpu_block_id) return self._cpu_content_cache.stage_store(content_hash, gpu_slot) - - def allocate_or_get_null_block(self) -> Block: - if self._null_block is None: - self._null_block = NullBlock( - self.allocate_mutable_block(None, Device.GPU)) - return self._null_block - - def allocate_mutable_block(self, prev_block: Optional[Block], - device: Device) -> Block: - """Allocates a new mutable block on the specified device. - - Args: - prev_block (Optional[Block]): The previous block to in the sequence. - Used for prefix hashing. - device (Device): The device on which to allocate the new block. - - Returns: - Block: The newly allocated mutable block. - """ - return self._allocators[device].allocate_mutable_block(prev_block) - - def allocate_immutable_blocks(self, prev_block: Optional[Block], - block_token_ids: List[List[int]], - device: Device) -> List[Block]: - """Allocates a new group of immutable blocks with the provided block - token IDs on the specified device. - - Args: - prev_block (Optional[Block]): The previous block in the sequence. - Used for prefix hashing. - block_token_ids (List[int]): The list of block token IDs to be - stored in the new blocks. - device (Device): The device on which to allocate the new block. - - Returns: - List[Block]: The newly allocated list of immutable blocks - containing the provided block token IDs. - """ - return self._allocators[device].allocate_immutable_blocks( - prev_block, block_token_ids) - - def allocate_immutable_block(self, prev_block: Optional[Block], - token_ids: List[int], - device: Device) -> Block: - """Allocates a new immutable block with the provided token IDs on the - specified device. - - Args: - prev_block (Optional[Block]): The previous block in the sequence. - Used for prefix hashing. - token_ids (List[int]): The list of token IDs to be stored in the new - block. - device (Device): The device on which to allocate the new block. - - Returns: - Block: The newly allocated immutable block containing the provided - token IDs. - """ - return self._allocators[device].allocate_immutable_block( - prev_block, token_ids) - - def free(self, block: Block) -> None: - """Frees the memory occupied by the given block. - - Args: - block (Block): The block to be freed. - """ - # Null block should never be freed - if isinstance(block, NullBlock): - return - block_id = block.block_id - assert block_id is not None - allocator = self._block_ids_to_allocator[block_id] - allocator.free(block) - - def fork(self, last_block: Block) -> List[Block]: - """Creates a new sequence of blocks that shares the same underlying - memory as the original sequence. - - Args: - last_block (Block): The last block in the original sequence. - - Returns: - List[Block]: A new list of blocks that shares the same memory as the - original sequence. - """ - # do not attempt to fork the null block - assert not isinstance(last_block, NullBlock) - block_id = last_block.block_id - assert block_id is not None - allocator = self._block_ids_to_allocator[block_id] - return allocator.fork(last_block) - - def get_num_free_blocks(self, device: Device) -> int: - """Returns the number of free blocks available on the specified device. - - Args: - device (Device): The device for which to query the number of free - blocks. AssertionError is raised if None is passed. - - Returns: - int: The number of free blocks available on the specified device. - """ - return self._allocators[device].get_num_free_blocks() - - def get_num_total_blocks(self, device: Device) -> int: - return self._allocators[device].get_num_total_blocks() - - def get_physical_block_id(self, device: Device, absolute_id: int) -> int: - """Returns the zero-offset block id on certain device given the - absolute block id. - - Args: - device (Device): The device for which to query relative block id. - absolute_id (int): The absolute block id for the block in - whole allocator. - - Returns: - int: The zero-offset block id on certain device. - """ - return self._allocators[device].get_physical_block_id(absolute_id) - + + def allocate_or_get_null_block(self) -> Block: + if self._null_block is None: + self._null_block = NullBlock( + self.allocate_mutable_block(None, Device.GPU)) + return self._null_block + + def allocate_mutable_block(self, prev_block: Optional[Block], + device: Device) -> Block: + """Allocates a new mutable block on the specified device. + + Args: + prev_block (Optional[Block]): The previous block to in the sequence. + Used for prefix hashing. + device (Device): The device on which to allocate the new block. + + Returns: + Block: The newly allocated mutable block. + """ + return self._allocators[device].allocate_mutable_block(prev_block) + + def allocate_immutable_blocks(self, prev_block: Optional[Block], + block_token_ids: List[List[int]], + device: Device) -> List[Block]: + """Allocates a new group of immutable blocks with the provided block + token IDs on the specified device. + + Args: + prev_block (Optional[Block]): The previous block in the sequence. + Used for prefix hashing. + block_token_ids (List[int]): The list of block token IDs to be + stored in the new blocks. + device (Device): The device on which to allocate the new block. + + Returns: + List[Block]: The newly allocated list of immutable blocks + containing the provided block token IDs. + """ + return self._allocators[device].allocate_immutable_blocks( + prev_block, block_token_ids) + + def allocate_immutable_block(self, prev_block: Optional[Block], + token_ids: List[int], + device: Device) -> Block: + """Allocates a new immutable block with the provided token IDs on the + specified device. + + Args: + prev_block (Optional[Block]): The previous block in the sequence. + Used for prefix hashing. + token_ids (List[int]): The list of token IDs to be stored in the new + block. + device (Device): The device on which to allocate the new block. + + Returns: + Block: The newly allocated immutable block containing the provided + token IDs. + """ + return self._allocators[device].allocate_immutable_block( + prev_block, token_ids) + + def free(self, block: Block) -> None: + """Frees the memory occupied by the given block. + + Args: + block (Block): The block to be freed. + """ + # Null block should never be freed + if isinstance(block, NullBlock): + return + block_id = block.block_id + assert block_id is not None + allocator = self._block_ids_to_allocator[block_id] + allocator.free(block) + + def fork(self, last_block: Block) -> List[Block]: + """Creates a new sequence of blocks that shares the same underlying + memory as the original sequence. + + Args: + last_block (Block): The last block in the original sequence. + + Returns: + List[Block]: A new list of blocks that shares the same memory as the + original sequence. + """ + # do not attempt to fork the null block + assert not isinstance(last_block, NullBlock) + block_id = last_block.block_id + assert block_id is not None + allocator = self._block_ids_to_allocator[block_id] + return allocator.fork(last_block) + + def get_num_free_blocks(self, device: Device) -> int: + """Returns the number of free blocks available on the specified device. + + Args: + device (Device): The device for which to query the number of free + blocks. AssertionError is raised if None is passed. + + Returns: + int: The number of free blocks available on the specified device. + """ + return self._allocators[device].get_num_free_blocks() + + def get_num_total_blocks(self, device: Device) -> int: + return self._allocators[device].get_num_total_blocks() + + def get_physical_block_id(self, device: Device, absolute_id: int) -> int: + """Returns the zero-offset block id on certain device given the + absolute block id. + + Args: + device (Device): The device for which to query relative block id. + absolute_id (int): The absolute block id for the block in + whole allocator. + + Returns: + int: The zero-offset block id on certain device. + """ + return self._allocators[device].get_physical_block_id(absolute_id) + def swap(self, blocks: List[Block], src_device: Device, dst_device: Device) -> Dict[int, int]: - """Execute the swap for the given blocks from source_device - on to dest_device, save the current swap mapping and append - them to the accumulated `self._swap_mapping` for each - scheduling move. - - Args: - blocks: List of blocks to be swapped. - src_device (Device): Device to swap the 'blocks' from. - dst_device (Device): Device to swap the 'blocks' to. - - Returns: - Dict[int, int]: Swap mapping from source_device - on to dest_device. - """ + """Execute the swap for the given blocks from source_device + on to dest_device, save the current swap mapping and append + them to the accumulated `self._swap_mapping` for each + scheduling move. + + Args: + blocks: List of blocks to be swapped. + src_device (Device): Device to swap the 'blocks' from. + dst_device (Device): Device to swap the 'blocks' to. + + Returns: + Dict[int, int]: Swap mapping from source_device + on to dest_device. + """ if self.content_offload_enabled: raise RuntimeError( "request-level preemption swap cannot share CPU slots with " "BI100_CPU_KV_OFFLOAD") src_block_ids = [block.block_id for block in blocks] - self._allocators[src_device].swap_out(blocks) - self._allocators[dst_device].swap_in(blocks) - dst_block_ids = [block.block_id for block in blocks] - - current_swap_mapping: Dict[int, int] = {} - for src_block_id, dst_block_id in zip(src_block_ids, dst_block_ids): - if src_block_id is not None and dst_block_id is not None: - self._swap_mapping[src_block_id] = dst_block_id - current_swap_mapping[src_block_id] = dst_block_id - return current_swap_mapping - - def get_num_full_blocks_touched(self, blocks: List[Block], - device: Device) -> int: - """Returns the number of full blocks that will be touched by - swapping in/out the given blocks on to the 'device'. - - Args: - blocks: List of blocks to be swapped. - device (Device): Device to swap the 'blocks' on. - - Returns: - int: the number of full blocks that will be touched by - swapping in/out the given blocks on to the 'device'. - Non full blocks are ignored when deciding the number - of blocks to touch. - """ - return self._allocators[device].get_num_full_blocks_touched(blocks) - - def clear_copy_on_writes(self) -> List[Tuple[int, int]]: - """Clears the copy-on-write (CoW) state and returns the mapping of - source to destination block IDs. - - Returns: - List[Tuple[int, int]]: A list mapping source block IDs to - destination block IDs. - """ - # CoW only supported on GPU - device = Device.GPU - return self._allocators[device].clear_copy_on_writes() - - def mark_blocks_as_accessed(self, block_ids: List[int], - now: float) -> None: - """Mark blocks as accessed, only use for prefix caching.""" - # Prefix caching only supported on GPU. - device = Device.GPU - return self._allocators[device].mark_blocks_as_accessed(block_ids, now) - - def mark_blocks_as_computed(self, block_ids: List[int]) -> None: - """Mark blocks as accessed, only use for prefix caching.""" - # Prefix caching only supported on GPU. - device = Device.GPU - return self._allocators[device].mark_blocks_as_computed(block_ids) - - def get_computed_block_ids(self, prev_computed_block_ids: List[int], - block_ids: List[int], - skip_last_block_id: bool) -> List[int]: - # Prefix caching only supported on GPU. - device = Device.GPU - return self._allocators[device].get_computed_block_ids( - prev_computed_block_ids, block_ids, skip_last_block_id) - - def get_common_computed_block_ids( - self, computed_seq_block_ids: List[List[int]]) -> List[int]: - # Prefix caching only supported on GPU. - device = Device.GPU - return self._allocators[device].get_common_computed_block_ids( - computed_seq_block_ids) - - @property - def all_block_ids(self) -> FrozenSet[int]: - return frozenset(self._block_ids_to_allocator.keys()) - - def get_prefix_cache_hit_rate(self, device: Device) -> float: - """Prefix cache hit rate. -1 means not supported or disabled.""" - assert device in self._allocators - return self._allocators[device].get_prefix_cache_hit_rate() - + self._allocators[src_device].swap_out(blocks) + self._allocators[dst_device].swap_in(blocks) + dst_block_ids = [block.block_id for block in blocks] + + current_swap_mapping: Dict[int, int] = {} + for src_block_id, dst_block_id in zip(src_block_ids, dst_block_ids): + if src_block_id is not None and dst_block_id is not None: + self._swap_mapping[src_block_id] = dst_block_id + current_swap_mapping[src_block_id] = dst_block_id + return current_swap_mapping + + def get_num_full_blocks_touched(self, blocks: List[Block], + device: Device) -> int: + """Returns the number of full blocks that will be touched by + swapping in/out the given blocks on to the 'device'. + + Args: + blocks: List of blocks to be swapped. + device (Device): Device to swap the 'blocks' on. + + Returns: + int: the number of full blocks that will be touched by + swapping in/out the given blocks on to the 'device'. + Non full blocks are ignored when deciding the number + of blocks to touch. + """ + return self._allocators[device].get_num_full_blocks_touched(blocks) + + def clear_copy_on_writes(self) -> List[Tuple[int, int]]: + """Clears the copy-on-write (CoW) state and returns the mapping of + source to destination block IDs. + + Returns: + List[Tuple[int, int]]: A list mapping source block IDs to + destination block IDs. + """ + # CoW only supported on GPU + device = Device.GPU + return self._allocators[device].clear_copy_on_writes() + + def mark_blocks_as_accessed(self, block_ids: List[int], + now: float) -> None: + """Mark blocks as accessed, only use for prefix caching.""" + # Prefix caching only supported on GPU. + device = Device.GPU + return self._allocators[device].mark_blocks_as_accessed(block_ids, now) + + def mark_blocks_as_computed(self, block_ids: List[int]) -> None: + """Mark blocks as accessed, only use for prefix caching.""" + # Prefix caching only supported on GPU. + device = Device.GPU + return self._allocators[device].mark_blocks_as_computed(block_ids) + + def get_computed_block_ids(self, prev_computed_block_ids: List[int], + block_ids: List[int], + skip_last_block_id: bool) -> List[int]: + # Prefix caching only supported on GPU. + device = Device.GPU + return self._allocators[device].get_computed_block_ids( + prev_computed_block_ids, block_ids, skip_last_block_id) + + def get_common_computed_block_ids( + self, computed_seq_block_ids: List[List[int]]) -> List[int]: + # Prefix caching only supported on GPU. + device = Device.GPU + return self._allocators[device].get_common_computed_block_ids( + computed_seq_block_ids) + + @property + def all_block_ids(self) -> FrozenSet[int]: + return frozenset(self._block_ids_to_allocator.keys()) + + def get_prefix_cache_hit_rate(self, device: Device) -> float: + """Prefix cache hit rate. -1 means not supported or disabled.""" + assert device in self._allocators + return self._allocators[device].get_prefix_cache_hit_rate() + def get_and_reset_swaps(self) -> List[Tuple[int, int]]: - """Returns and clears the mapping of source to destination block IDs. - Will be called after every swapping operations for now, and after every - schedule when BlockManagerV2 become default. Currently not useful. - - Returns: - List[Tuple[int, int]]: A mapping of source to destination block IDs. - """ - mapping = self._swap_mapping.copy() + """Returns and clears the mapping of source to destination block IDs. + Will be called after every swapping operations for now, and after every + schedule when BlockManagerV2 become default. Currently not useful. + + Returns: + List[Tuple[int, int]]: A mapping of source to destination block IDs. + """ + mapping = self._swap_mapping.copy() self._swap_mapping.clear() return list(mapping.items()) @@ -407,69 +407,69 @@ class CpuGpuBlockAllocator(DeviceAwareBlockAllocator): def begin_prefix_cache_step(self) -> None: if self._cpu_content_cache is not None: self._cpu_content_cache.begin_step() - - -class NullBlock(Block): - """ - Null blocks are used as a placeholders for KV cache blocks that have - been dropped due to sliding window. - This implementation just wraps an ordinary block and prevents it from - being modified. It also allows for testing if a block is NullBlock - via isinstance(). - """ - - def __init__(self, proxy: Block): - super().__init__() - self._proxy = proxy - - def append_token_ids(self, token_ids: List[BlockId]): - raise ValueError("null block should not be modified") - - @property - def block_id(self): - return self._proxy.block_id - - @block_id.setter - def block_id(self, value: Optional[BlockId]): - raise ValueError("null block should not be modified") - - @property - def token_ids(self) -> List[BlockId]: - return self._proxy.token_ids - - @property - def num_tokens_total(self) -> int: - raise NotImplementedError( - "num_tokens_total is not used for null block") - - @property - def num_empty_slots(self) -> BlockId: - return self._proxy.num_empty_slots - - @property - def is_full(self): - return self._proxy.is_full - - @property - def prev_block(self): - return self._proxy.prev_block - - @property - def computed(self): - return self._proxy.computed - - @computed.setter - def computed(self, value): - self._proxy.computed = value - - @property - def last_accessed(self) -> float: - return self._proxy.last_accessed - - @last_accessed.setter - def last_accessed(self, last_accessed_ts: float): - self._proxy.last_accessed = last_accessed_ts - - @property - def content_hash(self): - return self._proxy.content_hash + + +class NullBlock(Block): + """ + Null blocks are used as a placeholders for KV cache blocks that have + been dropped due to sliding window. + This implementation just wraps an ordinary block and prevents it from + being modified. It also allows for testing if a block is NullBlock + via isinstance(). + """ + + def __init__(self, proxy: Block): + super().__init__() + self._proxy = proxy + + def append_token_ids(self, token_ids: List[BlockId]): + raise ValueError("null block should not be modified") + + @property + def block_id(self): + return self._proxy.block_id + + @block_id.setter + def block_id(self, value: Optional[BlockId]): + raise ValueError("null block should not be modified") + + @property + def token_ids(self) -> List[BlockId]: + return self._proxy.token_ids + + @property + def num_tokens_total(self) -> int: + raise NotImplementedError( + "num_tokens_total is not used for null block") + + @property + def num_empty_slots(self) -> BlockId: + return self._proxy.num_empty_slots + + @property + def is_full(self): + return self._proxy.is_full + + @property + def prev_block(self): + return self._proxy.prev_block + + @property + def computed(self): + return self._proxy.computed + + @computed.setter + def computed(self, value): + self._proxy.computed = value + + @property + def last_accessed(self) -> float: + return self._proxy.last_accessed + + @last_accessed.setter + def last_accessed(self, last_accessed_ts: float): + self._proxy.last_accessed = last_accessed_ts + + @property + def content_hash(self): + return self._proxy.content_hash diff --git a/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py b/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py index 772fa71f..c2c38e38 100644 --- a/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py +++ b/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py @@ -4,105 +4,105 @@ import struct from os.path import commonprefix from typing import (Callable, Dict, FrozenSet, Iterable, List, Optional, Set, Tuple) - -from vllm.core.block.common import (CacheMetricData, CopyOnWriteTracker, - get_all_blocks_recursively) -from vllm.core.block.interfaces import Block, BlockAllocator, BlockId, Device -from vllm.core.block.naive_block import (BlockPool, NaiveBlock, - NaiveBlockAllocator) + +from vllm.core.block.common import (CacheMetricData, CopyOnWriteTracker, + get_all_blocks_recursively) +from vllm.core.block.interfaces import Block, BlockAllocator, BlockId, Device +from vllm.core.block.naive_block import (BlockPool, NaiveBlock, + NaiveBlockAllocator) from vllm.core.evictor_v2 import (EvictionPolicy, Evictor, eviction_policy_from_env, make_evictor) - + PrefixHash = bytes - -# By default, we init our block access time as _DEFAULT_LAST_ACCESSED_TIME -# so that if we find one block is still hold _DEFAULT_LAST_ACCESSED_TIME, -# then we know this block hasn't been accessed yet. -_DEFAULT_LAST_ACCESSED_TIME = -1 - - -class BlockTracker: - """Used to track the status of a block inside the prefix caching allocator - """ - __slots__ = ("active", "last_accessed", "computed") - - def reset(self): - self.last_accessed: float = _DEFAULT_LAST_ACCESSED_TIME - self.computed: bool = False - - def __init__(self): - self.active: bool = False - self.reset() - - def enable(self): - assert not self.active - self.active = True - self.reset() - - def disable(self): - assert self.active - self.active = False - self.reset() - - -class PrefixCachingBlockAllocator(BlockAllocator): - """A block allocator that implements prefix caching. - - The PrefixCachingBlockAllocator maintains a cache of blocks based on their - content hash. It reuses blocks with the same content hash to avoid redundant - memory allocation. The allocator also supports copy-on-write operations. - - Args: - num_blocks (int): The total number of blocks to manage. - block_size (int): The size of each block in tokens. - block_ids(Optional[Iterable[int]], optional): An optional iterable of - block IDs. If not provided, block IDs will be assigned sequentially - from 0 to num_blocks - 1. - """ - - def __init__( - self, - num_blocks: int, - block_size: int, - block_ids: Optional[Iterable[int]] = None, + +# By default, we init our block access time as _DEFAULT_LAST_ACCESSED_TIME +# so that if we find one block is still hold _DEFAULT_LAST_ACCESSED_TIME, +# then we know this block hasn't been accessed yet. +_DEFAULT_LAST_ACCESSED_TIME = -1 + + +class BlockTracker: + """Used to track the status of a block inside the prefix caching allocator + """ + __slots__ = ("active", "last_accessed", "computed") + + def reset(self): + self.last_accessed: float = _DEFAULT_LAST_ACCESSED_TIME + self.computed: bool = False + + def __init__(self): + self.active: bool = False + self.reset() + + def enable(self): + assert not self.active + self.active = True + self.reset() + + def disable(self): + assert self.active + self.active = False + self.reset() + + +class PrefixCachingBlockAllocator(BlockAllocator): + """A block allocator that implements prefix caching. + + The PrefixCachingBlockAllocator maintains a cache of blocks based on their + content hash. It reuses blocks with the same content hash to avoid redundant + memory allocation. The allocator also supports copy-on-write operations. + + Args: + num_blocks (int): The total number of blocks to manage. + block_size (int): The size of each block in tokens. + block_ids(Optional[Iterable[int]], optional): An optional iterable of + block IDs. If not provided, block IDs will be assigned sequentially + from 0 to num_blocks - 1. + """ + + def __init__( + self, + num_blocks: int, + block_size: int, + block_ids: Optional[Iterable[int]] = None, eviction_policy: Optional[EvictionPolicy] = None, - ): - if block_ids is None: - block_ids = range(num_blocks) - - self._block_size = block_size - + ): + if block_ids is None: + block_ids = range(num_blocks) + + self._block_size = block_size + # A mapping of prefix hash to block index. All blocks which have a # prefix hash will be in this dict, even if they have refcount 0. self._cached_blocks: Dict[PrefixHash, BlockId] = {} self._cache_namespace: Optional[bytes] = None - - # A list of immutable block IDs that have been touched by scheduler - # and should be marked as computed after an entire batch of sequences - # are scheduled. - self._touched_blocks: Set[BlockId] = set() - - # Used to track status of each physical block id - self._block_tracker: Dict[BlockId, BlockTracker] = {} - for block_id in block_ids: - self._block_tracker[block_id] = BlockTracker() - - # Pre-allocate "num_blocks * extra_factor" block objects. - # The "* extra_factor" is a buffer to allow more block objects - # than physical blocks - extra_factor = 4 - self._block_pool = BlockPool(self._block_size, self._create_block, - self, num_blocks * extra_factor) - - # An allocator for blocks that do not have prefix hashes. - self._hashless_allocator = NaiveBlockAllocator( - create_block=self._create_block, # type: ignore - num_blocks=num_blocks, - block_size=block_size, - block_ids=block_ids, - block_pool=self._block_pool, # Share block pool here - ) - + + # A list of immutable block IDs that have been touched by scheduler + # and should be marked as computed after an entire batch of sequences + # are scheduled. + self._touched_blocks: Set[BlockId] = set() + + # Used to track status of each physical block id + self._block_tracker: Dict[BlockId, BlockTracker] = {} + for block_id in block_ids: + self._block_tracker[block_id] = BlockTracker() + + # Pre-allocate "num_blocks * extra_factor" block objects. + # The "* extra_factor" is a buffer to allow more block objects + # than physical blocks + extra_factor = 4 + self._block_pool = BlockPool(self._block_size, self._create_block, + self, num_blocks * extra_factor) + + # An allocator for blocks that do not have prefix hashes. + self._hashless_allocator = NaiveBlockAllocator( + create_block=self._create_block, # type: ignore + num_blocks=num_blocks, + block_size=block_size, + block_ids=block_ids, + block_pool=self._block_pool, # Share block pool here + ) + if eviction_policy is None: eviction_policy = eviction_policy_from_env() self.eviction_policy = eviction_policy @@ -110,15 +110,15 @@ class PrefixCachingBlockAllocator(BlockAllocator): # Evitor used to maintain how we want to handle those computed blocks # if we find memory pressure is high. self.evictor: Evictor = make_evictor(eviction_policy) - - # We share the refcounter between allocators. This allows us to promote - # blocks originally allocated in the hashless allocator to immutable - # blocks. - self._refcounter = self._hashless_allocator.refcounter - - self._cow_tracker = CopyOnWriteTracker( - refcounter=self._refcounter.as_readonly()) - + + # We share the refcounter between allocators. This allows us to promote + # blocks originally allocated in the hashless allocator to immutable + # blocks. + self._refcounter = self._hashless_allocator.refcounter + + self._cow_tracker = CopyOnWriteTracker( + refcounter=self._refcounter.as_readonly()) + self.metric_data = CacheMetricData() self._external_cache_claim: Optional[ @@ -153,8 +153,8 @@ class PrefixCachingBlockAllocator(BlockAllocator): self._external_cache_load = load self._external_cache_cancel = cancel self._external_cache_store = store - - # Implements Block.Factory. + + # Implements Block.Factory. def _create_block( self, prev_block: Optional[Block], @@ -334,48 +334,48 @@ class PrefixCachingBlockAllocator(BlockAllocator): prev_block: Optional[Block], token_ids: List[int], device: Optional[Device] = None) -> Block: - """Allocates an immutable block with the given token IDs, reusing cached - blocks if possible. - - Args: - prev_block (Optional[Block]): The previous block in the sequence. - token_ids (List[int]): The token IDs to be stored in the block. - - Returns: - Block: The allocated immutable block. - """ + """Allocates an immutable block with the given token IDs, reusing cached + blocks if possible. + + Args: + prev_block (Optional[Block]): The previous block in the sequence. + token_ids (List[int]): The token IDs to be stored in the block. + + Returns: + Block: The allocated immutable block. + """ return self.allocate_immutable_block_with_cache_namespace( prev_block=prev_block, token_ids=token_ids, cache_namespace=b"", device=device) - - def allocate_immutable_blocks( - self, - prev_block: Optional[Block], - block_token_ids: List[List[int]], - device: Optional[Device] = None) -> List[Block]: - blocks = [] - for token_ids in block_token_ids: - prev_block = self.allocate_immutable_block(prev_block=prev_block, - token_ids=token_ids, - device=device) - blocks.append(prev_block) - return blocks - - def allocate_mutable_block(self, - prev_block: Optional[Block], - device: Optional[Device] = None) -> Block: - """Allocates a mutable block. If there are no free blocks, this will - evict unused cached blocks. - - Args: - prev_block (Block): The previous block in the sequence. - None is not allowed unlike it is super class. - - Returns: - Block: The allocated mutable block. - """ + + def allocate_immutable_blocks( + self, + prev_block: Optional[Block], + block_token_ids: List[List[int]], + device: Optional[Device] = None) -> List[Block]: + blocks = [] + for token_ids in block_token_ids: + prev_block = self.allocate_immutable_block(prev_block=prev_block, + token_ids=token_ids, + device=device) + blocks.append(prev_block) + return blocks + + def allocate_mutable_block(self, + prev_block: Optional[Block], + device: Optional[Device] = None) -> Block: + """Allocates a mutable block. If there are no free blocks, this will + evict unused cached blocks. + + Args: + prev_block (Block): The previous block in the sequence. + None is not allowed unlike it is super class. + + Returns: + Block: The allocated mutable block. + """ assert device is None assert_prefix_caching_block_or_none(prev_block) @@ -387,105 +387,105 @@ class PrefixCachingBlockAllocator(BlockAllocator): assert not block.computed assert block.content_hash is None return block - - def _incr_refcount_cached_block(self, block: Block) -> None: - # Set this block to be "computed" since it is pointing to a - # cached block id (which was already computed) - block.computed = True - - block_id = block.block_id - assert block_id is not None - - refcount = self._refcounter.incr(block_id) - if refcount == 1: - # In case a cached block was evicted, restore its tracking - if block_id in self.evictor: - self.evictor.remove(block_id) - - self._track_block_id(block_id, computed=True) - - def _decr_refcount_cached_block(self, block: Block) -> None: - # Ensure this is immutable/cached block - assert block.content_hash is not None - - block_id = block.block_id - assert block_id is not None - - refcount = self._refcounter.decr(block_id) - if refcount > 0: - block.block_id = None - return - else: - assert refcount == 0 - - # No longer used - assert block.content_hash in self._cached_blocks - - # Add the cached block to the evictor - # (This keeps the cached block around so it can be reused) - self.evictor.add(block_id, block.content_hash, block.num_tokens_total, - self._block_tracker[block_id].last_accessed) - - # Stop tracking the block - self._untrack_block_id(block_id) - - block.block_id = None - - def _decr_refcount_hashless_block(self, block: Block) -> None: - block_id = block.block_id - assert block_id is not None - - # We may have a fork case where block is shared, - # in which case, we cannot remove it from tracking - refcount = self._refcounter.get(block_id) - if refcount == 1: - self._untrack_block_id(block_id) - - # Decrement refcount of the block_id, but do not free the block object - # itself (will be handled by the caller) - self._hashless_allocator.free(block, keep_block_object=True) - - def _allocate_block_id(self) -> BlockId: - """First tries to allocate a block id from the hashless allocator, - and if there are no blocks, then tries to evict an unused cached block. - """ - hashless_block_id = self._maybe_allocate_hashless_block_id() - if hashless_block_id is not None: - return hashless_block_id - - evicted_block_id = self._maybe_allocate_evicted_block_id() - if evicted_block_id is not None: - return evicted_block_id - - # No block available in hashless allocator, nor in unused cache blocks. - raise BlockAllocator.NoFreeBlocksError() - - def _maybe_allocate_hashless_block_id(self) -> Optional[BlockId]: - try: - # Allocate mutable block and extract its block_id - block = self._hashless_allocator.allocate_mutable_block( - prev_block=None) - block_id = block.block_id - self._block_pool.free_block(block) - - self._track_block_id(block_id, computed=False) - return block_id - except BlockAllocator.NoFreeBlocksError: - return None - - def _maybe_allocate_evicted_block_id(self) -> Optional[BlockId]: - if self.evictor.num_blocks == 0: - return None - - # Here we get an evicted block, which is only added - # into evictor if its ref counter is 0 - # and since its content would be changed, we need - # to remove it from _cached_blocks's tracking list - block_id, content_hash_to_evict = self.evictor.evict() - - # Sanity checks - assert content_hash_to_evict in self._cached_blocks - _block_id = self._cached_blocks[content_hash_to_evict] + + def _incr_refcount_cached_block(self, block: Block) -> None: + # Set this block to be "computed" since it is pointing to a + # cached block id (which was already computed) + block.computed = True + + block_id = block.block_id + assert block_id is not None + + refcount = self._refcounter.incr(block_id) + if refcount == 1: + # In case a cached block was evicted, restore its tracking + if block_id in self.evictor: + self.evictor.remove(block_id) + + self._track_block_id(block_id, computed=True) + + def _decr_refcount_cached_block(self, block: Block) -> None: + # Ensure this is immutable/cached block + assert block.content_hash is not None + + block_id = block.block_id + assert block_id is not None + + refcount = self._refcounter.decr(block_id) + if refcount > 0: + block.block_id = None + return + else: + assert refcount == 0 + + # No longer used + assert block.content_hash in self._cached_blocks + + # Add the cached block to the evictor + # (This keeps the cached block around so it can be reused) + self.evictor.add(block_id, block.content_hash, block.num_tokens_total, + self._block_tracker[block_id].last_accessed) + + # Stop tracking the block + self._untrack_block_id(block_id) + + block.block_id = None + + def _decr_refcount_hashless_block(self, block: Block) -> None: + block_id = block.block_id + assert block_id is not None + + # We may have a fork case where block is shared, + # in which case, we cannot remove it from tracking + refcount = self._refcounter.get(block_id) + if refcount == 1: + self._untrack_block_id(block_id) + + # Decrement refcount of the block_id, but do not free the block object + # itself (will be handled by the caller) + self._hashless_allocator.free(block, keep_block_object=True) + + def _allocate_block_id(self) -> BlockId: + """First tries to allocate a block id from the hashless allocator, + and if there are no blocks, then tries to evict an unused cached block. + """ + hashless_block_id = self._maybe_allocate_hashless_block_id() + if hashless_block_id is not None: + return hashless_block_id + + evicted_block_id = self._maybe_allocate_evicted_block_id() + if evicted_block_id is not None: + return evicted_block_id + + # No block available in hashless allocator, nor in unused cache blocks. + raise BlockAllocator.NoFreeBlocksError() + + def _maybe_allocate_hashless_block_id(self) -> Optional[BlockId]: + try: + # Allocate mutable block and extract its block_id + block = self._hashless_allocator.allocate_mutable_block( + prev_block=None) + block_id = block.block_id + self._block_pool.free_block(block) + + self._track_block_id(block_id, computed=False) + return block_id + except BlockAllocator.NoFreeBlocksError: + return None + + def _maybe_allocate_evicted_block_id(self) -> Optional[BlockId]: + if self.evictor.num_blocks == 0: + return None + + # Here we get an evicted block, which is only added + # into evictor if its ref counter is 0 + # and since its content would be changed, we need + # to remove it from _cached_blocks's tracking list + block_id, content_hash_to_evict = self.evictor.evict() + + # Sanity checks + assert content_hash_to_evict in self._cached_blocks + _block_id = self._cached_blocks[content_hash_to_evict] assert self._refcounter.get(_block_id) == 0 assert _block_id == block_id @@ -493,242 +493,242 @@ class PrefixCachingBlockAllocator(BlockAllocator): self._external_cache_store(content_hash_to_evict, block_id) self._cached_blocks.pop(content_hash_to_evict) - - self._refcounter.incr(block_id) - self._track_block_id(block_id, computed=False) - - return block_id - - def _free_block_id(self, block: Block) -> None: - """Decrements the refcount of the block. The block may be in two - possible states: (1) immutable/cached or (2) mutable/hashless. - In the first case, the refcount is decremented directly and the block - may be possibly added to the evictor. In other case, hashless - allocator free(..) with keep_block_object=True is called to only free - the block id (since the block object may be reused by the caller) - """ - block_id = block.block_id - assert block_id is not None, "Freeing unallocated block is undefined" - - if block.content_hash is not None: - # Immutable: This type of block is always cached, and we want to - # keep it in the evictor for future reuse - self._decr_refcount_cached_block(block) - else: - # Mutable: This type of block is not cached, so we release it - # directly to the hashless allocator - self._decr_refcount_hashless_block(block) - - assert block.block_id is None - - def free(self, block: Block, keep_block_object: bool = False) -> None: - """Release the block (look at free_block_id(..) docs) - """ - # Release the physical block index - self._free_block_id(block) - - # Release the block object to the pool - if not keep_block_object: - self._block_pool.free_block(block) - - def fork(self, last_block: Block) -> List[Block]: - """Creates a new sequence of blocks that shares the same underlying - memory as the original sequence. - - Args: - last_block (Block): The last block in the original sequence. - - Returns: - List[Block]: The new sequence of blocks that shares the same memory - as the original sequence. - """ - source_blocks = get_all_blocks_recursively(last_block) - - forked_blocks: List[Block] = [] - prev_block = None - for block in source_blocks: - block_id = block.block_id - assert block_id is not None - - refcount = self._refcounter.incr(block_id) - assert refcount != 1, "can't fork free'd block_id = {}".format( - block_id) - + + self._refcounter.incr(block_id) + self._track_block_id(block_id, computed=False) + + return block_id + + def _free_block_id(self, block: Block) -> None: + """Decrements the refcount of the block. The block may be in two + possible states: (1) immutable/cached or (2) mutable/hashless. + In the first case, the refcount is decremented directly and the block + may be possibly added to the evictor. In other case, hashless + allocator free(..) with keep_block_object=True is called to only free + the block id (since the block object may be reused by the caller) + """ + block_id = block.block_id + assert block_id is not None, "Freeing unallocated block is undefined" + + if block.content_hash is not None: + # Immutable: This type of block is always cached, and we want to + # keep it in the evictor for future reuse + self._decr_refcount_cached_block(block) + else: + # Mutable: This type of block is not cached, so we release it + # directly to the hashless allocator + self._decr_refcount_hashless_block(block) + + assert block.block_id is None + + def free(self, block: Block, keep_block_object: bool = False) -> None: + """Release the block (look at free_block_id(..) docs) + """ + # Release the physical block index + self._free_block_id(block) + + # Release the block object to the pool + if not keep_block_object: + self._block_pool.free_block(block) + + def fork(self, last_block: Block) -> List[Block]: + """Creates a new sequence of blocks that shares the same underlying + memory as the original sequence. + + Args: + last_block (Block): The last block in the original sequence. + + Returns: + List[Block]: The new sequence of blocks that shares the same memory + as the original sequence. + """ + source_blocks = get_all_blocks_recursively(last_block) + + forked_blocks: List[Block] = [] + prev_block = None + for block in source_blocks: + block_id = block.block_id + assert block_id is not None + + refcount = self._refcounter.incr(block_id) + assert refcount != 1, "can't fork free'd block_id = {}".format( + block_id) + forked_block = self._init_block( prev_block=prev_block, token_ids=block.token_ids, block_size=self._block_size, physical_block_id=block_id, cache_namespace=block.cache_namespace) - - forked_blocks.append(forked_block) - prev_block = forked_blocks[-1] - - return forked_blocks - - def get_num_free_blocks(self, device: Optional[Device] = None) -> int: - assert device is None - # The number of free blocks is the number of hashless free blocks - # plus the number of blocks evictor could free from its list. - return self._hashless_allocator.get_num_free_blocks( - ) + self.evictor.num_blocks - - def get_num_total_blocks(self) -> int: - return self._hashless_allocator.get_num_total_blocks() - - def get_physical_block_id(self, absolute_id: int) -> int: - """Returns the zero-offset block id on certain block allocator - given the absolute block id. - - Args: - absolute_id (int): The absolute block id for the block - in whole allocator. - - Returns: - int: The rzero-offset block id on certain device. - """ - return sorted(self.all_block_ids).index(absolute_id) - - @property - def all_block_ids(self) -> FrozenSet[int]: - return self._hashless_allocator.all_block_ids - - def get_prefix_cache_hit_rate(self) -> float: - return self.metric_data.get_hit_rate() - - def is_block_cached(self, block: Block) -> bool: - assert block.content_hash is not None - return block.content_hash in self._cached_blocks - - def promote_to_immutable_block(self, block: Block) -> BlockId: - """Once a mutable block is full, it can be promoted to an immutable - block. This means that its content can be referenced by future blocks - having the same prefix. - - Note that if we already have a cached block with the same content, we - will replace the newly-promoted block's mapping with the existing cached - block id. - - Args: - block: The mutable block to be promoted. - - Returns: - BlockId: Either the original block index, or the block index of - the previously cached block matching the same content. - """ - # Ensure block can be promoted - assert block.content_hash is not None - assert block.block_id is not None - assert self._refcounter.get(block.block_id) > 0 - - if block.content_hash not in self._cached_blocks: - # No cached content hash => Set this block as cached. - # Note that this block cannot be marked as computed yet - # because other sequences in the same batch cannot reuse - # this block. - self._cached_blocks[block.content_hash] = block.block_id - # Mark this block as touched so that it can be marked as - # computed after the entire batch of sequences are scheduled. - self._touched_blocks.add(block.block_id) - return block.block_id - - # Reuse the cached content hash - self._decr_refcount_hashless_block(block) - block.block_id = self._cached_blocks[block.content_hash] - - # Increment refcount of the cached block and (possibly) restore - # it from the evictor. - # Note that in this case, the block is marked as computed - self._incr_refcount_cached_block(block) - - return block.block_id - - def cow_block_if_not_appendable(self, block: Block) -> BlockId: - """Performs a copy-on-write operation on the given block if it is not - appendable. - - Args: - block (Block): The block to check for copy-on-write. - - Returns: - BlockId: The block index of the new block if a copy-on-write - operation was performed, or the original block index if - no copy-on-write was necessary. - """ - src_block_id = block.block_id - assert src_block_id is not None - - if self._cow_tracker.is_appendable(block): - return src_block_id - - self._free_block_id(block) - trg_block_id = self._allocate_block_id() - - self._cow_tracker.record_cow(src_block_id, trg_block_id) - - return trg_block_id - - def clear_copy_on_writes(self) -> List[Tuple[BlockId, BlockId]]: - """Returns the copy-on-write source->destination mapping and clears it. - - Returns: - List[Tuple[BlockId, BlockId]]: A list mapping source - block indices to destination block indices. - """ - return self._cow_tracker.clear_cows() - - def mark_blocks_as_accessed(self, block_ids: List[int], - now: float) -> None: - """Mark blocks as accessed, used in prefix caching. - - If the block is added into evictor, we need to update corresponding - info in evictor's metadata. - """ - - for block_id in block_ids: - if self._block_tracker[block_id].active: - self._block_tracker[block_id].last_accessed = now - elif block_id in self.evictor: - self.evictor.update(block_id, now) - else: - raise ValueError( - "Mark block as accessed which is not belonged to GPU") - - def mark_blocks_as_computed(self, block_ids: List[int]) -> None: - # Mark all touched blocks as computed. - for block_id in self._touched_blocks: - self._block_tracker[block_id].computed = True - self._touched_blocks.clear() - - def _track_block_id(self, block_id: Optional[BlockId], - computed: bool) -> None: - assert block_id is not None - self._block_tracker[block_id].enable() - self._block_tracker[block_id].computed = computed - - def _untrack_block_id(self, block_id: Optional[BlockId]) -> None: - assert block_id is not None - self._block_tracker[block_id].disable() - - def block_is_computed(self, block_id: int) -> bool: - if self._block_tracker[block_id].active: - return self._block_tracker[block_id].computed - else: - return block_id in self.evictor - - def get_computed_block_ids(self, - prev_computed_block_ids: List[int], - block_ids: List[int], - skip_last_block_id: bool = True) -> List[int]: - prev_prefix_size = len(prev_computed_block_ids) - cur_size = len(block_ids) - if skip_last_block_id: - cur_size -= 1 - - # Sanity checks - assert cur_size >= 0 - assert prev_prefix_size <= cur_size - + + forked_blocks.append(forked_block) + prev_block = forked_blocks[-1] + + return forked_blocks + + def get_num_free_blocks(self, device: Optional[Device] = None) -> int: + assert device is None + # The number of free blocks is the number of hashless free blocks + # plus the number of blocks evictor could free from its list. + return self._hashless_allocator.get_num_free_blocks( + ) + self.evictor.num_blocks + + def get_num_total_blocks(self) -> int: + return self._hashless_allocator.get_num_total_blocks() + + def get_physical_block_id(self, absolute_id: int) -> int: + """Returns the zero-offset block id on certain block allocator + given the absolute block id. + + Args: + absolute_id (int): The absolute block id for the block + in whole allocator. + + Returns: + int: The rzero-offset block id on certain device. + """ + return sorted(self.all_block_ids).index(absolute_id) + + @property + def all_block_ids(self) -> FrozenSet[int]: + return self._hashless_allocator.all_block_ids + + def get_prefix_cache_hit_rate(self) -> float: + return self.metric_data.get_hit_rate() + + def is_block_cached(self, block: Block) -> bool: + assert block.content_hash is not None + return block.content_hash in self._cached_blocks + + def promote_to_immutable_block(self, block: Block) -> BlockId: + """Once a mutable block is full, it can be promoted to an immutable + block. This means that its content can be referenced by future blocks + having the same prefix. + + Note that if we already have a cached block with the same content, we + will replace the newly-promoted block's mapping with the existing cached + block id. + + Args: + block: The mutable block to be promoted. + + Returns: + BlockId: Either the original block index, or the block index of + the previously cached block matching the same content. + """ + # Ensure block can be promoted + assert block.content_hash is not None + assert block.block_id is not None + assert self._refcounter.get(block.block_id) > 0 + + if block.content_hash not in self._cached_blocks: + # No cached content hash => Set this block as cached. + # Note that this block cannot be marked as computed yet + # because other sequences in the same batch cannot reuse + # this block. + self._cached_blocks[block.content_hash] = block.block_id + # Mark this block as touched so that it can be marked as + # computed after the entire batch of sequences are scheduled. + self._touched_blocks.add(block.block_id) + return block.block_id + + # Reuse the cached content hash + self._decr_refcount_hashless_block(block) + block.block_id = self._cached_blocks[block.content_hash] + + # Increment refcount of the cached block and (possibly) restore + # it from the evictor. + # Note that in this case, the block is marked as computed + self._incr_refcount_cached_block(block) + + return block.block_id + + def cow_block_if_not_appendable(self, block: Block) -> BlockId: + """Performs a copy-on-write operation on the given block if it is not + appendable. + + Args: + block (Block): The block to check for copy-on-write. + + Returns: + BlockId: The block index of the new block if a copy-on-write + operation was performed, or the original block index if + no copy-on-write was necessary. + """ + src_block_id = block.block_id + assert src_block_id is not None + + if self._cow_tracker.is_appendable(block): + return src_block_id + + self._free_block_id(block) + trg_block_id = self._allocate_block_id() + + self._cow_tracker.record_cow(src_block_id, trg_block_id) + + return trg_block_id + + def clear_copy_on_writes(self) -> List[Tuple[BlockId, BlockId]]: + """Returns the copy-on-write source->destination mapping and clears it. + + Returns: + List[Tuple[BlockId, BlockId]]: A list mapping source + block indices to destination block indices. + """ + return self._cow_tracker.clear_cows() + + def mark_blocks_as_accessed(self, block_ids: List[int], + now: float) -> None: + """Mark blocks as accessed, used in prefix caching. + + If the block is added into evictor, we need to update corresponding + info in evictor's metadata. + """ + + for block_id in block_ids: + if self._block_tracker[block_id].active: + self._block_tracker[block_id].last_accessed = now + elif block_id in self.evictor: + self.evictor.update(block_id, now) + else: + raise ValueError( + "Mark block as accessed which is not belonged to GPU") + + def mark_blocks_as_computed(self, block_ids: List[int]) -> None: + # Mark all touched blocks as computed. + for block_id in self._touched_blocks: + self._block_tracker[block_id].computed = True + self._touched_blocks.clear() + + def _track_block_id(self, block_id: Optional[BlockId], + computed: bool) -> None: + assert block_id is not None + self._block_tracker[block_id].enable() + self._block_tracker[block_id].computed = computed + + def _untrack_block_id(self, block_id: Optional[BlockId]) -> None: + assert block_id is not None + self._block_tracker[block_id].disable() + + def block_is_computed(self, block_id: int) -> bool: + if self._block_tracker[block_id].active: + return self._block_tracker[block_id].computed + else: + return block_id in self.evictor + + def get_computed_block_ids(self, + prev_computed_block_ids: List[int], + block_ids: List[int], + skip_last_block_id: bool = True) -> List[int]: + prev_prefix_size = len(prev_computed_block_ids) + cur_size = len(block_ids) + if skip_last_block_id: + cur_size -= 1 + + # Sanity checks + assert cur_size >= 0 + assert prev_prefix_size <= cur_size + ret = prev_computed_block_ids for i in range(prev_prefix_size, cur_size): block_id = block_ids[i] @@ -736,73 +736,73 @@ class PrefixCachingBlockAllocator(BlockAllocator): break ret.append(block_id) return ret - - def get_common_computed_block_ids( - self, computed_seq_block_ids: List[List[int]]) -> List[int]: - """Return the block ids that are common for a given sequence group. - - Only those blocks that are immutable and already be marked - compyted would be taken consideration. - """ - - # NOTE We exclude the last block to avoid the case where the entire - # prompt is cached. This would cause erroneous behavior in model - # runner. - - # It returns a list of int although type annotation says list of string. - if len(computed_seq_block_ids) == 1: - return computed_seq_block_ids[0] - - return commonprefix([ - ids for ids in computed_seq_block_ids # type: ignore - if ids - ]) - - def get_num_full_blocks_touched(self, blocks: List[Block]) -> int: - """Returns the number of full blocks that will be touched by - swapping in/out. - - Args: - blocks: List of blocks to be swapped. - Returns: - int: the number of full blocks that will be touched by - swapping in/out the given blocks. Non full blocks are ignored - when deciding the number of blocks to touch. - """ - num_touched_blocks: int = 0 - for block in blocks: - # If the block has a match in the cache and the cached - # block is not referenced, then we still count it as a - # touched block - if block.is_full and (not self.is_block_cached(block) or \ - (block.content_hash is not None and \ - self._cached_blocks[block.content_hash] in \ - self.evictor)): - num_touched_blocks += 1 - return num_touched_blocks - - def swap_out(self, blocks: List[Block]) -> None: - """Execute the swap out actions. Basically just free the - given blocks. - - Args: - blocks: List of blocks to be swapped out. - """ - for block in blocks: - self._free_block_id(block) - - def swap_in(self, blocks: List[Block]) -> None: - """Execute the swap in actions. Change the block id from - old allocator to current allocator for each block to finish - the block table update. - - Args: - blocks: List of blocks to be swapped in. - """ - for block in blocks: - # Here we allocate either immutable or mutable block and then - # extract its block_id. Note that the block object is released - # and the block_id is assigned to "block" to allow reusing the + + def get_common_computed_block_ids( + self, computed_seq_block_ids: List[List[int]]) -> List[int]: + """Return the block ids that are common for a given sequence group. + + Only those blocks that are immutable and already be marked + compyted would be taken consideration. + """ + + # NOTE We exclude the last block to avoid the case where the entire + # prompt is cached. This would cause erroneous behavior in model + # runner. + + # It returns a list of int although type annotation says list of string. + if len(computed_seq_block_ids) == 1: + return computed_seq_block_ids[0] + + return commonprefix([ + ids for ids in computed_seq_block_ids # type: ignore + if ids + ]) + + def get_num_full_blocks_touched(self, blocks: List[Block]) -> int: + """Returns the number of full blocks that will be touched by + swapping in/out. + + Args: + blocks: List of blocks to be swapped. + Returns: + int: the number of full blocks that will be touched by + swapping in/out the given blocks. Non full blocks are ignored + when deciding the number of blocks to touch. + """ + num_touched_blocks: int = 0 + for block in blocks: + # If the block has a match in the cache and the cached + # block is not referenced, then we still count it as a + # touched block + if block.is_full and (not self.is_block_cached(block) or \ + (block.content_hash is not None and \ + self._cached_blocks[block.content_hash] in \ + self.evictor)): + num_touched_blocks += 1 + return num_touched_blocks + + def swap_out(self, blocks: List[Block]) -> None: + """Execute the swap out actions. Basically just free the + given blocks. + + Args: + blocks: List of blocks to be swapped out. + """ + for block in blocks: + self._free_block_id(block) + + def swap_in(self, blocks: List[Block]) -> None: + """Execute the swap in actions. Change the block id from + old allocator to current allocator for each block to finish + the block table update. + + Args: + blocks: List of blocks to be swapped in. + """ + for block in blocks: + # Here we allocate either immutable or mutable block and then + # extract its block_id. Note that the block object is released + # and the block_id is assigned to "block" to allow reusing the # existing "block" object if block.is_full: tmp_block = ( @@ -816,33 +816,33 @@ class PrefixCachingBlockAllocator(BlockAllocator): prev_block=block.prev_block, cache_namespace=block.cache_namespace)) tmp_block.append_token_ids(block.token_ids) - - block_id = tmp_block.block_id - self._block_pool.free_block(tmp_block) - - block.block_id = block_id # Assign block_id - - -class PrefixCachingBlock(Block): - """A block implementation that supports prefix caching. - - The PrefixCachingBlock class represents a block of token IDs with prefix - caching capabilities. It wraps a NaiveBlock internally and provides - additional functionality for content hashing and promoting immutable blocks - with the prefix caching allocator. - - Args: - prev_block (Optional[PrefixCachingBlock]): The previous block in the - sequence. - token_ids (List[int]): The initial token IDs to be stored in the block. - block_size (int): The maximum number of token IDs that can be stored in - the block. - allocator (BlockAllocator): The prefix - caching block allocator associated with this block. - block_id (Optional[int], optional): The physical block index - of this block. Defaults to None. - """ - + + block_id = tmp_block.block_id + self._block_pool.free_block(tmp_block) + + block.block_id = block_id # Assign block_id + + +class PrefixCachingBlock(Block): + """A block implementation that supports prefix caching. + + The PrefixCachingBlock class represents a block of token IDs with prefix + caching capabilities. It wraps a NaiveBlock internally and provides + additional functionality for content hashing and promoting immutable blocks + with the prefix caching allocator. + + Args: + prev_block (Optional[PrefixCachingBlock]): The previous block in the + sequence. + token_ids (List[int]): The initial token IDs to be stored in the block. + block_size (int): The maximum number of token IDs that can be stored in + the block. + allocator (BlockAllocator): The prefix + caching block allocator associated with this block. + block_id (Optional[int], optional): The physical block index + of this block. Defaults to None. + """ + def __init__( self, prev_block: Optional[Block], @@ -853,12 +853,12 @@ class PrefixCachingBlock(Block): computed: bool = False, cache_namespace: Optional[bytes] = None, ): - assert isinstance(allocator, PrefixCachingBlockAllocator), ( - "Currently this class is only tested with " - "PrefixCachingBlockAllocator. Got instead allocator = {}".format( - allocator)) - assert_prefix_caching_block_or_none(prev_block) - + assert isinstance(allocator, PrefixCachingBlockAllocator), ( + "Currently this class is only tested with " + "PrefixCachingBlockAllocator. Got instead allocator = {}".format( + allocator)) + assert_prefix_caching_block_or_none(prev_block) + self._prev_block = prev_block self._cached_content_hash: Optional[bytes] = None self._cache_namespace = cache_namespace or b"" @@ -868,16 +868,16 @@ class PrefixCachingBlock(Block): self._allocator = allocator self._last_accessed: float = _DEFAULT_LAST_ACCESSED_TIME self._computed = computed - - # On the first time, we create the block object, and next we only - # reinitialize it - if hasattr(self, "_block"): - self._block.__init__( # type: ignore[has-type] - prev_block=prev_block, - token_ids=token_ids, - block_size=block_size, - block_id=block_id, - allocator=self._allocator) + + # On the first time, we create the block object, and next we only + # reinitialize it + if hasattr(self, "_block"): + self._block.__init__( # type: ignore[has-type] + prev_block=prev_block, + token_ids=token_ids, + block_size=block_size, + block_id=block_id, + allocator=self._allocator) else: self._block = NaiveBlock(prev_block=prev_block, token_ids=token_ids, @@ -886,94 +886,94 @@ class PrefixCachingBlock(Block): allocator=self._allocator) self._update_num_tokens_total() - - def _update_num_tokens_total(self): - """Incrementally computes the number of tokens that there is - till the current block (included) - """ - res = 0 - - # Add all previous blocks - if self._prev_block is not None: - res += self._prev_block.num_tokens_total - - # Add current block - res += len(self.token_ids) - - self._cached_num_tokens_total = res - - @property - def computed(self) -> bool: - return self._computed - - @computed.setter - def computed(self, value) -> None: - self._computed = value - - @property - def last_accessed(self) -> float: - return self._last_accessed - - @last_accessed.setter - def last_accessed(self, last_accessed_ts: float): - self._last_accessed = last_accessed_ts - - def append_token_ids(self, token_ids: List[int]) -> None: - """Appends the given token IDs to the block and registers the block as - immutable if the block becomes full. - - Args: - token_ids (List[int]): The token IDs to be appended to the block. - """ - # Ensure this is mutable block (not promoted) - assert self.content_hash is None - assert not self.computed - - if len(token_ids) == 0: - return - - # Ensure there are input tokens - assert token_ids, "Got token_ids = {}".format(token_ids) - - # Naive block handles CoW. - self._block.append_token_ids(token_ids) - self._update_num_tokens_total() - - # If the content hash is present, then the block can be made immutable. - # Register ourselves with the allocator, potentially replacing the - # physical block index. - if self.content_hash is not None: - self.block_id = self._allocator.promote_to_immutable_block(self) - - @property - def block_id(self) -> Optional[int]: - return self._block.block_id - - @block_id.setter - def block_id(self, value) -> None: - self._block.block_id = value - - @property - def is_full(self) -> bool: - return self._block.is_full - - @property - def num_empty_slots(self) -> int: - return self._block.num_empty_slots - - @property - def num_tokens_total(self) -> int: - return self._cached_num_tokens_total - - @property - def block_size(self) -> int: - return self._block.block_size - - @property - def token_ids(self) -> List[int]: - return self._block.token_ids - - @property + + def _update_num_tokens_total(self): + """Incrementally computes the number of tokens that there is + till the current block (included) + """ + res = 0 + + # Add all previous blocks + if self._prev_block is not None: + res += self._prev_block.num_tokens_total + + # Add current block + res += len(self.token_ids) + + self._cached_num_tokens_total = res + + @property + def computed(self) -> bool: + return self._computed + + @computed.setter + def computed(self, value) -> None: + self._computed = value + + @property + def last_accessed(self) -> float: + return self._last_accessed + + @last_accessed.setter + def last_accessed(self, last_accessed_ts: float): + self._last_accessed = last_accessed_ts + + def append_token_ids(self, token_ids: List[int]) -> None: + """Appends the given token IDs to the block and registers the block as + immutable if the block becomes full. + + Args: + token_ids (List[int]): The token IDs to be appended to the block. + """ + # Ensure this is mutable block (not promoted) + assert self.content_hash is None + assert not self.computed + + if len(token_ids) == 0: + return + + # Ensure there are input tokens + assert token_ids, "Got token_ids = {}".format(token_ids) + + # Naive block handles CoW. + self._block.append_token_ids(token_ids) + self._update_num_tokens_total() + + # If the content hash is present, then the block can be made immutable. + # Register ourselves with the allocator, potentially replacing the + # physical block index. + if self.content_hash is not None: + self.block_id = self._allocator.promote_to_immutable_block(self) + + @property + def block_id(self) -> Optional[int]: + return self._block.block_id + + @block_id.setter + def block_id(self, value) -> None: + self._block.block_id = value + + @property + def is_full(self) -> bool: + return self._block.is_full + + @property + def num_empty_slots(self) -> int: + return self._block.num_empty_slots + + @property + def num_tokens_total(self) -> int: + return self._cached_num_tokens_total + + @property + def block_size(self) -> int: + return self._block.block_size + + @property + def token_ids(self) -> List[int]: + return self._block.token_ids + + @property def prev_block(self) -> Optional[Block]: return self._prev_block @@ -983,31 +983,31 @@ class PrefixCachingBlock(Block): @property def content_hash(self) -> Optional[bytes]: - """Return the content-based hash of the current block, or None if it is - not yet defined. - - For the content-based hash to be defined, the current block must be - full. - """ - # If the hash is already computed, return it. - if self._cached_content_hash is not None: - return self._cached_content_hash - - # We cannot compute a hash for the current block because it is not full. - if not self.is_full: - return None - + """Return the content-based hash of the current block, or None if it is + not yet defined. + + For the content-based hash to be defined, the current block must be + full. + """ + # If the hash is already computed, return it. + if self._cached_content_hash is not None: + return self._cached_content_hash + + # We cannot compute a hash for the current block because it is not full. + if not self.is_full: + return None + is_first_block = self._prev_block is None prev_block_hash = ( None if is_first_block else self._prev_block.content_hash # type: ignore ) - - # Previous block exists but does not yet have a hash. - # Return no hash in this case. - if prev_block_hash is None and not is_first_block: - return None - + + # Previous block exists but does not yet have a hash. + # Return no hash in this case. + if prev_block_hash is None and not is_first_block: + return None + self._cached_content_hash = PrefixCachingBlock.hash_block_tokens( is_first_block, prev_block_hash, @@ -1022,21 +1022,21 @@ class PrefixCachingBlock(Block): cur_block_token_ids: List[int], cache_namespace: Optional[bytes] = None, ) -> bytes: - """Computes a hash value corresponding to the contents of a block and - the contents of the preceding block(s). The hash value is used for - prefix caching. - - NOTE: Content-based hashing does not yet support LoRA. - - Parameters: - - is_first_block (bool): A flag indicating if the block is the first in - the sequence. - - prev_block_hash (Optional[int]): The hash of the previous block. None - if this is the first block. - - cur_block_token_ids (List[int]): A list of token ids in the current - block. The current block is assumed to be full. - - Returns: + """Computes a hash value corresponding to the contents of a block and + the contents of the preceding block(s). The hash value is used for + prefix caching. + + NOTE: Content-based hashing does not yet support LoRA. + + Parameters: + - is_first_block (bool): A flag indicating if the block is the first in + the sequence. + - prev_block_hash (Optional[int]): The hash of the previous block. None + if this is the first block. + - cur_block_token_ids (List[int]): A list of token ids in the current + block. The current block is assumed to be full. + + Returns: - bytes: The computed hash value for the block. """ assert (prev_block_hash is None) == is_first_block @@ -1053,131 +1053,131 @@ class PrefixCachingBlock(Block): f"!{len(cur_block_token_ids)}q", *(int(token_id) for token_id in cur_block_token_ids))) return digest.digest() - - -class ComputedBlocksTracker: - """Handles caching of per-sequence computed block ids. - When a sequence appears for the first time, it traverses all of the - blocks and detects the prefix of blocks that is computed. On the - subsequent times, it only traverses the new blocks that were added - and updates the already recorded prefix of blocks with the newly - computed blocks. - - To avoid redundant traversals, the algorithm also detects when there - is a "gap" in the computed prefix. For example, if we have blocks = - [1,2,3,4,5], and we have detected [1,2,3] as the computed prefix, then - we won't try to add more computed blocks to [1,2,3] in this sequence - iteration, and will add more computed blocks only after the sequence is - freed and reused again. - - Note that currently, for a given sequence, we also skip the last - block id for caching purposes, to avoid caching of a full sequence - """ - - def __init__(self, allocator): - self._allocator = allocator - self._cached_computed_seq_blocks: Dict[int, Tuple[List[int], - bool]] = {} - - def add_seq(self, seq_id: int) -> None: - """Start tracking seq_id - """ - assert seq_id not in self._cached_computed_seq_blocks - self._cached_computed_seq_blocks[seq_id] = ([], False) - - def remove_seq(self, seq_id: int) -> None: - """Stop tracking seq_id - """ - assert seq_id in self._cached_computed_seq_blocks - del self._cached_computed_seq_blocks[seq_id] - - def get_cached_computed_blocks_and_update( - self, seq_id: int, block_ids: List[int]) -> List[int]: - """ Look at the class documentation for details - """ - # Ensure seq_id is already tracked - assert seq_id in self._cached_computed_seq_blocks - - # Get cached data (may be empty on the first time) - prev_computed_block_ids, has_gap = self._cached_computed_seq_blocks[ - seq_id] - - if has_gap: - # When gap is detected, we do not add more computed blocks at this - # sequence iteration - return prev_computed_block_ids - - # We do not consider the last block id for caching purposes. - num_cur_blocks = len(block_ids) - 1 - assert num_cur_blocks >= 0 - - if len(prev_computed_block_ids) >= num_cur_blocks: - # Cache HIT - assert len(prev_computed_block_ids) == num_cur_blocks - return prev_computed_block_ids - - # If here, then we may possibly add more computed blocks. As a result, - # traverse the additional blocks after prev_computed_block_ids to - # detect more computed blocks and add them. - - # Incremental init for seq_id => Look only at the new blocks - computed_block_ids = self._allocator.get_computed_block_ids( # noqa: E501 - prev_computed_block_ids, - block_ids, - skip_last_block_id= - True, # We skip last block id to avoid caching of full seq - ) - - # Detect if there is a "gap" - has_gap = len(computed_block_ids) < num_cur_blocks - - # Record - self._cached_computed_seq_blocks[seq_id] = (computed_block_ids, - has_gap) - - return computed_block_ids - - -class LastAccessBlocksTracker: - """Manages the last access time of the tracked sequences, in order to allow - an efficient update of allocator's block last access times - """ - - def __init__(self, allocator): - self._allocator = allocator - self._seq_last_access: Dict[int, Optional[float]] = {} - - def add_seq(self, seq_id: int) -> None: - """Start tracking seq_id - """ - assert seq_id not in self._seq_last_access - self._seq_last_access[seq_id] = None - - def remove_seq(self, seq_id: int) -> None: - """Stop tracking seq_id - """ - assert seq_id in self._seq_last_access - del self._seq_last_access[seq_id] - - def update_last_access(self, seq_id: int, time: float) -> None: - assert seq_id in self._seq_last_access - self._seq_last_access[seq_id] = time - - def update_seq_blocks_last_access(self, seq_id: int, - block_ids: List[int]) -> None: - assert seq_id in self._seq_last_access - - ts = self._seq_last_access[seq_id] - - if ts is None: - # No last access was recorded, no need to update. - return - - self._allocator.mark_blocks_as_accessed(block_ids, ts) - - -def assert_prefix_caching_block_or_none(block: Optional[Block]): - if block is None: - return - assert isinstance(block, - PrefixCachingBlock), "Got block = {}".format(block) + + +class ComputedBlocksTracker: + """Handles caching of per-sequence computed block ids. + When a sequence appears for the first time, it traverses all of the + blocks and detects the prefix of blocks that is computed. On the + subsequent times, it only traverses the new blocks that were added + and updates the already recorded prefix of blocks with the newly + computed blocks. + + To avoid redundant traversals, the algorithm also detects when there + is a "gap" in the computed prefix. For example, if we have blocks = + [1,2,3,4,5], and we have detected [1,2,3] as the computed prefix, then + we won't try to add more computed blocks to [1,2,3] in this sequence + iteration, and will add more computed blocks only after the sequence is + freed and reused again. + + Note that currently, for a given sequence, we also skip the last + block id for caching purposes, to avoid caching of a full sequence + """ + + def __init__(self, allocator): + self._allocator = allocator + self._cached_computed_seq_blocks: Dict[int, Tuple[List[int], + bool]] = {} + + def add_seq(self, seq_id: int) -> None: + """Start tracking seq_id + """ + assert seq_id not in self._cached_computed_seq_blocks + self._cached_computed_seq_blocks[seq_id] = ([], False) + + def remove_seq(self, seq_id: int) -> None: + """Stop tracking seq_id + """ + assert seq_id in self._cached_computed_seq_blocks + del self._cached_computed_seq_blocks[seq_id] + + def get_cached_computed_blocks_and_update( + self, seq_id: int, block_ids: List[int]) -> List[int]: + """ Look at the class documentation for details + """ + # Ensure seq_id is already tracked + assert seq_id in self._cached_computed_seq_blocks + + # Get cached data (may be empty on the first time) + prev_computed_block_ids, has_gap = self._cached_computed_seq_blocks[ + seq_id] + + if has_gap: + # When gap is detected, we do not add more computed blocks at this + # sequence iteration + return prev_computed_block_ids + + # We do not consider the last block id for caching purposes. + num_cur_blocks = len(block_ids) - 1 + assert num_cur_blocks >= 0 + + if len(prev_computed_block_ids) >= num_cur_blocks: + # Cache HIT + assert len(prev_computed_block_ids) == num_cur_blocks + return prev_computed_block_ids + + # If here, then we may possibly add more computed blocks. As a result, + # traverse the additional blocks after prev_computed_block_ids to + # detect more computed blocks and add them. + + # Incremental init for seq_id => Look only at the new blocks + computed_block_ids = self._allocator.get_computed_block_ids( # noqa: E501 + prev_computed_block_ids, + block_ids, + skip_last_block_id= + True, # We skip last block id to avoid caching of full seq + ) + + # Detect if there is a "gap" + has_gap = len(computed_block_ids) < num_cur_blocks + + # Record + self._cached_computed_seq_blocks[seq_id] = (computed_block_ids, + has_gap) + + return computed_block_ids + + +class LastAccessBlocksTracker: + """Manages the last access time of the tracked sequences, in order to allow + an efficient update of allocator's block last access times + """ + + def __init__(self, allocator): + self._allocator = allocator + self._seq_last_access: Dict[int, Optional[float]] = {} + + def add_seq(self, seq_id: int) -> None: + """Start tracking seq_id + """ + assert seq_id not in self._seq_last_access + self._seq_last_access[seq_id] = None + + def remove_seq(self, seq_id: int) -> None: + """Stop tracking seq_id + """ + assert seq_id in self._seq_last_access + del self._seq_last_access[seq_id] + + def update_last_access(self, seq_id: int, time: float) -> None: + assert seq_id in self._seq_last_access + self._seq_last_access[seq_id] = time + + def update_seq_blocks_last_access(self, seq_id: int, + block_ids: List[int]) -> None: + assert seq_id in self._seq_last_access + + ts = self._seq_last_access[seq_id] + + if ts is None: + # No last access was recorded, no need to update. + return + + self._allocator.mark_blocks_as_accessed(block_ids, ts) + + +def assert_prefix_caching_block_or_none(block: Optional[Block]): + if block is None: + return + assert isinstance(block, + PrefixCachingBlock), "Got block = {}".format(block) diff --git a/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py b/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py index 321a60d2..5f9254f0 100644 --- a/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py +++ b/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py @@ -32,86 +32,86 @@ logger = init_logger(__name__) class BlockSpaceManagerV2(BlockSpaceManager): - """BlockSpaceManager which manages the allocation of KV cache. - - It owns responsibility for allocation, swapping, allocating memory for - autoregressively-generated tokens, and other advanced features such as - prefix caching, forking/copy-on-write, and sliding-window memory allocation. - - This class implements the design described in - https://github.com/vllm-project/vllm/pull/3492. - - Lookahead slots - The block manager has the notion of a "lookahead slot". These are slots - in the KV cache that are allocated for a sequence. Unlike the other - allocated slots, the content of these slots is undefined -- the worker - may use the memory allocations in any way. - - In practice, a worker could use these lookahead slots to run multiple - forward passes for a single scheduler invocation. Each successive - forward pass would write KV activations to the corresponding lookahead - slot. This allows low inter-token latency use-cases, where the overhead - of continuous batching scheduling is amortized over >1 generated tokens. - - Speculative decoding uses lookahead slots to store KV activations of - proposal tokens. - - See https://github.com/vllm-project/vllm/pull/3250 for more information - on lookahead scheduling. - - Args: - block_size (int): The size of each memory block. - num_gpu_blocks (int): The number of memory blocks allocated on GPU. - num_cpu_blocks (int): The number of memory blocks allocated on CPU. - watermark (float, optional): The threshold used for memory swapping. - Defaults to 0.01. - sliding_window (Optional[int], optional): The size of the sliding - window. Defaults to None. - enable_caching (bool, optional): Flag indicating whether caching is - enabled. Defaults to False. - """ - + """BlockSpaceManager which manages the allocation of KV cache. + + It owns responsibility for allocation, swapping, allocating memory for + autoregressively-generated tokens, and other advanced features such as + prefix caching, forking/copy-on-write, and sliding-window memory allocation. + + This class implements the design described in + https://github.com/vllm-project/vllm/pull/3492. + + Lookahead slots + The block manager has the notion of a "lookahead slot". These are slots + in the KV cache that are allocated for a sequence. Unlike the other + allocated slots, the content of these slots is undefined -- the worker + may use the memory allocations in any way. + + In practice, a worker could use these lookahead slots to run multiple + forward passes for a single scheduler invocation. Each successive + forward pass would write KV activations to the corresponding lookahead + slot. This allows low inter-token latency use-cases, where the overhead + of continuous batching scheduling is amortized over >1 generated tokens. + + Speculative decoding uses lookahead slots to store KV activations of + proposal tokens. + + See https://github.com/vllm-project/vllm/pull/3250 for more information + on lookahead scheduling. + + Args: + block_size (int): The size of each memory block. + num_gpu_blocks (int): The number of memory blocks allocated on GPU. + num_cpu_blocks (int): The number of memory blocks allocated on CPU. + watermark (float, optional): The threshold used for memory swapping. + Defaults to 0.01. + sliding_window (Optional[int], optional): The size of the sliding + window. Defaults to None. + enable_caching (bool, optional): Flag indicating whether caching is + enabled. Defaults to False. + """ + def __init__( self, block_size: int, num_gpu_blocks: int, num_cpu_blocks: int, - watermark: float = 0.01, - sliding_window: Optional[int] = None, - enable_caching: bool = False, - ) -> None: - self.block_size = block_size - self.num_total_gpu_blocks = num_gpu_blocks - self.num_total_cpu_blocks = num_cpu_blocks - - self.sliding_window = sliding_window - # max_block_sliding_window is the max number of blocks that need to be - # allocated - self.max_block_sliding_window = None - if sliding_window is not None: - # +1 here because // rounds down - num_blocks = sliding_window // block_size + 1 - # +1 here because the last block may not be full, - # and so the sequence stretches one more block at the beginning - # For example, if sliding_window is 3 and block_size is 4, - # we may need 2 blocks when the second block only holds 1 token. - self.max_block_sliding_window = num_blocks + 1 - - self.watermark = watermark - assert watermark >= 0.0 - - self.enable_caching = enable_caching - - self.watermark_blocks = int(watermark * num_gpu_blocks) - - self.block_allocator = CpuGpuBlockAllocator.create( - allocator_type="prefix_caching" if enable_caching else "naive", - num_gpu_blocks=num_gpu_blocks, - num_cpu_blocks=num_cpu_blocks, - block_size=block_size, - ) - - self.block_tables: Dict[SeqId, BlockTable] = {} + watermark: float = 0.01, + sliding_window: Optional[int] = None, + enable_caching: bool = False, + ) -> None: + self.block_size = block_size + self.num_total_gpu_blocks = num_gpu_blocks + self.num_total_cpu_blocks = num_cpu_blocks + + self.sliding_window = sliding_window + # max_block_sliding_window is the max number of blocks that need to be + # allocated + self.max_block_sliding_window = None + if sliding_window is not None: + # +1 here because // rounds down + num_blocks = sliding_window // block_size + 1 + # +1 here because the last block may not be full, + # and so the sequence stretches one more block at the beginning + # For example, if sliding_window is 3 and block_size is 4, + # we may need 2 blocks when the second block only holds 1 token. + self.max_block_sliding_window = num_blocks + 1 + + self.watermark = watermark + assert watermark >= 0.0 + + self.enable_caching = enable_caching + + self.watermark_blocks = int(watermark * num_gpu_blocks) + + self.block_allocator = CpuGpuBlockAllocator.create( + allocator_type="prefix_caching" if enable_caching else "naive", + num_gpu_blocks=num_gpu_blocks, + num_cpu_blocks=num_cpu_blocks, + block_size=block_size, + ) + + self.block_tables: Dict[SeqId, BlockTable] = {} self.cross_block_tables: Dict[EncoderSeqId, BlockTable] = {} self._warned_mm_namespace_requests = set[str]() self._request_local_namespace: Dict[str, bytes] = {} @@ -121,46 +121,46 @@ class BlockSpaceManagerV2(BlockSpaceManager): self.block_allocator) self._last_access_blocks_tracker = LastAccessBlocksTracker( self.block_allocator) - - def can_allocate(self, - seq_group: SequenceGroup, - num_lookahead_slots: int = 0) -> AllocStatus: - # FIXME(woosuk): Here we assume that all sequences in the group share - # the same prompt. This may not be true for preempted sequences. - - check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group) - - seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0] - num_required_blocks = BlockTable.get_num_required_blocks( - seq.get_token_ids(), - block_size=self.block_size, - num_lookahead_slots=num_lookahead_slots, - ) - - if seq_group.is_encoder_decoder(): - encoder_seq = seq_group.get_encoder_seq() - assert encoder_seq is not None - num_required_blocks += BlockTable.get_num_required_blocks( - encoder_seq.get_token_ids(), - block_size=self.block_size, - ) - - if self.max_block_sliding_window is not None: - num_required_blocks = min(num_required_blocks, - self.max_block_sliding_window) - - num_free_gpu_blocks = self.block_allocator.get_num_free_blocks( - device=Device.GPU) - - # Use watermark to avoid frequent cache eviction. - if (self.num_total_gpu_blocks - num_required_blocks < - self.watermark_blocks): - return AllocStatus.NEVER - if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks: - return AllocStatus.OK - else: - return AllocStatus.LATER - + + def can_allocate(self, + seq_group: SequenceGroup, + num_lookahead_slots: int = 0) -> AllocStatus: + # FIXME(woosuk): Here we assume that all sequences in the group share + # the same prompt. This may not be true for preempted sequences. + + check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group) + + seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0] + num_required_blocks = BlockTable.get_num_required_blocks( + seq.get_token_ids(), + block_size=self.block_size, + num_lookahead_slots=num_lookahead_slots, + ) + + if seq_group.is_encoder_decoder(): + encoder_seq = seq_group.get_encoder_seq() + assert encoder_seq is not None + num_required_blocks += BlockTable.get_num_required_blocks( + encoder_seq.get_token_ids(), + block_size=self.block_size, + ) + + if self.max_block_sliding_window is not None: + num_required_blocks = min(num_required_blocks, + self.max_block_sliding_window) + + num_free_gpu_blocks = self.block_allocator.get_num_free_blocks( + device=Device.GPU) + + # Use watermark to avoid frequent cache eviction. + if (self.num_total_gpu_blocks - num_required_blocks < + self.watermark_blocks): + return AllocStatus.NEVER + if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks: + return AllocStatus.OK + else: + return AllocStatus.LATER + def _allocate_sequence( self, seq: Sequence, @@ -181,10 +181,10 @@ class BlockSpaceManagerV2(BlockSpaceManager): def allocate(self, seq_group: SequenceGroup) -> None: # Allocate self-attention block tables for decoder sequences - waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING) - assert not (set(seq.seq_id for seq in waiting_seqs) - & self.block_tables.keys()), "block table already exists" - + waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING) + assert not (set(seq.seq_id for seq in waiting_seqs) + & self.block_tables.keys()), "block table already exists" + # NOTE: Here we assume that all sequences in the group have the same # prompt. seq = waiting_seqs[0] @@ -199,31 +199,31 @@ class BlockSpaceManagerV2(BlockSpaceManager): cache_namespace=cache_namespace, ) self.block_tables[seq.seq_id] = block_table - - # Track seq + + # Track seq self._computed_blocks_tracker.add_seq(seq.seq_id) self._last_access_blocks_tracker.add_seq(seq.seq_id) - - # Assign the block table for each sequence. - for seq in waiting_seqs[1:]: - self.block_tables[seq.seq_id] = block_table.fork() - - # Track seq - self._computed_blocks_tracker.add_seq(seq.seq_id) - self._last_access_blocks_tracker.add_seq(seq.seq_id) - - # Allocate cross-attention block table for encoder sequence - # - # NOTE: Here we assume that all sequences in the group have the same - # encoder prompt. + + # Assign the block table for each sequence. + for seq in waiting_seqs[1:]: + self.block_tables[seq.seq_id] = block_table.fork() + + # Track seq + self._computed_blocks_tracker.add_seq(seq.seq_id) + self._last_access_blocks_tracker.add_seq(seq.seq_id) + + # Allocate cross-attention block table for encoder sequence + # + # NOTE: Here we assume that all sequences in the group have the same + # encoder prompt. request_id = seq_group.request_id assert (request_id not in self.cross_block_tables), \ "block table already exists" - - check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group) - + + check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group) + if seq_group.is_encoder_decoder(): encoder_seq = seq_group.get_encoder_seq() assert encoder_seq is not None @@ -279,129 +279,129 @@ class BlockSpaceManagerV2(BlockSpaceManager): else: digest.update(b"text|") return digest.digest() - - def can_append_slots(self, seq_group: SequenceGroup, - num_lookahead_slots: int) -> bool: - """Determine if there is enough space in the GPU KV cache to continue - generation of the specified sequence group. - - We use a worst-case heuristic: assume each touched block will require a - new allocation (either via CoW or new block). We can append slots if the - number of touched blocks is less than the number of free blocks. - - "Lookahead slots" are slots that are allocated in addition to the slots - for known tokens. The contents of the lookahead slots are not defined. - This is used by speculative decoding when speculating future tokens. - """ - - num_touched_blocks = 0 - for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): - block_table = self.block_tables[seq.seq_id] - - num_touched_blocks += ( - block_table.get_num_blocks_touched_by_append_slots( - token_ids=block_table.get_unseen_token_ids( - seq.get_token_ids()), - num_lookahead_slots=num_lookahead_slots, - )) - - num_free_gpu_blocks = self.block_allocator.get_num_free_blocks( - Device.GPU) - return num_touched_blocks <= num_free_gpu_blocks - - def append_slots( - self, - seq: Sequence, - num_lookahead_slots: int, - ) -> List[Tuple[int, int]]: - - block_table = self.block_tables[seq.seq_id] - - block_table.append_token_ids( - token_ids=block_table.get_unseen_token_ids(seq.get_token_ids()), - num_lookahead_slots=num_lookahead_slots, - num_computed_slots=seq.data.get_num_computed_tokens(), - ) - # Return any new copy-on-writes. - new_cows = self.block_allocator.clear_copy_on_writes() - return new_cows - - def free(self, seq: Sequence) -> None: - seq_id = seq.seq_id - - if seq_id not in self.block_tables: - # Already freed or haven't been scheduled yet. - return - - # Update seq block ids with the latest access time - self._last_access_blocks_tracker.update_seq_blocks_last_access( - seq_id, self.block_tables[seq.seq_id].physical_block_ids) - - # Untrack seq - self._last_access_blocks_tracker.remove_seq(seq_id) - self._computed_blocks_tracker.remove_seq(seq_id) - - # Free table/blocks - self.block_tables[seq_id].free() - del self.block_tables[seq_id] - - def free_cross(self, seq_group: SequenceGroup) -> None: - request_id = seq_group.request_id - if request_id not in self.cross_block_tables: - # Already freed or hasn't been scheduled yet. - return - self.cross_block_tables[request_id].free() - del self.cross_block_tables[request_id] - + + def can_append_slots(self, seq_group: SequenceGroup, + num_lookahead_slots: int) -> bool: + """Determine if there is enough space in the GPU KV cache to continue + generation of the specified sequence group. + + We use a worst-case heuristic: assume each touched block will require a + new allocation (either via CoW or new block). We can append slots if the + number of touched blocks is less than the number of free blocks. + + "Lookahead slots" are slots that are allocated in addition to the slots + for known tokens. The contents of the lookahead slots are not defined. + This is used by speculative decoding when speculating future tokens. + """ + + num_touched_blocks = 0 + for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): + block_table = self.block_tables[seq.seq_id] + + num_touched_blocks += ( + block_table.get_num_blocks_touched_by_append_slots( + token_ids=block_table.get_unseen_token_ids( + seq.get_token_ids()), + num_lookahead_slots=num_lookahead_slots, + )) + + num_free_gpu_blocks = self.block_allocator.get_num_free_blocks( + Device.GPU) + return num_touched_blocks <= num_free_gpu_blocks + + def append_slots( + self, + seq: Sequence, + num_lookahead_slots: int, + ) -> List[Tuple[int, int]]: + + block_table = self.block_tables[seq.seq_id] + + block_table.append_token_ids( + token_ids=block_table.get_unseen_token_ids(seq.get_token_ids()), + num_lookahead_slots=num_lookahead_slots, + num_computed_slots=seq.data.get_num_computed_tokens(), + ) + # Return any new copy-on-writes. + new_cows = self.block_allocator.clear_copy_on_writes() + return new_cows + + def free(self, seq: Sequence) -> None: + seq_id = seq.seq_id + + if seq_id not in self.block_tables: + # Already freed or haven't been scheduled yet. + return + + # Update seq block ids with the latest access time + self._last_access_blocks_tracker.update_seq_blocks_last_access( + seq_id, self.block_tables[seq.seq_id].physical_block_ids) + + # Untrack seq + self._last_access_blocks_tracker.remove_seq(seq_id) + self._computed_blocks_tracker.remove_seq(seq_id) + + # Free table/blocks + self.block_tables[seq_id].free() + del self.block_tables[seq_id] + + def free_cross(self, seq_group: SequenceGroup) -> None: + request_id = seq_group.request_id + if request_id not in self.cross_block_tables: + # Already freed or hasn't been scheduled yet. + return + self.cross_block_tables[request_id].free() + del self.cross_block_tables[request_id] + def get_block_table(self, seq: Sequence) -> List[int]: block_ids = self.block_tables[seq.seq_id].physical_block_ids return block_ids # type: ignore def get_cross_block_table(self, seq_group: SequenceGroup) -> List[int]: - request_id = seq_group.request_id - assert request_id in self.cross_block_tables - block_ids = self.cross_block_tables[request_id].physical_block_ids - assert all(b is not None for b in block_ids) - return block_ids # type: ignore - - def access_all_blocks_in_seq(self, seq: Sequence, now: float): - if self.enable_caching: - # Record the latest access time for the sequence. The actual update - # of the block ids is deferred to the sequence free(..) call, since - # only during freeing of block ids, the blocks are actually added to - # the evictor (which is when the most updated time is required) - # (This avoids expensive calls to mark_blocks_as_accessed(..)) - self._last_access_blocks_tracker.update_last_access( - seq.seq_id, now) - - def mark_blocks_as_computed(self, seq_group: SequenceGroup, - token_chunk_size: int): - # If prefix caching is enabled, mark immutable blocks as computed - # right after they have been scheduled (for prefill). This assumes - # the scheduler is synchronous so blocks are actually computed when - # scheduling the next batch. - self.block_allocator.mark_blocks_as_computed([]) - + request_id = seq_group.request_id + assert request_id in self.cross_block_tables + block_ids = self.cross_block_tables[request_id].physical_block_ids + assert all(b is not None for b in block_ids) + return block_ids # type: ignore + + def access_all_blocks_in_seq(self, seq: Sequence, now: float): + if self.enable_caching: + # Record the latest access time for the sequence. The actual update + # of the block ids is deferred to the sequence free(..) call, since + # only during freeing of block ids, the blocks are actually added to + # the evictor (which is when the most updated time is required) + # (This avoids expensive calls to mark_blocks_as_accessed(..)) + self._last_access_blocks_tracker.update_last_access( + seq.seq_id, now) + + def mark_blocks_as_computed(self, seq_group: SequenceGroup, + token_chunk_size: int): + # If prefix caching is enabled, mark immutable blocks as computed + # right after they have been scheduled (for prefill). This assumes + # the scheduler is synchronous so blocks are actually computed when + # scheduling the next batch. + self.block_allocator.mark_blocks_as_computed([]) + def get_common_computed_block_ids( self, seqs: List[Sequence]) -> GenericSequence[int]: - """Determine which blocks for which we skip prefill. - - With prefix caching we can skip prefill for previously-generated blocks. - Currently, the attention implementation only supports skipping cached - blocks if they are a contiguous prefix of cached blocks. - - This method determines which blocks can be safely skipped for all - sequences in the sequence group. - """ - computed_seq_block_ids = [] - for seq in seqs: - computed_seq_block_ids.append( - self._computed_blocks_tracker. - get_cached_computed_blocks_and_update( - seq.seq_id, - self.block_tables[seq.seq_id].physical_block_ids)) - - # NOTE(sang): This assumes seq_block_ids doesn't contain any None. + """Determine which blocks for which we skip prefill. + + With prefix caching we can skip prefill for previously-generated blocks. + Currently, the attention implementation only supports skipping cached + blocks if they are a contiguous prefix of cached blocks. + + This method determines which blocks can be safely skipped for all + sequences in the sequence group. + """ + computed_seq_block_ids = [] + for seq in seqs: + computed_seq_block_ids.append( + self._computed_blocks_tracker. + get_cached_computed_blocks_and_update( + seq.seq_id, + self.block_tables[seq.seq_id].physical_block_ids)) + + # NOTE(sang): This assumes seq_block_ids doesn't contain any None. return self.block_allocator.get_common_computed_block_ids( computed_seq_block_ids) # type: ignore @@ -583,187 +583,187 @@ class BlockSpaceManagerV2(BlockSpaceManager): raise TypeError(f"Unsupported multimodal namespace value type {type(value)}") - - def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None: - if parent_seq.seq_id not in self.block_tables: - # Parent sequence has either been freed or never existed. - return - src_block_table = self.block_tables[parent_seq.seq_id] - self.block_tables[child_seq.seq_id] = src_block_table.fork() - - # Track child seq - self._computed_blocks_tracker.add_seq(child_seq.seq_id) - self._last_access_blocks_tracker.add_seq(child_seq.seq_id) - + + def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None: + if parent_seq.seq_id not in self.block_tables: + # Parent sequence has either been freed or never existed. + return + src_block_table = self.block_tables[parent_seq.seq_id] + self.block_tables[child_seq.seq_id] = src_block_table.fork() + + # Track child seq + self._computed_blocks_tracker.add_seq(child_seq.seq_id) + self._last_access_blocks_tracker.add_seq(child_seq.seq_id) + def can_swap_in(self, seq_group: SequenceGroup, num_lookahead_slots: int) -> AllocStatus: - """Returns the AllocStatus for the given sequence_group - with num_lookahead_slots. - - Args: - sequence_group (SequenceGroup): The sequence group to swap in. - num_lookahead_slots (int): Number of lookahead slots used in - speculative decoding, default to 0. - - Returns: - AllocStatus: The AllocStatus for the given sequence group. - """ + """Returns the AllocStatus for the given sequence_group + with num_lookahead_slots. + + Args: + sequence_group (SequenceGroup): The sequence group to swap in. + num_lookahead_slots (int): Number of lookahead slots used in + speculative decoding, default to 0. + + Returns: + AllocStatus: The AllocStatus for the given sequence group. + """ if self.block_allocator.content_offload_enabled: return AllocStatus.NEVER return self._can_swap(seq_group, Device.GPU, SequenceStatus.SWAPPED, num_lookahead_slots) - - def swap_in(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]: - """Returns the block id mapping (from CPU to GPU) generated by - swapping in the given seq_group with num_lookahead_slots. - - Args: - seq_group (SequenceGroup): The sequence group to swap in. - - Returns: - List[Tuple[int, int]]: The mapping of swapping block from CPU - to GPU. - """ - physical_block_id_mapping = [] - for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED): - blocks = self.block_tables[seq.seq_id].blocks - if len(blocks) == 0: - continue - - seq_swap_mapping = self.block_allocator.swap(blocks=blocks, - src_device=Device.CPU, - dst_device=Device.GPU) - - # Refresh the block ids of the table (post-swap) - self.block_tables[seq.seq_id].update(blocks) - - seq_physical_block_id_mapping = { - self.block_allocator.get_physical_block_id( - Device.CPU, cpu_block_id): - self.block_allocator.get_physical_block_id( - Device.GPU, gpu_block_id) - for cpu_block_id, gpu_block_id in seq_swap_mapping.items() - } - - physical_block_id_mapping.extend( - list(seq_physical_block_id_mapping.items())) - - return physical_block_id_mapping - + + def swap_in(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]: + """Returns the block id mapping (from CPU to GPU) generated by + swapping in the given seq_group with num_lookahead_slots. + + Args: + seq_group (SequenceGroup): The sequence group to swap in. + + Returns: + List[Tuple[int, int]]: The mapping of swapping block from CPU + to GPU. + """ + physical_block_id_mapping = [] + for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED): + blocks = self.block_tables[seq.seq_id].blocks + if len(blocks) == 0: + continue + + seq_swap_mapping = self.block_allocator.swap(blocks=blocks, + src_device=Device.CPU, + dst_device=Device.GPU) + + # Refresh the block ids of the table (post-swap) + self.block_tables[seq.seq_id].update(blocks) + + seq_physical_block_id_mapping = { + self.block_allocator.get_physical_block_id( + Device.CPU, cpu_block_id): + self.block_allocator.get_physical_block_id( + Device.GPU, gpu_block_id) + for cpu_block_id, gpu_block_id in seq_swap_mapping.items() + } + + physical_block_id_mapping.extend( + list(seq_physical_block_id_mapping.items())) + + return physical_block_id_mapping + def can_swap_out(self, seq_group: SequenceGroup) -> bool: - """Returns whether we can swap out the given sequence_group - with num_lookahead_slots. - - Args: - seq_group (SequenceGroup): The sequence group to swap in. - num_lookahead_slots (int): Number of lookahead slots used in - speculative decoding, default to 0. - - Returns: - bool: Whether it's possible to swap out current sequence group. - """ + """Returns whether we can swap out the given sequence_group + with num_lookahead_slots. + + Args: + seq_group (SequenceGroup): The sequence group to swap in. + num_lookahead_slots (int): Number of lookahead slots used in + speculative decoding, default to 0. + + Returns: + bool: Whether it's possible to swap out current sequence group. + """ if self.block_allocator.content_offload_enabled: return False alloc_status = self._can_swap(seq_group, Device.CPU, SequenceStatus.RUNNING) - return alloc_status == AllocStatus.OK - - def swap_out(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]: - """Returns the block id mapping (from GPU to CPU) generated by - swapping out the given sequence_group with num_lookahead_slots. - - Args: - sequence_group (SequenceGroup): The sequence group to swap in. - - Returns: - List[Tuple[int, int]]: The mapping of swapping block from - GPU to CPU. - """ - physical_block_id_mapping = [] - for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): - blocks = self.block_tables[seq.seq_id].blocks - if len(blocks) == 0: - continue - - seq_swap_mapping = self.block_allocator.swap(blocks=blocks, - src_device=Device.GPU, - dst_device=Device.CPU) - - # Refresh the block ids of the table (post-swap) - self.block_tables[seq.seq_id].update(blocks) - - seq_physical_block_id_mapping = { - self.block_allocator.get_physical_block_id( - Device.GPU, gpu_block_id): - self.block_allocator.get_physical_block_id( - Device.CPU, cpu_block_id) - for gpu_block_id, cpu_block_id in seq_swap_mapping.items() - } - - physical_block_id_mapping.extend( - list(seq_physical_block_id_mapping.items())) - - return physical_block_id_mapping - - def get_num_free_gpu_blocks(self) -> int: - return self.block_allocator.get_num_free_blocks(Device.GPU) - - def get_num_free_cpu_blocks(self) -> int: - return self.block_allocator.get_num_free_blocks(Device.CPU) - - def get_prefix_cache_hit_rate(self, device: Device) -> float: - return self.block_allocator.get_prefix_cache_hit_rate(device) - - def _can_swap(self, - seq_group: SequenceGroup, - device: Device, - status: SequenceStatus, - num_lookahead_slots: int = 0) -> AllocStatus: - """Returns the AllocStatus for swapping in/out the given sequence_group - on to the 'device'. - - Args: - sequence_group (SequenceGroup): The sequence group to swap in. - device (Device): device to swap the 'seq_group' on. - status (SequenceStatus): The status of sequence which is needed - for action. RUNNING for swap out and SWAPPED for swap in - num_lookahead_slots (int): Number of lookahead slots used in - speculative decoding, default to 0. - - Returns: - AllocStatus: The AllocStatus for swapping in/out the given - sequence_group on to the 'device'. - """ - # First determine the number of blocks that will be touched by this - # swap. Then verify if there are available blocks in the device - # to perform the swap. - num_blocks_touched = 0 - blocks: List[Block] = [] - for seq in seq_group.get_seqs(status=status): - block_table = self.block_tables[seq.seq_id] - if block_table.blocks is not None: - # Compute the number blocks to touch for the tokens to be - # appended. This does NOT include the full blocks that need - # to be touched for the swap. - num_blocks_touched += \ - block_table.get_num_blocks_touched_by_append_slots( - block_table.get_unseen_token_ids(seq.get_token_ids()), - num_lookahead_slots=num_lookahead_slots) - blocks.extend(block_table.blocks) - # Compute the number of full blocks to touch and add it to the - # existing count of blocks to touch. - num_blocks_touched += self.block_allocator.get_num_full_blocks_touched( - blocks, device=device) - - watermark_blocks = 0 - if device == Device.GPU: - watermark_blocks = self.watermark_blocks - - if self.block_allocator.get_num_total_blocks( - device) < num_blocks_touched: - return AllocStatus.NEVER - elif self.block_allocator.get_num_free_blocks( - device) - num_blocks_touched >= watermark_blocks: - return AllocStatus.OK - else: - return AllocStatus.LATER + return alloc_status == AllocStatus.OK + + def swap_out(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]: + """Returns the block id mapping (from GPU to CPU) generated by + swapping out the given sequence_group with num_lookahead_slots. + + Args: + sequence_group (SequenceGroup): The sequence group to swap in. + + Returns: + List[Tuple[int, int]]: The mapping of swapping block from + GPU to CPU. + """ + physical_block_id_mapping = [] + for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING): + blocks = self.block_tables[seq.seq_id].blocks + if len(blocks) == 0: + continue + + seq_swap_mapping = self.block_allocator.swap(blocks=blocks, + src_device=Device.GPU, + dst_device=Device.CPU) + + # Refresh the block ids of the table (post-swap) + self.block_tables[seq.seq_id].update(blocks) + + seq_physical_block_id_mapping = { + self.block_allocator.get_physical_block_id( + Device.GPU, gpu_block_id): + self.block_allocator.get_physical_block_id( + Device.CPU, cpu_block_id) + for gpu_block_id, cpu_block_id in seq_swap_mapping.items() + } + + physical_block_id_mapping.extend( + list(seq_physical_block_id_mapping.items())) + + return physical_block_id_mapping + + def get_num_free_gpu_blocks(self) -> int: + return self.block_allocator.get_num_free_blocks(Device.GPU) + + def get_num_free_cpu_blocks(self) -> int: + return self.block_allocator.get_num_free_blocks(Device.CPU) + + def get_prefix_cache_hit_rate(self, device: Device) -> float: + return self.block_allocator.get_prefix_cache_hit_rate(device) + + def _can_swap(self, + seq_group: SequenceGroup, + device: Device, + status: SequenceStatus, + num_lookahead_slots: int = 0) -> AllocStatus: + """Returns the AllocStatus for swapping in/out the given sequence_group + on to the 'device'. + + Args: + sequence_group (SequenceGroup): The sequence group to swap in. + device (Device): device to swap the 'seq_group' on. + status (SequenceStatus): The status of sequence which is needed + for action. RUNNING for swap out and SWAPPED for swap in + num_lookahead_slots (int): Number of lookahead slots used in + speculative decoding, default to 0. + + Returns: + AllocStatus: The AllocStatus for swapping in/out the given + sequence_group on to the 'device'. + """ + # First determine the number of blocks that will be touched by this + # swap. Then verify if there are available blocks in the device + # to perform the swap. + num_blocks_touched = 0 + blocks: List[Block] = [] + for seq in seq_group.get_seqs(status=status): + block_table = self.block_tables[seq.seq_id] + if block_table.blocks is not None: + # Compute the number blocks to touch for the tokens to be + # appended. This does NOT include the full blocks that need + # to be touched for the swap. + num_blocks_touched += \ + block_table.get_num_blocks_touched_by_append_slots( + block_table.get_unseen_token_ids(seq.get_token_ids()), + num_lookahead_slots=num_lookahead_slots) + blocks.extend(block_table.blocks) + # Compute the number of full blocks to touch and add it to the + # existing count of blocks to touch. + num_blocks_touched += self.block_allocator.get_num_full_blocks_touched( + blocks, device=device) + + watermark_blocks = 0 + if device == Device.GPU: + watermark_blocks = self.watermark_blocks + + if self.block_allocator.get_num_total_blocks( + device) < num_blocks_touched: + return AllocStatus.NEVER + elif self.block_allocator.get_num_free_blocks( + device) - num_blocks_touched >= watermark_blocks: + return AllocStatus.OK + else: + return AllocStatus.LATER diff --git a/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py b/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py index 393c8426..3ce89711 100644 --- a/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py +++ b/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py @@ -7,128 +7,128 @@ from typing import Dict, List, OrderedDict, Tuple ContentHash = bytes - - + + class EvictionPolicy(enum.Enum): - """Enum for eviction policy used by make_evictor to instantiate the correct - Evictor subclass. + """Enum for eviction policy used by make_evictor to instantiate the correct + Evictor subclass. """ LRU = enum.auto() FREQUENCY_AWARE = enum.auto() - - -class Evictor(ABC): - """The Evictor subclasses should be used by the BlockAllocator class to - handle eviction of freed PhysicalTokenBlocks. - """ - - @abstractmethod - def __init__(self): - pass - - @abstractmethod - def __contains__(self, block_id: int) -> bool: - pass - - @abstractmethod + + +class Evictor(ABC): + """The Evictor subclasses should be used by the BlockAllocator class to + handle eviction of freed PhysicalTokenBlocks. + """ + + @abstractmethod + def __init__(self): + pass + + @abstractmethod + def __contains__(self, block_id: int) -> bool: + pass + + @abstractmethod def evict(self) -> Tuple[int, ContentHash]: - """Runs the eviction algorithm and returns the evicted block's - content hash along with physical block id along with physical block id - """ - pass - - @abstractmethod + """Runs the eviction algorithm and returns the evicted block's + content hash along with physical block id along with physical block id + """ + pass + + @abstractmethod def add(self, block_id: int, content_hash: ContentHash, num_hashed_tokens: int, last_accessed: float): - """Adds block to the evictor, making it a candidate for eviction""" - pass - - @abstractmethod - def update(self, block_id: int, last_accessed: float): - """Update corresponding block's access time in metadata""" - pass - - @abstractmethod - def remove(self, block_id: int): - """Remove a given block id from the cache.""" - pass - - @property - @abstractmethod - def num_blocks(self) -> int: - pass - - -class BlockMetaData(): - """Data structure for storing key data describe cached block, so that - evitor could use to make its decision which one to choose for eviction - - Here we use physical block id as the dict key, as there maybe several - blocks with the same content hash, but their physical id is unique. - """ - + """Adds block to the evictor, making it a candidate for eviction""" + pass + + @abstractmethod + def update(self, block_id: int, last_accessed: float): + """Update corresponding block's access time in metadata""" + pass + + @abstractmethod + def remove(self, block_id: int): + """Remove a given block id from the cache.""" + pass + + @property + @abstractmethod + def num_blocks(self) -> int: + pass + + +class BlockMetaData(): + """Data structure for storing key data describe cached block, so that + evitor could use to make its decision which one to choose for eviction + + Here we use physical block id as the dict key, as there maybe several + blocks with the same content hash, but their physical id is unique. + """ + def __init__(self, content_hash: ContentHash, num_hashed_tokens: int, last_accessed: float): - self.content_hash = content_hash - self.num_hashed_tokens = num_hashed_tokens - self.last_accessed = last_accessed - - -class LRUEvictor(Evictor): - """Evicts in a least-recently-used order using the last_accessed timestamp - that's recorded in the PhysicalTokenBlock. If there are multiple blocks with - the same last_accessed time, then the one with the largest num_hashed_tokens - will be evicted. If two blocks each have the lowest last_accessed time and - highest num_hashed_tokens value, then one will be chose arbitrarily - """ - - def __init__(self): - self.free_table: OrderedDict[int, BlockMetaData] = OrderedDict() - - def __contains__(self, block_id: int) -> bool: - return block_id in self.free_table - + self.content_hash = content_hash + self.num_hashed_tokens = num_hashed_tokens + self.last_accessed = last_accessed + + +class LRUEvictor(Evictor): + """Evicts in a least-recently-used order using the last_accessed timestamp + that's recorded in the PhysicalTokenBlock. If there are multiple blocks with + the same last_accessed time, then the one with the largest num_hashed_tokens + will be evicted. If two blocks each have the lowest last_accessed time and + highest num_hashed_tokens value, then one will be chose arbitrarily + """ + + def __init__(self): + self.free_table: OrderedDict[int, BlockMetaData] = OrderedDict() + + def __contains__(self, block_id: int) -> bool: + return block_id in self.free_table + def evict(self) -> Tuple[int, ContentHash]: - if len(self.free_table) == 0: - raise ValueError("No usable cache memory left") - - evicted_block, evicted_block_id = None, None - # The blocks with the lowest timestamps should be placed consecutively - # at the start of OrderedDict. Loop through all these blocks to - # find the one with maximum number of hashed tokens. - for _id, block in self.free_table.items(): - if evicted_block is None: - evicted_block, evicted_block_id = block, _id - continue - if evicted_block.last_accessed < block.last_accessed: - break - if evicted_block.num_hashed_tokens < block.num_hashed_tokens: - evicted_block, evicted_block_id = block, _id - - assert evicted_block is not None - assert evicted_block_id is not None - self.free_table.pop(evicted_block_id) - - return evicted_block_id, evicted_block.content_hash - + if len(self.free_table) == 0: + raise ValueError("No usable cache memory left") + + evicted_block, evicted_block_id = None, None + # The blocks with the lowest timestamps should be placed consecutively + # at the start of OrderedDict. Loop through all these blocks to + # find the one with maximum number of hashed tokens. + for _id, block in self.free_table.items(): + if evicted_block is None: + evicted_block, evicted_block_id = block, _id + continue + if evicted_block.last_accessed < block.last_accessed: + break + if evicted_block.num_hashed_tokens < block.num_hashed_tokens: + evicted_block, evicted_block_id = block, _id + + assert evicted_block is not None + assert evicted_block_id is not None + self.free_table.pop(evicted_block_id) + + return evicted_block_id, evicted_block.content_hash + def add(self, block_id: int, content_hash: ContentHash, num_hashed_tokens: int, last_accessed: float): - self.free_table[block_id] = BlockMetaData(content_hash, - num_hashed_tokens, - last_accessed) - - def update(self, block_id: int, last_accessed: float): - self.free_table[block_id].last_accessed = last_accessed - - def remove(self, block_id: int): - if block_id not in self.free_table: - raise ValueError( - "Attempting to remove block that's not in the evictor") - self.free_table.pop(block_id) - - @property + self.free_table[block_id] = BlockMetaData(content_hash, + num_hashed_tokens, + last_accessed) + + def update(self, block_id: int, last_accessed: float): + self.free_table[block_id].last_accessed = last_accessed + + def remove(self, block_id: int): + if block_id not in self.free_table: + raise ValueError( + "Attempting to remove block that's not in the evictor") + self.free_table.pop(block_id) + + @property def num_blocks(self) -> int: return len(self.free_table) @@ -261,8 +261,8 @@ def eviction_policy_from_env( raise ValueError( "BI100_KV_EVICTION_POLICY must be one of: frequency, lru") return policies[value] - - + + def make_evictor(eviction_policy: EvictionPolicy) -> Evictor: if eviction_policy == EvictionPolicy.LRU: return LRUEvictor() diff --git a/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py b/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py index ef6db443..f47d35c4 100644 --- a/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py +++ b/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py @@ -1,1057 +1,1057 @@ -"""A layer that samples the next tokens from the model's outputs.""" -import itertools -import warnings -from dataclasses import dataclass -from importlib.util import find_spec -from math import inf -from typing import Dict, List, Optional, Tuple, Union - -import msgspec -import torch -import torch.nn as nn - -import vllm.envs as envs -from vllm.model_executor.sampling_metadata import (SamplingMetadata, - SamplingTensors, - SequenceGroupToSample) -from vllm.sampling_params import SamplingType -from vllm.sequence import (VLLM_INVALID_TOKEN_ID, - CompletionSequenceGroupOutput, Logprob, - PromptLogprobs, SampleLogprobs, SequenceOutput) -from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics - -if envs.VLLM_USE_FLASHINFER_SAMPLER and find_spec("flashinfer"): - import flashinfer.sampling - # yapf: disable - from flashinfer.sampling import ( - top_k_top_p_sampling_from_probs as flashinfer_top_k_top_p_sampling) - - # yapf: enable -else: - flashinfer_top_k_top_p_sampling = None - -# (num_token_ids, num_parent_ids) per sequence group. -SampleResultType = List[Tuple[List[int], List[int]]] - -# Types of temporary data structures used for -# computing sample_result -SampleMetadataType = Dict[SamplingType, Tuple[List[int], - List[SequenceGroupToSample]]] -MultinomialSamplesType = Dict[SamplingType, torch.Tensor] -SampleResultsDictType = Dict[int, Tuple[List[int], List[int]]] - - -# Encapsulates temporary data structures for computing -# sample_result. -# -# * For multi-step scheduling: must be returned -# by `Sampler.forward()` and used later to compute the pythonized -# sample_result -# -# * For single-step scheduling: consumed immediately -# inside `Sampler.forward()` to compute pythonized sample_result. -@dataclass -class SampleResultArgsType: - sample_metadata: SampleMetadataType - multinomial_samples: MultinomialSamplesType - sample_results_dict: SampleResultsDictType - sampling_metadata: SamplingMetadata - greedy_samples: Optional[torch.Tensor] - beam_search_logprobs: Optional[torch.Tensor] - - -# Union of non-deferred (single-step scheduling) -# vs deferred (multi-step scheduling) -# sample result types -MaybeDeferredSampleResultType = Union[SampleResultType, SampleResultArgsType] - -# Abbreviation of the _sample() return type -SampleReturnType = Tuple[MaybeDeferredSampleResultType, Optional[torch.Tensor]] - - -class SamplerOutput( - msgspec.Struct, - omit_defaults=True, # type: ignore[call-arg] - array_like=True): # type: ignore[call-arg] - """For each sequence group, we generate a list of SequenceOutput object, - each of which contains one possible candidate for the next token. - - This data structure implements methods, so it can be used like a list, but - also has optional fields for device tensors. - """ - - outputs: List[CompletionSequenceGroupOutput] - - # On-device tensor containing probabilities of each token. - sampled_token_probs: Optional[torch.Tensor] = None - - # On-device tensor containing the logprobs of each token. - logprobs: Optional["torch.Tensor"] = None - - # Holds either (1) the pythonized sampler result (single-step scheduling) - # or (2) what will be arguments for later deferred pythonization of the - # sampler result (muliti-step scheduling) - deferred_sample_results_args: Optional[SampleResultArgsType] = None - - # On-device tensor containing the sampled token ids. - sampled_token_ids: Optional[torch.Tensor] = None - # CPU tensor containing the sampled token ids. Used during multi-step to - # return the sampled token ids from last rank to AsyncLLMEngine to be - # 'broadcasted' to all other PP ranks for next step. - sampled_token_ids_cpu: Optional[torch.Tensor] = None - - # Spec decode metrics populated by workers. - spec_decode_worker_metrics: Optional[SpecDecodeWorkerMetrics] = None - - # Optional last hidden states from the model. - hidden_states: Optional[torch.Tensor] = None - - # Optional prefill hidden states from the model - # (used for models like EAGLE). - prefill_hidden_states: Optional[torch.Tensor] = None - - # Time taken in the forward pass for this across all workers - model_forward_time: Optional[float] = None - - # Time taken in the model execute function. This will include model forward, - # block/sync across workers, cpu-gpu sync time and sampling time. - model_execute_time: Optional[float] = None - - def __getitem__(self, idx: int): - return self.outputs[idx] - - def __setitem__(self, idx: int, value): - self.outputs[idx] = value - - def __len__(self): - return len(self.outputs) - - def __eq__(self, other: object): - return isinstance(other, - self.__class__) and self.outputs == other.outputs - - def __repr__(self) -> str: - """Show the shape of a tensor instead of its values to reduce noise. - """ - sampled_token_probs_repr = ("None" if self.sampled_token_probs is None - else self.sampled_token_probs.shape) - sampled_token_ids_repr = ("None" if self.sampled_token_ids is None else - self.sampled_token_ids.shape) - return ( - f"SamplerOutput(outputs={self.outputs}, " - f"sampled_token_probs={sampled_token_probs_repr}, " - f"sampled_token_ids={sampled_token_ids_repr}, " - f"spec_decode_worker_metrics={self.spec_decode_worker_metrics})") - - -class Sampler(nn.Module): - """Samples the next tokens from the model's outputs. - - This layer does the following: - 1. Discard the hidden states that are not used for sampling (i.e., all - tokens except the final one in each prompt). - 2. Compute the logits for the next tokens. - 3. Apply presence, frequency and repetition penalties. - 4. Apply temperature scaling. - 5. Apply top-p and top-k truncation. - 6. Sample the next tokens. - Here, each sequence group within the batch can have different sampling - parameters (e.g., sampling method, temperature, top-p, top-k, etc.). - - The structure of the logits tensor is coupled with the seq_groups in - sampling_metadata. Typically, each sequence in each seq_group has one row in - logits for the next token to be sampled; however, for a seq_group with a - prompt request with the prompt_logprobs sampling parameter, there are rows - in logits for each token in the input prompt. - """ - - def __init__(self): - super().__init__() - - # Whether or not the SamplerOutput should have on-device tensors - # containing the sampled token ids and probabilities. This is used by - # speculative decoding. - self.include_gpu_probs_tensor = False - self.should_modify_greedy_probs_inplace = False - - def _init_sampling_tensors( - self, - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, - ): - """The goal here is to reuse sampling tensors between similar decode - runs. This is possible because sampling logic does not change between - decodes of the same sequences. - """ - _, vocab_size = logits.shape - - # First free any existing stored sampling tensors. - # This is necessary because some sampling tensors may - # have pinned memory. - self._sampling_tensors = None - - # Initialize new sampling tensors - (sampling_tensors, do_penalties, do_top_p_top_k, - do_min_p) = SamplingTensors.from_sampling_metadata( - sampling_metadata, vocab_size, logits.device, logits.dtype) - - self._sampling_tensors = sampling_tensors - self._do_penalties = do_penalties - self._do_top_p_top_k = do_top_p_top_k - self._do_min_p = do_min_p - - def forward( - self, - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, - ) -> Optional[SamplerOutput]: - """ - Single-step scheduling: - * Perform GPU-side sampling computation & compute - GPU-side logprobs tensor - * Pythonize sampling result & logprobs tensor - - Multi-step scheduling: - * Perform GPU-side sampling computation & compute - GPU-side logprobs tensor - * Defer Pythonization of sampling result & logprobs - tensor - * Encapsulate arguments required for deferred Pythonization - in the :class:`SamplerOutput` structure - - Args: - logits: (num_tokens, vocab_size). - sampling_metadata: Metadata for sampling. - """ - assert logits is not None - _, vocab_size = logits.shape - - # Prepare sampling tensors with pinned memory to avoid blocking. - if not sampling_metadata.reuse_sampling_tensors: - self._init_sampling_tensors(logits, sampling_metadata) - elif self._do_penalties: - # In this case, the sampling tensors logic depends on - # "output_tokens" of a sequence. As a result, we cannot - # reuse sampling tensors, since "output_tokens" changes - # between decode runs. - self._init_sampling_tensors(logits, sampling_metadata) - - assert self._sampling_tensors is not None - sampling_tensors = self._sampling_tensors - do_penalties = self._do_penalties - do_top_p_top_k = self._do_top_p_top_k - do_min_p = self._do_min_p - - logits = _apply_min_tokens_penalty(logits, sampling_metadata) - - # Apply presence and frequency penalties. - if do_penalties: - logits = _apply_penalties(logits, sampling_tensors.prompt_tokens, - sampling_tensors.output_tokens, - sampling_tensors.presence_penalties, - sampling_tensors.frequency_penalties, - sampling_tensors.repetition_penalties) - - # Use float32 to apply temperature scaling. - # Use in-place division to avoid creating a new tensor. - logits = logits.to(torch.float) - logits.div_(sampling_tensors.temperatures.unsqueeze(dim=1)) - - if do_top_p_top_k and flashinfer_top_k_top_p_sampling is None: - logits = _apply_top_k_top_p(logits, sampling_tensors.top_ps, - sampling_tensors.top_ks) - - if do_min_p: - logits = _apply_min_p(logits, sampling_tensors.min_ps) - - # We use float32 for probabilities and log probabilities. - # Compute the probabilities. - probs = torch.softmax(logits, dim=-1, dtype=torch.float) - # Compute the log probabilities. - logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float) - - # Sample the next tokens. - maybe_deferred_sample_results, maybe_sampled_tokens_tensor = _sample( - probs, - logprobs, - sampling_metadata, - sampling_tensors, - include_gpu_probs_tensor=self.include_gpu_probs_tensor, - modify_greedy_probs=self._should_modify_greedy_probs_inplace, - ) - - if self.include_gpu_probs_tensor: - # Since we will defer sampler result Pythonization, - # preserve GPU-side tensors in support of later - # deferred pythonization of logprobs - assert maybe_sampled_tokens_tensor is not None - on_device_tensors = (probs, logprobs, maybe_sampled_tokens_tensor) - else: - # Since Pythonization has already happened, don't preserve - # GPU-side tensors. - on_device_tensors = None - - # Get the logprobs query results. - prompt_logprobs = None - sample_logprobs = None - if not sampling_metadata.skip_sampler_cpu_output: - # Pythonize logprobs now (GPU -> CPU); do not defer. - assert not isinstance(maybe_deferred_sample_results, - SampleResultArgsType) - prompt_logprobs, sample_logprobs = get_logprobs( - logprobs, sampling_metadata, maybe_deferred_sample_results) - - return _build_sampler_output( - maybe_deferred_sample_results, - sampling_metadata, - prompt_logprobs, - sample_logprobs, - on_device_tensors=on_device_tensors, - skip_sampler_cpu_output=sampling_metadata.skip_sampler_cpu_output) - - @property - def _should_modify_greedy_probs_inplace(self) -> bool: - """Whether or not the sampler should modify the probability distribution - of greedily-sampled tokens such that multinomial sampling would sample - the greedily-sampled token. - - In other words, if True then we set the probability of the greedily- - sampled token to 1. - - This is used by speculative decoding, which requires that the sampling - method be encoded into the probability distribution. - """ - return self.should_modify_greedy_probs_inplace - - -def _get_bin_counts_and_mask( - tokens: torch.Tensor, - vocab_size: int, - num_seqs: int, -) -> Tuple[torch.Tensor, torch.Tensor]: - # Compute the bin counts for the tokens. - # vocab_size + 1 for padding. - bin_counts = torch.zeros((num_seqs, vocab_size + 1), - dtype=torch.long, - device=tokens.device) - bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens)) - bin_counts = bin_counts[:, :vocab_size] - mask = bin_counts > 0 - - return bin_counts, mask - - -def _apply_min_tokens_penalty( - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, -) -> torch.Tensor: - """Apply min_tokens penalty which sets stop tokens to -inf if min_tokens - have not been generated yet - """ - # list of indices in logits that will be set to -inf - logits_to_penalize: List[Tuple[int, int]] = [] - logits_applied = 0 - for seq_group in sampling_metadata.seq_groups: - seq_ids = seq_group.seq_ids - sampling_params = seq_group.sampling_params - - sample_indices = seq_group.sample_indices - logits_applied += len(sample_indices) + len( - seq_group.prompt_logprob_indices) - if not seq_group.do_sample: - continue - - start_idx = sample_indices[0] - min_tokens = sampling_params.min_tokens - token_ids_to_penalize = sampling_params.all_stop_token_ids - if min_tokens > 0 and token_ids_to_penalize: - seqs_to_penalize: List[int] = [] - for j, seq_id in enumerate(seq_ids): - seq_data = seq_group.seq_data[seq_id] - if len(seq_data.output_token_ids_array) < min_tokens: - seqs_to_penalize.append(j) - - if seqs_to_penalize: - # convert to the index into logits - seqs_to_penalize = [start_idx + j for j in seqs_to_penalize] - # itertools.product pairs each seq index with every token id - logits_to_penalize.extend( - itertools.product(seqs_to_penalize, token_ids_to_penalize)) - - if logits_to_penalize: - # use zip and * to group indices along each dimension - # eg. [ (1,2), (1,3), (5,6) ] -> ( (1,1,5), (2,3,6) ) - logits[tuple(zip(*logits_to_penalize))] = -float("inf") - - # verifies that no rows in logits were missed unexpectedly - assert logits_applied == logits.shape[0] - return logits - - -def _apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor, - output_tokens_tensor: torch.Tensor, - presence_penalties: torch.Tensor, - frequency_penalties: torch.Tensor, - repetition_penalties: torch.Tensor) -> torch.Tensor: - num_seqs, vocab_size = logits.shape - _, prompt_mask = _get_bin_counts_and_mask(prompt_tokens_tensor, vocab_size, - num_seqs) - output_bin_counts, output_mask = _get_bin_counts_and_mask( - output_tokens_tensor, vocab_size, num_seqs) - - repetition_penalties = repetition_penalties[:, None].repeat(1, vocab_size) - repetition_penalties[~(prompt_mask | output_mask)] = 1.0 - logits = torch.where(logits > 0, logits / repetition_penalties, - logits * repetition_penalties) - - # We follow the definition in OpenAI API. - # Refer to https://platform.openai.com/docs/api-reference/parameter-details - logits -= frequency_penalties.unsqueeze_(dim=1) * output_bin_counts - logits -= presence_penalties.unsqueeze_(dim=1) * output_mask - return logits - - -def _apply_top_k_top_p( - logits: torch.Tensor, - p: torch.Tensor, - k: torch.Tensor, -) -> torch.Tensor: - logits_sort, logits_idx = logits.sort(dim=-1, descending=False) - - # Apply top-k. - top_k_mask = logits_sort.size(1) - k.to(torch.long) - # Get all the top_k values. - top_k_mask = logits_sort.gather(1, top_k_mask.unsqueeze(dim=1)) - top_k_mask = logits_sort < top_k_mask - logits_sort.masked_fill_(top_k_mask, -float("inf")) - - # Apply top-p. - probs_sort = logits_sort.softmax(dim=-1) - probs_sum = probs_sort.cumsum(dim=-1) - top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1) - # at least one - top_p_mask[:, -1] = False - logits_sort.masked_fill_(top_p_mask, -float("inf")) - - # Re-sort the probabilities. - logits = torch.empty_like(logits_sort).scatter_(dim=-1, - index=logits_idx, - src=logits_sort) - return logits - - -def _apply_min_p( - logits: torch.Tensor, - min_p: torch.Tensor, -) -> torch.Tensor: - """ - Adapted from - https://github.com/oobabooga/text-generation-webui/blob/3146124ec01f02c8fb1650a6517cf1b60b537aaf/modules/sampler_hijack.py#L16C17-L16C17 - """ - probs = torch.softmax(logits, dim=-1) - top_probs, _ = probs.max(dim=-1, keepdim=True) - scaled_min_p = min_p.unsqueeze_(dim=1) * top_probs - tokens_to_remove = probs < scaled_min_p - logits = logits.masked_fill_(tokens_to_remove, -float("inf")) - - return logits - - -def _greedy_sample( - selected_seq_groups: List[SequenceGroupToSample], - samples: torch.Tensor, -) -> SampleResultType: - """Run greedy sampling on a given samples. - - Args: - selected_seq_groups: A list of sequence groups batched. - samples: (num_selected_samples,) A tensor of samples. The length of - samples could be smaller than selected_seq_groups if - seq_group.do_sample is False. - Returns: - Tuple of (next_token_ids, parent_ids). The length of returned list is - same as the length of selected_seq_groups. If the corresponding - seq_group has do_sample=False, tuple contains ([], []) - """ - samples_lst = samples.tolist() - sample_idx = 0 - results: SampleResultType = [] - for seq_group in selected_seq_groups: - if not seq_group.do_sample: - results.append(([], [])) - continue - - seq_ids = seq_group.seq_ids - num_parent_seqs = len(seq_ids) - assert num_parent_seqs == 1, ( - "Greedy sampling should have only one seq.") - parent_ids = list(range(num_parent_seqs)) - next_token_ids = [samples_lst[sample_idx]] - results.append((next_token_ids, parent_ids)) - sample_idx += num_parent_seqs - return results - - -def _random_sample( - selected_seq_groups: List[SequenceGroupToSample], - random_samples: torch.Tensor, -) -> SampleResultType: - """Run random sampling on a given samples. - - Args: - selected_seq_groups: A list of sequence groups batched. - random_samples: (num_selected_samples,) A tensor of samples. The - length of samples could be smaller than selected_seq_groups if - seq_group.do_sample is False. - Returns: - Tuple of (next_token_ids, parent_ids). The length of returned list is - same as the length of selected_seq_groups. If the corresponding - seq_group has do_sample=False, tuple contains ([], []) - """ - # Find the maximum n value of the prompt phase requests. - random_samples = random_samples.cpu() - sample_idx = 0 - results: SampleResultType = [] - for seq_group in selected_seq_groups: - if not seq_group.do_sample: - results.append(([], [])) - continue - - seq_ids = seq_group.seq_ids - sampling_params = seq_group.sampling_params - is_prompt = seq_group.is_prompt - num_parent_seqs = len(seq_ids) - if is_prompt: - # Prompt phase. - parent_ids = [0] * sampling_params.n - next_token_ids = random_samples[ - sample_idx, :sampling_params.n].tolist() - else: - # Generation phase. - parent_ids = list(range(num_parent_seqs)) - next_token_ids = random_samples[sample_idx:sample_idx + - num_parent_seqs, 0].tolist() - results.append((next_token_ids, parent_ids)) - sample_idx += num_parent_seqs - return results - - -def _beam_search_sample( - selected_seq_groups: List[SequenceGroupToSample], - logprobs: torch.Tensor, -) -> SampleResultType: - """Run beam sampling on a given samples. - - Args: - selected_seq_groups: A list of sequence groups batched. - logprobs: (num_selected_samples, vocab_size,) A tensor of logprob - on selected sample indices. - Returns: - Tuple of (next_token_ids, parent_ids). The length of returned list is - same as the length of selected_seq_groups. If the corresponding - seq_group has do_sample=False, tuple contains ([], []) - """ - # We sample 2 * beam_width candidates to make sure that with high - # probability we can get `beam_width` candidates in addition to - # the finished sequences for the next iteration. See - # https://github.com/tensorflow/tensor2tensor/blob/bafdc1b67730430d38d6ab802cbd51f9d053ba2e/tensor2tensor/utils/beam_search.py#L557-L563 - # for details. See also HF reference: - # https://github.com/huggingface/transformers/blob/a4dd53d88e4852f023332d284ff07a01afcd5681/src/transformers/generation/utils.py#L3063-L3065 - # - # NOTE: Beam search is not vectorized, so its speed can be slower than - # other sampling methods. - sample_idx = 0 - results: SampleResultType = [] - for seq_group in selected_seq_groups: - if not seq_group.do_sample: - results.append(([], [])) - continue - - is_prompt = seq_group.is_prompt - seq_ids, sampling_params = seq_group.seq_ids, seq_group.sampling_params - num_parent_seqs = len(seq_ids) - beam_width = sampling_params.n - seq_group_logprobs = logprobs[sample_idx:sample_idx + num_parent_seqs] - if is_prompt: - # Prompt phase. - assert num_parent_seqs == 1, ( - "Prompt input should have only one seq.") - parent_ids = [0] * (2 * beam_width) - _, next_token_ids = torch.topk(seq_group_logprobs[0], - 2 * beam_width) - next_token_ids = next_token_ids.tolist() - else: - # Generation phase. - cumulative_logprobs: List[float] = [ - seq_group.seq_data[seq_id].cumulative_logprob - for seq_id in seq_ids - ] - cumulative_logprobs_tensor = torch.tensor( - cumulative_logprobs, - dtype=torch.float, - device=seq_group_logprobs.device) - seq_group_logprobs = (seq_group_logprobs + - cumulative_logprobs_tensor.unsqueeze(dim=1)) - _, topk_ids = torch.topk(seq_group_logprobs.flatten(), - 2 * beam_width) - topk_ids = topk_ids.tolist() - vocab_size = seq_group_logprobs.size(-1) - parent_ids = [i // vocab_size for i in topk_ids] - next_token_ids = [i % vocab_size for i in topk_ids] - results.append((next_token_ids, parent_ids)) - sample_idx += num_parent_seqs - assert sample_idx == logprobs.size(0) - return results - - -# torch.multinomial forces a GPU<->CPU sync. -# Therefore, we use an optimized implementation instead. -# Note that we always sample with replacement. -# probs will be modified in place, but this is fine, as we pass -# in a copy already. -def _multinomial( - probs: torch.Tensor, - num_samples: int, - seq_groups: Optional[List[SequenceGroupToSample]] = None, -) -> torch.Tensor: - if num_samples > 1: - probs = probs.repeat_interleave(num_samples, dim=0) - q = torch.empty_like(probs) - if seq_groups is None: - q.exponential_() - else: - sample_idx = 0 - for seq_group in seq_groups: - seq_ids = seq_group.seq_ids - stride = len(seq_ids) * num_samples - assert seq_group.generator is not None - q[sample_idx:sample_idx + - stride].exponential_(generator=seq_group.generator) - sample_idx += stride - return probs.div_(q).argmax(dim=1).view(-1, num_samples) - - -def _top_k_top_p_multinomial_with_flashinfer( - probs: torch.Tensor, top_ks: torch.Tensor, top_ps: torch.Tensor, - num_samples: int, seq_groups: Optional[List[SequenceGroupToSample]]): - max_top_k_round = 32 - if num_samples > 1: - probs = probs.repeat_interleave(num_samples, dim=0) - top_ks = top_ks.repeat_interleave(num_samples) - top_ps = top_ps.repeat_interleave(num_samples) - batch_size = probs.shape[0] - uniform_samples = torch.empty((max_top_k_round, batch_size), - device=probs.device) - if seq_groups is None: - uniform_samples.uniform_() - else: - sample_idx = 0 - for seq_group in seq_groups: - seq_ids = seq_group.seq_ids - stride = len(seq_ids) * num_samples - assert seq_group.generator is not None - uniform_samples[:, sample_idx:sample_idx + - stride].uniform_(generator=seq_group.generator) - sample_idx += stride - batch_next_token_ids, success = flashinfer_top_k_top_p_sampling( - probs, - uniform_samples, - top_ks, - top_ps, - ) - if not success.all(): - warnings.warn("FlashInfer rejection sampling failed, fallback.", - stacklevel=1) - probs = flashinfer.sampling.top_k_renorm_prob(probs, top_ks) - probs = flashinfer.sampling.top_p_renorm_prob(probs, top_ps) - batch_next_token_ids = flashinfer.sampling.sampling_from_probs( - probs, uniform_samples[0]) - return batch_next_token_ids.view(-1, num_samples) - - -def get_pythonized_sample_results( - sample_result_args: SampleResultArgsType) -> SampleResultType: - '''This function consumes GPU-side sampler results and computes - Pythonized CPU-side sampler results (GPU -> CPU sync.) - - Single-step scheduling: this function is invoked at sampling-time - for immediate Pythonization. - - Multi-step scheduling: Pythonization is deferred until after multiple - GPU-side steps have been completed. - - Args: - sample_result_args: GPU-side inputs to the Pythonization process - - Returns: - Pythonized sampler results - ''' - - ( - sample_metadata, - sampling_metadata, - greedy_samples, - multinomial_samples, - beam_search_logprobs, - sample_results_dict, - ) = ( - sample_result_args.sample_metadata, - sample_result_args.sampling_metadata, - sample_result_args.greedy_samples, - sample_result_args.multinomial_samples, - sample_result_args.beam_search_logprobs, - sample_result_args.sample_results_dict, - ) - - for sampling_type in SamplingType: - if sampling_type not in sample_metadata: - continue - (seq_group_id, seq_groups) = sample_metadata[sampling_type] - if sampling_type == SamplingType.GREEDY: - sample_results = _greedy_sample(seq_groups, greedy_samples) - elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED): - sample_results = _random_sample(seq_groups, - multinomial_samples[sampling_type]) - elif sampling_type == SamplingType.BEAM: - sample_results = _beam_search_sample(seq_groups, - beam_search_logprobs) - sample_results_dict.update(zip(seq_group_id, sample_results)) - - return [ - sample_results_dict.get(i, ([], [])) - for i in range(len(sampling_metadata.seq_groups)) - ] - - -def _sample_with_torch( - probs: torch.Tensor, - logprobs: torch.Tensor, - sampling_metadata: SamplingMetadata, - sampling_tensors: SamplingTensors, - include_gpu_probs_tensor: bool, - modify_greedy_probs: bool, -) -> SampleReturnType: - '''Torch-oriented _sample() implementation. - - Single-step scheduling: - * Perform GPU-side sampling computation - * Immediately Pythonize sampling result - - Multi-step scheduling: - * Perform GPU-side sampling computation - * Defer Pythonization & preserve GPU-side - tensors required for Pythonization - ''' - - categorized_seq_group_ids: Dict[SamplingType, - List[int]] = {t: [] - for t in SamplingType} - categorized_sample_indices = sampling_metadata.categorized_sample_indices - for i, seq_group in enumerate(sampling_metadata.seq_groups): - sampling_params = seq_group.sampling_params - sampling_type = sampling_params.sampling_type - categorized_seq_group_ids[sampling_type].append(i) - - sample_results_dict: SampleResultsDictType = {} - sample_metadata: SampleMetadataType = {} - multinomial_samples: MultinomialSamplesType = {} - greedy_samples: Optional[torch.Tensor] = None - beam_search_logprobs: Optional[torch.Tensor] = None - - # Create output tensor for sampled token ids. - if include_gpu_probs_tensor: - sampled_token_ids_tensor = torch.full((logprobs.shape[0], 1), - VLLM_INVALID_TOKEN_ID, - dtype=torch.long, - device=logprobs.device) - else: - sampled_token_ids_tensor = None - - # Counterintiutively, having two loops here is actually faster. - # The first loop can run without waiting on GPU<->CPU sync. - for sampling_type in SamplingType: - sample_indices = categorized_sample_indices[sampling_type] - num_tokens = len(sample_indices) - if num_tokens == 0: - continue - - seq_group_id = categorized_seq_group_ids[sampling_type] - seq_groups = [sampling_metadata.seq_groups[i] for i in seq_group_id] - sample_metadata[sampling_type] = (seq_group_id, seq_groups) - long_sample_indices = sample_indices.long() - if sampling_type == SamplingType.GREEDY: - greedy_samples = torch.argmax(logprobs[long_sample_indices], - dim=-1) - - if sampled_token_ids_tensor is not None: - # Store sampled tokens in output tensor. - sampled_token_ids_tensor[ - long_sample_indices] = greedy_samples.unsqueeze(-1) - - if modify_greedy_probs: - # If required, modify the probabilities such that sampling from - # the modified distribution would always sample the argmax - # token id. - _modify_greedy_probs_inplace(logprobs, probs, - long_sample_indices, - greedy_samples) - - elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED): - max_n_in_batch = 1 - for seq_group in seq_groups: - if seq_group.is_prompt: - sampling_params = seq_group.sampling_params - max_n_in_batch = max(max_n_in_batch, sampling_params.n) - seq_groups_arg = (None if sampling_type == SamplingType.RANDOM else - seq_groups) - - if flashinfer_top_k_top_p_sampling is not None: - multinomial_samples[ - sampling_type] = _top_k_top_p_multinomial_with_flashinfer( - probs[long_sample_indices], - sampling_tensors.top_ks[long_sample_indices], - sampling_tensors.top_ps[long_sample_indices], - max_n_in_batch, - seq_groups_arg, - ) - else: - multinomial_samples[sampling_type] = _multinomial( - probs[long_sample_indices], - max_n_in_batch, - seq_groups=seq_groups_arg) - - if sampled_token_ids_tensor is not None: - # Store sampled tokens in output tensor. - sampled_token_ids_tensor[long_sample_indices] = \ - multinomial_samples[sampling_type].to(torch.long) - - elif sampling_type == SamplingType.BEAM: - beam_search_logprobs = logprobs[sample_indices] - else: - raise ValueError(f"Unsupported sampling type: {sampling_type}") - - # Encapsulate arguments for computing Pythonized sampler - # results, whether deferred or otherwise. - maybe_deferred_args = SampleResultArgsType( - sampling_metadata=sampling_metadata, - sample_metadata=sample_metadata, - multinomial_samples=multinomial_samples, - greedy_samples=greedy_samples, - beam_search_logprobs=beam_search_logprobs, - sample_results_dict=sample_results_dict) - - if not sampling_metadata.skip_sampler_cpu_output: - # GPU<->CPU sync happens here. - # This also converts the sampler output to a Python object. - # Return Pythonized sampler result & sampled token ids - return get_pythonized_sample_results( - maybe_deferred_args), sampled_token_ids_tensor - else: - # Defer sampler result Pythonization; return deferred - # Pythonization args & sampled token ids - return ( - maybe_deferred_args, - sampled_token_ids_tensor, - ) - - -def _sample( - probs: torch.Tensor, - logprobs: torch.Tensor, - sampling_metadata: SamplingMetadata, - sampling_tensors: SamplingTensors, - include_gpu_probs_tensor: bool, - modify_greedy_probs: bool, -) -> SampleReturnType: - """ - Args: - probs: (num_query_tokens_in_batch, num_vocab) - logprobs: (num_query_tokens_in_batch, num_vocab) - sampling_metadata: The metadata for a batch for sampling. - sampling_tensors: Tensors that include sampling related metadata. - - Returns: - (next_token_ids, parent_seq_ids) for each seq group in a batch. - If sampling is skipped, it returns ([], []) - sampled_token_ids_tensor: A tensor of sampled token ids. - """ - return _sample_with_torch( - probs, - logprobs, - sampling_metadata, - sampling_tensors, - include_gpu_probs_tensor=include_gpu_probs_tensor, - modify_greedy_probs=modify_greedy_probs, - ) - - -def _get_ranks(x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: - """ - This function calculates the ranks of the chosen tokens in a logprob tensor. - - Args: - x (torch.Tensor): 2D logprob tensor of shape (N, M) - where N is the no. of tokens and M is the vocab dim. - indices (torch.Tensor): List of chosen token indices. - - Returns: - torch.Tensor: 1D tensor of shape (N,) where N is the no. of tokens. - Each element in the returned tensor represents the rank - of the chosen token in the input logprob tensor. - """ - vals = x[torch.arange(0, len(x), device=x.device, dtype=indices.dtype), - indices] - result = (x > vals[:, None]) - del vals - return result.sum(1).add_(1) - - -def get_logprobs( - logprobs: torch.Tensor, - sampling_metadata: SamplingMetadata, - sample_results: SampleResultType, -) -> Tuple[List[Optional[PromptLogprobs]], List[SampleLogprobs]]: - """Return sample logprobs and prompt logprobs. - - The logic consists of 3 parts. - - Select indices to compute logprob from, ranks of token ids, and - the top k token ids from logprobs. - - Compute prompt logprobs if required. - - Compute sample logprobs if required. - - Args: - logprobs: (num_query_tokens_across_batch, num_vocab). Each query token's - logprob per vocab. Sequence groups' query tokens are batched in a - single flattened tensor. For example, assuming there are N - seq groups, it is sorted by prefill tokens for seq_group_1 (if - prompt logprob is enabled), decode tokens for seq_group_1 (if - sampling is required), prefill tokens for seq_group_2, ... - sampling_metadata: The sampling metadata. - sample_results: (num_seq_groups) The tuple of (next_token_ids, - parent_ids) for each sequence group. When beam search is enabled, - sample_results can contain different number of seq_ids from - sampling_metadata.seq_groups. It is because beam search creates - 2 * BEAM_WIDTH number of samples (whereas there are only up to - BEAM_WIDTH number of seq_ids). - - Returns: - A tuple of prompt and sample logprobs per sequence group in a batch. - """ - # The index of query token to calculate logprobs. It includes both - # prompt and sample logprob indices. - query_indices: List[int] = [] - # The next token ids to get the logprob value from. - next_token_ids: List[int] = [] - # The largest requested number of logprobs. We find logprobs as many as the - # largest num logprobs in this API. If every logprobs is None, it will be - # set to -1. - largest_num_logprobs = -1 - - # Select indices to compute logprob from, ranks of token ids, and the top - # k token ids from logprobs. - for (seq_group, sample_result) in zip(sampling_metadata.seq_groups, - sample_results): - sampling_params = seq_group.sampling_params - - # Update indices and tokens for prompt logprobs. - if (seq_group.is_prompt - and sampling_params.prompt_logprobs is not None): - largest_num_logprobs = max(largest_num_logprobs, - sampling_params.prompt_logprobs) - next_prompt_tokens = _get_next_prompt_tokens(seq_group) - query_indices.extend(seq_group.prompt_logprob_indices) - next_token_ids.extend(next_prompt_tokens) - - # Update indices and next tokenes for sample logprob. - if seq_group.do_sample: - token_ids, parent_seq_ids = sample_result - # NOTE: We cannot directly use sample_indices because - # sample_indices only contain parent seq_ids of a previous step. - # The current step may have different number of seq_ids, and - # we can obtain it from `sample_result[1]`. - query_idx = seq_group.sample_indices[0] - query_indices.extend( - [query_idx + parent_id for parent_id in parent_seq_ids]) - next_token_ids.extend(token_ids) - - if sampling_params.logprobs is not None: - largest_num_logprobs = max(largest_num_logprobs, - sampling_params.logprobs) - - assert len(next_token_ids) == len(query_indices) - +"""A layer that samples the next tokens from the model's outputs.""" +import itertools +import warnings +from dataclasses import dataclass +from importlib.util import find_spec +from math import inf +from typing import Dict, List, Optional, Tuple, Union + +import msgspec +import torch +import torch.nn as nn + +import vllm.envs as envs +from vllm.model_executor.sampling_metadata import (SamplingMetadata, + SamplingTensors, + SequenceGroupToSample) +from vllm.sampling_params import SamplingType +from vllm.sequence import (VLLM_INVALID_TOKEN_ID, + CompletionSequenceGroupOutput, Logprob, + PromptLogprobs, SampleLogprobs, SequenceOutput) +from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics + +if envs.VLLM_USE_FLASHINFER_SAMPLER and find_spec("flashinfer"): + import flashinfer.sampling + # yapf: disable + from flashinfer.sampling import ( + top_k_top_p_sampling_from_probs as flashinfer_top_k_top_p_sampling) + + # yapf: enable +else: + flashinfer_top_k_top_p_sampling = None + +# (num_token_ids, num_parent_ids) per sequence group. +SampleResultType = List[Tuple[List[int], List[int]]] + +# Types of temporary data structures used for +# computing sample_result +SampleMetadataType = Dict[SamplingType, Tuple[List[int], + List[SequenceGroupToSample]]] +MultinomialSamplesType = Dict[SamplingType, torch.Tensor] +SampleResultsDictType = Dict[int, Tuple[List[int], List[int]]] + + +# Encapsulates temporary data structures for computing +# sample_result. +# +# * For multi-step scheduling: must be returned +# by `Sampler.forward()` and used later to compute the pythonized +# sample_result +# +# * For single-step scheduling: consumed immediately +# inside `Sampler.forward()` to compute pythonized sample_result. +@dataclass +class SampleResultArgsType: + sample_metadata: SampleMetadataType + multinomial_samples: MultinomialSamplesType + sample_results_dict: SampleResultsDictType + sampling_metadata: SamplingMetadata + greedy_samples: Optional[torch.Tensor] + beam_search_logprobs: Optional[torch.Tensor] + + +# Union of non-deferred (single-step scheduling) +# vs deferred (multi-step scheduling) +# sample result types +MaybeDeferredSampleResultType = Union[SampleResultType, SampleResultArgsType] + +# Abbreviation of the _sample() return type +SampleReturnType = Tuple[MaybeDeferredSampleResultType, Optional[torch.Tensor]] + + +class SamplerOutput( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + array_like=True): # type: ignore[call-arg] + """For each sequence group, we generate a list of SequenceOutput object, + each of which contains one possible candidate for the next token. + + This data structure implements methods, so it can be used like a list, but + also has optional fields for device tensors. + """ + + outputs: List[CompletionSequenceGroupOutput] + + # On-device tensor containing probabilities of each token. + sampled_token_probs: Optional[torch.Tensor] = None + + # On-device tensor containing the logprobs of each token. + logprobs: Optional["torch.Tensor"] = None + + # Holds either (1) the pythonized sampler result (single-step scheduling) + # or (2) what will be arguments for later deferred pythonization of the + # sampler result (muliti-step scheduling) + deferred_sample_results_args: Optional[SampleResultArgsType] = None + + # On-device tensor containing the sampled token ids. + sampled_token_ids: Optional[torch.Tensor] = None + # CPU tensor containing the sampled token ids. Used during multi-step to + # return the sampled token ids from last rank to AsyncLLMEngine to be + # 'broadcasted' to all other PP ranks for next step. + sampled_token_ids_cpu: Optional[torch.Tensor] = None + + # Spec decode metrics populated by workers. + spec_decode_worker_metrics: Optional[SpecDecodeWorkerMetrics] = None + + # Optional last hidden states from the model. + hidden_states: Optional[torch.Tensor] = None + + # Optional prefill hidden states from the model + # (used for models like EAGLE). + prefill_hidden_states: Optional[torch.Tensor] = None + + # Time taken in the forward pass for this across all workers + model_forward_time: Optional[float] = None + + # Time taken in the model execute function. This will include model forward, + # block/sync across workers, cpu-gpu sync time and sampling time. + model_execute_time: Optional[float] = None + + def __getitem__(self, idx: int): + return self.outputs[idx] + + def __setitem__(self, idx: int, value): + self.outputs[idx] = value + + def __len__(self): + return len(self.outputs) + + def __eq__(self, other: object): + return isinstance(other, + self.__class__) and self.outputs == other.outputs + + def __repr__(self) -> str: + """Show the shape of a tensor instead of its values to reduce noise. + """ + sampled_token_probs_repr = ("None" if self.sampled_token_probs is None + else self.sampled_token_probs.shape) + sampled_token_ids_repr = ("None" if self.sampled_token_ids is None else + self.sampled_token_ids.shape) + return ( + f"SamplerOutput(outputs={self.outputs}, " + f"sampled_token_probs={sampled_token_probs_repr}, " + f"sampled_token_ids={sampled_token_ids_repr}, " + f"spec_decode_worker_metrics={self.spec_decode_worker_metrics})") + + +class Sampler(nn.Module): + """Samples the next tokens from the model's outputs. + + This layer does the following: + 1. Discard the hidden states that are not used for sampling (i.e., all + tokens except the final one in each prompt). + 2. Compute the logits for the next tokens. + 3. Apply presence, frequency and repetition penalties. + 4. Apply temperature scaling. + 5. Apply top-p and top-k truncation. + 6. Sample the next tokens. + Here, each sequence group within the batch can have different sampling + parameters (e.g., sampling method, temperature, top-p, top-k, etc.). + + The structure of the logits tensor is coupled with the seq_groups in + sampling_metadata. Typically, each sequence in each seq_group has one row in + logits for the next token to be sampled; however, for a seq_group with a + prompt request with the prompt_logprobs sampling parameter, there are rows + in logits for each token in the input prompt. + """ + + def __init__(self): + super().__init__() + + # Whether or not the SamplerOutput should have on-device tensors + # containing the sampled token ids and probabilities. This is used by + # speculative decoding. + self.include_gpu_probs_tensor = False + self.should_modify_greedy_probs_inplace = False + + def _init_sampling_tensors( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ): + """The goal here is to reuse sampling tensors between similar decode + runs. This is possible because sampling logic does not change between + decodes of the same sequences. + """ + _, vocab_size = logits.shape + + # First free any existing stored sampling tensors. + # This is necessary because some sampling tensors may + # have pinned memory. + self._sampling_tensors = None + + # Initialize new sampling tensors + (sampling_tensors, do_penalties, do_top_p_top_k, + do_min_p) = SamplingTensors.from_sampling_metadata( + sampling_metadata, vocab_size, logits.device, logits.dtype) + + self._sampling_tensors = sampling_tensors + self._do_penalties = do_penalties + self._do_top_p_top_k = do_top_p_top_k + self._do_min_p = do_min_p + + def forward( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[SamplerOutput]: + """ + Single-step scheduling: + * Perform GPU-side sampling computation & compute + GPU-side logprobs tensor + * Pythonize sampling result & logprobs tensor + + Multi-step scheduling: + * Perform GPU-side sampling computation & compute + GPU-side logprobs tensor + * Defer Pythonization of sampling result & logprobs + tensor + * Encapsulate arguments required for deferred Pythonization + in the :class:`SamplerOutput` structure + + Args: + logits: (num_tokens, vocab_size). + sampling_metadata: Metadata for sampling. + """ + assert logits is not None + _, vocab_size = logits.shape + + # Prepare sampling tensors with pinned memory to avoid blocking. + if not sampling_metadata.reuse_sampling_tensors: + self._init_sampling_tensors(logits, sampling_metadata) + elif self._do_penalties: + # In this case, the sampling tensors logic depends on + # "output_tokens" of a sequence. As a result, we cannot + # reuse sampling tensors, since "output_tokens" changes + # between decode runs. + self._init_sampling_tensors(logits, sampling_metadata) + + assert self._sampling_tensors is not None + sampling_tensors = self._sampling_tensors + do_penalties = self._do_penalties + do_top_p_top_k = self._do_top_p_top_k + do_min_p = self._do_min_p + + logits = _apply_min_tokens_penalty(logits, sampling_metadata) + + # Apply presence and frequency penalties. + if do_penalties: + logits = _apply_penalties(logits, sampling_tensors.prompt_tokens, + sampling_tensors.output_tokens, + sampling_tensors.presence_penalties, + sampling_tensors.frequency_penalties, + sampling_tensors.repetition_penalties) + + # Use float32 to apply temperature scaling. + # Use in-place division to avoid creating a new tensor. + logits = logits.to(torch.float) + logits.div_(sampling_tensors.temperatures.unsqueeze(dim=1)) + + if do_top_p_top_k and flashinfer_top_k_top_p_sampling is None: + logits = _apply_top_k_top_p(logits, sampling_tensors.top_ps, + sampling_tensors.top_ks) + + if do_min_p: + logits = _apply_min_p(logits, sampling_tensors.min_ps) + + # We use float32 for probabilities and log probabilities. + # Compute the probabilities. + probs = torch.softmax(logits, dim=-1, dtype=torch.float) + # Compute the log probabilities. + logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float) + + # Sample the next tokens. + maybe_deferred_sample_results, maybe_sampled_tokens_tensor = _sample( + probs, + logprobs, + sampling_metadata, + sampling_tensors, + include_gpu_probs_tensor=self.include_gpu_probs_tensor, + modify_greedy_probs=self._should_modify_greedy_probs_inplace, + ) + + if self.include_gpu_probs_tensor: + # Since we will defer sampler result Pythonization, + # preserve GPU-side tensors in support of later + # deferred pythonization of logprobs + assert maybe_sampled_tokens_tensor is not None + on_device_tensors = (probs, logprobs, maybe_sampled_tokens_tensor) + else: + # Since Pythonization has already happened, don't preserve + # GPU-side tensors. + on_device_tensors = None + + # Get the logprobs query results. + prompt_logprobs = None + sample_logprobs = None + if not sampling_metadata.skip_sampler_cpu_output: + # Pythonize logprobs now (GPU -> CPU); do not defer. + assert not isinstance(maybe_deferred_sample_results, + SampleResultArgsType) + prompt_logprobs, sample_logprobs = get_logprobs( + logprobs, sampling_metadata, maybe_deferred_sample_results) + + return _build_sampler_output( + maybe_deferred_sample_results, + sampling_metadata, + prompt_logprobs, + sample_logprobs, + on_device_tensors=on_device_tensors, + skip_sampler_cpu_output=sampling_metadata.skip_sampler_cpu_output) + + @property + def _should_modify_greedy_probs_inplace(self) -> bool: + """Whether or not the sampler should modify the probability distribution + of greedily-sampled tokens such that multinomial sampling would sample + the greedily-sampled token. + + In other words, if True then we set the probability of the greedily- + sampled token to 1. + + This is used by speculative decoding, which requires that the sampling + method be encoded into the probability distribution. + """ + return self.should_modify_greedy_probs_inplace + + +def _get_bin_counts_and_mask( + tokens: torch.Tensor, + vocab_size: int, + num_seqs: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + # Compute the bin counts for the tokens. + # vocab_size + 1 for padding. + bin_counts = torch.zeros((num_seqs, vocab_size + 1), + dtype=torch.long, + device=tokens.device) + bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens)) + bin_counts = bin_counts[:, :vocab_size] + mask = bin_counts > 0 + + return bin_counts, mask + + +def _apply_min_tokens_penalty( + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, +) -> torch.Tensor: + """Apply min_tokens penalty which sets stop tokens to -inf if min_tokens + have not been generated yet + """ + # list of indices in logits that will be set to -inf + logits_to_penalize: List[Tuple[int, int]] = [] + logits_applied = 0 + for seq_group in sampling_metadata.seq_groups: + seq_ids = seq_group.seq_ids + sampling_params = seq_group.sampling_params + + sample_indices = seq_group.sample_indices + logits_applied += len(sample_indices) + len( + seq_group.prompt_logprob_indices) + if not seq_group.do_sample: + continue + + start_idx = sample_indices[0] + min_tokens = sampling_params.min_tokens + token_ids_to_penalize = sampling_params.all_stop_token_ids + if min_tokens > 0 and token_ids_to_penalize: + seqs_to_penalize: List[int] = [] + for j, seq_id in enumerate(seq_ids): + seq_data = seq_group.seq_data[seq_id] + if len(seq_data.output_token_ids_array) < min_tokens: + seqs_to_penalize.append(j) + + if seqs_to_penalize: + # convert to the index into logits + seqs_to_penalize = [start_idx + j for j in seqs_to_penalize] + # itertools.product pairs each seq index with every token id + logits_to_penalize.extend( + itertools.product(seqs_to_penalize, token_ids_to_penalize)) + + if logits_to_penalize: + # use zip and * to group indices along each dimension + # eg. [ (1,2), (1,3), (5,6) ] -> ( (1,1,5), (2,3,6) ) + logits[tuple(zip(*logits_to_penalize))] = -float("inf") + + # verifies that no rows in logits were missed unexpectedly + assert logits_applied == logits.shape[0] + return logits + + +def _apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor, + output_tokens_tensor: torch.Tensor, + presence_penalties: torch.Tensor, + frequency_penalties: torch.Tensor, + repetition_penalties: torch.Tensor) -> torch.Tensor: + num_seqs, vocab_size = logits.shape + _, prompt_mask = _get_bin_counts_and_mask(prompt_tokens_tensor, vocab_size, + num_seqs) + output_bin_counts, output_mask = _get_bin_counts_and_mask( + output_tokens_tensor, vocab_size, num_seqs) + + repetition_penalties = repetition_penalties[:, None].repeat(1, vocab_size) + repetition_penalties[~(prompt_mask | output_mask)] = 1.0 + logits = torch.where(logits > 0, logits / repetition_penalties, + logits * repetition_penalties) + + # We follow the definition in OpenAI API. + # Refer to https://platform.openai.com/docs/api-reference/parameter-details + logits -= frequency_penalties.unsqueeze_(dim=1) * output_bin_counts + logits -= presence_penalties.unsqueeze_(dim=1) * output_mask + return logits + + +def _apply_top_k_top_p( + logits: torch.Tensor, + p: torch.Tensor, + k: torch.Tensor, +) -> torch.Tensor: + logits_sort, logits_idx = logits.sort(dim=-1, descending=False) + + # Apply top-k. + top_k_mask = logits_sort.size(1) - k.to(torch.long) + # Get all the top_k values. + top_k_mask = logits_sort.gather(1, top_k_mask.unsqueeze(dim=1)) + top_k_mask = logits_sort < top_k_mask + logits_sort.masked_fill_(top_k_mask, -float("inf")) + + # Apply top-p. + probs_sort = logits_sort.softmax(dim=-1) + probs_sum = probs_sort.cumsum(dim=-1) + top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1) + # at least one + top_p_mask[:, -1] = False + logits_sort.masked_fill_(top_p_mask, -float("inf")) + + # Re-sort the probabilities. + logits = torch.empty_like(logits_sort).scatter_(dim=-1, + index=logits_idx, + src=logits_sort) + return logits + + +def _apply_min_p( + logits: torch.Tensor, + min_p: torch.Tensor, +) -> torch.Tensor: + """ + Adapted from + https://github.com/oobabooga/text-generation-webui/blob/3146124ec01f02c8fb1650a6517cf1b60b537aaf/modules/sampler_hijack.py#L16C17-L16C17 + """ + probs = torch.softmax(logits, dim=-1) + top_probs, _ = probs.max(dim=-1, keepdim=True) + scaled_min_p = min_p.unsqueeze_(dim=1) * top_probs + tokens_to_remove = probs < scaled_min_p + logits = logits.masked_fill_(tokens_to_remove, -float("inf")) + + return logits + + +def _greedy_sample( + selected_seq_groups: List[SequenceGroupToSample], + samples: torch.Tensor, +) -> SampleResultType: + """Run greedy sampling on a given samples. + + Args: + selected_seq_groups: A list of sequence groups batched. + samples: (num_selected_samples,) A tensor of samples. The length of + samples could be smaller than selected_seq_groups if + seq_group.do_sample is False. + Returns: + Tuple of (next_token_ids, parent_ids). The length of returned list is + same as the length of selected_seq_groups. If the corresponding + seq_group has do_sample=False, tuple contains ([], []) + """ + samples_lst = samples.tolist() + sample_idx = 0 + results: SampleResultType = [] + for seq_group in selected_seq_groups: + if not seq_group.do_sample: + results.append(([], [])) + continue + + seq_ids = seq_group.seq_ids + num_parent_seqs = len(seq_ids) + assert num_parent_seqs == 1, ( + "Greedy sampling should have only one seq.") + parent_ids = list(range(num_parent_seqs)) + next_token_ids = [samples_lst[sample_idx]] + results.append((next_token_ids, parent_ids)) + sample_idx += num_parent_seqs + return results + + +def _random_sample( + selected_seq_groups: List[SequenceGroupToSample], + random_samples: torch.Tensor, +) -> SampleResultType: + """Run random sampling on a given samples. + + Args: + selected_seq_groups: A list of sequence groups batched. + random_samples: (num_selected_samples,) A tensor of samples. The + length of samples could be smaller than selected_seq_groups if + seq_group.do_sample is False. + Returns: + Tuple of (next_token_ids, parent_ids). The length of returned list is + same as the length of selected_seq_groups. If the corresponding + seq_group has do_sample=False, tuple contains ([], []) + """ + # Find the maximum n value of the prompt phase requests. + random_samples = random_samples.cpu() + sample_idx = 0 + results: SampleResultType = [] + for seq_group in selected_seq_groups: + if not seq_group.do_sample: + results.append(([], [])) + continue + + seq_ids = seq_group.seq_ids + sampling_params = seq_group.sampling_params + is_prompt = seq_group.is_prompt + num_parent_seqs = len(seq_ids) + if is_prompt: + # Prompt phase. + parent_ids = [0] * sampling_params.n + next_token_ids = random_samples[ + sample_idx, :sampling_params.n].tolist() + else: + # Generation phase. + parent_ids = list(range(num_parent_seqs)) + next_token_ids = random_samples[sample_idx:sample_idx + + num_parent_seqs, 0].tolist() + results.append((next_token_ids, parent_ids)) + sample_idx += num_parent_seqs + return results + + +def _beam_search_sample( + selected_seq_groups: List[SequenceGroupToSample], + logprobs: torch.Tensor, +) -> SampleResultType: + """Run beam sampling on a given samples. + + Args: + selected_seq_groups: A list of sequence groups batched. + logprobs: (num_selected_samples, vocab_size,) A tensor of logprob + on selected sample indices. + Returns: + Tuple of (next_token_ids, parent_ids). The length of returned list is + same as the length of selected_seq_groups. If the corresponding + seq_group has do_sample=False, tuple contains ([], []) + """ + # We sample 2 * beam_width candidates to make sure that with high + # probability we can get `beam_width` candidates in addition to + # the finished sequences for the next iteration. See + # https://github.com/tensorflow/tensor2tensor/blob/bafdc1b67730430d38d6ab802cbd51f9d053ba2e/tensor2tensor/utils/beam_search.py#L557-L563 + # for details. See also HF reference: + # https://github.com/huggingface/transformers/blob/a4dd53d88e4852f023332d284ff07a01afcd5681/src/transformers/generation/utils.py#L3063-L3065 + # + # NOTE: Beam search is not vectorized, so its speed can be slower than + # other sampling methods. + sample_idx = 0 + results: SampleResultType = [] + for seq_group in selected_seq_groups: + if not seq_group.do_sample: + results.append(([], [])) + continue + + is_prompt = seq_group.is_prompt + seq_ids, sampling_params = seq_group.seq_ids, seq_group.sampling_params + num_parent_seqs = len(seq_ids) + beam_width = sampling_params.n + seq_group_logprobs = logprobs[sample_idx:sample_idx + num_parent_seqs] + if is_prompt: + # Prompt phase. + assert num_parent_seqs == 1, ( + "Prompt input should have only one seq.") + parent_ids = [0] * (2 * beam_width) + _, next_token_ids = torch.topk(seq_group_logprobs[0], + 2 * beam_width) + next_token_ids = next_token_ids.tolist() + else: + # Generation phase. + cumulative_logprobs: List[float] = [ + seq_group.seq_data[seq_id].cumulative_logprob + for seq_id in seq_ids + ] + cumulative_logprobs_tensor = torch.tensor( + cumulative_logprobs, + dtype=torch.float, + device=seq_group_logprobs.device) + seq_group_logprobs = (seq_group_logprobs + + cumulative_logprobs_tensor.unsqueeze(dim=1)) + _, topk_ids = torch.topk(seq_group_logprobs.flatten(), + 2 * beam_width) + topk_ids = topk_ids.tolist() + vocab_size = seq_group_logprobs.size(-1) + parent_ids = [i // vocab_size for i in topk_ids] + next_token_ids = [i % vocab_size for i in topk_ids] + results.append((next_token_ids, parent_ids)) + sample_idx += num_parent_seqs + assert sample_idx == logprobs.size(0) + return results + + +# torch.multinomial forces a GPU<->CPU sync. +# Therefore, we use an optimized implementation instead. +# Note that we always sample with replacement. +# probs will be modified in place, but this is fine, as we pass +# in a copy already. +def _multinomial( + probs: torch.Tensor, + num_samples: int, + seq_groups: Optional[List[SequenceGroupToSample]] = None, +) -> torch.Tensor: + if num_samples > 1: + probs = probs.repeat_interleave(num_samples, dim=0) + q = torch.empty_like(probs) + if seq_groups is None: + q.exponential_() + else: + sample_idx = 0 + for seq_group in seq_groups: + seq_ids = seq_group.seq_ids + stride = len(seq_ids) * num_samples + assert seq_group.generator is not None + q[sample_idx:sample_idx + + stride].exponential_(generator=seq_group.generator) + sample_idx += stride + return probs.div_(q).argmax(dim=1).view(-1, num_samples) + + +def _top_k_top_p_multinomial_with_flashinfer( + probs: torch.Tensor, top_ks: torch.Tensor, top_ps: torch.Tensor, + num_samples: int, seq_groups: Optional[List[SequenceGroupToSample]]): + max_top_k_round = 32 + if num_samples > 1: + probs = probs.repeat_interleave(num_samples, dim=0) + top_ks = top_ks.repeat_interleave(num_samples) + top_ps = top_ps.repeat_interleave(num_samples) + batch_size = probs.shape[0] + uniform_samples = torch.empty((max_top_k_round, batch_size), + device=probs.device) + if seq_groups is None: + uniform_samples.uniform_() + else: + sample_idx = 0 + for seq_group in seq_groups: + seq_ids = seq_group.seq_ids + stride = len(seq_ids) * num_samples + assert seq_group.generator is not None + uniform_samples[:, sample_idx:sample_idx + + stride].uniform_(generator=seq_group.generator) + sample_idx += stride + batch_next_token_ids, success = flashinfer_top_k_top_p_sampling( + probs, + uniform_samples, + top_ks, + top_ps, + ) + if not success.all(): + warnings.warn("FlashInfer rejection sampling failed, fallback.", + stacklevel=1) + probs = flashinfer.sampling.top_k_renorm_prob(probs, top_ks) + probs = flashinfer.sampling.top_p_renorm_prob(probs, top_ps) + batch_next_token_ids = flashinfer.sampling.sampling_from_probs( + probs, uniform_samples[0]) + return batch_next_token_ids.view(-1, num_samples) + + +def get_pythonized_sample_results( + sample_result_args: SampleResultArgsType) -> SampleResultType: + '''This function consumes GPU-side sampler results and computes + Pythonized CPU-side sampler results (GPU -> CPU sync.) + + Single-step scheduling: this function is invoked at sampling-time + for immediate Pythonization. + + Multi-step scheduling: Pythonization is deferred until after multiple + GPU-side steps have been completed. + + Args: + sample_result_args: GPU-side inputs to the Pythonization process + + Returns: + Pythonized sampler results + ''' + + ( + sample_metadata, + sampling_metadata, + greedy_samples, + multinomial_samples, + beam_search_logprobs, + sample_results_dict, + ) = ( + sample_result_args.sample_metadata, + sample_result_args.sampling_metadata, + sample_result_args.greedy_samples, + sample_result_args.multinomial_samples, + sample_result_args.beam_search_logprobs, + sample_result_args.sample_results_dict, + ) + + for sampling_type in SamplingType: + if sampling_type not in sample_metadata: + continue + (seq_group_id, seq_groups) = sample_metadata[sampling_type] + if sampling_type == SamplingType.GREEDY: + sample_results = _greedy_sample(seq_groups, greedy_samples) + elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED): + sample_results = _random_sample(seq_groups, + multinomial_samples[sampling_type]) + elif sampling_type == SamplingType.BEAM: + sample_results = _beam_search_sample(seq_groups, + beam_search_logprobs) + sample_results_dict.update(zip(seq_group_id, sample_results)) + + return [ + sample_results_dict.get(i, ([], [])) + for i in range(len(sampling_metadata.seq_groups)) + ] + + +def _sample_with_torch( + probs: torch.Tensor, + logprobs: torch.Tensor, + sampling_metadata: SamplingMetadata, + sampling_tensors: SamplingTensors, + include_gpu_probs_tensor: bool, + modify_greedy_probs: bool, +) -> SampleReturnType: + '''Torch-oriented _sample() implementation. + + Single-step scheduling: + * Perform GPU-side sampling computation + * Immediately Pythonize sampling result + + Multi-step scheduling: + * Perform GPU-side sampling computation + * Defer Pythonization & preserve GPU-side + tensors required for Pythonization + ''' + + categorized_seq_group_ids: Dict[SamplingType, + List[int]] = {t: [] + for t in SamplingType} + categorized_sample_indices = sampling_metadata.categorized_sample_indices + for i, seq_group in enumerate(sampling_metadata.seq_groups): + sampling_params = seq_group.sampling_params + sampling_type = sampling_params.sampling_type + categorized_seq_group_ids[sampling_type].append(i) + + sample_results_dict: SampleResultsDictType = {} + sample_metadata: SampleMetadataType = {} + multinomial_samples: MultinomialSamplesType = {} + greedy_samples: Optional[torch.Tensor] = None + beam_search_logprobs: Optional[torch.Tensor] = None + + # Create output tensor for sampled token ids. + if include_gpu_probs_tensor: + sampled_token_ids_tensor = torch.full((logprobs.shape[0], 1), + VLLM_INVALID_TOKEN_ID, + dtype=torch.long, + device=logprobs.device) + else: + sampled_token_ids_tensor = None + + # Counterintiutively, having two loops here is actually faster. + # The first loop can run without waiting on GPU<->CPU sync. + for sampling_type in SamplingType: + sample_indices = categorized_sample_indices[sampling_type] + num_tokens = len(sample_indices) + if num_tokens == 0: + continue + + seq_group_id = categorized_seq_group_ids[sampling_type] + seq_groups = [sampling_metadata.seq_groups[i] for i in seq_group_id] + sample_metadata[sampling_type] = (seq_group_id, seq_groups) + long_sample_indices = sample_indices.long() + if sampling_type == SamplingType.GREEDY: + greedy_samples = torch.argmax(logprobs[long_sample_indices], + dim=-1) + + if sampled_token_ids_tensor is not None: + # Store sampled tokens in output tensor. + sampled_token_ids_tensor[ + long_sample_indices] = greedy_samples.unsqueeze(-1) + + if modify_greedy_probs: + # If required, modify the probabilities such that sampling from + # the modified distribution would always sample the argmax + # token id. + _modify_greedy_probs_inplace(logprobs, probs, + long_sample_indices, + greedy_samples) + + elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED): + max_n_in_batch = 1 + for seq_group in seq_groups: + if seq_group.is_prompt: + sampling_params = seq_group.sampling_params + max_n_in_batch = max(max_n_in_batch, sampling_params.n) + seq_groups_arg = (None if sampling_type == SamplingType.RANDOM else + seq_groups) + + if flashinfer_top_k_top_p_sampling is not None: + multinomial_samples[ + sampling_type] = _top_k_top_p_multinomial_with_flashinfer( + probs[long_sample_indices], + sampling_tensors.top_ks[long_sample_indices], + sampling_tensors.top_ps[long_sample_indices], + max_n_in_batch, + seq_groups_arg, + ) + else: + multinomial_samples[sampling_type] = _multinomial( + probs[long_sample_indices], + max_n_in_batch, + seq_groups=seq_groups_arg) + + if sampled_token_ids_tensor is not None: + # Store sampled tokens in output tensor. + sampled_token_ids_tensor[long_sample_indices] = \ + multinomial_samples[sampling_type].to(torch.long) + + elif sampling_type == SamplingType.BEAM: + beam_search_logprobs = logprobs[sample_indices] + else: + raise ValueError(f"Unsupported sampling type: {sampling_type}") + + # Encapsulate arguments for computing Pythonized sampler + # results, whether deferred or otherwise. + maybe_deferred_args = SampleResultArgsType( + sampling_metadata=sampling_metadata, + sample_metadata=sample_metadata, + multinomial_samples=multinomial_samples, + greedy_samples=greedy_samples, + beam_search_logprobs=beam_search_logprobs, + sample_results_dict=sample_results_dict) + + if not sampling_metadata.skip_sampler_cpu_output: + # GPU<->CPU sync happens here. + # This also converts the sampler output to a Python object. + # Return Pythonized sampler result & sampled token ids + return get_pythonized_sample_results( + maybe_deferred_args), sampled_token_ids_tensor + else: + # Defer sampler result Pythonization; return deferred + # Pythonization args & sampled token ids + return ( + maybe_deferred_args, + sampled_token_ids_tensor, + ) + + +def _sample( + probs: torch.Tensor, + logprobs: torch.Tensor, + sampling_metadata: SamplingMetadata, + sampling_tensors: SamplingTensors, + include_gpu_probs_tensor: bool, + modify_greedy_probs: bool, +) -> SampleReturnType: + """ + Args: + probs: (num_query_tokens_in_batch, num_vocab) + logprobs: (num_query_tokens_in_batch, num_vocab) + sampling_metadata: The metadata for a batch for sampling. + sampling_tensors: Tensors that include sampling related metadata. + + Returns: + (next_token_ids, parent_seq_ids) for each seq group in a batch. + If sampling is skipped, it returns ([], []) + sampled_token_ids_tensor: A tensor of sampled token ids. + """ + return _sample_with_torch( + probs, + logprobs, + sampling_metadata, + sampling_tensors, + include_gpu_probs_tensor=include_gpu_probs_tensor, + modify_greedy_probs=modify_greedy_probs, + ) + + +def _get_ranks(x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + """ + This function calculates the ranks of the chosen tokens in a logprob tensor. + + Args: + x (torch.Tensor): 2D logprob tensor of shape (N, M) + where N is the no. of tokens and M is the vocab dim. + indices (torch.Tensor): List of chosen token indices. + + Returns: + torch.Tensor: 1D tensor of shape (N,) where N is the no. of tokens. + Each element in the returned tensor represents the rank + of the chosen token in the input logprob tensor. + """ + vals = x[torch.arange(0, len(x), device=x.device, dtype=indices.dtype), + indices] + result = (x > vals[:, None]) + del vals + return result.sum(1).add_(1) + + +def get_logprobs( + logprobs: torch.Tensor, + sampling_metadata: SamplingMetadata, + sample_results: SampleResultType, +) -> Tuple[List[Optional[PromptLogprobs]], List[SampleLogprobs]]: + """Return sample logprobs and prompt logprobs. + + The logic consists of 3 parts. + - Select indices to compute logprob from, ranks of token ids, and + the top k token ids from logprobs. + - Compute prompt logprobs if required. + - Compute sample logprobs if required. + + Args: + logprobs: (num_query_tokens_across_batch, num_vocab). Each query token's + logprob per vocab. Sequence groups' query tokens are batched in a + single flattened tensor. For example, assuming there are N + seq groups, it is sorted by prefill tokens for seq_group_1 (if + prompt logprob is enabled), decode tokens for seq_group_1 (if + sampling is required), prefill tokens for seq_group_2, ... + sampling_metadata: The sampling metadata. + sample_results: (num_seq_groups) The tuple of (next_token_ids, + parent_ids) for each sequence group. When beam search is enabled, + sample_results can contain different number of seq_ids from + sampling_metadata.seq_groups. It is because beam search creates + 2 * BEAM_WIDTH number of samples (whereas there are only up to + BEAM_WIDTH number of seq_ids). + + Returns: + A tuple of prompt and sample logprobs per sequence group in a batch. + """ + # The index of query token to calculate logprobs. It includes both + # prompt and sample logprob indices. + query_indices: List[int] = [] + # The next token ids to get the logprob value from. + next_token_ids: List[int] = [] + # The largest requested number of logprobs. We find logprobs as many as the + # largest num logprobs in this API. If every logprobs is None, it will be + # set to -1. + largest_num_logprobs = -1 + + # Select indices to compute logprob from, ranks of token ids, and the top + # k token ids from logprobs. + for (seq_group, sample_result) in zip(sampling_metadata.seq_groups, + sample_results): + sampling_params = seq_group.sampling_params + + # Update indices and tokens for prompt logprobs. + if (seq_group.is_prompt + and sampling_params.prompt_logprobs is not None): + largest_num_logprobs = max(largest_num_logprobs, + sampling_params.prompt_logprobs) + next_prompt_tokens = _get_next_prompt_tokens(seq_group) + query_indices.extend(seq_group.prompt_logprob_indices) + next_token_ids.extend(next_prompt_tokens) + + # Update indices and next tokenes for sample logprob. + if seq_group.do_sample: + token_ids, parent_seq_ids = sample_result + # NOTE: We cannot directly use sample_indices because + # sample_indices only contain parent seq_ids of a previous step. + # The current step may have different number of seq_ids, and + # we can obtain it from `sample_result[1]`. + query_idx = seq_group.sample_indices[0] + query_indices.extend( + [query_idx + parent_id for parent_id in parent_seq_ids]) + next_token_ids.extend(token_ids) + + if sampling_params.logprobs is not None: + largest_num_logprobs = max(largest_num_logprobs, + sampling_params.logprobs) + + assert len(next_token_ids) == len(query_indices) + selected_logprobs, ranks = None, None - top_logprobs, top_token_ids = None, None - - # If largest_num_logprobs == -1, i.e. no logprobs are requested, we can - # skip the whole logprob calculation. + top_logprobs, top_token_ids = None, None + + # If largest_num_logprobs == -1, i.e. no logprobs are requested, we can + # skip the whole logprob calculation. if query_indices and largest_num_logprobs >= 0: - query_indices_gpu = torch.tensor(query_indices, device=logprobs.device) - next_token_ids_gpu = torch.tensor(next_token_ids, - device=logprobs.device) - - # (num_selected_query_tokens, num_logprobs). Note that query_indices can - # contain duplicates if beam search is enabled. - selected_logprobs = logprobs[[ - query_indices_gpu, - next_token_ids_gpu, - ]] - ranks = _get_ranks( - logprobs[query_indices_gpu], - next_token_ids_gpu, - ) - assert selected_logprobs.shape[0] == ranks.shape[0] - - # We need to compute top k only if there exists logprobs > 0. - if largest_num_logprobs > 0: - # Logprobs of topk tokens for a batch of sequence groups. - # (num_query_tokens_across_batch). - top_logprobs, top_token_ids = torch.topk(logprobs, - largest_num_logprobs, - dim=-1) - top_logprobs = top_logprobs.to('cpu') - top_token_ids = top_token_ids.to('cpu') - - selected_logprobs = selected_logprobs.to('cpu') - ranks = ranks.to('cpu') - - # Find prompt/sample logprobs. - prompt_logprobs_per_seq_group: List[Optional[PromptLogprobs]] = [] - sample_logprobs_per_seq_group: List[SampleLogprobs] = [] - top_logprob_idx = 0 - selected_logprobs_idx = 0 - - for seq_group, sample_result in zip(sampling_metadata.seq_groups, - sample_results): - (prompt_logprobs, top_logprob_idx, - selected_logprobs_idx) = _get_prompt_logprob_if_needed( - seq_group, selected_logprobs, ranks, top_token_ids, top_logprobs, - selected_logprobs_idx, top_logprob_idx) - prompt_logprobs_per_seq_group.append(prompt_logprobs) - - (sampled_logprobs, top_logprob_idx, - selected_logprobs_idx) = _get_sampled_logprob_if_needed( - seq_group, sample_result, selected_logprobs, ranks, top_token_ids, - top_logprobs, selected_logprobs_idx, top_logprob_idx) - sample_logprobs_per_seq_group.append(sampled_logprobs) - - return prompt_logprobs_per_seq_group, sample_logprobs_per_seq_group - - -def _get_prompt_logprob_if_needed( - seq_group: SequenceGroupToSample, - selected_logprobs: torch.Tensor, - ranks: torch.Tensor, - top_token_ids: torch.Tensor, - top_logprobs: torch.Tensor, - selected_logprobs_idx: int, - top_logprob_idx: int, -): - """Compute the prompt logprob from a sequence group if needed.""" - sampling_params = seq_group.sampling_params - is_prompt = seq_group.is_prompt - - # Find prompt logprobs + query_indices_gpu = torch.tensor(query_indices, device=logprobs.device) + next_token_ids_gpu = torch.tensor(next_token_ids, + device=logprobs.device) + + # (num_selected_query_tokens, num_logprobs). Note that query_indices can + # contain duplicates if beam search is enabled. + selected_logprobs = logprobs[[ + query_indices_gpu, + next_token_ids_gpu, + ]] + ranks = _get_ranks( + logprobs[query_indices_gpu], + next_token_ids_gpu, + ) + assert selected_logprobs.shape[0] == ranks.shape[0] + + # We need to compute top k only if there exists logprobs > 0. + if largest_num_logprobs > 0: + # Logprobs of topk tokens for a batch of sequence groups. + # (num_query_tokens_across_batch). + top_logprobs, top_token_ids = torch.topk(logprobs, + largest_num_logprobs, + dim=-1) + top_logprobs = top_logprobs.to('cpu') + top_token_ids = top_token_ids.to('cpu') + + selected_logprobs = selected_logprobs.to('cpu') + ranks = ranks.to('cpu') + + # Find prompt/sample logprobs. + prompt_logprobs_per_seq_group: List[Optional[PromptLogprobs]] = [] + sample_logprobs_per_seq_group: List[SampleLogprobs] = [] + top_logprob_idx = 0 + selected_logprobs_idx = 0 + + for seq_group, sample_result in zip(sampling_metadata.seq_groups, + sample_results): + (prompt_logprobs, top_logprob_idx, + selected_logprobs_idx) = _get_prompt_logprob_if_needed( + seq_group, selected_logprobs, ranks, top_token_ids, top_logprobs, + selected_logprobs_idx, top_logprob_idx) + prompt_logprobs_per_seq_group.append(prompt_logprobs) + + (sampled_logprobs, top_logprob_idx, + selected_logprobs_idx) = _get_sampled_logprob_if_needed( + seq_group, sample_result, selected_logprobs, ranks, top_token_ids, + top_logprobs, selected_logprobs_idx, top_logprob_idx) + sample_logprobs_per_seq_group.append(sampled_logprobs) + + return prompt_logprobs_per_seq_group, sample_logprobs_per_seq_group + + +def _get_prompt_logprob_if_needed( + seq_group: SequenceGroupToSample, + selected_logprobs: torch.Tensor, + ranks: torch.Tensor, + top_token_ids: torch.Tensor, + top_logprobs: torch.Tensor, + selected_logprobs_idx: int, + top_logprob_idx: int, +): + """Compute the prompt logprob from a sequence group if needed.""" + sampling_params = seq_group.sampling_params + is_prompt = seq_group.is_prompt + + # Find prompt logprobs prompt_logprobs: Optional[PromptLogprobs] = None if is_prompt and sampling_params.prompt_logprobs is not None: query_len = seq_group.query_len @@ -1093,243 +1093,243 @@ def _get_prompt_logprob_if_needed( for idx, (token_id, output_index) in enumerate(zip( next_prompt_tokens, seq_group.prompt_logprob_output_indices)): - # Calculate the prompt logprob of the real prompt tokens. - # {token_id: (logprob, rank_from_vocab)} - prompt_logprobs_dict: Dict[int, Tuple[float, int]] = { - token_id: (selected_logprob_items[idx], rank_items[idx]) - } - + # Calculate the prompt logprob of the real prompt tokens. + # {token_id: (logprob, rank_from_vocab)} + prompt_logprobs_dict: Dict[int, Tuple[float, int]] = { + token_id: (selected_logprob_items[idx], rank_items[idx]) + } + # Add top K prompt logprobs along with its rank. if num_logprobs > 0: assert top_token_ids is not None assert top_logprobs is not None top_ids = top_token_ids[ - top_logprob_idx, :num_logprobs].tolist() - top_probs = top_logprobs[ - top_logprob_idx, :num_logprobs].tolist() - # Top K is already sorted by rank, so we can use 1 ~ - # num_logprobs + 1 for rank. - top_ranks = range(1, num_logprobs + 1) - prompt_logprobs_dict.update({ - top_id: (top_prob, rank) - for top_id, top_prob, rank in zip(top_ids, top_probs, - top_ranks) - }) + top_logprob_idx, :num_logprobs].tolist() + top_probs = top_logprobs[ + top_logprob_idx, :num_logprobs].tolist() + # Top K is already sorted by rank, so we can use 1 ~ + # num_logprobs + 1 for rank. + top_ranks = range(1, num_logprobs + 1) + prompt_logprobs_dict.update({ + top_id: (top_prob, rank) + for top_id, top_prob, rank in zip(top_ids, top_probs, + top_ranks) + }) prompt_logprobs[output_index] = { token_id: Logprob(*logprob_and_rank) for token_id, logprob_and_rank in prompt_logprobs_dict.items() } - # + 1 to go to the next prompt token. - top_logprob_idx += 1 - - # + len(next_prompt_tokens) to go to the next prompt. + # + 1 to go to the next prompt token. + top_logprob_idx += 1 + + # + len(next_prompt_tokens) to go to the next prompt. selected_logprobs_idx += len(next_prompt_tokens) - return prompt_logprobs, top_logprob_idx, selected_logprobs_idx - - -def _get_sampled_logprob_if_needed( - seq_group: SequenceGroupToSample, - sample_result: Tuple[List[int], List[int]], - selected_logprobs: torch.Tensor, - ranks: torch.Tensor, - top_token_ids: torch.Tensor, - top_logprobs: torch.Tensor, - selected_logprobs_idx: int, - top_logprob_idx: int, -): - """Compute the sample logprob if needed.""" - seq_ids = seq_group.seq_ids - num_logprobs = seq_group.sampling_params.logprobs - sampled_logprobs: SampleLogprobs = [] - next_token_ids, parent_seq_ids = sample_result - - if seq_group.do_sample: - assert len(next_token_ids) > 0 - if num_logprobs is None: - for next_token_id in next_token_ids: - # Use a dummy logprob - sampled_logprobs.append({next_token_id: Logprob(inf)}) - else: - # Pre-select items from tensor. tolist() is faster than repetitive - # `.item()` calls. - selected_logprob_items = selected_logprobs[ - selected_logprobs_idx:selected_logprobs_idx + - len(next_token_ids)].tolist() - rank_items = ranks[selected_logprobs_idx:selected_logprobs_idx + - len(next_token_ids)].tolist() - for idx, (next_token_id, parent_id) in enumerate( - zip(next_token_ids, parent_seq_ids)): - # Get the logprob of a sampled token. - sampled_logprobs_dict = { - next_token_id: - (selected_logprob_items[idx], rank_items[idx]) - } - if num_logprobs is not None and num_logprobs > 0: - # Get top K logprobs. - top_ids = top_token_ids[top_logprob_idx + - parent_id, :num_logprobs].tolist() - top_probs = top_logprobs[ - top_logprob_idx + parent_id, :num_logprobs].tolist() - # Top K is already sorted by rank, so we can use 1 ~ - # num_logprobs + 1 for rank. - top_ranks = range(1, num_logprobs + 1) - sampled_logprobs_dict.update({ - top_id: (top_prob, rank) - for top_id, top_prob, rank in zip( - top_ids, top_probs, top_ranks) - }) - - sampled_logprobs.append({ - token_id: Logprob(*logprob_and_rank) - for token_id, logprob_and_rank in - sampled_logprobs_dict.items() - }) - - # NOTE: This part of code is not intuitive. `selected_logprobs` include - # logprobs for the current step, which has len(next_token_ids) tokens - # per sequence group. `logprobs` includes logprobs from the previous - # steps, which has len(seq_ids) tokens per sequence group. - - # Iterate to the next sequence group in a batch. - selected_logprobs_idx += len(next_token_ids) - # Iterate to the next sequence group in a batch. - top_logprob_idx += len(seq_ids) - return sampled_logprobs, top_logprob_idx, selected_logprobs_idx - - -def _modify_greedy_probs_inplace(logprobs: torch.Tensor, probs: torch.Tensor, - sample_indices: torch.Tensor, - greedy_samples: torch.Tensor) -> None: - """Modify the probability distributions of the greedily-sampled tokens such - that each sampled token has a "probability" of 1.0. This is required by - speculative decoding, which depends on the sampling method being encoded - within the probability distribution for correctness. - - # Why do we only need to do this for greedy sampling? - - vLLM's sampler performs the following steps for greedy or multinomial - (random) sampling: - 1. Get logits from model. - 2. Modify logits according to per-sequence sampling parameters. - - Multiply by temperature, top-k and top-p masking, penalize tokens - according to their frequency, etc. - 3. Sample a token. - - Random sampling simply samples from the modified probability - distribution. - - Greedy sampling performs `argmax` to obtain the token with the - highest likelihood. - - Ignoring greedy sampling for a moment, we find that the computed probability - distribution has the following property: we can sample from it independently - and find that the token sampled by the Sampler has a frequency corresponding - to how often we see it in our sampling. In other words, for tokens sampled - with vLLM's random SamplingType, the computed probability distribution - encodes the sampling methodology completely. - - Greedy sampling does not normally have this property. vLLM modifies logits - according to sampling params, then performs `argmax`, then returns the - sampled token and the computed probability distribution. If we sample from - the distribution, we'll find the likelihood of the greedily-sampled token - is not always 1.0. - - Since lossless speculative decoding requires that the sampling methodology - be encoded within the probability distribution, we are motivated to modify - the probability distribution such that the sampled token has probability 1 - when speculative decoding is used. - - NOTE: Alternatively, we could use an extremely low temperature to achieve - greedy sampling using multinomial computation and unite the codepaths. This - has implications on the overall design of the sampler, e.g. how to record - accurate logprobs for the user, so this improvement is deferred to later. - """ - # NOTE: logprobs are not modified so they can be returned to the user. - probs[sample_indices, :] = 0 - probs[sample_indices, greedy_samples] = 1.0 - - -def _build_sampler_output( - maybe_deferred_sample_results: MaybeDeferredSampleResultType, - sampling_metadata: SamplingMetadata, - prompt_logprobs: Optional[List[Optional[PromptLogprobs]]], - sample_logprobs: Optional[List[SampleLogprobs]], - on_device_tensors: Optional[Tuple[torch.Tensor, torch.Tensor, - torch.Tensor]], - skip_sampler_cpu_output: bool = False, -) -> SamplerOutput: - """Construct Python objects with the output of sampling. - - Args: - on_device_tensors: Tuple containing on-device tensors with the - probabilities used in sampling and the sampled token ids. This - allows post-processing without copies to CPU/serialization, e.g. in - speculative decoding rejection sampling. - """ - sampler_output: List[CompletionSequenceGroupOutput] = [] - - if skip_sampler_cpu_output: - assert isinstance(maybe_deferred_sample_results, SampleResultArgsType) - deferred_sample_results_args = maybe_deferred_sample_results - else: - assert prompt_logprobs is not None - assert sample_logprobs is not None - assert not isinstance(maybe_deferred_sample_results, - SampleResultArgsType) - deferred_sample_results_args = None - - for (seq_group, sample_result, group_prompt_logprobs, - group_sample_logprobs) in zip(sampling_metadata.seq_groups, - maybe_deferred_sample_results, - prompt_logprobs, sample_logprobs): - seq_ids = seq_group.seq_ids - next_token_ids, parent_ids = sample_result - seq_outputs: List[SequenceOutput] = [] - for parent_id, next_token_id, logprobs in zip( - parent_ids, next_token_ids, group_sample_logprobs): - seq_outputs.append( - SequenceOutput(seq_ids[parent_id], next_token_id, - logprobs)) - sampler_output.append( - CompletionSequenceGroupOutput(seq_outputs, - group_prompt_logprobs)) - - # If not specified, store None values in SamplerOutput. - if on_device_tensors is not None: - (sampled_token_probs, logprobs_tensor, - sampled_token_ids) = on_device_tensors - else: - sampled_token_probs, logprobs_tensor, sampled_token_ids = (None, None, - None) - - return SamplerOutput( - outputs=sampler_output, - sampled_token_probs=sampled_token_probs, - sampled_token_ids=sampled_token_ids, - logprobs=logprobs_tensor, - deferred_sample_results_args=deferred_sample_results_args) - - -def _get_next_prompt_tokens(seq_group: SequenceGroupToSample) -> List[int]: - """Get a list of next prompt tokens to compute logprob from a - given sequence group. - - It is used to compute prompt logprob. Imagine you have logprob for each - query token. Query token needs to know the next prompt token id to compute - prompt logprob. This is a helper to obtain next prompt token ids. - - This API has to be used only when the caller knows seq_group is in prefill - stage. - - Returns: - A list of next prompt tokens to compute logprob. - """ - assert seq_group.is_prompt, ( - "Caller should ensure the sequence group is in a prefill stage.") - seq_ids = seq_group.seq_ids - query_len = seq_group.query_len - assert query_len is not None - # prompt has only 1 seq id. - assert len(seq_ids) == 1 - seq_data = seq_group.seq_data[seq_ids[0]] + return prompt_logprobs, top_logprob_idx, selected_logprobs_idx + + +def _get_sampled_logprob_if_needed( + seq_group: SequenceGroupToSample, + sample_result: Tuple[List[int], List[int]], + selected_logprobs: torch.Tensor, + ranks: torch.Tensor, + top_token_ids: torch.Tensor, + top_logprobs: torch.Tensor, + selected_logprobs_idx: int, + top_logprob_idx: int, +): + """Compute the sample logprob if needed.""" + seq_ids = seq_group.seq_ids + num_logprobs = seq_group.sampling_params.logprobs + sampled_logprobs: SampleLogprobs = [] + next_token_ids, parent_seq_ids = sample_result + + if seq_group.do_sample: + assert len(next_token_ids) > 0 + if num_logprobs is None: + for next_token_id in next_token_ids: + # Use a dummy logprob + sampled_logprobs.append({next_token_id: Logprob(inf)}) + else: + # Pre-select items from tensor. tolist() is faster than repetitive + # `.item()` calls. + selected_logprob_items = selected_logprobs[ + selected_logprobs_idx:selected_logprobs_idx + + len(next_token_ids)].tolist() + rank_items = ranks[selected_logprobs_idx:selected_logprobs_idx + + len(next_token_ids)].tolist() + for idx, (next_token_id, parent_id) in enumerate( + zip(next_token_ids, parent_seq_ids)): + # Get the logprob of a sampled token. + sampled_logprobs_dict = { + next_token_id: + (selected_logprob_items[idx], rank_items[idx]) + } + if num_logprobs is not None and num_logprobs > 0: + # Get top K logprobs. + top_ids = top_token_ids[top_logprob_idx + + parent_id, :num_logprobs].tolist() + top_probs = top_logprobs[ + top_logprob_idx + parent_id, :num_logprobs].tolist() + # Top K is already sorted by rank, so we can use 1 ~ + # num_logprobs + 1 for rank. + top_ranks = range(1, num_logprobs + 1) + sampled_logprobs_dict.update({ + top_id: (top_prob, rank) + for top_id, top_prob, rank in zip( + top_ids, top_probs, top_ranks) + }) + + sampled_logprobs.append({ + token_id: Logprob(*logprob_and_rank) + for token_id, logprob_and_rank in + sampled_logprobs_dict.items() + }) + + # NOTE: This part of code is not intuitive. `selected_logprobs` include + # logprobs for the current step, which has len(next_token_ids) tokens + # per sequence group. `logprobs` includes logprobs from the previous + # steps, which has len(seq_ids) tokens per sequence group. + + # Iterate to the next sequence group in a batch. + selected_logprobs_idx += len(next_token_ids) + # Iterate to the next sequence group in a batch. + top_logprob_idx += len(seq_ids) + return sampled_logprobs, top_logprob_idx, selected_logprobs_idx + + +def _modify_greedy_probs_inplace(logprobs: torch.Tensor, probs: torch.Tensor, + sample_indices: torch.Tensor, + greedy_samples: torch.Tensor) -> None: + """Modify the probability distributions of the greedily-sampled tokens such + that each sampled token has a "probability" of 1.0. This is required by + speculative decoding, which depends on the sampling method being encoded + within the probability distribution for correctness. + + # Why do we only need to do this for greedy sampling? + + vLLM's sampler performs the following steps for greedy or multinomial + (random) sampling: + 1. Get logits from model. + 2. Modify logits according to per-sequence sampling parameters. + - Multiply by temperature, top-k and top-p masking, penalize tokens + according to their frequency, etc. + 3. Sample a token. + - Random sampling simply samples from the modified probability + distribution. + - Greedy sampling performs `argmax` to obtain the token with the + highest likelihood. + + Ignoring greedy sampling for a moment, we find that the computed probability + distribution has the following property: we can sample from it independently + and find that the token sampled by the Sampler has a frequency corresponding + to how often we see it in our sampling. In other words, for tokens sampled + with vLLM's random SamplingType, the computed probability distribution + encodes the sampling methodology completely. + + Greedy sampling does not normally have this property. vLLM modifies logits + according to sampling params, then performs `argmax`, then returns the + sampled token and the computed probability distribution. If we sample from + the distribution, we'll find the likelihood of the greedily-sampled token + is not always 1.0. + + Since lossless speculative decoding requires that the sampling methodology + be encoded within the probability distribution, we are motivated to modify + the probability distribution such that the sampled token has probability 1 + when speculative decoding is used. + + NOTE: Alternatively, we could use an extremely low temperature to achieve + greedy sampling using multinomial computation and unite the codepaths. This + has implications on the overall design of the sampler, e.g. how to record + accurate logprobs for the user, so this improvement is deferred to later. + """ + # NOTE: logprobs are not modified so they can be returned to the user. + probs[sample_indices, :] = 0 + probs[sample_indices, greedy_samples] = 1.0 + + +def _build_sampler_output( + maybe_deferred_sample_results: MaybeDeferredSampleResultType, + sampling_metadata: SamplingMetadata, + prompt_logprobs: Optional[List[Optional[PromptLogprobs]]], + sample_logprobs: Optional[List[SampleLogprobs]], + on_device_tensors: Optional[Tuple[torch.Tensor, torch.Tensor, + torch.Tensor]], + skip_sampler_cpu_output: bool = False, +) -> SamplerOutput: + """Construct Python objects with the output of sampling. + + Args: + on_device_tensors: Tuple containing on-device tensors with the + probabilities used in sampling and the sampled token ids. This + allows post-processing without copies to CPU/serialization, e.g. in + speculative decoding rejection sampling. + """ + sampler_output: List[CompletionSequenceGroupOutput] = [] + + if skip_sampler_cpu_output: + assert isinstance(maybe_deferred_sample_results, SampleResultArgsType) + deferred_sample_results_args = maybe_deferred_sample_results + else: + assert prompt_logprobs is not None + assert sample_logprobs is not None + assert not isinstance(maybe_deferred_sample_results, + SampleResultArgsType) + deferred_sample_results_args = None + + for (seq_group, sample_result, group_prompt_logprobs, + group_sample_logprobs) in zip(sampling_metadata.seq_groups, + maybe_deferred_sample_results, + prompt_logprobs, sample_logprobs): + seq_ids = seq_group.seq_ids + next_token_ids, parent_ids = sample_result + seq_outputs: List[SequenceOutput] = [] + for parent_id, next_token_id, logprobs in zip( + parent_ids, next_token_ids, group_sample_logprobs): + seq_outputs.append( + SequenceOutput(seq_ids[parent_id], next_token_id, + logprobs)) + sampler_output.append( + CompletionSequenceGroupOutput(seq_outputs, + group_prompt_logprobs)) + + # If not specified, store None values in SamplerOutput. + if on_device_tensors is not None: + (sampled_token_probs, logprobs_tensor, + sampled_token_ids) = on_device_tensors + else: + sampled_token_probs, logprobs_tensor, sampled_token_ids = (None, None, + None) + + return SamplerOutput( + outputs=sampler_output, + sampled_token_probs=sampled_token_probs, + sampled_token_ids=sampled_token_ids, + logprobs=logprobs_tensor, + deferred_sample_results_args=deferred_sample_results_args) + + +def _get_next_prompt_tokens(seq_group: SequenceGroupToSample) -> List[int]: + """Get a list of next prompt tokens to compute logprob from a + given sequence group. + + It is used to compute prompt logprob. Imagine you have logprob for each + query token. Query token needs to know the next prompt token id to compute + prompt logprob. This is a helper to obtain next prompt token ids. + + This API has to be used only when the caller knows seq_group is in prefill + stage. + + Returns: + A list of next prompt tokens to compute logprob. + """ + assert seq_group.is_prompt, ( + "Caller should ensure the sequence group is in a prefill stage.") + seq_ids = seq_group.seq_ids + query_len = seq_group.query_len + assert query_len is not None + # prompt has only 1 seq id. + assert len(seq_ids) == 1 + seq_data = seq_group.seq_data[seq_ids[0]] computed_len = seq_data.get_num_computed_tokens() prompt_tokens = seq_data.prompt_token_ids next_prompt_tokens = [] diff --git a/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py b/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py index 56785643..98ffffdc 100644 --- a/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py +++ b/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py @@ -1,58 +1,58 @@ -from array import array -from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple - -import torch - -from vllm.sampling_params import SamplingParams, SamplingType -from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE, SequenceData, - SequenceGroupMetadata) -from vllm.utils import (PyObjectCache, async_tensor_h2d, - is_pin_memory_available, make_tensor_with_pad) - -_SAMPLING_EPS = 1e-5 - - -@dataclass -class SequenceGroupToSample: - # |---------- N-1 iteration --------| - # |---------------- N iteration ---------------------| - # |- tokenA -|......................|-- newTokens ---| - # |---------- context_len ----------| - # |-------------------- seq_len ----------------------| - # |-- query_len ---| - - # Sequence ids for the sequence group in a previous step. - seq_ids: List[int] - sampling_params: SamplingParams - # seq_id -> sequence data. - seq_data: Dict[int, SequenceData] - # The length of the sequence (all tokens seen in the past + new token to - # compute attention) of the sequence group. None if it is in a decode - # stage. - seq_len: Optional[int] - # The length of new query tokens to compute in the current step. None if it - # is in a decode stage. The length of query_len <= seq_len if chunked - # prefill is enabled. - query_len: Optional[int] - # A random number generator for sampling. - generator: Optional[torch.Generator] - # True if the sequence group is in prefill stage. False if it is in a - # decode stage. - is_prompt: bool +from array import array +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import torch + +from vllm.sampling_params import SamplingParams, SamplingType +from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE, SequenceData, + SequenceGroupMetadata) +from vllm.utils import (PyObjectCache, async_tensor_h2d, + is_pin_memory_available, make_tensor_with_pad) + +_SAMPLING_EPS = 1e-5 + + +@dataclass +class SequenceGroupToSample: + # |---------- N-1 iteration --------| + # |---------------- N iteration ---------------------| + # |- tokenA -|......................|-- newTokens ---| + # |---------- context_len ----------| + # |-------------------- seq_len ----------------------| + # |-- query_len ---| + + # Sequence ids for the sequence group in a previous step. + seq_ids: List[int] + sampling_params: SamplingParams + # seq_id -> sequence data. + seq_data: Dict[int, SequenceData] + # The length of the sequence (all tokens seen in the past + new token to + # compute attention) of the sequence group. None if it is in a decode + # stage. + seq_len: Optional[int] + # The length of new query tokens to compute in the current step. None if it + # is in a decode stage. The length of query_len <= seq_len if chunked + # prefill is enabled. + query_len: Optional[int] + # A random number generator for sampling. + generator: Optional[torch.Generator] + # True if the sequence group is in prefill stage. False if it is in a + # decode stage. + is_prompt: bool # Query token indices from logits. to compute prompt logprob. Empty if # prompt logprob is not required. prompt_logprob_indices: List[int] # Output offsets within this prefill chunk. Sparse diagnostic requests use # this to retain the standard full-length prompt-logprob response shape. prompt_logprob_output_indices: List[int] - # Sample token indices from logits. Empty if sampling is not required. - sample_indices: List[int] - - @property - def do_sample(self): - return len(self.sample_indices) > 0 - + # Sample token indices from logits. Empty if sampling is not required. + sample_indices: List[int] + + @property + def do_sample(self): + return len(self.sample_indices) > 0 + def __post_init__(self): if len(self.prompt_logprob_indices) > 0: assert self.sampling_params.prompt_logprobs is not None @@ -66,142 +66,142 @@ class SequenceGroupToSample: assert all( 0 <= index < self.query_len for index in self.prompt_logprob_output_indices) - - -def gen_seq_group_to_sample_builder(num_seqs: int): - return lambda: SequenceGroupToSample( - seq_ids=[0] * num_seqs, - sampling_params=None, - seq_data=None, # type: ignore - seq_len=0, - query_len=0, - generator=None, + + +def gen_seq_group_to_sample_builder(num_seqs: int): + return lambda: SequenceGroupToSample( + seq_ids=[0] * num_seqs, + sampling_params=None, + seq_data=None, # type: ignore + seq_len=0, + query_len=0, + generator=None, is_prompt=True, prompt_logprob_indices=[], prompt_logprob_output_indices=[], sample_indices=[], - ) - - -class SamplingMetadataCache: - """Used to cache SamplingMetadata objects between scheduler iterations""" - - def __init__(self): - self._seq_group_to_sample_cache: Dict[int, PyObjectCache] = {} - - def get_cached_seq_group_to_sample(self, num_seqs): - if num_seqs not in self._seq_group_to_sample_cache: - self._seq_group_to_sample_cache[num_seqs] = PyObjectCache( - gen_seq_group_to_sample_builder(num_seqs)) - - obj = self._seq_group_to_sample_cache[num_seqs].get_object() - return obj - - def reset(self): - for cache in self._seq_group_to_sample_cache.values(): - cache.reset() - - -class SamplingMetadata: - """Metadata for input sequences. Used in sampler. - - The usage is as follow; - ``` - hidden_states = execute_model(...) - logits = hidden_states[sampling_metadata.selected_token_indices] - sample(logits) - - def sample(logits): - # Use categorized_sample_indices for sampling.... - ``` - - Args: - seq_groups: List of batched sequence groups. - selected_token_indices: (num_query_tokens_to_logprob). Indices to find - logits from the initial model output hidden states. - categorized_sample_indices: SamplingType -> token indices to sample. - Each token indices is 2D tensor of (num_indices, num_indices) where - the first item means the sample index within the returned logit - (before pruning padding), and the second item means the sample - index after pruning using selected_token_indices. - For example, if the returned logit is [1, 2, 3], and we select - [1, 2] for sampling, the pruned logit will be [2, 3]. In this case, - The first tuple is [1, 2] (sampled index within original logit), - and the second tuple is [0, 1] (sampled index within pruned logit). - num_prompts: Number of prompt sequence groups in seq_groups. - skip_sampler_cpu_output: Indicates if we want to skip the GPU=>CPU - serialization of token outputs. - reuse_sampling_tensors: Indicates if we want to reuse sampling - tensors that are part of the sampler forward pass. Currently, - it is mainly used for multi-step decode. - - """ - - def __init__( - self, - seq_groups: List[SequenceGroupToSample], - selected_token_indices: torch.Tensor, - categorized_sample_indices: Dict[SamplingType, torch.Tensor], - num_prompts: int, - skip_sampler_cpu_output: bool = False, - reuse_sampling_tensors: bool = False, - ) -> None: - self.seq_groups = seq_groups - self.selected_token_indices = selected_token_indices - self.categorized_sample_indices = categorized_sample_indices - self.num_prompts = num_prompts - self.skip_sampler_cpu_output = skip_sampler_cpu_output - self.reuse_sampling_tensors = reuse_sampling_tensors - - @staticmethod - def prepare( - seq_group_metadata_list: List[SequenceGroupMetadata], - seq_lens: List[int], - query_lens: List[int], - device: str, - pin_memory: bool, - generators: Optional[Dict[str, torch.Generator]] = None, - cache: Optional[SamplingMetadataCache] = None, - ) -> "SamplingMetadata": - ( - seq_groups, - selected_token_indices, - categorized_sample_indices, - num_prompts, - ) = _prepare_seq_groups(seq_group_metadata_list, seq_lens, query_lens, - device, generators, cache) - selected_token_indices = async_tensor_h2d( - selected_token_indices, - dtype=torch.long, - target_device=device, - pin_memory=pin_memory, - ) - categorized_sample_indices = { - t: async_tensor_h2d( - seq_ids, - dtype=torch.int, - target_device=device, - pin_memory=pin_memory, - ) - for t, seq_ids in categorized_sample_indices.items() - } - - sampling_metadata = SamplingMetadata( - seq_groups=seq_groups, - selected_token_indices=selected_token_indices, - categorized_sample_indices=categorized_sample_indices, - num_prompts=num_prompts, - ) - return sampling_metadata - - def __repr__(self) -> str: - return ( - "SamplingMetadata(" - f"seq_groups={self.seq_groups}, " - f"selected_token_indices={self.selected_token_indices}, " - f"categorized_sample_indices={self.categorized_sample_indices}), ") - - + ) + + +class SamplingMetadataCache: + """Used to cache SamplingMetadata objects between scheduler iterations""" + + def __init__(self): + self._seq_group_to_sample_cache: Dict[int, PyObjectCache] = {} + + def get_cached_seq_group_to_sample(self, num_seqs): + if num_seqs not in self._seq_group_to_sample_cache: + self._seq_group_to_sample_cache[num_seqs] = PyObjectCache( + gen_seq_group_to_sample_builder(num_seqs)) + + obj = self._seq_group_to_sample_cache[num_seqs].get_object() + return obj + + def reset(self): + for cache in self._seq_group_to_sample_cache.values(): + cache.reset() + + +class SamplingMetadata: + """Metadata for input sequences. Used in sampler. + + The usage is as follow; + ``` + hidden_states = execute_model(...) + logits = hidden_states[sampling_metadata.selected_token_indices] + sample(logits) + + def sample(logits): + # Use categorized_sample_indices for sampling.... + ``` + + Args: + seq_groups: List of batched sequence groups. + selected_token_indices: (num_query_tokens_to_logprob). Indices to find + logits from the initial model output hidden states. + categorized_sample_indices: SamplingType -> token indices to sample. + Each token indices is 2D tensor of (num_indices, num_indices) where + the first item means the sample index within the returned logit + (before pruning padding), and the second item means the sample + index after pruning using selected_token_indices. + For example, if the returned logit is [1, 2, 3], and we select + [1, 2] for sampling, the pruned logit will be [2, 3]. In this case, + The first tuple is [1, 2] (sampled index within original logit), + and the second tuple is [0, 1] (sampled index within pruned logit). + num_prompts: Number of prompt sequence groups in seq_groups. + skip_sampler_cpu_output: Indicates if we want to skip the GPU=>CPU + serialization of token outputs. + reuse_sampling_tensors: Indicates if we want to reuse sampling + tensors that are part of the sampler forward pass. Currently, + it is mainly used for multi-step decode. + + """ + + def __init__( + self, + seq_groups: List[SequenceGroupToSample], + selected_token_indices: torch.Tensor, + categorized_sample_indices: Dict[SamplingType, torch.Tensor], + num_prompts: int, + skip_sampler_cpu_output: bool = False, + reuse_sampling_tensors: bool = False, + ) -> None: + self.seq_groups = seq_groups + self.selected_token_indices = selected_token_indices + self.categorized_sample_indices = categorized_sample_indices + self.num_prompts = num_prompts + self.skip_sampler_cpu_output = skip_sampler_cpu_output + self.reuse_sampling_tensors = reuse_sampling_tensors + + @staticmethod + def prepare( + seq_group_metadata_list: List[SequenceGroupMetadata], + seq_lens: List[int], + query_lens: List[int], + device: str, + pin_memory: bool, + generators: Optional[Dict[str, torch.Generator]] = None, + cache: Optional[SamplingMetadataCache] = None, + ) -> "SamplingMetadata": + ( + seq_groups, + selected_token_indices, + categorized_sample_indices, + num_prompts, + ) = _prepare_seq_groups(seq_group_metadata_list, seq_lens, query_lens, + device, generators, cache) + selected_token_indices = async_tensor_h2d( + selected_token_indices, + dtype=torch.long, + target_device=device, + pin_memory=pin_memory, + ) + categorized_sample_indices = { + t: async_tensor_h2d( + seq_ids, + dtype=torch.int, + target_device=device, + pin_memory=pin_memory, + ) + for t, seq_ids in categorized_sample_indices.items() + } + + sampling_metadata = SamplingMetadata( + seq_groups=seq_groups, + selected_token_indices=selected_token_indices, + categorized_sample_indices=categorized_sample_indices, + num_prompts=num_prompts, + ) + return sampling_metadata + + def __repr__(self) -> str: + return ( + "SamplingMetadata(" + f"seq_groups={self.seq_groups}, " + f"selected_token_indices={self.selected_token_indices}, " + f"categorized_sample_indices={self.categorized_sample_indices}), ") + + def _get_prompt_logprob_output_indices( sampling_params: SamplingParams, seq_data: SequenceData, @@ -231,106 +231,106 @@ def _get_prompt_logprob_output_indices( def _prepare_seq_groups( - seq_group_metadata_list: List[SequenceGroupMetadata], - seq_lens: List[int], - query_lens: List[int], - device: str, - generators: Optional[Dict[str, torch.Generator]] = None, - cache: Optional[SamplingMetadataCache] = None, -) -> Tuple[List[SequenceGroupToSample], List[int], Dict[SamplingType, - List[int]], int, ]: - """Prepare sequence groups and indices for sampling. - - Args: - seq_group_metadata_list: A list of sequence group to batch. - seq_lens: A list of sequence lens per sequence group. - Index of prompt len should match with seq_group_metadata_list. - query_lens: A list of query lengths. Prompt lens include the length - of entire prompt tokens, and it could be shorter. - device: A device to use for random number generators, - `SequenceGroupToSample.generator`. - generators: A store of per-request random number generators used - for seeded requests. - - Returns: - seq_groups: A list of sequence group to sample. - selected_token_indices: See the definition from `SamplingMetadata`. - categorized_sample_indices: See the definition from `SamplingMetadata`. - num_prompts: Total number of prompts from `seq_group_metadata_list`. - """ - # Batched sequence groups for the current model forward stsep. - seq_groups: List[SequenceGroupToSample] = [] - # A list of token indices to sample/compute logprob. It is used to - # prune the outcome logits from the model for the performance. - selected_token_indices: List[int] = [] - # Used for selected_token_indices. - model_output_idx = 0 - - # Sampling type -> ( - # indices to sample/prompt logprob within pruned output logits, - # indices to sample within pruned logits) - categorized_sample_indices: Dict[SamplingType, List[int]] = { - t: [] - for t in SamplingType - } - # Index of logits to compute logprob. Logits include both prompt logprob - # and sample logprob indices. - logit_idx = 0 - # Total number of prompts from given sequence groups. + seq_group_metadata_list: List[SequenceGroupMetadata], + seq_lens: List[int], + query_lens: List[int], + device: str, + generators: Optional[Dict[str, torch.Generator]] = None, + cache: Optional[SamplingMetadataCache] = None, +) -> Tuple[List[SequenceGroupToSample], List[int], Dict[SamplingType, + List[int]], int, ]: + """Prepare sequence groups and indices for sampling. + + Args: + seq_group_metadata_list: A list of sequence group to batch. + seq_lens: A list of sequence lens per sequence group. + Index of prompt len should match with seq_group_metadata_list. + query_lens: A list of query lengths. Prompt lens include the length + of entire prompt tokens, and it could be shorter. + device: A device to use for random number generators, + `SequenceGroupToSample.generator`. + generators: A store of per-request random number generators used + for seeded requests. + + Returns: + seq_groups: A list of sequence group to sample. + selected_token_indices: See the definition from `SamplingMetadata`. + categorized_sample_indices: See the definition from `SamplingMetadata`. + num_prompts: Total number of prompts from `seq_group_metadata_list`. + """ + # Batched sequence groups for the current model forward stsep. + seq_groups: List[SequenceGroupToSample] = [] + # A list of token indices to sample/compute logprob. It is used to + # prune the outcome logits from the model for the performance. + selected_token_indices: List[int] = [] + # Used for selected_token_indices. + model_output_idx = 0 + + # Sampling type -> ( + # indices to sample/prompt logprob within pruned output logits, + # indices to sample within pruned logits) + categorized_sample_indices: Dict[SamplingType, List[int]] = { + t: [] + for t in SamplingType + } + # Index of logits to compute logprob. Logits include both prompt logprob + # and sample logprob indices. + logit_idx = 0 + # Total number of prompts from given sequence groups. num_prompts = 0 - - for i, seq_group_metadata in enumerate(seq_group_metadata_list): - seq_ids = seq_group_metadata.seq_data.keys() - - if cache is not None: - sample_obj = cache.get_cached_seq_group_to_sample(len(seq_ids)) - - for j, seq_id in enumerate(seq_ids): - sample_obj.seq_ids[j] = seq_id - + + for i, seq_group_metadata in enumerate(seq_group_metadata_list): + seq_ids = seq_group_metadata.seq_data.keys() + + if cache is not None: + sample_obj = cache.get_cached_seq_group_to_sample(len(seq_ids)) + + for j, seq_id in enumerate(seq_ids): + sample_obj.seq_ids[j] = seq_id + sample_obj.prompt_logprob_indices.clear() sample_obj.prompt_logprob_output_indices.clear() sample_obj.sample_indices.clear() - - sampling_params = seq_group_metadata.sampling_params - is_prompt = seq_group_metadata.is_prompt - generator: Optional[torch.Generator] = None - # If the current seq group is in decode stage, it is None. - seq_len: Optional[int] = None - query_len: Optional[int] = None + + sampling_params = seq_group_metadata.sampling_params + is_prompt = seq_group_metadata.is_prompt + generator: Optional[torch.Generator] = None + # If the current seq group is in decode stage, it is None. + seq_len: Optional[int] = None + query_len: Optional[int] = None prompt_logprob_indices: List[int] = (sample_obj.prompt_logprob_indices if cache is not None else []) prompt_logprob_output_indices: List[int] = ( sample_obj.prompt_logprob_output_indices if cache is not None else []) - sample_indices: List[int] = (sample_obj.sample_indices - if cache is not None else []) - do_sample = seq_group_metadata.do_sample - - if seq_group_metadata.is_prompt: - if sampling_params.seed is not None: - generator = torch.Generator(device=device).manual_seed( - sampling_params.seed) - if generators is not None: - generators[seq_group_metadata.request_id] = generator - - num_prompts += 1 - num_prefill_sample = len(seq_ids) - assert num_prefill_sample == 1 - assert query_lens is not None and seq_lens is not None - query_len, seq_len = query_lens[i], seq_lens[i] - # If we need sampling, exclude num_prefill_sample tokens from - # prompt logprob. - prompt_logprob_len = (query_len - num_prefill_sample - if do_sample else query_len) - sample_len = num_prefill_sample if do_sample else 0 + sample_indices: List[int] = (sample_obj.sample_indices + if cache is not None else []) + do_sample = seq_group_metadata.do_sample + + if seq_group_metadata.is_prompt: + if sampling_params.seed is not None: + generator = torch.Generator(device=device).manual_seed( + sampling_params.seed) + if generators is not None: + generators[seq_group_metadata.request_id] = generator + + num_prompts += 1 + num_prefill_sample = len(seq_ids) + assert num_prefill_sample == 1 + assert query_lens is not None and seq_lens is not None + query_len, seq_len = query_lens[i], seq_lens[i] + # If we need sampling, exclude num_prefill_sample tokens from + # prompt logprob. + prompt_logprob_len = (query_len - num_prefill_sample + if do_sample else query_len) + sample_len = num_prefill_sample if do_sample else 0 else: - # Decode - prompt_logprob_len = 0 - query_len = query_lens[i] if query_lens is not None else 1 - sample_len = len(seq_ids) * query_len if do_sample else 0 - - if sampling_params.seed is not None and generators is not None: + # Decode + prompt_logprob_len = 0 + query_len = query_lens[i] if query_lens is not None else 1 + sample_len = len(seq_ids) * query_len if do_sample else 0 + + if sampling_params.seed is not None and generators is not None: generator = generators.get(seq_group_metadata.request_id) seq_data = next(iter(seq_group_metadata.seq_data.values())) @@ -340,65 +340,65 @@ def _prepare_seq_groups( seq_data, prompt_logprob_len, )) - - # Update indices to select from the model output. - """ - This blocks computes selected_token_indices which is used in the - following way. - - hidden_states = model(...) - logits = hidden_states[selected_token_indices] - """ - + + # Update indices to select from the model output. + """ + This blocks computes selected_token_indices which is used in the + following way. + + hidden_states = model(...) + logits = hidden_states[selected_token_indices] + """ + if sampling_params.prompt_logprobs is not None: selected_token_indices.extend( model_output_idx + output_index for output_index in prompt_logprob_output_indices) - model_output_idx += prompt_logprob_len - if do_sample: - selected_token_indices.extend( - range(model_output_idx, model_output_idx + sample_len)) - model_output_idx += sample_len - - # We now find indices for logprob computation and sampling. - """ - This block computes categorized_sample_indices which is used in the - following way. - - hidden_states = model(...) - logits = hidden_states[selected_token_indices] - def sample(logits): - # Use categorized_sample_indices for sampling. - # prompt_logprob_indices to find prompt logprob indices. - # sample_indices to find sample indices. - """ - + model_output_idx += prompt_logprob_len + if do_sample: + selected_token_indices.extend( + range(model_output_idx, model_output_idx + sample_len)) + model_output_idx += sample_len + + # We now find indices for logprob computation and sampling. + """ + This block computes categorized_sample_indices which is used in the + following way. + + hidden_states = model(...) + logits = hidden_states[selected_token_indices] + def sample(logits): + # Use categorized_sample_indices for sampling. + # prompt_logprob_indices to find prompt logprob indices. + # sample_indices to find sample indices. + """ + if sampling_params.prompt_logprobs is not None: prompt_logprob_indices.extend( range(logit_idx, logit_idx + len(prompt_logprob_output_indices))) logit_idx += len(prompt_logprob_output_indices) - if do_sample: - sample_indices.extend(range(logit_idx, logit_idx + sample_len)) - categorized_sample_indices[sampling_params.sampling_type].extend( - list(range(logit_idx, logit_idx + sample_len))) - logit_idx += sample_len - + if do_sample: + sample_indices.extend(range(logit_idx, logit_idx + sample_len)) + categorized_sample_indices[sampling_params.sampling_type].extend( + list(range(logit_idx, logit_idx + sample_len))) + logit_idx += sample_len + if cache is not None: sample_obj.sampling_params = sampling_params sample_obj.seq_data = seq_group_metadata.seq_data - sample_obj.seq_len = seq_len - sample_obj.query_len = query_len + sample_obj.seq_len = seq_len + sample_obj.query_len = query_len sample_obj.generator = generator sample_obj.is_prompt = is_prompt else: - sample_obj = SequenceGroupToSample( - seq_ids=list(seq_ids), - sampling_params=sampling_params, - seq_data=seq_group_metadata.seq_data, - seq_len=seq_len, - query_len=query_len, - generator=generator, + sample_obj = SequenceGroupToSample( + seq_ids=list(seq_ids), + sampling_params=sampling_params, + seq_data=seq_group_metadata.seq_data, + seq_len=seq_len, + query_len=query_len, + generator=generator, is_prompt=is_prompt, prompt_logprob_indices=list(prompt_logprob_indices), prompt_logprob_output_indices=list( @@ -409,236 +409,236 @@ def _prepare_seq_groups( assert (len(sample_obj.prompt_logprob_indices) == len(sample_obj.prompt_logprob_output_indices)) seq_groups.append(sample_obj) - - if cache is not None: - cache.reset() - - return (seq_groups, selected_token_indices, categorized_sample_indices, - num_prompts) - - -@dataclass -class SamplingTensors: - """Tensors for sampling.""" - - temperatures: torch.Tensor - top_ps: torch.Tensor - top_ks: torch.Tensor - min_ps: torch.Tensor - presence_penalties: torch.Tensor - frequency_penalties: torch.Tensor - repetition_penalties: torch.Tensor - prompt_tokens: torch.Tensor - output_tokens: torch.Tensor - - @classmethod - def from_sampling_metadata( - cls, - sampling_metadata: "SamplingMetadata", - vocab_size: int, - device: torch.device, - dtype: torch.dtype, - ) -> Tuple["SamplingTensors", bool, bool, bool]: - prompt_tokens: List[array] = [] - output_tokens: List[array] = [] - top_ks: List[int] = [] - temperatures: List[float] = [] - top_ps: List[float] = [] - min_ps: List[float] = [] - presence_penalties: List[float] = [] - frequency_penalties: List[float] = [] - repetition_penalties: List[float] = [] - do_penalties = False - do_top_p_top_k = False - do_min_p = False - - assert sampling_metadata.seq_groups is not None - for seq_group in sampling_metadata.seq_groups: - seq_ids = seq_group.seq_ids - sampling_params = seq_group.sampling_params - temperature = sampling_params.temperature - p = sampling_params.presence_penalty - f = sampling_params.frequency_penalty - r = sampling_params.repetition_penalty - top_p = sampling_params.top_p - min_p = sampling_params.min_p - - # k should not be greater than the vocab size. - top_k = min(sampling_params.top_k, vocab_size) - top_k = vocab_size if top_k == -1 else top_k - if temperature < _SAMPLING_EPS: - # NOTE: Zero temperature means deterministic sampling - # (i.e., greedy sampling or beam search). - # Set the temperature to 1 to avoid division by zero. - temperature = 1.0 - if not do_top_p_top_k and (top_p < 1.0 - _SAMPLING_EPS - or top_k != vocab_size): - do_top_p_top_k = True - if not do_min_p and min_p > _SAMPLING_EPS: - do_min_p = True - if not do_penalties and (abs(p) >= _SAMPLING_EPS - or abs(f) >= _SAMPLING_EPS - or abs(r - 1.0) >= _SAMPLING_EPS): - do_penalties = True - - is_prompt = seq_group.is_prompt - if is_prompt and sampling_params.prompt_logprobs is not None: - # For tokens in the prompt that we only need to get - # their logprobs - query_len = seq_group.query_len - assert query_len is not None - prefill_len = len(seq_group.prompt_logprob_indices) - temperatures += [temperature] * prefill_len - top_ps += [top_p] * prefill_len - top_ks += [top_k] * prefill_len - min_ps += [min_p] * prefill_len - presence_penalties += [0] * prefill_len - frequency_penalties += [0] * prefill_len - repetition_penalties += [1] * prefill_len - - if seq_group.do_sample: - sample_lens = len(seq_group.sample_indices) - assert sample_lens >= len(seq_ids) - temperatures += [temperature] * sample_lens - top_ps += [top_p] * sample_lens - top_ks += [top_k] * sample_lens - min_ps += [min_p] * sample_lens - presence_penalties += [p] * sample_lens - frequency_penalties += [f] * sample_lens - repetition_penalties += [r] * sample_lens - - if do_penalties: - for seq_group in sampling_metadata.seq_groups: - seq_ids = seq_group.seq_ids - if (seq_group.is_prompt - and sampling_params.prompt_logprobs is not None): - prefill_len = len(seq_group.prompt_logprob_indices) - prompt_tokens.extend( - array(VLLM_TOKEN_ID_ARRAY_TYPE) - for _ in range(prefill_len)) - output_tokens.extend( - array(VLLM_TOKEN_ID_ARRAY_TYPE) - for _ in range(prefill_len)) - if seq_group.do_sample: - for seq_id in seq_ids: - seq_data = seq_group.seq_data[seq_id] - prompt_tokens.append(seq_data.prompt_token_ids_array) - output_tokens.append(seq_data.output_token_ids_array) - - sampling_tensors = SamplingTensors.from_lists( - temperatures, - top_ps, - top_ks, - min_ps, - presence_penalties, - frequency_penalties, - repetition_penalties, - prompt_tokens, - output_tokens, - vocab_size, - device, - dtype, - ) - return (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p) - - @classmethod - def from_lists( - cls, - temperatures: List[float], - top_ps: List[float], - top_ks: List[int], - min_ps: List[float], - presence_penalties: List[float], - frequency_penalties: List[float], - repetition_penalties: List[float], - prompt_tokens: List[array], - output_tokens: List[array], - vocab_size: int, - device: torch.device, - dtype: torch.dtype, - ) -> "SamplingTensors": - # Note that the performance will be very bad without - # pinned memory. - pin_memory = is_pin_memory_available() - - do_penalties = prompt_tokens or output_tokens - - if do_penalties: - prompt_t = make_tensor_with_pad( - prompt_tokens, - vocab_size, - device="cpu", - dtype=torch.int64, - pin_memory=pin_memory, - ) - output_t = make_tensor_with_pad( - output_tokens, - vocab_size, - device="cpu", - dtype=torch.int64, - pin_memory=pin_memory, - ) - else: - empty_tensor = torch.empty(0, device=device, dtype=torch.long) - prompt_t = empty_tensor - output_t = empty_tensor - - temperatures_t = torch.tensor( - temperatures, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - top_ps_t = torch.tensor( - top_ps, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - min_ps_t = torch.tensor( - min_ps, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - presence_penalties_t = torch.tensor( - presence_penalties, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - frequency_penalties_t = torch.tensor( - frequency_penalties, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - repetition_penalties_t = torch.tensor( - repetition_penalties, - device="cpu", - dtype=dtype, - pin_memory=pin_memory, - ) - top_ks_t = torch.tensor( - top_ks, - device="cpu", - dtype=torch.int, - pin_memory=pin_memory, - ) - # Because the memory is pinned, we can do non-blocking - # transfer to device. - - return cls( - temperatures=temperatures_t.to(device=device, non_blocking=True), - top_ps=top_ps_t.to(device=device, non_blocking=True), - top_ks=top_ks_t.to(device=device, non_blocking=True), - min_ps=min_ps_t.to(device=device, non_blocking=True), - presence_penalties=presence_penalties_t.to(device=device, - non_blocking=True), - frequency_penalties=frequency_penalties_t.to(device=device, - non_blocking=True), - repetition_penalties=repetition_penalties_t.to(device=device, - non_blocking=True), - prompt_tokens=prompt_t.to(device=device, non_blocking=True), - output_tokens=output_t.to(device=device, non_blocking=True), - ) + + if cache is not None: + cache.reset() + + return (seq_groups, selected_token_indices, categorized_sample_indices, + num_prompts) + + +@dataclass +class SamplingTensors: + """Tensors for sampling.""" + + temperatures: torch.Tensor + top_ps: torch.Tensor + top_ks: torch.Tensor + min_ps: torch.Tensor + presence_penalties: torch.Tensor + frequency_penalties: torch.Tensor + repetition_penalties: torch.Tensor + prompt_tokens: torch.Tensor + output_tokens: torch.Tensor + + @classmethod + def from_sampling_metadata( + cls, + sampling_metadata: "SamplingMetadata", + vocab_size: int, + device: torch.device, + dtype: torch.dtype, + ) -> Tuple["SamplingTensors", bool, bool, bool]: + prompt_tokens: List[array] = [] + output_tokens: List[array] = [] + top_ks: List[int] = [] + temperatures: List[float] = [] + top_ps: List[float] = [] + min_ps: List[float] = [] + presence_penalties: List[float] = [] + frequency_penalties: List[float] = [] + repetition_penalties: List[float] = [] + do_penalties = False + do_top_p_top_k = False + do_min_p = False + + assert sampling_metadata.seq_groups is not None + for seq_group in sampling_metadata.seq_groups: + seq_ids = seq_group.seq_ids + sampling_params = seq_group.sampling_params + temperature = sampling_params.temperature + p = sampling_params.presence_penalty + f = sampling_params.frequency_penalty + r = sampling_params.repetition_penalty + top_p = sampling_params.top_p + min_p = sampling_params.min_p + + # k should not be greater than the vocab size. + top_k = min(sampling_params.top_k, vocab_size) + top_k = vocab_size if top_k == -1 else top_k + if temperature < _SAMPLING_EPS: + # NOTE: Zero temperature means deterministic sampling + # (i.e., greedy sampling or beam search). + # Set the temperature to 1 to avoid division by zero. + temperature = 1.0 + if not do_top_p_top_k and (top_p < 1.0 - _SAMPLING_EPS + or top_k != vocab_size): + do_top_p_top_k = True + if not do_min_p and min_p > _SAMPLING_EPS: + do_min_p = True + if not do_penalties and (abs(p) >= _SAMPLING_EPS + or abs(f) >= _SAMPLING_EPS + or abs(r - 1.0) >= _SAMPLING_EPS): + do_penalties = True + + is_prompt = seq_group.is_prompt + if is_prompt and sampling_params.prompt_logprobs is not None: + # For tokens in the prompt that we only need to get + # their logprobs + query_len = seq_group.query_len + assert query_len is not None + prefill_len = len(seq_group.prompt_logprob_indices) + temperatures += [temperature] * prefill_len + top_ps += [top_p] * prefill_len + top_ks += [top_k] * prefill_len + min_ps += [min_p] * prefill_len + presence_penalties += [0] * prefill_len + frequency_penalties += [0] * prefill_len + repetition_penalties += [1] * prefill_len + + if seq_group.do_sample: + sample_lens = len(seq_group.sample_indices) + assert sample_lens >= len(seq_ids) + temperatures += [temperature] * sample_lens + top_ps += [top_p] * sample_lens + top_ks += [top_k] * sample_lens + min_ps += [min_p] * sample_lens + presence_penalties += [p] * sample_lens + frequency_penalties += [f] * sample_lens + repetition_penalties += [r] * sample_lens + + if do_penalties: + for seq_group in sampling_metadata.seq_groups: + seq_ids = seq_group.seq_ids + if (seq_group.is_prompt + and sampling_params.prompt_logprobs is not None): + prefill_len = len(seq_group.prompt_logprob_indices) + prompt_tokens.extend( + array(VLLM_TOKEN_ID_ARRAY_TYPE) + for _ in range(prefill_len)) + output_tokens.extend( + array(VLLM_TOKEN_ID_ARRAY_TYPE) + for _ in range(prefill_len)) + if seq_group.do_sample: + for seq_id in seq_ids: + seq_data = seq_group.seq_data[seq_id] + prompt_tokens.append(seq_data.prompt_token_ids_array) + output_tokens.append(seq_data.output_token_ids_array) + + sampling_tensors = SamplingTensors.from_lists( + temperatures, + top_ps, + top_ks, + min_ps, + presence_penalties, + frequency_penalties, + repetition_penalties, + prompt_tokens, + output_tokens, + vocab_size, + device, + dtype, + ) + return (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p) + + @classmethod + def from_lists( + cls, + temperatures: List[float], + top_ps: List[float], + top_ks: List[int], + min_ps: List[float], + presence_penalties: List[float], + frequency_penalties: List[float], + repetition_penalties: List[float], + prompt_tokens: List[array], + output_tokens: List[array], + vocab_size: int, + device: torch.device, + dtype: torch.dtype, + ) -> "SamplingTensors": + # Note that the performance will be very bad without + # pinned memory. + pin_memory = is_pin_memory_available() + + do_penalties = prompt_tokens or output_tokens + + if do_penalties: + prompt_t = make_tensor_with_pad( + prompt_tokens, + vocab_size, + device="cpu", + dtype=torch.int64, + pin_memory=pin_memory, + ) + output_t = make_tensor_with_pad( + output_tokens, + vocab_size, + device="cpu", + dtype=torch.int64, + pin_memory=pin_memory, + ) + else: + empty_tensor = torch.empty(0, device=device, dtype=torch.long) + prompt_t = empty_tensor + output_t = empty_tensor + + temperatures_t = torch.tensor( + temperatures, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + top_ps_t = torch.tensor( + top_ps, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + min_ps_t = torch.tensor( + min_ps, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + presence_penalties_t = torch.tensor( + presence_penalties, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + frequency_penalties_t = torch.tensor( + frequency_penalties, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + repetition_penalties_t = torch.tensor( + repetition_penalties, + device="cpu", + dtype=dtype, + pin_memory=pin_memory, + ) + top_ks_t = torch.tensor( + top_ks, + device="cpu", + dtype=torch.int, + pin_memory=pin_memory, + ) + # Because the memory is pinned, we can do non-blocking + # transfer to device. + + return cls( + temperatures=temperatures_t.to(device=device, non_blocking=True), + top_ps=top_ps_t.to(device=device, non_blocking=True), + top_ks=top_ks_t.to(device=device, non_blocking=True), + min_ps=min_ps_t.to(device=device, non_blocking=True), + presence_penalties=presence_penalties_t.to(device=device, + non_blocking=True), + frequency_penalties=frequency_penalties_t.to(device=device, + non_blocking=True), + repetition_penalties=repetition_penalties_t.to(device=device, + non_blocking=True), + prompt_tokens=prompt_t.to(device=device, non_blocking=True), + output_tokens=output_t.to(device=device, non_blocking=True), + ) diff --git a/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py b/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py index 884d07e0..58790be1 100644 --- a/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py +++ b/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py @@ -1,391 +1,391 @@ -"""Sampling parameters for text generation.""" -import copy -from dataclasses import dataclass -from enum import Enum, IntEnum -from functools import cached_property -from typing import Any, Callable, Dict, List, Optional, Set, Union - -import msgspec -import torch -from pydantic import BaseModel -from typing_extensions import Annotated - -from vllm.logger import init_logger - -logger = init_logger(__name__) - -_SAMPLING_EPS = 1e-5 -_MAX_TEMP = 1e-2 - - -class SamplingType(IntEnum): - GREEDY = 0 - RANDOM = 1 - RANDOM_SEED = 2 - - -LogitsProcessor = Union[Callable[[List[int], torch.Tensor], torch.Tensor], - Callable[[List[int], List[int], torch.Tensor], - torch.Tensor]] -"""LogitsProcessor is a function that takes a list -of previously generated tokens, the logits tensor -for the next token and, optionally, prompt tokens as a -first argument, and returns a modified tensor of logits -to sample from.""" - - -# maybe make msgspec? -@dataclass -class GuidedDecodingParams: - """One of these fields will be used to build a logit processor.""" - json: Optional[Union[str, Dict]] = None - regex: Optional[str] = None - choice: Optional[List[str]] = None - grammar: Optional[str] = None - json_object: Optional[bool] = None - """These are other options that can be set""" - backend: Optional[str] = None - whitespace_pattern: Optional[str] = None - - @staticmethod - def from_optional( - json: Optional[Union[Dict, BaseModel, str]], - regex: Optional[str] = None, - choice: Optional[List[str]] = None, - grammar: Optional[str] = None, - json_object: Optional[bool] = None, - backend: Optional[str] = None, - whitespace_pattern: Optional[str] = None, - ) -> "GuidedDecodingParams": - # Extract json schemas from pydantic models - if isinstance(json, (BaseModel, type(BaseModel))): - json = json.model_json_schema() - return GuidedDecodingParams( - json=json, - regex=regex, - choice=choice, - grammar=grammar, - json_object=json_object, - backend=backend, - whitespace_pattern=whitespace_pattern, - ) - - def __post_init__(self): - """Validate that some fields are mutually exclusive.""" - guide_count = sum([ - self.json is not None, self.regex is not None, self.choice - is not None, self.grammar is not None, self.json_object is not None - ]) - if guide_count > 1: - raise ValueError( - "You can only use one kind of guided decoding but multiple are " - f"specified: {self.__dict__}") - - -class RequestOutputKind(Enum): - # Return entire output so far in every RequestOutput - CUMULATIVE = 0 - # Return only deltas in each RequestOutput - DELTA = 1 - # Do not return intermediate RequestOuputs - FINAL_ONLY = 2 - - -class SamplingParams( - msgspec.Struct, - omit_defaults=True, # type: ignore[call-arg] - # required for @cached_property. - dict=True): # type: ignore[call-arg] - """Sampling parameters for text generation. - - Overall, we follow the sampling parameters from the OpenAI text completion - API (https://platform.openai.com/docs/api-reference/completions/create). - In addition, we support beam search, which is not supported by OpenAI. - - Args: - n: Number of output sequences to return for the given prompt. - best_of: Number of output sequences that are generated from the prompt. - From these `best_of` sequences, the top `n` sequences are returned. - `best_of` must be greater than or equal to `n`. By default, - `best_of` is set to `n`. - presence_penalty: Float that penalizes new tokens based on whether they - appear in the generated text so far. Values > 0 encourage the model - to use new tokens, while values < 0 encourage the model to repeat - tokens. - frequency_penalty: Float that penalizes new tokens based on their - frequency in the generated text so far. Values > 0 encourage the - model to use new tokens, while values < 0 encourage the model to - repeat tokens. - repetition_penalty: Float that penalizes new tokens based on whether - they appear in the prompt and the generated text so far. Values > 1 - encourage the model to use new tokens, while values < 1 encourage - the model to repeat tokens. - temperature: Float that controls the randomness of the sampling. Lower - values make the model more deterministic, while higher values make - the model more random. Zero means greedy sampling. - top_p: Float that controls the cumulative probability of the top tokens - to consider. Must be in (0, 1]. Set to 1 to consider all tokens. - top_k: Integer that controls the number of top tokens to consider. Set - to -1 to consider all tokens. - min_p: Float that represents the minimum probability for a token to be - considered, relative to the probability of the most likely token. - Must be in [0, 1]. Set to 0 to disable this. - seed: Random seed to use for the generation. - stop: List of strings that stop the generation when they are generated. - The returned output will not contain the stop strings. - stop_token_ids: List of tokens that stop the generation when they are - generated. The returned output will contain the stop tokens unless - the stop tokens are special tokens. - include_stop_str_in_output: Whether to include the stop strings in - output text. Defaults to False. - ignore_eos: Whether to ignore the EOS token and continue generating - tokens after the EOS token is generated. - max_tokens: Maximum number of tokens to generate per output sequence. - min_tokens: Minimum number of tokens to generate per output sequence - before EOS or stop_token_ids can be generated - logprobs: Number of log probabilities to return per output token. - When set to None, no probability is returned. If set to a non-None - value, the result includes the log probabilities of the specified - number of most likely tokens, as well as the chosen tokens. - Note that the implementation follows the OpenAI API: The API will - always return the log probability of the sampled token, so there - may be up to `logprobs+1` elements in the response. - prompt_logprobs: Number of log probabilities to return per prompt token. - detokenize: Whether to detokenize the output. Defaults to True. - skip_special_tokens: Whether to skip special tokens in the output. - spaces_between_special_tokens: Whether to add spaces between special - tokens in the output. Defaults to True. - logits_processors: List of functions that modify logits based on - previously generated tokens, and optionally prompt tokens as - a first argument. - truncate_prompt_tokens: If set to an integer k, will use only the last k - tokens from the prompt (i.e., left truncation). Defaults to None - (i.e., no truncation). - guided_decoding: If provided, the engine will construct a guided - decoding logits processor from these parameters. Defaults to None. - logit_bias: If provided, the engine will construct a logits processor - that applies these logit biases. Defaults to None. +"""Sampling parameters for text generation.""" +import copy +from dataclasses import dataclass +from enum import Enum, IntEnum +from functools import cached_property +from typing import Any, Callable, Dict, List, Optional, Set, Union + +import msgspec +import torch +from pydantic import BaseModel +from typing_extensions import Annotated + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_SAMPLING_EPS = 1e-5 +_MAX_TEMP = 1e-2 + + +class SamplingType(IntEnum): + GREEDY = 0 + RANDOM = 1 + RANDOM_SEED = 2 + + +LogitsProcessor = Union[Callable[[List[int], torch.Tensor], torch.Tensor], + Callable[[List[int], List[int], torch.Tensor], + torch.Tensor]] +"""LogitsProcessor is a function that takes a list +of previously generated tokens, the logits tensor +for the next token and, optionally, prompt tokens as a +first argument, and returns a modified tensor of logits +to sample from.""" + + +# maybe make msgspec? +@dataclass +class GuidedDecodingParams: + """One of these fields will be used to build a logit processor.""" + json: Optional[Union[str, Dict]] = None + regex: Optional[str] = None + choice: Optional[List[str]] = None + grammar: Optional[str] = None + json_object: Optional[bool] = None + """These are other options that can be set""" + backend: Optional[str] = None + whitespace_pattern: Optional[str] = None + + @staticmethod + def from_optional( + json: Optional[Union[Dict, BaseModel, str]], + regex: Optional[str] = None, + choice: Optional[List[str]] = None, + grammar: Optional[str] = None, + json_object: Optional[bool] = None, + backend: Optional[str] = None, + whitespace_pattern: Optional[str] = None, + ) -> "GuidedDecodingParams": + # Extract json schemas from pydantic models + if isinstance(json, (BaseModel, type(BaseModel))): + json = json.model_json_schema() + return GuidedDecodingParams( + json=json, + regex=regex, + choice=choice, + grammar=grammar, + json_object=json_object, + backend=backend, + whitespace_pattern=whitespace_pattern, + ) + + def __post_init__(self): + """Validate that some fields are mutually exclusive.""" + guide_count = sum([ + self.json is not None, self.regex is not None, self.choice + is not None, self.grammar is not None, self.json_object is not None + ]) + if guide_count > 1: + raise ValueError( + "You can only use one kind of guided decoding but multiple are " + f"specified: {self.__dict__}") + + +class RequestOutputKind(Enum): + # Return entire output so far in every RequestOutput + CUMULATIVE = 0 + # Return only deltas in each RequestOutput + DELTA = 1 + # Do not return intermediate RequestOuputs + FINAL_ONLY = 2 + + +class SamplingParams( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + # required for @cached_property. + dict=True): # type: ignore[call-arg] + """Sampling parameters for text generation. + + Overall, we follow the sampling parameters from the OpenAI text completion + API (https://platform.openai.com/docs/api-reference/completions/create). + In addition, we support beam search, which is not supported by OpenAI. + + Args: + n: Number of output sequences to return for the given prompt. + best_of: Number of output sequences that are generated from the prompt. + From these `best_of` sequences, the top `n` sequences are returned. + `best_of` must be greater than or equal to `n`. By default, + `best_of` is set to `n`. + presence_penalty: Float that penalizes new tokens based on whether they + appear in the generated text so far. Values > 0 encourage the model + to use new tokens, while values < 0 encourage the model to repeat + tokens. + frequency_penalty: Float that penalizes new tokens based on their + frequency in the generated text so far. Values > 0 encourage the + model to use new tokens, while values < 0 encourage the model to + repeat tokens. + repetition_penalty: Float that penalizes new tokens based on whether + they appear in the prompt and the generated text so far. Values > 1 + encourage the model to use new tokens, while values < 1 encourage + the model to repeat tokens. + temperature: Float that controls the randomness of the sampling. Lower + values make the model more deterministic, while higher values make + the model more random. Zero means greedy sampling. + top_p: Float that controls the cumulative probability of the top tokens + to consider. Must be in (0, 1]. Set to 1 to consider all tokens. + top_k: Integer that controls the number of top tokens to consider. Set + to -1 to consider all tokens. + min_p: Float that represents the minimum probability for a token to be + considered, relative to the probability of the most likely token. + Must be in [0, 1]. Set to 0 to disable this. + seed: Random seed to use for the generation. + stop: List of strings that stop the generation when they are generated. + The returned output will not contain the stop strings. + stop_token_ids: List of tokens that stop the generation when they are + generated. The returned output will contain the stop tokens unless + the stop tokens are special tokens. + include_stop_str_in_output: Whether to include the stop strings in + output text. Defaults to False. + ignore_eos: Whether to ignore the EOS token and continue generating + tokens after the EOS token is generated. + max_tokens: Maximum number of tokens to generate per output sequence. + min_tokens: Minimum number of tokens to generate per output sequence + before EOS or stop_token_ids can be generated + logprobs: Number of log probabilities to return per output token. + When set to None, no probability is returned. If set to a non-None + value, the result includes the log probabilities of the specified + number of most likely tokens, as well as the chosen tokens. + Note that the implementation follows the OpenAI API: The API will + always return the log probability of the sampled token, so there + may be up to `logprobs+1` elements in the response. + prompt_logprobs: Number of log probabilities to return per prompt token. + detokenize: Whether to detokenize the output. Defaults to True. + skip_special_tokens: Whether to skip special tokens in the output. + spaces_between_special_tokens: Whether to add spaces between special + tokens in the output. Defaults to True. + logits_processors: List of functions that modify logits based on + previously generated tokens, and optionally prompt tokens as + a first argument. + truncate_prompt_tokens: If set to an integer k, will use only the last k + tokens from the prompt (i.e., left truncation). Defaults to None + (i.e., no truncation). + guided_decoding: If provided, the engine will construct a guided + decoding logits processor from these parameters. Defaults to None. + logit_bias: If provided, the engine will construct a logits processor + that applies these logit biases. Defaults to None. allowed_token_ids: If provided, the engine will construct a logits processor which only retains scores for the given token ids. Defaults to None. prompt_logprob_positions: Optional prompt-token positions whose logits should be materialized. None preserves the standard all-position prompt-logprob behavior. - """ - - n: int = 1 - best_of: Optional[int] = None - _real_n: Optional[int] = None - presence_penalty: float = 0.0 - frequency_penalty: float = 0.0 - repetition_penalty: float = 1.0 - temperature: float = 1.0 - top_p: float = 1.0 - top_k: int = -1 - min_p: float = 0.0 - seed: Optional[int] = None - stop: Optional[Union[str, List[str]]] = None - stop_token_ids: Optional[List[int]] = None - ignore_eos: bool = False - max_tokens: Optional[int] = 16 - min_tokens: int = 0 - logprobs: Optional[int] = None - prompt_logprobs: Optional[int] = None - # NOTE: This parameter is only exposed at the engine level for now. - # It is not exposed in the OpenAI API server, as the OpenAI API does - # not support returning only a list of token IDs. - detokenize: bool = True - skip_special_tokens: bool = True - spaces_between_special_tokens: bool = True - # Optional[List[LogitsProcessor]] type. We use Any here because - # Optional[List[LogitsProcessor]] type is not supported by msgspec. - logits_processors: Optional[Any] = None - include_stop_str_in_output: bool = False - truncate_prompt_tokens: Optional[Annotated[int, msgspec.Meta(ge=1)]] = None - output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE - - # The below fields are not supposed to be used as an input. - # They are set in post_init. - output_text_buffer_length: int = 0 - _all_stop_token_ids: Set[int] = msgspec.field(default_factory=set) - - # Fields used to construct logits processors + """ + + n: int = 1 + best_of: Optional[int] = None + _real_n: Optional[int] = None + presence_penalty: float = 0.0 + frequency_penalty: float = 0.0 + repetition_penalty: float = 1.0 + temperature: float = 1.0 + top_p: float = 1.0 + top_k: int = -1 + min_p: float = 0.0 + seed: Optional[int] = None + stop: Optional[Union[str, List[str]]] = None + stop_token_ids: Optional[List[int]] = None + ignore_eos: bool = False + max_tokens: Optional[int] = 16 + min_tokens: int = 0 + logprobs: Optional[int] = None + prompt_logprobs: Optional[int] = None + # NOTE: This parameter is only exposed at the engine level for now. + # It is not exposed in the OpenAI API server, as the OpenAI API does + # not support returning only a list of token IDs. + detokenize: bool = True + skip_special_tokens: bool = True + spaces_between_special_tokens: bool = True + # Optional[List[LogitsProcessor]] type. We use Any here because + # Optional[List[LogitsProcessor]] type is not supported by msgspec. + logits_processors: Optional[Any] = None + include_stop_str_in_output: bool = False + truncate_prompt_tokens: Optional[Annotated[int, msgspec.Meta(ge=1)]] = None + output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE + + # The below fields are not supposed to be used as an input. + # They are set in post_init. + output_text_buffer_length: int = 0 + _all_stop_token_ids: Set[int] = msgspec.field(default_factory=set) + + # Fields used to construct logits processors guided_decoding: Optional[GuidedDecodingParams] = None logit_bias: Optional[Dict[int, float]] = None allowed_token_ids: Optional[List[int]] = None prompt_logprob_positions: Optional[List[int]] = None - - @staticmethod - def from_optional( - n: Optional[int] = 1, - best_of: Optional[int] = None, - presence_penalty: Optional[float] = 0.0, - frequency_penalty: Optional[float] = 0.0, - repetition_penalty: Optional[float] = 1.0, - temperature: Optional[float] = 1.0, - top_p: Optional[float] = 1.0, - top_k: int = -1, - min_p: float = 0.0, - seed: Optional[int] = None, - stop: Optional[Union[str, List[str]]] = None, - stop_token_ids: Optional[List[int]] = None, - include_stop_str_in_output: bool = False, - ignore_eos: bool = False, - max_tokens: Optional[int] = 16, - min_tokens: int = 0, - logprobs: Optional[int] = None, - prompt_logprobs: Optional[int] = None, - detokenize: bool = True, - skip_special_tokens: bool = True, - spaces_between_special_tokens: bool = True, - logits_processors: Optional[List[LogitsProcessor]] = None, - truncate_prompt_tokens: Optional[Annotated[int, - msgspec.Meta(ge=1)]] = None, - output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE, + + @staticmethod + def from_optional( + n: Optional[int] = 1, + best_of: Optional[int] = None, + presence_penalty: Optional[float] = 0.0, + frequency_penalty: Optional[float] = 0.0, + repetition_penalty: Optional[float] = 1.0, + temperature: Optional[float] = 1.0, + top_p: Optional[float] = 1.0, + top_k: int = -1, + min_p: float = 0.0, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + include_stop_str_in_output: bool = False, + ignore_eos: bool = False, + max_tokens: Optional[int] = 16, + min_tokens: int = 0, + logprobs: Optional[int] = None, + prompt_logprobs: Optional[int] = None, + detokenize: bool = True, + skip_special_tokens: bool = True, + spaces_between_special_tokens: bool = True, + logits_processors: Optional[List[LogitsProcessor]] = None, + truncate_prompt_tokens: Optional[Annotated[int, + msgspec.Meta(ge=1)]] = None, + output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE, guided_decoding: Optional[GuidedDecodingParams] = None, logit_bias: Optional[Union[Dict[int, float], Dict[str, float]]] = None, allowed_token_ids: Optional[List[int]] = None, prompt_logprob_positions: Optional[List[int]] = None, - ) -> "SamplingParams": - if logit_bias is not None: - logit_bias = { - int(token): bias - for token, bias in logit_bias.items() - } - - return SamplingParams( - n=1 if n is None else n, - best_of=best_of, - presence_penalty=0.0 - if presence_penalty is None else presence_penalty, - frequency_penalty=0.0 - if frequency_penalty is None else frequency_penalty, - repetition_penalty=1.0 - if repetition_penalty is None else repetition_penalty, - temperature=1.0 if temperature is None else temperature, - top_p=1.0 if top_p is None else top_p, - top_k=top_k, - min_p=min_p, - seed=seed, - stop=stop, - stop_token_ids=stop_token_ids, - include_stop_str_in_output=include_stop_str_in_output, - ignore_eos=ignore_eos, - max_tokens=max_tokens, - min_tokens=min_tokens, - logprobs=logprobs, - prompt_logprobs=prompt_logprobs, - detokenize=detokenize, - skip_special_tokens=skip_special_tokens, - spaces_between_special_tokens=spaces_between_special_tokens, - logits_processors=logits_processors, - truncate_prompt_tokens=truncate_prompt_tokens, - output_kind=output_kind, + ) -> "SamplingParams": + if logit_bias is not None: + logit_bias = { + int(token): bias + for token, bias in logit_bias.items() + } + + return SamplingParams( + n=1 if n is None else n, + best_of=best_of, + presence_penalty=0.0 + if presence_penalty is None else presence_penalty, + frequency_penalty=0.0 + if frequency_penalty is None else frequency_penalty, + repetition_penalty=1.0 + if repetition_penalty is None else repetition_penalty, + temperature=1.0 if temperature is None else temperature, + top_p=1.0 if top_p is None else top_p, + top_k=top_k, + min_p=min_p, + seed=seed, + stop=stop, + stop_token_ids=stop_token_ids, + include_stop_str_in_output=include_stop_str_in_output, + ignore_eos=ignore_eos, + max_tokens=max_tokens, + min_tokens=min_tokens, + logprobs=logprobs, + prompt_logprobs=prompt_logprobs, + detokenize=detokenize, + skip_special_tokens=skip_special_tokens, + spaces_between_special_tokens=spaces_between_special_tokens, + logits_processors=logits_processors, + truncate_prompt_tokens=truncate_prompt_tokens, + output_kind=output_kind, guided_decoding=guided_decoding, logit_bias=logit_bias, allowed_token_ids=allowed_token_ids, prompt_logprob_positions=prompt_logprob_positions, ) - - def __post_init__(self) -> None: - # how we deal with `best_of``: - # if `best_of`` is not set, we default to `n`; - # if `best_of`` is set, we set `n`` to `best_of`, - # and set `_real_n`` to the original `n`. - # when we return the result, we will check - # if we need to return `n` or `_real_n` results - if self.best_of: - if self.best_of < self.n: - raise ValueError( - f"best_of must be greater than or equal to n, " - f"got n={self.n} and best_of={self.best_of}.") - self._real_n = self.n - self.n = self.best_of - if 0 < self.temperature < _MAX_TEMP: - logger.warning( - "temperature %s is less than %s, which may cause numerical " - "errors nan or inf in tensors. We have maxed it out to %s.", - self.temperature, _MAX_TEMP, _MAX_TEMP) - self.temperature = max(self.temperature, _MAX_TEMP) - if self.seed == -1: - self.seed = None - else: - self.seed = self.seed - if self.stop is None: - self.stop = [] - elif isinstance(self.stop, str): - self.stop = [self.stop] - else: - self.stop = list(self.stop) - if self.stop_token_ids is None: - self.stop_token_ids = [] - else: - self.stop_token_ids = list(self.stop_token_ids) - self.logprobs = 1 if self.logprobs is True else self.logprobs + + def __post_init__(self) -> None: + # how we deal with `best_of``: + # if `best_of`` is not set, we default to `n`; + # if `best_of`` is set, we set `n`` to `best_of`, + # and set `_real_n`` to the original `n`. + # when we return the result, we will check + # if we need to return `n` or `_real_n` results + if self.best_of: + if self.best_of < self.n: + raise ValueError( + f"best_of must be greater than or equal to n, " + f"got n={self.n} and best_of={self.best_of}.") + self._real_n = self.n + self.n = self.best_of + if 0 < self.temperature < _MAX_TEMP: + logger.warning( + "temperature %s is less than %s, which may cause numerical " + "errors nan or inf in tensors. We have maxed it out to %s.", + self.temperature, _MAX_TEMP, _MAX_TEMP) + self.temperature = max(self.temperature, _MAX_TEMP) + if self.seed == -1: + self.seed = None + else: + self.seed = self.seed + if self.stop is None: + self.stop = [] + elif isinstance(self.stop, str): + self.stop = [self.stop] + else: + self.stop = list(self.stop) + if self.stop_token_ids is None: + self.stop_token_ids = [] + else: + self.stop_token_ids = list(self.stop_token_ids) + self.logprobs = 1 if self.logprobs is True else self.logprobs self.prompt_logprobs = (1 if self.prompt_logprobs is True else self.prompt_logprobs) if self.prompt_logprob_positions is not None: self.prompt_logprob_positions = list( self.prompt_logprob_positions) - - # Number of characters to hold back for stop string evaluation - # until sequence is finished. - if self.stop and not self.include_stop_str_in_output: - self.output_text_buffer_length = max(len(s) for s in self.stop) - 1 - - self._verify_args() - - if self.temperature < _SAMPLING_EPS: - # Zero temperature means greedy sampling. - self.top_p = 1.0 - self.top_k = -1 - self.min_p = 0.0 - self._verify_greedy_sampling() - # eos_token_id is added to this by the engine - self._all_stop_token_ids = set(self.stop_token_ids) - - def _verify_args(self) -> None: - if not isinstance(self.n, int): - raise ValueError(f"n must be an int, but is of " - f"type {type(self.n)}") - if self.n < 1: - raise ValueError(f"n must be at least 1, got {self.n}.") - if not -2.0 <= self.presence_penalty <= 2.0: - raise ValueError("presence_penalty must be in [-2, 2], got " - f"{self.presence_penalty}.") - if not -2.0 <= self.frequency_penalty <= 2.0: - raise ValueError("frequency_penalty must be in [-2, 2], got " - f"{self.frequency_penalty}.") - if not 0.0 < self.repetition_penalty <= 2.0: - raise ValueError("repetition_penalty must be in (0, 2], got " - f"{self.repetition_penalty}.") - if self.temperature < 0.0: - raise ValueError( - f"temperature must be non-negative, got {self.temperature}.") - if not 0.0 < self.top_p <= 1.0: - raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.") - if self.top_k < -1 or self.top_k == 0: - raise ValueError(f"top_k must be -1 (disable), or at least 1, " - f"got {self.top_k}.") - if not isinstance(self.top_k, int): - raise TypeError( - f"top_k must be an integer, got {type(self.top_k).__name__}") - if not 0.0 <= self.min_p <= 1.0: - raise ValueError("min_p must be in [0, 1], got " - f"{self.min_p}.") - if self.max_tokens is not None and self.max_tokens < 1: - raise ValueError( - f"max_tokens must be at least 1, got {self.max_tokens}.") - if self.min_tokens < 0: - raise ValueError(f"min_tokens must be greater than or equal to 0, " - f"got {self.min_tokens}.") - if self.max_tokens is not None and self.min_tokens > self.max_tokens: - raise ValueError( - f"min_tokens must be less than or equal to " - f"max_tokens={self.max_tokens}, got {self.min_tokens}.") - if self.logprobs is not None and self.logprobs < 0: - raise ValueError( - f"logprobs must be non-negative, got {self.logprobs}.") + + # Number of characters to hold back for stop string evaluation + # until sequence is finished. + if self.stop and not self.include_stop_str_in_output: + self.output_text_buffer_length = max(len(s) for s in self.stop) - 1 + + self._verify_args() + + if self.temperature < _SAMPLING_EPS: + # Zero temperature means greedy sampling. + self.top_p = 1.0 + self.top_k = -1 + self.min_p = 0.0 + self._verify_greedy_sampling() + # eos_token_id is added to this by the engine + self._all_stop_token_ids = set(self.stop_token_ids) + + def _verify_args(self) -> None: + if not isinstance(self.n, int): + raise ValueError(f"n must be an int, but is of " + f"type {type(self.n)}") + if self.n < 1: + raise ValueError(f"n must be at least 1, got {self.n}.") + if not -2.0 <= self.presence_penalty <= 2.0: + raise ValueError("presence_penalty must be in [-2, 2], got " + f"{self.presence_penalty}.") + if not -2.0 <= self.frequency_penalty <= 2.0: + raise ValueError("frequency_penalty must be in [-2, 2], got " + f"{self.frequency_penalty}.") + if not 0.0 < self.repetition_penalty <= 2.0: + raise ValueError("repetition_penalty must be in (0, 2], got " + f"{self.repetition_penalty}.") + if self.temperature < 0.0: + raise ValueError( + f"temperature must be non-negative, got {self.temperature}.") + if not 0.0 < self.top_p <= 1.0: + raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.") + if self.top_k < -1 or self.top_k == 0: + raise ValueError(f"top_k must be -1 (disable), or at least 1, " + f"got {self.top_k}.") + if not isinstance(self.top_k, int): + raise TypeError( + f"top_k must be an integer, got {type(self.top_k).__name__}") + if not 0.0 <= self.min_p <= 1.0: + raise ValueError("min_p must be in [0, 1], got " + f"{self.min_p}.") + if self.max_tokens is not None and self.max_tokens < 1: + raise ValueError( + f"max_tokens must be at least 1, got {self.max_tokens}.") + if self.min_tokens < 0: + raise ValueError(f"min_tokens must be greater than or equal to 0, " + f"got {self.min_tokens}.") + if self.max_tokens is not None and self.min_tokens > self.max_tokens: + raise ValueError( + f"min_tokens must be less than or equal to " + f"max_tokens={self.max_tokens}, got {self.min_tokens}.") + if self.logprobs is not None and self.logprobs < 0: + raise ValueError( + f"logprobs must be non-negative, got {self.logprobs}.") if self.prompt_logprobs is not None and self.prompt_logprobs < 0: raise ValueError(f"prompt_logprobs must be non-negative, got " f"{self.prompt_logprobs}.") @@ -407,114 +407,114 @@ class SamplingParams( raise ValueError( "prompt_logprob_positions must be a sorted unique list " "of positive integers.") - if (self.truncate_prompt_tokens is not None - and self.truncate_prompt_tokens < 1): - raise ValueError(f"truncate_prompt_tokens must be >= 1, " - f"got {self.truncate_prompt_tokens}") - assert isinstance(self.stop, list) - if any(not stop_str for stop_str in self.stop): - raise ValueError("stop cannot contain an empty string.") - if self.stop and not self.detokenize: - raise ValueError( - "stop strings are only supported when detokenize is True. " - "Set detokenize=True to use stop.") - if self.best_of != self._real_n and self.output_kind == ( - RequestOutputKind.DELTA): - raise ValueError("best_of must equal n to use output_kind=DELTA") - - def _verify_greedy_sampling(self) -> None: - if self.n > 1: - raise ValueError("n must be 1 when using greedy sampling, " - f"got {self.n}.") - - def update_from_generation_config( - self, - generation_config: Dict[str, Any], - model_eos_token_id: Optional[int] = None) -> None: - """Update if there are non-default values from generation_config""" - - if model_eos_token_id is not None: - # Add the eos token id into the sampling_params to support - # min_tokens processing. - self._all_stop_token_ids.add(model_eos_token_id) - - # Update eos_token_id for generation - if (eos_ids := generation_config.get("eos_token_id")) is not None: - # it can be either int or list of int - eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids) - if model_eos_token_id is not None: - # We don't need to include the primary eos_token_id in - # stop_token_ids since it's handled separately for stopping - # purposes. - eos_ids.discard(model_eos_token_id) - if eos_ids: - self._all_stop_token_ids.update(eos_ids) - if not self.ignore_eos: - eos_ids.update(self.stop_token_ids) - self.stop_token_ids = list(eos_ids) - - @cached_property - def sampling_type(self) -> SamplingType: - if self.temperature < _SAMPLING_EPS: - return SamplingType.GREEDY - if self.seed is not None: - return SamplingType.RANDOM_SEED - return SamplingType.RANDOM - - @property - def all_stop_token_ids(self) -> Set[int]: - return self._all_stop_token_ids - - def clone(self) -> "SamplingParams": - """Deep copy excluding LogitsProcessor objects. - - LogitsProcessor objects are excluded because they may contain an - arbitrary, nontrivial amount of data. - See https://github.com/vllm-project/vllm/issues/3087 - """ - - logit_processor_refs = None if self.logits_processors is None else { - id(lp): lp - for lp in self.logits_processors - } - return copy.deepcopy(self, memo=logit_processor_refs) - - def __repr__(self) -> str: - return ( - f"SamplingParams(n={self.n}, " - f"presence_penalty={self.presence_penalty}, " - f"frequency_penalty={self.frequency_penalty}, " - f"repetition_penalty={self.repetition_penalty}, " - f"temperature={self.temperature}, " - f"top_p={self.top_p}, " - f"top_k={self.top_k}, " - f"min_p={self.min_p}, " - f"seed={self.seed}, " - f"stop={self.stop}, " - f"stop_token_ids={self.stop_token_ids}, " - f"include_stop_str_in_output={self.include_stop_str_in_output}, " - f"ignore_eos={self.ignore_eos}, " - f"max_tokens={self.max_tokens}, " - f"min_tokens={self.min_tokens}, " + if (self.truncate_prompt_tokens is not None + and self.truncate_prompt_tokens < 1): + raise ValueError(f"truncate_prompt_tokens must be >= 1, " + f"got {self.truncate_prompt_tokens}") + assert isinstance(self.stop, list) + if any(not stop_str for stop_str in self.stop): + raise ValueError("stop cannot contain an empty string.") + if self.stop and not self.detokenize: + raise ValueError( + "stop strings are only supported when detokenize is True. " + "Set detokenize=True to use stop.") + if self.best_of != self._real_n and self.output_kind == ( + RequestOutputKind.DELTA): + raise ValueError("best_of must equal n to use output_kind=DELTA") + + def _verify_greedy_sampling(self) -> None: + if self.n > 1: + raise ValueError("n must be 1 when using greedy sampling, " + f"got {self.n}.") + + def update_from_generation_config( + self, + generation_config: Dict[str, Any], + model_eos_token_id: Optional[int] = None) -> None: + """Update if there are non-default values from generation_config""" + + if model_eos_token_id is not None: + # Add the eos token id into the sampling_params to support + # min_tokens processing. + self._all_stop_token_ids.add(model_eos_token_id) + + # Update eos_token_id for generation + if (eos_ids := generation_config.get("eos_token_id")) is not None: + # it can be either int or list of int + eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids) + if model_eos_token_id is not None: + # We don't need to include the primary eos_token_id in + # stop_token_ids since it's handled separately for stopping + # purposes. + eos_ids.discard(model_eos_token_id) + if eos_ids: + self._all_stop_token_ids.update(eos_ids) + if not self.ignore_eos: + eos_ids.update(self.stop_token_ids) + self.stop_token_ids = list(eos_ids) + + @cached_property + def sampling_type(self) -> SamplingType: + if self.temperature < _SAMPLING_EPS: + return SamplingType.GREEDY + if self.seed is not None: + return SamplingType.RANDOM_SEED + return SamplingType.RANDOM + + @property + def all_stop_token_ids(self) -> Set[int]: + return self._all_stop_token_ids + + def clone(self) -> "SamplingParams": + """Deep copy excluding LogitsProcessor objects. + + LogitsProcessor objects are excluded because they may contain an + arbitrary, nontrivial amount of data. + See https://github.com/vllm-project/vllm/issues/3087 + """ + + logit_processor_refs = None if self.logits_processors is None else { + id(lp): lp + for lp in self.logits_processors + } + return copy.deepcopy(self, memo=logit_processor_refs) + + def __repr__(self) -> str: + return ( + f"SamplingParams(n={self.n}, " + f"presence_penalty={self.presence_penalty}, " + f"frequency_penalty={self.frequency_penalty}, " + f"repetition_penalty={self.repetition_penalty}, " + f"temperature={self.temperature}, " + f"top_p={self.top_p}, " + f"top_k={self.top_k}, " + f"min_p={self.min_p}, " + f"seed={self.seed}, " + f"stop={self.stop}, " + f"stop_token_ids={self.stop_token_ids}, " + f"include_stop_str_in_output={self.include_stop_str_in_output}, " + f"ignore_eos={self.ignore_eos}, " + f"max_tokens={self.max_tokens}, " + f"min_tokens={self.min_tokens}, " f"logprobs={self.logprobs}, " f"prompt_logprobs={self.prompt_logprobs}, " "prompt_logprob_positions=" f"{self.prompt_logprob_positions}, " - f"skip_special_tokens={self.skip_special_tokens}, " - "spaces_between_special_tokens=" - f"{self.spaces_between_special_tokens}, " - f"truncate_prompt_tokens={self.truncate_prompt_tokens}), " - f"guided_decoding={self.guided_decoding}") - - -class BeamSearchParams( - msgspec.Struct, - omit_defaults=True, # type: ignore[call-arg] - # required for @cached_property. - dict=True): # type: ignore[call-arg] - """Beam search parameters for text generation.""" - beam_width: int - max_tokens: int - ignore_eos: bool = False - temperature: float = 0.0 - length_penalty: float = 1.0 + f"skip_special_tokens={self.skip_special_tokens}, " + "spaces_between_special_tokens=" + f"{self.spaces_between_special_tokens}, " + f"truncate_prompt_tokens={self.truncate_prompt_tokens}), " + f"guided_decoding={self.guided_decoding}") + + +class BeamSearchParams( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + # required for @cached_property. + dict=True): # type: ignore[call-arg] + """Beam search parameters for text generation.""" + beam_width: int + max_tokens: int + ignore_eos: bool = False + temperature: float = 0.0 + length_penalty: float = 1.0