| """OpenAI/vLLM-compatible serving shim for the DiffusionGemma engine. |
| |
| The published deployment is vLLM plus its structured decision server. This shim speaks the small part of vLLM's |
| HTTP API that server uses — chat/completions with ``vllm_xargs`` diffusion requests — so the unmodified upstream |
| server runs on top of the MLX engine. That keeps the decision layer (prompt, templates, canvas, slot reads, |
| adaptive thinking) bit-for-bit the published one instead of a re-implementation. |
| |
| Diffusion xargs understood (as sent by runtime/structured_server.py): |
| diffusion_seed_canvas initial canvas token ids |
| diffusion_canvas_length canvas width |
| diffusion_max_steps denoising steps |
| diffusion_read_only return logprobs at every canvas position, do not sample |
| diffusion_pinned canvas positions held at their seed values |
| diffusion_constrained the unembedding is read over ``logprob_token_ids`` only |
| diffusion_samples number of draws |
| and ``djev_bidir`` = flat [a1, b1, a2, b2, ...] prompt ranges that attend bidirectionally (patch_vllm_bidir.py). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import time |
| from typing import Any |
|
|
| import mlx.core as mx |
| import numpy as np |
| from fastapi import FastAPI |
| from fastapi.responses import JSONResponse |
| from pydantic import BaseModel |
|
|
| from .engine import CANVAS_LENGTH, Engine |
|
|
| app = FastAPI() |
| ENGINE: Engine | None = None |
| OPTS: dict[str, Any] = {"canvas": CANVAS_LENGTH, "max_prompt": 8192, "model": "super-decider"} |
|
|
|
|
| class ChatRequest(BaseModel): |
| model: str | None = None |
| messages: list[dict] | None = None |
| prompt: Any = None |
| max_tokens: int | None = None |
| logprobs: Any = None |
| top_logprobs: int | None = None |
| logprob_token_ids: list[int] | None = None |
| return_tokens_as_token_ids: bool = False |
| chat_template_kwargs: dict | None = None |
| vllm_xargs: dict | None = None |
| stop_token_ids: list[int] | None = None |
| seed: int | None = None |
| temperature: float | None = None |
| stream: bool = False |
|
|
|
|
| def _tok(t: int, as_id: bool) -> str: |
| return f"token_id:{t}" if as_id else ENGINE.tokenizer.decode([t]) |
|
|
|
|
| def _block_ids(n_prompt: int, ranges) -> mx.array | None: |
| """Per-token bidirectional group ids from the flat ``djev_bidir`` ranges, -1 outside them.""" |
| if not ranges: |
| return None |
| ids = np.full(n_prompt, -1, dtype=np.int32) |
| for k, i in enumerate(range(0, len(ranges) - 1, 2)): |
| a, b = int(ranges[i]), int(ranges[i + 1]) |
| ids[max(a, 0) : min(b, n_prompt)] = k |
| return mx.array(ids) |
|
|
|
|
| def _usage(prompt_tokens: int, completion_tokens: int) -> dict: |
| return { |
| "prompt_tokens": prompt_tokens, |
| "completion_tokens": completion_tokens, |
| "total_tokens": prompt_tokens + completion_tokens, |
| } |
|
|
|
|
| def _read(req: ChatRequest, prompt_ids: list[int]): |
| """A diffusion read: label logits at every canvas position, plus the canvas tokens themselves.""" |
| x = req.vllm_xargs or {} |
| canvas = [int(t) for t in x.get("diffusion_seed_canvas") or []] |
| canvas_len = int(x.get("diffusion_canvas_length") or OPTS["canvas"]) |
| steps = int(x.get("diffusion_max_steps") or 1) |
| pinned_pos = [int(p) for p in x.get("diffusion_pinned") or []] |
| samples = int(x.get("diffusion_samples") or 1) |
| label_ids = [int(t) for t in (req.logprob_token_ids or [])] |
| if not canvas: |
| raise ValueError("a diffusion request needs diffusion_seed_canvas") |
| canvas = (canvas + [0] * canvas_len)[:canvas_len] |
| block_ids = _block_ids(len(prompt_ids), x.get("djev_bidir")) |
| pinned = {p: canvas[p] for p in pinned_pos if 0 <= p < canvas_len} |
| probs = ENGINE.read( |
| prompt_ids, |
| canvas, |
| list(range(canvas_len)), |
| label_ids, |
| block_ids=block_ids, |
| steps=steps, |
| pinned=pinned, |
| seed=int(req.seed or 0), |
| samples=samples, |
| ) |
| return canvas, label_ids, probs |
|
|
|
|
| def _top_rows(probs_row: np.ndarray, label_ids, as_id: bool, k: int): |
| order = np.argsort(-probs_row)[: max(k, 1)] |
| return [ |
| {"token": _tok(label_ids[i], as_id), "logprob": float(probs_row[i]), "bytes": None} |
| for i in order |
| ] |
|
|
|
|
| def _read_content(canvas, label_ids, probs, as_id: bool, top_k: int): |
| """One logprobs entry per canvas position (``logprobs.content`` / ``top_logprobs``).""" |
| rows = [] |
| for pos, tok in enumerate(canvas): |
| tops = _top_rows(probs[pos], label_ids, as_id, top_k) if label_ids else [] |
| best = tops[0]["logprob"] if tops else 0.0 |
| rows.append( |
| {"token": _tok(tok, as_id), "logprob": best, "bytes": None, "top_logprobs": tops} |
| ) |
| return rows |
|
|
|
|
| def _token_logprobs(tokens: list[int], lps, as_id: bool, chat: bool): |
| """vLLM's legacy per-token logprobs object (``content`` for chat, ``tokens``/``token_logprobs`` for completions). |
| |
| ``top_logprobs`` stays empty: the model reads a label set, not a token ranking, so nothing here asks for it. |
| """ |
| if lps is None: |
| return None |
| vals = [float(v) for v in lps] |
| if chat: |
| return { |
| "content": [ |
| {"token": _tok(t, as_id), "logprob": v, "bytes": None, "top_logprobs": []} |
| for t, v in zip(tokens, vals) |
| ], |
| "refusal": None, |
| } |
| return { |
| "tokens": [_tok(t, as_id) for t in tokens], |
| "token_logprobs": vals, |
| "top_logprobs": [{} for _ in tokens], |
| "text_offset": list(range(len(tokens))), |
| } |
|
|
|
|
| def _generate(req: ChatRequest, prompt_ids: list[int]): |
| x = req.vllm_xargs or {} |
| want = req.logprobs is not None |
| out = ENGINE.generate( |
| prompt_ids, |
| max_new_tokens=req.max_tokens, |
| max_denoising_steps=int(x.get("diffusion_max_steps") or 0) or None, |
| stop_ids=req.stop_token_ids, |
| block_ids=_block_ids(len(prompt_ids), x.get("djev_bidir")), |
| token_logprobs=want, |
| ) |
| gen, lps = out if want else (out, None) |
| if req.stop_token_ids: |
| for k, t in enumerate(gen): |
| if t in req.stop_token_ids: |
| gen = gen[: k + 1] |
| break |
| if req.max_tokens: |
| gen, lps = gen[: req.max_tokens], (lps[: req.max_tokens] if lps else lps) |
| return gen, lps |
|
|
|
|
| @app.get("/v1/models") |
| @app.get("/health") |
| def models(): |
| return {"object": "list", "data": [{"id": OPTS["model"], "object": "model", "owned_by": "local"}]} |
|
|
|
|
| def _handle(req: ChatRequest, chat: bool): |
| if chat: |
| prompt_ids = ENGINE.chat_ids( |
| req.messages or [], |
| bool((req.chat_template_kwargs or {}).get("enable_thinking", False)), |
| ) |
| else: |
| prompt_ids = [int(t) for t in (req.prompt or [])] |
| if len(prompt_ids) > OPTS["max_prompt"]: |
| |
| |
| return JSONResponse( |
| status_code=400, |
| content={ |
| "error": { |
| "code": 400, |
| "message": f"state too long: {len(prompt_ids)} tokens, this build serves " |
| f"{OPTS['max_prompt']} (raise --max-prompt; attention cost grows quadratically)", |
| "type": "BadRequestError", |
| } |
| }, |
| ) |
| as_id = bool(req.return_tokens_as_token_ids) |
| diffusion = bool((req.vllm_xargs or {}).get("diffusion_seed_canvas")) |
| created = int(time.time()) |
| if diffusion: |
| canvas, label_ids, probs = _read(req, prompt_ids) |
| |
| choices = [] |
| for k in range(probs.shape[0]): |
| content = _read_content(canvas, label_ids, probs[k], as_id, int(req.top_logprobs or 20)) |
| choice = { |
| "index": k, |
| "text": "", |
| "finish_reason": "stop", |
| "logprobs": { |
| "tokens": [c["token"] for c in content], |
| "token_logprobs": [c["logprob"] for c in content], |
| "top_logprobs": [{t["token"]: t["logprob"] for t in c["top_logprobs"]} for c in content], |
| "text_offset": list(range(len(content))), |
| }, |
| } |
| if chat: |
| choice["message"] = {"role": "assistant", "content": ""} |
| choice["logprobs"] = {"content": content, "refusal": None} |
| choices.append(choice) |
| body = { |
| "id": "djev-local", |
| "object": "chat.completion" if chat else "text_completion", |
| "created": created, |
| "model": OPTS["model"], |
| "choices": choices, |
| "usage": _usage(len(prompt_ids), len(canvas)), |
| "system_fingerprint": "super-decider-metal", |
| } |
| return body |
|
|
| gen, lps = _generate(req, prompt_ids) |
| text = ENGINE.tokenizer.decode(gen) |
| choice = {"index": 0, "text": text, "finish_reason": "stop"} |
| if chat: |
| choice["message"] = {"role": "assistant", "content": text} |
| choice["logprobs"] = _token_logprobs(gen, lps, as_id, chat) |
| body = { |
| "id": "djev-local", |
| "object": "chat.completion" if chat else "text_completion", |
| "created": created, |
| "model": OPTS["model"], |
| "choices": [choice], |
| "usage": _usage(len(prompt_ids), len(gen)), |
| "system_fingerprint": "super-decider-metal", |
| } |
| return body |
|
|
|
|
| @app.post("/v1/chat/completions") |
| async def chat_completions(req: ChatRequest): |
| return _handle(req, chat=True) |
|
|
|
|
| @app.post("/v1/completions") |
| async def completions(req: ChatRequest): |
| return _handle(req, chat=False) |
|
|
|
|
| def main(): |
| global ENGINE |
| ap = argparse.ArgumentParser(description="super-decider-metal engine server (vLLM-compatible subset)") |
| ap.add_argument("--model", default="artifacts/mlx-nvfp4") |
| ap.add_argument("--name", default="super-decider") |
| ap.add_argument("--host", default="127.0.0.1") |
| ap.add_argument("--port", type=int, default=8000) |
| ap.add_argument("--canvas", type=int, default=CANVAS_LENGTH) |
| ap.add_argument("--max-prompt", type=int, default=8192) |
| args = ap.parse_args() |
| OPTS.update(canvas=args.canvas, max_prompt=args.max_prompt, model=args.name) |
| ENGINE = Engine(args.model, canvas_length=args.canvas) |
| import uvicorn |
|
|
| uvicorn.run(app, host=args.host, port=args.port, log_level="warning") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
