From 9d118ec5fcb643858b0b908154647d05d55d743d Mon Sep 17 00:00:00 2001 From: CearX <56916338+CearX@users.noreply.github.com> Date: Mon, 13 Jul 2026 01:54:36 +0800 Subject: [PATCH] perf(server): optimize concurrent long-output serving --- python/infinilm/base_config.py | 7 + python/infinilm/config/engine_config.py | 1 + python/infinilm/llm/llm.py | 66 ++++++- .../infinilm/llm/model_runner/model_runner.py | 1 + python/infinilm/llm/request.py | 43 +++++ python/infinilm/llm/scheduler.py | 67 ++++++- .../processors/basic_llm_processor.py | 6 +- python/infinilm/server/inference_server.py | 81 +++++--- scripts/competition/official_client_bench.py | 182 ++++++++++++++++++ 9 files changed, 409 insertions(+), 45 deletions(-) create mode 100644 scripts/competition/official_client_bench.py diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 3b8f6533..fc48210d 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -112,6 +112,7 @@ def __init__(self): self.port = self.args.port self.endpoint = self.args.endpoint self.ignore_eos = self.args.ignore_eos + self.stream_interval = self.args.stream_interval # PD separation (KV transfer) self.kv_transfer_config = self.args.kv_transfer_config @@ -269,6 +270,12 @@ def _add_common_args(self): default=8, help="maximum batch size for server", ) + self.parser.add_argument( + "--stream-interval", + type=int, + default=1, + help="emit first token immediately, then coalesce decode steps per SSE chunk", + ) self.parser.add_argument( "--input-len", type=parse_list, default=10, help="input sequence length" ) diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index 1bb1f733..851f68d1 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -47,6 +47,7 @@ class EngineConfig: max_tokens: int = 4096 num_blocks: int = 512 block_size: int = 256 + kv_cache_dtype: Optional[str] = None max_cache_len: int = 4096 temperature: float = 1.0 top_p: float = 0.8 diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59e9e94b..6b48ff3c 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -24,6 +24,7 @@ FinishReason, InferenceRequest, RequestOutput, + RequestStatus, TokenOutput, ) from infinilm.llm.sampling_params import SamplingParams @@ -177,6 +178,7 @@ def _update_requests( case _: raise ValueError(f"Unsupported cache_type: {self.cache_type}") pending = [] + step_complete_requests = [] for req, token_ids in zip(requests, sampled_tokens): if req.is_aborted(): logger.info( @@ -186,8 +188,24 @@ def _update_requests( # (status still RUNNING). if not req.is_finished(): req.mark_canceled() + step_complete_requests.append(req) continue + if is_prefill and not getattr(req, "_prefill_chunk_final", True): + # Intermediate chunk: KV has been written, but the sampled token + # is not a model output because the prompt is not complete yet. + req.num_computed_tokens = req._prefill_chunk_end + req.num_local_cached_tokens = req._prefill_chunk_end + req.status = RequestStatus.WAITING + self.scheduler.waiting_queue.sync_q.put(req) + continue + + step_complete_requests.append(req) + + if is_prefill: + req.num_computed_tokens = req.get_prompt_length() + req.num_local_cached_tokens = req.get_prompt_length() + if not isinstance(token_ids, list): token_ids = [token_ids] @@ -195,6 +213,27 @@ def _update_requests( if req.is_finished(): break req.generated_token_ids.append(token_id) + + # Non-streaming benchmark/API path: avoid per-token queue traffic and + # incremental detokenization. With ignore_eos=true and no stop strings, + # the only normal finish condition is max_tokens, so final text can be + # decoded once when the request completes. + fast_non_stream = ( + req._output_queue is None + and req.sampling_params.ignore_eos + and not req.sampling_params.stop + ) + if fast_non_stream: + is_finished = self._check_request_finished(req, token_id) + if is_finished: + req.generated_text = self.tokenizer.decode( + req.generated_token_ids + ) + req._token_decode_offset = len(req.generated_token_ids) + req._text_output_offset = len(req.generated_text) + req.mark_finished(req.finish_reason) + continue + pending_tokens = req.generated_token_ids[req._token_decode_offset :] delta = self.tokenizer.decode(pending_tokens) holds_back = bool(delta) and delta.endswith("\ufffd") @@ -214,12 +253,22 @@ def _update_requests( req.mark_finished(req.finish_reason) else: + stream_interval = max( + 1, int((req.request_data or {}).get("stream_interval", 1)) + ) + should_emit = ( + is_finished + or req.get_num_generated_tokens() == 1 + or req.get_num_generated_tokens() % stream_interval == 0 + ) + if not should_emit: + continue + if holds_back and not is_finished: token_text = "" else: if is_finished and req.finish_reason in ( FinishReason.EOS_TOKEN, - FinishReason.LENGTH, FinishReason.STOP_STRING, ): token_text = "" @@ -246,7 +295,7 @@ def _update_requests( continue pending.append((req.output_queue.async_q, output)) - self.scheduler.complete_requests(requests) + self.scheduler.complete_requests(step_complete_requests) return pending def _check_request_finished(self, req: InferenceRequest, token_id: int) -> bool: @@ -323,6 +372,7 @@ def __init__( max_tokens: int = 4096, num_blocks: int = 512, block_size: int = 256, + kv_cache_dtype: Optional[str] = None, max_cache_len: int = 4096, temperature: float = 1.0, top_p: float = 0.8, @@ -369,6 +419,7 @@ def __init__( max_tokens=max_tokens, num_blocks=num_blocks, block_size=block_size, + kv_cache_dtype=kv_cache_dtype, max_cache_len=max_cache_len, temperature=temperature, top_p=top_p, @@ -530,6 +581,7 @@ def __init__( max_tokens: int = 512, num_blocks: int = 512, block_size: int = 256, + kv_cache_dtype: Optional[str] = None, max_cache_len: int = 4096, temperature: float = 1.0, top_p: float = 0.8, @@ -579,6 +631,7 @@ def __init__( max_tokens=max_tokens, num_blocks=num_blocks, block_size=block_size, + kv_cache_dtype=kv_cache_dtype, max_cache_len=max_cache_len, temperature=temperature, top_p=top_p, @@ -721,6 +774,7 @@ def add_request( request_id: Optional[str] = None, # For server use request_data: Optional[dict] = None, + init_output_queue: bool = True, ) -> InferenceRequest: """Add a request to the engine. @@ -816,8 +870,10 @@ def add_request( kv_params = request_data["kv_transfer_params"] request.kv_transfer_params = kv_params - # Initialize output queue for streaming - _ = request.output_queue + # Initialize output queue only when the caller needs token streaming. + # Non-streaming HTTP responses can wait for request completion and decode once. + if init_output_queue: + _ = request.output_queue self.engine.add_request(request) return request @@ -829,6 +885,7 @@ def add_chat_request( request_id: Optional[str] = None, request_data: Optional[dict] = None, add_generation_prompt: bool = True, + init_output_queue: bool = True, **kwargs, ) -> InferenceRequest: """Add a chat request to the engine. @@ -850,6 +907,7 @@ def add_chat_request( sampling_params=sampling_params, request_id=request_id, request_data=request_data, + init_output_queue=init_output_queue, ) async def stream_request( diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index 45f1eaee..84dfc48e 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -77,6 +77,7 @@ def __init__(self, config: EngineConfig): cache_config=cache_config, enable_graph_compiling=config.enable_graph, attention_backend=config.attn_backend, + kv_cache_dtype=config.kv_cache_dtype, use_mla=config.use_mla, weight_load_mode=config.weight_load_mode, skip_legacy_moe=config.skip_legacy_moe, diff --git a/python/infinilm/llm/request.py b/python/infinilm/llm/request.py index abccfe53..369b1ac5 100644 --- a/python/infinilm/llm/request.py +++ b/python/infinilm/llm/request.py @@ -181,6 +181,8 @@ def __init__( self._aborted: bool = False self._text_output_offset: int = 0 self._token_decode_offset: int = 0 + self._completion_loop: Optional[asyncio.AbstractEventLoop] = None + self._completion_event: Optional[asyncio.Event] = None @property def output_queue(self) -> janus.Queue: @@ -222,6 +224,43 @@ def is_finished(self) -> bool: RequestStatus.TIMEOUT, ] + def bind_completion_event(self, loop: asyncio.AbstractEventLoop): + """Bind an asyncio event used by non-streaming HTTP waiters.""" + self._completion_loop = loop + self._completion_event = asyncio.Event() + if self.is_finished(): + self._completion_event.set() + + async def wait_finished(self, timeout: Optional[float] = None): + """Wait until the request reaches a terminal state.""" + if self.is_finished(): + return True + if self._completion_event is None: + start = time.time() + while not self.is_finished(): + if timeout is not None and time.time() - start > timeout: + return False + await asyncio.sleep(0.001) + return True + try: + if timeout is None: + await self._completion_event.wait() + else: + await asyncio.wait_for(self._completion_event.wait(), timeout=timeout) + return True + except asyncio.TimeoutError: + return False + + def _notify_finished(self): + event = self._completion_event + loop = self._completion_loop + if event is None: + return + if loop is not None and loop.is_running(): + loop.call_soon_threadsafe(event.set) + else: + event.set() + def abort(self): """Signal that the request has been aborted and should stop generation.""" self._aborted = True @@ -235,6 +274,7 @@ def mark_finished(self, reason: FinishReason): self.status = RequestStatus.FINISHED self.finish_reason = reason self.finished_time = time.time() + self._notify_finished() def mark_failed(self, reason: FinishReason = FinishReason.ERROR): """Mark the request as failed.""" @@ -242,6 +282,7 @@ def mark_failed(self, reason: FinishReason = FinishReason.ERROR): self.status = RequestStatus.FAILED self.finish_reason = reason self.finished_time = time.time() + self._notify_finished() def mark_canceled(self): """Mark the request as canceled.""" @@ -249,6 +290,7 @@ def mark_canceled(self): self.status = RequestStatus.CANCELED self.finish_reason = FinishReason.CANCELED self.finished_time = time.time() + self._notify_finished() def mark_timeout(self): """Mark the request as timed out.""" @@ -256,6 +298,7 @@ def mark_timeout(self): self.status = RequestStatus.TIMEOUT self.finish_reason = FinishReason.TIMEOUT self.finished_time = time.time() + self._notify_finished() async def close(self): """Close the output queue and clean up resources.""" diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index a1153e91..20b2c3cf 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -3,6 +3,7 @@ """ import logging +import os import queue from typing import List, Optional @@ -99,6 +100,13 @@ def __init__( self.speculative_cache_ops = SpeculativeCacheOps(self.cache_manager) self.block_size = block_size self.max_num_batched_tokens = max_num_batched_tokens + self.enable_chunked_prefill = os.getenv( + "INFINILM_ENABLE_CHUNKED_PREFILL", "0" + ).lower() in {"1", "true", "yes", "on"} + self.kv_safety_blocks = max(0, int(os.getenv("INFINILM_KV_SAFETY_BLOCKS", "0"))) + self.long_prompt_kv_safety_blocks = max( + 0, int(os.getenv("INFINILM_LONG_PROMPT_KV_SAFETY_BLOCKS", "0")) + ) self.connector = connector def add_request(self, request: InferenceRequest): @@ -179,14 +187,19 @@ def schedule(self) -> Optional[SchedulerOutput]: ) if load_kv_async: num_computed_tokens -= 1 - num_new_tokens = req.get_prompt_length() - num_computed_tokens + remaining_prompt_tokens = req.get_prompt_length() - num_computed_tokens + if self.enable_chunked_prefill: + available_tokens = max( + 1, self.max_num_batched_tokens - current_num_batched_tokens + ) + num_new_tokens = min(remaining_prompt_tokens, available_tokens) + else: + num_new_tokens = remaining_prompt_tokens + num_tokens_this_step = num_new_tokens # Early token budget check: skip can_accept_request and allocate_slots # for requests that would exceed the per-schedule token budget. if not load_kv_async: - num_tokens_this_step = ( - req.get_prompt_length() - num_local_computed_tokens - ) if self._exceeds_token_budget( current_num_batched_tokens, num_tokens_this_step, @@ -201,6 +214,7 @@ def schedule(self) -> Optional[SchedulerOutput]: req, num_local_computed_tokens, current_prefill_extra_blocks, + len(scheduled_requests), ): logger.warning( "Insufficient KV cache blocks for request %s, deferring.", @@ -257,8 +271,16 @@ def schedule(self) -> Optional[SchedulerOutput]: ) else: load_kv_async = False + remaining_prompt_tokens = ( + req.get_prompt_length() - req.num_computed_tokens + ) + available_tokens = max( + 1, self.max_num_batched_tokens - current_num_batched_tokens + ) num_tokens_this_step = ( - req.get_prompt_length() - req.num_local_cached_tokens + min(remaining_prompt_tokens, available_tokens) + if self.enable_chunked_prefill + else remaining_prompt_tokens ) if self._exceeds_token_budget( current_num_batched_tokens, @@ -267,9 +289,23 @@ def schedule(self) -> Optional[SchedulerOutput]: ): deferred_requests.append(req) break - self.cache_manager.update_blocks_hash( - req.block_table, req.num_local_cached_tokens + req_blocks, slot_mapping = self.cache_manager.allocate_slots( + req.get_input_tokens(), + num_tokens_this_step, + num_computed_tokens=req.num_computed_tokens, + cached_block_table=req.block_table, ) + if req_blocks is None: + deferred_requests.append(req) + break + req.block_table = req_blocks + req.slot_mapping = slot_mapping + req.num_blocks = len(req_blocks) + req.num_local_cached_tokens = req.num_computed_tokens + + prefill_end = req.num_computed_tokens + num_tokens_this_step + req._prefill_chunk_end = prefill_end + req._prefill_chunk_final = prefill_end >= req.get_prompt_length() if load_kv_async: req.status = RequestStatus.WAITING_FOR_REMOTE_KVS @@ -282,7 +318,6 @@ def schedule(self) -> Optional[SchedulerOutput]: current_prefill_extra_blocks += self._get_prefill_extra_blocks(req) scheduled_requests.append(req) - num_tokens_this_step = req.get_prompt_length() - req.num_local_cached_tokens current_num_batched_tokens += num_tokens_this_step req.status = RequestStatus.RUNNING @@ -464,6 +499,7 @@ def can_accept_request( request: InferenceRequest, num_local_computed_tokens: int, current_prefill_extra_blocks: int = 0, + current_prefill_requests: int = 0, ) -> bool: if ( self.mamba_cache_manager is not None @@ -500,8 +536,21 @@ def can_accept_request( # Include decode headroom for requests accepted earlier in this batch. total_required_blocks += current_prefill_extra_blocks + # Longer prompts leave more partially filled tail blocks across + # staggered prefill waves. Reserve one fragmentation block per active + # sequence so decode never reaches an unrecoverable zero-block state. + fragmentation_headroom = 0 + if request.get_prompt_length() > self.block_size: + fragmentation_headroom = max( + self.long_prompt_kv_safety_blocks, + running_queue_size + current_prefill_requests + 1, + ) + # Compare with total usable blocks in cache manager - return total_required_blocks <= self.cache_manager.get_total_usable_blocks() + return ( + total_required_blocks + self.kv_safety_blocks + fragmentation_headroom + <= self.cache_manager.get_total_usable_blocks() + ) def _get_prefill_extra_blocks(self, request: InferenceRequest) -> int: total_length = request.get_prompt_length() diff --git a/python/infinilm/processors/basic_llm_processor.py b/python/infinilm/processors/basic_llm_processor.py index a6fbc33a..b735cbfc 100644 --- a/python/infinilm/processors/basic_llm_processor.py +++ b/python/infinilm/processors/basic_llm_processor.py @@ -207,11 +207,11 @@ def _build_model_input_from_batch_scheduler_output( if scheduler_output.is_prefill: # Prefill phase req_tokens = req.get_input_tokens() - tokens_to_compute = req_tokens[num_cached:] + compute_len = len(req.slot_mapping) + tokens_to_compute = req_tokens[num_cached : num_cached + compute_len] tokens.extend(tokens_to_compute) - compute_len = len(tokens_to_compute) - seq_len = len(req_tokens) + seq_len = num_cached + compute_len seq_lens.append(seq_len) current_offset += compute_len diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 242593e2..4e55cd74 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -22,8 +22,8 @@ logger = logging.getLogger(__name__) -DEFAULT_STREAM_TIMEOUT = 100.0 -DEFAULT_REQUEST_TIMEOUT = 1000.0 +DEFAULT_STREAM_TIMEOUT = 3600.0 +DEFAULT_REQUEST_TIMEOUT = 3600.0 def chunk_json( @@ -105,6 +105,7 @@ def __init__( max_batch_size: int = 16, num_blocks: int = 512, block_size: int = 256, + kv_cache_dtype: Optional[str] = None, max_cache_len: int = 4096, temperature: float = 1.0, top_p: float = 0.8, @@ -116,6 +117,7 @@ def __init__( use_mla: bool = False, weight_load_mode: str = "async", ignore_eos: bool = False, + stream_interval: int = 1, kv_transfer_config: Optional[KVTransferConfig] = None, ): """Initialize inference server. @@ -158,6 +160,7 @@ def __init__( self.max_batch_size = max_batch_size self.num_blocks = num_blocks self.block_size = block_size + self.kv_cache_dtype = kv_cache_dtype self.max_cache_len = max_cache_len self.temperature = temperature self.top_p = top_p @@ -169,6 +172,7 @@ def __init__( self.use_mla = use_mla self.weight_load_mode = weight_load_mode self.ignore_eos = ignore_eos + self.stream_interval = max(1, stream_interval) self.kv_transfer_config = kv_transfer_config self.engine: AsyncLLMEngine = None @@ -197,6 +201,7 @@ async def lifespan(app: FastAPI): max_tokens=self.max_tokens, num_blocks=self.num_blocks, block_size=self.block_size, + kv_cache_dtype=self.kv_cache_dtype, max_cache_len=self.max_cache_len, temperature=self.temperature, top_p=self.top_p, @@ -341,10 +346,21 @@ def pick(key: str, default): return default # Accept common alias - max_tokens = pick("max_tokens", self.max_tokens) - if max_tokens is None: - # Some clients use max_new_tokens - max_tokens = pick("max_new_tokens", self.max_tokens) + # OpenAI clients increasingly use max_completion_tokens. Keep legacy + # aliases for compatibility, with explicit top-level values winning. + if data.get("max_tokens") is not None: + max_tokens = data["max_tokens"] + elif data.get("max_completion_tokens") is not None: + max_tokens = data["max_completion_tokens"] + elif data.get("max_new_tokens") is not None: + max_tokens = data["max_new_tokens"] + else: + max_tokens = sp.get( + "max_tokens", + sp.get( + "max_completion_tokens", sp.get("max_new_tokens", self.max_tokens) + ), + ) stop = pick("stop", None) if isinstance(stop, str): @@ -363,6 +379,7 @@ async def _stream_chat(self, request_id: str, data: dict, http_request: Request) """Handle streaming chat request.""" req = None _abort_reason = FinishReason.CANCELED + data.setdefault("stream_interval", self.stream_interval) try: messages = data.get("messages", []) @@ -435,6 +452,24 @@ async def _stream_chat(self, request_id: str, data: dict, http_request: Request) ensure_ascii=False, ) yield f"data: {chunk}\n\n" + stream_options = data.get("stream_options") or {} + if stream_options.get("include_usage"): + usage_chunk = json.dumps( + { + "id": request_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": self.model_id, + "choices": [], + "usage": { + "prompt_tokens": req.get_prompt_length(), + "completion_tokens": req.get_num_generated_tokens(), + "total_tokens": req.get_total_length(), + }, + }, + ensure_ascii=False, + ) + yield f"data: {usage_chunk}\n\n" break except asyncio.CancelledError: @@ -480,37 +515,23 @@ async def _chat(self, request_id: str, data: dict, http_request: Request): request_data=data, add_generation_prompt=bool(data.get("add_generation_prompt", True)), chat_template_kwargs=data.get("chat_template_kwargs") or {}, + init_output_queue=False, ) + req.bind_completion_event(asyncio.get_running_loop()) - # Collect all generated tokens - output_text = "" - async for token_output in self.engine.stream_request( - req, - timeout=DEFAULT_STREAM_TIMEOUT, - request_timeout=DEFAULT_REQUEST_TIMEOUT, - ): - # Check client disconnect + start_time = time.time() + while not req.is_finished(): if await http_request.is_disconnected(): logger.info(f"Client disconnected for request {request_id}") break - - # Request-level timeout is handled inside stream_request. - if token_output.finish_reason == FinishReason.TIMEOUT: + if time.time() - start_time > DEFAULT_REQUEST_TIMEOUT: logger.warning(f"Request {request_id} timed out") + self.engine.add_aborted_req(req, FinishReason.TIMEOUT) break + remaining = DEFAULT_REQUEST_TIMEOUT - (time.time() - start_time) + await req.wait_finished(timeout=min(0.1, max(0.0, remaining))) - # Skip EOS token text for OpenAI API compatibility - # Check if this token is an EOS token by comparing token_id with eos_token_ids - eos_token_ids = self.engine.engine.eos_token_ids - is_eos_token = eos_token_ids and token_output.token_id in eos_token_ids - - if not is_eos_token and token_output.token_text: - output_text += token_output.token_text - - if token_output.finished: - break - - output_text = output_text.strip() + output_text = req.generated_text.strip() finish_reason = self._convert_finish_reason(req.finish_reason) response = completion_json( @@ -612,6 +633,7 @@ def main(): max_batch_size=cfg.max_batch_size, num_blocks=cfg.num_blocks, block_size=cfg.block_size, + kv_cache_dtype=cfg.kv_cache_dtype, max_cache_len=cfg.max_cache_len, temperature=cfg.temperature, top_p=cfg.top_p, @@ -623,6 +645,7 @@ def main(): use_mla=cfg.use_mla, weight_load_mode=cfg.weight_load_mode, ignore_eos=cfg.ignore_eos, + stream_interval=cfg.stream_interval, kv_transfer_config=kv_transfer_config, ) server.start() diff --git a/scripts/competition/official_client_bench.py b/scripts/competition/official_client_bench.py new file mode 100644 index 00000000..e98eacde --- /dev/null +++ b/scripts/competition/official_client_bench.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +"""Dependency-light equivalent of `vllm bench serve --backend openai-chat`. + +It preserves vLLM's random dataset construction, burst submission, streaming +OpenAI request shape, and tokenizer-based output accounting without importing +vLLM's CUDA engine extensions. +""" + +import argparse +import asyncio +import json +import time +from dataclasses import asdict, dataclass, field + +import aiohttp +import numpy as np +from transformers import AutoTokenizer + + +@dataclass +class Result: + success: bool = False + text: str = "" + latency: float = 0.0 + ttft: float = 0.0 + itl: list[float] = field(default_factory=list) + error: str = "" + prompt_len: int = 0 + output_tokens: int = 0 + + +def make_prompts(tokenizer, count: int, input_len: int, seed: int): + rng = np.random.default_rng(seed) + special = int(tokenizer.num_special_tokens_to_add()) + inner_len = max(0, input_len - special) + offsets = rng.integers(0, tokenizer.vocab_size, size=count) + prompts = [] + for i, offset in enumerate(offsets): + ids = ((int(offset) + i + np.arange(inner_len)) % tokenizer.vocab_size).tolist() + prompt = tokenizer.decode(ids) + ids = tokenizer.encode(prompt, add_special_tokens=False)[:inner_len] + prompts.append((tokenizer.decode(ids), len(ids) + special)) + return prompts + + +async def one(session, sem, url, model, prompt, prompt_len, output_len): + result = Result(prompt_len=prompt_len) + payload = { + "model": model, + "messages": [{"role": "user", "content": [{"type": "text", "text": prompt}]}], + "temperature": 0.0, + "max_completion_tokens": output_len, + "max_tokens": output_len, + "stream": True, + "stream_options": {"include_usage": True}, + } + async with sem: + start = time.perf_counter() + recent = start + buffer = "" + try: + async with session.post(url, json=payload) as response: + if response.status != 200: + result.error = f"HTTP {response.status}: {await response.text()}" + return result + async for raw in response.content.iter_any(): + buffer += raw.decode("utf-8") + while "\n\n" in buffer: + message, buffer = buffer.split("\n\n", 1) + message = message.strip() + if not message or message.startswith(":"): + continue + chunk = message.removeprefix("data: ") + if chunk == "[DONE]": + continue + data = json.loads(chunk) + now = time.perf_counter() + if choices := data.get("choices"): + content = choices[0].get("delta", {}).get("content") or "" + if result.ttft == 0.0: + result.ttft = now - start + else: + result.itl.append(now - recent) + result.text += content + recent = now + elif usage := data.get("usage"): + result.output_tokens = usage.get("completion_tokens") or 0 + result.latency = recent - start + result.success = result.ttft > 0.0 + except Exception as exc: + result.error = repr(exc) + return result + + +async def run(args): + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + prompts = make_prompts( + tokenizer, args.num_prompts, args.random_input_len, args.seed + ) + connector = aiohttp.TCPConnector( + limit=args.max_concurrency, limit_per_host=args.max_concurrency + ) + timeout = aiohttp.ClientTimeout(total=6 * 60 * 60) + sem = asyncio.Semaphore(args.max_concurrency) + async with aiohttp.ClientSession( + connector=connector, timeout=timeout, trust_env=False + ) as session: + # Match vLLM's readiness request, excluded from measured duration. + warm = await one( + session, sem, args.url, args.model, *prompts[0], args.random_output_len + ) + if not warm.success: + raise RuntimeError(f"warmup failed: {warm.error}") + start = time.perf_counter() + tasks = [ + one(session, sem, args.url, args.model, p, n, args.random_output_len) + for p, n in prompts + ] + results = await asyncio.gather(*tasks) + duration = time.perf_counter() - start + + actual = [] + for r in results: + if r.success: + actual.append( + r.output_tokens + or len(tokenizer(r.text, add_special_tokens=False).input_ids) + ) + else: + actual.append(0) + good = [r for r in results if r.success] + total_input = sum(r.prompt_len for r in good) + total_output = sum(actual) + ttfts = [r.ttft * 1000 for r in good] + tpots = [ + (r.latency - r.ttft) / (n - 1) * 1000 + for r, n in zip(results, actual) + if r.success and n > 1 + ] + report = { + "config": vars(args), + "duration": duration, + "completed": len(good), + "total_input_tokens": total_input, + "total_output_tokens": total_output, + "request_throughput": len(good) / duration, + "output_throughput": total_output / duration, + "total_token_throughput": (total_input + total_output) / duration, + "mean_ttft_ms": float(np.mean(ttfts or [0])), + "p99_ttft_ms": float(np.percentile(ttfts or [0], 99)), + "mean_tpot_ms": float(np.mean(tpots or [0])), + "actual_output_lens": actual, + "actual_prompt_lens": [r.prompt_len for r in results], + "errors": [r.error for r in results if not r.success], + } + print(json.dumps(report, ensure_ascii=False, indent=2)) + if args.save_result: + with open(args.save_result, "w", encoding="utf-8") as f: + json.dump( + {**report, "results": [asdict(r) for r in results]}, + f, + ensure_ascii=False, + indent=2, + ) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--url", default="http://127.0.0.1:8102/v1/chat/completions") + p.add_argument("--model", default="Qwen3-8B") + p.add_argument("--tokenizer", required=True) + p.add_argument("--num-prompts", type=int, required=True) + p.add_argument("--max-concurrency", type=int, required=True) + p.add_argument("--random-input-len", type=int, required=True) + p.add_argument("--random-output-len", type=int, required=True) + p.add_argument("--seed", type=int, default=0) + p.add_argument("--save-result") + asyncio.run(run(p.parse_args())) + + +if __name__ == "__main__": + main()