Feat: Enable structured output in Mini-SGLang

.gitignore

@@ -31,3 +31,4 @@ env/
 # .vscode
 .vscode/
 .coverage
+.idea

benchmark/offline/bench_json.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import argparse
+import time
+from random import seed
+
+import torch
+from minisgl.benchmark.json import (
+    collect_filtered_json_samples,
+    render_json_prompt_ids,
+    validate_json_output,
+)
+from minisgl.core import SamplingParams
+from minisgl.llm import LLM
+from transformers import AutoTokenizer
+
+
+def print_len_stats(name: str, lengths: list[int]) -> None:
+    if not lengths:
+        print(f"{name}: no data")
+        return
+    arr = sorted(lengths)
+    n = len(arr)
+    print(
+        f"{name}: count={n}, min={arr[0]}, p50={arr[int(0.50*n)]}, "
+        f"p90={arr[int(0.90*n)]}, p99={arr[int(0.99*n)]}, max={arr[-1]}"
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default="Qwen/Qwen2-0.5B")
+    parser.add_argument(
+        "--mode",
+        choices=["constrained", "unconstrained"],
+        default="constrained",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    seed(0)
+    MODEL = args.model
+    NUM_SEQS = 100
+    MAX_OUTPUT_LEN = 4096
+    IGNORE_EOS = False
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL)
+    samples = collect_filtered_json_samples(NUM_SEQS)
+    prompt_token_ids = [render_json_prompt_ids(tokenizer, sample) for sample in samples]
+
+    assert prompt_token_ids, "No valid json-mode-eval samples found"
+
+    sampling_params = []
+    for sample in samples:
+        json_schema = sample.json_schema if args.mode == "constrained" else None
+        sampling_params.append(
+            SamplingParams(
+                temperature=0.0,
+                top_k=1,
+                ignore_eos=IGNORE_EOS,
+                max_tokens=MAX_OUTPUT_LEN,
+                json_schema=json_schema,
+            )
+        )
+
+    llm: LLM | None = None
+    try:
+        llm = LLM(MODEL)
+
+        warmup_result = llm.generate(
+            [prompt_token_ids[-1]],
+            sampling_params[-1],
+        )[0]
+        templated_input_preview = tokenizer.decode(
+            prompt_token_ids[-1],
+            skip_special_tokens=False,
+        )
+        templated_input_preview = templated_input_preview.replace("\n", "\\n")
+        warmup_token_ids = warmup_result["token_ids"]
+        warmup_text = warmup_result["text"]
+        print(
+            "Warmup sample: "
+            f"model={MODEL}, "
+            f"mode={args.mode}, "
+            f"input={len(prompt_token_ids[-1])}tok, "
+            f"templated_input_preview='{templated_input_preview}', "
+            f"output={len(warmup_token_ids)}tok, "
+            f"preview='{warmup_text}'"
+        )
+
+        torch.cuda.synchronize(llm.device)
+        t = time.time()
+        bench_results = llm.generate(prompt_token_ids, sampling_params)
+        torch.cuda.synchronize(llm.device)
+        t = time.time() - t
+    finally:
+        if llm is not None:
+            llm.shutdown()
+
+    output_lens = []
+    parse_ok = 0
+    schema_ok = 0
+    schema_checked = 0
+    for sample, result in zip(samples, bench_results):
+        token_ids = result["token_ids"]
+        output_lens.append(len(token_ids))
+        parsed, valid = validate_json_output(result["text"], sample.json_schema)
+        parse_ok += int(parsed)
+        if valid is not None:
+            schema_checked += 1
+            schema_ok += int(valid)
+
+    total_output_budget = sum(sp.max_tokens for sp in sampling_params)
+    total_output_tokens = sum(output_lens)
+
+    print(f"Mode: {args.mode}")
+    print_len_stats("Input length", [len(x) for x in prompt_token_ids])
+    print_len_stats("Output length", output_lens)
+    print(f"Bench requests: {len(prompt_token_ids)}")
+    print(f"Output budget: {total_output_budget}tok, " f"Actual output: {total_output_tokens}tok")
+    print(f"JSON parse: {parse_ok}/{len(bench_results)}")
+    print(f"Schema valid: {schema_ok}/{schema_checked}")
+    throughput = total_output_tokens / t if t > 0 else 0.0
+    print(f"Total: {total_output_tokens}tok, Time: {t:.2f}s, " f"Throughput: {throughput:.2f}tok/s")
+
+
+if __name__ == "__main__":
+    main()

