diff --git a/.gitignore b/.gitignore index eaf21add..911a3d4b 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,4 @@ env/ # .vscode .vscode/ .coverage +.idea diff --git a/benchmark/offline/bench_json.py b/benchmark/offline/bench_json.py new file mode 100644 index 00000000..ba98c17f --- /dev/null +++ b/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() diff --git a/benchmark/online/bench_qwen.py b/benchmark/online/bench_qwen.py index 1233011c..bbc9bb12 100644 --- a/benchmark/online/bench_qwen.py +++ b/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.") diff --git a/pyproject.toml b/pyproject.toml index f75470bc..15028910 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/python/minisgl/benchmark/client.py b/python/minisgl/benchmark/client.py index adc82ef1..38e732fe 100644 --- a/python/minisgl/benchmark/client.py +++ b/python/minisgl/benchmark/client.py @@ -18,9 +18,12 @@ @dataclass(frozen=True) class BenchmarkTrace: timestamp: float # unit (second) - message: str + message: str | List[Dict[str, str]] output_length: int # output length in tokens input_length: int | None = None # input length in tokens, optional + ignore_eos: bool = True + response_format: Dict[str, Any] | None = None + json_schema: str | None = None @dataclass(frozen=True) @@ -43,8 +46,9 @@ def from_json(raw: List[float]) -> BenchOneResult: class RawResult: input_len: int | None output_len: int - message: str + message: str | List[Dict[str, str]] tics: List[float] + output_text: str = "" @dataclass @@ -201,50 +205,56 @@ def generate_prompt(tokenizer: Any, n: int) -> str: async def benchmark_one( client: OpenAI, - prompt: str, + prompt: str | List[Dict[str, str]], output_length: int, model: str, *, pbar: Console | bool = True, extra_body: Dict[str, Any] | None = None, input_length: int | None = None, # a hack to force input length + ignore_eos: bool = True, + response_format: Dict[str, Any] | None = None, ) -> RawResult: if isinstance(pbar, bool): pbar = make_console(1, output_length, use_pbar=pbar) with pbar.inflight(1): kwargs = { - "ignore_eos": True, + "ignore_eos": ignore_eos, "top_k": 1, } # this is an internal kwargs that might work for our system if input_length is not None: kwargs["input_length_override"] = input_length kwargs.update(extra_body or {}) # can override kwargs - response = await client.chat.completions.create( - model=model, - stream=True, - messages=[ - { - "role": "user", - "content": prompt, - }, - ], - max_tokens=output_length, - temperature=0.0, - extra_body=kwargs, - ) + messages = prompt if isinstance(prompt, list) else [{"role": "user", "content": prompt}] + + requests = { + "model": model, + "stream": True, + "messages": messages, + "max_tokens": output_length, + "temperature": 0.0, + "extra_body": kwargs, + } + if response_format is not None: + requests["response_format"] = response_format + response = await client.chat.completions.create(**requests) tics = [time.perf_counter()] - async for _ in response: + chunks: List[str] = [] + async for chunk in response: tics.append(time.perf_counter()) if len(tics) == 2: pbar.update_prefill() elif len(tics) <= output_length + 1: pbar.update_decode() + if delta := chunk.choices[0].delta.content: + chunks.append(delta) return RawResult( input_len=input_length, output_len=output_length, message=prompt, tics=tics, + output_text="".join(chunks), ) @@ -257,6 +267,7 @@ async def benchmark_one_batch( extra_body: Dict[str, Any] | None = None, input_lengths: List[int | None] | None = None, pbar: Console | bool = True, + ignore_eos: bool = True, ) -> List[RawResult]: if isinstance(output_lengths, int): output_lengths = [output_lengths] * len(prompts) @@ -275,6 +286,7 @@ async def benchmark_one_batch( pbar=pbar, extra_body=extra_body, input_length=input_length, + ignore_eos=ignore_eos, ) for prompt, output_length, input_length in zip( prompts, output_lengths, input_lengths, strict=True @@ -301,7 +313,14 @@ async def benchmark_timed(msg: BenchmarkTrace): target = start + msg.timestamp - offset await asyncio.sleep(max(0, target - time.perf_counter())) return await benchmark_one( - client, msg.message, msg.output_length, model, pbar=pbar, input_length=msg.input_length + client, + msg.message, + msg.output_length, + model, + pbar=pbar, + input_length=msg.input_length, + ignore_eos=msg.ignore_eos, + response_format=msg.response_format, ) tasks = [benchmark_timed(msg) for msg in msgs] @@ -408,7 +427,9 @@ def read_qwen_trace( file_path: str, tokenizer: Any, n: int | None = None, - dummy: bool = False, + prompt_mode: str = "dummy", + max_new_tokens: int = 4096, + json_mode: str = "constrained", ) -> List[BenchmarkTrace]: class JSONInput(BaseModel): chat_id: int @@ -425,12 +446,30 @@ class JSONInput(BaseModel): if n is not None: lines = lines[:n] objs = [JSONInput.model_validate_json(line) for line in lines] - if dummy: + if prompt_mode == "dummy": prompt = generate_prompt(tokenizer, max(obj.input_length for obj in objs)) ids = tokenizer.encode(prompt, add_special_tokens=False) _get_prompt = lambda obj: tokenizer.decode(ids[: obj.input_length]) - else: + elif prompt_mode == "random": _get_prompt = lambda obj: generate_prompt(tokenizer, obj.input_length) + elif prompt_mode == "json": + from minisgl.benchmark.json import collect_repeated_json_samples, render_json_prompt_ids + + samples = collect_repeated_json_samples(len(objs)) + return [ + BenchmarkTrace( + timestamp=obj.timestamp, + message=sample.messages, + input_length=len(render_json_prompt_ids(tokenizer, sample)), + output_length=max_new_tokens, + ignore_eos=False, + response_format=sample.response_format if json_mode == "constrained" else None, + json_schema=sample.json_schema, + ) + for obj, sample in zip(objs, samples, strict=True) + ] + else: + raise ValueError(f"Unknown prompt_mode: {prompt_mode}") return [ BenchmarkTrace( timestamp=obj.timestamp, @@ -469,8 +508,8 @@ class JSONInput(BaseModel): BenchmarkTrace( timestamp=obj.timestamp / 1000, message=_get_prompt(obj), - input_length=obj.input_length, output_length=obj.output_length, + input_length=obj.input_length, ) for obj in objs ] @@ -488,6 +527,9 @@ def scale_traces( message=trace.message, input_length=trace.input_length, output_length=trace.output_length, + ignore_eos=trace.ignore_eos, + response_format=trace.response_format, + json_schema=trace.json_schema, ) for trace in traces ], diff --git a/python/minisgl/benchmark/json.py b/python/minisgl/benchmark/json.py new file mode 100644 index 00000000..d1720c2d --- /dev/null +++ b/python/minisgl/benchmark/json.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import json +import re +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Iterator + +if TYPE_CHECKING: + from transformers import PreTrainedTokenizerBase + +DEFAULT_JSON_DATA_PATH = "NousResearch/json-mode-eval" + + +@dataclass(frozen=True) +class JsonBenchSample: + system: str + user: str + json_schema: str + + @property + def messages(self) -> list[dict[str, str]]: + return [ + {"role": "system", "content": self.system}, + {"role": "user", "content": self.user}, + ] + + @property + def response_format(self) -> dict[str, object]: + return { + "type": "json_schema", + "json_schema": {"schema": json.loads(self.json_schema)}, + } + + +def _normalize_content(content: Any) -> str: + if isinstance(content, str): + return content + return json.dumps(content, ensure_ascii=False) + + +def iter_filtered_json_samples( + data_path: str = DEFAULT_JSON_DATA_PATH, +) -> Iterator[JsonBenchSample]: + from datasets import load_dataset + + raw_dataset = load_dataset(data_path) + for data in raw_dataset["train"]: + messages = data["prompt"] + schema = json.loads(data["schema"]) + + if schema.get("type") is None: + continue + + assert len(messages) == 2, "invalid message length" + system = messages[0] + user = messages[1] + assert system["role"] == "system", "invalid role" + assert user["role"] == "user", "invalid role" + yield JsonBenchSample( + system=_normalize_content(system["content"]), + user=_normalize_content(user["content"]), + json_schema=json.dumps(schema, separators=(",", ":"), sort_keys=True), + ) + + +def collect_filtered_json_samples( + n: int | None = None, + data_path: str = DEFAULT_JSON_DATA_PATH, +) -> list[JsonBenchSample]: + samples = [] + for sample in iter_filtered_json_samples(data_path): + samples.append(sample) + if n is not None and len(samples) >= n: + break + return samples + + +def collect_repeated_json_samples( + n: int, + data_path: str = DEFAULT_JSON_DATA_PATH, +) -> list[JsonBenchSample]: + if n <= 0: + return [] + + samples = collect_filtered_json_samples(data_path=data_path) + if not samples: + return [] + return [samples[i % len(samples)] for i in range(n)] + + +def render_json_prompt_ids( + tokenizer: PreTrainedTokenizerBase | Any, + sample: JsonBenchSample, +) -> list[int]: + return tokenizer.apply_chat_template( + sample.messages, + tokenize=True, + add_generation_prompt=True, + ) + + +def normalize_json_output(output: str) -> str: + # Reasoning models may emit analysis before the final answer. + if "" in output: + output = output.split("", 1)[1] + output = output.strip() + + # Some models wrap the final JSON in one fenced Markdown block. + matched = re.fullmatch(r"```(?:json)?[ \t]*\r?\n(.*)\r?\n```", output, re.S | re.I) + if matched is not None: + output = matched.group(1).strip() + return output + + +def validate_json_output(output: str, json_schema: str) -> tuple[bool, bool | None]: + output = normalize_json_output(output) + + try: + obj = json.loads(output) + except Exception: + return False, False + + try: + import jsonschema + except ImportError: + return True, None + + try: + schema = json.loads(json_schema) + validator_cls = jsonschema.validators.validator_for(schema) + validator_cls.check_schema(schema) + validator = validator_cls(schema, format_checker=jsonschema.FormatChecker()) + validator.validate(obj) + except Exception: + return True, False + return True, True diff --git a/python/minisgl/constrained/__init__.py b/python/minisgl/constrained/__init__.py new file mode 100644 index 00000000..ba77a5ab --- /dev/null +++ b/python/minisgl/constrained/__init__.py @@ -0,0 +1,20 @@ +from .base import ( + INVALID_GRAMMAR_OBJ, + BaseGrammarBackend, + BaseGrammarObject, + GrammarKey, + GrammarValue, + create_grammar_backend, +) +from .reasoner_backend import ReasonerGrammarBackend, ReasonerGrammarObject + +__all__ = [ + "BaseGrammarBackend", + "BaseGrammarObject", + "GrammarKey", + "GrammarValue", + "INVALID_GRAMMAR_OBJ", + "ReasonerGrammarBackend", + "ReasonerGrammarObject", + "create_grammar_backend", +] diff --git a/python/minisgl/constrained/base.py b/python/minisgl/constrained/base.py new file mode 100644 index 00000000..3bfbd340 --- /dev/null +++ b/python/minisgl/constrained/base.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from concurrent.futures import Future, ThreadPoolExecutor +from typing import TypeAlias + +import torch + +GrammarKey: TypeAlias = tuple[str, str] + + +class BaseGrammarObject: + def __init__(self) -> None: + self._finished = False + self.current_token: int | None = None + + def accept_token(self, token: int) -> None: + raise NotImplementedError() + + def rollback(self, k: int) -> None: + raise NotImplementedError() + + def is_terminated(self) -> bool: + return False + + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device: torch.device + ) -> torch.Tensor: + raise NotImplementedError() + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + raise NotImplementedError() + + @staticmethod + def move_vocab_mask(vocab_mask: torch.Tensor, device: torch.device) -> torch.Tensor: + raise NotImplementedError() + + @staticmethod + def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: + raise NotImplementedError() + + def copy(self) -> BaseGrammarObject: + return self + + @property + def finished(self) -> bool: + return self._finished + + @finished.setter + def finished(self, finished: bool) -> None: + self._finished = finished + + +GrammarFuture: TypeAlias = Future[BaseGrammarObject | None] +GrammarValue: TypeAlias = BaseGrammarObject | GrammarFuture + + +INVALID_GRAMMAR_OBJ = BaseGrammarObject() + + +class BaseGrammarBackend: + def __init__(self) -> None: + self.executor = ThreadPoolExecutor() + self.cache: dict[GrammarKey, BaseGrammarObject] = {} + + def dispatch_json(self, key_string: str) -> BaseGrammarObject | None: + raise NotImplementedError() + + def _init_value_dispatch(self, key: GrammarKey) -> BaseGrammarObject | None: + key_type, key_string = key + if key_type == "json": + return self.dispatch_json(key_string) + raise NotImplementedError(f"Structured output format is not implemented: {key_type}") + + def get_cached_or_future_value(self, key: GrammarKey) -> tuple[GrammarValue, bool]: + value = self.cache.get(key) + if value is not None: + return value.copy(), True + return self.executor.submit(self._init_value_dispatch, key), False + + def set_cache(self, key: GrammarKey, value: BaseGrammarObject) -> None: + self.cache[key] = value + + def reset(self) -> None: + self.cache.clear() + + def shutdown(self) -> None: + self.reset() + self.executor.shutdown(wait=False, cancel_futures=True) + + +def create_grammar_backend( + tokenizer, vocab_size: int, eos_token_ids: int | list[int] | set[int] | None = None +) -> BaseGrammarBackend | None: + from .reasoner_backend import ReasonerGrammarBackend + from .xgrammar_backend import XGrammarBackend + + try: + backend: BaseGrammarBackend = XGrammarBackend(tokenizer, vocab_size, eos_token_ids) + except Exception: + return None + think_end_id = getattr(tokenizer, "think_end_id", None) + if think_end_id is not None: + backend = ReasonerGrammarBackend(backend, think_end_id) + return backend diff --git a/python/minisgl/constrained/reasoner_backend.py b/python/minisgl/constrained/reasoner_backend.py new file mode 100644 index 00000000..0a89a299 --- /dev/null +++ b/python/minisgl/constrained/reasoner_backend.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +import torch + +from .base import INVALID_GRAMMAR_OBJ, BaseGrammarBackend, BaseGrammarObject + + +class ReasonerGrammarObject(BaseGrammarObject): + def __init__(self, grammar: BaseGrammarObject, think_end_id: int) -> None: + super().__init__() + self.grammar = grammar + self.think_end_id = think_end_id + self.tokens_after_think_end = -1 + + def _advance_state(self, token: int) -> None: + if self.tokens_after_think_end == -1 and token == self.think_end_id: + self.tokens_after_think_end = 0 + elif self.tokens_after_think_end >= 0: + self.tokens_after_think_end += 1 + + def _rollback_state(self) -> None: + if self.tokens_after_think_end == 0: + self.tokens_after_think_end = -1 + elif self.tokens_after_think_end > 0: + self.tokens_after_think_end -= 1 + + def accept_token(self, token: int) -> None: + self.current_token = token + if self.tokens_after_think_end >= 0: + self.grammar.accept_token(token) + self._advance_state(token) + + def rollback(self, k: int) -> None: + steps_after_think = min(k, self.tokens_after_think_end) + if steps_after_think > 0: + self.grammar.rollback(steps_after_think) + for _ in range(k): + self._rollback_state() + self.current_token = None + + def is_terminated(self) -> bool: + return self.grammar.is_terminated() + + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device: torch.device + ) -> torch.Tensor: + return self.grammar.allocate_vocab_mask(vocab_size, batch_size, device) + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + if self.tokens_after_think_end >= 0: + self.grammar.fill_vocab_mask(vocab_mask, idx) + + def move_vocab_mask(self, vocab_mask: torch.Tensor, device: torch.device) -> torch.Tensor: + return self.grammar.move_vocab_mask(vocab_mask, device) + + def apply_vocab_mask(self, logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: + self.grammar.apply_vocab_mask(logits, vocab_mask) + + def copy(self) -> BaseGrammarObject: + return ReasonerGrammarObject(self.grammar.copy(), self.think_end_id) + + @property + def finished(self) -> bool: + return self.grammar.finished + + @finished.setter + def finished(self, finished: bool) -> None: + self.grammar.finished = finished + + +class ReasonerGrammarBackend(BaseGrammarBackend): + def __init__(self, grammar_backend: BaseGrammarBackend, think_end_id: int) -> None: + super().__init__() + self.grammar_backend = grammar_backend + self.think_end_id = think_end_id + + def dispatch_json(self, key_string: str) -> BaseGrammarObject | None: + ret = self.grammar_backend.dispatch_json(key_string) + if ret is None or ret is INVALID_GRAMMAR_OBJ: + return ret + return ReasonerGrammarObject(ret, self.think_end_id) + + def reset(self) -> None: + super().reset() + self.grammar_backend.reset() + + def shutdown(self) -> None: + super().shutdown() + self.grammar_backend.shutdown() diff --git a/python/minisgl/constrained/xgrammar_backend.py b/python/minisgl/constrained/xgrammar_backend.py new file mode 100644 index 00000000..d7f9a73c --- /dev/null +++ b/python/minisgl/constrained/xgrammar_backend.py @@ -0,0 +1,164 @@ +from __future__ import annotations + +import json + +import torch +from xgrammar import ( + CompiledGrammar, + GrammarCompiler, + GrammarMatcher, + TokenizerInfo, + allocate_token_bitmask, +) + +from .base import INVALID_GRAMMAR_OBJ, BaseGrammarBackend, BaseGrammarObject + +MAX_ROLLBACK_TOKENS = 200 + + +def _apply_vocab_mask_torch(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: + # NOTE: logits shape: [batch, vocab], vocab_mask shape: [batch, ceil(vocab / 32)]. + vocab_size = min(logits.shape[1], vocab_mask.shape[1] * 32) + shifts = torch.arange(32, dtype=torch.int32, device=vocab_mask.device) + unpacked = ((vocab_mask.unsqueeze(-1) >> shifts) & 1).reshape(vocab_mask.shape[0], -1) + logits[:, :vocab_size].masked_fill_(~unpacked[:, :vocab_size].to(torch.bool), float("-inf")) + + +class XGrammarGrammar(BaseGrammarObject): + def __init__( + self, + matcher: GrammarMatcher, + vocab_size: int, + ctx: CompiledGrammar, + override_stop_tokens: list[int] | int | None, + key_string: str | None = None, + ) -> None: + super().__init__() + self.matcher = matcher + self.vocab_size = vocab_size + self.ctx = ctx + self.override_stop_tokens = override_stop_tokens + self.key_string = key_string + self.accepted_tokens: list[int] = [] + + def accept_token(self, token: int) -> None: + if self.is_terminated(): + return + self.current_token = token + if not self.matcher.accept_token(token): + raise ValueError( + f"Token {token} not accepted by grammar {self.key_string!r}: " + f"{self.accepted_tokens=}" + ) + self.accepted_tokens.append(token) + + def rollback(self, k: int) -> None: + self.matcher.rollback(k) + if k > 0: + self.accepted_tokens = self.accepted_tokens[:-k] + self.current_token = self.accepted_tokens[-1] if self.accepted_tokens else None + + def is_terminated(self) -> bool: + return self.matcher.is_terminated() + + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device: torch.device + ) -> torch.Tensor: + return allocate_token_bitmask(batch_size, vocab_size) + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + self.matcher.fill_next_token_bitmask(vocab_mask, idx) + + @staticmethod + def move_vocab_mask(vocab_mask: torch.Tensor, device: torch.device) -> torch.Tensor: + return vocab_mask.to(device, non_blocking=True) + + @staticmethod + def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: + _apply_vocab_mask_torch(logits, vocab_mask) + + def copy(self) -> BaseGrammarObject: + matcher = GrammarMatcher( + self.ctx, + max_rollback_tokens=MAX_ROLLBACK_TOKENS, + override_stop_tokens=self.override_stop_tokens, + ) + return XGrammarGrammar( + matcher, + self.vocab_size, + self.ctx, + self.override_stop_tokens, + self.key_string, + ) + + +class TokenizerNotSupportedError(Exception): + pass + + +class XGrammarBackend(BaseGrammarBackend): + def __init__( + self, + tokenizer, + vocab_size: int, + eos_token_ids: int | list[int] | set[int] | None = None, + any_whitespace: bool = True, + ) -> None: + super().__init__() + try: + tokenizer_info = TokenizerInfo.from_huggingface( + tokenizer, + vocab_size=vocab_size, + stop_token_ids=self._normalize_eos_token_ids(eos_token_ids), + ) + except Exception as e: + raise TokenizerNotSupportedError( + f"Failed to create XGrammar TokenizerInfo from tokenizer: {e}" + ) from e + override_stop_tokens = None + + self.vocab_size = vocab_size + self.any_whitespace = any_whitespace + self.override_stop_tokens = override_stop_tokens + self.grammar_compiler = GrammarCompiler(tokenizer_info=tokenizer_info) + + @staticmethod + def _normalize_eos_token_ids( + eos_token_ids: int | list[int] | set[int] | None, + ) -> list[int] | None: + if eos_token_ids is None: + return None + if isinstance(eos_token_ids, int): + return [eos_token_ids] + return list(eos_token_ids) + + def _from_context(self, ctx: CompiledGrammar, key_string: str) -> XGrammarGrammar: + matcher = GrammarMatcher( + ctx, + max_rollback_tokens=MAX_ROLLBACK_TOKENS, + override_stop_tokens=self.override_stop_tokens, + ) + return XGrammarGrammar( + matcher, + self.vocab_size, + ctx, + self.override_stop_tokens, + key_string, + ) + + def dispatch_json(self, key_string: str) -> BaseGrammarObject | None: + try: + if key_string == "$$ANY$$": + ctx = self.grammar_compiler.compile_builtin_json_grammar() + else: + ctx = self.grammar_compiler.compile_json_schema( + schema=key_string, + any_whitespace=self.any_whitespace, + ) + except (RuntimeError, json.decoder.JSONDecodeError, UnicodeDecodeError): + return INVALID_GRAMMAR_OBJ + return self._from_context(ctx, key_string) + + def reset(self) -> None: + super().reset() + self.grammar_compiler.clear_cache() diff --git a/python/minisgl/core.py b/python/minisgl/core.py index be4d643e..e6982187 100644 --- a/python/minisgl/core.py +++ b/python/minisgl/core.py @@ -8,6 +8,7 @@ if TYPE_CHECKING: from minisgl.attention import BaseAttnBackend, BaseAttnMetadata + from minisgl.constrained import GrammarKey, GrammarValue from minisgl.kvcache import BaseCacheHandle, BaseKVCachePool from minisgl.moe import BaseMoeBackend @@ -19,12 +20,20 @@ class SamplingParams: top_p: float = 1.0 ignore_eos: bool = False max_tokens: int = 1024 + json_schema: str | None = None @property def is_greedy(self) -> bool: return (self.temperature <= 0.0 or self.top_k == 1) and self.top_p == 1.0 +@dataclass +class ReqConstraintState: + grammar_key: GrammarKey | None = None + grammar: GrammarValue | None = None + grammar_wait_ct: int = 0 + + @dataclass(eq=False) class Req: input_ids: torch.Tensor # cpu tensor @@ -34,6 +43,8 @@ class Req: uid: int sampling_params: SamplingParams cache_handle: BaseCacheHandle + constraint: ReqConstraintState | None = None + finished: bool = False def __post_init__(self) -> None: assert self.input_ids.is_cpu @@ -49,16 +60,34 @@ def remain_len(self) -> int: def extend_len(self) -> int: return self.device_len - self.cached_len - def complete_one(self) -> None: + def commit_forward(self) -> None: self.cached_len = self.device_len + + def advance_token(self) -> None: self.device_len += 1 def append_host(self, next_token: torch.Tensor) -> None: self.input_ids = torch.cat([self.input_ids, next_token]) + def should_finish( + self, + next_token: int, + eos_token_id: int, + *, + grammar_terminated: bool = False, + ) -> bool: + finished = not self.can_decode or grammar_terminated + if not self.sampling_params.ignore_eos: + finished |= next_token == eos_token_id + return finished + @property def can_decode(self) -> bool: - return self.remain_len > 0 + return not self.finished and self.remain_len > 0 + + @property + def is_constrained(self) -> bool: + return self.constraint is not None def __repr__(self) -> str: return ( diff --git a/python/minisgl/distributed/impl.py b/python/minisgl/distributed/impl.py index 5c63f4a8..088a77f0 100644 --- a/python/minisgl/distributed/impl.py +++ b/python/minisgl/distributed/impl.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING, List, Literal import torch import torch.distributed as dist @@ -11,11 +11,21 @@ from minisgl.distributed import DistributedInfo from minisgl.kernel import PyNCCLCommunicator +ReduceOp = Literal["sum", "prod", "max", "min", "avg"] + +TORCH_REDUCE_OPS = { + "sum": dist.ReduceOp.SUM, + "prod": dist.ReduceOp.PRODUCT, + "max": dist.ReduceOp.MAX, + "min": dist.ReduceOp.MIN, + "avg": dist.ReduceOp.AVG, +} + @dataclass class DistributedImpl(ABC): @abstractmethod - def all_reduce(self, x: torch.Tensor) -> torch.Tensor: ... + def all_reduce(self, x: torch.Tensor, op: ReduceOp = "sum") -> torch.Tensor: ... @abstractmethod def all_gather(self, x: torch.Tensor) -> torch.Tensor: ... @@ -23,11 +33,14 @@ def all_gather(self, x: torch.Tensor) -> torch.Tensor: ... @dataclass class TorchDistributedImpl(DistributedImpl): - def all_reduce(self, x: torch.Tensor) -> torch.Tensor: + def all_reduce(self, x: torch.Tensor, op: ReduceOp = "sum") -> torch.Tensor: tp_size = dist.get_world_size() if tp_size == 1: return x - dist.all_reduce(x, op=dist.ReduceOp.SUM) + reduce_op = TORCH_REDUCE_OPS.get(op) + if reduce_op is None: + raise ValueError(f"Unsupported reduce op: {op}") + dist.all_reduce(x, op=reduce_op) return x def all_gather(self, x: torch.Tensor) -> torch.Tensor: @@ -45,8 +58,8 @@ def all_gather(self, x: torch.Tensor) -> torch.Tensor: class PyNCCLDistributedImpl(DistributedImpl): comm: PyNCCLCommunicator - def all_reduce(self, x: torch.Tensor) -> torch.Tensor: - self.comm.all_reduce(x, "sum") + def all_reduce(self, x: torch.Tensor, op: ReduceOp = "sum") -> torch.Tensor: + self.comm.all_reduce(x, op) return x def all_gather(self, x: torch.Tensor) -> torch.Tensor: @@ -63,8 +76,8 @@ def all_gather(self, x: torch.Tensor) -> torch.Tensor: class DistributedCommunicator: plugins: List[DistributedImpl] = [TorchDistributedImpl()] - def all_reduce(self, x: torch.Tensor) -> torch.Tensor: - return self.plugins[-1].all_reduce(x) + def all_reduce(self, x: torch.Tensor, op: ReduceOp = "sum") -> torch.Tensor: + return self.plugins[-1].all_reduce(x, op) def all_gather(self, x: torch.Tensor) -> torch.Tensor: return self.plugins[-1].all_gather(x) diff --git a/python/minisgl/engine/__init__.py b/python/minisgl/engine/__init__.py index cf74975d..3ae06630 100644 --- a/python/minisgl/engine/__init__.py +++ b/python/minisgl/engine/__init__.py @@ -1,5 +1,5 @@ from .config import EngineConfig -from .engine import Engine, ForwardOutput +from .engine import Engine, ForwardOutput, ModelForwardOutput from .sample import BatchSamplingArgs -__all__ = ["Engine", "EngineConfig", "ForwardOutput", "BatchSamplingArgs"] +__all__ = ["Engine", "EngineConfig", "ForwardOutput", "ModelForwardOutput", "BatchSamplingArgs"] diff --git a/python/minisgl/engine/engine.py b/python/minisgl/engine/engine.py index ea29a96b..bc7971dd 100644 --- a/python/minisgl/engine/engine.py +++ b/python/minisgl/engine/engine.py @@ -6,7 +6,13 @@ import torch from minisgl.attention import create_attention_backend from minisgl.core import Batch, Context, Req, set_global_ctx -from minisgl.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info +from minisgl.distributed import ( + DistributedCommunicator, + destroy_distributed, + enable_pynccl_distributed, + set_tp_info, +) +from minisgl.env import ENV from minisgl.kvcache import create_kvcache_pool from minisgl.layers import set_rope_device from minisgl.models import create_model, load_weight @@ -26,9 +32,15 @@ class ForwardOutput(NamedTuple): copy_done_event: torch.cuda.Event +class ModelForwardOutput(NamedTuple): + logits: torch.Tensor + + class Engine: def __init__(self, config: EngineConfig): assert not torch.cuda.is_initialized() + torch.set_num_threads(ENV.TORCH_NUM_THREADS.value) # reduce thread conflicts + set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size) _adjust_config(config) @@ -38,6 +50,8 @@ def __init__(self, config: EngineConfig): self.stream = torch.cuda.Stream() torch.cuda.set_stream(self.stream) self.dtype = config.dtype + self.tp_size = config.tp_info.size + self.tp_comm = DistributedCommunicator() self.ctx = Context(config.page_size) set_global_ctx(self.ctx) @@ -188,18 +202,31 @@ def _sync_get_memory(self) -> Tuple[int, int]: return min_free_memory, max_free_memory - def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: + def _sync_next_tokens(self, next_tokens: torch.Tensor, args: BatchSamplingArgs) -> None: + if self.tp_size == 1 or not args.has_grammar: + return + + self.tp_comm.all_reduce(next_tokens, "min") + + def forward_batch(self, batch: Batch) -> ModelForwardOutput: assert torch.cuda.current_stream() == self.stream with self.ctx.forward_batch(batch): if self.graph_runner.can_use_cuda_graph(batch): - logits = self.graph_runner.replay(batch) - else: - logits = self.model.forward() - + return ModelForwardOutput(self.graph_runner.replay(batch)) + return ModelForwardOutput(self.model.forward()) + + def sample_batch( + self, + batch: Batch, + forward_output: ModelForwardOutput, + args: BatchSamplingArgs, + ) -> ForwardOutput: + next_tokens_gpu = self.sampler.sample(forward_output.logits[: batch.size], args).to( + torch.int32 + ) + self._sync_next_tokens(next_tokens_gpu, args) for req in batch.reqs: - req.complete_one() - - next_tokens_gpu = self.sampler.sample(logits[: batch.size], args).to(torch.int32) + req.advance_token() next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True) copy_done_event = torch.cuda.Event() copy_done_event.record(self.stream) diff --git a/python/minisgl/engine/sample.py b/python/minisgl/engine/sample.py index cb6c7ee7..6b67891e 100644 --- a/python/minisgl/engine/sample.py +++ b/python/minisgl/engine/sample.py @@ -7,6 +7,7 @@ from minisgl.utils import is_sm90_supported, nvtx_annotate if TYPE_CHECKING: + from minisgl.constrained import BaseGrammarObject from minisgl.core import Batch @@ -15,6 +16,11 @@ class BatchSamplingArgs: temperatures: torch.Tensor | None top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None + grammars: List[BaseGrammarObject | None] | None = None + + @property + def has_grammar(self) -> bool: + return self.grammars is not None def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor: @@ -50,10 +56,34 @@ class Sampler: device: torch.device vocab_size: int + def _apply_grammar_mask(self, logits: torch.Tensor, args: BatchSamplingArgs) -> None: + grammars = args.grammars + if grammars is None: + return + + first_grammar = next((grammar for grammar in grammars if grammar is not None), None) + if first_grammar is None: + return + + vocab_mask = first_grammar.allocate_vocab_mask( + vocab_size=self.vocab_size, + batch_size=len(grammars), + device=logits.device, + ) + for i, grammar in enumerate(grammars): + if grammar and not grammar.finished and not grammar.is_terminated(): + grammar.fill_vocab_mask(vocab_mask, i) + vocab_mask = first_grammar.move_vocab_mask(vocab_mask, logits.device) + first_grammar.apply_vocab_mask(logits, vocab_mask) + def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] + grammars = None + if any(r.is_constrained for r in batch.reqs): + grammars = [r.constraint.grammar if r.constraint else None for r in batch.reqs] + if all(p.is_greedy for p in params): - return BatchSamplingArgs(temperatures=None) + return BatchSamplingArgs(temperatures=None, grammars=grammars) MIN_P = MIN_T = 1e-6 ts = [max(0.0 if p.is_greedy else p.temperature, MIN_T) for p in params] @@ -65,11 +95,12 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: top_k = make_device_tensor(top_ks, torch.int32, self.device) if any(p < 1.0 for p in top_ps): top_p = make_device_tensor(top_ps, torch.float32, self.device) - return BatchSamplingArgs(temperatures, top_k=top_k, top_p=top_p) + return BatchSamplingArgs(temperatures, top_k=top_k, top_p=top_p, grammars=grammars) @nvtx_annotate("Sampler") def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): + self._apply_grammar_mask(logits, args) if args.temperatures is None: # greedy sampling return torch.argmax(logits, dim=-1) return sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) diff --git a/python/minisgl/env.py b/python/minisgl/env.py index 1f7edf19..6ec6b8df 100644 --- a/python/minisgl/env.py +++ b/python/minisgl/env.py @@ -67,7 +67,10 @@ class EnvClassSingleton: # backend runtime FLASHINFER_USE_TENSOR_CORES = EnvOption() DISABLE_OVERLAP_SCHEDULING = EnvBool(False) + TORCH_NUM_THREADS = EnvInt(1) PYNCCL_MAX_BUFFER_SIZE = EnvMem(1024**3) + GRAMMAR_POLL_INTERVAL = EnvFloat(0.005) + GRAMMAR_MAX_POLL_ITERATIONS = EnvInt(10000) def __new__(cls): # single instance diff --git a/python/minisgl/kernel/csrc/src/pynccl.cu b/python/minisgl/kernel/csrc/src/pynccl.cu index cb1cd9df..1de69b17 100644 --- a/python/minisgl/kernel/csrc/src/pynccl.cu +++ b/python/minisgl/kernel/csrc/src/pynccl.cu @@ -60,6 +60,7 @@ const auto kNCCLDtypeMap = std::unordered_map{ {{DLDataTypeCode::kDLFloat, 16, 1}, ncclFloat16}, {{DLDataTypeCode::kDLBfloat, 16, 1}, ncclBfloat16}, + {{DLDataTypeCode::kDLInt, 32, 1}, ncclInt32}, }; using std::shared_ptr; diff --git a/python/minisgl/kernel/pynccl.py b/python/minisgl/kernel/pynccl.py index 9c59502d..762bea91 100644 --- a/python/minisgl/kernel/pynccl.py +++ b/python/minisgl/kernel/pynccl.py @@ -7,6 +7,8 @@ from .utils import load_aot +ReduceOp = Literal["sum", "prod", "max", "min", "avg"] + if TYPE_CHECKING: from abc import abstractmethod @@ -15,7 +17,7 @@ class PyNCCLCommunicator: @abstractmethod - def all_reduce(self, input: torch.Tensor, op: Literal["sum"]) -> None: ... + def all_reduce(self, input: torch.Tensor, op: ReduceOp = "sum") -> None: ... @abstractmethod def all_gather(self, output: torch.Tensor, input: torch.Tensor) -> None: ... @abstractmethod diff --git a/python/minisgl/scheduler/grammar.py b/python/minisgl/scheduler/grammar.py new file mode 100644 index 00000000..f4c0a09d --- /dev/null +++ b/python/minisgl/scheduler/grammar.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import time +from concurrent import futures +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Literal, Set, Tuple, TypeAlias + +import torch +from minisgl.constrained import INVALID_GRAMMAR_OBJ, create_grammar_backend +from minisgl.env import ENV + +from .utils import PendingReq + +if TYPE_CHECKING: + from minisgl.constrained import BaseGrammarObject, GrammarKey + from transformers import PreTrainedTokenizerBase + +GRAMMAR_JSON = "json" +GRAMMAR_READY = "ready" +GRAMMAR_FAILED = "failed" +GRAMMAR_QUEUED = "queued" +GrammarSubmitStatus: TypeAlias = Literal["ready", "failed", "queued"] + + +@dataclass +class GrammarPollResult: + ready_reqs: List[PendingReq] + failed_uids: List[int] + + +class GrammarManager: + def __init__( + self, + tokenizer: PreTrainedTokenizerBase, + vocab_size: int, + eos_token_id: int, + tp_cpu_group: torch.distributed.ProcessGroup, + ) -> None: + self.grammar_queue: List[PendingReq] = [] + self.grammar_backend = create_grammar_backend( + tokenizer, + vocab_size, + eos_token_id, + ) + self.tp_cpu_group = tp_cpu_group + self.tp_size = torch.distributed.get_world_size(group=tp_cpu_group) + self.poll_interval = ENV.GRAMMAR_POLL_INTERVAL.value + self.max_poll_iterations = ENV.GRAMMAR_MAX_POLL_ITERATIONS.value + + @property + def runnable(self) -> bool: + return len(self.grammar_queue) > 0 + + def _get_grammar_key(self, req: PendingReq) -> GrammarKey: + if (json_schema := req.sampling_params.json_schema) is not None: + return (GRAMMAR_JSON, json_schema) + raise NotImplementedError("Structured output format is not implemented") + + def submit(self, req: PendingReq) -> GrammarSubmitStatus: + assert req.constraint is not None + key = self._get_grammar_key(req) + req.constraint.grammar_key = key + + backend = self.grammar_backend + if backend is None: + req.constraint.grammar = None + return GRAMMAR_FAILED + + value, cache_hit = backend.get_cached_or_future_value(key) + req.constraint.grammar = value + if cache_hit: + if value is INVALID_GRAMMAR_OBJ: + req.constraint.grammar = None + return GRAMMAR_FAILED + return GRAMMAR_READY + + self.grammar_queue.append(req) + return GRAMMAR_QUEUED + + def poll_ready(self) -> GrammarPollResult: + ready_uids: Set[int] = set() + failed_uids: Set[int] = set() + ready_values: dict[int, BaseGrammarObject] = {} + + deadline = time.perf_counter() + self.poll_interval + while True: + timeout = time.perf_counter() >= deadline + for req in self.grammar_queue: + uid = req.uid + if uid in ready_uids or uid in failed_uids: + continue + + assert req.constraint is not None + grammar = req.constraint.grammar + assert isinstance(grammar, futures.Future) + if grammar.done(): + value = grammar.result() + if value is INVALID_GRAMMAR_OBJ: + failed_uids.add(uid) + else: + ready_uids.add(uid) + ready_values[uid] = value + elif timeout: + req.constraint.grammar_wait_ct += 1 + if req.constraint.grammar_wait_ct >= self.max_poll_iterations: + failed_uids.add(uid) + + if timeout: + break + time.sleep(self.poll_interval / 10) + + if self.tp_size > 1: + gathered: List[Tuple[Set[int], Set[int]] | None] = [None] * self.tp_size + torch.distributed.all_gather_object( + gathered, + (ready_uids, failed_uids), + group=self.tp_cpu_group, + ) + ready_uids = set.intersection(*(x[0] for x in gathered if x is not None)) + failed_uids = set.union(*(x[1] for x in gathered if x is not None)) + + ready_uids -= failed_uids + backend = self.grammar_backend + ready_reqs: List[PendingReq] = [] + failed_list: List[int] = [] + next_queue: List[PendingReq] = [] + + for req in self.grammar_queue: + uid = req.uid + assert req.constraint is not None + key = req.constraint.grammar_key + grammar = req.constraint.grammar + + if uid in failed_uids: + if isinstance(grammar, futures.Future): + grammar.cancel() + if backend is not None and key is not None: + backend.set_cache(key, INVALID_GRAMMAR_OBJ) + req.constraint.grammar = None + failed_list.append(uid) + continue + + if uid in ready_uids: + value = ready_values[uid] + req.constraint.grammar = value + if backend is not None and key is not None: + backend.set_cache(key, value.copy()) + ready_reqs.append(req) + continue + + next_queue.append(req) + + self.grammar_queue = next_queue + return GrammarPollResult(ready_reqs=ready_reqs, failed_uids=failed_list) + + def abort_req(self, uid: int) -> bool: + for i, req in enumerate(self.grammar_queue): + if req.uid != uid: + continue + grammar = req.constraint.grammar + if isinstance(grammar, futures.Future): + grammar.cancel() + self.grammar_queue.pop(i) + return True + return False + + def shutdown(self) -> None: + for req in self.grammar_queue: + grammar = req.constraint.grammar + if isinstance(grammar, futures.Future): + grammar.cancel() + self.grammar_queue.clear() + if self.grammar_backend is not None: + self.grammar_backend.shutdown() diff --git a/python/minisgl/scheduler/prefill.py b/python/minisgl/scheduler/prefill.py index 3ffe3a56..cef3e983 100644 --- a/python/minisgl/scheduler/prefill.py +++ b/python/minisgl/scheduler/prefill.py @@ -11,7 +11,6 @@ if TYPE_CHECKING: from minisgl.kvcache import BaseCacheHandle - from minisgl.message import UserMsg from .cache import CacheManager from .decode import DecodeManager @@ -87,6 +86,7 @@ def _add_one_req( uid=pending_req.uid, cache_handle=cache_handle, sampling_params=pending_req.sampling_params, + constraint=pending_req.constraint, ) def try_add_one(self, pending_req: PendingReq) -> Req | None: @@ -120,8 +120,8 @@ class PrefillManager: decode_manager: DecodeManager pending_list: List[PendingReq] = field(default_factory=list) - def add_one_req(self, req: UserMsg) -> None: - self.pending_list.append(PendingReq(req.uid, req.input_ids, req.sampling_params)) + def enqueue(self, req: PendingReq) -> None: + self.pending_list.append(req) def schedule_next_batch(self, prefill_budget: int) -> Batch | None: if len(self.pending_list) == 0: diff --git a/python/minisgl/scheduler/scheduler.py b/python/minisgl/scheduler/scheduler.py index c0cda624..5ceaeaa7 100644 --- a/python/minisgl/scheduler/scheduler.py +++ b/python/minisgl/scheduler/scheduler.py @@ -1,6 +1,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, List, NamedTuple, NoReturn, Set, Tuple, TypeAlias +from concurrent import futures +from typing import TYPE_CHECKING, List, NamedTuple, NoReturn, Tuple, TypeAlias import torch from minisgl.core import Batch, Req @@ -18,12 +19,14 @@ from .cache import CacheManager from .config import SchedulerConfig from .decode import DecodeManager +from .grammar import GRAMMAR_FAILED, GRAMMAR_READY, GrammarManager from .io import SchedulerIOMixin from .prefill import ChunkedReq, PrefillManager from .table import TableManager +from .utils import PendingReq if TYPE_CHECKING: - from minisgl.engine import BatchSamplingArgs, ForwardOutput + from minisgl.engine import BatchSamplingArgs, ForwardOutput, ModelForwardOutput logger = init_logger(__name__) @@ -40,6 +43,7 @@ class ForwardInput(NamedTuple): ForwardData: TypeAlias = "Tuple[ForwardInput, ForwardOutput]" +PendingData: TypeAlias = "Tuple[ForwardInput, ModelForwardOutput]" class Scheduler(SchedulerIOMixin): @@ -65,12 +69,17 @@ def __init__(self, config: SchedulerConfig): ) # some alias for easy access - self.finished_reqs: Set[Req] = set() self.tokenizer = load_tokenizer(config.model_path) self.eos_token_id = self.tokenizer.eos_token_id self.token_pool = self.table_manager.token_pool self.prefill_budget = config.max_extend_tokens # self.config = config + self.grammar_manager = GrammarManager( + self.tokenizer, + self.engine.sampler.vocab_size, + self.eos_token_id, + self.engine.tp_cpu_group, + ) # Initialize the I/O mixin super().__init__(config, self.engine.tp_cpu_group) @@ -91,22 +100,37 @@ def overlap_loop(self, last_data: ForwardData | None) -> ForwardData | None: last_data is not None # don't block if we have a batch to be processed or self.prefill_manager.runnable or self.decode_manager.runnable + or self.grammar_manager.runnable ) for msg in self.receive_msg(blocking=blocking): self._process_one_msg(msg) forward_input = self._schedule_next_batch() - ongoing_data = None + ongoing_data: ForwardData | None = None + pending_data: PendingData | None = None if forward_input is not None: with self.engine_stream_ctx: # run the batch in the engine's stream self.engine.stream.wait_stream(self.stream) - ongoing_data = (forward_input, self._forward(forward_input)) + if self._should_delay_sampling(forward_input, last_data): + pending_data = (forward_input, self._forward_batch(forward_input)) + else: + ongoing_data = (forward_input, self._forward(forward_input)) self._process_last_data(last_data) + if pending_data is not None: + with self.engine_stream_ctx: + ongoing_data = ( + pending_data[0], + self._finalize_batch(pending_data[0], pending_data[1]), + ) return ongoing_data def normal_loop(self) -> None: - blocking = not (self.prefill_manager.runnable or self.decode_manager.runnable) + blocking = not ( + self.prefill_manager.runnable + or self.decode_manager.runnable + or self.grammar_manager.runnable + ) for msg in self.receive_msg(blocking=blocking): self._process_one_msg(msg) @@ -133,8 +157,28 @@ def run_forever(self) -> NoReturn: def shutdown(self) -> None: torch.cuda.synchronize(self.device) self.sync_all_ranks() + self.grammar_manager.shutdown() self.engine.shutdown() + def _accept_grammar_token(self, req: Req, next_token: int) -> bool: + grammar = req.constraint.grammar + assert grammar is not None + assert not isinstance(grammar, futures.Future) + + try: + grammar.accept_token(next_token) + except ValueError: + grammar.finished = True + raise + + finished = req.should_finish( + next_token, + self.eos_token_id, + grammar_terminated=grammar.is_terminated(), + ) + grammar.finished = finished + return finished + def _process_last_data(self, last_data: ForwardData | None) -> None: if last_data is None: return @@ -142,28 +186,39 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: batch, (_, next_tokens_cpu, copy_done) = last_data[0].batch, last_data[1] copy_done.synchronize() reply: List[DetokenizeMsg] = [] - new_finished_reqs: Set[Req] = set() with self.cache_manager.lazy_free_region(): for i, req in enumerate(batch.reqs): + if req.finished: + continue if isinstance(req, ChunkedReq): continue + next_token = next_tokens_cpu[i] - req.append_host(next_token.unsqueeze(0)) - next_token = int(next_token.item()) - finished = not req.can_decode - if not req.sampling_params.ignore_eos: - finished |= next_token == self.eos_token_id - reply.append(DetokenizeMsg(uid=req.uid, next_token=next_token, finished=finished)) - - # NOTE: overlap scheduling may make the request freed twice, skip second free - if finished and req not in self.finished_reqs: - self.decode_manager.remove_req(req) - self._free_req_resources(req) - new_finished_reqs.add(req) + next_token_id = int(next_token.item()) + reply_token_id = next_token_id + append_token = True + + if req.is_constrained: + try: + finished = self._accept_grammar_token(req, next_token_id) + except ValueError: + finished = True + reply_token_id = self.eos_token_id + append_token = False + else: + finished = req.should_finish(next_token_id, self.eos_token_id) + + if append_token: + req.append_host(next_token.unsqueeze(0)) + + reply.append( + DetokenizeMsg(uid=req.uid, next_token=reply_token_id, finished=finished) + ) + + if finished: + self._finish_req(req) elif batch.is_prefill: # for prefill, non-chunk req, cache the prefix self.cache_manager.cache_req(req, finished=False) - - self.finished_reqs = new_finished_reqs self.send_result(reply) def _process_one_msg(self, msg: BaseBackendMsg) -> None: @@ -186,17 +241,25 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None: logger.warning_rank0( f"Adjust max_tokens to {max_output_len} for request {msg.uid}." ) - self.prefill_manager.add_one_req(msg) + self._add_pending_req(PendingReq(msg.uid, msg.input_ids, msg.sampling_params)) elif isinstance(msg, AbortBackendMsg): logger.debug_rank0("Aborting request %d", msg.uid) + self.grammar_manager.abort_req(msg.uid) req_to_free = self.prefill_manager.abort_req(msg.uid) req_to_free = req_to_free or self.decode_manager.abort_req(msg.uid) if req_to_free is not None: - self._free_req_resources(req_to_free) + self._finish_req(req_to_free) else: logger.error(f"Unknown message type: {type(msg)}") raise NotImplementedError + def _finish_req(self, req: Req) -> None: + if req.finished: + return + req.finished = True + self.decode_manager.remove_req(req) + self._free_req_resources(req) + def _free_req_resources(self, req: Req) -> None: self.table_manager.free(req.table_idx) self.cache_manager.cache_req(req, finished=True) @@ -217,20 +280,81 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput: ) def _schedule_next_batch(self) -> ForwardInput | None: - # TODO: support other policies: e.g. DECODE first + self._poll_grammar_queue() batch = ( self.prefill_manager.schedule_next_batch(self.prefill_budget) or self.decode_manager.schedule_next_batch() ) return self._prepare_batch(batch) if batch else None + def _should_delay_sampling( + self, + forward_input: ForwardInput, + last_data: ForwardData | None, + ) -> bool: + if last_data is None or not forward_input.sample_args.has_grammar: + return False + + last_reqs = set(last_data[0].batch.reqs) + return any(req.is_constrained and req in last_reqs for req in forward_input.batch.reqs) + def _forward(self, forward_input: ForwardInput) -> ForwardOutput: - batch, sample_args, input_mapping, output_mapping = forward_input + return self._finalize_batch(forward_input, self._forward_batch(forward_input)) + + def _forward_batch(self, forward_input: ForwardInput) -> ModelForwardOutput: + batch, _, input_mapping, _ = forward_input batch.input_ids = self.token_pool[input_mapping] - forward_output = self.engine.forward_batch(batch, sample_args) - self.token_pool[output_mapping] = forward_output.next_tokens_gpu - self.decode_manager.filter_reqs(forward_input.batch.reqs) - return forward_output + output = self.engine.forward_batch(batch) + for req in batch.reqs: + req.commit_forward() + return output + + def _finalize_batch( + self, + forward_input: ForwardInput, + forward_output: ModelForwardOutput, + ) -> ForwardOutput: + batch, sample_args, _, output_mapping = forward_input + output = self.engine.sample_batch(batch, forward_output, sample_args) + self.token_pool[output_mapping] = output.next_tokens_gpu + self.decode_manager.filter_reqs(batch.reqs) + return output + + def _add_pending_req(self, req: PendingReq) -> None: + if not req.is_constrained: + self.prefill_manager.enqueue(req) + return + + status = self.grammar_manager.submit(req) + if status == GRAMMAR_READY: + self.prefill_manager.enqueue(req) + elif status == GRAMMAR_FAILED: + logger.warning_rank0( + "Grammar preprocessing failed for request %d. Finishing without output.", + req.uid, + ) + self.send_result( + [DetokenizeMsg(uid=req.uid, next_token=self.eos_token_id, finished=True)] + ) + + def _poll_grammar_queue(self) -> None: + if not self.grammar_manager.runnable: + return + + result = self.grammar_manager.poll_ready() + for req in result.ready_reqs: + self.prefill_manager.enqueue(req) + if result.failed_uids: + logger.warning_rank0( + "Grammar preprocessing failed for requests %s. Finishing without output.", + result.failed_uids, + ) + self.send_result( + [ + DetokenizeMsg(uid=uid, next_token=self.eos_token_id, finished=True) + for uid in result.failed_uids + ] + ) def _make_positions(batch: Batch, device: torch.device) -> torch.Tensor: diff --git a/python/minisgl/scheduler/utils.py b/python/minisgl/scheduler/utils.py index 5060a74f..5f2a8e86 100644 --- a/python/minisgl/scheduler/utils.py +++ b/python/minisgl/scheduler/utils.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, List import torch +from minisgl.core import ReqConstraintState if TYPE_CHECKING: from minisgl.core import SamplingParams @@ -16,8 +17,13 @@ class PendingReq: uid: int input_ids: torch.Tensor sampling_params: SamplingParams + constraint: ReqConstraintState | None = None chunked_req: ChunkedReq | None = None + def __post_init__(self) -> None: + if self.constraint is None and self.sampling_params.json_schema is not None: + self.constraint = ReqConstraintState() + @property def input_len(self) -> int: return len(self.input_ids) @@ -26,6 +32,10 @@ def input_len(self) -> int: def output_len(self) -> int: return self.sampling_params.max_tokens + @property + def is_constrained(self) -> bool: + return self.constraint is not None + @dataclass class ScheduleResult: diff --git a/python/minisgl/server/api_server.py b/python/minisgl/server/api_server.py index 9318ccdd..3abb2d88 100644 --- a/python/minisgl/server/api_server.py +++ b/python/minisgl/server/api_server.py @@ -51,6 +51,15 @@ def _unwrap_msg(msg: BaseFrontendMsg) -> List[UserReply]: return [msg] +def _normalize_json_schema(response_format: ResponseFormat | None) -> str | None: + if response_format is None: + return None + schema = response_format.json_schema.schema_ + if isinstance(schema, str): + return schema + return json.dumps(schema, separators=(",", ":"), sort_keys=True) + + class GenerateRequest(BaseModel): prompt: str max_tokens: int @@ -62,6 +71,18 @@ class Message(BaseModel): content: str +class JsonSchemaFormat(BaseModel): + name: str | None = None + description: str | None = None + schema_: Dict[str, object] | str = Field(alias="schema") + strict: bool | None = False + + +class ResponseFormat(BaseModel): + type: Literal["json_schema"] + json_schema: JsonSchemaFormat + + class OpenAICompletionRequest(BaseModel): """Unified request model for OpenAI-style completions and chat-completions.""" @@ -80,6 +101,7 @@ class OpenAICompletionRequest(BaseModel): stop: List[str] = [] presence_penalty: float = 0.0 frequency_penalty: float = 0.0 + response_format: ResponseFormat | None = None ignore_eos: bool = False @@ -274,6 +296,7 @@ async def v1_completions(req: OpenAICompletionRequest, request: Request): temperature=req.temperature, top_k=req.top_k, top_p=req.top_p, + json_schema=_normalize_json_schema(req.response_format), ), ) ) @@ -307,6 +330,7 @@ async def shell_completion(req: OpenAICompletionRequest): temperature=req.temperature, top_k=req.top_k, top_p=req.top_p, + json_schema=_normalize_json_schema(req.response_format), ), ) ) diff --git a/python/minisgl/utils/hf.py b/python/minisgl/utils/hf.py index 5dbabbd6..178dc22a 100644 --- a/python/minisgl/utils/hf.py +++ b/python/minisgl/utils/hf.py @@ -7,6 +7,13 @@ from tqdm.asyncio import tqdm from transformers import AutoConfig, AutoTokenizer, PretrainedConfig, PreTrainedTokenizerBase +_THINK_END_TOKENS = { + # NOTE: register more reasoning model here + "Qwen3ForCausalLM": "", + "Qwen3MoeForCausalLM": "", +} + + class DisabledTqdm(tqdm): def __init__(self, *args, **kwargs): kwargs.pop("name", None) @@ -24,6 +31,17 @@ def load_tokenizer(model_path: str) -> PreTrainedTokenizerBase: tokenizer.chat_template = json.load(f)["chat_template"] except Exception: pass + for arch in cached_load_hf_config(model_path).architectures: + if arch not in _THINK_END_TOKENS: + continue + tokenizer.think_end_token = _THINK_END_TOKENS[arch] + token_ids = tokenizer.encode(tokenizer.think_end_token, add_special_tokens=False) + if len(token_ids) != 1: + raise ValueError( + f"{tokenizer.think_end_token!r} must map to exactly one token, got {token_ids}" + ) + tokenizer.think_end_id = token_ids[0] + break return tokenizer @@ -49,4 +67,4 @@ def download_hf_weight(model_path: str) -> str: except Exception as e: raise ValueError( f"Model path '{model_path}' is neither a local directory nor a valid model ID: {e}" - ) \ No newline at end of file + )