diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 716e4c27a2..b9adba13e6 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -12,6 +12,35 @@ from loggers import get_logger logger = get_logger(__name__) +def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): + """Map mlx_lm / mlx_vlm stream stats onto the usage/timings shape + llama-server emits so the chat speed popover renders the same.""" + prompt_n = int(prompt_n or 0) + gen_n = int(gen_n or 0) + prompt_tps = float(prompt_tps or 0.0) + gen_tps = float(gen_tps or 0.0) + prompt_ms = (prompt_n / prompt_tps * 1000.0) if prompt_tps > 0 else 0.0 + predicted_ms = (gen_n / gen_tps * 1000.0) if gen_tps > 0 else 0.0 + return { + "usage": { + "prompt_tokens": prompt_n, + "completion_tokens": gen_n, + "total_tokens": prompt_n + gen_n, + }, + "timings": { + "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": prompt_tps, + "predicted_n": gen_n, + "predicted_ms": predicted_ms, + "predicted_per_token_ms": (predicted_ms / gen_n) if gen_n > 0 else 0.0, + "predicted_per_second": gen_tps, + "cache_n": 0, + }, + } + + class MLXInferenceBackend: def __init__(self): self.models = {} @@ -20,6 +49,8 @@ class MLXInferenceBackend: self.loaded_local_models = [] self.device = "mlx" self._generation_lock = threading.Lock() + # usage/timings of the latest generation; shipped on gen_done. + self.last_generation_stats = None # MLX state self._model = None @@ -258,6 +289,9 @@ class MLXInferenceBackend: if self._model is None: raise RuntimeError("No model loaded") + # Reset so a failed run cannot surface stale stats. + self.last_generation_stats = None + # Build messages with system prompt full_messages = [] if system_prompt: @@ -380,6 +414,7 @@ class MLXInferenceBackend: type(self._tokenizer).__name__, ) with self._generation_lock: + final_response = None try: gen_kwargs = dict( prompt = prompt, @@ -393,6 +428,7 @@ class MLXInferenceBackend: self._tokenizer, **gen_kwargs, ): + final_response = response token_ids.append(response.token) # Decode full sequence with skip_special_tokens — same as GPU cumulative = self._tokenizer.decode( @@ -408,6 +444,15 @@ class MLXInferenceBackend: logger.error("stream_generate failed:\n%s", traceback.format_exc()) raise + finally: + # Latch final cumulative stats for the usage/timings chunk. + if final_response is not None: + self.last_generation_stats = _build_generation_stats( + getattr(final_response, "prompt_tokens", 0), + getattr(final_response, "prompt_tps", 0.0), + getattr(final_response, "generation_tokens", 0), + getattr(final_response, "generation_tps", 0.0), + ) def _generate_vlm( self, @@ -483,20 +528,32 @@ class MLXInferenceBackend: vlm_kwargs["repetition_penalty"] = float(repetition_penalty) with self._generation_lock: - for response in vlm_stream( - self._model, - self._processor, - prompt, - images, - **vlm_kwargs, - ): - token_text = ( - response.text if hasattr(response, "text") else str(response) - ) - cumulative += token_text - yield cumulative - if cancel_event and cancel_event.is_set(): - break + final_response = None + try: + for response in vlm_stream( + self._model, + self._processor, + prompt, + images, + **vlm_kwargs, + ): + final_response = response + token_text = ( + response.text if hasattr(response, "text") else str(response) + ) + cumulative += token_text + yield cumulative + if cancel_event and cancel_event.is_set(): + break + finally: + # mlx_vlm exposes the same stats fields as mlx_lm. + if final_response is not None: + self.last_generation_stats = _build_generation_stats( + getattr(final_response, "prompt_tokens", 0), + getattr(final_response, "prompt_tps", 0.0), + getattr(final_response, "generation_tokens", 0), + getattr(final_response, "generation_tps", 0.0), + ) def generate_with_adapter_control( self, use_adapter = None, cancel_event = None, **gen_kwargs diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 7e7d7026f6..6321ca892d 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -453,6 +453,7 @@ class InferenceOrchestrator: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + stats_holder: Optional[dict] = None, ) -> Generator[str, None, None]: """Dispatched generation — sends command without holding _gen_lock. @@ -544,6 +545,8 @@ class InferenceOrchestrator: yield resp.get("text", "") elif rtype == "gen_done": + if stats_holder is not None: + stats_holder["stats"] = resp.get("stats") return elif rtype == "gen_error": @@ -793,6 +796,7 @@ class InferenceOrchestrator: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + stats_holder: Optional[dict] = None, ) -> Generator[str, None, None]: """Generate response, streaming tokens from subprocess. @@ -800,6 +804,10 @@ class InferenceOrchestrator: ``preserve_thinking`` kwargs are forwarded into the worker so ``tokenizer.apply_chat_template`` can render tool schemas and reasoning controls when the template understands them. + + ``stats_holder``: caller-owned dict; on gen_done its "stats" key + receives the worker's usage/timings. Request-scoped by design so + concurrent streams cannot read each other's stats. """ yield from self._generate_inner( messages = messages, @@ -817,6 +825,7 @@ class InferenceOrchestrator: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + stats_holder = stats_holder, ) def generate_chat_completion_with_tools( @@ -839,6 +848,7 @@ class InferenceOrchestrator: tool_call_timeout: int = 300, session_id: Optional[str] = None, use_adapter: Optional[Union[bool, str]] = None, + stats_holder: Optional[dict] = None, **_unused, ): """Run the safetensors agentic tool loop in this (parent) @@ -872,6 +882,8 @@ class InferenceOrchestrator: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + # last turn wins, same as the GGUF tool loop's metadata + stats_holder = stats_holder, ) if use_adapter is not None: yield from self.generate_with_adapter_control( @@ -901,6 +913,7 @@ class InferenceOrchestrator: self, use_adapter: Optional[Union[bool, str]] = None, cancel_event = None, + stats_holder: Optional[dict] = None, **gen_kwargs, ) -> Generator[str, None, None]: """Generate with adapter control, streaming tokens from subprocess. @@ -912,6 +925,7 @@ class InferenceOrchestrator: yield from self._generate_dispatched( use_adapter = use_adapter, cancel_event = cancel_event, + stats_holder = stats_holder, **gen_kwargs, ) @@ -932,6 +946,7 @@ class InferenceOrchestrator: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + stats_holder: Optional[dict] = None, ) -> Generator[str, None, None]: """Inner generation logic — sends command to subprocess, yields tokens. @@ -972,6 +987,7 @@ class InferenceOrchestrator: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + stats_holder = stats_holder, ) def _generate_locked( @@ -991,6 +1007,7 @@ class InferenceOrchestrator: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + stats_holder: Optional[dict] = None, ) -> Generator[str, None, None]: """Actual generation logic — must be called under _gen_lock.""" request_id = str(uuid.uuid4()) @@ -1069,6 +1086,8 @@ class InferenceOrchestrator: yield resp.get("text", "") elif rtype == "gen_done": + if stats_holder is not None: + stats_holder["stats"] = resp.get("stats") return elif rtype == "gen_error": diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 20a7d2d16c..d01cce4fc1 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -481,6 +481,8 @@ def _handle_generate( { "type": "gen_done", "request_id": request_id, + # usage/timings from the MLX backend (None elsewhere). + "stats": getattr(backend, "last_generation_stats", None), "ts": time.time(), }, ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1d505e31f6..d3a279319a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3426,6 +3426,9 @@ async def openai_chat_completions( else: _sf_chat_messages.append(_msg) + # Request-scoped usage/timings receptacle (filled at gen_done). + _sf_stats_holder: dict = {} + def sf_generate_with_tools(): return backend.generate_chat_completion_with_tools( messages = _sf_chat_messages, @@ -3450,6 +3453,7 @@ async def openai_chat_completions( else 300, session_id = payload.session_id, use_adapter = payload.use_adapter, + stats_holder = _sf_stats_holder, ) _sf_tool_sentinel = object() @@ -3537,6 +3541,25 @@ async def openai_chat_completions( ], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" + # Usage chunk from the last turn, same shape as the + # GGUF tool loop's metadata. Request-scoped holder, so + # concurrent streams cannot read each other's stats. + _stats = _sf_stats_holder.get("stats") + if _stats: + _stream_usage = _stats.get("usage") or {} + usage_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [], + usage = CompletionUsage( + prompt_tokens = _stream_usage.get("prompt_tokens", 0), + completion_tokens = _stream_usage.get("completion_tokens", 0), + total_tokens = _stream_usage.get("total_tokens", 0), + ), + timings = _stats.get("timings"), + ) + yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -3627,19 +3650,25 @@ async def openai_chat_completions( if payload.preserve_thinking is not None: gen_kwargs["preserve_thinking"] = payload.preserve_thinking + # Request-scoped usage/timings receptacle (filled at gen_done). + stats_holder: dict = {} + if payload.use_adapter is not None: def generate(): return backend.generate_with_adapter_control( use_adapter = payload.use_adapter, cancel_event = cancel_event, + stats_holder = stats_holder, **gen_kwargs, ) else: def generate(): return backend.generate_chat_response( - cancel_event = cancel_event, **gen_kwargs + cancel_event = cancel_event, + stats_holder = stats_holder, + **gen_kwargs, ) # ── Streaming response ──────────────────────────────────────── @@ -3716,6 +3745,26 @@ async def openai_chat_completions( ], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" + # Usage chunk (choices=[], usage set), same shape as the + # GGUF path so the speed popover works for MLX too. + # Request-scoped holder, so concurrent streams cannot + # read each other's stats. + _stats = stats_holder.get("stats") + if _stats: + _stream_usage = _stats.get("usage") or {} + usage_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [], + usage = CompletionUsage( + prompt_tokens = _stream_usage.get("prompt_tokens", 0), + completion_tokens = _stream_usage.get("completion_tokens", 0), + total_tokens = _stream_usage.get("total_tokens", 0), + ), + timings = _stats.get("timings"), + ) + yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" yield "data: [DONE]\n\n" except asyncio.CancelledError: