diff --git a/docs/cli.md b/docs/cli.md index 350c03d45..cb618b443 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -116,8 +116,23 @@ Speculative decoding needs `--tp 1` and runs without overlap scheduling; targets | `--tool-call-parser` | auto | Tool-call format; auto-inferred from the model family | | `--reasoning-parser` | auto | Splits chain-of-thought into `reasoning_content`; auto-inferred; `off` disables | | `--enable-cache-report` | off | Report prefix-cache hits in each response's usage block | +| `--enable-timings-report` | off | Attach a llama.cpp-shaped per-request `timings` object to responses | | `--anthropic-inline-system` | auto | Placement of late Anthropic system instructions: `auto`, `preserve`, or `fold` | +With `--enable-timings-report` every chat completion, completion and Responses response +also carries a llama.cpp-shaped `timings` object next to `usage` (`cache_n`, `prompt_n`, +`prompt_ms`, `prompt_per_token_ms`, `prompt_per_second`, `predicted_n`, `predicted_ms`, +`predicted_per_token_ms`, `predicted_per_second`, plus `queue_ms` as a FreeToken +extension and `draft_n` / `draft_n_accepted` under speculative decoding), so frontends +like llama-swap and Open Web UI show per-request prefill/decode rates. All spans are +measured in the scheduler: `queue_ms` covers admission to the first prefill chunk, +`prompt_ms` the first chunk to the first sampled token (chunk gaps included), +`predicted_ms` the first to the last committed token; `prompt_n` counts only the +evaluated tokens and the decode rates divide by `predicted_n - 1`, as in llama.cpp. On +streams the object rides the final chunk; a multi-prompt `/v1/completions` request +reports the first prompt's spans, as llama.cpp does. `/v1/messages` never carries it +(llama.cpp emits no timings there either). + ### Image input Experimental. Needs a checkpoint whose family registers a vision encoder ([models.md](models.md#image-input) lists them and how each one diff --git a/python/freetoken/core.py b/python/freetoken/core.py index a7e41dd9a..0a928cba1 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -32,6 +32,48 @@ def is_greedy(self) -> bool: return self.temperature <= 0.0 or self.top_k == 1 +@dataclass +class ReqStats: + """Per-request timing and draft counters, stamped by the scheduler with + time.perf_counter(). Spans follow llama.cpp's server_slot_stats plus a queue + span: queue = waiting queue -> first prefill chunk launch, prompt = first + chunk launch -> first sampled token on the host (chunk gaps included), + predicted = first token -> last committed token. One instance is shared by + every chunk Req of a request, so the stamps survive chunked prefill.""" + + enqueued_at: float | None = None # entered the scheduler's waiting queue + first_chunk_at: float | None = None # first prefill chunk launched + first_token_at: float | None = None # first sampled token back on the host + last_token_at: float | None = None # last committed token + cache_n: int = 0 # prompt tokens served from the prefix cache at admission + prompt_n: int = 0 # prompt tokens actually computed + draft_n: int = 0 # drafted tokens summed over verify steps + draft_n_accepted: int = 0 # accepted drafts that were committed, after stop truncation + + def stamp_token(self, now: float) -> None: + if self.first_token_at is None: + self.first_token_at = now + self.last_token_at = now + + def finish(self, completion_tokens: int) -> dict | None: + """Raw spans and counts for the terminal reply; None when no prefill ever + launched or no token was committed (errors and aborts emit nothing).""" + if self.first_chunk_at is None or self.first_token_at is None: + return None + return { + "queue_ms": (self.first_chunk_at - self.enqueued_at) * 1000.0 + if self.enqueued_at is not None else 0.0, + "cache_n": self.cache_n, + "prompt_n": self.prompt_n, + "prompt_ms": (self.first_token_at - self.first_chunk_at) * 1000.0, + "predicted_n": completion_tokens, + # 1 us floor as in llama.cpp: the first token can land in the same instant + "predicted_ms": max((self.last_token_at - self.first_token_at) * 1000.0, 0.001), + "draft_n": self.draft_n, + "draft_n_accepted": self.draft_n_accepted, + } + + @dataclass(eq=False) class Req: input_ids: torch.Tensor # cpu tensor @@ -68,6 +110,9 @@ class Req: # handler must not free resources under an in-flight forward; it sets this flag and # _process_last_data frees the request when the batch drains (after copy_done.synchronize). aborted: bool = False + # Timing/draft counters shared with every other chunk Req of this request (set by + # the PrefillManager from the PendingReq's instance; a default blank otherwise). + stats: ReqStats = field(default_factory=ReqStats, init=False, repr=False) output_token_counts: torch.Tensor | None = field(default=None, init=False, repr=False) def __post_init__(self) -> None: diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py index abf692643..6a1d0e87d 100644 --- a/python/freetoken/message/frontend.py +++ b/python/freetoken/message/frontend.py @@ -45,6 +45,10 @@ class UserReply(BaseFrontendMsg): swa_total_tokens: int = 0 # Bytes the engine process holds on the GPU (torch reserved pool). 0 when not reported. gpu_mem_bytes: int = 0 + # Scheduler-measured per-request timings (queue/prompt/predicted spans, cache/prompt/ + # predicted counts, draft counters), passthrough of DetokenizeMsg.timings: present on + # the terminal reply of a served request, None otherwise. + timings: dict | None = None # Set (with finished=True) when a request failed before producing output — e.g. a chat # template that the tokenizer cannot render, or a prompt that exceeds the KV budget the # scheduler can serve. Carries a human-readable reason. Without this, such a request would diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index 56308d6ec..e2b57be0f 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -49,6 +49,9 @@ class DetokenizeMsg(BaseTokenizerMsg): swa_total_tokens: int = 0 # Bytes this engine process holds on the GPU (torch reserved pool). 0 on CPU. gpu_mem_bytes: int = 0 + # Scheduler-measured per-request timings (ReqStats.finish) on the terminal reply + # only; None on intermediate replies, errors and aborts. + timings: dict | None = None @dataclass diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index 94e708fb1..89f6b6471 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -1,11 +1,12 @@ from __future__ import annotations import math +import time from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Tuple import torch -from freetoken.core import Batch, Req +from freetoken.core import Batch, Req, ReqStats from freetoken.utils import align_down, div_ceil, init_logger from .mm import mm_chunk_end, mm_rows_after @@ -222,6 +223,9 @@ def _add_one_req( req.mm_items = pending_req.mm_items req.mrope_positions_full = pending_req.mrope_positions_full req.mrope_delta = pending_req.mrope_delta + # one ReqStats across all chunk Reqs of the request: the stamps the scheduler + # takes on the final chunk's Req must survive the intermediate ones + req.stats = pending_req.stats # Hybrid GDN per-request state slots (None for non-hybrid). On a fresh admit these are # freshly allocated; on a chunked continuation they are inherited from the prior chunk. req.linear_slot_idx = linear_slot_idx @@ -283,6 +287,7 @@ class PrefillManager: pending_list: List[PendingReq] = field(default_factory=list) def add_one_req(self, req: UserMsg) -> None: + stats = ReqStats(enqueued_at=time.perf_counter()) self.pending_list.append( PendingReq( req.uid, @@ -291,6 +296,7 @@ def add_one_req(self, req: UserMsg) -> None: mm_items=req.mm_items, mrope_positions_full=req.mrope_positions, mrope_delta=req.mrope_delta, + stats=stats, ) ) @@ -331,6 +337,11 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None: prompt_admissions.append( (req.uid, pending_req.input_len, req.cache_handle.cached_len) ) + # The first prefill chunk is launching: queue time ends here and the + # prompt span starts (ReqStats carries both to the final chunk's Req). + pending_req.stats.first_chunk_at = time.perf_counter() + pending_req.stats.cache_n = req.cache_handle.cached_len + pending_req.stats.prompt_n = pending_req.input_len - req.cache_handle.cached_len if pending_req.mm_items and self.encoder_cache is not None: # claim the rows every chunk of this request will gather; the entry outlives the chunks for item in pending_req.mm_items: diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 023051ff2..4bdf2631d 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -336,6 +336,8 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: batch, output = last_data[0].batch, last_data[1] _, next_tokens_cpu, copy_done = output[:3] copy_done.synchronize() + # tokens are on the host from here: the per-request timings' token stamps + now = time.perf_counter() reply: List[DetokenizeMsg] = [] new_finished_reqs: Set[Req] = set() generated = None @@ -370,7 +372,10 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: # client's terminal reply. continue msg, _ = self._commit_tokens(req, next_tokens_cpu[i : i + 1]) + req.stats.stamp_token(now) reply.append(msg) + if msg.finished: + msg.timings = req.stats.finish(req.input_ids.numel() - req.prompt_len) finished = msg.finished # NOTE: overlap scheduling may make the request freed twice, skip second free @@ -424,12 +429,24 @@ def _drain_verify(self, batch: Batch, output, reply: List[DetokenizeMsg], new_fi """Commit each request's accepted tokens, return the slots its rejected rows took, then the target state; returns the tokens committed.""" k = self.config.decode_rows_per_req - 1 counts, finished = [], [] + now = time.perf_counter() for i, req in enumerate(batch.reqs): # spec forces normal_loop, which frees an aborted request between steps; overlap would need rollback_tail here and must not donate the GDN slot the verify rolled assert not req.aborted, "a verify batch drained an aborted request" cached = req.cached_len - msg, n = self._commit_tokens(req, output.next_tokens_cpu[i, : int(output.emitted_cpu[i]) + 1]) + emitted = int(output.emitted_cpu[i]) + msg, n = self._commit_tokens(req, output.next_tokens_cpu[i, : emitted + 1]) + # k drafts were proposed; the step hands _commit_tokens the emitted accepted + # drafts plus the target's own token, so when a stop lands on a draft every + # committed token is an accepted draft: count min(n, emitted), with n already + # post stop-truncation. Never exceeds predicted_n - 1: the request's first + # committed token came from prefill, outside every verify step. + req.stats.stamp_token(now) + req.stats.draft_n += k + req.stats.draft_n_accepted += min(n, emitted) reply.append(msg) + if msg.finished: + msg.timings = req.stats.finish(req.input_ids.numel() - req.prompt_len) req.commit_pending(n) self.cache_manager.rollback_tail(req, req.cached_len, cached + k + 1) counts.append(n) diff --git a/python/freetoken/scheduler/utils.py b/python/freetoken/scheduler/utils.py index 83b79f2a3..2d660db58 100644 --- a/python/freetoken/scheduler/utils.py +++ b/python/freetoken/scheduler/utils.py @@ -1,10 +1,12 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, List import torch +from freetoken.core import ReqStats + if TYPE_CHECKING: from freetoken.core import SamplingParams @@ -20,6 +22,9 @@ class PendingReq: mm_items: list | None = None mrope_positions_full: torch.Tensor | None = None mrope_delta: int = 0 + # timing/draft counters; add_one_req stamps enqueued_at on it, and every chunk + # Req built from this pending request shares the same instance + stats: ReqStats = field(default_factory=ReqStats) @property def input_len(self) -> int: diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index a3293d14a..3d75143e2 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -61,6 +61,10 @@ class ServerArgs(SchedulerConfig): # prompt_tokens_details.cached_tokens, Anthropic cache_read_input_tokens, Responses # input_tokens_details.cached_tokens). Mirrors sglang's --enable-cache-report. enable_cache_report: bool = False + # Attach llama.cpp's per-request `timings` object to OpenAI chat/completions and + # Responses responses (plus queue_ms and, under speculative decoding, draft_n / + # draft_n_accepted). Opt-in like vLLM's --enable-per-request-metrics. + enable_timings_report: bool = False anthropic_inline_system: str = "auto" # Comma-separated hostname allowlist for client-supplied image URLs; empty admits any domain. allowed_media_domains: str = "" @@ -584,6 +588,18 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--enable-timings-report", + action="store_true", + default=ServerArgs.enable_timings_report, + help=( + "Attach a llama.cpp-shaped per-request timings object to every chat completion, " + "completion and Responses response (and to the final stream chunk), so frontends " + "like llama-swap and Open Web UI can show prefill/decode rates and cache hits. " + "Spans are measured in the scheduler; queue_ms is a FreeToken extension." + ), + ) + parser.add_argument( "--anthropic-inline-system", choices=("auto", "preserve", "fold"), diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index ef90402aa..6a7d9122b 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -115,6 +115,7 @@ class GenDone: completion_tokens: int matched_stop: str | None = None cached_tokens: int = 0 + timings: dict[str, Any] | None = None GenEvent = ReasoningDelta | ContentDelta | ToolCallStart | ToolCallArgsDelta | ToolCallsDelta | GenDone @@ -130,6 +131,46 @@ class GenResult: completion_tokens: int matched_stop: str | None = None cached_tokens: int = 0 + timings: dict[str, Any] | None = None + + +def timings_block(state: Any, timings: dict[str, Any] | None) -> dict[str, Any] | None: + """llama.cpp's per-request timings object (server_slot_stats::to_json, b10408+) as + frontends like llama-swap read it, plus queue_ms as a FreeToken extension. The spans + and counts are measured in the scheduler (ReqStats) and arrive on the terminal ack: + queue_ms covers waiting-queue admission to the first prefill chunk, prompt_ms the + first chunk to the first sampled token on the host, predicted_ms the first to the + last committed token. prompt_n counts only the evaluated tokens and the decode rates + divide by predicted_n - 1 steps (the first token rides the last prompt batch), + matching llama.cpp's arithmetic. Values are not rounded. Off unless + --enable-timings-report: usage.prompt_tokens - prompt_n would otherwise still + reveal the prefix-cache hit that --enable-cache-report hides.""" + if timings is None: + return None + config = getattr(state, "config", None) if state is not None else None + if not getattr(config, "enable_timings_report", False): + return None + prompt_n = timings["prompt_n"] + predicted_n = timings["predicted_n"] + prompt_ms = timings["prompt_ms"] + predicted_ms = timings["predicted_ms"] + gen_steps = max(predicted_n - 1, 0) + block: dict[str, Any] = { + "cache_n": timings["cache_n"], + "prompt_n": prompt_n, + "prompt_ms": prompt_ms, + "prompt_per_token_ms": prompt_ms / prompt_n if prompt_n > 0 else 0.0, + "prompt_per_second": 1e3 / prompt_ms * prompt_n if prompt_ms > 0 else 0.0, + "predicted_n": predicted_n, + "predicted_ms": predicted_ms, + "predicted_per_token_ms": predicted_ms / gen_steps if gen_steps > 0 else 0.0, + "predicted_per_second": 1e3 / predicted_ms * gen_steps if predicted_ms > 0 else 0.0, + "queue_ms": timings["queue_ms"], + } + if timings.get("draft_n", 0) > 0: + block["draft_n"] = timings["draft_n"] + block["draft_n_accepted"] = timings["draft_n_accepted"] + return block @dataclass @@ -699,12 +740,15 @@ def _route_tool_text(piece: str) -> list[GenEvent]: engine_finish_reason: str | None = None engine_matched_stop: str | None = None + timings: dict[str, Any] | None = None async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): raise GenerationError(ack.error, getattr(ack, "error_code", None)) prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens + if getattr(ack, "timings", None) is not None: + timings = ack.timings content_delta = ack.incremental_output if reasoning_parser is not None and content_delta: reasoning_delta, content_delta = reasoning_parser.parse_stream_chunk(content_delta) @@ -787,6 +831,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]: yield GenDone( finish_reason, prompt_tokens, completion_tokens, matched_stop=engine_matched_stop, cached_tokens=cached_tokens, + timings=timings, ) @@ -797,6 +842,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + timings: dict[str, Any] | None = None engine_finish_reason: str | None = None engine_matched_stop: str | None = None async for ack in state.wait_for_ack(uid): @@ -805,6 +851,8 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens + if getattr(ack, "timings", None) is not None: + timings = ack.timings full_content += ack.incremental_output if ack.finished: engine_finish_reason = getattr(ack, "finish_reason", None) @@ -831,4 +879,5 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: completion_tokens=completion_tokens, matched_stop=engine_matched_stop, cached_tokens=cached_tokens, + timings=timings, ) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 15992930e..4442a8490 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -37,6 +37,7 @@ prerender_error, render_messages, resolve_sampling, + timings_block, submit_generation, ) @@ -219,7 +220,7 @@ async def handle_chat_completion( if result.tool_calls: message["tool_calls"] = _tool_calls_to_openai(result.tool_calls) - return { + response: dict[str, Any] = { "id": f"chatcmpl-{uid}", "object": "chat.completion", "created": int(time.time()), @@ -237,6 +238,9 @@ async def handle_chat_completion( _reported_cached(state, result.cached_tokens), ), } + if (timings := timings_block(state, result.timings)) is not None: + response["timings"] = timings + return response async def stream_chat_completion_chunks( @@ -259,6 +263,10 @@ async def stream_chat_completion_chunks( prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + timings: dict[str, Any] | None = None + # timings ride the stream's LAST chunk only: the usage chunk when the client asked + # for one, otherwise the finish chunk + include_usage = bool(req.stream_options and req.stream_options.include_usage) tool_calls_sent = 0 open_tool: dict[str, Any] | None = None events = generate_events(uid, spec, state, source="/v1/chat/completions") @@ -368,21 +376,26 @@ async def stream_chat_completion_chunks( prompt_tokens = ev.prompt_tokens completion_tokens = ev.completion_tokens cached_tokens = ev.cached_tokens - yield _sse(_chat_chunk(req, uid, [{"delta": {}, "index": 0, "finish_reason": ev.finish_reason}])) - - if req.stream_options and req.stream_options.include_usage: - yield _sse( - { - "id": f"chatcmpl-{uid}", - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [], - "usage": _usage( - prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens) - ), - } - ) + timings = ev.timings + done = _chat_chunk(req, uid, [{"delta": {}, "index": 0, "finish_reason": ev.finish_reason}]) + if not include_usage and (t := timings_block(state, timings)) is not None: + done["timings"] = t + yield _sse(done) + + if include_usage: + final: dict[str, Any] = { + "id": f"chatcmpl-{uid}", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": req.model, + "choices": [], + "usage": _usage( + prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens) + ), + } + if (t := timings_block(state, timings)) is not None: + final["timings"] = t + yield _sse(final) yield b"data: [DONE]\n\n" @@ -426,6 +439,10 @@ async def handle_completion( prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + # with several prompts the timings are the FIRST prompt's own spans, as llama.cpp + # returns arr[0]; one clock across the sequential loop would mix the later prompts' + # queue and prefill into the rates + first_timings: dict[str, Any] | None = None for index, prompt in enumerate(prompts): uid = state.new_user() await state.send_one( @@ -447,11 +464,13 @@ async def handle_completion( cached_tokens += ack.cached_tokens text += ack.incremental_output if ack.finished: + if index == 0: + first_timings = getattr(ack, "timings", None) finish_reason = getattr(ack, "finish_reason", None) or "stop" break choices.append({"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None}) - return { + response: dict[str, Any] = { "id": f"cmpl-{uuid.uuid4().hex}", "object": "text_completion", "created": int(time.time()), @@ -459,12 +478,17 @@ async def handle_completion( "choices": choices, "usage": _usage(prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)), } + if (timings := timings_block(state, first_timings)) is not None: + response["timings"] = timings + return response async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) -> AsyncIterator[bytes]: prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + timings: dict[str, Any] | None = None + include_usage = bool(req.stream_options and req.stream_options.include_usage) finish_reason = "stop" async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): @@ -474,6 +498,8 @@ async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens + if getattr(ack, "timings", None) is not None: + timings = ack.timings if ack.incremental_output: yield _sse( { @@ -495,28 +521,32 @@ async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) finish_reason = getattr(ack, "finish_reason", None) or "stop" break - yield _sse( - { + done_chunk: dict[str, Any] = { + "id": f"cmpl-{uid}", + "object": "text_completion.chunk", + "created": int(time.time()), + "model": req.model, + "choices": [{"text": "", "index": 0, "finish_reason": finish_reason, "logprobs": None}], + } + if not include_usage and (t := timings_block(state, timings)) is not None: + # as on the chat stream: the finish chunk is the last one only when no usage + # chunk follows + done_chunk["timings"] = t + yield _sse(done_chunk) + if include_usage: + final: dict[str, Any] = { "id": f"cmpl-{uid}", "object": "text_completion.chunk", "created": int(time.time()), "model": req.model, - "choices": [{"text": "", "index": 0, "finish_reason": finish_reason, "logprobs": None}], + "choices": [], + "usage": _usage( + prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens) + ), } - ) - if req.stream_options and req.stream_options.include_usage: - yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [], - "usage": _usage( - prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens) - ), - } - ) + if (t := timings_block(state, timings)) is not None: + final["timings"] = t + yield _sse(final) yield b"data: [DONE]\n\n" diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py index 71d6eed87..621e01c25 100644 --- a/python/freetoken/server/responses_api.py +++ b/python/freetoken/server/responses_api.py @@ -75,6 +75,7 @@ resolve_sampling, split_tool_lists, submit_generation, + timings_block, with_keepalive, ) from .request_logger import log_request @@ -166,7 +167,7 @@ async def handle_responses( if req.stream: events = responses_stream_generator( generate_events(uid, spec, state, source="/v1/responses"), req, response_id, created, - cache_report=cache_report, + cache_report=cache_report, state=state, ) if request is not None: events = state.stream_with_cancellation(events, request, uid) @@ -176,7 +177,9 @@ async def handle_responses( result = await generate_full(uid, spec, state, source="/v1/responses") except GenerationError as exc: return _error_response(400, str(exc), exc.code) - response = build_responses_response(result, req, response_id, created, cache_report=cache_report) + response = build_responses_response( + result, req, response_id, created, cache_report=cache_report, state=state + ) return JSONResponse(content=response.model_dump(mode="json")) @@ -410,6 +413,7 @@ def build_responses_response( response_id: str, created: int, cache_report: bool = False, + state: Any = None, ) -> Response: truncated = result.finish_reason == "length" item_status = "incomplete" if truncated else "completed" @@ -454,14 +458,17 @@ def build_responses_response( result.cached_tokens if cache_report else 0, ), incomplete_reason="max_output_tokens" if truncated else None, + timings=timings_block(state, result.timings), ) def _response_obj( response_id: str, created: int, model: str, output: list[Any], *, status: str, usage: ResponseUsage | None, error: ResponseError | None = None, - incomplete_reason: str | None = None, + incomplete_reason: str | None = None, timings: dict[str, Any] | None = None, ) -> Response: + # Response allows extra fields, so timings rides along without forking the schema. + extra: dict[str, Any] = {"timings": timings} if timings is not None else {} return Response( id=response_id, created_at=created, @@ -475,6 +482,7 @@ def _response_obj( parallel_tool_calls=True, tool_choice="auto", tools=[], + **extra, ) @@ -499,6 +507,7 @@ async def responses_stream_generator( response_id: str, created: int, cache_report: bool = False, + state: Any = None, ) -> AsyncIterator[str]: seq = _Seq() @@ -707,6 +716,7 @@ def args_delta_frame(fragment: str) -> str: finish_reason = ev.finish_reason usage_pt, usage_ct = ev.prompt_tokens, ev.completion_tokens usage_cached = ev.cached_tokens if cache_report else 0 + final_timings = timings_block(state, ev.timings) for f in close_current(): yield f if finish_reason == "length": @@ -717,6 +727,7 @@ def args_delta_frame(fragment: str) -> str: usage=_usage(usage_pt, usage_ct, usage_cached), incomplete_reason="max_output_tokens", ), + **({"timings": final_timings} if final_timings is not None else {}), )) else: yield _sse(ResponseCompletedEvent( @@ -725,6 +736,7 @@ def args_delta_frame(fragment: str) -> str: "completed", output_items, usage=_usage(usage_pt, usage_ct, usage_cached), ), + **({"timings": final_timings} if final_timings is not None else {}), )) except GenerationError as exc: # Request failed before/instead of producing output (template render error, over-length diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index 988871605..816e723b6 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -244,6 +244,7 @@ def tokenize_worker( swa_used_tokens=msg.swa_used_tokens, swa_total_tokens=msg.swa_total_tokens, gpu_mem_bytes=msg.gpu_mem_bytes, + timings=msg.timings, ) for msg, reply in zip(detokenize_msg, replies, strict=True) ] diff --git a/tests/scheduler/test_scheduler_chunked_prefill.py b/tests/scheduler/test_scheduler_chunked_prefill.py index a4c7e0659..2fff1758e 100644 --- a/tests/scheduler/test_scheduler_chunked_prefill.py +++ b/tests/scheduler/test_scheduler_chunked_prefill.py @@ -152,6 +152,46 @@ def test_batched_prefill_carries_each_new_prompt_admission(): assert batch.prompt_admissions == [(1, 3, 0), (2, 5, 0)] +def test_chunked_reqs_share_one_stats_stamped_at_the_first_chunk(): + """The queue span ends and the prompt span starts when the FIRST chunk launches, and + every chunk Req shares one ReqStats, so the token stamps the final chunk's drain + takes land on the instance the admission stamps populated. A per-chunk restamp would + hide the queue wait and the earlier chunks' prefill time.""" + from types import SimpleNamespace + + from freetoken.core import SamplingParams + + cm, _tm, _dm, pm = _build_managers(num_pages=64) + prompt_len = 3 * CHUNK + pm.add_one_req(SimpleNamespace( + uid=UID, input_ids=torch.arange(prompt_len, dtype=torch.int32), + sampling_params=SamplingParams(max_tokens=4), mm_items=None, + mrope_positions=None, mrope_delta=0, + )) + pending = pm.pending_list[0] + assert pending.stats.enqueued_at is not None + + seen_first_chunks = [] + while pm.runnable: + batch = pm.schedule_next_batch(CHUNK) + assert batch is not None + cm.allocate_paged(batch.reqs) + for req in batch.reqs: + req.complete_one() + seen_first_chunks.append(req.stats.first_chunk_at) + assert req.stats is pending.stats + if pm.pending_list: # the final chunk leaves no continuation behind + assert pm.pending_list[0].stats is pending.stats + + assert pending.stats.first_chunk_at is not None + assert pending.stats.first_chunk_at >= pending.stats.enqueued_at + assert pending.stats.cache_n == 0 + assert pending.stats.prompt_n == prompt_len + # stamped once, on the first chunk: every later chunk reads the same value + assert len(seen_first_chunks) == 3 + assert all(f == pending.stats.first_chunk_at for f in seen_first_chunks) + + def test_chunk_ends_before_an_image_it_would_cut_on_bidirectional_models(): from freetoken.core import SamplingParams from freetoken.message import MMItem diff --git a/tests/scheduler/test_spec_drain.py b/tests/scheduler/test_spec_drain.py index d2c9dda4d..89c7730df 100644 --- a/tests/scheduler/test_spec_drain.py +++ b/tests/scheduler/test_spec_drain.py @@ -108,3 +108,42 @@ def test_drain_commits_the_accepted_tokens_and_returns_the_rejected_slots(): for req in (a, c): stub._free_req_resources(req) cm.check_integrity() + + +def test_drain_counts_drafts_and_attaches_timings(): + """The verify drain feeds ReqStats: k drafts per step, accepted = min(n, emitted) + (a stop landing on a draft makes every committed token an accepted draft), and the + terminal reply carries the finished request's spans.""" + pool, cm, tm, dm, stub, events = _setup() + reqs = [_decoding_req(pool, cm, tm, 1, 16), _decoding_req(pool, cm, tm, 2, 3)] + dm.filter_reqs(reqs) + for req in reqs: + req.widen_pending(K) + cm.allocate_paged(reqs) + # spans the prefill drain would have stamped: enqueued -> first chunk -> first token + for req in reqs: + req.stats.enqueued_at = 1.0 + req.stats.first_chunk_at = 2.0 + req.stats.first_token_at = 3.0 + req.stats.prompt_n = PROMPT + out = torch.tensor([[11, 12, 13, 14], [21, 22, 23, 24]]) + output = SimpleNamespace(next_tokens_gpu=out, next_tokens_cpu=out, emitted_cpu=torch.tensor([3, 3])) + batch = Batch(reqs=reqs, phase="verify", rows_per_req=K + 1) + reply, finished = [], set() + + Scheduler._drain_verify(stub, batch, output, reply, finished) + + # req 1: every draft accepted, n = emitted + 1 -> 3 accepted drafts + assert reqs[0].stats.draft_n == K and reqs[0].stats.draft_n_accepted == 3 + # req 2: all 3 drafts accepted but the length budget cuts the step to n = 2, and + # both committed tokens ARE accepted drafts (the target's own never landed): 2 + assert reqs[1].stats.draft_n == K and reqs[1].stats.draft_n_accepted == 2 + # only the terminal reply carries timings + assert reply[0].timings is None + t = reply[1].timings + assert t["queue_ms"] == 1000.0 # enqueued -> first chunk + assert t["prompt_ms"] == 1000.0 # first chunk -> first token + assert t["prompt_n"] == PROMPT + assert t["predicted_n"] == 3 # first token + the two committed here + assert t["predicted_ms"] > 0.0 # first -> last committed token + assert t["draft_n"] == K and t["draft_n_accepted"] == 2 diff --git a/tests/server/test_generation_accounting.py b/tests/server/test_generation_accounting.py index a6d2bdb91..76971fab2 100644 --- a/tests/server/test_generation_accounting.py +++ b/tests/server/test_generation_accounting.py @@ -24,6 +24,7 @@ GenSpec, generate_events, generate_full, + timings_block, ) @@ -63,7 +64,10 @@ async def wait_for_ack(self, uid: int): yield reply -def _ack(prompt: int = 0, completion: int = 0, out: str = "", finished: bool = False) -> UserReply: +def _ack( + prompt: int = 0, completion: int = 0, out: str = "", finished: bool = False, + timings: dict | None = None, +) -> UserReply: return UserReply( uid=42, incremental_output=out, @@ -71,9 +75,22 @@ def _ack(prompt: int = 0, completion: int = 0, out: str = "", finished: bool = F prompt_tokens_delta=prompt, completion_tokens_delta=completion, finish_reason="stop" if finished else None, + timings=timings, ) +def _sched_timings(**over) -> dict: + """A scheduler ReqStats.finish payload: raw spans and counts, no derived rates.""" + t = dict(queue_ms=0.0, cache_n=0, prompt_n=0, prompt_ms=0.0, predicted_n=0, + predicted_ms=0.0, draft_n=0, draft_n_accepted=0) + t.update(over) + return t + + +def _timings_state() -> SimpleNamespace: + return SimpleNamespace(config=SimpleNamespace(enable_timings_report=True)) + + def _spec() -> GenSpec: return GenSpec(messages=[{"role": "user", "content": "hi"}], sampling_params=SimpleNamespace()) @@ -206,3 +223,69 @@ def test_ttft_mean_is_zero_without_samples(): request_ring.reset() request_ring.record_request(_row(ttft_ms=None)) assert request_ring.requests_ttft_mean_ms() == 0 + + +def test_timings_block_follows_llama_cpp_semantics(): + t = timings_block(_timings_state(), _sched_timings( + queue_ms=500.0, cache_n=800, prompt_n=200, prompt_ms=2500.0, + predicted_n=21, predicted_ms=2000.0, + )) + # prompt_n excludes the cache; decode steps are n - 1 (llama.cpp semantics). + assert t["cache_n"] == 800 + assert t["prompt_n"] == 200 + assert t["prompt_per_second"] == 80.0 # 200 evaluated tokens / 2.5 s + assert t["predicted_n"] == 21 + assert t["predicted_per_second"] == 10.0 # 20 decode steps / 2.0 s + assert t["prompt_ms"] == 2500.0 and t["predicted_ms"] == 2000.0 + assert t["queue_ms"] == 500.0 # FreeToken extension + assert "draft_n" not in t and "draft_n_accepted" not in t # no drafting + + +def test_timings_block_is_opt_in_and_needs_measured_spans(): + raw = _sched_timings(prompt_n=10, prompt_ms=100.0, predicted_n=2, predicted_ms=50.0) + # off without the flag: usage.prompt_tokens - prompt_n would still leak the cache hit + assert timings_block(SimpleNamespace(config=SimpleNamespace()), raw) is None + # the responses helpers may pass no state at all + assert timings_block(None, raw) is None + # errors and aborts deliver no scheduler timings + assert timings_block(_timings_state(), None) is None + + +def test_timings_block_rates_survive_the_edge_shapes(): + state = _timings_state() + # one output token: the first rides the last prompt batch, so no decode steps + t = timings_block(state, _sched_timings(prompt_n=10, prompt_ms=100.0, predicted_n=1, predicted_ms=0.001)) + assert t["predicted_per_second"] == 0.0 and t["predicted_per_token_ms"] == 0.0 + # fully cached prompt: nothing was evaluated + t = timings_block(state, _sched_timings(prompt_n=0, prompt_ms=100.0, predicted_n=2, predicted_ms=50.0)) + assert t["prompt_per_second"] == 0.0 and t["prompt_per_token_ms"] == 0.0 + # values are passed through unrounded + t = timings_block(state, _sched_timings(prompt_n=3, prompt_ms=1234.5678, predicted_n=2, predicted_ms=50.0)) + assert t["prompt_ms"] == 1234.5678 and t["prompt_per_token_ms"] == 1234.5678 / 3 + # speculative decoding adds the draft counters, and only then + t = timings_block(state, _sched_timings(prompt_n=3, prompt_ms=300.0, predicted_n=5, + predicted_ms=400.0, draft_n=64, draft_n_accepted=40)) + assert t["draft_n"] == 64 and t["draft_n_accepted"] == 40 + + +def test_gen_done_and_result_carry_the_scheduler_timings(): + raw = _sched_timings(queue_ms=5.0, prompt_n=7, prompt_ms=70.0, predicted_n=5, predicted_ms=80.0) + st = FakeState([ + _ack(prompt=7, completion=1, out="x"), + _ack(completion=4, out="yz", finished=True, timings=raw), + ]) + + async def drain(): + done = None + async for ev in generate_events(42, _spec(), st): + if isinstance(ev, GenDone): + done = ev + return done + + done = asyncio.run(drain()) + assert done.timings is raw # the terminal ack's scheduler spans, unchanged + + result = asyncio.run(generate_full(42, _spec(), FakeState( + [_ack(prompt=5, completion=1, out="a"), _ack(completion=2, out="bc", finished=True, timings=raw)] + ))) + assert result.timings is raw diff --git a/tests/server/test_message_wire.py b/tests/server/test_message_wire.py index 76d8ecc39..cf1ea1d9f 100644 --- a/tests/server/test_message_wire.py +++ b/tests/server/test_message_wire.py @@ -118,6 +118,26 @@ def test_detokenize_msg_carries_kv_usage_round_trip(): assert (decoded.swa_used_tokens, decoded.swa_total_tokens) == (8448, 76800) +def test_terminal_replies_carry_scheduler_timings_over_the_wire(): + timings = {"queue_ms": 5.5, "cache_n": 8, "prompt_n": 2, "prompt_ms": 100.0, + "predicted_n": 4, "predicted_ms": 250.0, "draft_n": 6, "draft_n_accepted": 3} + msg = DetokenizeMsg(uid=3, next_tokens=[42], finished=True, timings=timings) + decoded = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg)) + assert isinstance(decoded, DetokenizeMsg) + assert decoded.timings == timings + + reply = UserReply(uid=3, incremental_output="", finished=True, + completion_tokens_delta=1, timings=timings) + out = BaseFrontendMsg.decoder(BaseFrontendMsg.encoder(reply)) + assert isinstance(out, UserReply) + assert out.timings == timings + # intermediate replies and aborts/errors never carry timings + plain = BaseFrontendMsg.decoder(BaseFrontendMsg.encoder( + UserReply(uid=3, incremental_output="x", finished=False) + )) + assert plain.timings is None + + def test_client_dicts_with_the_wire_tag_key_survive_intact(): """Tool JSON Schemas and chat_template_kwargs are free-form client data. A field literally named ``__type__`` (a common discriminator) must not be read back as a serialized class -- diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index 3dcef25f3..899f9e878 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -411,6 +411,80 @@ def test_completion_forwards_length_finish_reason(): assert response["choices"][0]["finish_reason"] == "length" +class _PerPromptState(FakeState): + """new_user hands out one uid per prompt; each uid drains its own canned replies.""" + + def __init__(self, replies_by_uid: dict[int, list[UserReply]]) -> None: + super().__init__([]) + self.replies_by_uid = replies_by_uid + self._next_uid = 42 + + def new_user(self) -> int: + uid = self._next_uid + self._next_uid += 1 + return uid + + async def wait_for_ack(self, uid: int): + for reply in self.replies_by_uid[uid]: + yield reply + + +def _sched_timings(**over) -> dict: + t = dict(queue_ms=0.0, cache_n=0, prompt_n=0, prompt_ms=0.0, predicted_n=0, + predicted_ms=0.0, draft_n=0, draft_n_accepted=0) + t.update(over) + return t + + +def test_multi_prompt_completion_reports_the_first_prompts_own_timings(): + """llama.cpp returns arr[0]'s slot stats for a multi-prompt completion; one clock + spanning the sequential loop would fold the later prompts' queue and prefill into + the rates while prompt_n still sums every prompt.""" + first = _sched_timings(queue_ms=10.0, prompt_n=6, prompt_ms=300.0, + predicted_n=3, predicted_ms=150.0) + second = _sched_timings(queue_ms=2500.0, prompt_n=9, prompt_ms=4000.0, + predicted_n=4, predicted_ms=800.0) + state = _PerPromptState({ + 42: [UserReply(uid=42, incremental_output="a", finished=True, prompt_tokens_delta=6, + completion_tokens_delta=3, timings=first)], + 43: [UserReply(uid=43, incremental_output="b", finished=True, prompt_tokens_delta=9, + completion_tokens_delta=4, timings=second)], + }) + state.config.enable_timings_report = True + + response = run( + handle_completion( + CompletionRequest(model="client-model", prompt=["first prompt", "second prompt"]), + request=None, + state=state, + model_sampling={}, + ) + ) + + # usage still sums both prompts; the timings are the FIRST prompt's spans alone + assert response["usage"]["prompt_tokens"] == 15 + assert response["timings"]["queue_ms"] == 10.0 + assert response["timings"]["prompt_n"] == 6 + assert response["timings"]["prompt_per_second"] == 20.0 # 6 tokens / 0.3 s + assert response["timings"]["predicted_n"] == 3 + + +def test_completion_and_chat_omit_timings_without_the_flag(): + state = FakeState([UserReply(uid=42, incremental_output="hi", finished=True, + prompt_tokens_delta=5, completion_tokens_delta=1, + timings=_sched_timings(prompt_n=5, prompt_ms=100.0, + predicted_n=1, predicted_ms=0.001))]) + response = run(handle_completion( + CompletionRequest(model="client-model", prompt="say hi"), + request=None, state=state, model_sampling={}, + )) + assert "timings" not in response + chat = run(handle_chat_completion( + chat_request(tools=None), request=None, state=state, model_sampling={}, + )) + assert "timings" not in chat + + def test_omitted_max_tokens_honors_server_default(): from freetoken.server.generation import DEFAULT_MAX_OUTPUT_TOKENS @@ -693,6 +767,59 @@ async def collect(): assert usage["prompt_tokens_details"] == {"cached_tokens": 3} +def _timings_stream_state() -> FakeState: + state = FakeState([UserReply( + uid=42, incremental_output="hi", finished=True, prompt_tokens_delta=5, + completion_tokens_delta=2, + timings=_sched_timings(prompt_n=5, prompt_ms=100.0, predicted_n=2, predicted_ms=50.0), + )]) + state.config.enable_timings_report = True + return state + + +def _timings_stream_events(req): + state = _timings_stream_state() + + async def collect(): + return [chunk async for chunk in stream_chat_completion_chunks(42, req, state)] + + return [e for e in parse_sse(run(collect())) if isinstance(e, dict)] + + +def test_stream_timings_ride_the_usage_chunk_when_requested(): + events = _timings_stream_events(chat_request(tools=None, stream=True, + stream_options={"include_usage": True})) + carriers = [e for e in events if "timings" in e] + # the usage chunk is the stream's last chunk, so the timings appear exactly once, there + assert len(carriers) == 1 + assert carriers[0].get("usage") is not None + assert carriers[0]["timings"]["prompt_per_second"] == 50.0 # 5 tokens / 0.1 s + + +def test_stream_timings_ride_the_finish_chunk_without_usage(): + events = _timings_stream_events(chat_request(tools=None, stream=True)) + carriers = [e for e in events if "timings" in e] + # no usage chunk follows, so the finish chunk is the last and the only carrier + assert len(carriers) == 1 + assert carriers[0]["choices"] and carriers[0]["choices"][0]["finish_reason"] is not None + assert all("usage" not in e for e in events) + + +def test_stream_timings_stay_off_without_the_flag(): + state = FakeState([UserReply( + uid=42, incremental_output="hi", finished=True, prompt_tokens_delta=5, + completion_tokens_delta=2, + timings=_sched_timings(prompt_n=5, prompt_ms=100.0, predicted_n=2, predicted_ms=50.0), + )]) + req = chat_request(tools=None, stream=True, stream_options={"include_usage": True}) + + async def collect(): + return [chunk async for chunk in stream_chat_completion_chunks(42, req, state)] + + events = parse_sse(run(collect())) + assert not [e for e in events if isinstance(e, dict) and "timings" in e] + + # --------------------------------------------------------------- minimax think def test_minimax_http_non_stream_forces_implicit_reasoning_without_request_knob(): state = FakeState( diff --git a/tests/server/test_responses_api.py b/tests/server/test_responses_api.py index 887e78e57..4c5a636c7 100644 --- a/tests/server/test_responses_api.py +++ b/tests/server/test_responses_api.py @@ -306,6 +306,33 @@ def test_stream_text_events(): assert completed["usage"]["input_tokens"] == 4 +def test_stream_timings_ride_the_completed_event_data_not_the_response(): + req = ResponsesRequest.model_validate({"model": "gpt-x", "input": "hi", "stream": True}) + timings = {"queue_ms": 5.0, "cache_n": 0, "prompt_n": 4, "prompt_ms": 200.0, + "predicted_n": 2, "predicted_ms": 100.0, "draft_n": 0, "draft_n_accepted": 0} + events = [ContentDelta("Hi"), GenDone("stop", 4, 2, timings=timings)] + state = SimpleNamespace(config=SimpleNamespace(enable_timings_report=True)) + + async def run(): + out = [] + async for frame in RP.responses_stream_generator( + _aiter(events), req, "resp_1", 0, state=state + ): + for line in frame.split("\n"): + if line.startswith("data:"): + out.append(json.loads(line[len("data:"):].strip())) + return out + + data_events = asyncio.run(run()) + completed = data_events[-1] + assert completed["type"] == "response.completed" + # the timings sit on the event's data, not nested in the response object + assert completed["timings"]["prompt_per_second"] == 20.0 # 4 tokens / 0.2 s + assert "timings" not in completed["response"] + # every earlier event stays clean + assert all("timings" not in e for e in data_events[:-1]) + + def test_stream_tool_call_events(): req = ResponsesRequest.model_validate({"model": "gpt-x", "input": "weather?", "stream": True}) args = '{"city": "SF"}'