benchmark/online/bench_qwen.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import argparse
 import asyncio
 import os
 import random
@@ -12,6 +13,7 @@
     read_qwen_trace,
     scale_traces,
 )
+from minisgl.benchmark.json import validate_json_output
 from minisgl.utils import init_logger
 from openai import AsyncOpenAI as OpenAI
 from transformers import AutoTokenizer
@@ -34,20 +36,78 @@ def download_qwen_trace(url: str) -> str:
     return str(file_path)
 
 
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Benchmark MiniSGL with Qwen trace replay.")
+    parser.add_argument(
+        "--prompt-mode",
+        choices=["dummy", "json"],
+        default="dummy",
+        help="Prompt source mode.",
+    )
+    parser.add_argument(
+        "--json-mode",
+        choices=["constrained", "unconstrained"],
+        default="constrained",
+        help="Only takes effect when --prompt-mode=json.",
+    )
+    parser.add_argument("--N", type=int, default=1000)
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=4096,
+        help="Only takes effect when --prompt-mode=json.",
+    )
+    return parser.parse_args()
+
+
+def process_json_correctness(results, traces) -> None:
+    total = 0
+    parse_ok = 0
+    schema_ok = 0
+    schema_checked = 0
+    for trace, result in zip(traces, results, strict=True):
+        if trace.json_schema is None:
+            continue
+        total += 1
+        parsed, valid = validate_json_output(result.output_text, trace.json_schema)
+        parse_ok += int(parsed)
+        if valid is not None:
+            schema_checked += 1
+            schema_ok += int(valid)
+
+    logger.info(f"JSON parse: {parse_ok}/{total}")
+    logger.info(f"Schema valid: {schema_ok}/{schema_checked}")
+
+
 async def main():
+    args = parse_args()
     random.seed(42)  # reproducibility
     PORT = 1919
-    N = 1000
     SCALES = [0.4, 0.5, 0.6, 0.7, 0.8, 1.6]  # from fast to slow
     async with OpenAI(base_url=f"http://127.0.0.1:{PORT}/v1", api_key="") as client:
         MODEL = await get_model_name(client)
         tokenizer = AutoTokenizer.from_pretrained(MODEL)
-        TRACES = read_qwen_trace(download_qwen_trace(URL), tokenizer, n=N, dummy=True)
-        logger.info(f"Start benchmarking with {N} requests using model {MODEL}...")
+        traces = read_qwen_trace(
+            download_qwen_trace(URL),
+            tokenizer,
+            n=args.N,
+            prompt_mode=args.prompt_mode,
+            max_new_tokens=args.max_new_tokens,
+            json_mode=args.json_mode,
+        )
+
+        logger.info(f"Start benchmarking with {len(traces)} requests using model {MODEL}...")
+
         for scale in SCALES:
-            traces = scale_traces(TRACES, scale)
-            results = await benchmark_trace(client, traces, MODEL)
+            scaled_traces = scale_traces(traces, scale)
+            results = await benchmark_trace(
+                client,
+                scaled_traces,
+                MODEL,
+            )
             process_benchmark_results(results)
+            if args.prompt_mode == "json":
+                process_json_correctness(results, scaled_traces)
         logger.info("Benchmarking completed.")
 
 

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "torch<2.10.0",
     "transformers>=4.56.0,<=4.57.3",
     "flashinfer-python>=0.5.3",
+    "xgrammar==0.1.27",
     "pyzmq",
     "uvicorn",
     "fastapi",
@@ -47,6 +48,8 @@ dev = [
     "mypy>=0.950",
     "pre-commit>=3.0.0",
     "ruff>=0.11.0",
+    "datasets",
+    "jsonschema",
     "matplotlib>=3.10.5",
     "pyarrow",
 ]