diff --git a/benchmark/uno/README.md b/benchmark/uno/README.md new file mode 100644 index 000000000..2036235c8 --- /dev/null +++ b/benchmark/uno/README.md @@ -0,0 +1,148 @@ +# UNO full-dataset math evaluation + +`run_math_eval.py` evaluates AR, UNO, DFLASH, EAGLE, or EAGLE3 with identical +datasets, prompts, sampling parameters, and grading. It creates an in-process +`sgl.Engine`; there is no separate server process. Engine startup is excluded +from the timed interval, and no additional request warmup is run. + +The runner downloads pinned revisions of GSM8K, MATH-500, AIME 2024, AIME +2025, and AIME 2026. It applies the same boxed-answer instruction and Qwen +reasoning chat template to every engine, then grades with `math_verify`. + +Install SGLang with its evaluation dependencies: + +```bash +pip install -e "python[test]" +``` + +## Reproduce the H200 table + +Run from the SGLang repository root. Each invocation below produces one row of +the PR table. GSM8K and MATH-500 use one sample per problem; AIME 2025 uses ten +samples per problem, or 300 completions. + +```bash +export MODEL_PATH=Qwen/Qwen3-8B +export TOKENIZER_PATH=Qwen/Qwen3-8B +export UNO_LORA_PATH=s-sahoo/uno-qwen3-8B +export DATA_ROOT=/path/to/math-eval-data +export RESULT_ROOT=/path/to/math-eval-results + +COMMON_ARGS=( + --model-path "$MODEL_PATH" + --tokenizer-path "$TOKENIZER_PATH" + --data-root "$DATA_ROOT" + --context-length 40960 + --max-tokens 32768 + --temperature 1 + --top-k 50 + --top-p 0.95 + --random-seed 42 +) + +run_ar() { + local benchmark=$1 samples=$2 requests=$3 output_name=$4 + PYTHONPATH=python python -m benchmark.uno.run_math_eval \ + "${COMMON_ARGS[@]}" \ + --benchmark "$benchmark" \ + --num-samples "$samples" \ + --max-running-requests "$requests" \ + --output-dir "$RESULT_ROOT/$output_name" +} + +run_linear_uno() { + local benchmark=$1 samples=$2 requests=$3 output_name=$4 + PYTHONPATH=python python -m benchmark.uno.run_math_eval \ + "${COMMON_ARGS[@]}" \ + --benchmark "$benchmark" \ + --num-samples "$samples" \ + --max-running-requests "$requests" \ + --output-dir "$RESULT_ROOT/$output_name" \ + --speculative-algorithm UNO \ + --uno-lora-path "$UNO_LORA_PATH" \ + --speculative-num-steps 1 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 8 +} + +run_tree_uno() { + local benchmark=$1 samples=$2 requests=$3 output_name=$4 + PYTHONPATH=python python -m benchmark.uno.run_math_eval \ + "${COMMON_ARGS[@]}" \ + --benchmark "$benchmark" \ + --num-samples "$samples" \ + --max-running-requests "$requests" \ + --output-dir "$RESULT_ROOT/$output_name" \ + --speculative-algorithm UNO \ + --uno-lora-path "$UNO_LORA_PATH" \ + --speculative-num-steps 15 \ + --speculative-eagle-topk 32 \ + --speculative-num-draft-tokens 32 +} +``` + +Run the six batch-64 AR and linear `B/K/V = 8/1/8` rows: + +```bash +run_ar gsm8k 1 64 ar-gsm8k-c64 +run_linear_uno gsm8k 1 64 uno-linear-b8-k1-v8-gsm8k-c64 +run_ar math500 1 64 ar-math500-c64 +run_linear_uno math500 1 64 uno-linear-b8-k1-v8-math500-c64 +run_ar aime25 10 64 ar-aime25-c64 +run_linear_uno aime25 10 64 uno-linear-b8-k1-v8-aime25-c64 +``` + +Run the six batch-1 AR and tree `B/K/V = 16/32/32` rows: + +```bash +run_ar gsm8k 1 1 ar-gsm8k-c1 +run_tree_uno gsm8k 1 1 uno-tree-b16-k32-v32-gsm8k-c1 +run_ar math500 1 1 ar-math500-c1 +run_tree_uno math500 1 1 uno-tree-b16-k32-v32-math500-c1 +run_ar aime25 10 1 ar-aime25-c1 +run_tree_uno aime25 10 1 uno-tree-b16-k32-v32-aime25-c1 +``` + +Each output directory contains raw generations, per-answer grades, and +`summary.json` and `summary.md`. AR TPF is one. UNO TPF counts both full +target-model forwards in each cycle: the diffusion-pathway draft and +AR-pathway verification forwards. + +## Other speculative decoders + +The runner uses the same public option names as `sglang serve`. For example, +DFLASH can be evaluated with: + +```bash +PYTHONPATH=python python -m benchmark.uno.run_math_eval \ + "${COMMON_ARGS[@]}" \ + --benchmark math500 \ + --num-samples 1 \ + --output-dir "$RESULT_ROOT/dflash-b8-math500-c64" \ + --max-running-requests 64 \ + --speculative-algorithm DFLASH \ + --speculative-draft-model-path z-lab/Qwen3-8B-DFlash-b16 \ + --speculative-dflash-block-size 8 \ + --speculative-draft-attention-backend fa3 +``` + +EAGLE or EAGLE3 can be evaluated with the corresponding draft model: + +```bash +export EAGLE_DRAFT_MODEL=/path/to/compatible-eagle-draft-model + +PYTHONPATH=python python -m benchmark.uno.run_math_eval \ + "${COMMON_ARGS[@]}" \ + --benchmark math500 \ + --num-samples 1 \ + --output-dir "$RESULT_ROOT/eagle3-b8-math500-c64" \ + --max-running-requests 64 \ + --speculative-algorithm EAGLE3 \ + --speculative-draft-model-path "$EAGLE_DRAFT_MODEL" \ + --speculative-num-steps 7 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 8 +``` + +For EAGLE and DFLASH, TPF follows SGLang's acceptance-length convention and +counts generated tokens per target verification forward. diff --git a/benchmark/uno/__init__.py b/benchmark/uno/__init__.py new file mode 100644 index 000000000..f67261e49 --- /dev/null +++ b/benchmark/uno/__init__.py @@ -0,0 +1 @@ +"""Speculative-decoding benchmark utilities.""" diff --git a/benchmark/uno/math_data.py b/benchmark/uno/math_data.py new file mode 100644 index 000000000..a3eef51fc --- /dev/null +++ b/benchmark/uno/math_data.py @@ -0,0 +1,159 @@ +"""Pinned dataset preparation for speculative-decoding math evaluation.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, NamedTuple + +MATH_INSTRUCTION = "Please reason step by step and put your final answer in \\boxed{}." + + +class BenchmarkConfig(NamedTuple): + name: str + expected_rows: int + instruction: str + chat_template_kwargs: dict[str, object] + + +BENCHMARKS = { + name: BenchmarkConfig( + name=name, + expected_rows=expected_rows, + instruction=MATH_INSTRUCTION, + chat_template_kwargs={"reasoning_effort": "high"}, + ) + for name, expected_rows in ( + ("gsm8k", 1319), + ("math500", 500), + ("aime24", 30), + ("aime25", 30), + ("aime26", 30), + ) +} + +DATASET_REVISIONS = { + "openai/gsm8k": "740312add88f781978c0658806c59bc2815b9866", + "HuggingFaceH4/MATH-500": "6e4ed1a2a79af7d8630a6b768ec859cb5af4d3be", + "hypaai/Hypa_AIME2024": "11ab79f0eed5f4fdf3d469b466663ab86bbd77c8", + "math-ai/aime25": "563bb8404243c5f09de6ec262f2db674fe5bce9b", + "math-ai/aime26": "79037aebdb6580008fb960d17cb21fd3099083e3", +} + + +def get_benchmark(name: str) -> BenchmarkConfig: + try: + return BENCHMARKS[name] + except KeyError as exc: + available = ", ".join(BENCHMARKS) + raise KeyError(f"Unknown benchmark {name!r}. Available: {available}") from exc + + +def _load_dataset( + repo_id: str, + config_name: str | None = None, + *, + split: str, +) -> list[dict[str, Any]]: + from datasets import load_dataset + + dataset = load_dataset( + repo_id, + config_name, + split=split, + revision=DATASET_REVISIONS[repo_id], + ) + return [dict(row) for row in dataset] + + +def _prepare_gsm8k() -> list[dict[str, Any]]: + rows = _load_dataset("openai/gsm8k", "main", split="test") + records = [] + for index, row in enumerate(rows): + prompt = f"Q: {row['question']}\nA: Let's think step by step." + records.append( + { + "row": index, + "ground_truth": row["answer"], + "chat_input": [{"role": "user", "content": prompt}], + } + ) + return records + + +def _prepare_math500() -> list[dict[str, Any]]: + rows = _load_dataset("HuggingFaceH4/MATH-500", split="test") + return [ + { + "row": index, + "ground_truth": row["answer"], + "chat_input": [{"role": "user", "content": row["problem"]}], + } + for index, row in enumerate(rows) + ] + + +def _prepare_aime(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + records = [] + for index, row in enumerate(rows): + prompt = f"Question: {row['problem']}\nAnswer:" + records.append( + { + "row": index, + "ground_truth": str(row.get("answer", row.get("solution"))).strip(), + "chat_input": [{"role": "user", "content": prompt}], + } + ) + return records + + +def _prepare_aime24() -> list[dict[str, Any]]: + return _prepare_aime(_load_dataset("hypaai/Hypa_AIME2024", split="english")) + + +def _prepare_aime25() -> list[dict[str, Any]]: + return _prepare_aime(_load_dataset("math-ai/aime25", split="test")) + + +def _prepare_aime26() -> list[dict[str, Any]]: + return _prepare_aime(_load_dataset("math-ai/aime26", split="test")) + + +BUILDERS = { + "gsm8k": _prepare_gsm8k, + "math500": _prepare_math500, + "aime24": _prepare_aime24, + "aime25": _prepare_aime25, + "aime26": _prepare_aime26, +} + + +def _row_count(path: Path) -> int: + with path.open("rb") as handle: + return sum(bool(line.strip()) for line in handle) + + +def prepare_benchmark_data(name: str, *, output_dir: Path) -> Path: + benchmark = get_benchmark(name) + output = output_dir / f"{name}.jsonl" + if output.is_file(): + count = _row_count(output) + if count == benchmark.expected_rows: + return output + raise ValueError( + f"{output} has {count} rows; expected {benchmark.expected_rows}" + ) + + records = BUILDERS[name]() + if len(records) != benchmark.expected_rows: + raise ValueError( + f"Expected {benchmark.expected_rows} rows for {name}, got {len(records)}" + ) + + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_suffix(".jsonl.tmp") + with temporary.open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + temporary.replace(output) + return output diff --git a/benchmark/uno/math_grader.py b/benchmark/uno/math_grader.py new file mode 100644 index 000000000..2c1247d04 --- /dev/null +++ b/benchmark/uno/math_grader.py @@ -0,0 +1,310 @@ +"""Minimal math scorer adapted from Nano-vLLM-UNO's Eval360 grader.""" + +from __future__ import annotations + +import re +from typing import Any + +try: + from math_verify import parse, verify +except ModuleNotFoundError: # Allow CLI help without optional eval dependencies. + parse = None + verify = None + + +def _require_math_verify() -> None: + if parse is None or verify is None: + raise RuntimeError("math scoring requires the 'math_verify' package") + + +def extract_last_boxed_content(text: str | None) -> str | None: + if not text: + return None + matches = list(re.finditer(r"\\(?:boxed|fbox)\s*\{", text, re.IGNORECASE)) + for match in reversed(matches): + start = match.end() + depth = 1 + index = start + while index < len(text) and depth: + depth += (text[index] == "{") - (text[index] == "}") + index += 1 + if depth == 0: + return text[start : index - 1].strip().strip("$").strip() + return None + + +def _generation_list(row: dict[str, Any]) -> list[str]: + if isinstance(row.get("generations"), list): + return [str(value or "") for value in row["generations"]] + return [str(row.get("generation") or "")] + + +def _get_accuracy(correct: list[bool]) -> float: + return sum(correct) / len(correct) if correct else float("nan") + + +def _source_key(row: dict[str, Any], fallback: int) -> str: + return str(row.get("source_row", row.get("row", fallback))) + + +def normalize_answer_text(text: str | None) -> str: + if text is None: + return "" + text = str(text).strip().strip("$").strip() + gsm8k_answer = re.search(r"####\s*([^\n]+)", text) + if gsm8k_answer: + text = gsm8k_answer.group(1).strip() + text = re.sub(r"\\(?:boxed|fbox)\s*\{(.+)\}\s*$", r"\1", text) + text = text.replace("\\dfrac", "\\frac").replace("\\tfrac", "\\frac") + text = re.sub(r"\\(?:left|right|bigl|bigr|Bigl|Bigr|big|Big)", "", text) + text = text.replace("\\displaystyle", "") + text = re.sub(r"\\mathbf\s*\{([^{}]+)\}", r"\1", text) + text = re.sub(r"\\mathbf\s+([A-Za-z])", r"\1", text) + text = re.sub(r"\\\\\s*\[[^\]]+\]", r"\\\\", text) + text = re.sub(r"\\frac\s*\{([^{}]+)\}\s*\{([^{}]+)\}", r"\1/\2", text) + text = re.sub(r"\\frac\s*([+-]?\d+)\s*([+-]?\d+)", r"\1/\2", text) + text = text.replace("{,}", "") + text = re.sub(r"\\(?:,|;|!|:)", "", text) + text = text.replace("\\%", "%") + text = text.replace("\\$", "").replace("$", "") + text = re.sub(r"\\(?:text|mathrm)\s*\{([^{}]*)\}", r"\1", text) + text = re.sub(r"\^\{([^{}])\}", r"^\1", text) + text = re.sub(r"_\{([^{}])\}", r"_\1", text) + text = re.sub(r"\s+", "", text) + text = re.sub(r"(?<=\d),(?=\d{3}(?:\D|$))", "", text) + return re.sub(r"^[A-Za-z]+=(?=\\begin\{pmatrix\})", "", text) + + +def _has_plain_variable(text: str) -> bool: + text = re.sub(r"\\(?:text|mathrm)\s*\{[^{}]*\}", "", text) + text = re.sub(r"\\[A-Za-z]+", "", text) + return bool(re.search(r"[A-Za-z]", text)) + + +def _parse_boxed_content(boxed: str) -> Any: + _require_math_verify() + boxed = normalize_answer_text(boxed) + leading_number = re.match(r"\s*([-+]?[0-9][0-9,]*(?:\.[0-9]+)?)", boxed) + if leading_number and re.fullmatch(r"[A-Za-z]+", boxed[leading_number.end() :]): + answer = leading_number.group(1).replace(",", "") + try: + return parse(answer) + except Exception: + return answer + if _has_plain_variable(boxed): + return boxed + try: + parsed = parse(f"${boxed}$") or parse(boxed) + if parsed: + return parsed + except Exception: + pass + if leading_number: + answer = leading_number.group(1).replace(",", "") + try: + return parse(answer) + except Exception: + return answer + return boxed + + +def _parse_unboxed_answer(text: str) -> Any: + _require_math_verify() + try: + parsed = parse(f"${text}$") or parse(text) + if parsed: + return parsed + except Exception: + pass + + patterns = [ + r"The answer is:?\s*\$?([\-0-9\.,]+)", + r"#### ?\$?([\-0-9\.,]+)", + r"Therefore,? the answer is:?\s*\$?([\-0-9\.,]+)", + r"So,? the answer is:?\s*\$?([\-0-9\.,]+)", + r"Thus,? the answer is:?\s*\$?([\-0-9\.,]+)", + r"Hence,? the answer is:?\s*\$?([\-0-9\.,]+)", + r"Final answer:?\s*\$?([\-0-9\.,]+)", + r"The final answer is:?\s*\$?([\-0-9\.,]+)", + r"The answer is:?\s*\$?([\-0-9\.,]+)\s*(?:miles?|minutes?|hours?|dollars?|GB)?", + ( + r"=\s*\$?([\-0-9\.,]+)" + r"\s*(?:miles?|minutes?|hours?|dollars?|GB)?" + r"\.?\s*(?:The answer|$)" + ), + ] + for pattern in patterns: + matches = re.findall(pattern, text, re.IGNORECASE) + if matches: + answer = matches[-1].replace(",", "").strip().rstrip(".") + try: + return parse(answer) + except Exception: + return answer + + sentence_end = ( + r"(?:is|are|equals?|makes?|has|have|gets?|arrives?|covers?|travels?)" + r"\s+\$?([\-0-9\.,]+)(?:\s*(?:miles?|minutes?|hours?|dollars?|GB))?" + r"\.?\s*$" + ) + match = re.search(sentence_end, text, re.MULTILINE | re.IGNORECASE) + if match: + answer = match.group(1).replace(",", "").strip().rstrip(".") + try: + return parse(answer) + except Exception: + return answer + + for sentence in reversed(text.split(".")): + if "Human:" in sentence or "Assistant:" in sentence: + continue + numbers = re.findall(r"[-+]?[0-9]*\.?[0-9]+", sentence) + if numbers: + answer = numbers[-1].lstrip("0") or "0" + try: + return parse(answer) + except Exception: + return answer + return None + + +def _parse_answer(text: str) -> Any: + boxed = extract_last_boxed_content(text) + return _parse_boxed_content(boxed) if boxed else _parse_unboxed_answer(text) + + +def _answer_candidates(text: str) -> list[Any]: + boxed = extract_last_boxed_content(text) + candidates = ( + [_parse_boxed_content(boxed), normalize_answer_text(boxed), boxed] + if boxed + else [_parse_unboxed_answer(text)] + ) + unique = [] + seen = set() + for candidate in candidates: + if candidate is not None and repr(candidate) not in seen: + seen.add(repr(candidate)) + unique.append(candidate) + return unique + + +def _vector_components(text: str) -> tuple[str, ...] | None: + matrix = re.fullmatch(r"\\begin\{pmatrix\}(.+)\\end\{pmatrix\}", text) + if matrix: + return tuple(part for part in matrix.group(1).split(r"\\") if part) + tuple_match = re.fullmatch(r"\(([^()]+)\)", text) + if tuple_match and "," in tuple_match.group(1): + return tuple(part for part in tuple_match.group(1).split(",") if part) + return None + + +def _text_answers_match(answer: str, gold: str) -> bool: + answer_norm = normalize_answer_text(answer) + gold_norm = normalize_answer_text(gold) + if answer_norm == gold_norm: + return True + try: + if float(answer_norm.rstrip("%")) == float(gold_norm.rstrip("%")): + return True + except ValueError: + pass + if re.fullmatch(r"\([A-Za-z]\)", gold_norm) and answer_norm == gold_norm[1:-1]: + return True + answer_parts = _vector_components(answer_norm) + gold_parts = _vector_components(gold_norm) + return bool(answer_parts and answer_parts == gold_parts) + + +def _compare_answers(answer: Any, gold: str | None) -> bool: + _require_math_verify() + if not answer or gold is None: + return False + if isinstance(answer, str) and _text_answers_match(answer, gold): + return True + try: + if verify(answer, gold): + return True + gold_answer = _parse_answer(gold) + if gold_answer: + return bool(verify(gold_answer, answer)) + except Exception: + return isinstance(answer, str) and _text_answers_match(answer, gold) + + +def grade_math_row(row: dict[str, Any]) -> dict[str, Any]: + expected = row.get("ground_truth") + expected_values = expected if isinstance(expected, list) else [expected] + correct = [] + parsed_generations = [] + for generation in _generation_list(row): + candidates = _answer_candidates(generation) + parsed_generations.append([str(candidate) for candidate in candidates]) + correct.append( + any( + _compare_answers(candidate, str(gold)) + for candidate in candidates + for gold in expected_values + if gold is not None + ) + ) + return { + **row, + "parsed_generations": parsed_generations, + "correct": correct, + "accuracy": _get_accuracy(correct), + "grader": "math", + } + + +def score_math( + rows: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + graded = [grade_math_row(row) for row in rows] + correct = [bool(value) for row in graded for value in row["correct"]] + by_source: dict[str, list[tuple[int, bool]]] = {} + for index, row in enumerate(graded): + try: + sample_index = int(row.get("sample_index", 0)) + except (TypeError, ValueError): + sample_index = 0 + samples = by_source.setdefault(_source_key(row, index), []) + samples.extend( + (sample_index + offset, bool(value)) + for offset, value in enumerate(row["correct"]) + ) + + per_problem = [ + [value for _, value in sorted(samples)] + for samples in by_source.values() + if samples + ] + samples_per_problem = max((len(samples) for samples in per_problem), default=0) + sample0 = [samples[0] for samples in per_problem] + summary = { + "grader": "math", + "num_rows": len(graded), + "num_problems": len(per_problem), + "samples_per_problem": samples_per_problem, + "num_correct": sum(correct), + "accuracy": _get_accuracy(correct), + } + if per_problem: + summary.update( + avg_at_1=_get_accuracy(sample0), + pass_at_1=_get_accuracy(sample0), + num_correct_at_1=sum(sample0), + ) + if samples_per_problem > 1: + summary.update( + { + f"avg_at_{samples_per_problem}": _get_accuracy(correct), + f"pass_at_{samples_per_problem}": _get_accuracy( + [any(samples) for samples in per_problem] + ), + f"num_pass_at_{samples_per_problem}": sum( + any(samples) for samples in per_problem + ), + } + ) + return graded, summary diff --git a/benchmark/uno/run_math_eval.py b/benchmark/uno/run_math_eval.py new file mode 100644 index 000000000..1827b5c2f --- /dev/null +++ b/benchmark/uno/run_math_eval.py @@ -0,0 +1,517 @@ +"""Run full-dataset AR or speculative-decoding math evaluation offline.""" + +from __future__ import annotations + +import argparse +import json +from collections import defaultdict +from pathlib import Path +from time import perf_counter +from typing import Any + +from benchmark.uno.math_data import ( + BENCHMARKS, + BenchmarkConfig, + get_benchmark, + prepare_benchmark_data, +) +from benchmark.uno.math_grader import score_math + +CONTEXT_LENGTH = 40960 +MAX_TOKENS = 2**15 +SPECULATIVE_ALGORITHMS = ("EAGLE", "EAGLE3", "DFLASH", "UNO") +SPECULATIVE_OPTION_NAMES = ( + "speculative_algorithm", + "speculative_draft_model_path", + "speculative_draft_model_revision", + "speculative_num_steps", + "speculative_eagle_topk", + "speculative_num_draft_tokens", + "speculative_dflash_block_size", + "speculative_draft_attention_backend", + "uno_lora_path", +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-path", default="Qwen/Qwen3-8B") + parser.add_argument("--tokenizer-path") + parser.add_argument("--revision") + parser.add_argument("--dtype", default="bfloat16") + parser.add_argument("--attention-backend", default="fa3") + parser.add_argument("--data-root", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument( + "--speculative-algorithm", + type=str.upper, + choices=SPECULATIVE_ALGORITHMS, + ) + parser.add_argument("--speculative-draft-model-path") + parser.add_argument("--speculative-draft-model-revision") + parser.add_argument("--speculative-num-steps", type=int) + parser.add_argument("--speculative-eagle-topk", type=int) + parser.add_argument("--speculative-num-draft-tokens", type=int) + parser.add_argument("--speculative-dflash-block-size", type=int) + parser.add_argument("--speculative-draft-attention-backend") + parser.add_argument("--uno-lora-path") + parser.add_argument( + "--benchmark", + action="append", + choices=tuple(BENCHMARKS), + help="Benchmark to run; repeat to select multiple (default: all).", + ) + parser.add_argument("--limit", type=int, help="Maximum problems per benchmark.") + parser.add_argument("--num-samples", type=int, default=1) + parser.add_argument("--max-running-requests", type=int, default=4) + parser.add_argument("--context-length", type=int, default=CONTEXT_LENGTH) + parser.add_argument("--max-tokens", type=int, default=MAX_TOKENS) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--top-k", type=int, default=50) + parser.add_argument("--top-p", type=float, default=0.95) + parser.add_argument("--random-seed", type=int, default=42) + args = parser.parse_args() + _validate_args(parser=parser, args=args) + return args + + +def _validate_args( + *, parser: argparse.ArgumentParser, args: argparse.Namespace +) -> None: + for name in ( + "num_samples", + "max_running_requests", + "context_length", + "max_tokens", + "top_k", + ): + if getattr(args, name) <= 0: + parser.error(f"--{name.replace('_', '-')} must be positive") + if args.limit is not None and args.limit <= 0: + parser.error("--limit must be positive") + if not 0 < args.top_p <= 1: + parser.error("--top-p must be in (0, 1]") + if args.temperature < 0: + parser.error("--temperature must be non-negative") + + +def _context_reserve(args: argparse.Namespace) -> int: + width = args.speculative_num_draft_tokens + if width is None: + width = args.speculative_dflash_block_size or 1 + if args.speculative_algorithm == "DFLASH": + return 2 * width + if args.speculative_algorithm == "UNO" and (args.speculative_eagle_topk or 1) > 1: + draft_width = (args.speculative_num_steps or 0) + 1 + return max(width, draft_width) + 1 + return width + + +def _model_ref(value: str) -> str: + path = Path(value).expanduser() + return str(path.resolve()) if path.exists() else value + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + with path.open(encoding="utf-8") as handle: + return [json.loads(line) for line in handle if line.strip()] + + +def _first_value(row: dict[str, Any], *names: str) -> Any | None: + return next((row[name] for name in names if row.get(name) is not None), None) + + +def _format_prompt( + *, + tokenizer: Any, + messages: list[dict[str, Any]], + benchmark: BenchmarkConfig, +) -> tuple[list[int], str]: + messages = [dict(message) for message in messages] + system_message = next( + (message for message in messages if message.get("role") == "system"), + None, + ) + if system_message is None: + messages.insert(0, {"role": "system", "content": benchmark.instruction}) + elif benchmark.instruction not in system_message["content"]: + existing = system_message["content"].strip() + system_message["content"] = f"{benchmark.instruction}\n\n{existing}" + + rendered = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + **benchmark.chat_template_kwargs, + ) + token_ids = tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_dict=False, + **benchmark.chat_template_kwargs, + ) + return list(token_ids), str(rendered) + + +def _prepare_prompts( + *, + args: argparse.Namespace, + tokenizer: Any, + context_reserve: int, +) -> list[dict[str, Any]]: + prompts = [] + for benchmark_name in args.benchmark or BENCHMARKS: + benchmark = get_benchmark(benchmark_name) + data_path = prepare_benchmark_data( + benchmark.name, + output_dir=args.data_root, + ) + rows = _read_jsonl(data_path) + rows = rows[: args.limit] if args.limit is not None else rows + for row_index, row in enumerate(rows): + input_ids, rendered = _format_prompt( + tokenizer=tokenizer, + messages=row["chat_input"], + benchmark=benchmark, + ) + available = args.context_length - len(input_ids) - context_reserve + if available < args.max_tokens: + source = _first_value(row, "id", "problem_id", "index", "row") + source = row_index if source is None else source + raise ValueError( + f"{benchmark.name}:{source} has only {max(0, available)} " + "completion tokens available; increase --context-length" + ) + source = _first_value(row, "id", "problem_id", "index", "row") + source = row_index if source is None else source + for sample_index in range(args.num_samples): + prompt_id = f"{source}:sample{sample_index}" + prompts.append( + { + "id": prompt_id, + "benchmark": benchmark.name, + "input_ids": input_ids, + "row": { + "id": prompt_id, + "source_row": source, + "sample_index": sample_index, + "problem": rendered, + "ground_truth": row["ground_truth"], + "prompt_token_count": len(input_ids), + "resolved_max_tokens": args.max_tokens, + }, + } + ) + if not prompts: + raise ValueError("No prompts were prepared") + return prompts + + +def _engine_options( + *, + args: argparse.Namespace, +) -> dict[str, Any]: + tokenizer_path = _model_ref(args.tokenizer_path or args.model_path) + options: dict[str, Any] = { + "model_path": _model_ref(args.model_path), + "tokenizer_path": tokenizer_path, + "skip_tokenizer_init": True, + "context_length": args.context_length, + "dtype": args.dtype, + "random_seed": args.random_seed, + "max_running_requests": args.max_running_requests, + "attention_backend": args.attention_backend, + "log_level": "info", + } + if args.revision is not None: + options["revision"] = args.revision + for name in SPECULATIVE_OPTION_NAMES: + value = getattr(args, name) + if value is not None: + if name in ( + "speculative_draft_model_path", + "uno_lora_path", + ): + value = _model_ref(value) + options[name] = value + return options + + +def _generate( + *, + args: argparse.Namespace, + prompts: list[dict[str, Any]], + engine_options: dict[str, Any], +) -> tuple[list[dict[str, Any]], float]: + import sglang as sgl + + engine = sgl.Engine(**engine_options) + sampling = [ + { + "max_new_tokens": args.max_tokens, + "temperature": args.temperature, + "top_k": args.top_k, + "top_p": args.top_p, + } + for _ in prompts + ] + try: + start = perf_counter() + outputs = engine.generate( + input_ids=[prompt["input_ids"] for prompt in prompts], + sampling_params=sampling, + rid=[f"{prompt['benchmark']}:{prompt['id']}" for prompt in prompts], + ) + elapsed = perf_counter() - start + finally: + engine.shutdown() + return outputs, elapsed + + +def _build_rows( + *, + prompts: list[dict[str, Any]], + outputs: list[dict[str, Any]], + tokenizer: Any, + speculative_algorithm: str | None, +) -> list[dict[str, Any]]: + rows = [] + for prompt, output in zip(prompts, outputs, strict=True): + token_ids = list(output["output_ids"]) + metadata = output["meta_info"] + verify_forwards = int(metadata.get("spec_verify_ct", 0)) + if speculative_algorithm == "UNO": + # Both UNO pathways are full target-model forward passes. + num_forwards = 2 * verify_forwards + elif speculative_algorithm is not None: + # EAGLE and DFLASH TPF conventionally count target verification. + num_forwards = verify_forwards + else: + num_forwards = len(token_ids) + rows.append( + { + "benchmark": prompt["benchmark"], + **prompt["row"], + "output_ids": token_ids, + "generation": tokenizer.decode(token_ids, skip_special_tokens=True), + "num_tokens": len(token_ids), + "num_forwards": num_forwards, + "tokens_per_forward": _divide(len(token_ids), num_forwards), + "sglang_meta_info": metadata, + } + ) + return rows + + +def _divide(numerator: int | float, denominator: int | float) -> float | None: + return numerator / denominator if denominator else None + + +def _write_json(path: Path, value: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + for row in rows: + handle.write(json.dumps(row, ensure_ascii=False) + "\n") + + +def _benchmark_summary(rows: list[dict[str, Any]]) -> dict[str, Any]: + tokens = sum(row["num_tokens"] for row in rows) + forwards = sum(row["num_forwards"] for row in rows) + tpfs = [ + row["tokens_per_forward"] + for row in rows + if row["tokens_per_forward"] is not None + ] + return { + "num_tokens": tokens, + "num_forwards": forwards, + "tokens_per_forward": _divide(tokens, forwards), + "unweighted_mean_tokens_per_forward": (sum(tpfs) / len(tpfs) if tpfs else None), + } + + +def _write_markdown(path: Path, summary: dict[str, Any]) -> None: + lines = [ + "| Dataset | Accuracy | TPF | tok/s | tok/s/request |", + "| --- | ---: | ---: | ---: | ---: |", + ] + for name, metrics in summary["by_benchmark"].items(): + tokens_per_second = metrics.get("tokens_per_second") + per_request = metrics.get("tokens_per_second_per_request") + tokens_per_forward = metrics.get("tokens_per_forward") + tps = f"{tokens_per_second:.2f}" if tokens_per_second is not None else "—" + request_tps = f"{per_request:.2f}" if per_request is not None else "—" + tpf = f"{tokens_per_forward:.3f}" if tokens_per_forward is not None else "—" + lines.append( + f"| {name} | {metrics['accuracy']:.2%} | {tpf} | {tps} | {request_tps} |" + ) + lines.extend( + [ + "| **Average** | " + f"**{summary['macro_average_accuracy']:.2%}** | " + f"**{summary['macro_average_tokens_per_forward']:.3f}** | " + f"**{summary['tokens_per_second']:.2f}** | " + f"**{summary['tokens_per_second_per_request']:.2f}** |", + "", + "Accuracy and TPF in the Average row are unweighted means across " + "datasets. Throughput is aggregate output tokens divided by timed " + "generation seconds; tok/s/request divides it by max running requests.", + ] + ) + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def _write_results( + *, + args: argparse.Namespace, + rows: list[dict[str, Any]], + generation_seconds: float, + engine_options: dict[str, Any], +) -> None: + rows_by_benchmark: dict[str, list[dict[str, Any]]] = defaultdict(list) + for row in rows: + output_row = {key: value for key, value in row.items() if key != "benchmark"} + rows_by_benchmark[row["benchmark"]].append(output_row) + + benchmark_metrics = {} + totals = {"prompts": 0, "completions": 0, "correct": 0, "tokens": 0, "forwards": 0} + request_tpfs = [] + for name, benchmark_rows in rows_by_benchmark.items(): + output_dir = args.output_dir / name + _write_jsonl(output_dir / "generations.jsonl", benchmark_rows) + graded, scores = score_math(benchmark_rows) + _write_jsonl(output_dir / "grades.jsonl", graded) + _write_json(output_dir / "scores.json", scores) + + metrics = { + "num_prompts": scores["num_problems"], + "num_completions": scores["num_rows"], + "num_correct": scores["num_correct"], + "accuracy": scores["accuracy"], + **_benchmark_summary(benchmark_rows), + } + benchmark_metrics[name] = metrics + totals["prompts"] += scores["num_problems"] + totals["completions"] += scores["num_rows"] + totals["correct"] += scores["num_correct"] + totals["tokens"] += metrics["num_tokens"] + totals["forwards"] += metrics["num_forwards"] + request_tpfs.extend( + row["tokens_per_forward"] + for row in benchmark_rows + if row["tokens_per_forward"] is not None + ) + + tokens_per_second = totals["tokens"] / generation_seconds + tokens_per_second_per_request = tokens_per_second / args.max_running_requests + if len(benchmark_metrics) == 1: + only_metrics = next(iter(benchmark_metrics.values())) + only_metrics["tokens_per_second"] = tokens_per_second + only_metrics["tokens_per_second_per_request"] = tokens_per_second_per_request + + dataset_accuracies = [value["accuracy"] for value in benchmark_metrics.values()] + dataset_tpfs = [ + value["tokens_per_forward"] + for value in benchmark_metrics.values() + if value["tokens_per_forward"] is not None + ] + summary = { + "engine": "sglang-offline", + "mode": _mode_name(args), + "model_path": engine_options["model_path"], + "tokenizer_path": engine_options["tokenizer_path"], + "revision": args.revision, + "dtype": args.dtype, + "attention_backend": args.attention_backend, + "speculative_parameters": { + name: engine_options[name] + for name in SPECULATIVE_OPTION_NAMES + if name in engine_options + }, + "max_running_requests": args.max_running_requests, + "benchmarks": list(rows_by_benchmark), + "context_length": args.context_length, + "max_tokens": args.max_tokens, + "num_samples": args.num_samples, + "random_seed": args.random_seed, + "sampling": { + "temperature": args.temperature, + "top_k": args.top_k, + "top_p": args.top_p, + }, + "num_prompts": totals["prompts"], + "num_completions": totals["completions"], + "num_correct": totals["correct"], + "accuracy": _divide(totals["correct"], totals["completions"]), + "macro_average_accuracy": sum(dataset_accuracies) / len(dataset_accuracies), + "generation_seconds": generation_seconds, + "num_tokens": totals["tokens"], + "tokens_per_second": tokens_per_second, + "tokens_per_second_per_request": tokens_per_second_per_request, + "num_forwards": totals["forwards"], + "tokens_per_forward": _divide(totals["tokens"], totals["forwards"]), + "macro_average_tokens_per_forward": sum(dataset_tpfs) / len(dataset_tpfs), + "unweighted_mean_tokens_per_forward": ( + sum(request_tpfs) / len(request_tpfs) if request_tpfs else None + ), + "by_benchmark": benchmark_metrics, + } + _write_json(args.output_dir / "summary.json", summary) + _write_markdown(args.output_dir / "summary.md", summary) + print(json.dumps(summary, indent=2)) + + +def _mode_name(args: argparse.Namespace) -> str: + algorithm = args.speculative_algorithm + if algorithm is None: + return "ar" + if algorithm == "UNO": + return "uno-tree" if (args.speculative_eagle_topk or 1) > 1 else "uno-linear" + return algorithm.lower() + + +def main() -> None: + args = parse_args() + from transformers import AutoTokenizer + + engine_options = _engine_options(args=args) + tokenizer = AutoTokenizer.from_pretrained( + engine_options["tokenizer_path"], + use_fast=True, + trust_remote_code=True, + ) + prompts = _prepare_prompts( + args=args, + tokenizer=tokenizer, + context_reserve=_context_reserve(args), + ) + outputs, elapsed = _generate( + args=args, + prompts=prompts, + engine_options=engine_options, + ) + rows = _build_rows( + prompts=prompts, + outputs=outputs, + tokenizer=tokenizer, + speculative_algorithm=args.speculative_algorithm, + ) + _write_results( + args=args, + rows=rows, + generation_seconds=elapsed, + engine_options=engine_options, + ) + + +if __name__ == "__main__": + main() diff --git a/docs/docs/advanced_features/speculative_decoding.mdx b/docs/docs/advanced_features/speculative_decoding.mdx index 084a2260e..b0a058a6f 100644 --- a/docs/docs/advanced_features/speculative_decoding.mdx +++ b/docs/docs/advanced_features/speculative_decoding.mdx @@ -1,9 +1,9 @@ --- title: "Speculative Decoding" metatags: - description: "SGLang speculative decoding: EAGLE-2/EAGLE-3, MTP, DFLASH, draft model configuration, and overlap-scheduler guidance." + description: "SGLang speculative decoding: EAGLE-2/EAGLE-3, MTP, UNO, DFLASH, draft model configuration, and overlap-scheduler guidance." --- -SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3, MTP, DFLASH, classic draft-model decoding, and an NGRAM-based variant. Our implementation aims to maximize speed and efficiency and is considered to be among the fastest in open-source LLM engines. +SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3, MTP, UNO, DFLASH, classic draft-model decoding, and an NGRAM-based variant. Our implementation aims to maximize speed and efficiency and is considered to be among the fastest in open-source LLM engines. ## Summary @@ -15,6 +15,7 @@ SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3, - [EAGLE-2 Decoding via Frequency-Ranked Speculative Sampling](#eagle-2-decoding-via-frequency-ranked-speculative-sampling) - [EAGLE-3 Decoding](#eagle-3-decoding) - [Multi Token Prediction](#multi-token-prediction) +- [UNO decoding](#uno-decoding) - [DFlash Decoding](#dflash-decoding) - [Standalone Speculative Decoding (Small Draft Model)](#standalone-speculative-decoding-small-draft-model) - [Speculative Decoding V2 (Overlap Scheduler)](#speculative-decoding-v2-overlap-scheduler) @@ -30,6 +31,7 @@ SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3, - **Workload acceptance changes over time**: Use [**Adaptive speculative decoding**](./adaptive_speculative_decoding) on top of **EAGLE** with `--speculative-eagle-topk 1`. - **Lower `lm_head` overhead for EAGLE-2**: Enable **FR-Spec** with `--speculative-token-map`. - **Model is MTP-enabled**: Use **MTP via speculative decoding** (often with small `speculative_num_steps/topk/num_draft_tokens`, see the example section). +- **You have a UNO adapter for the target model**: Use **UNO** with `--speculative-algorithm UNO` and `--uno-lora-path ...`. - **You have a DFlash draft checkpoint**: Use **DFLASH** with `--speculative-algorithm DFLASH` and `--speculative-draft-model-path ...`. - **You have a smaller draft LLM**: Use **STANDALONE** (`--speculative-algorithm STANDALONE`). - **No extra model available**: Use **NGRAM** (`--speculative-algorithm NGRAM`, CUDA-only). @@ -86,6 +88,13 @@ SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3, See Multi Token Prediction section Uses speculative workflow; draft path may be auto-handled for some models + + UNO + Target model with a trained UNO LoRA (linear chain or tree) + No + --speculative-algorithm UNO + --uno-lora-path ... + CUDA and FA3; currently requires TP=PP=1 + DFLASH DFlash draft model (linear block verification) @@ -434,6 +443,88 @@ print(response.json()) --- +## UNO decoding + +UNO reuses the target transformer for both passes of each speculative decode cycle instead of loading a separate draft model. During the draft forward, the first row of each request uses the base weights and the remaining `B - 1` rows use a trained UNO LoRA. Target verification is entirely base-only. + +UNO provides two sampling modes: + +- **Linear** constructs and verifies one chain of draft tokens, similar to DFlash's proposal layout. +- **Tree** expands multiple candidates at each draft depth and uses SGLang's EAGLE tree-verification path. + +For UNO modes, the three numbers are `B/K/V`: draft-forward width, candidates kept per expansion, and verification width. Thus, `Linear 8/1/8` uses an 8-token draft width, one candidate per expansion, and an 8-token verification width. UNO's `K` controls proposal-tree breadth and is unrelated to the request sampling parameter `top_k`. The command-line mapping depends on the mode: + +- In linear mode, `B` and `V` are both `--speculative-num-draft-tokens`; set `--speculative-num-steps 1` and `--speculative-eagle-topk 1`. +- In tree mode, `B` is `--speculative-num-steps + 1`, `K` is `--speculative-eagle-topk`, and `V` is `--speculative-num-draft-tokens`. + +Use an adapter trained for the exact target-model checkpoint. The current integration has been validated with `Qwen/Qwen3-8B`; unsupported LoRA target layers are rejected at startup. Set the adapter's local path or Hugging Face repository before starting the server: + +### Linear 8/1/8 + +```bash Command +export UNO_LORA_PATH="s-sahoo/uno-qwen3-8B" + +sglang serve \ + --model-path Qwen/Qwen3-8B \ + --speculative-algorithm UNO \ + --uno-lora-path "$UNO_LORA_PATH" \ + --speculative-num-steps 1 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 8 \ + --attention-backend fa3 \ + --tp 1 \ + --host 0.0.0.0 \ + --port 30000 +``` + +### Tree 8/32/8 + +The following command selects tree mode because `K` is greater than one: + +```bash Command +export UNO_LORA_PATH="s-sahoo/uno-qwen3-8B" + +sglang serve \ + --model-path Qwen/Qwen3-8B \ + --speculative-algorithm UNO \ + --uno-lora-path "$UNO_LORA_PATH" \ + --speculative-num-steps 7 \ + --speculative-eagle-topk 32 \ + --speculative-num-draft-tokens 8 \ + --attention-backend fa3 \ + --tp 1 \ + --host 0.0.0.0 \ + --port 30000 +``` + +Send requests through the standard OpenAI-compatible API: + +```python Example +import openai + +client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="None") + +response = client.chat.completions.create( + model="Qwen/Qwen3-8B", + messages=[{"role": "user", "content": "Explain speculative decoding."}], + max_tokens=128, +) + +print(response.choices[0].message.content) +``` + +### Key requirements and limitations + +- UNO requires CUDA with FA3 for prefill and decode, tensor and pipeline parallel sizes of 1, and no DP attention or context parallelism. +- `--uno-lora-path` loads UNO's fixed internal adapter and cannot be combined with request-selectable Multi-LoRA serving. +- UNO manages its own stochastic verification; do not set `--speculative-use-rejection-sampling`. +- Grammar decoding, returned logprobs or hidden states, sampling penalties, `min_p`, logit bias, custom logit processors, strict thinking, and deterministic inference are not yet supported. +- In tree mode, both speculative acceptance thresholds must remain at 1.0, `V` must be at least `B` and at most 128, and `V * K` must not exceed 2048. SGLang validates the remaining tree-capacity and EAGLE parent-representation constraints at startup. +- Mixed chunked prefill is disabled for UNO. +- Ordinary overlap scheduling is supported. Tree mode does not yet support PDMux or the separate `--enable-two-batch-overlap` feature. + +--- + ## DFlash Decoding SGLang also supports **DFLASH** speculative decoding using a dedicated draft model checkpoint. Compared with EAGLE-style tree verification, DFLASH verifies a linear draft block and is configured around a block size / draft window. This path is useful when the target model has a matching DFlash draft checkpoint, such as `meta-llama/Llama-3.1-8B-Instruct` with `z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat`. @@ -753,7 +844,7 @@ Below is a comprehensive list of all speculative decoding parameters available i --speculative-algorithm str None - Algorithm to use: DFLASH, EAGLE, EAGLE3, STANDALONE, NGRAM, NEXTN (alias of EAGLE) + Algorithm to use: UNO, DFLASH, EAGLE, EAGLE3, STANDALONE, NGRAM, NEXTN (alias of EAGLE) --speculative-draft-model-path @@ -761,6 +852,12 @@ Below is a comprehensive list of all speculative decoding parameters available i None Path to the draft model weights + + --uno-lora-path + str + None + UNO-only path or Hugging Face repository for the draft LoRA checkpoint + --speculative-draft-model-revision str @@ -777,19 +874,19 @@ Below is a comprehensive list of all speculative decoding parameters available i --speculative-num-steps int None (auto-chosen when omitted) - Autoregressive drafting depth + Autoregressive drafting depth. Some algorithms auto-tune it; tree UNO requires it and uses B = speculative_num_steps + 1. --speculative-eagle-topk int None (auto-chosen when omitted) - Branching factor per drafting step + Branching factor per drafting step. Some algorithms auto-tune it; UNO resolves an omitted value to 1, and values greater than one select UNO tree mode. --speculative-num-draft-tokens int None (auto-chosen when omitted) - Maximum number of draft tokens for verification + Maximum number of draft tokens for verification. Some algorithms auto-tune it; UNO requires it and uses it as linear width or tree verification width V. --speculative-dflash-block-size diff --git a/python/sglang/kernels/ops/attention/__init__.py b/python/sglang/kernels/ops/attention/__init__.py index b7516f92d..6350324fc 100644 --- a/python/sglang/kernels/ops/attention/__init__.py +++ b/python/sglang/kernels/ops/attention/__init__.py @@ -16,6 +16,7 @@ _TRITON_KERNELS = [ ("extend_attention", "build_unified_kv_indices"), ("prefill_attention", "context_attention_fwd"), ("merge_state", "merge_state_triton"), + ("suffix_attention_merge", "merge_suffix_attention_in_place"), ("metadata", "get_num_kv_splits_triton"), ("metadata", "prepare_swa_spec_page_table_triton"), ("metadata", "normal_decode_set_metadata"), diff --git a/python/sglang/kernels/ops/attention/suffix_attention_merge.py b/python/sglang/kernels/ops/attention/suffix_attention_merge.py new file mode 100644 index 000000000..2b7a58596 --- /dev/null +++ b/python/sglang/kernels/ops/attention/suffix_attention_merge.py @@ -0,0 +1,203 @@ +"""Fused attention over a short sparse suffix and merge with a prefix state.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +def can_use_fused_suffix_attention_merge( + *, + layer, + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + extra_kwargs: dict, +) -> bool: + """Whether attention can use the specialized suffix merge.""" + return bool( + q.dtype in (torch.float16, torch.bfloat16) + and key_cache.dtype == q.dtype + and value_cache.dtype == q.dtype + and layer.head_dim == layer.v_head_dim + and not layer.is_cross_attention + and not layer.logit_cap + and not extra_kwargs + ) + + +@triton.jit +def _fused_suffix_attention_merge_kernel( + q_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + suffix_seqlens_ptr, + prefix_ptr, + prefix_lse_ptr, + scale, + q_stride_t, + q_stride_h, + q_stride_d, + k_stride_t, + k_stride_h, + k_stride_d, + v_stride_t, + v_stride_h, + v_stride_d, + page_stride_t, + page_stride_s, + prefix_stride_t, + prefix_stride_h, + prefix_stride_d, + prefix_lse_stride_h, + prefix_lse_stride_t, + NUM_Q_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_SUFFIX: tl.constexpr, + BLOCK_D: tl.constexpr, +): + token = tl.program_id(0) + q_head = tl.program_id(1) + kv_head = q_head // (NUM_Q_HEADS // NUM_KV_HEADS) + + suffix_offsets = tl.arange(0, BLOCK_SUFFIX) + suffix_length = tl.load(suffix_seqlens_ptr + token) + suffix_valid = suffix_offsets < suffix_length + slots = tl.load( + page_table_ptr + token * page_stride_t + suffix_offsets * page_stride_s, + mask=suffix_valid, + other=0, + ).to(tl.int64) + + dims = tl.arange(0, BLOCK_D) + dim_valid = dims < HEAD_DIM + q = tl.load( + q_ptr + token * q_stride_t + q_head * q_stride_h + dims * q_stride_d, + mask=dim_valid, + other=0.0, + ).to(tl.float32) + k = tl.load( + k_cache_ptr + + slots[:, None] * k_stride_t + + kv_head * k_stride_h + + dims[None, :] * k_stride_d, + mask=suffix_valid[:, None] & dim_valid[None, :], + other=0.0, + ).to(tl.float32) + scores = tl.sum(k * q[None, :], axis=1) * scale + scores = tl.where(suffix_valid, scores, -float("inf")) + suffix_max = tl.max(scores, axis=0) + + prefix_lse = tl.load( + prefix_lse_ptr + q_head * prefix_lse_stride_h + token * prefix_lse_stride_t + ).to(tl.float32) + global_max = tl.maximum(prefix_lse, suffix_max) + prefix_weight = tl.exp(prefix_lse - global_max) + suffix_weights = tl.exp(scores - global_max) + denominator = prefix_weight + tl.sum(suffix_weights, axis=0) + + v = tl.load( + v_cache_ptr + + slots[:, None] * v_stride_t + + kv_head * v_stride_h + + dims[None, :] * v_stride_d, + mask=suffix_valid[:, None] & dim_valid[None, :], + other=0.0, + ).to(tl.float32) + suffix_numerator = tl.sum(suffix_weights[:, None] * v, axis=0) + prefix = tl.load( + prefix_ptr + + token * prefix_stride_t + + q_head * prefix_stride_h + + dims * prefix_stride_d, + mask=dim_valid, + other=0.0, + ).to(tl.float32) + output = (prefix * prefix_weight + suffix_numerator) / denominator + tl.store( + prefix_ptr + + token * prefix_stride_t + + q_head * prefix_stride_h + + dims * prefix_stride_d, + output, + mask=dim_valid, + ) + + +def merge_suffix_attention_in_place( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + suffix_page_table: torch.Tensor, + suffix_cache_seqlens: torch.Tensor, + prefix: torch.Tensor, + prefix_lse: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Compute a short sparse suffix and merge it into ``prefix`` in place. + + ``prefix_lse`` uses FlashAttention's varlen layout ``[num_q_heads, + num_queries]``. The suffix page table contains physical token slots, one + row per query, and only its first ``suffix_cache_seqlens[row]`` entries are + visible. + """ + if q.ndim != 3: + raise ValueError("q must have shape [num_queries, num_q_heads, head_dim]") + num_queries, num_q_heads, head_dim = q.shape + if prefix.shape != q.shape: + raise ValueError("prefix output must have the same shape as q") + if prefix_lse.shape != (num_q_heads, num_queries): + raise ValueError("prefix_lse must have shape [num_q_heads, num_queries]") + if k_cache.ndim != 3 or v_cache.ndim != 3: + raise ValueError("flattened KV caches must have shape [slots, heads, dim]") + if k_cache.shape != v_cache.shape: + raise ValueError("K and V caches must have matching shapes") + num_kv_heads = k_cache.shape[1] + if k_cache.shape[2] != head_dim: + raise ValueError("K/V and query head dimensions must match") + if num_q_heads % num_kv_heads: + raise ValueError("query heads must be divisible by KV heads") + if suffix_page_table.ndim != 2 or suffix_page_table.shape[0] != num_queries: + raise ValueError("suffix page table must have one row per query") + if suffix_cache_seqlens.numel() != num_queries: + raise ValueError("suffix cache lengths must have one value per query") + if suffix_page_table.shape[1] == 0 or num_queries == 0: + return prefix + + _fused_suffix_attention_merge_kernel[(num_queries, num_q_heads)]( + q, + k_cache, + v_cache, + suffix_page_table, + suffix_cache_seqlens, + prefix, + prefix_lse, + softmax_scale, + q.stride(0), + q.stride(1), + q.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + suffix_page_table.stride(0), + suffix_page_table.stride(1), + prefix.stride(0), + prefix.stride(1), + prefix.stride(2), + prefix_lse.stride(0), + prefix_lse.stride(1), + NUM_Q_HEADS=num_q_heads, + NUM_KV_HEADS=num_kv_heads, + HEAD_DIM=head_dim, + BLOCK_SUFFIX=triton.next_power_of_2(suffix_page_table.shape[1]), + BLOCK_D=triton.next_power_of_2(head_dim), + num_warps=4, + num_stages=1, + ) + return prefix diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 3726bf44e..4dd09822f 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -327,6 +327,171 @@ def _handle_dflash(server_args: ServerArgs) -> None: ) +def _handle_uno(server_args: ServerArgs) -> None: + cfg = resolving_view(server_args) + + if not cfg.device.startswith("cuda"): + raise ValueError("UNO only supports CUDA.") + if cfg.speculative_draft_model_path is not None: + raise ValueError( + "UNO reuses the target model and does not accept " + "--speculative-draft-model-path." + ) + if cfg.uno_lora_path is None: + raise ValueError("UNO requires --uno-lora-path.") + if cfg.enable_deterministic_inference: + raise ValueError( + "UNO does not support --enable-deterministic-inference because its " + "sampling path does not use per-request seeds." + ) + if cfg.enable_strict_thinking: + raise ValueError( + "UNO does not support --enable-strict-thinking because it requires " + "grammar decoding." + ) + + verify_width = cfg.speculative_num_draft_tokens + if verify_width is None or int(verify_width) < 1: + raise ValueError( + "UNO requires --speculative-num-draft-tokens to be a positive " + "integer denoting the linear width or tree verify width Q." + ) + verify_width = int(verify_width) + declare_resolution( + server_args, + "_handle_uno", + speculative_num_draft_tokens=verify_width, + ) + + candidate_top_k = ( + 1 if cfg.speculative_eagle_topk is None else int(cfg.speculative_eagle_topk) + ) + if candidate_top_k < 1: + raise ValueError( + "UNO requires --speculative-eagle-topk to be at least 1, " + f"got {candidate_top_k}." + ) + + if candidate_top_k > 1: + if cfg.speculative_num_steps is None: + raise ValueError( + "UNO tree mode requires --speculative-num-steps so its draft " + "width F can be derived as speculative_num_steps + 1." + ) + speculative_num_steps = int(cfg.speculative_num_steps) + if speculative_num_steps < 1: + raise ValueError( + "UNO tree mode requires --speculative-num-steps to be positive, " + f"got {speculative_num_steps}." + ) + + draft_width = speculative_num_steps + 1 + if verify_width < draft_width: + raise ValueError( + f"UNO tree mode requires Q >= F; got Q={verify_width}, F={draft_width}." + ) + if verify_width > 128: + raise ValueError( + "UNO tree mode currently supports at most Q=128 verify nodes, " + f"got Q={verify_width}." + ) + frontier_slots = verify_width * candidate_top_k + if frontier_slots > 2048: + raise ValueError( + "UNO tree mode currently supports Q*K <= 2048; got " + f"Q*K={verify_width}*{candidate_top_k}={frontier_slots}." + ) + + tree_capacity = 1 + nodes_at_depth = 1 + for _ in range(speculative_num_steps): + nodes_at_depth *= candidate_top_k + tree_capacity += nodes_at_depth + if tree_capacity >= verify_width: + break + if verify_width > tree_capacity: + raise ValueError( + "UNO tree mode cannot build the requested Q from F and K: " + f"Q={verify_width} exceeds capacity={tree_capacity} for " + f"F={draft_width}, K={candidate_top_k}." + ) + + parent_width = candidate_top_k * max(speculative_num_steps - 1, 0) + 1 + if verify_width - 1 > parent_width: + raise ValueError( + "UNO tree mode cannot represent the requested Q in EAGLE's " + "parent-list ABI: " + f"Q-1={verify_width - 1} exceeds " + f"K*(F-2)+1={parent_width} for " + f"F={draft_width}, K={candidate_top_k}." + ) + + if cfg.enable_pdmux: + raise ValueError("UNO tree mode does not yet support PDMux.") + if cfg.enable_two_batch_overlap: + raise ValueError("UNO tree mode does not yet support two-batch overlap.") + if ( + cfg.speculative_accept_threshold_single != 1.0 + or cfg.speculative_accept_threshold_acc != 1.0 + ): + raise ValueError( + "UNO tree mode reuses EAGLE target-only sampling and requires " + "both speculative accept thresholds to be 1.0." + ) + declare_resolution( + server_args, + "_handle_uno", + speculative_num_steps=speculative_num_steps, + speculative_eagle_topk=candidate_top_k, + ) + else: + for field in ("speculative_num_steps", "speculative_eagle_topk"): + old_value = getattr(cfg, field) + if old_value not in (None, 1): + logger.warning("UNO uses %s=1; overriding %s.", field, old_value) + declare_resolution( + server_args, + "_handle_uno", + speculative_num_steps=1, + speculative_eagle_topk=1, + ) + + if (cfg.tp_size, cfg.pp_size) != (1, 1): + raise ValueError("UNO requires TP=PP=1.") + if cfg.enable_dp_attention or cfg.attn_cp_size != 1: + raise ValueError("UNO does not support DP attention or context parallelism.") + if cfg.enable_lora or cfg.lora_paths: + raise ValueError("UNO does not support public Multi-LoRA serving.") + declare_resolution( + server_args, + "_handle_uno", + enable_lora_overlap_loading=False, + lora_strict_loading=True, + ) + + if cfg.speculative_use_rejection_sampling: + raise ValueError( + "UNO manages its own stochastic verification and does not use " + "--speculative-use-rejection-sampling." + ) + if cfg.enable_mixed_chunk: + declare_resolution( + server_args, + "_handle_uno", + enable_mixed_chunk=False, + ) + logger.warning( + "Mixed chunked prefill is disabled for UNO speculative decoding." + ) + + prefill_backend, decode_backend = attention_backends_of(resolved_view(server_args)) + if (prefill_backend, decode_backend) != ("fa3", "fa3"): + raise ValueError( + "UNO requires FA3 for both prefill and decode attention; " + f"got prefill={prefill_backend!r}, decode={decode_backend!r}." + ) + + def _target_checkpoint_bundles_dspark_draft(server_args: ServerArgs) -> bool: from sglang.srt.speculative.dspark_components.dspark_config import ( checkpoint_bundles_dspark_draft, diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index d228b4048..32d6df547 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -12,6 +12,10 @@ from sglang.kernels.ops.attention.metadata import ( prepare_swa_spec_page_table_triton, ) from sglang.kernels.ops.attention.pa_page_table import _build_pa_page_table +from sglang.kernels.ops.attention.suffix_attention_merge import ( + can_use_fused_suffix_attention_merge, + merge_suffix_attention_in_place, +) from sglang.kernels.ops.attention.utils import assert_buffer_fits from sglang.kernels.ops.kvcache.trtllm_mha_page_table import ( build_trtllm_mha_page_table, @@ -1579,14 +1583,36 @@ class FlashAttentionBackend(AttentionBackend): if use_cascade_attn: o, softmax_lse, *rest = result + if ( + use_cascade_attn + and forward_batch.spec_algorithm.is_uno() + and can_use_fused_suffix_attention_merge( + layer=layer, + q=q, + key_cache=key_cache, + value_cache=value_cache, + extra_kwargs=kwargs, + ) + ): + suffix_metadata = self.forward_metadata_spec_decode_expand + o = merge_suffix_attention_in_place( + q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + k_cache=key_cache.view(-1, layer.tp_k_head_num, layer.head_dim), + v_cache=value_cache.view(-1, layer.tp_v_head_num, layer.v_head_dim), + suffix_page_table=suffix_metadata.page_table, + suffix_cache_seqlens=suffix_metadata.cache_seqlens_int32, + prefix=o, + prefix_lse=softmax_lse, + softmax_scale=layer.scaling, + ) + elif use_cascade_attn: o_expand, softmax_lse_expand, *rest_expand = flash_attn_with_kvcache( q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), - # Here metadata_expand.page_table is not divided with page_size. - # This is because we loose the fine control of what token to attend, - # but has to attend to some block completely. + # The suffix table stores physical token slots, so expose + # the paged cache as page-size-one blocks. k_cache=key_cache.view(-1, 1, layer.tp_k_head_num, layer.head_dim), v_cache=value_cache.view( - -1, 1, layer.tp_v_head_num, layer.head_dim + -1, 1, layer.tp_v_head_num, layer.v_head_dim ), page_table=self.forward_metadata_spec_decode_expand.page_table, cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32, diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py index 6d3cf22f8..29723d929 100644 --- a/python/sglang/srt/lora/backend/base_backend.py +++ b/python/sglang/srt/lora/backend/base_backend.py @@ -23,6 +23,9 @@ class BaseLoRABackend(LoRABackendLmHeadMixing): device: the device where the backend runs. """ + supports_lora_a_overlap = False + skip_inactive_lora_batches = False + # Supporting backends implement init_prefill_cuda_graph_batch_info() and # honor use_prefill_cuda_graph in prepare_lora_batch(). supports_prefill_cuda_graph: bool = False @@ -52,6 +55,14 @@ class BaseLoRABackend(LoRABackendLmHeadMixing): self.lm_head_pass_batch_infos = None self._lm_head_pass_idx = None + def validate_lora_targets( + self, + base_model: torch.nn.Module, + target_modules: set[str], + ) -> None: + """Raise before wrapping when this backend cannot execute its targets.""" + pass + def run_lora_a_embedding( self, input_ids: torch.Tensor, @@ -351,6 +362,20 @@ class BaseLoRABackend(LoRABackendLmHeadMixing): """ pass + def prepare_lora_token_segments( + self, + *, + segment_lens: list[int], + weight_indices: list[int], + lora_ranks: list[int], + scalings: list[float], + ) -> None: + """Prepare explicit eager token-row LoRA segments.""" + raise NotImplementedError( + f"LoRA backend {type(self).__name__} does not support explicit " + "token segments." + ) + @triton.jit def _compute_moe_lora_info_kernel( diff --git a/python/sglang/srt/lora/backend/lora_registry.py b/python/sglang/srt/lora/backend/lora_registry.py index 15b531ab3..40447d6e4 100644 --- a/python/sglang/srt/lora/backend/lora_registry.py +++ b/python/sglang/srt/lora/backend/lora_registry.py @@ -44,6 +44,13 @@ def create_torch_native_backend(): return TorchNativeLoRABackend +@register_lora_backend("uno_cublas") +def create_uno_cublas_backend(): + from sglang.srt.lora.backend.uno_cublas_backend import UnoCublasLoRABackend + + return UnoCublasLoRABackend + + @register_lora_backend("flashinfer") def create_flashinfer_backend(): raise ValueError( diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index 53a413275..12656ecbe 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/python/sglang/srt/lora/backend/triton_backend.py @@ -362,6 +362,44 @@ class TritonLoRABackend(BaseLoRABackend): self._prepare_lm_head_batch_info(forward_batch, weight_indices, batch_info) ) + def prepare_lora_token_segments( + self, + *, + segment_lens: list[int], + weight_indices: list[int], + lora_ranks: list[int], + scalings: list[float], + ) -> None: + """Install explicit eager token-row routing metadata.""" + self.reset_batch_state() + + segment_lens_tensor = torch.tensor( + segment_lens, dtype=torch.int32, device=self.device + ) + segment_indptr = torch.zeros( + len(segment_lens) + 1, dtype=torch.int32, device=self.device + ) + segment_indptr[1:] = torch.cumsum(segment_lens_tensor, dim=0) + + self.batch_info = LoRABatchInfo( + use_cuda_graph=False, + bs=len(segment_lens), + num_segments=len(segment_lens), + seg_lens=segment_lens_tensor, + seg_indptr=segment_indptr, + max_len=max(segment_lens), + weight_indices=torch.tensor( + weight_indices, dtype=torch.int32, device=self.device + ), + lora_ranks=torch.tensor(lora_ranks, dtype=torch.int64, device=self.device), + scalings=torch.tensor(scalings, dtype=torch.float, device=self.device), + permutation=None, + expected_tokens=sum(segment_lens), + ) + + # These segments already describe physical token-row order. + self.sgemm_batch_info = None + def _prepare_lm_head_batch_info( self, forward_batch: ForwardBatch, diff --git a/python/sglang/srt/lora/backend/uno_cublas_backend.py b/python/sglang/srt/lora/backend/uno_cublas_backend.py new file mode 100644 index 000000000..e11b9f00c --- /dev/null +++ b/python/sglang/srt/lora/backend/uno_cublas_backend.py @@ -0,0 +1,455 @@ +"""Single-adapter LoRA backend for UNO draft forwards. + +Parallel-linear layers overlap LoRA-A with the base GEMM on an auxiliary CUDA +stream. Other dense LoRA layers use the inherited Triton implementation. +Single-request cuBLAS batches operate only on draft rows; larger batches use +one ``mm``/``addmm_`` over all rows and zero the seed-row LoRA hidden states +between the two GEMMs. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from sglang.srt.lora.backend.triton_backend import TritonLoRABackend +from sglang.srt.lora.utils import LoRABatchInfo + + +@dataclass(frozen=True) +class _UnoSingleAdapterRoute: + weight_index: int + rank: int + scaling: float + batch_size: int + forward_width: int + total_rows: int + active_rows: int + + @property + def lora_rows(self) -> int: + # For C=1, skipping the seed row makes both GEMMs smaller. For C>1, + # one GEMM across C*F rows is faster than C tiny (F-1)-row GEMMs. + return self.active_rows if self.batch_size == 1 else self.total_rows + + +@dataclass(frozen=True) +class _PendingLoRAA: + output: torch.Tensor + producer_stream: torch.cuda.Stream + + +class UnoCublasLoRABackend(TritonLoRABackend): + """Fast LoRA backend for UNO's draft forwards. + Use cuBLAS and CUDA streams. + + Reuse multi-LoRA batch metadata from the Triton parent. + """ + + name = "uno_cublas" + supports_lora_a_overlap = True + # K2's runners prepare base-only LoRA metadata whenever any internal + # manager exists. UNO never exposes request-selectable adapters, so those + # prefill/warmup batches must stay on the plain base-model path. + skip_inactive_lora_batches = True + + def __init__( + self, + max_loras_per_batch: int, + device: torch.device, + **kwargs, + ): + super().__init__(max_loras_per_batch, device, **kwargs) + # Different CUDA-graph runners may capture and replay on different main + # streams. Give each main stream its own LoRA-A side stream so concurrently + # replayed graphs do not serialize or interfere through one shared stream. + self._lora_a_streams: dict[torch.cuda.Stream, torch.cuda.Stream] = {} + self._pending_lora_a: Optional[_PendingLoRAA] = None + self._use_cublas_lora_b = False + + def reset_batch_state(self): + self._pending_lora_a = None + self._use_cublas_lora_b = False + super().reset_batch_state() + + def validate_lora_targets( + self, + base_model: torch.nn.Module, + target_modules: set[str], + ) -> None: + """Reject target layers that cannot honor UNO's token-row routing.""" + + from sglang.srt.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, + ) + from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE + from sglang.srt.layers.utils import get_layer_id + from sglang.srt.models.inkling_common.dense_mlp import InklingBatchDenseMLP + + unsupported: list[str] = [] + # Embedding and LM-head wrappers use the inherited Triton kernels and + # are handled separately by LoRAManager. Decoder-layer projections use + # either the overlapped cuBLAS path or the Triton ReplicatedLinear path. + supported = (ColumnParallelLinear, RowParallelLinear, ReplicatedLinear) + target_moe = {"gate_up_proj", "down_proj"}.issubset(target_modules) + for module_name, module in base_model.named_modules(): + parts = module_name.split(".") + named_target = bool(parts) and ( + parts[-1] in target_modules or ".".join(parts[-2:]) in target_modules + ) + special_moe_target = target_moe and isinstance( + module, (FusedMoE, InklingBatchDenseMLP) + ) + if not (named_target or special_moe_target): + continue + if get_layer_id(module_name) is None: + continue + if not isinstance(module, supported): + unsupported.append(f"{module_name} ({type(module).__name__})") + + if unsupported: + raise ValueError( + "UNO's LoRA backend cannot execute these target modules: " + + ", ".join(sorted(set(unsupported))) + ) + + def prepare_lora_token_segments( + self, + *, + segment_lens: list[int], + weight_indices: list[int], + lora_ranks: list[int], + scalings: list[float], + ) -> None: + super().prepare_lora_token_segments( + segment_lens=segment_lens, + weight_indices=weight_indices, + lora_ranks=lora_ranks, + scalings=scalings, + ) + + route: Optional[_UnoSingleAdapterRoute] = None + if ( + len(segment_lens) >= 2 + and len(segment_lens) % 2 == 0 + and len(weight_indices) == len(segment_lens) + and segment_lens[0] == 1 + and segment_lens[1] > 0 + ): + batch_size = len(segment_lens) // 2 + forward_width = segment_lens[1] + 1 + base_index, adapter_index = weight_indices[:2] + adapter_rank = lora_ranks[adapter_index] + if ( + base_index != adapter_index + and lora_ranks[base_index] == 0 + and adapter_rank > 0 + and all( + segment_lens[2 * index] == 1 + and segment_lens[2 * index + 1] == forward_width - 1 + and weight_indices[2 * index] == base_index + and weight_indices[2 * index + 1] == adapter_index + for index in range(batch_size) + ) + ): + route = _UnoSingleAdapterRoute( + weight_index=adapter_index, + rank=adapter_rank, + scaling=float(scalings[adapter_index]), + batch_size=batch_size, + forward_width=forward_width, + total_rows=sum(segment_lens), + active_rows=batch_size * (forward_width - 1), + ) + + # Each CUDA-graph bucket retains its own batch_info. Store the immutable + # UNO route there so switching back to a captured bucket restores the + # route corresponding to that bucket, rather than using metadata last + # written by another bucket. + self.batch_info.uno_single_adapter_route = route + + def _route(self) -> _UnoSingleAdapterRoute: + route = getattr(self.batch_info, "uno_single_adapter_route", None) + if route is None: + raise RuntimeError("UNO cuBLAS execution requires an active UNO route.") + return route + + @staticmethod + def _output_offsets(output_offset_cpu, output_offset) -> list[int]: + offsets = output_offset_cpu if output_offset_cpu is not None else output_offset + return [int(offset) for offset in offsets.tolist()] + + @staticmethod + def _lora_a_input( + x: torch.Tensor, + route: _UnoSingleAdapterRoute, + ) -> torch.Tensor: + if route.batch_size == 1: + return x[1:] + return x + + def _compute_lora_a( + self, + lora_input: torch.Tensor, + active_a: torch.Tensor, + route: _UnoSingleAdapterRoute, + *, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + hidden = torch.mm(lora_input, active_a.t(), out=output) + if route.batch_size > 1: + # LoRA must not affect each request's seed token. Zeroing C rows + # is cheaper than masking every hidden element or running C tiny + # strided-batched GEMMs. + hidden.view(route.batch_size, route.forward_width, -1)[:, 0].zero_() + return hidden + + def _accumulate_lora_b( + self, + *, + hidden: torch.Tensor, + active_b: torch.Tensor, + base_output: torch.Tensor, + output_start: int, + output_end: int, + route: _UnoSingleAdapterRoute, + ) -> None: + if route.batch_size == 1: + base_output[1:, output_start:output_end].addmm_( + hidden, + active_b.t(), + beta=1.0, + alpha=route.scaling, + ) + return + + base_output[:, output_start:output_end].addmm_( + hidden, + active_b.t(), + beta=1.0, + alpha=route.scaling, + ) + + def _run_lora_b( + self, + x: torch.Tensor, + weights: torch.Tensor, + base_output: torch.Tensor, + ) -> torch.Tensor: + route = self._route() + active_b = weights[route.weight_index, :, : route.rank] + self._accumulate_lora_b( + hidden=x[:, : route.rank], + active_b=active_b, + base_output=base_output, + output_start=0, + output_end=active_b.shape[0], + route=route, + ) + return base_output + + def run_lora_a_sgemm( + self, + x: torch.Tensor, + weights: torch.Tensor, + pruned_batch_info: LoRABatchInfo = None, + stack_num: int = 1, + *args, + **kwargs, + ) -> torch.Tensor: + pending = self._pending_lora_a + if pending is None: + self._use_cublas_lora_b = False + return super().run_lora_a_sgemm( + x, + weights, + pruned_batch_info, + stack_num, + *args, + **kwargs, + ) + + output = self._consume_lora_a_overlap(pending) + self._use_cublas_lora_b = True + return output + + def run_lora_b_sgemm( + self, + x: torch.Tensor, + weights: torch.Tensor, + base_output: torch.Tensor = None, + pruned_batch_info: LoRABatchInfo = None, + *args, + **kwargs, + ) -> torch.Tensor: + if not self._use_cublas_lora_b: + return super().run_lora_b_sgemm( + x, + weights, + base_output, + pruned_batch_info, + *args, + **kwargs, + ) + + self._use_cublas_lora_b = False + return self._run_lora_b(x, weights, base_output) + + def _run_stacked_lora( + self, + *, + lora_b: torch.Tensor, + base_output: torch.Tensor, + output_offset, + output_offset_cpu, + num_slices: int, + pending: _PendingLoRAA, + ) -> torch.Tensor: + route = self._route() + hidden = self._consume_lora_a_overlap(pending) + offsets = self._output_offsets(output_offset_cpu, output_offset) + for slice_index in range(num_slices): + input_start = slice_index * route.rank + input_end = input_start + route.rank + output_start = offsets[slice_index] + output_end = offsets[slice_index + 1] + active_b = lora_b[ + route.weight_index, + output_start:output_end, + : route.rank, + ] + self._accumulate_lora_b( + hidden=hidden[:, input_start:input_end], + active_b=active_b, + base_output=base_output, + output_start=output_start, + output_end=output_end, + route=route, + ) + return base_output + + def run_qkv_lora( + self, + x: torch.Tensor, + qkv_lora_a: torch.Tensor, + qkv_lora_b: torch.Tensor, + output_offset: torch.Tensor, + max_qkv_out_dim: int, + base_output: torch.Tensor = None, + n_slices: int = 3, + *args, + output_offset_cpu=None, + **kwargs, + ) -> torch.Tensor: + pending = self._pending_lora_a + if pending is None: + return super().run_qkv_lora( + x, + qkv_lora_a, + qkv_lora_b, + output_offset, + max_qkv_out_dim, + base_output, + n_slices, + *args, + **kwargs, + ) + + return self._run_stacked_lora( + lora_b=qkv_lora_b, + base_output=base_output, + output_offset=output_offset, + output_offset_cpu=output_offset_cpu, + num_slices=n_slices, + pending=pending, + ) + + def run_gate_up_lora( + self, + x: torch.Tensor, + gate_up_lora_a: torch.Tensor, + gate_up_lora_b: torch.Tensor, + base_output: torch.Tensor = None, + *args, + output_offset=None, + output_offset_cpu=None, + **kwargs, + ) -> torch.Tensor: + pending = self._pending_lora_a + if pending is None: + return super().run_gate_up_lora( + x, + gate_up_lora_a, + gate_up_lora_b, + base_output, + *args, + **kwargs, + ) + + return self._run_stacked_lora( + lora_b=gate_up_lora_b, + base_output=base_output, + output_offset=output_offset, + output_offset_cpu=output_offset_cpu, + num_slices=2, + pending=pending, + ) + + def start_lora_a_overlap( + self, + x: torch.Tensor, + weights: torch.Tensor, + *, + num_slices: int = 1, + ) -> None: + """Launch LoRA-A on the auxiliary stream before the base GEMM.""" + + if self._pending_lora_a is not None: + raise RuntimeError("Previous UNO LoRA-A overlap was not consumed.") + + route = self._route() + out_dim = num_slices * route.rank + lora_input = self._lora_a_input(x, route) + active_a = weights[route.weight_index, :out_dim] + + main_stream = torch.cuda.current_stream(x.device) + stream = self._lora_a_streams.get(main_stream) + if stream is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "UNO LoRA overlap side stream was not created during " + "CUDA-graph warmup." + ) + stream = torch.cuda.Stream(device=x.device) + self._lora_a_streams[main_stream] = stream + + # Allocate on the main/consumer stream before handing the buffer to + # the auxiliary stream. This keeps CUDA-graph allocator ownership and + # the eventual LoRA-B consumer on the same stream. + output = torch.empty( + (route.lora_rows, out_dim), + dtype=x.dtype, + device=x.device, + ) + stream.wait_stream(main_stream) + with torch.cuda.stream(stream): + self._compute_lora_a(lora_input, active_a, route, output=output) + + self._pending_lora_a = _PendingLoRAA( + output=output, + producer_stream=stream, + ) + + def _consume_lora_a_overlap( + self, + pending: _PendingLoRAA, + ) -> torch.Tensor: + self._pending_lora_a = None + torch.cuda.current_stream(pending.output.device).wait_stream( + pending.producer_stream + ) + return pending.output diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index 173e1000e..c5777e6df 100644 --- a/python/sglang/srt/lora/layers.py +++ b/python/sglang/srt/lora/layers.py @@ -66,7 +66,15 @@ class BaseLayerWithLoRA(nn.Module): has LoRA batch metadata. batch_info is None on DP-attention idle forwards (see LoRAManager.prepare_lora_batch), so idle forwards take the base path.""" - return self.set_lora and self.lora_backend.batch_info is not None + batch_info = self.lora_backend.batch_info + return ( + self.set_lora + and batch_info is not None + and ( + not self.lora_backend.skip_inactive_lora_batches + or batch_info.has_active_lora + ) + ) def set_lora_info(self, *args): pass @@ -482,14 +490,22 @@ class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA): ) return lora_output + def start_lora_a_overlap(self, x: torch.Tensor) -> None: + if self.lora_backend.supports_lora_a_overlap: + self.lora_backend.start_lora_a_overlap(x, self.A_buffer) + def forward(self, input_: torch.Tensor): # duplicate the logic in ColumnParallelLinear + lora_active = self.lora_active + if lora_active: + self.start_lora_a_overlap(input_) + bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None output_parallel = self.base_layer.quant_method.apply( self.base_layer, input_, bias ) - if self.lora_active: + if lora_active: output_parallel = self.apply_lora(output_parallel, input_) if self.base_layer.gather_output: @@ -596,6 +612,12 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): ) return lora_output + def start_lora_a_overlap(self, x: torch.Tensor) -> None: + if self.lora_backend.supports_lora_a_overlap: + self.lora_backend.start_lora_a_overlap( + x, self.A_buffer, num_slices=self._get_lora_n_slices() + ) + def slice_lora_a_weights(self, A: torch.Tensor): return A @@ -703,6 +725,10 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): return lora_output + def start_lora_a_overlap(self, x: torch.Tensor) -> None: + if self.lora_backend.supports_lora_a_overlap: + self.lora_backend.start_lora_a_overlap(x, self.A_buffer_qkv, num_slices=3) + def slice_lora_a_weights(self, A: torch.Tensor): return A @@ -773,6 +799,10 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA): ) return lora_output + def start_lora_a_overlap(self, x: torch.Tensor) -> None: + if self.lora_backend.supports_lora_a_overlap: + self.lora_backend.start_lora_a_overlap(x, self.A_buffer) + def forward(self, input_: torch.Tensor, skip_all_reduce=False, forward_batch=None): if self.base_layer.input_is_parallel: input_parallel = input_ @@ -783,6 +813,10 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA): ) input_parallel = splitted_input[tp_rank].contiguous() + lora_active = self.lora_active + if lora_active: + self.start_lora_a_overlap(input_parallel) + bias_ = ( None if (self.base_layer.tp_rank > 0 or self.base_layer.skip_bias_add) @@ -806,7 +840,6 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA): all_reduce = get_parallel().attn_tp_group.all_reduce else: all_reduce = tensor_model_parallel_all_reduce - lora_active = self.lora_active if lora_active and should_reduce: lora_a_output = self.lora_backend.run_lora_a_sgemm( input_parallel, self.A_buffer diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index fc94f8b2c..564a0aa1d 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -17,7 +17,7 @@ import logging import re -from typing import Dict, Iterable, List, Optional +from typing import Dict, Iterable, List, Optional, Sequence import torch @@ -431,6 +431,16 @@ class LoRAManager: self.lora_backend.reset_batch_state() def prepare_lora_batch(self, forward_batch: ForwardBatch): + # Some internal-only backends (currently UNO) use explicit token-row + # routing for their adapted forwards and want all-base batches to run + # through the plain model path. Clear any routing retained by the + # preceding adapted forward before inspecting CUDA-graph metadata. + if self.lora_backend.skip_inactive_lora_batches and not any( + uid is not None for uid in forward_batch.lora_ids + ): + self.reset_lora_batch() + return + # set up batch info shared by all lora modules bs = forward_batch.batch_size @@ -470,6 +480,39 @@ class LoRAManager: lora_ranks[wi] > 0 for wi in weight_indices ) + def prepare_lora_token_segments( + self, + *, + lora_ids: Sequence[Optional[str]], + segment_lens: Sequence[int], + ) -> None: + """Prepare eager LoRA routing independently of request batching.""" + lora_ids = list(lora_ids) + segment_lens = list(segment_lens) + if len(lora_ids) != len(segment_lens): + raise ValueError("LoRA ids and segment lengths must have equal length.") + + weight_indices = [] + lora_ranks = [0] * self.max_loras_per_batch + scalings = [0.0] * self.max_loras_per_batch + for lora_id in lora_ids: + weight_index = self.memory_pool.get_buffer_id(lora_id) + weight_indices.append(weight_index) + if lora_id is not None: + lora = self.loras[lora_id] + lora_ranks[weight_index] = lora.config.r + scalings[weight_index] = lora.scaling + + self.lora_backend.prepare_lora_token_segments( + segment_lens=segment_lens, + weight_indices=weight_indices, + lora_ranks=lora_ranks, + scalings=scalings, + ) + self.lora_backend.batch_info.has_active_lora = any( + lora_ranks[index] > 0 for index in weight_indices + ) + def update_lora_info(self): """ Update all LoRA modules to associate them with the latest memory buffer. @@ -579,6 +622,10 @@ class LoRAManager: max_lora_rank=max_lora_rank, target_modules=target_modules, ) + self.lora_backend.validate_lora_targets( + base_model=self.base_model, + target_modules=self.target_modules, + ) if self._experts_shared_outer_override is not None: self.experts_shared_outer_loras = self._experts_shared_outer_override diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index cf0f12b17..00a6844f0 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -308,6 +308,7 @@ from sglang.srt.speculative.eagle_utils import ( get_draft_recurrent_hidden_state_spec_from_config, ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.speculative.uno_validation import validate_uno_request from sglang.srt.utils import ( DynamicGradMode, configure_gc_logger, @@ -2813,6 +2814,14 @@ class Scheduler( self._add_request_to_queue(req) return + if self.spec_algorithm.is_uno(): + error_msg = validate_uno_request(req) + if error_msg is not None: + req.set_finish_with_abort(error_msg) + self.init_req_max_new_tokens(req) + self._add_request_to_queue(req) + return + if ( req.return_sampling_mask and self.disaggregation_mode != DisaggregationMode.NULL @@ -4484,7 +4493,7 @@ class Scheduler( self.decode_moment_totals, batch_size, step_us, - batch_size + result.num_correct_drafts, + result.get_num_generated_tokens(batch_size), ) def maybe_send_health_check_signal(self): diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 5aeec8a2f..eb6cf4a8f 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -78,6 +78,16 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def _get_speculative_output_stride(result: GenerationBatchResult) -> int: + """Return the padded per-request width in flattened speculative output.""" + stride = result.speculative_output_stride + if stride is None: + stride = result.speculative_num_draft_tokens + if stride is None or stride < 1: + raise RuntimeError("speculative result is missing a positive output row stride") + return stride + + @dataclass(kw_only=True, slots=True, frozen=True) class SchedulerBatchResultProcessor: is_generation: bool @@ -711,8 +721,12 @@ class SchedulerBatchResultProcessor: next_token_ids = result.next_token_ids.tolist() accept_lens = result.accept_lens.tolist() - result.num_correct_drafts = sum(accept_lens) - len(batch.reqs) - result.num_correct_drafts_per_req_cpu = [x - 1 for x in accept_lens] + stride = _get_speculative_output_stride(result) + num_non_draft = result.num_non_draft_tokens_per_req + result.num_correct_drafts_per_req_cpu = [ + length - num_non_draft for length in accept_lens + ] + result.num_correct_drafts = sum(result.num_correct_drafts_per_req_cpu) block_accept_lens = ( result.block_accept_lens.tolist() @@ -740,11 +754,6 @@ class SchedulerBatchResultProcessor: self.advance_grammar_fsm(result, batch) predict_tokens = [] - # In adaptive spec-v2, the worker state may already have switched when this - # delayed result is processed. Use the draft token count recorded on result. - stride = result.speculative_num_draft_tokens - assert stride is not None, "spec-v2 result missing speculative_num_draft_tokens" - for i, req in enumerate(batch.reqs): accept_tokens = next_token_ids[i * stride : i * stride + accept_lens[i]] @@ -853,8 +862,7 @@ class SchedulerBatchResultProcessor: if result.accept_lens is None: return accept_lens = result.accept_lens.tolist() - stride = result.speculative_num_draft_tokens - assert stride is not None, "spec-v2 result missing speculative_num_draft_tokens" + stride = _get_speculative_output_stride(result) retained = [None] * len(batch.reqs) for i, req in enumerate(batch.reqs): if req.grammar is None or req.is_retracted or req.finished(): @@ -907,11 +915,14 @@ class SchedulerBatchResultProcessor: next_token_ids=next_token_ids, ) - self.metrics_reporter.num_generated_tokens += len(batch.reqs) + batch_size = batch.batch_size() + num_generated_tokens = result.get_num_generated_tokens(batch_size) + self.metrics_reporter.num_generated_tokens += num_generated_tokens if not batch.spec_algorithm.is_none(): self.metrics_reporter.update_spec_metrics( - batch.batch_size(), + batch_size, result.num_correct_drafts, + num_accept_tokens=num_generated_tokens, num_block_accept_tokens=result.num_block_accept_tokens, num_cap_tokens=result.num_cap_tokens, ) @@ -982,8 +993,8 @@ class SchedulerBatchResultProcessor: if req.return_hidden_states and logits_output.hidden_states is not None: # hidden_states is [bs * stride, hidden_dim], one row per emitted - # token; stride = speculative_num_draft_tokens for spec, 1 for non-spec. - stride = result.speculative_num_draft_tokens or 1 + # token; speculative workers record their padded row width. + stride = _get_speculative_output_stride(result) if is_spec else 1 accept_len = len(next_token_id) start = i * stride self._append_decode_hidden_states( @@ -1016,7 +1027,7 @@ class SchedulerBatchResultProcessor: self.metrics_reporter.report_decode_stats( can_run_cuda_graph, running_batch=batch, - num_correct_drafts=result.num_correct_drafts, + num_generated_tokens=num_generated_tokens, ) def _normalize_decode_outputs( diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 092521c40..f57ee42a1 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -175,9 +175,10 @@ class SchedulerMetricsReporter: }.get(getattr(self.scheduler, "device", ""), "cuda graph") # Cumulative spec-decoding counters (reset every decode_log_interval). - # Each update adds (num_correct_drafts + bs, bs). - # `*_accept_tokens` = drafts + bonus; `*_correct_drafts` = drafts-only. + # `*_accept_tokens` includes accepted drafts and non-draft output tokens; + # `*_correct_drafts` counts accepted draft proposals only. self.spec_num_accept_tokens = 0 # per-log-interval + self.spec_num_correct_drafts = 0 self.spec_num_forward_ct = 0 self.spec_total_num_accept_tokens = 0 # lifetime self.spec_total_num_forward_ct = 0 @@ -398,17 +399,16 @@ class SchedulerMetricsReporter: self, bs: int, num_correct_drafts: int, + num_accept_tokens: int, num_block_accept_tokens: int = 0, num_cap_tokens: int = 0, ): - self.spec_num_accept_tokens += num_correct_drafts + bs + self.spec_num_accept_tokens += num_accept_tokens + self.spec_num_correct_drafts += num_correct_drafts self.spec_num_forward_ct += bs self.spec_num_block_accept_tokens += num_block_accept_tokens self.spec_num_cap_tokens += num_cap_tokens - # Bonus tokens updated elsewhere - self.num_generated_tokens += num_correct_drafts - def _init_estimated_perf_constants(self) -> None: model_config = self.scheduler.model_config hf_text_config = model_config.hf_text_config @@ -572,6 +572,7 @@ class SchedulerMetricsReporter: self.forward_ct_decode = 0 self.num_generated_tokens = 0 self.spec_num_accept_tokens = 0 + self.spec_num_correct_drafts = 0 self.spec_num_forward_ct = 0 self.spec_total_num_accept_tokens = 0 self.spec_total_num_forward_ct = 0 @@ -757,13 +758,13 @@ class SchedulerMetricsReporter: self, can_run_cuda_graph: bool, running_batch: ScheduleBatch = None, - num_correct_drafts: int = 0, + num_generated_tokens: int = 0, ): batch = running_batch or self.scheduler.running_batch # Every-iteration work: realtime token counting + status logger if self.current_scheduler_metrics_enabled: - decode_tokens = batch.batch_size() + num_correct_drafts + decode_tokens = num_generated_tokens self.metrics_collector.increment_realtime_tokens( # TODO unify this w/ the bumping logic in `Scheduler.num_generated_tokens` accumulator decode_tokens=decode_tokens, @@ -826,7 +827,7 @@ class SchedulerMetricsReporter: spec_block_accept_length = 0 else: spec_accept_length = self.spec_num_accept_tokens / self.spec_num_forward_ct - num_correct_drafts = self.spec_num_accept_tokens - self.spec_num_forward_ct + num_correct_drafts = self.spec_num_correct_drafts if get_spec().speculative_num_draft_tokens: draft_per_round = get_spec().speculative_num_draft_tokens - 1 else: @@ -853,7 +854,8 @@ class SchedulerMetricsReporter: ) self.spec_total_num_accept_tokens += self.spec_num_accept_tokens self.spec_total_num_forward_ct += self.spec_num_forward_ct - self.spec_num_accept_tokens = self.spec_num_forward_ct = 0 + self.spec_num_accept_tokens = self.spec_num_correct_drafts = 0 + self.spec_num_forward_ct = 0 self.spec_num_block_accept_tokens = 0 self.spec_num_cap_tokens = 0 msg += f"accept len: {spec_accept_length:.2f}, accept rate: {spec_accept_rate:.2f}, " diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index ee08f42bf..2695df7d7 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -22,7 +22,7 @@ from sglang.srt.state_capturer.base import TopkCaptureOutput if TYPE_CHECKING: from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.sampling.sampling_observer import HostAuxiliaryOutput - from sglang.srt.speculative.eagle_info import EagleDraftInput + from sglang.srt.speculative.spec_info import SpecInput logger = logging.getLogger(__name__) @@ -71,6 +71,12 @@ class GenerationBatchResult: delay_sample_func: Optional[callable] = None future_indices: Optional[torch.Tensor] = None speculative_num_draft_tokens: Optional[int] = None + # Padded row width in flattened speculative output. Existing algorithms + # default to speculative_num_draft_tokens; linear UNO emits F + 1 columns. + speculative_output_stride: Optional[int] = None + # Valid output tokens that are not accepted draft proposals. Existing + # algorithms have one bonus token; UNO also emits its clean root. + num_non_draft_tokens_per_req: int = 1 # Grammar FSM advance memoization (spec-v2 overlap). advance_grammar_fsm sets # these once — eagerly via the scheduler's grammar barrier inside verify(), or @@ -91,7 +97,7 @@ class GenerationBatchResult: new_seq_lens: Optional[torch.Tensor] = None # relay path: forward stream -> next step forward - next_draft_input: Optional[EagleDraftInput] = None + next_draft_input: Optional[SpecInput] = None # Refs the worker wants scheduler to keep alive for the same 2-iter window # as batch_record_buf. Used for cross-stream tensor lifetime (e.g. a spec @@ -117,6 +123,9 @@ class GenerationBatchResult: this rank/split (a non-last PP rank or a non-final prefill split).""" return isinstance(self.next_token_ids, torch.Tensor) + def get_num_generated_tokens(self, batch_size: int) -> int: + return self.num_correct_drafts + batch_size * self.num_non_draft_tokens_per_req + @torch.profiler.record_function("copy_result_to_cpu") def copy_to_cpu(self, return_logprob: bool, return_hidden_states: bool = True): """Copy tensors to CPU in overlap scheduling. diff --git a/python/sglang/srt/mem_cache/allocation_sizing.py b/python/sglang/srt/mem_cache/allocation_sizing.py index bcb476f92..f41a0b1c4 100644 --- a/python/sglang/srt/mem_cache/allocation_sizing.py +++ b/python/sglang/srt/mem_cache/allocation_sizing.py @@ -33,6 +33,11 @@ def get_alloc_len_per_decode() -> int: from sglang.srt.speculative.spec_info import SpeculativeAlgorithm spec_algo = SpeculativeAlgorithm.from_string(spec.speculative_algorithm) + if spec_algo.is_uno(): + if spec_tokens is None: + raise RuntimeError("UNO requires speculative_num_draft_tokens") + # UNO retains an additional clean-root position beside Q/F draft slots. + return spec_tokens + 1 if page_size == 1 or spec_topk == 1 or not spec_algo.has_draft_kv(): return max(spec_steps * spec_topk, spec_tokens) else: @@ -88,9 +93,13 @@ def get_req_to_token_extra_context_len() -> int: # FIXME(lsyin): temporary fix for the context length issue under spec decoding extra = 4 + (max_speculative_num_draft_tokens() or 0) page_size = get_alloc_page_size() - if get_spec().speculative_algorithm is not None and page_size > 1: - # kv_allocated_len is page-aligned (eagle_prepare_for_decode), so near - # the context limit the aligned reserve can overshoot by page_size - 1; - # without the headroom the row write silently lands in the neighbor row. - extra = max(extra, get_alloc_reserve_per_decode() + page_size - 1) + spec_algorithm = get_spec().speculative_algorithm + if spec_algorithm is not None: + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + spec_algo = SpeculativeAlgorithm.from_string(spec_algorithm) + if page_size > 1 or spec_algo.is_uno(): + # UNO's double-buffer reserve applies at every page size. Larger + # pages may additionally round the allocation up by page_size - 1. + extra = max(extra, get_alloc_reserve_per_decode() + page_size - 1) return extra diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ace5fd98a..cd539390b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -777,6 +777,12 @@ class ModelRunner: self.apply_torch_tp() def maybe_init_lora_manager(self): + if self.spec_algorithm.is_uno(): + from sglang.srt.speculative.uno_lora import init_uno_lora_manager + + self.lora_manager, self.uno_lora_id = init_uno_lora_manager(self) + return + # Adapters apply to the target model only; the draft runs unadapted. if get_lora().enable_lora and not self.is_draft_worker: self.init_lora_manager() @@ -1430,6 +1436,13 @@ class ModelRunner: Subclasses can override this to install specialized decode graph runners. """ + if self.spec_algorithm.is_uno(): + from sglang.srt.speculative.uno_cuda_graph_runner import ( + UnoDecodeCudaGraphRunner, + ) + + return UnoDecodeCudaGraphRunner + from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( DecodeCudaGraphRunner, ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 94392accd..b534f9d74 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2067,7 +2067,12 @@ class ServerArgs: # ------------------------------------------------------------------------- speculative_algorithm: A[ Optional[str], - "Speculative algorithm. Builtins: EAGLE, EAGLE3, NEXTN, STANDALONE, NGRAM, DFLASH, DSPARK. Or any name registered via `SpeculativeAlgorithm.register`.", + "Speculative algorithm. Builtins: EAGLE, EAGLE3, NEXTN, STANDALONE, NGRAM, DFLASH, DSPARK, UNO. Or any name registered via `SpeculativeAlgorithm.register`.", + NS("spec"), + ] = None + uno_lora_path: A[ + Optional[str], + "Path to the UNO draft LoRA checkpoint.", NS("spec"), ] = None speculative_draft_model_path: A[ diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 368c18035..f679b5050 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -663,11 +663,30 @@ def _verify_coins( return coins, coins_for_final_sampling +def _can_use_sparse_uno_tree_target_sampling( + max_top_k: Optional[int], + sampling_info: SamplingBatchInfo, +) -> bool: + if max_top_k is None: + return False + + from sglang.srt.speculative.uno_utils import _SPARSE_TOP_K_LIMIT + + return bool( + _is_cuda + and max_top_k <= _SPARSE_TOP_K_LIMIT + and sampling_info.sampling_seed is None + and not sampling_info.need_min_p_sampling + and not get_spec().speculative_use_rejection_sampling + ) + + def eagle_sample( verify_input: EagleVerifyInput, batch: ScheduleBatch, logits_output: LogitsProcessorOutput, grammar_mask: Optional[GrammarMask] = None, + uno_target_max_top_k: Optional[int] = None, ): """ Verify and find accepted tokens based on logits output and batch @@ -769,6 +788,42 @@ def eagle_sample( tp_group.broadcast(predict, src=0) tp_group.broadcast(accept_index, src=0) tp_group.broadcast(num_correct_drafts, src=0) + elif _can_use_sparse_uno_tree_target_sampling( + uno_target_max_top_k, + sampling_info, + ): + from sglang.srt.speculative.uno_utils import ( + sample_uno_tree_target_tokens, + ) + + target_predict = sample_uno_tree_target_tokens( + next_token_logits=next_token_logits, + sampling_info=sampling_info, + batch_size=bs, + verify_width=verify_input.draft_token_num, + max_top_k=uno_target_max_top_k, + ) + predict, accept_index, num_correct_drafts = verify_tree_greedy_func( + predicts=predict, + accept_index=accept_index, + accept_token_num=num_correct_drafts, + candidates=candidates, + retrieve_index=verify_input.retrieve_index, + retrieve_next_token=verify_input.retrieve_next_token, + retrieve_next_sibling=verify_input.retrieve_next_sibling, + target_predict=target_predict, + topk=verify_input.tree_topk, + ) + + tp_group = ( + get_parallel().attn_tp_group + if is_dp_attention_enabled() + else get_tp_group() + ) + if tp_group.world_size > 1: + tp_group.broadcast(predict, src=0) + tp_group.broadcast(accept_index, src=0) + tp_group.broadcast(num_correct_drafts, src=0) else: from sgl_kernel import ( top_k_renorm_prob, diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py index 91ce8a147..afc843408 100644 --- a/python/sglang/srt/speculative/eagle_worker_common.py +++ b/python/sglang/srt/speculative/eagle_worker_common.py @@ -472,6 +472,7 @@ def run_eagle_verify( metadata_ready_pre_pad: bool, finalize_tree_path: bool, grammar_barrier=None, + uno_target_max_top_k: Optional[int] = None, ) -> GenerationBatchResult: """Shared verify step: target-verify forward, sampling, acceptance bookkeeping. @@ -584,7 +585,13 @@ def run_eagle_verify( predict, accept_lens, accept_index, - ) = eagle_sample(verify_input, batch, logits_output, grammar_mask) + ) = eagle_sample( + verify_input, + batch, + logits_output, + grammar_mask, + uno_target_max_top_k=uno_target_max_top_k, + ) new_seq_lens = batch.seq_lens + accept_lens clear_unaccepted_c128 = getattr( token_to_kv_pool_allocator.get_kvcache(), diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 0cca5c007..30f211fe5 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -37,6 +37,7 @@ class SpeculativeAlgorithm(Enum): """ DFLASH = auto() + UNO = auto() DSPARK = auto() EAGLE = auto() EAGLE3 = auto() @@ -114,6 +115,9 @@ class SpeculativeAlgorithm(Enum): def is_dflash(self) -> bool: return self == SpeculativeAlgorithm.DFLASH + def is_uno(self) -> bool: + return self == SpeculativeAlgorithm.UNO + def is_dspark(self) -> bool: return self == SpeculativeAlgorithm.DSPARK @@ -220,6 +224,7 @@ class SpeculativeAlgorithm(Enum): _handle_eagle_family, _handle_frozen_kv_mtp, _handle_ngram, + _handle_uno, ) # Validate for every algorithm at startup: the metrics paths read the @@ -230,6 +235,8 @@ class SpeculativeAlgorithm(Enum): if self.is_dflash(): _handle_dflash(server_args) + elif self.is_uno(): + _handle_uno(server_args) elif self.is_dspark(): _handle_dspark(server_args) elif self.is_frozen_kv_mtp(): @@ -304,6 +311,11 @@ class SpeculativeAlgorithm(Enum): return DFlashWorkerV2 + if self.is_uno(): + from sglang.srt.speculative.uno_worker_v2 import UnoWorkerV2 + + return UnoWorkerV2 + if self.is_dspark(): from sglang.srt.speculative.dspark_components.dspark_worker_v2 import ( DSparkWorkerV2, @@ -356,6 +368,9 @@ class SpecInputType(IntEnum): DFLASH_DRAFT = auto() DFLASH_VERIFY = auto() NGRAM_VERIFY = auto() + UNO_STATE = auto() + UNO_DRAFT = auto() + UNO_VERIFY = auto() class SpecInput(ABC): @@ -393,6 +408,7 @@ class SpecInput(ABC): SpecInputType.EAGLE_DRAFT_EXTEND, SpecInputType.FROZEN_KV_MTP_DRAFT, SpecInputType.DFLASH_DRAFT, + SpecInputType.UNO_DRAFT, } def is_verify_input(self) -> bool: @@ -401,6 +417,7 @@ class SpecInput(ABC): SpecInputType.FROZEN_KV_MTP_VERIFY, SpecInputType.DFLASH_VERIFY, SpecInputType.NGRAM_VERIFY, + SpecInputType.UNO_VERIFY, } diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py index 62a371813..f7e486712 100644 --- a/python/sglang/srt/speculative/spec_registry.py +++ b/python/sglang/srt/speculative/spec_registry.py @@ -82,6 +82,9 @@ class CustomSpecAlgo: def is_dflash(self) -> bool: return False + def is_uno(self) -> bool: + return False + def is_dspark(self) -> bool: return False diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 464ff501a..95dfb753f 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -1051,6 +1051,12 @@ def spec_prepare_for_decode(batch: ScheduleBatch) -> None: ) if batch.spec_algorithm.is_dflash_family(): batch.spec_info.prepare_for_decode(batch) + elif batch.spec_algorithm.is_uno(): + from sglang.srt.speculative.uno_info import UnoDraftInput + + if not isinstance(batch.spec_info, UnoDraftInput): + raise RuntimeError("UNO decode preparation requires UnoDraftInput") + batch.spec_info.prepare_for_decode(batch) else: from sglang.srt.speculative.eagle_utils import eagle_prepare_for_decode diff --git a/python/sglang/srt/speculative/uno_cuda_graph_runner.py b/python/sglang/srt/speculative/uno_cuda_graph_runner.py new file mode 100644 index 000000000..921353f89 --- /dev/null +++ b/python/sglang/srt/speculative/uno_cuda_graph_runner.py @@ -0,0 +1,214 @@ +import torch + +from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode +from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( + DecodeCudaGraphRunner, +) +from sglang.srt.runtime_context import get_spec +from sglang.srt.speculative.eagle_info import EagleVerifyInput +from sglang.srt.speculative.spec_info import SpecInputType +from sglang.srt.speculative.uno_info import UnoForwardInput +from sglang.srt.speculative.uno_lora import UnoCudaGraphLoRAState + + +class UnoDecodeCudaGraphRunner(DecodeCudaGraphRunner): + """Decode graph runner for linear UNO and both tree forward roles. + + Linear UNO uses two-variant F-wide capture. Tree UNO uses + separate runner instances for its F-wide LoRA draft + and native Q/K EAGLE target verification. + """ + + def __init__( + self, + model_runner, + *, + tree_draft_attn_backend=None, + tree_draft_width=None, + **kwargs, + ): + candidate_top_k = get_spec().speculative_eagle_topk + self._tree_mode = candidate_top_k > 1 + self._tree_draft_mode = tree_draft_width is not None + + if self._tree_draft_mode: + self.record_nolora_graph = False + self._capture_spec_input_type = SpecInputType.UNO_DRAFT + self._lora_state = UnoCudaGraphLoRAState( + model_runner.lora_manager, + model_runner.uno_lora_id, + tree_draft_width, + ) + model_runner.lora_manager.reset_lora_batch() + kwargs.update( + attn_backend=tree_draft_attn_backend, + speculative_num_steps=1, + speculative_num_draft_tokens=tree_draft_width, + ) + super().__init__(model_runner, **kwargs) + model_runner.lora_manager.reset_lora_batch() + return + + if self._tree_mode: + # Capture exactly one base-model target graph. The internal UNO + # adapter is active only in the rejected F-wide draft phase. + self.record_nolora_graph = False + model_runner.lora_manager.reset_lora_batch() + kwargs.update( + attn_backend=model_runner.attn_backend, + speculative_num_steps=get_spec().speculative_num_steps, + speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens, + ) + super().__init__(model_runner, **kwargs) + model_runner.lora_manager.reset_lora_batch() + return + + forward_width = model_runner.decode_num_tokens_per_req() + self.record_nolora_graph = forward_width > 1 + self._capture_spec_input_type = SpecInputType.UNO_VERIFY + self._lora_state = UnoCudaGraphLoRAState( + model_runner.lora_manager, + model_runner.uno_lora_id, + forward_width, + ) + model_runner.lora_manager.reset_lora_batch() + super().__init__(model_runner, **kwargs) + + def capture_prepare(self, size, stream_idx=None, num_tokens=None): + forward_batch, attn_backend, pp_proxy_tensors = super().capture_prepare( + size, stream_idx=stream_idx, num_tokens=num_tokens + ) + # UNO owns token-row routing directly. K2's generic graph runner now + # keys LoRA setup off lora_manager presence, so suppress its synthetic + # request-level base-adapter routing during UNO capture. + forward_batch.lora_ids = None + return forward_batch, attn_backend, pp_proxy_tensors + + def can_run_graph(self, forward_batch): + spec_info = forward_batch.spec_info + if self._tree_draft_mode: + if not isinstance(spec_info, UnoForwardInput): + return False + if spec_info.spec_input_type != SpecInputType.UNO_DRAFT: + return False + return super().can_run_graph(forward_batch) + + if self._tree_mode: + if not isinstance(spec_info, EagleVerifyInput): + return False + return super().can_run_graph(forward_batch) + + if not isinstance(spec_info, UnoForwardInput): + return False + # At F=1 both phases are base-only and share variant_label=None. + if spec_info.spec_input_type not in { + SpecInputType.UNO_DRAFT, + SpecInputType.UNO_VERIFY, + }: + return False + return super().can_run_graph(forward_batch) + + def _resolve_lora_variant(self, forward_batch): + """borrowed technique from multi-LoRA serving + to capture separate graph for each step.""" + if self._tree_mode: + # Tree verification is always the clean target model. Keeping the + # graph keys unlabeled is sufficient because draft and verify own + # separate runners. + return None + if not self.record_nolora_graph: + return None + if forward_batch.spec_info.spec_input_type == SpecInputType.UNO_DRAFT: + return "lora" + return "nolora" + + def capture_one_shape( + self, + size, + forward, + stream_idx=None, + variant_label=None, + dsa_variant=None, + ): + """capture one CUDA graph with/out UNO LoRA.""" + if self._tree_draft_mode: + self._lora_state.capture_draft(size) + try: + return super().capture_one_shape( + size, + forward, + stream_idx, + None, + dsa_variant, + ) + finally: + self._lora_state.reset() + + if self._tree_mode: + self.model_runner.lora_manager.reset_lora_batch() + try: + return super().capture_one_shape( + size, + forward, + stream_idx, + None, + dsa_variant, + ) + finally: + self.model_runner.lora_manager.reset_lora_batch() + + if variant_label == "lora": + self._capture_spec_input_type = SpecInputType.UNO_DRAFT + self._lora_state.capture_draft(size) + else: + self._capture_spec_input_type = SpecInputType.UNO_VERIFY + self._lora_state.reset() + + super().capture_one_shape( + size, + forward, + stream_idx, + variant_label, + dsa_variant, + ) + self._lora_state.reset() + + def get_spec_info(self, num_tokens: int): + if self._tree_draft_mode: + return UnoForwardInput( + spec_input_type=SpecInputType.UNO_DRAFT, + positions=self.buffers.positions[:num_tokens], + draft_token_num=self.captured_req_width, + ) + + if self._tree_mode: + # This deliberately mirrors DecodeCudaGraphRunner's EAGLE capture + # input. Current eagle_prepare_for_verify requests FULL hidden + # capture for every non-STANDALONE algorithm, including UNO. + spec_info = EagleVerifyInput( + draft_token=None, + custom_mask=self.buffers.custom_mask, + positions=None, + retrieve_index=None, + retrieve_next_token=None, + retrieve_next_sibling=None, + retrieve_cum_len=None, + spec_steps=self.speculative_num_steps, + topk=get_spec().speculative_eagle_topk, + draft_token_num=self.speculative_num_draft_tokens, + capture_hidden_mode=CaptureHiddenMode.FULL, + seq_lens_sum=None, + seq_lens_cpu=None, + ) + spec_info.hidden_states = torch.zeros( + (num_tokens, self.model_runner.model_config.hidden_size), + dtype=self.model_runner.dtype, + device=self.model_runner.device, + ) + return spec_info + + return UnoForwardInput( + spec_input_type=self._capture_spec_input_type, + positions=self.buffers.positions[:num_tokens], + draft_token_num=self.captured_req_width, + ) diff --git a/python/sglang/srt/speculative/uno_info.py b/python/sglang/srt/speculative/uno_info.py new file mode 100644 index 000000000..445ccd69c --- /dev/null +++ b/python/sglang/srt/speculative/uno_info.py @@ -0,0 +1,276 @@ +from dataclasses import dataclass, field +from typing import ClassVar, List, Optional + +import torch + +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.mem_cache.allocation import alloc_for_spec_decode +from sglang.srt.mem_cache.allocation_sizing import ( + get_alloc_reserve_per_decode, + page_aligned_decode_alloc_lens, +) +from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode +from sglang.srt.speculative.spec_info import SpecInput, SpecInputType + + +@dataclass +class UnoDraftInput(SpecInput): + """UNO state carried from one engine iteration to the next.""" + + # Previously emitted token at logical position C. It has no KV yet. + bonus_tokens: torch.Tensor + + # Target-correct KV frontier C. + new_seq_lens: torch.Tensor + + # Number of queries in each UNO forward. + forward_width: int + + # FutureMap compatibility. UNO relays only bonus_tokens, but the generic + # relay payload reads these optional Eagle-shaped fields. + topk_p: ClassVar[Optional[torch.Tensor]] = None + topk_index: ClassVar[Optional[torch.Tensor]] = None + hidden_states: ClassVar[Optional[torch.Tensor]] = None + + # Filled by the scheduler after an overlapped dispatch. + future_indices: Optional[torch.Tensor] = None + + # Host-side upper bound for the allocator mapping prepared for this step. + reserved_seq_lens_cpu: Optional[torch.Tensor] = None + reserved_seq_lens_sum: Optional[int] = None + + # Sampling metadata prepared before either internal forward. + max_top_k: int = 1 + uniform_top_k_value: Optional[int] = None + + def __post_init__(self): + super().__init__(SpecInputType.UNO_STATE) + + if self.forward_width < 1: + raise ValueError("UNO forward_width must be positive.") + + # The carried state represents one request row, not an F-row forward. + self.num_tokens_per_req = 1 + self.num_tokens_for_logprob_per_req = 1 + + @property + def tail_width(self) -> int: + return self.forward_width + 1 + + @classmethod + def create_idle_input( + cls, + *, + device, + forward_width: int, + ) -> "UnoDraftInput": + return cls( + bonus_tokens=torch.empty( + (0,), + dtype=torch.int64, + device=device, + ), + new_seq_lens=torch.empty( + (0,), + dtype=torch.int64, + device=device, + ), + forward_width=forward_width, + ) + + def prepare_for_decode(self, batch: ScheduleBatch) -> None: + batch.maybe_evict_swa() + + batch_size = batch.batch_size() + if batch_size == 0: + return + + if self.future_indices is None: + if self.bonus_tokens.numel() != batch_size: + raise RuntimeError("UNO seed count does not match decode batch size.") + + if self.new_seq_lens.numel() != batch_size: + raise RuntimeError( + "UNO frontier count does not match decode batch size." + ) + elif self.future_indices.numel() != batch_size: + raise RuntimeError( + "UNO future-index count does not match decode batch size." + ) + + committed_lengths: list[int] = [] + reserve_width = int(get_alloc_reserve_per_decode()) + + max_top_k = 1 + uniform_top_k_value = None + uniform_top_k = True + + for index, req in enumerate(batch.reqs): + if req.kv is None: + raise RuntimeError("UNO decode request has no KV allocation.") + + committed = int(req.kv.kv_committed_len) + allocated = int(req.kv.kv_allocated_len) + + if allocated < committed: + raise RuntimeError( + "UNO encountered an invalid KV watermark: " + f"committed={committed}, allocated={allocated}." + ) + + committed_lengths.append(committed) + + top_k = int(req.sampling_params.top_k) + max_top_k = max(max_top_k, top_k) + if index == 0: + uniform_top_k_value = top_k + elif uniform_top_k and top_k != uniform_top_k_value: + uniform_top_k = False + + self.max_top_k = max_top_k + self.uniform_top_k_value = uniform_top_k_value if uniform_top_k else None + + page_size = batch.token_to_kv_pool_allocator.page_size + current_lengths, next_lengths, num_needed_tokens = ( + page_aligned_decode_alloc_lens( + batch.reqs, + reserve=reserve_width, + page_size=page_size, + ) + ) + row_width = int(batch.req_to_token_pool.req_to_token.shape[1]) + if max(next_lengths) > row_width: + raise RuntimeError( + "UNO allocation exceeds the req_to_token row: " + f"needed={max(next_lengths)}, available={row_width}." + ) + + current_cpu = torch.tensor( + current_lengths, + dtype=torch.int32, + device="cpu", + ) + next_cpu = torch.tensor( + next_lengths, + dtype=torch.int32, + device="cpu", + ) + current_device = current_cpu.to( + batch.device, + non_blocking=True, + ) + next_device = next_cpu.to( + batch.device, + non_blocking=True, + ) + + alloc_for_spec_decode( + batch.tree_cache, + batch.req_to_token_pool, + reqs=batch.reqs, + req_pool_indices=batch.req_pool_indices, + cur_kv_lens=current_device, + cur_kv_lens_cpu=current_cpu, + nxt_kv_lens=next_device, + nxt_kv_lens_cpu=next_cpu, + num_needed_tokens=num_needed_tokens, + batch=batch, + ) + + for req in batch.reqs: + req.decode_batch_idx += 1 + + batch.seq_lens_cpu = torch.tensor( + committed_lengths, + dtype=torch.int64, + device="cpu", + ) + batch.seq_lens_sum = sum(committed_lengths) + self.reserved_seq_lens_cpu = next_cpu + self.reserved_seq_lens_sum = sum(next_lengths) + + def filter_batch( + self, + new_indices: torch.Tensor, + new_indices_cpu: Optional[List[int]] = None, + ) -> None: + if self.reserved_seq_lens_cpu is not None: + host_indices = ( + new_indices_cpu if new_indices_cpu is not None else new_indices.cpu() + ) + self.reserved_seq_lens_cpu = self.reserved_seq_lens_cpu[host_indices] + self.reserved_seq_lens_sum = int(self.reserved_seq_lens_cpu.sum().item()) + + if self.future_indices is not None: + self.future_indices = self.future_indices[new_indices] + return + + self.bonus_tokens = self.bonus_tokens[new_indices] + self.new_seq_lens = self.new_seq_lens[new_indices] + + def merge_batch(self, other: "UnoDraftInput") -> None: + if not isinstance(other, UnoDraftInput): + raise TypeError(f"Cannot merge UnoDraftInput with {type(other).__name__}.") + + if self.forward_width != other.forward_width: + raise RuntimeError("Cannot merge UNO states with different forward widths.") + + self_has_reservation = self.reserved_seq_lens_cpu is not None + other_has_reservation = other.reserved_seq_lens_cpu is not None + if self_has_reservation != other_has_reservation: + raise RuntimeError("Cannot merge prepared and unprepared UNO states.") + + if self_has_reservation: + self.reserved_seq_lens_cpu = torch.cat( + ( + self.reserved_seq_lens_cpu, + other.reserved_seq_lens_cpu, + ) + ) + self.reserved_seq_lens_sum = int(self.reserved_seq_lens_cpu.sum().item()) + + if self.future_indices is not None: + assert other.future_indices is not None + self.future_indices = torch.cat((self.future_indices, other.future_indices)) + return + + self.bonus_tokens = torch.cat((self.bonus_tokens, other.bonus_tokens)) + self.new_seq_lens = torch.cat((self.new_seq_lens, other.new_seq_lens)) + + +@dataclass +class UnoForwardInput(SpecInput): + """Metadata for one fixed-width UNO forward.""" + + # constructor + spec_input_type: SpecInputType + positions: torch.Tensor + # For UNO, this is forward width F, not proposal count F - 1. + draft_token_num: int + + # expected by interface + custom_mask: Optional[torch.Tensor] = None + capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.NULL + hidden_states: Optional[torch.Tensor] = None + + # derived + num_tokens_per_req: int = field(init=False) + num_tokens_for_logprob_per_req: int = field(init=False) + + def __post_init__(self): + if self.spec_input_type not in { + SpecInputType.UNO_DRAFT, + SpecInputType.UNO_VERIFY, + }: + raise ValueError(f"Invalid UNO input type: {self.spec_input_type}") + + if self.draft_token_num < 1: + raise ValueError("UNO forward width must be positive.") + + # Dataclass-generated __init__ does not call the non-dataclass base + # initializer. This currently reassigns the same field, while also + # preserving the SpecInput initialization contract. + super().__init__(self.spec_input_type) + + self.num_tokens_per_req = self.draft_token_num + self.num_tokens_for_logprob_per_req = self.draft_token_num diff --git a/python/sglang/srt/speculative/uno_lora.py b/python/sglang/srt/speculative/uno_lora.py new file mode 100644 index 000000000..09ebd8750 --- /dev/null +++ b/python/sglang/srt/speculative/uno_lora.py @@ -0,0 +1,97 @@ +"""Loading for UNO's draft LoRA.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from sglang.srt.lora.lora_manager import LoRAManager +from sglang.srt.lora.lora_registry import LoRARef +from sglang.srt.runtime_context import get_spec + +if TYPE_CHECKING: + from sglang.srt.model_executor.model_runner import ModelRunner + + +# This name is internal to the model-execution process. It is never exposed as +# a request-selectable serving adapter. +_UNO_INTERNAL_LORA_NAME = "__uno_draft__" + +# LoRA pool capacity includes the base-model slot. +_UNO_LORA_POOL_CAPACITY = 2 + + +def init_uno_lora_manager( + model_runner: ModelRunner, +) -> tuple[LoRAManager, str]: + """Load and pin the single UNO draft adapter.""" + + lora_path = get_spec().uno_lora_path + + uno_ref = LoRARef( + lora_id=LoRARef.deterministic_id( + _UNO_INTERNAL_LORA_NAME, + lora_path, + ), + lora_name=_UNO_INTERNAL_LORA_NAME, + lora_path=lora_path, + pinned=True, + ) + + manager = LoRAManager( + base_model=model_runner.model, + base_hf_config=model_runner.model_config.hf_config, + max_loras_per_batch=_UNO_LORA_POOL_CAPACITY, + load_config=model_runner.load_config, + dtype=model_runner.dtype, + server_args=model_runner.server_args, + lora_backend="uno_cublas", # fast path + tp_size=model_runner.ps.tp_size, + tp_rank=model_runner.ps.tp_rank, + # Infer these from the one trained adapter. + max_lora_rank=None, + target_modules=None, + lora_paths=[uno_ref], + ) + + # LoRAManager construction initially makes only the base slot resident. + # UNO always needs both fixed choices resident: + # + # None -> base model + # uno_ref.lora_id -> base model + UNO draft LoRA + manager.fetch_new_loras({None, uno_ref.lora_id}) + + return manager, uno_ref.lora_id + + +class UnoCudaGraphLoRAState: + """Retained token-row LoRA routing for UNO draft graph buckets.""" + + def __init__( + self, + manager: LoRAManager, + uno_lora_id: str, + forward_width: int, + ): + self.manager = manager + self.uno_lora_id = uno_lora_id + self.forward_width = forward_width + self._draft_batch_infos = {} + + def capture_draft(self, batch_size: int) -> None: + if batch_size not in self._draft_batch_infos: + self.manager.prepare_lora_token_segments( + lora_ids=[None, self.uno_lora_id] * batch_size, + segment_lens=[1, self.forward_width - 1] * batch_size, + ) + batch_info = self.manager.lora_backend.batch_info + batch_info.use_cuda_graph = True + self._draft_batch_infos[batch_size] = batch_info + + self.activate_draft(batch_size) + + def activate_draft(self, batch_size: int) -> None: + self.manager.reset_lora_batch() + self.manager.lora_backend.batch_info = self._draft_batch_infos[batch_size] + + def reset(self) -> None: + self.manager.reset_lora_batch() diff --git a/python/sglang/srt/speculative/uno_tree.py b/python/sglang/srt/speculative/uno_tree.py new file mode 100644 index 000000000..aa9890482 --- /dev/null +++ b/python/sglang/srt/speculative/uno_tree.py @@ -0,0 +1,584 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from nano-vllm-uno's fixed-budget draft-tree builder for SGLang. + +"""GPU-native UNO proposal-tree construction for SGLang spec-v2. + +UNO owns only proposal ranking and fixed-budget best-first selection. The +result is expressed directly in EAGLE's candidate-lineage ABI so the existing +EAGLE implementation can build masks, positions, traversal links, verify the +tree, sample the accepted path, and compact KV state. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl +from flashinfer import top_k as _flashinfer_top_k +from torch import Tensor + + +@dataclass(frozen=True) +class UnoTreeProposal: + """EAGLE-native representation of one fixed-width UNO tree batch. + + For non-root node ``i``, ``top_scores_index[:, i - 1]`` is the implicit + candidate edge ``parent_node * candidate_top_k + candidate_rank``. + ``parent_list`` maps an EAGLE candidate row back to the edge that created + that row's node. Its rows use EAGLE's native + ``candidate_top_k * (max_depth - 1) + 1`` stride, including unused + padding. The tensors can therefore be passed directly to + ``build_tree_kernel_efficient`` without constructing direct parent arrays. + + Tensors may be backed by the caller's reusable workspace and remain valid + only until that workspace is reused. + """ + + root_tokens: Tensor + draft_tokens: Tensor + parent_list: Tensor + top_scores_index: Tensor + candidate_top_k: int + max_depth: int + + @property + def num_verify_tokens(self) -> int: + return int(self.draft_tokens.shape[1]) + 1 + + +@triton.jit +def _candidate_lse_partials_kernel( + logits, + partial_max, + partial_sum, + temperature_values, + inverse_temperature, + BATCH_STRIDE: tl.constexpr, + DEPTH_STRIDE: tl.constexpr, + VOCAB_STRIDE: tl.constexpr, + NUM_DEPTHS: tl.constexpr, + VOCAB_SIZE: tl.constexpr, + NUM_BLOCKS: tl.constexpr, + BLOCK_VOCAB: tl.constexpr, + TEMPERATURE_BATCH_STRIDE: tl.constexpr, + TEMPERATURE_IS_TENSOR: tl.constexpr, +): + batch = tl.program_id(0) + depth = tl.program_id(1) + block = tl.program_id(2) + row = batch * NUM_DEPTHS + depth + offsets = block * BLOCK_VOCAB + tl.arange(0, BLOCK_VOCAB) + values = tl.load( + logits + batch * BATCH_STRIDE + depth * DEPTH_STRIDE + offsets * VOCAB_STRIDE, + mask=offsets < VOCAB_SIZE, + other=-float("inf"), + ).to(tl.float32) + if TEMPERATURE_IS_TENSOR: + row_temperature = tl.load( + temperature_values + batch * TEMPERATURE_BATCH_STRIDE + ).to(tl.float32) + row_inverse_temperature = tl.where( + row_temperature > 0.0, + 1.0 / row_temperature, + 1.0, + ) + else: + row_inverse_temperature = inverse_temperature + values *= row_inverse_temperature + maximum = tl.max(values, axis=0) + total = tl.sum(tl.exp(values - maximum), axis=0) + output = row * NUM_BLOCKS + block + tl.store(partial_max + output, maximum) + tl.store(partial_sum + output, total) + + +@triton.jit +def _candidate_lse_finalize_kernel( + top_values, + partial_max, + partial_sum, + top_log_probs, + temperature_values, + inverse_temperature, + K: tl.constexpr, + NUM_DEPTHS: tl.constexpr, + NUM_BLOCKS: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_PARTIALS: tl.constexpr, + TEMPERATURE_BATCH_STRIDE: tl.constexpr, + TEMPERATURE_IS_TENSOR: tl.constexpr, +): + row = tl.program_id(0) + batch = row // NUM_DEPTHS + blocks = tl.arange(0, BLOCK_PARTIALS) + block_max = tl.load( + partial_max + row * NUM_BLOCKS + blocks, + mask=blocks < NUM_BLOCKS, + other=-float("inf"), + ) + maximum = tl.max(block_max, axis=0) + block_sum = tl.load( + partial_sum + row * NUM_BLOCKS + blocks, + mask=blocks < NUM_BLOCKS, + other=0.0, + ) + total = tl.sum(block_sum * tl.exp(block_max - maximum), axis=0) + normalizer = maximum + tl.log(total) + + ranks = tl.arange(0, BLOCK_K) + candidates = tl.load( + top_values + row * K + ranks, + mask=ranks < K, + other=-float("inf"), + ).to(tl.float32) + if TEMPERATURE_IS_TENSOR: + row_temperature = tl.load( + temperature_values + batch * TEMPERATURE_BATCH_STRIDE + ).to(tl.float32) + row_inverse_temperature = tl.where( + row_temperature > 0.0, + 1.0 / row_temperature, + 1.0, + ) + else: + row_inverse_temperature = inverse_temperature + + log_probs = candidates * row_inverse_temperature - normalizer + # Proposal mass affects efficiency, not correctness. A malformed row + # must not leave the best-first frontier without a deterministic winner. + log_probs = tl.where(log_probs == log_probs, log_probs, -float("inf")) + log_probs = tl.where(log_probs > 0.0, 0.0, log_probs) + tl.store( + top_log_probs + row * K + ranks, + log_probs, + mask=ranks < K, + ) + + +def _temperature_rows( + temperature: float | Tensor, + *, + batch_size: int, + device: torch.device, +) -> tuple[Tensor, int] | None: + if not isinstance(temperature, Tensor): + return None + if temperature.device != device: + raise ValueError("temperature and logits must share a device") + if not temperature.is_floating_point(): + raise TypeError("temperature tensor must use a floating dtype") + if temperature.ndim == 0: + return temperature.reshape(1), 0 + if temperature.shape not in ((batch_size,), (batch_size, 1)): + raise ValueError( + "temperature tensor must be scalar, [B], or [B, 1]; got " + f"{tuple(temperature.shape)}" + ) + values = temperature.reshape(batch_size) + return values, values.stride(0) + + +def _build_candidate_log_probs( + logits: Tensor, + top_values: Tensor, + top_log_probs: Tensor, + partial_max: Tensor, + partial_sum: Tensor, + temperature: float | Tensor, +) -> None: + """Normalize selected logits over the full vocabulary in FP32.""" + + batch_size, num_depths, vocab_size = logits.shape + num_rows = batch_size * num_depths + candidate_top_k = int(top_values.size(-1)) + block_vocab = 8192 + num_blocks = triton.cdiv(vocab_size, block_vocab) + temperature_rows = _temperature_rows( + temperature, + batch_size=batch_size, + device=logits.device, + ) + if temperature_rows is None: + temperature_values = logits + temperature_batch_stride = 0 + temperature_is_tensor = False + scalar_temperature = float(temperature) + inverse_temperature = ( + 1.0 / scalar_temperature if scalar_temperature > 0.0 else 1.0 + ) + else: + temperature_values, temperature_batch_stride = temperature_rows + temperature_is_tensor = True + inverse_temperature = 1.0 + + _candidate_lse_partials_kernel[(batch_size, num_depths, num_blocks)]( + logits, + partial_max, + partial_sum, + temperature_values, + inverse_temperature, + BATCH_STRIDE=logits.stride(0), + DEPTH_STRIDE=logits.stride(1), + VOCAB_STRIDE=logits.stride(-1), + NUM_DEPTHS=num_depths, + VOCAB_SIZE=vocab_size, + NUM_BLOCKS=num_blocks, + BLOCK_VOCAB=block_vocab, + TEMPERATURE_BATCH_STRIDE=temperature_batch_stride, + TEMPERATURE_IS_TENSOR=temperature_is_tensor, + num_warps=4, + num_stages=1, + ) + _candidate_lse_finalize_kernel[(num_rows,)]( + top_values, + partial_max, + partial_sum, + top_log_probs, + temperature_values, + inverse_temperature, + K=candidate_top_k, + NUM_DEPTHS=num_depths, + NUM_BLOCKS=num_blocks, + BLOCK_K=triton.next_power_of_2(candidate_top_k), + BLOCK_PARTIALS=triton.next_power_of_2(num_blocks), + TEMPERATURE_BATCH_STRIDE=temperature_batch_stride, + TEMPERATURE_IS_TENSOR=temperature_is_tensor, + num_warps=1, + num_stages=1, + ) + + +@triton.jit +def _build_tree_kernel( + top_token_ids, + top_log_probs, + draft_tokens, + selected_edges, + parent_list, + search_depths, + search_log_masses, + TOKEN_BATCH_STRIDE: tl.constexpr, + TOKEN_DEPTH_STRIDE: tl.constexpr, + TOKEN_RANK_STRIDE: tl.constexpr, + PROB_BATCH_STRIDE: tl.constexpr, + PROB_DEPTH_STRIDE: tl.constexpr, + PROB_RANK_STRIDE: tl.constexpr, + PARENT_BATCH_STRIDE: tl.constexpr, + NUM_DEPTHS: tl.constexpr, + K: tl.constexpr, + Q: tl.constexpr, + BLOCK_CANDIDATES: tl.constexpr, +): + batch = tl.program_id(0) + search_offset = batch * Q + edge_output_offset = batch * (Q - 1) + parent_output_offset = batch * PARENT_BATCH_STRIDE + candidate_slots = tl.arange(0, BLOCK_CANDIDATES) + candidate_parents = candidate_slots // K + candidate_ranks = candidate_slots % K + candidate_in_bounds = candidate_slots < Q * K + used = tl.zeros((BLOCK_CANDIDATES,), dtype=tl.int1) + + # The root token itself is returned by reference. Only its private search + # state is stored here; EAGLE prepends it as the verify-tree root later. + tl.store( + parent_list + parent_output_offset + candidate_slots, + -1, + mask=candidate_slots < PARENT_BATCH_STRIDE, + ) + tl.store(search_depths + search_offset, 0) + tl.store(search_log_masses + search_offset, 0.0) + tl.debug_barrier() + + for node_index in range(1, Q): + parent_valid = candidate_in_bounds & (candidate_parents < node_index) + safe_parent = tl.where(parent_valid, candidate_parents, 0) + parent_depth = tl.load( + search_depths + search_offset + safe_parent, + mask=parent_valid, + other=NUM_DEPTHS, + ) + parent_mass = tl.load( + search_log_masses + search_offset + safe_parent, + mask=parent_valid, + other=-float("inf"), + ).to(tl.float32) + valid = parent_valid & (parent_depth < NUM_DEPTHS) & ~used + safe_depth = tl.where(valid, parent_depth, 0) + token = tl.load( + top_token_ids + + batch * TOKEN_BATCH_STRIDE + + safe_depth * TOKEN_DEPTH_STRIDE + + candidate_ranks * TOKEN_RANK_STRIDE, + mask=valid, + other=0, + ) + log_prob = tl.load( + top_log_probs + + batch * PROB_BATCH_STRIDE + + safe_depth * PROB_DEPTH_STRIDE + + candidate_ranks * PROB_RANK_STRIDE, + mask=valid, + other=-float("inf"), + ).to(tl.float32) + mass = parent_mass + log_prob + child_depth = parent_depth + 1 + + # Deterministic best-first order: mass, shallower depth, lower rank, + # lower token ID, then lower parent node. + best_mass = tl.max(tl.where(valid, mass, -float("inf")), axis=0) + winner = valid & (mass == best_mass) + best_depth = tl.min(tl.where(winner, child_depth, 1 << 30), axis=0) + winner &= child_depth == best_depth + best_rank = tl.min(tl.where(winner, candidate_ranks, 1 << 30), axis=0) + winner &= candidate_ranks == best_rank + best_token = tl.min(tl.where(winner, token, 1 << 30), axis=0) + winner &= token == best_token + best_parent = tl.min(tl.where(winner, candidate_parents, 1 << 30), axis=0) + winner &= candidate_parents == best_parent + best_slot = tl.min(tl.where(winner, candidate_slots, 1 << 30), axis=0) + + selected_parent = best_slot // K + selected_rank = best_slot % K + selected_depth = tl.load(search_depths + search_offset + selected_parent) + selected_mass = tl.load(search_log_masses + search_offset + selected_parent).to( + tl.float32 + ) + tl.load( + top_log_probs + + batch * PROB_BATCH_STRIDE + + selected_depth * PROB_DEPTH_STRIDE + + selected_rank * PROB_RANK_STRIDE + ).to(tl.float32) + selected_token = tl.load( + top_token_ids + + batch * TOKEN_BATCH_STRIDE + + selected_depth * TOKEN_DEPTH_STRIDE + + selected_rank * TOKEN_RANK_STRIDE + ) + + output_index = edge_output_offset + node_index - 1 + tl.store(draft_tokens + output_index, selected_token) + # This is already EAGLE's implicit selected-edge encoding. + tl.store(selected_edges + output_index, best_slot) + if node_index < Q - 1: + # EAGLE shifts each selected edge by one candidate row. Fusing + # this write avoids separate fill/copy launches on every step. + tl.store( + parent_list + parent_output_offset + node_index, + best_slot, + ) + tl.store( + search_depths + search_offset + node_index, + selected_depth + 1, + ) + tl.store( + search_log_masses + search_offset + node_index, + selected_mass, + ) + used |= candidate_slots == best_slot + tl.debug_barrier() + + +def _candidate_tree_capacity( + num_depths: int, + candidate_top_k: int, + stop_at: int, +) -> int: + capacity = 1 + width = 1 + for _ in range(num_depths): + width *= candidate_top_k + capacity += width + if capacity >= stop_at: + break + return capacity + + +@torch.inference_mode() +def build_uno_tree_proposal( + root_tokens: Tensor, + draft_logits: Tensor, + *, + max_nodes: int, + candidate_top_k: int, + temperature: float | Tensor, + workspace: dict[str, Tensor] | None = None, +) -> UnoTreeProposal: + """Build fixed-``Q`` UNO trees directly in EAGLE's proposal ABI. + + ``root_tokens`` is ``[B]`` and ``draft_logits`` is ``[B, F-1, V]``. + Candidate log probabilities are normalized over the full vocabulary. No + CPU reference/fallback, direct parent array, attention mask, traversal + structure, acceptance walk, or KV operation is implemented here. + """ + + if root_tokens.ndim != 1: + raise ValueError( + f"root_tokens must have shape [B], got {tuple(root_tokens.shape)}" + ) + if draft_logits.ndim != 3 or draft_logits.size(0) != root_tokens.size(0): + raise ValueError( + "draft_logits must have shape [B, depth, vocab] with the same B " + f"as root_tokens; got {tuple(draft_logits.shape)}" + ) + if not root_tokens.is_cuda or not draft_logits.is_cuda: + raise ValueError("UNO tree construction requires CUDA tensors") + if root_tokens.device != draft_logits.device: + raise ValueError("tree roots and draft logits must share a device") + if root_tokens.dtype not in (torch.int32, torch.int64): + raise TypeError("tree roots must use an integer dtype") + if not draft_logits.is_floating_point(): + raise TypeError("draft logits must use a floating dtype") + if root_tokens.numel() == 0: + raise ValueError("UNO tree construction requires a non-empty batch") + if max_nodes < 1: + raise ValueError("max_nodes must include at least the root") + if candidate_top_k < 1: + raise ValueError("candidate_top_k must be >= 1") + + batch_size = int(root_tokens.size(0)) + num_depths = int(draft_logits.size(1)) + vocab_size = int(draft_logits.size(2)) + candidate_top_k = int(candidate_top_k) + draft_width = num_depths + 1 + parent_width = candidate_top_k * max(num_depths - 1, 0) + 1 + if max_nodes < draft_width: + raise ValueError( + f"max_nodes Q must be >= draft width F; got Q={max_nodes}, F={draft_width}" + ) + if candidate_top_k > vocab_size: + raise ValueError( + f"candidate_top_k ({candidate_top_k}) exceeds vocabulary size " + f"({vocab_size})" + ) + if max_nodes > 128 or max_nodes * candidate_top_k > 2048: + raise ValueError( + "the initial single-program UNO builder requires Q <= 128 and " + f"Q*K <= 2048; got Q={max_nodes}, K={candidate_top_k}" + ) + capacity = _candidate_tree_capacity( + num_depths, + candidate_top_k, + max_nodes, + ) + if capacity < max_nodes: + raise ValueError( + f"candidate set can produce only {capacity} tree nodes, but " + f"fixed tree verification requires {max_nodes}" + ) + if max_nodes - 1 > parent_width: + raise ValueError( + "EAGLE's parent-list ABI cannot represent this UNO tree: " + f"Q-1={max_nodes - 1} exceeds K*(depth-1)+1={parent_width}" + ) + _temperature_rows( + temperature, + batch_size=batch_size, + device=draft_logits.device, + ) + + def buffer( + name: str, + shape: tuple[int, ...], + dtype: torch.dtype, + ) -> Tensor: + if workspace is None: + return torch.empty( + shape, + dtype=dtype, + device=draft_logits.device, + ) + value = workspace.get(name) + if ( + value is None + or value.shape != shape + or value.dtype != dtype + or value.device != draft_logits.device + ): + value = torch.empty( + shape, + dtype=dtype, + device=draft_logits.device, + ) + workspace[name] = value + return value + + edge_shape = (batch_size, max_nodes - 1) + draft_tokens = buffer("draft_tokens", edge_shape, torch.long) + selected_edges = buffer("top_scores_index", edge_shape, torch.long) + parent_list = buffer( + "parent_list", + (batch_size, parent_width), + torch.long, + ) + + if max_nodes == 1: + parent_list.fill_(-1) + return UnoTreeProposal( + root_tokens=root_tokens, + draft_tokens=draft_tokens, + parent_list=parent_list, + top_scores_index=selected_edges, + candidate_top_k=candidate_top_k, + max_depth=num_depths, + ) + + candidate_shape = (batch_size, num_depths, candidate_top_k) + flat_logits = draft_logits.contiguous().view(batch_size * num_depths, vocab_size) + flat_top_values, flat_top_token_ids = _flashinfer_top_k( + flat_logits, + candidate_top_k, + sorted=True, + deterministic=False, + ) + top_values = flat_top_values.view(candidate_shape) + top_token_ids = buffer("top_token_ids", candidate_shape, torch.long) + top_token_ids.copy_(flat_top_token_ids.view(candidate_shape)) + top_log_probs = buffer("top_log_probs", candidate_shape, torch.float32) + num_partial_blocks = (vocab_size + 8191) // 8192 + partial_shape = (batch_size * num_depths, num_partial_blocks) + _build_candidate_log_probs( + draft_logits, + top_values, + top_log_probs, + buffer("partial_lse_max", partial_shape, torch.float32), + buffer("partial_lse_sum", partial_shape, torch.float32), + temperature, + ) + + search_shape = (batch_size, max_nodes) + search_depths = buffer("search_depths", search_shape, torch.int32) + search_log_masses = buffer("search_log_masses", search_shape, torch.float32) + _build_tree_kernel[(batch_size,)]( + top_token_ids, + top_log_probs, + draft_tokens, + selected_edges, + parent_list, + search_depths, + search_log_masses, + TOKEN_BATCH_STRIDE=top_token_ids.stride(0), + TOKEN_DEPTH_STRIDE=top_token_ids.stride(1), + TOKEN_RANK_STRIDE=top_token_ids.stride(2), + PROB_BATCH_STRIDE=top_log_probs.stride(0), + PROB_DEPTH_STRIDE=top_log_probs.stride(1), + PROB_RANK_STRIDE=top_log_probs.stride(2), + PARENT_BATCH_STRIDE=parent_list.stride(0), + NUM_DEPTHS=num_depths, + K=candidate_top_k, + Q=max_nodes, + BLOCK_CANDIDATES=triton.next_power_of_2(max_nodes * candidate_top_k), + num_warps=8, + num_stages=1, + ) + + return UnoTreeProposal( + root_tokens=root_tokens, + draft_tokens=draft_tokens, + parent_list=parent_list, + top_scores_index=selected_edges, + candidate_top_k=candidate_top_k, + max_depth=num_depths, + ) diff --git a/python/sglang/srt/speculative/uno_utils.py b/python/sglang/srt/speculative/uno_utils.py new file mode 100644 index 000000000..b91ec88e7 --- /dev/null +++ b/python/sglang/srt/speculative/uno_utils.py @@ -0,0 +1,587 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +from flashinfer import top_k as _flashinfer_top_k + +from sglang.kernels.ops.speculative.reject_sampling import ( + chain_speculative_sampling_triton, +) +from sglang.srt.speculative.dflash_utils import ( + _get_or_create_chain_verify_buffers, + build_dflash_verify_target_probs, +) +from sglang.srt.speculative.spec_utils import fast_sample + +_SPARSE_TOP_K_LIMIT = 128 + + +def _normalize_sparse_topk_probs( + topk_logits: torch.Tensor, + temperatures: torch.Tensor, + valid: torch.Tensor, + top_ps: torch.Tensor, +) -> torch.Tensor: + """Normalize a compact top-k support with top-k-first top-p semantics.""" + scaled = topk_logits.float() / temperatures + scaled = scaled.masked_fill(~valid, float("-inf")) + probs = torch.softmax(scaled, dim=-1) + cdf = torch.cumsum(probs, dim=-1) + probs = probs.masked_fill((cdf - probs) > top_ps, 0.0) + return probs / probs.sum(dim=-1, keepdim=True).clamp_min(1e-12) + + +@torch.compile(dynamic=True) +def _sparse_rejection_from_support( + candidates: torch.Tensor, + target_ids: torch.Tensor, + target_probs: torch.Tensor, + draft_ids: torch.Tensor, + draft_probs: torch.Tensor, + accept_uniforms: torch.Tensor, + final_uniforms: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run exact p/q rejection sampling on compact linear-chain supports.""" + batch_size, forward_width = candidates.shape + num_proposals = forward_width - 1 + + if num_proposals == 0: + final_ids = target_ids[:, 0] + final_probs = target_probs[:, 0] + accepted_counts = torch.zeros( + batch_size, + dtype=torch.int32, + device=candidates.device, + ) + else: + proposal_ids = candidates[:, 1:] + p_proposal = torch.where( + target_ids[:, :num_proposals].eq(proposal_ids.unsqueeze(-1)), + target_probs[:, :num_proposals], + torch.zeros( + (), + dtype=target_probs.dtype, + device=target_probs.device, + ), + ).sum(dim=-1) + q_proposal = torch.where( + draft_ids.eq(proposal_ids.unsqueeze(-1)), + draft_probs, + torch.zeros( + (), + dtype=draft_probs.dtype, + device=draft_probs.device, + ), + ).sum(dim=-1) + ratios = torch.where( + q_proposal > 0, + p_proposal / q_proposal, + torch.zeros_like(p_proposal), + ).clamp_(max=1.0) + accepted_flags = accept_uniforms[:, :num_proposals] < ratios + accepted_counts = accepted_flags.to(torch.int32).cumprod(dim=1).sum(dim=1) + + batch_indices = torch.arange(batch_size, device=candidates.device) + final_rows = accepted_counts.to(torch.long) + final_ids = target_ids[batch_indices, final_rows] + final_probs = target_probs[batch_indices, final_rows] + + rejected = final_rows < num_proposals + draft_rows = final_rows.clamp(max=num_proposals - 1) + final_draft_ids = draft_ids[batch_indices, draft_rows] + final_draft_probs = draft_probs[batch_indices, draft_rows] + q_on_target = torch.where( + final_ids.unsqueeze(2).eq(final_draft_ids.unsqueeze(1)), + final_draft_probs.unsqueeze(1), + torch.zeros( + (), + dtype=final_draft_probs.dtype, + device=final_draft_probs.device, + ), + ).sum(dim=2) + correction_probs = (final_probs - q_on_target).clamp_min_(0.0) + correction_sum = correction_probs.sum(dim=1, keepdim=True) + correction_probs = torch.where( + correction_sum > 0, + correction_probs / correction_sum.clamp_min(1e-12), + final_probs, + ) + final_probs = torch.where( + rejected[:, None], + correction_probs, + final_probs, + ) + + cdf = torch.cumsum(final_probs, dim=-1) + thresholds = final_uniforms * final_probs.sum(dim=-1) + sampled_offsets = (cdf <= thresholds[:, None]).sum(dim=-1) + sampled_offsets.clamp_(max=final_ids.shape[-1] - 1) + bonus = ( + final_ids.gather(1, sampled_offsets[:, None]).squeeze(1).to(candidates.dtype) + ) + return accepted_counts, bonus + + +@torch.compile(dynamic=True) +def _build_sparse_target_support_tensors( + next_token_logits: torch.Tensor, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + batch_size: int, + forward_width: int, + max_top_k: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build compact support using the fastest available top-k primitive.""" + rows = batch_size * forward_width + topk_logits, topk_ids = _flashinfer_top_k( + next_token_logits.contiguous(), + max_top_k, + sorted=True, + deterministic=False, + ) + + expanded_temperatures = torch.repeat_interleave( + temperatures, + forward_width, + dim=0, + ).reshape(rows, -1) + expanded_top_ks = torch.repeat_interleave( + top_ks, + forward_width, + dim=0, + ).reshape(rows, 1) + expanded_top_ps = torch.repeat_interleave( + top_ps, + forward_width, + dim=0, + ).reshape(rows, 1) + ranks = torch.arange( + max_top_k, + dtype=expanded_top_ks.dtype, + device=next_token_logits.device, + )[None, :] + probs = _normalize_sparse_topk_probs( + topk_logits, + expanded_temperatures, + ranks < expanded_top_ks, + expanded_top_ps, + ) + return ( + topk_ids.view(batch_size, forward_width, max_top_k), + probs.view(batch_size, forward_width, max_top_k), + ) + + +def _build_sparse_target_support( + *, + next_token_logits: torch.Tensor, + sampling_info: Any, + batch_size: int, + forward_width: int, + max_top_k: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return compact target token IDs/probabilities without a dense scatter.""" + if bool(getattr(sampling_info, "need_top_p_sampling", False)): + top_ps = sampling_info.top_ps + else: + top_ps = torch.ones( + (batch_size,), + dtype=torch.float32, + device=next_token_logits.device, + ) + + return _build_sparse_target_support_tensors( + next_token_logits, + sampling_info.temperatures, + sampling_info.top_ks, + top_ps, + batch_size, + forward_width, + max_top_k, + ) + + +def _sample_from_support( + support_ids: torch.Tensor, + support_probs: torch.Tensor, +) -> torch.Tensor: + """Sample one token from every compact support row.""" + flat_ids = support_ids.flatten(0, 1) + flat_probs = support_probs.flatten(0, 1) + _, offsets = fast_sample(flat_probs) + return flat_ids.gather(1, offsets).view(support_ids.shape[:2]) + + +@dataclass(frozen=True) +class UnoDraftDistribution: + """The exact q used to sample future UNO proposal rows.""" + + probs: torch.Tensor + token_ids: torch.Tensor | None = None + + +def _run_sparse_rejection( + *, + candidates: torch.Tensor, + next_token_logits: torch.Tensor, + sampling_info: Any, + max_top_k: int, + draft_distribution: UnoDraftDistribution, +) -> tuple[torch.Tensor, torch.Tensor]: + if draft_distribution.token_ids is None: + raise RuntimeError("Sparse UNO verification requires sparse draft q.") + + batch_size, forward_width = candidates.shape + support_ids, support_probs = _build_sparse_target_support( + next_token_logits=next_token_logits, + sampling_info=sampling_info, + batch_size=batch_size, + forward_width=forward_width, + max_top_k=max_top_k, + ) + + # Preserve the legacy path's two RNG draws and tensor shapes. The final + # uniforms select the correction/bonus; the first F - 1 acceptance coins + # are consumed by a linear chain. + accept_uniforms = torch.rand( + (batch_size, forward_width), + dtype=torch.float32, + device=next_token_logits.device, + ) + final_uniforms = torch.rand( + (batch_size,), + dtype=torch.float32, + device=next_token_logits.device, + ) + return _sparse_rejection_from_support( + candidates, + support_ids, + support_probs, + draft_distribution.token_ids, + draft_distribution.probs, + accept_uniforms, + final_uniforms, + ) + + +def _build_dense_probs( + *, + next_token_logits: torch.Tensor, + sampling_info: Any, + batch_size: int, + forward_width: int, + max_top_k: int, + uniform_top_k_value: int | None, +) -> torch.Tensor: + """Build the dense sampling distribution used by SGLang verification.""" + return build_dflash_verify_target_probs( + next_token_logits=next_token_logits, + sampling_info=sampling_info, + draft_token_num=forward_width, + bs=batch_size, + max_top_k=max_top_k, + uniform_top_k_value=uniform_top_k_value, + use_sparse_topk=True, + ) + + +def _sample_from_dense_probs(probs: torch.Tensor) -> torch.Tensor: + """Sample one token from every dense distribution row.""" + _, token_ids = fast_sample(probs.flatten(0, 1)) + return token_ids.view(probs.shape[:2]) + + +def _run_dense_rejection( + *, + candidates: torch.Tensor, + next_token_logits: torch.Tensor, + sampling_info: Any, + max_top_k: int, + uniform_top_k_value: int | None, + draft_distribution: UnoDraftDistribution, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run SGLang's fused linear-chain p/q rejection kernel.""" + if draft_distribution.token_ids is not None: + raise RuntimeError("Dense UNO verification requires dense draft q.") + + batch_size, forward_width = candidates.shape + target_probs = _build_dense_probs( + next_token_logits=next_token_logits, + sampling_info=sampling_info, + batch_size=batch_size, + forward_width=forward_width, + max_top_k=max_top_k, + uniform_top_k_value=uniform_top_k_value, + ) + accept_uniforms = torch.rand( + (batch_size, forward_width), + dtype=torch.float32, + device=next_token_logits.device, + ) + final_uniforms = torch.rand( + (batch_size,), + dtype=torch.float32, + device=next_token_logits.device, + ) + ( + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + predicts, + accept_index, + accepted_counts, + ) = _get_or_create_chain_verify_buffers( + bs=batch_size, + draft_token_num=forward_width, + device=next_token_logits.device, + ) + chain_speculative_sampling_triton( + predicts=predicts, + accept_index=accept_index, + accept_token_num=accepted_counts, + candidates=candidates, + retrive_index=retrieve_index, + retrive_next_token=retrieve_next_token, + retrive_next_sibling=retrieve_next_sibling, + uniform_samples=accept_uniforms, + uniform_samples_for_final_sampling=final_uniforms, + target_probs=target_probs, + draft_probs=draft_distribution.probs, + threshold_single=1.0, + threshold_acc=1.0, + deterministic=True, + ) + + rows = torch.arange(batch_size, device=candidates.device) + bonus_positions = accept_index[ + rows, + accepted_counts.to(torch.long), + ].to(torch.long) + bonus = predicts[bonus_positions].to(candidates.dtype) + return accepted_counts, bonus + + +@dataclass +class UnoSamplingResult: + output_ids: torch.Tensor + accept_lens: torch.Tensor + new_seq_lens: torch.Tensor + next_seed_tokens: torch.Tensor + + +@dataclass +class UnoTreeSamplingResult: + output_ids: torch.Tensor + accept_lens: torch.Tensor + + +def build_uno_draft_input( + *, + seed_tokens: torch.Tensor, + forward_width: int, + vocab_size: int, + noise_tokens: torch.Tensor | None = None, # for testing +) -> torch.Tensor: + """Build one ``[seed, uniform noise...]`` row per request. + + Random noise is sampled independently from ``[0, vocab_size)``. + + Supplying ``noise_tokens`` bypasses random generation. This is used + by deterministic tests and must have shape + ``(batch_size, forward_width - 1)``. + + ``forward_width == 1`` never generates noise or consumes RNG state. + """ + seed_tokens = seed_tokens.reshape(-1).to(dtype=torch.int64) + batch_size = seed_tokens.numel() + noise_shape = (batch_size, forward_width - 1) + + if forward_width == 1: + return seed_tokens[:, None] + + if noise_tokens is None: + noise_tokens = torch.randint( + low=0, + high=vocab_size, + size=noise_shape, + dtype=torch.int64, + device=seed_tokens.device, + ) + else: + noise_tokens = noise_tokens.to( + device=seed_tokens.device, + dtype=torch.int64, + ) + + draft_input_ids = seed_tokens.new_empty((batch_size, forward_width)) + draft_input_ids[:, 0].copy_(seed_tokens) + draft_input_ids[:, 1:].copy_(noise_tokens) + + return draft_input_ids + + +def sample_uno_candidates( + *, + draft_logits: torch.Tensor, # [B, F, V] + sampling_info: Any, + max_top_k: int, + uniform_top_k_value: int | None = None, +) -> tuple[torch.Tensor, UnoDraftDistribution]: + """Sample 1 clean token from seed and F-1 draft tokens. + Depending on max_top_k value, may use sparse representations for efficiency. + candidates: [B, F] sampled tokens including clean and draft. + draft_distribution: probabilities of the draft tokens for rejection sampling. + """ + batch_size, forward_width, vocab_size = draft_logits.shape + flat_logits = draft_logits.reshape(-1, vocab_size) + if max_top_k <= _SPARSE_TOP_K_LIMIT: + support_ids, support_probs = _build_sparse_target_support( + next_token_logits=flat_logits, + sampling_info=sampling_info, + batch_size=batch_size, + forward_width=forward_width, + max_top_k=max_top_k, + ) + candidates = _sample_from_support(support_ids, support_probs) + draft_distribution = UnoDraftDistribution( + token_ids=support_ids[:, 1:], + probs=support_probs[:, 1:], + ) + else: + probs = _build_dense_probs( + next_token_logits=flat_logits, + sampling_info=sampling_info, + batch_size=batch_size, + forward_width=forward_width, + max_top_k=max_top_k, + uniform_top_k_value=uniform_top_k_value, + ) + candidates = _sample_from_dense_probs(probs) + draft_distribution = UnoDraftDistribution(probs=probs[:, 1:]) + return candidates, draft_distribution + + +def sample_uno_clean_root( + *, + seed_tokens: torch.Tensor, + draft_logits: torch.Tensor, + sampling_info: Any, + max_top_k: int, + uniform_top_k_value: int | None = None, +) -> torch.Tensor: + """Sample the clean root using the current UNO sampling path.""" + del seed_tokens + candidates, _ = sample_uno_candidates( + draft_logits=draft_logits[:, :1, :].contiguous(), + sampling_info=sampling_info, + max_top_k=max_top_k, + uniform_top_k_value=uniform_top_k_value, + ) + return candidates[:, 0] + + +def sample_uno_tree_target_tokens( + *, + next_token_logits: torch.Tensor, + sampling_info: Any, + batch_size: int, + verify_width: int, + max_top_k: int, +) -> torch.Tensor: + """Sample one target token per verify node from compact top-k support.""" + support_ids, support_probs = _build_sparse_target_support( + next_token_logits=next_token_logits, + sampling_info=sampling_info, + batch_size=batch_size, + forward_width=verify_width, + max_top_k=max_top_k, + ) + return _sample_from_support(support_ids, support_probs) + + +def pack_uno_tree_result( + *, + clean_root_tokens: torch.Tensor, + eagle_predict: torch.Tensor, + eagle_accept_lens: torch.Tensor, + draft_width: int, +) -> UnoTreeSamplingResult: + """Convert an internal EAGLE tree result into UNO's public row.""" + clean_root_tokens = clean_root_tokens.reshape(-1).to(dtype=torch.int64) + batch_size = clean_root_tokens.numel() + verify_width = eagle_predict.numel() // batch_size + predict_rows = eagle_predict.reshape(batch_size, verify_width) + + output_ids = clean_root_tokens.new_zeros((batch_size, draft_width + 1)) + output_ids[:, 0].copy_(clean_root_tokens) + output_ids[:, 1:].copy_(predict_rows[:, :draft_width].to(dtype=output_ids.dtype)) + return UnoTreeSamplingResult( + output_ids=output_ids, + accept_lens=eagle_accept_lens + 1, + ) + + +def pack_uno_result( + *, + candidates: torch.Tensor, # [B, F] + accepted_proposal_counts: torch.Tensor, # [B] + bonus_tokens: torch.Tensor, # [B] + committed_frontiers: torch.Tensor, # [B] +) -> UnoSamplingResult: + """Pack acceptance into fixed-width UNO output rows.""" + batch_size, forward_width = candidates.shape + output_ids = candidates.new_zeros((batch_size, forward_width + 1)) + output_ids[:, :forward_width].copy_(candidates) + output_ids.scatter_( + 1, + (accepted_proposal_counts.to(torch.long) + 1)[:, None], + bonus_tokens[:, None], + ) + + accept_lens = accepted_proposal_counts + 2 + new_seq_lens = committed_frontiers + accept_lens.to(committed_frontiers.dtype) + return UnoSamplingResult( + output_ids=output_ids, + accept_lens=accept_lens, + new_seq_lens=new_seq_lens, + next_seed_tokens=bonus_tokens, + ) + + +def run_uno_sampling( + *, + candidates: torch.Tensor, # [B, F] + next_token_logits: torch.Tensor, # [B x F, V] + sampling_info: Any, + committed_frontiers: torch.Tensor, # [B] + draft_distribution: UnoDraftDistribution, + max_top_k: int, + uniform_top_k_value: int | None = None, +) -> UnoSamplingResult: + """Verify sampled UNO proposals against target p and pack the result.""" + if draft_distribution.token_ids is not None: + accepted, bonus = _run_sparse_rejection( + candidates=candidates, + next_token_logits=next_token_logits, + sampling_info=sampling_info, + max_top_k=max_top_k, + draft_distribution=draft_distribution, + ) + else: + accepted, bonus = _run_dense_rejection( + candidates=candidates, + next_token_logits=next_token_logits, + sampling_info=sampling_info, + draft_distribution=draft_distribution, + max_top_k=max_top_k, + uniform_top_k_value=uniform_top_k_value, + ) + return pack_uno_result( + candidates=candidates, + accepted_proposal_counts=accepted, + bonus_tokens=bonus, + committed_frontiers=committed_frontiers, + ) diff --git a/python/sglang/srt/speculative/uno_validation.py b/python/sglang/srt/speculative/uno_validation.py new file mode 100644 index 000000000..436491fbd --- /dev/null +++ b/python/sglang/srt/speculative/uno_validation.py @@ -0,0 +1,50 @@ +"""Request-admission validation for UNO speculative decoding.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + + +def validate_uno_request(req: Req) -> Optional[str]: + """Return an error for request features that UNO cannot execute.""" + + sampling_params = req.sampling_params + + if sampling_params.min_p > 0.0: + return "UNO speculative decoding does not support min_p sampling." + + has_grammar = req.grammar is not None or any( + getattr(sampling_params, field) is not None + for field in ("json_schema", "regex", "ebnf", "structural_tag") + ) + if has_grammar: + return "UNO speculative decoding does not support grammar decoding." + + if req.return_logprob: + return "UNO speculative decoding does not support returned logprobs." + + if req.return_hidden_states_mode.need_capture(): + return "UNO speculative decoding does not support return_hidden_states." + + has_penalties = ( + sampling_params.frequency_penalty != 0.0 + or sampling_params.presence_penalty != 0.0 + or sampling_params.repetition_penalty != 1.0 + or sampling_params.min_new_tokens > 0 + ) + if has_penalties: + return "UNO speculative decoding does not support sampling penalties." + + if sampling_params.logit_bias is not None: + return "UNO speculative decoding does not support logit_bias." + + if req.custom_logit_processor: + return "UNO speculative decoding does not support custom logit processors." + + if req.lora_id is not None: + return "UNO speculative decoding does not support request-selectable LoRA." + + return None diff --git a/python/sglang/srt/speculative/uno_worker_v2.py b/python/sglang/srt/speculative/uno_worker_v2.py new file mode 100644 index 000000000..ce87e4741 --- /dev/null +++ b/python/sglang/srt/speculative/uno_worker_v2.py @@ -0,0 +1,867 @@ +from __future__ import annotations + +import contextlib +import copy +import logging +import time +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.managers.utils import GenerationBatchResult +from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + ForwardBatch, + ForwardMode, +) +from sglang.srt.model_executor.forward_context import ( + ForwardContext, + forward_context, +) +from sglang.srt.runtime_context import get_schedule, get_spec +from sglang.srt.speculative.base_spec_worker import BaseSpecWorker +from sglang.srt.speculative.eagle_info import EagleDraftInput +from sglang.srt.speculative.eagle_utils import default_tree_mask_mode +from sglang.srt.speculative.eagle_worker_common import ( + build_eagle_verify_input, + run_eagle_verify, +) +from sglang.srt.speculative.spec_info import SpecInputType, SpeculativeAlgorithm +from sglang.srt.speculative.spec_utils import get_plan_stream +from sglang.srt.speculative.uno_cuda_graph_runner import ( + UnoDecodeCudaGraphRunner, +) +from sglang.srt.speculative.uno_info import UnoDraftInput, UnoForwardInput +from sglang.srt.speculative.uno_tree import build_uno_tree_proposal +from sglang.srt.speculative.uno_utils import ( + build_uno_draft_input, + pack_uno_tree_result, + run_uno_sampling, + sample_uno_candidates, + sample_uno_clean_root, +) +from sglang.srt.utils.common import ( + get_available_gpu_memory, + log_info_on_rank0, +) + +if TYPE_CHECKING: + from sglang.srt.distributed.parallel_state_wrapper import ParallelState + from sglang.srt.managers.schedule_batch import ScheduleBatch + from sglang.srt.managers.tp_worker import TpModelWorker + from sglang.srt.server_args import ServerArgs + + +logger = logging.getLogger(__name__) + + +class UnoWorkerV2(BaseSpecWorker): + """Single-model UNO worker with linear and native-EAGLE tree decode.""" + + def __init__( + self, + server_args: ServerArgs, + gpu_id: int, + ps: ParallelState, + nccl_port: int, + target_worker: TpModelWorker, + ): + super().__init__() + + self.server_args = server_args + self.gpu_id = gpu_id + self.ps = ps + self.nccl_port = nccl_port + + self._target_worker = target_worker + self._draft_worker = None + + self.model_runner = target_worker.model_runner + self.lora_manager = self.model_runner.lora_manager + self.uno_lora_id = self.model_runner.uno_lora_id + self.device = target_worker.device + + self.enable_overlap = not get_schedule().disable_overlap_schedule + configured_topk = int(get_spec().speculative_eagle_topk or 1) + self.tree_mode = configured_topk > 1 + # Linear UNO stores F in speculative_num_draft_tokens. Tree UNO reuses + # EAGLE's native dimensions: F=steps+1, K=eagle_topk, Q=draft_tokens. + default_forward_width = ( + int(get_spec().speculative_num_steps) + 1 + if self.tree_mode + else int(get_spec().speculative_num_draft_tokens) + ) + self.forward_width = default_forward_width + self.verify_width = int(get_spec().speculative_num_draft_tokens) + self.candidate_top_k = ( + int(get_spec().speculative_eagle_topk) if self.tree_mode else 1 + ) + self.tree_depth = int(get_spec().speculative_num_steps) if self.tree_mode else 1 + self.num_speculative_proposals = self.forward_width - 1 + self.tail_width = self.forward_width + 1 + + # Compatibility fields read by speculative infrastructure. + self.speculative_num_draft_tokens = self.verify_width + self.speculative_num_steps = self.tree_depth + self.topk = self.candidate_top_k + + # Ordinary scheduler overlap still serializes model work on the + # forward stream, so one persistent proposal workspace is sufficient: + # the next reuse is ordered after this step's tree build and verify. + self._uno_tree_workspace = {} if self.tree_mode else None + # The scheduler constructs speculative workers before allocating the + # target KV pools. Build the private tree-draft backend later, from + # init_attention_backends(), after those pools exist. + self._uno_draft_attn_backend = None + self._uno_draft_cuda_graph_runner = None + self.plan_stream, self.plan_stream_ctx = ( + get_plan_stream(self.device) + if self.tree_mode + else (None, contextlib.nullcontext()) + ) + + self._tail_offsets = torch.arange( + self.tail_width, + dtype=torch.int64, + device=self.device, + ) + + def _build_uno_draft_attn_backend(self): + """Build only the F/1 backend absent from the native Q/K target role.""" + + model_runner = self.model_runner + original_workspace_flag = model_runner.init_new_workspace + try: + with get_spec().override( + speculative_num_steps=1, + speculative_eagle_topk=1, + speculative_num_draft_tokens=self.forward_width, + ): + return model_runner._get_attention_backend(init_new_workspace=True) + finally: + model_runner.init_new_workspace = original_workspace_flag + + def init_attention_backends(self): + """Initialize only UNO's private backend after target pool allocation.""" + + if self.tree_mode: + self._uno_draft_attn_backend = self._build_uno_draft_attn_backend() + + def init_cuda_graphs(self): + """Capture only the private F-wide tree-draft graph.""" + + self._uno_draft_cuda_graph_runner = None + if not self.tree_mode or self.model_runner.decode_cuda_graph_runner is None: + return None + if self._uno_draft_attn_backend is None: + raise RuntimeError( + "UNO tree draft graph capture requires its attention backend." + ) + + tic = time.perf_counter() + before_mem = get_available_gpu_memory( + self.device, + self.gpu_id, + empty_cache=False, + ) + log_info_on_rank0( + logger, + "Capture UNO tree draft CUDA graph begin. " + f"num_tokens_per_req={self.forward_width}, " + f"avail mem={before_mem:.2f} GB", + ) + with self._bind_uno_draft_runtime(): + self._uno_draft_cuda_graph_runner = UnoDecodeCudaGraphRunner( + self.model_runner, + tree_draft_attn_backend=self._uno_draft_attn_backend, + tree_draft_width=self.forward_width, + ) + + after_mem = get_available_gpu_memory( + self.device, + self.gpu_id, + empty_cache=False, + ) + capture_time = time.perf_counter() - tic + self._additional_graph_memory_usage["draft_decode"] = before_mem - after_mem + self._additional_graph_time_usage["draft_decode"] = capture_time + log_info_on_rank0( + logger, + "Capture UNO tree draft CUDA graph end. " + f"elapsed={capture_time:.2f} s, " + f"mem usage={(before_mem - after_mem):.2f} GB, " + f"avail mem={after_mem:.2f} GB.", + ) + + return None + + @property + def draft_worker(self): + # Both passes use the target runner and its KV pool. + return None + + @property + def last_shared_read_runner(self): + # The target verify is the final phase that reads shared scheduler + # buffers, so its runner owns the WAR-barrier completion event. + return self._target_worker.model_runner + + @property + def spec_v2_attn_backends(self) -> tuple: + """Return every attention backend touched by one UNO step. + + Linear UNO uses only the target runner's native backend. Tree UNO adds + one private F-wide draft backend before finishing on the native Q-wide + target backend. The scheduler ORs these capabilities when deciding + whether FutureMap must carry a CPU sequence-length mirror. + """ + + target_backend = self._target_worker.model_runner.attn_backend + if not self.tree_mode: + return (target_backend,) + return (target_backend, self._uno_draft_attn_backend) + + def __getattr__(self, name): + # Scheduler-facing methods not implemented by this wrapper belong to + # the target worker. Guard initialization to avoid recursive lookup. + if name == "_target_worker": + raise AttributeError(name) + return getattr(self.target_worker, name) + + def _validate_batch(self, batch: ScheduleBatch) -> None: + if batch.forward_mode.is_idle(): + raise NotImplementedError("UNO does not support idle batches.") + + if batch.forward_mode.is_mixed() or ( + batch.forward_mode.is_decode() and batch.is_extend_in_batch + ): + raise NotImplementedError( + "UNO does not support mixed extend/decode batches." + ) + + if not batch.spec_algorithm.is_uno(): + raise RuntimeError( + "UnoWorkerV2 received a batch whose speculative algorithm is not UNO." + ) + + sampling_info = batch.sampling_info + if sampling_info is None: + raise RuntimeError("UNO requires sampling metadata.") + + if sampling_info.need_min_p_sampling: + raise NotImplementedError("UNO does not support min-p sampling.") + + if batch.has_grammar: + raise NotImplementedError("UNO does not support grammar decoding.") + + if batch.return_logprob: + raise NotImplementedError("UNO does not support returned logprobs.") + + if batch.return_hidden_states: + raise NotImplementedError("UNO does not support returned hidden states.") + + penalizer = sampling_info.penalizer_orchestrator + penalties_active = ( + (penalizer is not None and penalizer.is_required) + or sampling_info.acc_additive_penalties is not None + or sampling_info.acc_scaling_penalties is not None + ) + if penalties_active: + raise NotImplementedError("UNO does not support sampling penalties.") + + if sampling_info.logit_bias is not None: + raise NotImplementedError("UNO does not support logit bias.") + + if sampling_info.has_custom_logit_processor: + raise NotImplementedError("UNO does not support custom logit processors.") + + if any(req.lora_id is not None for req in batch.reqs): + raise NotImplementedError("UNO does not support multi-LoRA.") + + def _make_forward_batch( + self, + *, + spec_input_type: SpecInputType, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + prefix_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + seq_lens_sum: Optional[int], + req_pool_indices: torch.Tensor, + ) -> ForwardBatch: + input_ids = input_ids.reshape(-1) + positions = positions.reshape(-1) + out_cache_loc = out_cache_loc.reshape(-1) + + spec_info = UnoForwardInput( + spec_input_type=spec_input_type, + positions=positions, + draft_token_num=self.forward_width, + ) + + return ForwardBatch( + forward_mode=ForwardMode.TARGET_VERIFY, + batch_size=len(prefix_lens), + input_ids=input_ids, + req_pool_indices=req_pool_indices, + seq_lens=prefix_lens, + out_cache_loc=out_cache_loc, + seq_lens_sum=seq_lens_sum, + seq_lens_cpu=seq_lens_cpu, + positions=positions, + spec_algorithm=SpeculativeAlgorithm.UNO, + spec_info=spec_info, + capture_hidden_mode=CaptureHiddenMode.NULL, + return_hidden_states_before_norm=False, + ) + + def _run_target_block( + self, + forward_batch: ForwardBatch, + *, + need_top1: bool = True, + ) -> tuple: + result = self.target_worker.forward_batch_generation( + batch=None, + forward_batch=forward_batch, + is_verify=True, + ) + + if result.logits_output is None: + raise RuntimeError("UNO target block returned no logits output.") + + logits = result.logits_output.next_token_logits + if logits is None: + raise RuntimeError("UNO target block returned no next-token logits.") + + expected_rows = forward_batch.batch_size * self.forward_width + if logits.ndim != 2 or logits.shape[0] != expected_rows: + raise RuntimeError( + "UNO target block returned an invalid logits shape: " + f"expected ({expected_rows}, vocab_size), got " + f"{tuple(logits.shape)}." + ) + + # DFlash consumes logits directly; only greedy acceptance needs top-1. + if not need_top1: + return result, None + + predictions = torch.argmax(logits, dim=-1).view( + forward_batch.batch_size, + self.forward_width, + ) + return result, predictions + + @staticmethod + def _accept_and_pack( + *, + candidates: torch.Tensor, + target_top1: torch.Tensor, + committed_seq_lens: torch.Tensor, + ) -> tuple: + if candidates.ndim != 2: + raise RuntimeError( + f"UNO candidates must be rank 2, got shape={tuple(candidates.shape)}." + ) + if target_top1.shape != candidates.shape: + raise RuntimeError( + "UNO candidate and target shapes differ: " + f"{tuple(candidates.shape)} versus {tuple(target_top1.shape)}." + ) + + batch_size, forward_width = candidates.shape + device = candidates.device + + # forward_width is static configuration, so this branch does not inspect + # or synchronize a device tensor. + if forward_width == 1: + accepted_specs = torch.zeros( + batch_size, + dtype=torch.int32, + device=device, + ) + else: + matches = candidates[:, 1:] == target_top1[:, :-1] + accepted_specs = ( + matches.to(torch.int32).cumprod(dim=1).sum(dim=1).to(torch.int32) + ) + + accepted_specs_long = accepted_specs.to(torch.int64) + correction = target_top1.gather( + 1, + accepted_specs_long[:, None], + ).squeeze(1) + + output_ids = torch.zeros( + (batch_size, forward_width + 1), + dtype=torch.int64, + device=device, + ) + output_ids[:, :forward_width].copy_(candidates) + output_ids.scatter_( + 1, + (accepted_specs_long + 1)[:, None], + correction[:, None], + ) + + accept_lens = accepted_specs + 2 + new_seq_lens = committed_seq_lens + accept_lens.to(committed_seq_lens.dtype) + + return output_ids, accept_lens, new_seq_lens, correction + + def _forward_prefill( + self, + batch: ScheduleBatch, + on_publish, + ) -> GenerationBatchResult: + result = self.target_worker.forward_batch_generation(batch) + if not isinstance(result.next_token_ids, torch.Tensor): + raise RuntimeError("UNO target prefill returned no sampled seed tensor.") + + seed_tokens = result.next_token_ids.reshape(-1) + if seed_tokens.shape[0] != len(batch.reqs): + raise RuntimeError( + "UNO prefill seed count does not match batch size: " + f"{seed_tokens.shape[0]} versus {len(batch.reqs)}." + ) + + result.new_seq_lens = batch.seq_lens + result.next_draft_input = UnoDraftInput( + bonus_tokens=seed_tokens, + new_seq_lens=batch.seq_lens, + forward_width=self.forward_width, + ) + + if on_publish is not None: + on_publish(result.new_seq_lens) + return result + + def _forward_decode_tree( + self, + batch: ScheduleBatch, + on_publish, + ) -> GenerationBatchResult: + """Run UNO's F-wide proposal pass, then native EAGLE Q-node verify.""" + + draft_state = batch.spec_info + + if batch.seq_lens.is_cuda: + batch.seq_lens.record_stream( + torch.get_device_module(self.device).current_stream() + ) + + batch_size = len(batch.seq_lens) + committed_seq_lens = batch.seq_lens.clone() + seed_tokens = draft_state.bonus_tokens.reshape(-1).to( + device=self.device, + dtype=torch.int64, + ) + + committed_seq_lens_cpu = None + if batch.seq_lens_cpu is not None: + committed_seq_lens_cpu = batch.seq_lens_cpu.to( + device="cpu", + dtype=torch.int64, + ) + draft_seq_lens_cpu = committed_seq_lens_cpu + self.forward_width + draft_seq_lens_sum = int(draft_seq_lens_cpu.sum()) + elif draft_state.reserved_seq_lens_cpu is not None: + # This host tensor is a planning upper bound only. The device + # frontier below remains the exact committed length. + draft_seq_lens_cpu = draft_state.reserved_seq_lens_cpu + draft_seq_lens_sum = draft_state.reserved_seq_lens_sum + else: + draft_seq_lens_cpu = None + draft_seq_lens_sum = None + + draft_positions = ( + committed_seq_lens.to(torch.int64)[:, None] + + (self._tail_offsets[None, : self.forward_width]) + ) + req_pool_indices_long = batch.req_pool_indices.to(torch.int64) + req_to_token = self.model_runner.req_to_token_pool.req_to_token + draft_locs = req_to_token[ + req_pool_indices_long[:, None], + draft_positions, + ].to(torch.int64) + + draft_input_ids = build_uno_draft_input( + seed_tokens=seed_tokens, + forward_width=self.forward_width, + vocab_size=self.model_runner.model_config.vocab_size, + ) + draft_forward_batch = self._make_forward_batch( + spec_input_type=SpecInputType.UNO_DRAFT, + input_ids=draft_input_ids, + positions=draft_positions, + out_cache_loc=draft_locs, + prefix_lens=committed_seq_lens, + seq_lens_cpu=draft_seq_lens_cpu, + seq_lens_sum=draft_seq_lens_sum, + req_pool_indices=batch.req_pool_indices, + ) + draft_result, _ = self._run_draft_block( + draft_forward_batch, + need_top1=False, + ) + draft_logits = draft_result.logits_output.next_token_logits.reshape( + batch_size, + self.forward_width, + -1, + ) + + sampling_info = batch.sampling_info + if sampling_info.is_all_greedy: + clean_root_tokens = torch.argmax( + draft_logits[:, 0, :], + dim=-1, + ) + else: + clean_root_tokens = sample_uno_clean_root( + seed_tokens=seed_tokens, + draft_logits=draft_logits, + sampling_info=sampling_info, + max_top_k=draft_state.max_top_k, + uniform_top_k_value=draft_state.uniform_top_k_value, + ) + + proposal = build_uno_tree_proposal( + clean_root_tokens, + draft_logits[:, 1:, :], + max_nodes=self.verify_width, + candidate_top_k=self.candidate_top_k, + temperature=sampling_info.temperatures, + workspace=self._uno_tree_workspace, + ) + + # The first pass wrote the carried seed at C. Give EAGLE a shallow + # batch whose KV-ready prefix is therefore C+1; its existing allocator + # assigns all Q tree slots at C+1 and its compactor can stay unchanged. + verify_batch = copy.copy(batch) + verify_batch.seq_lens = committed_seq_lens + 1 + if committed_seq_lens_cpu is None: + verify_batch.seq_lens_cpu = None + verify_batch.seq_lens_sum = None + else: + verify_batch.seq_lens_cpu = committed_seq_lens_cpu + 1 + verify_batch.seq_lens_sum = int(verify_batch.seq_lens_cpu.sum()) + + verify_input = build_eagle_verify_input( + verify_batch, + EagleDraftInput(bonus_tokens=proposal.root_tokens), + proposal.parent_list, + proposal.top_scores_index, + proposal.draft_tokens, + None, + target_worker=self.target_worker, + topk=self.candidate_top_k, + num_steps=self.tree_depth, + num_draft_tokens=self.verify_width, + tree_mask_mode=default_tree_mask_mode(), + device=self.device, + ) + verify_batch.spec_info = verify_input + if self.plan_stream is not None: + # C+1 was produced on the forward stream immediately above. The + # generic EAGLE path receives an older, already-visible frontier; + # UNO must explicitly order its freshly derived tensor before the + # plan stream assigns Q verify cache locations from it. + self.plan_stream.wait_stream( + torch.get_device_module(self.device).current_stream() + ) + eagle_result = run_eagle_verify( + verify_batch, + target_worker=self.target_worker, + req_to_token_pool=self.model_runner.req_to_token_pool, + token_to_kv_pool_allocator=(self.model_runner.token_to_kv_pool_allocator), + plan_stream=self.plan_stream, + plan_stream_ctx=self.plan_stream_ctx, + topk=self.candidate_top_k, + num_draft_tokens=self.verify_width, + device=self.device, + metadata_ready_pre_pad=False, + finalize_tree_path=True, + uno_target_max_top_k=draft_state.max_top_k, + ) + + packed = pack_uno_tree_result( + clean_root_tokens=clean_root_tokens, + eagle_predict=eagle_result.next_token_ids, + eagle_accept_lens=eagle_result.accept_lens, + draft_width=self.forward_width, + ) + new_seq_lens = eagle_result.new_seq_lens + next_draft_input = UnoDraftInput( + bonus_tokens=eagle_result.next_draft_input.bonus_tokens, + new_seq_lens=new_seq_lens, + forward_width=self.forward_width, + ) + + if on_publish is not None: + on_publish(new_seq_lens) + + # Preserve EAGLE's verify ForwardBatch keep-alive refs verbatim. FutureMap + # relays the wrapped bonus; on_publish above relays the new frontier. + return GenerationBatchResult( + logits_output=eagle_result.logits_output, + next_token_ids=packed.output_ids.reshape(-1), + accept_lens=packed.accept_lens, + next_draft_input=next_draft_input, + speculative_num_draft_tokens=self.forward_width, + speculative_output_stride=self.forward_width + 1, + num_non_draft_tokens_per_req=2, + new_seq_lens=new_seq_lens, + can_run_cuda_graph=eagle_result.can_run_cuda_graph, + routed_experts_output=eagle_result.routed_experts_output, + indexer_topk_output=eagle_result.indexer_topk_output, + extra_keep_alive_refs=eagle_result.extra_keep_alive_refs, + ) + + def _forward_decode( + self, + batch: ScheduleBatch, + on_publish, + ) -> GenerationBatchResult: + if self.tree_mode: + return self._forward_decode_tree(batch, on_publish) + + draft_state = batch.spec_info + + if batch.seq_lens.is_cuda: + batch.seq_lens.record_stream( + torch.get_device_module(self.device).current_stream() + ) + + batch_size = len(batch.seq_lens) + committed_seq_lens = batch.seq_lens.clone() + + seed_tokens = draft_state.bonus_tokens.reshape(-1).to( + device=self.device, + dtype=torch.int64, + ) + + if batch.seq_lens_cpu is not None: + committed_seq_lens_cpu = batch.seq_lens_cpu.to( + device="cpu", + dtype=torch.int32, + ) + draft_seq_lens_cpu = committed_seq_lens_cpu + self.forward_width + verify_seq_lens_cpu = committed_seq_lens_cpu + self.tail_width + draft_seq_lens_sum = int(draft_seq_lens_cpu.sum()) + verify_seq_lens_sum = int(verify_seq_lens_cpu.sum()) + elif draft_state.reserved_seq_lens_cpu is not None: + # Triton only needs a safe host planning bound. The allocator's + # retained reservation avoids a D2H copy when FutureMap keeps the + # exact committed frontier on GPU. + draft_seq_lens_cpu = draft_state.reserved_seq_lens_cpu + verify_seq_lens_cpu = draft_state.reserved_seq_lens_cpu + draft_seq_lens_sum = draft_state.reserved_seq_lens_sum + verify_seq_lens_sum = draft_state.reserved_seq_lens_sum + else: + draft_seq_lens_cpu = None + verify_seq_lens_cpu = None + draft_seq_lens_sum = None + verify_seq_lens_sum = None + + logical_positions = ( + committed_seq_lens.to(torch.int64)[:, None] + (self._tail_offsets[None, :]) + ) + req_pool_indices_long = batch.req_pool_indices.to(torch.int64) + req_to_token = self.model_runner.req_to_token_pool.req_to_token + tail_locs = req_to_token[ + req_pool_indices_long[:, None], + logical_positions, + ].to(torch.int64) + + draft_input_ids = build_uno_draft_input( + seed_tokens=seed_tokens, + forward_width=self.forward_width, + vocab_size=self.model_runner.model_config.vocab_size, + ) + draft_forward_batch = self._make_forward_batch( + spec_input_type=SpecInputType.UNO_DRAFT, + input_ids=draft_input_ids, + positions=logical_positions[:, : self.forward_width], + out_cache_loc=tail_locs[:, : self.forward_width], + prefix_lens=committed_seq_lens, + seq_lens_cpu=draft_seq_lens_cpu, + seq_lens_sum=draft_seq_lens_sum, + req_pool_indices=batch.req_pool_indices, + ) + sampling_info = batch.sampling_info + all_greedy = sampling_info.is_all_greedy + max_top_k = draft_state.max_top_k + uniform_top_k_value = draft_state.uniform_top_k_value + draft_result, candidates = self._run_draft_block( + draft_forward_batch, + need_top1=all_greedy, + ) + if not all_greedy: + draft_logits = draft_result.logits_output.next_token_logits.reshape( + batch_size, + self.forward_width, + -1, + ) + candidates, draft_distribution = sample_uno_candidates( + draft_logits=draft_logits, + sampling_info=sampling_info, + max_top_k=max_top_k, + uniform_top_k_value=uniform_top_k_value, + ) + + verify_prefix_lens = committed_seq_lens + 1 + verify_forward_batch = self._make_forward_batch( + spec_input_type=SpecInputType.UNO_VERIFY, + input_ids=candidates, + positions=logical_positions[:, 1:], + out_cache_loc=tail_locs[:, 1:], + prefix_lens=verify_prefix_lens, + seq_lens_cpu=verify_seq_lens_cpu, + seq_lens_sum=verify_seq_lens_sum, + req_pool_indices=batch.req_pool_indices, + ) + verify_result, target_top1 = self._run_target_block( + verify_forward_batch, + need_top1=all_greedy, + ) + + if all_greedy: + output_ids, accept_lens, new_seq_lens, correction = self._accept_and_pack( + candidates=candidates, + target_top1=target_top1, + committed_seq_lens=committed_seq_lens, + ) + else: + sampling_result = run_uno_sampling( + candidates=candidates, + next_token_logits=verify_result.logits_output.next_token_logits, + sampling_info=sampling_info, + committed_frontiers=committed_seq_lens, + draft_distribution=draft_distribution, + max_top_k=max_top_k, + uniform_top_k_value=uniform_top_k_value, + ) + output_ids = sampling_result.output_ids + accept_lens = sampling_result.accept_lens + new_seq_lens = sampling_result.new_seq_lens + correction = sampling_result.next_seed_tokens + + next_draft_input = UnoDraftInput( + bonus_tokens=correction, + new_seq_lens=new_seq_lens, + forward_width=self.forward_width, + ) + + if on_publish is not None: + on_publish(new_seq_lens) + + return GenerationBatchResult( + logits_output=verify_result.logits_output, + next_token_ids=output_ids.reshape(-1), + accept_lens=accept_lens, + next_draft_input=next_draft_input, + speculative_num_draft_tokens=self.forward_width, + speculative_output_stride=self.tail_width, + num_non_draft_tokens_per_req=2, + new_seq_lens=new_seq_lens, + can_run_cuda_graph=verify_result.can_run_cuda_graph, + routed_experts_output=verify_result.routed_experts_output, + indexer_topk_output=verify_result.indexer_topk_output, + ) + + def forward_batch_generation( + self, + batch: ScheduleBatch, + on_publish=None, + grammar_barrier=None, + ) -> GenerationBatchResult: + del grammar_barrier + self._validate_batch(batch) + + if batch.forward_mode == ForwardMode.EXTEND: + return self._forward_prefill(batch, on_publish) + + if batch.forward_mode == ForwardMode.DECODE: + return self._forward_decode(batch, on_publish) + + raise RuntimeError( + f"UNO expected an EXTEND or DECODE batch, got {batch.forward_mode}." + ) + + def update_weights_from_disk(self, recv_req): + # The scheduler updates the target worker before calling the spec worker. + return True, "UNO has no separate draft weights." + + def update_weights_from_ipc(self, recv_req): + # The scheduler updates the target worker before calling the spec worker. + return True, "UNO has no separate draft weights." + + def update_weights_from_tensor(self, recv_req): + # This update route selects the spec worker instead of updating both. + return self.target_worker.update_weights_from_tensor(recv_req) + + @contextlib.contextmanager + def _bind_uno_draft_runtime(self): + target_attn_backend = self.model_runner.attn_backend + target_graph_runner = self.model_runner.decode_cuda_graph_runner + self.model_runner.attn_backend = self._uno_draft_attn_backend + self.model_runner.decode_cuda_graph_runner = self._uno_draft_cuda_graph_runner + try: + yield + finally: + self.model_runner.attn_backend = target_attn_backend + self.model_runner.decode_cuda_graph_runner = target_graph_runner + + def _run_draft_block( + self, + forward_batch: ForwardBatch, + *, + need_top1: bool = True, + ): + batch_size = forward_batch.batch_size + self.lora_manager.reset_lora_batch() + backend_context = ( + self._bind_uno_draft_runtime() + if self.tree_mode + else contextlib.nullcontext() + ) + attn_context = ( + forward_context(ForwardContext(attn_backend=self._uno_draft_attn_backend)) + if self.tree_mode + else contextlib.nullcontext() + ) + # The eager runner plans through model_runner.attn_backend, while + # attention layers execute through ForwardContext. Tree UNO binds both + # to the same private F/1 backend for this draft pass. + with backend_context, attn_context: + if self.num_speculative_proposals == 0: + return self._run_target_block( + forward_batch, + need_top1=need_top1, + ) + + graph_runner = getattr( + self.model_runner, + "decode_cuda_graph_runner", + None, + ) + # If cuda-graph is on, reuse its captured LoRA routing + # and replay it in _run_target_block. + # Else, prepare routing for eager. + if not ( + forward_batch.forward_mode.is_cuda_graph() + and graph_runner is not None + and graph_runner.can_run_graph(forward_batch) + ): + self.lora_manager.prepare_lora_token_segments( + lora_ids=[None, self.uno_lora_id] * batch_size, + segment_lens=[1, self.num_speculative_proposals] * batch_size, + ) + result = self._run_target_block( + forward_batch, + need_top1=need_top1, + ) + # Clear LoRA routing before verification + self.lora_manager.reset_lora_batch() + return result diff --git a/test/registered/kernels/ops/attention/test_suffix_attention_merge.py b/test/registered/kernels/ops/attention/test_suffix_attention_merge.py new file mode 100644 index 000000000..b5b426327 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_suffix_attention_merge.py @@ -0,0 +1,121 @@ +"""CUDA correctness tests for the fused suffix-attention merge.""" + +import math +import unittest + +import torch + +from sglang.kernels.ops.attention.suffix_attention_merge import ( + merge_suffix_attention_in_place, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci( + est_time=15, + stage="base-b-kernel-unit", + runner_config="1-gpu-large", +) + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +class TestSuffixAttentionMerge(CustomTestCase): + def _case(self, *, num_queries: int, head_dim: int, dtype: torch.dtype): + torch.manual_seed(7) + device = torch.device("cuda") + num_q_heads = 8 + num_kv_heads = 2 + num_slots = 2 * num_queries + 11 + + q = torch.randn(num_queries, num_q_heads, head_dim, device=device, dtype=dtype) + k_cache = torch.randn( + num_slots, num_kv_heads, head_dim, device=device, dtype=dtype + ) + v_cache = torch.randn_like(k_cache) + page_table = torch.stack( + [ + torch.randperm(num_slots, device=device)[:num_queries] + for _ in range(num_queries) + ] + ).to(torch.int32) + suffix_lengths = ( + torch.arange(num_queries, device=device, dtype=torch.int32) + .remainder(num_queries) + .add_(1) + ) + prefix = torch.randn_like(q) + prefix_lse = torch.randn( + num_q_heads, num_queries, device=device, dtype=torch.float32 + ) + scale = 1.0 / math.sqrt(head_dim) + + reference = prefix.float().clone() + heads_per_kv = num_q_heads // num_kv_heads + kv_heads = torch.arange(num_q_heads, device=device) // heads_per_kv + for token in range(num_queries): + length = int(suffix_lengths[token]) + slots = page_table[token, :length].long() + keys = k_cache[slots][:, kv_heads].float() + values = v_cache[slots][:, kv_heads].float() + scores = torch.einsum("lhd,hd->lh", keys, q[token].float()) * scale + maximum = torch.maximum(prefix_lse[:, token], scores.max(dim=0).values) + prefix_weight = torch.exp(prefix_lse[:, token] - maximum) + suffix_weights = torch.exp(scores - maximum) + reference[token] = ( + reference[token] * prefix_weight[:, None] + + torch.einsum("lh,lhd->hd", suffix_weights, values) + ) / (prefix_weight + suffix_weights.sum(dim=0))[:, None] + + static_prefix = prefix.clone() + merge_suffix_attention_in_place( + q, + k_cache, + v_cache, + page_table, + suffix_lengths, + static_prefix, + prefix_lse, + scale, + ) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_prefix.copy_(prefix) + merge_suffix_attention_in_place( + q, + k_cache, + v_cache, + page_table, + suffix_lengths, + static_prefix, + prefix_lse, + scale, + ) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close( + static_prefix.float(), reference, rtol=2e-2, atol=2e-2 + ) + + def test_representative_shapes(self): + cases = ( + (16, 64, torch.float16), + (60, 128, torch.bfloat16), + ) + for num_queries, head_dim, dtype in cases: + with self.subTest( + num_queries=num_queries, + head_dim=head_dim, + dtype=dtype, + ): + self._case( + num_queries=num_queries, + head_dim=head_dim, + dtype=dtype, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/uno/test_uno.py b/test/registered/spec/uno/test_uno.py new file mode 100644 index 000000000..0e98d5543 --- /dev/null +++ b/test/registered/spec/uno/test_uno.py @@ -0,0 +1,267 @@ +"""End-to-end CUDA-graph coverage for linear and tree UNO decoding. + +The test runs both modes on the same prompts. Linear UNO alternates +LoRA-draft and clean-target variants in one graph runner. Tree UNO uses a +private LoRA-draft runner before native EAGLE tree verification. Besides the +generation contract, short greedy comparisons guard lossless output parity +with autoregressive decoding, and the stochastic comparison guards that tree +search improves TPF over the linear proposal on a small, fixed GSM8K sample. +""" + +import os +import unittest +from typing import NamedTuple + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci( + est_time=480, + stage="base-b", + runner_config="1-gpu-large", +) + +MODEL = "Qwen/Qwen3-8B" +DEFAULT_UNO_LORA = "s-sahoo/uno-qwen3-8B" +LORA_PATH_ENV = "SGLANG_TEST_UNO_LORA_PATH" +MAX_NEW_TOKENS = 128 +# AR decode and UNO verification use different kernel shapes, so compare a +# bounded greedy prefix instead of requiring full-sequence bitwise identity. +PARITY_TOKENS = 32 +# One LoRA draft forward plus one clean verification forward. +FORWARDS_PER_UNO_CYCLE = 2 +PROMPTS = ( + ( + "Question: Janet's ducks lay 16 eggs per day. She eats three for " + "breakfast every morning and bakes muffins for her friends every day " + "with four. She sells the remainder at the farmers' market daily for " + "$2 per fresh duck egg. How much in dollars does she make every day " + "at the farmers' market?\nAnswer:" + ), + ( + "Question: A robe takes 2 bolts of blue fiber and half that much " + "white fiber. How many bolts in total does it take?\nAnswer:" + ), + ( + "Question: Josh decides to try flipping a house. He buys a house for " + "$80,000 and then puts in $50,000 in repairs. This increased the value " + "of the house by 150%. How much profit did he make?\nAnswer:" + ), +) + + +class _UnoConfig(NamedTuple): + name: str + speculative_num_steps: int + speculative_eagle_topk: int + speculative_num_draft_tokens: int + + +LINEAR_CONFIG = _UnoConfig( + name="linear", + speculative_num_steps=1, + speculative_eagle_topk=1, + speculative_num_draft_tokens=8, # F = 8 +) +TREE_CONFIG = _UnoConfig( + name="tree", + speculative_num_steps=7, # F = 8 + speculative_eagle_topk=16, + speculative_num_draft_tokens=8, # Q = 8 +) + + +class TestUnoCudaGraph(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.base_url = DEFAULT_URL_FOR_TEST + cls.adapter_path = os.environ.get(LORA_PATH_ENV, DEFAULT_UNO_LORA) + + def _server_args(self, config: _UnoConfig | None) -> list[str]: + args = [ + "--dtype", + "bfloat16", + "--attention-backend", + "fa3", + "--max-running-requests", + str(len(PROMPTS)), + "--cuda-graph-max-bs-decode", + str(len(PROMPTS)), + "--mem-fraction-static", + "0.7", + "--page-size", + "1", + "--disable-radix-cache", + "--random-seed", + "17", + ] + if config is not None: + args.extend( + [ + "--speculative-algorithm", + "UNO", + "--uno-lora-path", + self.adapter_path, + "--speculative-num-steps", + str(config.speculative_num_steps), + "--speculative-eagle-topk", + str(config.speculative_eagle_topk), + "--speculative-num-draft-tokens", + str(config.speculative_num_draft_tokens), + ] + ) + return args + + def _run_ar_reference(self) -> list[list[int]]: + process = None + try: + process = popen_launch_server( + MODEL, + self.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=self._server_args(None), + ) + return self._run_greedy_output_ids() + finally: + if process is not None: + kill_process_tree(process.pid) + + def _run_config(self, config: _UnoConfig) -> tuple[float, list[list[int]]]: + process = None + try: + process = popen_launch_server( + MODEL, + self.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=self._server_args(config), + ) + greedy_output_ids = self._run_greedy_output_ids() + tpf = self._run_generation_contract(config) + return tpf, greedy_output_ids + finally: + if process is not None: + kill_process_tree(process.pid) + + def _run_greedy_output_ids(self) -> list[list[int]]: + # A list-valued request can be admitted with different prefill batch + # shapes across server launches. Run each parity prompt at BS1 so the + # AR and UNO comparisons use the same execution shape. + output_ids = [] + for prompt in PROMPTS: + response = requests.post( + self.base_url + "/generate", + json={ + "text": prompt, + "sampling_params": { + "temperature": 0, + "max_new_tokens": PARITY_TOKENS, + "ignore_eos": True, + }, + }, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + ) + self.assertEqual(response.status_code, 200, response.text) + + result = response.json() + self.assertIn("output_ids", result, result) + self.assertEqual( + len(result["output_ids"]), + PARITY_TOKENS, + f"Wrong greedy output length for prompt {prompt!r}", + ) + output_ids.append(result["output_ids"]) + return output_ids + + def _assert_ar_parity( + self, + mode: str, + actual: list[list[int]], + expected: list[list[int]], + ) -> None: + for prompt, actual_ids, expected_ids in zip(PROMPTS, actual, expected): + self.assertEqual( + actual_ids, + expected_ids, + f"{mode} UNO diverged from AR within the first " + f"{PARITY_TOKENS} tokens for prompt {prompt!r}", + ) + + def _run_generation_contract(self, config: _UnoConfig) -> float: + server_info = requests.get(self.base_url + "/server_info", timeout=30).json() + self.assertEqual( + server_info["speculative_eagle_topk"], config.speculative_eagle_topk + ) + self.assertEqual( + server_info["speculative_num_steps"], config.speculative_num_steps + ) + self.assertEqual( + server_info["speculative_num_draft_tokens"], + config.speculative_num_draft_tokens, + ) + + response = requests.post( + self.base_url + "/generate", + json={ + "text": PROMPTS, + "sampling_params": { + "temperature": 0.7, + "top_k": 50, + "top_p": 0.95, + "max_new_tokens": MAX_NEW_TOKENS, + "ignore_eos": True, + }, + }, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + ) + self.assertEqual(response.status_code, 200, response.text) + + results = response.json() + self.assertEqual(len(results), len(PROMPTS)) + total_completion_tokens = 0 + total_verify_ct = 0 + for result in results: + self.assertTrue(result["text"].strip()) + meta_info = result["meta_info"] + self.assertEqual(meta_info["completion_tokens"], MAX_NEW_TOKENS) + total_completion_tokens += meta_info["completion_tokens"] + total_verify_ct += meta_info.get("spec_verify_ct", 0) + + self.assertGreater( + total_verify_ct, 0, f"{config.name} performed no verify steps" + ) + total_forwards = FORWARDS_PER_UNO_CYCLE * total_verify_ct + tpf = total_completion_tokens / total_forwards + self.assertGreater( + tpf, + 1.5, + f"{config.name} did not advance beyond autoregressive decoding: {tpf=}", + ) + return tpf + + def test_ar_parity_and_tree_tpf_exceeds_linear(self): + ar_output_ids = self._run_ar_reference() + + linear_tpf, linear_output_ids = self._run_config(LINEAR_CONFIG) + self._assert_ar_parity("Linear", linear_output_ids, ar_output_ids) + + tree_tpf, tree_output_ids = self._run_config(TREE_CONFIG) + self._assert_ar_parity("Tree", tree_output_ids, ar_output_ids) + + print(f"UNO GSM8K sample: {linear_tpf=:.3f}, {tree_tpf=:.3f}") + self.assertGreater( + tree_tpf, + linear_tpf, + f"Tree UNO did not improve TPF: {linear_tpf=:.3f}, {tree_tpf=:.3f}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/lora/test_uno_inactive_lora_batch.py b/test/registered/unit/lora/test_uno_inactive_lora_batch.py new file mode 100644 index 000000000..ade5ebae8 --- /dev/null +++ b/test/registered/unit/lora/test_uno_inactive_lora_batch.py @@ -0,0 +1,41 @@ +"""Regression test for UNO's base-only LoRA routing fast path.""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.lora.lora_manager import LoRAManager +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class _InactiveSkippingBackend: + skip_inactive_lora_batches = True + + def __init__(self): + self.batch_info = object() + self.prepare_called = False + + def reset_batch_state(self): + self.batch_info = None + + def prepare_lora_batch(self, *args, **kwargs): + self.prepare_called = True + + +class TestUnoInactiveLoRABatch(CustomTestCase): + def test_all_base_batch_clears_stale_routing_before_graph_metadata(self): + backend = _InactiveSkippingBackend() + manager = LoRAManager.__new__(LoRAManager) + manager.lora_backend = backend + forward_batch = SimpleNamespace(lora_ids=[None], batch_size=1) + + manager.prepare_lora_batch(forward_batch) + + self.assertIsNone(backend.batch_info) + self.assertFalse(backend.prepare_called) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/lora/test_uno_lora_targets.py b/test/registered/unit/lora/test_uno_lora_targets.py new file mode 100644 index 000000000..e33500cf0 --- /dev/null +++ b/test/registered/unit/lora/test_uno_lora_targets.py @@ -0,0 +1,236 @@ +"""Target-layer validation for UNO's specialized LoRA backend.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import torch + +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.lora.backend.triton_backend import TritonLoRABackend +from sglang.srt.lora.backend.uno_cublas_backend import UnoCublasLoRABackend +from sglang.srt.lora.lora_manager import LoRAManager +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestUnoLoRATargets(CustomTestCase): + def setUp(self): + self.backend = UnoCublasLoRABackend.__new__(UnoCublasLoRABackend) + self.backend._pending_lora_a = None + self.backend._use_cublas_lora_b = False + + @staticmethod + def _model(modules, **attributes): + return SimpleNamespace( + named_modules=lambda: modules, + **attributes, + ) + + def test_supported_decoder_targets_are_accepted(self): + modules = [ + ( + "model.layers.0.qkv_proj", + ColumnParallelLinear.__new__(ColumnParallelLinear), + ), + ( + "model.layers.0.o_proj", + RowParallelLinear.__new__(RowParallelLinear), + ), + ( + "model.layers.0.fused_qkv_a_proj_with_mqa", + ReplicatedLinear.__new__(ReplicatedLinear), + ), + ] + self.backend.validate_lora_targets( + base_model=self._model(modules), + target_modules={ + "qkv_proj", + "o_proj", + "fused_qkv_a_proj_with_mqa", + }, + ) + + def test_unsupported_targets_are_rejected(self): + cases = { + "unknown decoder layer": ( + self._model( + [ + ( + "model.layers.0.custom_proj", + torch.nn.Linear(2, 2), + ) + ] + ), + {"custom_proj"}, + "Linear", + ), + "fused MoE": ( + self._model( + [ + ( + "model.layers.0.mlp", + FusedMoE.__new__(FusedMoE), + ) + ] + ), + {"gate_up_proj", "down_proj"}, + "FusedMoE", + ), + } + + for name, (model, targets, expected) in cases.items(): + with self.subTest(name=name), self.assertRaisesRegex(ValueError, expected): + self.backend.validate_lora_targets( + base_model=model, + target_modules=targets, + ) + + def test_nonoverlap_dense_calls_fall_back_to_triton(self): + x = object() + weights = object() + hidden = object() + base_output = object() + pruned_batch_info = object() + expected = object() + + with ( + patch.object( + TritonLoRABackend, + "run_lora_a_sgemm", + return_value=hidden, + ) as run_lora_a, + patch.object( + TritonLoRABackend, + "run_lora_b_sgemm", + return_value=expected, + ) as run_lora_b, + ): + actual_hidden = self.backend.run_lora_a_sgemm( + x, + weights, + pruned_batch_info=pruned_batch_info, + ) + actual = self.backend.run_lora_b_sgemm( + actual_hidden, + weights, + base_output=base_output, + pruned_batch_info=pruned_batch_info, + ) + + self.assertIs(actual_hidden, hidden) + self.assertIs(actual, expected) + run_lora_a.assert_called_once_with( + x, + weights, + pruned_batch_info, + 1, + ) + run_lora_b.assert_called_once_with( + hidden, + weights, + base_output, + pruned_batch_info, + ) + + def test_overlap_launch_selects_cublas(self): + pending = object() + x = object() + weights = object() + hidden = object() + base_output = object() + expected = object() + self.backend._pending_lora_a = pending + self.backend._consume_lora_a_overlap = MagicMock(return_value=hidden) + self.backend._run_lora_b = MagicMock(return_value=expected) + + with ( + patch.object(TritonLoRABackend, "run_lora_a_sgemm") as run_lora_a, + patch.object(TritonLoRABackend, "run_lora_b_sgemm") as run_lora_b, + ): + actual_hidden = self.backend.run_lora_a_sgemm(x, weights) + actual = self.backend.run_lora_b_sgemm( + actual_hidden, + weights, + base_output=base_output, + ) + + self.assertIs(actual_hidden, hidden) + self.assertIs(actual, expected) + self.backend._consume_lora_a_overlap.assert_called_once_with(pending) + self.backend._run_lora_b.assert_called_once_with( + hidden, + weights, + base_output, + ) + self.assertFalse(self.backend._use_cublas_lora_b) + run_lora_a.assert_not_called() + run_lora_b.assert_not_called() + + def test_nonoverlap_qkv_call_falls_back_to_triton(self): + expected = object() + args = { + "x": object(), + "qkv_lora_a": object(), + "qkv_lora_b": object(), + "output_offset": object(), + "output_offset_cpu": object(), + "max_qkv_out_dim": 128, + "base_output": object(), + "n_slices": 2, + } + + with patch.object( + TritonLoRABackend, + "run_qkv_lora", + return_value=expected, + ) as run_qkv_lora: + actual = self.backend.run_qkv_lora(**args) + + self.assertIs(actual, expected) + run_qkv_lora.assert_called_once_with( + args["x"], + args["qkv_lora_a"], + args["qkv_lora_b"], + args["output_offset"], + 128, + args["base_output"], + 2, + ) + + def test_manager_preflights_targets_before_wrapping(self): + manager = LoRAManager.__new__(LoRAManager) + manager.base_model = object() + manager.lora_backend = MagicMock() + manager._experts_shared_outer_override = None + manager.init_lora_adapters = MagicMock() + manager.init_lora_shapes = MagicMock( + side_effect=lambda **_: setattr(manager, "target_modules", {"qkv_proj"}) + ) + manager._detect_shared_outer_loras = MagicMock(return_value=False) + manager.init_lora_modules = MagicMock() + manager.init_memory_pool = MagicMock() + manager.update_lora_info = MagicMock() + manager.lora_backend.validate_lora_targets.side_effect = ValueError( + "unsupported target" + ) + + with self.assertRaisesRegex(ValueError, "unsupported target"): + manager.init_state(max_lora_rank=1, target_modules={"q_proj"}) + + manager.lora_backend.validate_lora_targets.assert_called_once_with( + base_model=manager.base_model, + target_modules={"qkv_proj"}, + ) + manager.init_lora_modules.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py index 8a0fa3827..f273100de 100644 --- a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py +++ b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py @@ -7,6 +7,7 @@ import torch from sglang.srt.managers.scheduler_components.batch_result_processor import ( SchedulerBatchResultProcessor, ) +from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci @@ -178,17 +179,8 @@ class TestDecodeHiddenStateRetention(CustomTestCase): second_step = torch.arange(16, dtype=torch.float32).view(8, 2)[4:] def result(hidden_states): - return SimpleNamespace( - copy_done=None, - auxiliary_host_output=None, - routed_experts_output=None, - indexer_topk_output=None, + return GenerationBatchResult( logits_output=SimpleNamespace(hidden_states=hidden_states), - next_token_ids=None, - can_run_cuda_graph=False, - num_correct_drafts=0, - num_block_accept_tokens=0, - num_cap_tokens=0, speculative_num_draft_tokens=4, ) diff --git a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py index 7ec18c12a..1f9754e68 100644 --- a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py +++ b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py @@ -10,6 +10,7 @@ from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.scheduler_components.batch_result_processor import ( SchedulerBatchResultProcessor, ) +from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.runtime_context import get_context from sglang.srt.sampling.sampling_params import SamplingParams from sglang.test.ci.ci_register import register_cpu_ci @@ -78,17 +79,9 @@ def _make_processor() -> SchedulerBatchResultProcessor: def _make_result(): - return SimpleNamespace( - copy_done=None, - auxiliary_host_output=None, - routed_experts_output=None, - indexer_topk_output=None, + return GenerationBatchResult( logits_output=SimpleNamespace(hidden_states=None, customized_info=None), next_token_ids=[4], - can_run_cuda_graph=False, - num_correct_drafts=0, - num_block_accept_tokens=0, - num_cap_tokens=0, speculative_num_draft_tokens=0, ) diff --git a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py index f8547049e..796c7a99e 100644 --- a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py +++ b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py @@ -14,6 +14,7 @@ from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.scheduler_components.batch_result_processor import ( SchedulerBatchResultProcessor, ) +from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.sampling.sampling_params import ( REQUEST_REASONING_END_TOKEN_IDS_KEY, SamplingParams, @@ -101,16 +102,10 @@ def _make_req(terminate_after: int) -> Req: def _make_result(num_draft_tokens, accept_lens, flat_tokens): - return SimpleNamespace( + return GenerationBatchResult( next_token_ids=torch.tensor(flat_tokens, dtype=torch.long), accept_lens=torch.tensor(accept_lens, dtype=torch.long), speculative_num_draft_tokens=num_draft_tokens, - num_correct_drafts=None, - num_correct_drafts_per_req_cpu=None, - block_accept_lens=None, - cap_lens=None, - copy_done=None, - grammar_advanced=False, ) diff --git a/test/registered/unit/managers/test_scheduler_uno_request_validation.py b/test/registered/unit/managers/test_scheduler_uno_request_validation.py new file mode 100644 index 000000000..d008b3a1d --- /dev/null +++ b/test/registered/unit/managers/test_scheduler_uno_request_validation.py @@ -0,0 +1,67 @@ +"""Scheduler containment for unsupported UNO requests.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.managers.scheduler import Scheduler + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class TestSchedulerUnoRequestValidation(CustomTestCase): + def test_invalid_request_is_aborted_before_scheduler_admission(self): + scheduler = Scheduler.__new__(Scheduler) + scheduler.enable_session_radix_cache = False + scheduler.model_config = SimpleNamespace( + hf_eos_token_id={1}, + vocab_size=128, + ) + scheduler.disaggregation_mode = DisaggregationMode.NULL + scheduler.metrics_reporter = SimpleNamespace(enable_metrics=False) + scheduler.tokenizer = None + scheduler.dllm_config = None + scheduler._maybe_namespace_elastic_radix_cache = MagicMock() + scheduler.spec_algorithm = SimpleNamespace( + is_dflash_family=lambda: False, + is_uno=lambda: True, + ) + scheduler.init_req_max_new_tokens = MagicMock() + scheduler._add_request_to_queue = MagicMock() + + recv_req = MagicMock( + session_params=None, + session_id=None, + input_embeds=None, + bootstrap_port=1, + ) + req = MagicMock() + error = "UNO request is unsupported." + + with ( + patch( + "sglang.srt.managers.scheduler.BeamCoordinator.request_beam_width", + return_value=1, + ), + patch("sglang.srt.managers.scheduler.Req", return_value=req), + patch( + "sglang.srt.managers.scheduler.validate_uno_request", + return_value=error, + ) as validate_uno_request, + ): + scheduler.handle_generate_request(recv_req) + + validate_uno_request.assert_called_once_with(req) + req.set_finish_with_abort.assert_called_once_with(error) + scheduler.init_req_max_new_tokens.assert_called_once_with(req) + scheduler._add_request_to_queue.assert_called_once_with(req) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_uno_token_accounting.py b/test/registered/unit/managers/test_uno_token_accounting.py new file mode 100644 index 000000000..c450c5c99 --- /dev/null +++ b/test/registered/unit/managers/test_uno_token_accounting.py @@ -0,0 +1,74 @@ +"""CPU regressions for UNO aggregate token accounting.""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.scheduler_components.metrics_reporter import ( + SchedulerMetricsReporter, +) +from sglang.srt.managers.utils import GenerationBatchResult +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestUnoTokenAccounting(CustomTestCase): + def setUp(self): + self.batch_size = 2 + self.result = GenerationBatchResult( + num_correct_drafts=3, + num_non_draft_tokens_per_req=2, + ) + + def test_generated_token_count_includes_both_non_draft_tokens(self): + self.assertEqual(self.result.get_num_generated_tokens(self.batch_size), 7) + self.assertEqual( + GenerationBatchResult(num_correct_drafts=3).get_num_generated_tokens( + self.batch_size + ), + 5, + ) + + def test_spec_metrics_keep_generated_and_draft_counts_separate(self): + reporter = SchedulerMetricsReporter.__new__(SchedulerMetricsReporter) + reporter.spec_num_accept_tokens = 0 + reporter.spec_num_correct_drafts = 0 + reporter.spec_num_forward_ct = 0 + reporter.spec_num_block_accept_tokens = 0 + reporter.spec_num_cap_tokens = 0 + + reporter.update_spec_metrics( + self.batch_size, + self.result.num_correct_drafts, + num_accept_tokens=self.result.get_num_generated_tokens(self.batch_size), + ) + + self.assertEqual(reporter.spec_num_accept_tokens, 7) + self.assertEqual(reporter.spec_num_correct_drafts, 3) + self.assertEqual(reporter.spec_num_forward_ct, 2) + + def test_decode_moment_receives_full_generated_token_count(self): + scheduler = Scheduler.__new__(Scheduler) + scheduler._prev_step = (1, 10.0, False) + scheduler.decode_moment_totals = [0.0] * 6 + batch = SimpleNamespace( + forward_mode=SimpleNamespace( + is_extend_without_speculative=lambda: False, + is_decode=lambda: True, + is_target_verify=lambda: False, + ), + reqs=[SimpleNamespace(rid="req-0"), SimpleNamespace(rid="req-1")], + forward_iter=2, + launch_ts=10.001, + after_idle_gap=False, + ) + + scheduler._record_step_counters(batch, self.result) + + self.assertEqual(scheduler.decode_moment_totals[5], 7) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_uno_allocation_sizing.py b/test/registered/unit/mem_cache/test_uno_allocation_sizing.py new file mode 100644 index 000000000..c06724d2f --- /dev/null +++ b/test/registered/unit/mem_cache/test_uno_allocation_sizing.py @@ -0,0 +1,36 @@ +"""Unit tests for UNO allocation sizing.""" + +import unittest + +from sglang.srt.mem_cache.allocation_sizing import ( + get_alloc_len_per_decode, + get_alloc_reserve_per_decode, + get_req_to_token_extra_context_len, +) +from sglang.srt.runtime_context import get_context, get_parallel +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestUnoAllocationSizing(CustomTestCase): + def test_page_size_one_row_covers_decode_reserve(self): + with ( + get_context().override_server_args( + speculative_algorithm="UNO", + speculative_num_draft_tokens=8, + page_size=1, + ), + get_parallel().override(attn_dcp_size=1), + ): + self.assertEqual(get_alloc_len_per_decode(), 9) + self.assertEqual(get_alloc_reserve_per_decode(), 18) + self.assertGreaterEqual( + get_req_to_token_extra_context_len(), + get_alloc_reserve_per_decode(), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/test_draft_runner_skips_lora.py b/test/registered/unit/model_executor/test_draft_runner_skips_lora.py index 59f7f52a2..31ba12d42 100644 --- a/test/registered/unit/model_executor/test_draft_runner_skips_lora.py +++ b/test/registered/unit/model_executor/test_draft_runner_skips_lora.py @@ -14,6 +14,7 @@ import unittest from unittest.mock import patch from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -25,6 +26,7 @@ class TestDraftRunnerSkipsLoRA(CustomTestCase): runner = ModelRunner.__new__(ModelRunner) runner.is_draft_worker = is_draft_worker runner.lora_manager = None + runner.spec_algorithm = SpeculativeAlgorithm.NONE with patch.object(ModelRunner, "init_lora_manager") as init_lora: with patch("sglang.srt.model_executor.model_runner.get_lora") as get_lora: get_lora.return_value.enable_lora = enable_lora diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index 01d687cca..a7e70a0b2 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py @@ -45,6 +45,10 @@ _DFLASH_DECODE = ( "speculative/dflash_info_v2.py", "DFlashDraftInputV2.prepare_for_decode", ) +_UNO_DECODE = ( + "speculative/uno_info.py", + "UnoDraftInput.prepare_for_decode", +) _RESOLVE = ( "managers/scheduler_components/batch_result_processor.py", "SchedulerBatchResultProcessor._resolve_spec_v2_tokens", @@ -71,6 +75,8 @@ _OWNER_SITES = { # one of these two owners for each speculative decode iteration. (*_DFLASH_DECODE, "decode_batch_idx"): 1, (*_DFLASH_DECODE, "evict"): 1, + (*_UNO_DECODE, "decode_batch_idx"): 1, + (*_UNO_DECODE, "evict"): 1, ( "mem_cache/allocation.py", "alloc_for_spec_decode", diff --git a/test/registered/unit/spec/test_suffix_attention_merge_dispatch.py b/test/registered/unit/spec/test_suffix_attention_merge_dispatch.py new file mode 100644 index 000000000..ff8958671 --- /dev/null +++ b/test/registered/unit/spec/test_suffix_attention_merge_dispatch.py @@ -0,0 +1,72 @@ +"""CPU contracts for the fused suffix-attention merge dispatch guard.""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.kernels.ops.attention.suffix_attention_merge import ( + can_use_fused_suffix_attention_merge, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestSuffixAttentionMergeDispatch(CustomTestCase): + def _inputs(self): + layer = SimpleNamespace( + head_dim=64, + v_head_dim=64, + is_cross_attention=False, + logit_cap=0.0, + ) + q = torch.empty((16, 8 * 64), dtype=torch.bfloat16) + key_cache = torch.empty((8, 16, 2, 64), dtype=torch.bfloat16) + value_cache = torch.empty_like(key_cache) + return layer, q, key_cache, value_cache + + def _eligible(self, **overrides): + layer, q, key_cache, value_cache = self._inputs() + arguments = dict( + layer=layer, + q=q, + key_cache=key_cache, + value_cache=value_cache, + extra_kwargs={}, + ) + arguments.update(overrides) + return can_use_fused_suffix_attention_merge(**arguments) + + def test_standard_attention_is_eligible(self): + self.assertTrue(self._eligible()) + + def test_special_attention_features_fall_back(self): + self.assertFalse(self._eligible(extra_kwargs={"sinks": object()})) + + layer, _, _, _ = self._inputs() + layer.is_cross_attention = True + self.assertFalse(self._eligible(layer=layer)) + + layer, _, _, _ = self._inputs() + layer.logit_cap = 20.0 + self.assertFalse(self._eligible(layer=layer)) + + def test_unsupported_tensor_layout_falls_back(self): + layer, _, _, _ = self._inputs() + layer.v_head_dim = 32 + self.assertFalse(self._eligible(layer=layer)) + + _, q, key_cache, value_cache = self._inputs() + self.assertFalse( + self._eligible( + q=q.float(), + key_cache=key_cache.float(), + value_cache=value_cache.float(), + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/spec/test_uno_request_validation.py b/test/registered/unit/spec/test_uno_request_validation.py new file mode 100644 index 000000000..8e9b7fc97 --- /dev/null +++ b/test/registered/unit/spec/test_uno_request_validation.py @@ -0,0 +1,93 @@ +"""Request-admission validation for UNO speculative decoding.""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.speculative.uno_validation import validate_uno_request +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +def _make_request(**overrides): + sampling_params = SimpleNamespace( + min_p=0.0, + json_schema=None, + regex=None, + ebnf=None, + structural_tag=None, + frequency_penalty=0.0, + presence_penalty=0.0, + repetition_penalty=1.0, + min_new_tokens=0, + logit_bias=None, + ) + request = SimpleNamespace( + sampling_params=sampling_params, + grammar=None, + return_logprob=False, + return_hidden_states_mode=SimpleNamespace(need_capture=lambda: False), + custom_logit_processor=None, + lora_id=None, + ) + for name, value in overrides.items(): + target, field = name.split("__", maxsplit=1) + owner = sampling_params if target == "sampling_params" else request + setattr(owner, field, value) + return request + + +class TestUnoRequestValidation(CustomTestCase): + def test_supported_request_is_accepted(self): + self.assertIsNone(validate_uno_request(_make_request())) + + def test_unsupported_request_features_are_rejected(self): + cases = { + "min_p": ({"sampling_params__min_p": 0.1}, "min_p"), + "grammar": ({"sampling_params__regex": "[0-9]+"}, "grammar"), + "logprobs": ({"request__return_logprob": True}, "logprobs"), + "hidden states": ( + { + "request__return_hidden_states_mode": SimpleNamespace( + need_capture=lambda: True + ) + }, + "return_hidden_states", + ), + "frequency penalty": ( + {"sampling_params__frequency_penalty": 0.1}, + "penalties", + ), + "presence penalty": ( + {"sampling_params__presence_penalty": 0.1}, + "penalties", + ), + "repetition penalty": ( + {"sampling_params__repetition_penalty": 1.1}, + "penalties", + ), + "minimum new tokens": ( + {"sampling_params__min_new_tokens": 1}, + "penalties", + ), + "logit bias": ( + {"sampling_params__logit_bias": {1: 0.5}}, + "logit_bias", + ), + "custom processor": ( + {"request__custom_logit_processor": "processor"}, + "custom logit processors", + ), + "public LoRA": ({"request__lora_id": "adapter"}, "LoRA"), + } + + for name, (overrides, expected) in cases.items(): + with self.subTest(name=name): + error = validate_uno_request(_make_request(**overrides)) + self.assertIsNotNone(error) + self.assertIn(expected, error) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/spec/test_uno_tree_config.py b/test/registered/unit/spec/test_uno_tree_config.py new file mode 100644 index 000000000..3434dd701 --- /dev/null +++ b/test/registered/unit/spec/test_uno_tree_config.py @@ -0,0 +1,73 @@ +"""Startup validation for UNO configuration.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.srt.arg_groups.speculative_hook import _handle_uno +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestUnoTreeConfig(CustomTestCase): + def test_unsupported_runtime_modes_are_rejected_at_startup(self): + cases = { + "deterministic inference": ( + {"enable_deterministic_inference": True}, + "enable-deterministic-inference", + ), + "strict thinking": ( + {"enable_strict_thinking": True}, + "enable-strict-thinking", + ), + } + + for name, (overrides, expected) in cases.items(): + with self.subTest(name=name): + values = { + "device": "cuda", + "speculative_draft_model_path": None, + "uno_lora_path": "/tmp/uno-lora", + "enable_deterministic_inference": False, + "enable_strict_thinking": False, + } + values.update(overrides) + server_args = SimpleNamespace(**values) + with ( + patch( + "sglang.srt.arg_groups.speculative_hook.resolving_view", + side_effect=lambda args: args, + ), + self.assertRaisesRegex(ValueError, expected), + ): + _handle_uno(server_args) + + def test_parent_list_overflow_is_rejected_at_startup(self): + """An invalid tree must not survive startup and crash on first decode.""" + + server_args = SimpleNamespace( + device="cuda", + enable_deterministic_inference=False, + enable_strict_thinking=False, + speculative_draft_model_path=None, + uno_lora_path="/tmp/uno-lora", + speculative_num_draft_tokens=8, + speculative_num_steps=3, + speculative_eagle_topk=2, + ) + + with ( + patch( + "sglang.srt.arg_groups.speculative_hook.resolving_view", + side_effect=lambda args: args, + ), + patch("sglang.srt.arg_groups.speculative_hook.declare_resolution"), + self.assertRaisesRegex(ValueError, "parent-list ABI"), + ): + _handle_uno(server_args) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/spec/test_uno_tree_sparse_sampling.py b/test/registered/unit/spec/test_uno_tree_sparse_sampling.py new file mode 100644 index 000000000..0ba50cffc --- /dev/null +++ b/test/registered/unit/spec/test_uno_tree_sparse_sampling.py @@ -0,0 +1,105 @@ +"""CPU contracts for UNO tree compact target sampling.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.speculative.eagle_utils import ( + _can_use_sparse_uno_tree_target_sampling, +) +from sglang.srt.speculative.uno_utils import sample_uno_tree_target_tokens +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestUnoTreeSparseSampling(CustomTestCase): + def test_sparse_dispatch_guard(self): + sampling_info = SimpleNamespace( + sampling_seed=None, + need_min_p_sampling=False, + ) + spec_config = SimpleNamespace( + speculative_use_rejection_sampling=False, + ) + with ( + patch("sglang.srt.speculative.eagle_utils._is_cuda", True), + patch( + "sglang.srt.speculative.eagle_utils.get_spec", + return_value=spec_config, + ), + ): + self.assertTrue( + _can_use_sparse_uno_tree_target_sampling(128, sampling_info) + ) + self.assertFalse( + _can_use_sparse_uno_tree_target_sampling(None, sampling_info) + ) + self.assertFalse( + _can_use_sparse_uno_tree_target_sampling(129, sampling_info) + ) + + sampling_info.sampling_seed = torch.tensor([1]) + self.assertFalse( + _can_use_sparse_uno_tree_target_sampling(128, sampling_info) + ) + sampling_info.sampling_seed = None + sampling_info.need_min_p_sampling = True + self.assertFalse( + _can_use_sparse_uno_tree_target_sampling(128, sampling_info) + ) + sampling_info.need_min_p_sampling = False + spec_config.speculative_use_rejection_sampling = True + self.assertFalse( + _can_use_sparse_uno_tree_target_sampling(128, sampling_info) + ) + + def test_targets_are_sampled_from_compact_support(self): + support_ids = torch.tensor( + [ + [[10, 11], [20, 21], [30, 31]], + [[40, 41], [50, 51], [60, 61]], + ], + dtype=torch.int64, + ) + support_probs = torch.full((2, 3, 2), 0.5) + sampled_offsets = torch.tensor( + [[0], [1], [0], [1], [0], [1]], + dtype=torch.long, + ) + sampling_info = SimpleNamespace() + next_token_logits = torch.empty((6, 100)) + + with ( + patch( + "sglang.srt.speculative.uno_utils._build_sparse_target_support", + return_value=(support_ids, support_probs), + ) as build_support, + patch( + "sglang.srt.speculative.uno_utils.fast_sample", + return_value=(torch.empty((6, 1)), sampled_offsets), + ), + ): + targets = sample_uno_tree_target_tokens( + next_token_logits=next_token_logits, + sampling_info=sampling_info, + batch_size=2, + verify_width=3, + max_top_k=2, + ) + + self.assertEqual(targets.tolist(), [[10, 21, 30], [41, 50, 61]]) + build_support.assert_called_once_with( + next_token_logits=next_token_logits, + sampling_info=sampling_info, + batch_size=2, + forward_width=3, + max_top_k=2, + ) + + +if __name__ == "__main__": + unittest.main()