#!/usr/bin/env python3 """OpenAI-compatible, microbatched RivetCoder FP8 serving entry point.""" from __future__ import annotations import argparse import asyncio import concurrent.futures import json import queue import threading import time import uuid from collections import deque from dataclasses import dataclass from pathlib import Path from typing import Any import torch from transformers import AutoModelForCausalLM, AutoTokenizer @dataclass(slots=True) class PendingCompletion: messages: list[dict[str, Any]] max_tokens: int temperature: float top_p: float future: concurrent.futures.Future[str] @property def batch_key(self) -> tuple[int, float, float]: return self.max_tokens, self.temperature, self.top_p class MicrobatchEngine: """Collect compatible requests briefly, then generate them as one batch.""" def __init__( self, model: Any, tokenizer: Any, *, max_batch_size: int, batch_wait_ms: float, ) -> None: self.model = model self.tokenizer = tokenizer self.max_batch_size = int(max_batch_size) self.batch_wait_seconds = float(batch_wait_ms) / 1000.0 self.incoming: queue.Queue[PendingCompletion | None] = queue.Queue() self.deferred: deque[PendingCompletion] = deque() self.thread = threading.Thread(target=self._worker, name="rivetcoder-gpu", daemon=True) self.thread.start() def submit( self, messages: list[dict[str, Any]], *, max_tokens: int, temperature: float, top_p: float, ) -> concurrent.futures.Future[str]: future: concurrent.futures.Future[str] = concurrent.futures.Future() self.incoming.put( PendingCompletion( messages=messages, max_tokens=max_tokens, temperature=temperature, top_p=top_p, future=future, ) ) return future def close(self) -> None: self.incoming.put(None) def _next_request(self) -> PendingCompletion | None: if self.deferred: return self.deferred.popleft() return self.incoming.get() def _collect_batch(self, first: PendingCompletion) -> list[PendingCompletion]: batch = [first] key = first.batch_key deadline = time.perf_counter() + self.batch_wait_seconds while len(batch) < self.max_batch_size: remaining = deadline - time.perf_counter() if remaining <= 0: break try: item = self.incoming.get(timeout=remaining) except queue.Empty: break if item is None: self.incoming.put(None) break if item.batch_key == key: batch.append(item) else: self.deferred.append(item) return batch def _worker(self) -> None: while True: first = self._next_request() if first is None: return batch = self._collect_batch(first) try: results = self._generate(batch) except BaseException as error: for item in batch: item.future.set_exception(error) continue for item, text in zip(batch, results, strict=True): item.future.set_result(text) def _generate(self, batch: list[PendingCompletion]) -> list[str]: rendered = [ self.tokenizer.apply_chat_template( item.messages, add_generation_prompt=True, tokenize=False, ) for item in batch ] encoded = self.tokenizer( rendered, add_special_tokens=False, padding=True, return_tensors="pt", ).to("cuda") prompt_width = encoded["input_ids"].shape[-1] temperature = batch[0].temperature generation_kwargs = { "max_new_tokens": batch[0].max_tokens, "do_sample": temperature > 0, "use_cache": True, "logits_to_keep": 1, "pad_token_id": self.tokenizer.pad_token_id, "eos_token_id": self.tokenizer.eos_token_id, } if temperature > 0: generation_kwargs.update(temperature=temperature, top_p=batch[0].top_p) with torch.no_grad(): generated = self.model.generate(**encoded, **generation_kwargs) return self.tokenizer.batch_decode( generated[:, prompt_width:], skip_special_tokens=True, ) def load_runtime(args: argparse.Namespace) -> tuple[Any, Any, dict[str, Any]]: tokenizer = AutoTokenizer.from_pretrained( args.model, trust_remote_code=True, local_files_only=args.local_files_only, ) tokenizer.padding_side = "left" model = AutoModelForCausalLM.from_pretrained( args.model, trust_remote_code=True, local_files_only=args.local_files_only, dtype=torch.bfloat16, device_map=0, attn_implementation="sdpa", ).eval() model.config.use_cache = True model.set_coding_enabled(True) if not hasattr(model, "enable_fast_fp8_serving"): raise RuntimeError( "The model package does not contain the grouped-FP8 runtime. " "Use the updated RivetCoder FP8 package." ) report = model.enable_fast_fp8_serving() return tokenizer, model, report def warmup(model: Any, tokenizer: Any, batch_sizes: list[int]) -> None: text = tokenizer.apply_chat_template( [{"role": "user", "content": "Return the integer 1."}], add_generation_prompt=True, tokenize=False, ) for batch_size in batch_sizes: encoded = tokenizer( [text] * batch_size, add_special_tokens=False, padding=True, return_tensors="pt", ).to("cuda") with torch.no_grad(): model.generate( **encoded, max_new_tokens=2, do_sample=False, use_cache=True, logits_to_keep=1, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, ) def build_app(engine: MicrobatchEngine, runtime_report: dict[str, Any], model_name: str) -> Any: try: from fastapi import FastAPI, HTTPException except ImportError as error: raise RuntimeError("Serving requires fastapi and uvicorn") from error app = FastAPI(title="RivetCoder FP8 Server") created = int(time.time()) @app.get("/health") async def health() -> dict[str, Any]: return { "status": "ok", "model": model_name, "runtime": runtime_report, "max_batch_size": engine.max_batch_size, "batch_wait_ms": engine.batch_wait_seconds * 1000.0, "cuda_allocated_gib": torch.cuda.memory_allocated() / 1024**3, } @app.get("/v1/models") async def models() -> dict[str, Any]: return { "object": "list", "data": [{"id": model_name, "object": "model", "created": created, "owned_by": "HCHs"}], } @app.post("/v1/chat/completions") async def chat_completions(payload: dict[str, Any]) -> dict[str, Any]: if payload.get("stream", False): raise HTTPException(status_code=400, detail="Streaming is not implemented in the microbatch server") messages = payload.get("messages") if not isinstance(messages, list) or not messages: raise HTTPException(status_code=400, detail="messages must be a non-empty list") max_tokens = int(payload.get("max_tokens", 512)) if max_tokens < 1 or max_tokens > 4096: raise HTTPException(status_code=400, detail="max_tokens must be between 1 and 4096") temperature = float(payload.get("temperature", 0.2)) top_p = float(payload.get("top_p", 0.95)) future = engine.submit( messages, max_tokens=max_tokens, temperature=temperature, top_p=top_p, ) try: text = await asyncio.wrap_future(future) except Exception as error: raise HTTPException(status_code=500, detail=str(error)) from error completion_id = f"chatcmpl-{uuid.uuid4().hex}" return { "id": completion_id, "object": "chat.completion", "created": int(time.time()), "model": model_name, "choices": [ { "index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop", } ], } return app def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--model", default=str(Path("RivetCoder-9B-A4B-FP8")), help="Local FP8 model directory or Hugging Face model id", ) parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--port", type=int, default=8000) parser.add_argument("--max-batch-size", type=int, default=16) parser.add_argument("--batch-wait-ms", type=float, default=3.0) parser.add_argument("--warmup-batches", default="1,8,16") parser.add_argument("--local-files-only", action=argparse.BooleanOptionalAction, default=True) return parser.parse_args() def main() -> int: args = parse_args() if not torch.cuda.is_available(): raise SystemExit("CUDA is required") torch.set_float32_matmul_precision("high") tokenizer, model, runtime_report = load_runtime(args) warmup_batches = [int(value) for value in args.warmup_batches.split(",") if value] warmup(model, tokenizer, warmup_batches) engine = MicrobatchEngine( model, tokenizer, max_batch_size=args.max_batch_size, batch_wait_ms=args.batch_wait_ms, ) app = build_app(engine, runtime_report, str(args.model)) print(json.dumps({"runtime": runtime_report, "listen": f"http://{args.host}:{args.port}"}, indent=2)) import uvicorn uvicorn.run(app, host=args.host, port=args.port, workers=1) return 0 if __name__ == "__main__": raise SystemExit(main())