ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: emit usage and timings for MLX generation speed stats (#6068)

Browse source 
* Studio: emit usage and timings for MLX generation speed stats

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: make MLX generation stats request scoped

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Michael Han 2026-06-07 01:57:52 -07:00 • committed by GitHub
commit 1b588cd141
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 142 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

85
studio/backend/core/inference/mlx_inference.py
View file

@ -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

19
studio/backend/core/inference/orchestrator.py
View file

@ -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":

2
studio/backend/core/inference/worker.py
View file

@ -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(),
},
)

51
studio/backend/routes/inference.py
View file

@ -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: