diff --git a/python/sglang/benchmark/dspark_sps_profiler.py b/python/sglang/benchmark/dspark_sps_profiler.py new file mode 100644 index 000000000..1f914fa7c --- /dev/null +++ b/python/sglang/benchmark/dspark_sps_profiler.py @@ -0,0 +1,1577 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["msgspec", "requests", "plotly", "kaleido", "numpy"] +# /// +from __future__ import annotations + +import argparse +import importlib.util +import json +import logging +import math +import random +import statistics +import threading +import time +from pathlib import Path +from typing import Optional + +import msgspec +import numpy as np +import requests + +logger = logging.getLogger(__name__) + + +def _load_module_by_path(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +try: + from sglang.benchmark.one_batch_server import ( + DEFAULT_TIMEOUT, + should_skip_due_to_max_running_requests, + should_skip_due_to_token_capacity, + ) + from sglang.benchmark.utils import get_tokenizer + from sglang.srt.speculative.dspark_components.dspark_sps import ( + SpsAdditiveCostTable, + load_sps_table_from_path, + profile_sps_table, + ) +except ImportError as exc: + logger.warning( + "Full sglang runtime unavailable (%s); using a torch-free fallback import. " + "The 'fit' subcommand works; 'run' requires the full sglang install.", + exc, + ) + _benchmark_dir = Path(__file__).resolve().parent + _table_module = _load_module_by_path( + "dspark_sps", + _benchmark_dir.parent / "srt/speculative/dspark_components/dspark_sps.py", + ) + SpsAdditiveCostTable = _table_module.SpsAdditiveCostTable + load_sps_table_from_path = _table_module.load_sps_table_from_path + profile_sps_table = _table_module.profile_sps_table + DEFAULT_TIMEOUT = 60 + get_tokenizer = None + should_skip_due_to_max_running_requests = None + should_skip_due_to_token_capacity = None + +DEFAULT_OUT = "dspark_sps.json" +DEFAULT_MAX_BATCH_SIZE = 256 +DEFAULT_INPUT_LEN = 16 +DEFAULT_TEMPERATURE = 1.0 +DEFAULT_MIN_STEADY_STEPS = 32 +DEFAULT_MIN_STEADY_SECONDS = 10.0 +DEFAULT_ROUND_TIMEOUT_SECONDS = 300.0 +ROUND_WARMUP_STEPS = 8 +ROUND_STEP_SLACK = 64 +STEP_TIME_FLOOR_SECONDS = 0.02 +WARMUP_ROUND_STEADY_STEPS = 16 +POLL_INTERVAL_SECONDS = 2.0 +LOAD_JOIN_TIMEOUT_SECONDS = 60.0 +MATCH_FRACTION_WARN = 0.9 +MATCH_FRACTION_ERROR = 0.5 +PROFILE_SEED = 42 +REQUIRED_SIMULATE_ACC_LEN = 1.0 +RANDOM_TOKEN_LOW = 1000 +RANDOM_TOKEN_HIGH_MARGIN = 1000 + +STATIC_CONDITIONING_CAVEAT = ( + "Profiled with SGLANG_RAGGED_VERIFY_MODE=static: a verify step of B tokens " + "comes from B/(gamma+1) requests, the fewest possible for that B. A " + "compact-mode step with the same B usually spans more requests and reads " + "more KV history, so the table slightly over-estimates steps_per_sec and " + "the scheduler may admit slightly more than optimal. Much smaller bias " + "than the retired non-spec decode proxy, and in the opposite direction." +) +CONVERSION_FORMULA = ( + "batch_tokens = num_running_reqs_per_rank * verify_num_draft_tokens; " + "steps_per_sec = 1 / median(server-side step_time over aligned steady steps)" +) + + +class SpsRow(msgspec.Struct, frozen=True): + forward_ct: int + num_running_reqs: int + num_verify_tokens: int + step_time: float + + +# Server-side per-step record feed: the DsparkInfoDumper 'core' + +# 'step_cpu_time' components exported on /server_info. +INFO_RECORD_PAYLOAD_KEY = "dspark_info_record" +INFO_RECORD_ENABLE_HINT = ( + "SGLANG_DSPARK_ENABLE_SPS_RECORD=1 (or SGLANG_DSPARK_DEBUG_DUMP=" + "core,step_cpu_time) and SGLANG_RAGGED_VERIFY_MODE=static" +) + + +class ServerContext(msgspec.Struct, frozen=True): + base_url: str + tokenizer_path: str + tp_size: int + dp_size: int + verify_num_draft_tokens: int + simulate_acc_len: float + cuda_graph_max_bs: Optional[int] + skip_max_running_requests_threshold: float + skip_token_capacity_threshold: float + + +class RoundSettings(msgspec.Struct, frozen=True): + input_len: int + temperature: float + min_steady_steps: int + min_steady_seconds: float + round_timeout_seconds: float + ramp_token_slack: int = 0 + + +class LoadInfo(msgspec.Struct, frozen=True): + num_requests: int + max_new_tokens: int + wall_seconds: float + reached_target: bool + + +class RoundOutcome(msgspec.Struct, frozen=True): + batch_size: int + batch_size_per_rank: int + batch_tokens: int + steps_per_sec: float + num_steady_steps: int + match_fraction: float + per_rank_median_step_time: list[float] + rank_rows: list[list[SpsRow]] + load_info: LoadInfo + frac: Optional[float] = None + + +def out_paths(*, out: str) -> dict[str, Path]: + out_path = Path(out).expanduser() + return { + "table": out_path, + "records": out_path.with_name(out_path.stem + ".records.jsonl"), + "rounds": out_path.with_name(out_path.stem + ".rounds.jsonl"), + "manifest": out_path.with_name(out_path.name + ".manifest.json"), + "plot": out_path.with_name(out_path.stem + ".plot.png"), + } + + +def run_profile( + *, + base_url: str, + batch_sizes: list[int], + settings: RoundSettings, + out: str, + repeats: int, + local_tokenizer_path: Optional[str], + fracs: Optional[list[float]], +) -> None: + if get_tokenizer is None: + raise RuntimeError( + "'run' needs the full sglang runtime (torch, tokenizers, ...), but this " + "process loaded the torch-free fallback. Run 'run' where sglang is " + "installed; 'fit' works in either environment." + ) + if not base_url: + raise ValueError( + "dspark_sps_profiler connects to an already-running DSpark server " + "(SGLANG_RAGGED_VERIFY_MODE=static, SGLANG_DSPARK_ENABLE_SPS_RECORD=1); " + "pass --base-url (it never launches a server)." + ) + + offdiag = fracs is not None + if offdiag: + for frac in fracs: + if not 0.0 < frac <= 1.0: + raise ValueError( + f"--fracs values must be in (0, 1], got {frac}. The off-diagonal " + "budget pin runs frac * full verify, so frac <= 1.0 keeps M below " + "the full uniform tier and inside the captured cuda graphs." + ) + + paths = out_paths(out=out) + paths["table"].parent.mkdir(parents=True, exist_ok=True) + for path in (paths["records"], paths["rounds"]): + if path.exists(): + path.unlink() + + context = fetch_server_context( + base_url=base_url, + local_tokenizer_path=local_tokenizer_path, + allowed_modes=("compact", "cap-accept") if offdiag else ("static",), + ) + vocab_size = len(get_tokenizer(context.tokenizer_path)) + batch_sizes = sorted(set(batch_sizes)) + validate_sweep_against_server(context=context, batch_sizes=batch_sizes) + rng = random.Random(PROFILE_SEED) + + frac_sweep: list[Optional[float]] = sorted(fracs) if offdiag else [None] + + run_warmup_round( + context=context, + vocab_size=vocab_size, + batch_sizes=batch_sizes, + settings=settings, + rng=rng, + frac=frac_sweep[-1], + ) + + rounds: list[RoundOutcome] = [] + for repeat in range(max(1, repeats)): + for batch_size_per_rank in batch_sizes: + for frac in frac_sweep: + outcome = run_one_round( + context=context, + vocab_size=vocab_size, + batch_size_per_rank=batch_size_per_rank, + settings=settings, + rng=rng, + frac=frac, + ) + if outcome is None: + continue + logger.info( + "Round bs=%s (per-rank %s, frac=%s, batch_tokens=%s) " + "repeat=%s/%s: steps_per_sec=%.3f over %s steady steps " + "(match_fraction=%.2f, wall=%.1fs, per-rank median " + "step_time=%s)", + outcome.batch_size, + outcome.batch_size_per_rank, + outcome.frac, + outcome.batch_tokens, + repeat + 1, + max(1, repeats), + outcome.steps_per_sec, + outcome.num_steady_steps, + outcome.match_fraction, + outcome.load_info.wall_seconds, + ["%.4f" % value for value in outcome.per_rank_median_step_time], + ) + append_round_files( + records_path=paths["records"], + rounds_path=paths["rounds"], + outcome=outcome, + repeat=repeat, + ) + rounds.append(outcome) + + if not rounds: + raise RuntimeError( + "No usable rounds (all were skipped by capacity guards or failed); " + "check the batch-size sweep against the server's " + "max_running_requests / KV capacity." + ) + + write_manifest( + manifest_path=paths["manifest"], + records_path=paths["records"], + rounds_path=paths["rounds"], + context=context, + batch_sizes=batch_sizes, + settings=settings, + repeats=repeats, + rounds=rounds, + fracs=fracs, + ) + logger.info( + "Collected %s rounds; wrote %s, %s and %s", + len(rounds), + paths["rounds"].name, + paths["records"].name, + paths["manifest"].name, + ) + + +def fit_profile( + *, + out: str, + max_batch_tokens: Optional[int], + self_check: bool, + plot: bool, +) -> None: + paths = out_paths(out=out) + if not paths["rounds"].exists(): + raise FileNotFoundError( + f"No rounds file at {paths['rounds']}; run the 'run' subcommand first " + "(or point --out at a prior run's table path)." + ) + + summaries = load_round_summaries(rounds_path=paths["rounds"]) + if not summaries: + raise RuntimeError(f"{paths['rounds']} has no rounds to fit.") + offdiag = any(summary.get("frac") is not None for summary in summaries) + + table = build_table_from_summaries( + summaries=summaries, max_batch_tokens=max_batch_tokens, offdiag=offdiag + ) + paths["table"].write_text(table.to_json(), encoding="utf-8") + if offdiag: + logger.info( + "Fit SpsAdditiveCostTable (%s bs probes x %s M probes) -> %s", + len(table.bs_probes), + len(table.m_probes), + paths["table"], + ) + else: + logger.info( + "Fit SpsCostTable (%s probes) -> %s", + len(table.sample_batch_tokens), + paths["table"], + ) + + if plot: + plot_fit( + cells=summaries_to_cells(summaries=summaries), + table=table, + plot_path=paths["plot"], + ) + + if self_check: + run_self_check(out_path=paths["table"], offdiag=offdiag) + + +def profile_all( + *, + base_url: str, + batch_sizes: list[int], + settings: RoundSettings, + out: str, + max_batch_tokens: Optional[int], + repeats: int, + self_check: bool, + local_tokenizer_path: Optional[str], + fracs: Optional[list[float]], + plot: bool, +) -> None: + run_profile( + base_url=base_url, + batch_sizes=batch_sizes, + settings=settings, + out=out, + repeats=repeats, + local_tokenizer_path=local_tokenizer_path, + fracs=fracs, + ) + fit_profile( + out=out, + max_batch_tokens=max_batch_tokens, + self_check=self_check, + plot=plot, + ) + + +def load_round_summaries(*, rounds_path: Path) -> list[dict]: + return [ + json.loads(line) + for line in rounds_path.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + + +def summaries_to_cells(*, summaries: list[dict]) -> list[dict]: + return [ + { + "bs": summary["batch_size_per_rank"], + "M": summary["batch_tokens"], + "T": 1.0 / summary["steps_per_sec"], + "frac": summary.get("frac"), + } + for summary in summaries + ] + + +def build_table_from_summaries( + *, summaries: list[dict], max_batch_tokens: Optional[int], offdiag: bool +): + if offdiag: + return build_additive_table_from_cells( + cells=summaries_to_cells(summaries=summaries) + ) + + by_batch_tokens: dict[int, list[float]] = {} + for summary in summaries: + by_batch_tokens.setdefault(summary["batch_tokens"], []).append( + summary["steps_per_sec"] + ) + probes = [ + (batch_tokens, statistics.median(values)) + for batch_tokens, values in sorted(by_batch_tokens.items()) + ] + return profile_sps_table(probes=probes, max_batch_tokens=max_batch_tokens) + + +def fetch_server_context( + *, + base_url: str, + local_tokenizer_path: Optional[str], + allowed_modes: tuple[str, ...] = ("static",), +) -> ServerContext: + response = requests.get(base_url + "/server_info", timeout=DEFAULT_TIMEOUT) + response.raise_for_status() + info = response.json() + + speculative_algorithm = info.get("speculative_algorithm") + if speculative_algorithm != "DSPARK": + raise ValueError( + f"Profile against a DSpark server: {base_url} reports " + f"speculative_algorithm={speculative_algorithm!r}. The SPS table is " + "measured from real static-mode DSpark verify steps; relaunch with " + "--speculative-algorithm DSPARK and SGLANG_RAGGED_VERIFY_MODE=static." + ) + if info.get("disable_cuda_graph"): + raise ValueError( + "The server runs with --disable-cuda-graph; an SPS table measured " + "without cuda graphs is uselessly slow. Relaunch with cuda graphs " + "enabled." + ) + + internal_states = info.get("internal_states") or [] + if not internal_states: + raise RuntimeError(f"{base_url}/server_info returned no internal_states.") + sps_payloads = [state.get(INFO_RECORD_PAYLOAD_KEY) for state in internal_states] + for rank_index, payload in enumerate(sps_payloads): + if payload is None: + raise ValueError( + f"DP rank {rank_index} reports no {INFO_RECORD_PAYLOAD_KEY}; " + f"launch the server with {INFO_RECORD_ENABLE_HINT}." + ) + if payload.get("mode") not in allowed_modes: + raise ValueError( + f"{INFO_RECORD_PAYLOAD_KEY}.mode must be one of {allowed_modes}, " + f"got {payload.get('mode')!r} on DP rank {rank_index}." + ) + components = payload.get("components") or [] + missing = {"core", "step_cpu_time"} - set(components) + if missing: + raise ValueError( + f"DP rank {rank_index} {INFO_RECORD_PAYLOAD_KEY} is missing " + f"component(s) {sorted(missing)}; launch with " + f"{INFO_RECORD_ENABLE_HINT}." + ) + if payload.get("simulate_acc_len") != REQUIRED_SIMULATE_ACC_LEN: + raise ValueError( + f"DP rank {rank_index} reports simulate_acc_len=" + f"{payload.get('simulate_acc_len')!r}, but SPS profiling " + f"requires exactly SGLANG_SIMULATE_ACC_LEN=" + f"{REQUIRED_SIMULATE_ACC_LEN} (spec fully ineffective: every " + "step advances every request by exactly the bonus token, so " + "the per-step KV conditioning is deterministic instead of " + "drifting with the model's accept behavior)." + ) + verify_num_draft_tokens = { + int(payload["verify_num_draft_tokens"]) for payload in sps_payloads + } + if len(verify_num_draft_tokens) != 1: + raise RuntimeError( + "DP ranks disagree on verify_num_draft_tokens: " + f"{sorted(verify_num_draft_tokens)}." + ) + + tokenizer_path = local_tokenizer_path or info.get("tokenizer_path") + if not tokenizer_path: + raise RuntimeError( + "Could not resolve a tokenizer path from /server_info; pass " + "--local-tokenizer-path explicitly." + ) + + internal_state = internal_states[0] + dp_size = int(internal_state.get("dp_size") or 1) + cuda_graph_max_bs = resolve_cuda_graph_max_bs(internal_state=internal_state) + max_running_per_dp = internal_state.get("effective_max_running_requests_per_dp", -1) + if max_running_per_dp and max_running_per_dp > 0: + skip_max_running = float(max_running_per_dp * dp_size) + else: + logger.warning( + "Server did not report effective_max_running_requests_per_dp (%s); " + "not clamping the batch-size sweep against the running cap.", + max_running_per_dp, + ) + skip_max_running = float("inf") + + skip_token_capacity = 0.0 + for state in internal_states: + skip_token_capacity += state.get("memory_usage", {}).get( + "token_capacity", 1_000_000_000 + ) + + return ServerContext( + base_url=base_url, + tokenizer_path=tokenizer_path, + tp_size=int(info.get("tp_size", 1) or 1), + dp_size=dp_size, + verify_num_draft_tokens=verify_num_draft_tokens.pop(), + simulate_acc_len=REQUIRED_SIMULATE_ACC_LEN, + cuda_graph_max_bs=cuda_graph_max_bs, + skip_max_running_requests_threshold=skip_max_running, + skip_token_capacity_threshold=skip_token_capacity, + ) + + +def resolve_cuda_graph_max_bs(*, internal_state: dict) -> Optional[int]: + cuda_graph_config = internal_state.get("cuda_graph_config") + if not isinstance(cuda_graph_config, dict): + return None + decode_config = cuda_graph_config.get("decode") + if not isinstance(decode_config, dict): + return None + captured_bs = decode_config.get("bs") + if isinstance(captured_bs, list) and captured_bs: + return int(max(captured_bs)) + max_bs = decode_config.get("max_bs") + if max_bs is not None: + return int(max_bs) + return None + + +def validate_sweep_against_server( + *, context: ServerContext, batch_sizes: list[int] +) -> None: + if context.cuda_graph_max_bs is None: + raise ValueError( + "Could not resolve the server's captured cuda-graph max batch size " + "from /server_info, so the sweep cannot be confirmed to stay inside " + "the captured decode graphs. Steps that fall back to eager silently " + "poison the table with a different perf regime. Relaunch the server " + "so it reports cuda_graph_config in /server_info, or -- if you really " + "want to profile anyway -- delete this raise, but be careful." + ) + max_per_rank = max(batch_sizes) + if max_per_rank > context.cuda_graph_max_bs: + raise ValueError( + f"The sweep reaches {max_per_rank} running requests per DP rank but " + "the server captured decode cuda graphs only up to bs=" + f"{context.cuda_graph_max_bs}; steps beyond it run eager and poison " + "the table. Relaunch the server with a larger --cuda-graph-max-bs-decode " + "or shrink --max-batch-size. If you really want to profile anyway, " + "delete this raise, but be careful." + ) + + +def build_request_count_sweep(max_num_reqs: int) -> list[int]: + if max_num_reqs < 1: + raise ValueError(f"max_num_reqs must be >= 1, got {max_num_reqs}.") + raw = [ + 1, + 2, + 4, + 8, + *range(16, 64, 8), + *range(64, 128, 16), + *range(128, 256, 32), + *range(256, max_num_reqs + 1, 64), + ] + sweep = sorted({value for value in raw if 1 <= value <= max_num_reqs}) + if sweep[-1] != max_num_reqs: + sweep.append(max_num_reqs) + return sweep + + +def round_max_new_tokens(*, settings: RoundSettings) -> int: + # simulate_acc_len is pinned to exactly 1.0 (enforced in + # fetch_server_context), so every step commits exactly the bonus token + # per request. + commit_tokens_per_step = 1 + steady_steps_budget = max( + settings.min_steady_steps, + math.ceil(settings.min_steady_seconds / STEP_TIME_FLOOR_SECONDS), + ) + total_steps = ROUND_WARMUP_STEPS + steady_steps_budget + ROUND_STEP_SLACK + return total_steps * commit_tokens_per_step + settings.ramp_token_slack + + +def run_warmup_round( + *, + context: ServerContext, + vocab_size: int, + batch_sizes: list[int], + settings: RoundSettings, + rng: random.Random, + frac: Optional[float] = None, +) -> None: + warmup_settings = RoundSettings( + input_len=settings.input_len, + temperature=settings.temperature, + min_steady_steps=WARMUP_ROUND_STEADY_STEPS, + min_steady_seconds=0.0, + round_timeout_seconds=settings.round_timeout_seconds, + ramp_token_slack=settings.ramp_token_slack, + ) + try: + run_one_round( + context=context, + vocab_size=vocab_size, + batch_size_per_rank=min(8, max(batch_sizes)), + settings=warmup_settings, + rng=rng, + frac=frac, + ) + except Exception: + logger.warning("Warmup round failed; continuing.", exc_info=True) + + +def run_one_round( + *, + context: ServerContext, + vocab_size: int, + batch_size_per_rank: int, + settings: RoundSettings, + rng: random.Random, + frac: Optional[float] = None, +) -> Optional[RoundOutcome]: + batch_size = batch_size_per_rank * context.dp_size + max_new_tokens = round_max_new_tokens(settings=settings) + if should_skip_due_to_max_running_requests( + batch_size, context.skip_max_running_requests_threshold + ) or should_skip_due_to_token_capacity( + batch_size, + settings.input_len, + max_new_tokens, + context.skip_token_capacity_threshold, + ): + return None + + if frac is not None: + set_forced_budget_frac(base_url=context.base_url, frac=frac) + + flush_cache(base_url=context.base_url) + watermarks = [ + max((row.forward_ct for row in rows), default=-1) + for rows in fetch_rank_rows(base_url=context.base_url) + ] + + start_time = time.monotonic() + load_thread = start_load( + base_url=context.base_url, + num_requests=batch_size, + input_len=settings.input_len, + max_new_tokens=max_new_tokens, + temperature=settings.temperature, + vocab_size=vocab_size, + rng=rng, + ) + reached_target = wait_for_aligned_steps( + context=context, + watermarks=watermarks, + batch_size_per_rank=batch_size_per_rank, + min_steady_steps=settings.min_steady_steps, + min_steady_seconds=settings.min_steady_seconds, + timeout_seconds=settings.round_timeout_seconds, + ) + abort_all_requests(base_url=context.base_url) + load_thread.join(timeout=LOAD_JOIN_TIMEOUT_SECONDS) + if load_thread.is_alive(): + logger.warning( + "Load batch for bs=%s did not return within %.0fs after abort; " + "continuing with the collected records.", + batch_size, + LOAD_JOIN_TIMEOUT_SECONDS, + ) + wall_seconds = time.monotonic() - start_time + if not reached_target: + logger.warning( + "Round bs=%s hit the %.0fs timeout before both gates (>=%s steady " + "steps and >=%.1fs) were met; proceeding with what was collected.", + batch_size, + settings.round_timeout_seconds, + settings.min_steady_steps, + settings.min_steady_seconds, + ) + + rank_rows = fetch_rank_rows(base_url=context.base_url) + if len(rank_rows) != len(watermarks): + raise RuntimeError( + f"DP rank count changed mid-profile: {len(watermarks)} -> " + f"{len(rank_rows)}." + ) + new_rank_rows = [ + [row for row in rows if row.forward_ct > watermark] + for rows, watermark in zip(rank_rows, watermarks) + ] + return postprocess_round( + rank_rows=new_rank_rows, + batch_size_per_rank=batch_size_per_rank, + dp_size=context.dp_size, + verify_num_draft_tokens=context.verify_num_draft_tokens, + min_steady_steps=settings.min_steady_steps, + load_info=LoadInfo( + num_requests=batch_size, + max_new_tokens=max_new_tokens, + wall_seconds=round(wall_seconds, 3), + reached_target=reached_target, + ), + frac=frac, + ) + + +def start_load( + *, + base_url: str, + num_requests: int, + input_len: int, + max_new_tokens: int, + temperature: float, + vocab_size: int, + rng: random.Random, +) -> threading.Thread: + token_high = vocab_size - RANDOM_TOKEN_HIGH_MARGIN + if token_high <= RANDOM_TOKEN_LOW: + raise ValueError(f"vocab_size={vocab_size} too small for random prompts.") + input_ids = [ + [rng.randrange(RANDOM_TOKEN_LOW, token_high) for _ in range(input_len)] + for _ in range(num_requests) + ] + payload = { + "input_ids": input_ids, + "sampling_params": { + "temperature": temperature, + "max_new_tokens": max_new_tokens, + "ignore_eos": True, + }, + "stream": False, + } + + def _post() -> None: + try: + requests.post(base_url + "/generate", json=payload, timeout=DEFAULT_TIMEOUT) + except Exception: + logger.warning( + "Load batch POST /generate failed (expected on abort for some " + "server versions).", + exc_info=True, + ) + + thread = threading.Thread(target=_post, daemon=True) + thread.start() + return thread + + +def wait_for_aligned_steps( + *, + context: ServerContext, + watermarks: list[int], + batch_size_per_rank: int, + min_steady_steps: int, + min_steady_seconds: float, + timeout_seconds: float, +) -> bool: + deadline = time.monotonic() + timeout_seconds + steady_start: Optional[float] = None + while time.monotonic() < deadline: + time.sleep(POLL_INTERVAL_SECONDS) + try: + rank_rows = fetch_rank_rows(base_url=context.base_url) + except Exception: + logger.warning("Polling /server_info failed; retrying.", exc_info=True) + continue + new_rank_rows = [ + [row for row in rows if row.forward_ct > watermark] + for rows, watermark in zip(rank_rows, watermarks) + ] + if len(new_rank_rows) != len(watermarks): + continue + aligned = count_aligned_steps( + rank_rows=new_rank_rows, batch_size_per_rank=batch_size_per_rank + ) + if aligned >= ROUND_WARMUP_STEPS and steady_start is None: + steady_start = time.monotonic() + steady_steps = max(0, aligned - ROUND_WARMUP_STEPS) + steady_seconds = ( + time.monotonic() - steady_start if steady_start is not None else 0.0 + ) + logger.debug( + "Aligned-step poll: %d aligned (%d/%d steady steps, %.1f/%.1fs)", + aligned, + steady_steps, + min_steady_steps, + steady_seconds, + min_steady_seconds, + ) + if steady_steps >= min_steady_steps and steady_seconds >= min_steady_seconds: + return True + return False + + +def count_aligned_steps( + *, rank_rows: list[list[SpsRow]], batch_size_per_rank: int +) -> int: + if any(not rows for rows in rank_rows): + return 0 + by_ct_per_rank = [{row.forward_ct: row for row in rows} for rows in rank_rows] + common_cts = set(by_ct_per_rank[0]) + for by_ct in by_ct_per_rank[1:]: + common_cts &= set(by_ct) + return sum( + 1 + for ct in common_cts + if all( + by_ct[ct].num_running_reqs == batch_size_per_rank + for by_ct in by_ct_per_rank + ) + ) + + +def abort_all_requests(*, base_url: str) -> None: + response = requests.post( + base_url + "/abort_request", + json={"abort_all": True}, + timeout=DEFAULT_TIMEOUT, + ) + response.raise_for_status() + + +def flush_cache(*, base_url: str) -> None: + try: + requests.post(base_url + "/flush_cache", timeout=DEFAULT_TIMEOUT) + except Exception: + logger.warning("POST /flush_cache failed; continuing.", exc_info=True) + + +def set_forced_budget_frac(*, base_url: str, frac: Optional[float]) -> None: + response = requests.post( + base_url + "/set_internal_state", + json={"server_args": {"dspark_force_budget_frac": frac}}, + timeout=DEFAULT_TIMEOUT, + ).json() + outs = response if isinstance(response, list) else [response] + + def _ok(out) -> bool: + return bool(out.get("updated")) if isinstance(out, dict) else bool(out) + + if not outs or not all(_ok(out) for out in outs): + raise RuntimeError( + f"set dspark_force_budget_frac={frac} rejected by server: {response}" + ) + + +def fetch_rank_rows(*, base_url: str) -> list[list[SpsRow]]: + response = requests.get(base_url + "/server_info", timeout=DEFAULT_TIMEOUT) + response.raise_for_status() + internal_states = response.json().get("internal_states") or [] + rank_rows: list[list[SpsRow]] = [] + for state in internal_states: + payload = state.get(INFO_RECORD_PAYLOAD_KEY) or {} + rows: list[SpsRow] = [] + for record in payload.get("records", []): + step_cpu_ms = record.get("step_cpu_ms") + if step_cpu_ms is None: + continue + rows.append( + SpsRow( + forward_ct=int(record["forward_ct"]), + num_running_reqs=int(record["num_running_reqs"]), + num_verify_tokens=int(record["num_verify_tokens"]), + step_time=float(step_cpu_ms) / 1000.0, + ) + ) + rank_rows.append(rows) + return rank_rows + + +def postprocess_round( + *, + rank_rows: list[list[SpsRow]], + batch_size_per_rank: int, + dp_size: int, + verify_num_draft_tokens: int, + min_steady_steps: int, + load_info: LoadInfo, + frac: Optional[float] = None, +) -> RoundOutcome: + offdiag = frac is not None + batch_size = batch_size_per_rank * dp_size + expected_tokens = batch_size_per_rank * verify_num_draft_tokens + + if len(rank_rows) != dp_size: + raise RuntimeError( + f"Expected records from {dp_size} DP ranks, got {len(rank_rows)}." + ) + + by_ct_per_rank: list[dict[int, SpsRow]] = [] + for rank_index, rows in enumerate(rank_rows): + if not rows: + raise RuntimeError( + f"DP rank {rank_index} produced no new decode-step records this " + "round; the load generator did not reach it (DP imbalance or " + "the round was too short)." + ) + by_ct_per_rank.append({row.forward_ct: row for row in rows}) + + common_cts = set(by_ct_per_rank[0]) + for by_ct in by_ct_per_rank[1:]: + common_cts &= set(by_ct) + if not common_cts: + raise RuntimeError( + "DP ranks share no common forward_ct in this round; their step " + "counters are misaligned, so per-step cross-rank checks are " + "impossible. This breaks the uniformity assumption of the table." + ) + + aligned_cts: list[int] = [] + aligned_verify_tokens: set[int] = set() + for ct in sorted(common_cts): + rows_at_ct = [by_ct[ct] for by_ct in by_ct_per_rank] + if all(row.num_running_reqs == batch_size_per_rank for row in rows_at_ct): + for rank_index, row in enumerate(rows_at_ct): + if not offdiag and row.num_verify_tokens < expected_tokens: + raise RuntimeError( + f"DP rank {rank_index} at forward_ct={ct} reports " + f"num_verify_tokens={row.num_verify_tokens}, expected at " + f"least {expected_tokens} (= {batch_size_per_rank} reqs x " + f"{verify_num_draft_tokens}); ranks are not running the " + "uniform static verify the table assumes. The recorded " + "count is the replayed graph tier, which may exceed the " + "candidate count when a bs is not an exact capture tier." + ) + aligned_verify_tokens.add(row.num_verify_tokens) + aligned_cts.append(ct) + + if len(aligned_cts) < ROUND_WARMUP_STEPS + min_steady_steps: + raise RuntimeError( + f"Round bs={batch_size} never stabilized: only {len(aligned_cts)} " + f"of {len(common_cts)} common decode steps had every rank at the " + f"target {batch_size_per_rank} requests (need at least " + f"{ROUND_WARMUP_STEPS + min_steady_steps}). Increase " + "--round-timeout / --target-steady-steps, or inspect the raw " + "records for retraction / DP imbalance." + ) + + window_cts = [ + ct for ct in sorted(common_cts) if aligned_cts[0] <= ct <= aligned_cts[-1] + ] + match_fraction = len(aligned_cts) / len(window_cts) + if match_fraction < MATCH_FRACTION_ERROR: + raise RuntimeError( + f"Round bs={batch_size} is unstable mid-round: only " + f"{match_fraction:.0%} of the {len(window_cts)} decode steps inside " + "the steady window ran at the target per-rank batch (retraction or " + "DP imbalance, not just ramp-in/drain). Inspect the raw records." + ) + if match_fraction < MATCH_FRACTION_WARN: + logger.warning( + "Round bs=%s: %.0f%% of %s steady-window decode steps ran at the " + "target per-rank batch; treat this probe with suspicion.", + batch_size, + match_fraction * 100.0, + len(window_cts), + ) + + steady_cts = aligned_cts[ROUND_WARMUP_STEPS:] + per_ct_step_times = [ + statistics.fmean(by_ct[ct].step_time for by_ct in by_ct_per_rank) + for ct in steady_cts + ] + per_rank_median_step_time = [ + statistics.median(by_ct[ct].step_time for ct in steady_cts) + for by_ct in by_ct_per_rank + ] + median_step_time = statistics.median(per_ct_step_times) + + if offdiag: + if len(aligned_verify_tokens) != 1: + raise RuntimeError( + f"Round bs={batch_size} frac={frac}: aligned steps ran at " + f"differing num_verify_tokens {sorted(aligned_verify_tokens)}; the " + "budget pin did not hold a single graph tier across all " + "ranks/steps, so the measurement is ambiguous. Inspect the raw " + "records." + ) + graph_tier = aligned_verify_tokens.pop() + budget = int(frac * batch_size_per_rank * (verify_num_draft_tokens - 1)) + batch_tokens = batch_size_per_rank + budget + if graph_tier < batch_tokens: + raise RuntimeError( + f"Round bs={batch_size} frac={frac}: replayed graph tier " + f"{graph_tier} is smaller than the pinned M={batch_tokens} " + f"(= {batch_size_per_rank} + int({frac} * {batch_size_per_rank} " + f"* {verify_num_draft_tokens - 1})); the budget pin did not take." + ) + else: + batch_tokens = expected_tokens + + return RoundOutcome( + batch_size=batch_size, + batch_size_per_rank=batch_size_per_rank, + batch_tokens=batch_tokens, + steps_per_sec=1.0 / median_step_time, + num_steady_steps=len(steady_cts), + match_fraction=match_fraction, + per_rank_median_step_time=per_rank_median_step_time, + rank_rows=rank_rows, + load_info=load_info, + frac=frac, + ) + + +def fitted_step_time(*, table, bs: int, m: int) -> float: + if isinstance(table, SpsAdditiveCostTable): + return table.step_time(num_reqs=bs, budget=max(0, m - bs)) + return 1.0 / table.lookup(m) + + +def bs_color_map(*, batch_sizes: list[int]) -> dict[int, str]: + span = max(len(batch_sizes) - 1, 1) + colors = {} + for index, bs in enumerate(batch_sizes): + hue = 240.0 * (1.0 - index / span) + colors[bs] = f"hsl({hue:.0f}, 70%, 50%)" + return colors + + +def plot_fit(*, cells: list[dict], table, plot_path: Path) -> None: + try: + import plotly.graph_objects as go + from plotly.subplots import make_subplots + except ImportError: + logger.warning( + "plotly not installed; skipping the fit plot. Install plotly + " + "kaleido to render %s.", + plot_path.name, + ) + return + + batch_sizes = sorted({cell["bs"] for cell in cells}) + color_of = bs_color_map(batch_sizes=batch_sizes) + fig = make_subplots( + rows=1, + cols=3, + horizontal_spacing=0.07, + subplot_titles=( + "step time", + "throughput = M / T", + "raw (circle) vs fit (square)", + ), + ) + for bs in batch_sizes: + points = sorted( + (cell for cell in cells if cell["bs"] == bs), key=lambda c: c["M"] + ) + m_values = [cell["M"] for cell in points] + t_ms = [cell["T"] * 1e3 for cell in points] + color = color_of[bs] + fig.add_trace( + go.Scatter( + x=m_values, + y=t_ms, + mode="markers+lines", + name=f"bs={bs}", + legendgroup=f"bs={bs}", + marker=dict(color=color, size=7), + line=dict(color=color, width=1), + ), + row=1, + col=1, + ) + fig.add_trace( + go.Scatter( + x=m_values, + y=[cell["M"] / cell["T"] for cell in points], + mode="markers+lines", + legendgroup=f"bs={bs}", + showlegend=False, + marker=dict(color=color, size=7), + line=dict(color=color, width=1), + ), + row=1, + col=2, + ) + fig.add_trace( + go.Scatter( + x=m_values, + y=t_ms, + mode="markers", + legendgroup=f"bs={bs}", + showlegend=False, + marker=dict(color=color, size=8, symbol="circle"), + ), + row=1, + col=3, + ) + fig.add_trace( + go.Scatter( + x=m_values, + y=[ + fitted_step_time(table=table, bs=bs, m=cell["M"]) * 1e3 + for cell in points + ], + mode="markers", + legendgroup=f"bs={bs}", + showlegend=False, + marker=dict( + color=color, size=9, symbol="square-open", line=dict(width=2) + ), + ), + row=1, + col=3, + ) + fig.update_xaxes( + title_text="M = num total verify tokens", rangemode="tozero", row=1, col=1 + ) + fig.update_xaxes( + title_text="M = num total verify tokens", rangemode="tozero", row=1, col=2 + ) + fig.update_xaxes( + title_text="M = num total verify tokens", rangemode="tozero", row=1, col=3 + ) + fig.update_yaxes(title_text="T = step time (ms)", rangemode="tozero", row=1, col=1) + fig.update_yaxes( + title_text="throughput (tokens/s)", rangemode="tozero", row=1, col=2 + ) + fig.update_yaxes(title_text="T = step time (ms)", rangemode="tozero", row=1, col=3) + fig.update_layout( + title="DSpark SPS profiler: raw cells vs additive fit", + legend_title="batch size", + template="plotly_white", + width=1700, + height=640, + ) + try: + fig.write_image(str(plot_path), scale=2) + except Exception: + logger.warning( + "Failed to render %s (kaleido missing?); skipping plot.", + plot_path.name, + exc_info=True, + ) + return + logger.info("Wrote fit plot to %s", plot_path) + + +def build_additive_table_from_cells(*, cells: list[dict]) -> SpsAdditiveCostTable: + if len(cells) < 4: + raise RuntimeError( + f"Off-diagonal fit needs at least 4 cells, got {len(cells)}." + ) + bias, alpha, theta, _rel, _stats = ols_resid_backfit(cells) + bs_probes = sorted(alpha) + m_probes = sorted(theta) + return SpsAdditiveCostTable( + bias_seconds=bias, + bs_probes=bs_probes, + alpha_seconds=[alpha[b] for b in bs_probes], + m_probes=m_probes, + theta_seconds=[theta[m] for m in m_probes], + ) + + +def ols_resid_backfit(cells: list, mbin_w: int = 64): + def mbin(m): + return round(m / mbin_w) * mbin_w + + bslist = sorted({c["bs"] for c in cells}) + mbins = sorted({mbin(c["M"]) for c in cells}) + bs_col = {b: i for i, b in enumerate(bslist[1:])} + m_col = {m: i for i, m in enumerate(mbins[1:])} + num_cols = 1 + len(bs_col) + len(m_col) + + design = np.zeros((len(cells), num_cols)) + target = np.array([c["T"] for c in cells], dtype=float) + for row, c in enumerate(cells): + design[row, 0] = 1.0 + if c["bs"] in bs_col: + design[row, 1 + bs_col[c["bs"]]] = 1.0 + if mbin(c["M"]) in m_col: + design[row, 1 + len(bs_col) + m_col[mbin(c["M"])]] = 1.0 + + beta, _, _, _ = np.linalg.lstsq(design, target, rcond=None) + bias = float(beta[0]) + alpha = {bslist[0]: 0.0} + for b in bslist[1:]: + alpha[b] = float(beta[1 + bs_col[b]]) + theta = {mbins[0]: 0.0} + for m in mbins[1:]: + theta[m] = float(beta[1 + len(bs_col) + m_col[m]]) + + resid = [c["T"] - (bias + alpha[c["bs"]] + theta[mbin(c["M"])]) for c in cells] + rel = [abs(r) / c["T"] * 100 for r, c in zip(resid, cells)] + rms = (sum(r * r for r in resid) / len(resid)) ** 0.5 + tbar = statistics.fmean(c["T"] for c in cells) + ss_tot = sum((c["T"] - tbar) ** 2 for c in cells) + r2 = 1.0 - sum(r * r for r in resid) / ss_tot if ss_tot > 0 else float("nan") + + def probe_se(pred): + se = {} + for key in sorted({pred(c) for c in cells}): + rs = [r for r, c in zip(resid, cells) if pred(c) == key] + se[key] = (statistics.pstdev(rs) / (len(rs) ** 0.5)) if len(rs) > 1 else 0.0 + return se + + stats = { + "rms_ms": rms * 1e3, + "r2": r2, + "n": len(cells), + "alpha_se": probe_se(lambda c: c["bs"]), + "theta_se": probe_se(lambda c: mbin(c["M"])), + } + return bias, alpha, theta, rel, stats + + +def round_summary_dict(*, outcome: RoundOutcome, repeat: int) -> dict: + return { + "repeat": repeat, + "batch_size": outcome.batch_size, + "batch_size_per_rank": outcome.batch_size_per_rank, + "frac": outcome.frac, + "batch_tokens": outcome.batch_tokens, + "steps_per_sec": outcome.steps_per_sec, + "num_steady_steps": outcome.num_steady_steps, + "match_fraction": outcome.match_fraction, + "per_rank_median_step_time": outcome.per_rank_median_step_time, + "load_info": msgspec.to_builtins(outcome.load_info), + } + + +def append_round_files( + *, + records_path: Path, + rounds_path: Path, + outcome: RoundOutcome, + repeat: int, +) -> None: + with records_path.open("a", encoding="utf-8") as fout: + for rank_index, rows in enumerate(outcome.rank_rows): + for row in rows: + fout.write( + json.dumps( + { + "repeat": repeat, + "batch_size": outcome.batch_size, + "batch_size_per_rank": outcome.batch_size_per_rank, + "dp_rank": rank_index, + "forward_ct": row.forward_ct, + "num_running_reqs": row.num_running_reqs, + "num_verify_tokens": row.num_verify_tokens, + "step_time": row.step_time, + } + ) + + "\n" + ) + with rounds_path.open("a", encoding="utf-8") as fout: + fout.write( + json.dumps(round_summary_dict(outcome=outcome, repeat=repeat)) + "\n" + ) + + +def write_manifest( + *, + manifest_path: Path, + records_path: Path, + rounds_path: Path, + context: ServerContext, + batch_sizes: list[int], + settings: RoundSettings, + repeats: int, + rounds: list[RoundOutcome], + fracs: Optional[list[float]], +) -> None: + manifest = { + "base_url": context.base_url, + "tp_size": context.tp_size, + "dp_size": context.dp_size, + "verify_num_draft_tokens": context.verify_num_draft_tokens, + "simulate_acc_len": context.simulate_acc_len, + "batch_size_per_rank_sweep": batch_sizes, + "fracs": fracs, + "settings": msgspec.to_builtins(settings), + "repeats": repeats, + "seed": PROFILE_SEED, + "timestamp": time.time(), + "timestamp_iso": time.strftime("%Y-%m-%dT%H:%M:%S%z", time.localtime()), + "conversion_formula": CONVERSION_FORMULA, + "static_conditioning_caveat": STATIC_CONDITIONING_CAVEAT, + "records_jsonl": records_path.name, + "rounds_jsonl": rounds_path.name, + "round_summaries": [ + round_summary_dict(outcome=outcome, repeat=0) for outcome in rounds + ], + } + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") + + +def run_self_check(*, out_path: Path, offdiag: bool) -> None: + table = load_sps_table_from_path(str(out_path)) + if offdiag: + run_additive_self_check(table=table) + return + if len(table.sample_batch_tokens) != len(table.sample_steps_per_sec): + raise RuntimeError("Reloaded table has mismatched probe / SPS lengths.") + + previous_sps: Optional[float] = None + for batch_tokens in table.sample_batch_tokens: + looked_up = table.lookup(batch_tokens) + if looked_up <= 0: + raise RuntimeError( + f"Reloaded table lookup at batch_tokens={batch_tokens} returned " + f"non-positive SPS {looked_up}." + ) + if previous_sps is not None and looked_up > previous_sps * 1.10: + logger.warning( + "Non-monotone SPS across probes: batch_tokens=%s SPS=%.3f rose " + "above the previous probe's SPS=%.3f by >10%%; verify the server " + "was at steady state (no co-tenants, steady clocks).", + batch_tokens, + looked_up, + previous_sps, + ) + previous_sps = looked_up + + below_floor = table.lookup(table.sample_batch_tokens[0] - 1) + if below_floor != table.sample_steps_per_sec[0]: + raise RuntimeError( + "Reloaded table lookup below the smallest probe did not clamp to the " + f"first SPS ({below_floor} != {table.sample_steps_per_sec[0]})." + ) + logger.info( + "Self-check passed: reloaded %s probes, all lookups positive and " + "below-floor clamp holds.", + len(table.sample_batch_tokens), + ) + + +def run_additive_self_check(*, table: SpsAdditiveCostTable) -> None: + for num_reqs in table.bs_probes: + for budget in (0, max(table.m_probes) - table.bs_probes[0]): + value = table.step_time(num_reqs=num_reqs, budget=max(0, budget)) + if not value > 0: + raise RuntimeError( + f"Reloaded additive table step_time(num_reqs={num_reqs}, " + f"budget={budget}) is non-positive ({value})." + ) + logger.info( + "Self-check passed: reloaded additive table (%s bs probes x %s M " + "probes), all step_time lookups positive.", + len(table.bs_probes), + len(table.m_probes), + ) + + +def add_out_arg(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--out", + type=str, + default=DEFAULT_OUT, + help="Output JSON path for the SPS table. Raw per-step records land next " + "to it as .records.jsonl, one line per (bs, frac) cell as " + ".rounds.jsonl, .manifest.json ties everything together, and " + "the fit plot as .plot.png.", + ) + parser.add_argument( + "--log-level", + type=str, + default="info", + help="Python logging level for the profiler.", + ) + + +def add_run_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--base-url", + type=str, + default="", + help="Base URL of the already-running DSpark server, e.g. " + "http://localhost:30000. The profiler never launches a server.", + ) + parser.add_argument( + "--fracs", + type=float, + nargs="+", + default=None, + help="Off-diagonal K-fraction sweep in (0, 1]. When given, the server " + "must run SGLANG_RAGGED_VERIFY_MODE=compact and each (bs, frac) cell " + "pins dspark_force_budget_frac to profile T(bs, M); the fit is a 2D " + "SpsAdditiveCostTable. When omitted, the diagonal static sweep runs and " + "the fit is a 1D SpsCostTable.", + ) + parser.add_argument( + "--batch-size", + type=int, + nargs="+", + default=None, + help="Explicit PER-DP-RANK running-request counts to sweep; the load " + "generator sends value * dp_size requests so every rank (GPU group) " + "sits at the given batch. Overrides --max-batch-size when given.", + ) + parser.add_argument( + "--max-batch-size", + type=int, + default=DEFAULT_MAX_BATCH_SIZE, + help="Upper bound of the auto-generated tapered PER-DP-RANK " + "request-count sweep (used only when --batch-size is not given), so " + "per-rank token coverage is identical for any dp_size.", + ) + parser.add_argument( + "--input-len", + type=int, + default=DEFAULT_INPUT_LEN, + help="Prompt length per request. Short: the table is conditioned on the " + "decode-heavy regime.", + ) + parser.add_argument( + "--temperature", + type=float, + default=DEFAULT_TEMPERATURE, + help="Sampling temperature for the load requests; default 1.0 to hit " + "the same accept/sampling kernels as real serving.", + ) + parser.add_argument( + "--min-steady-steps", + type=int, + default=DEFAULT_MIN_STEADY_STEPS, + help="Stop a round only after at least this many aligned steady steps " + "(and --min-steady-seconds) have been collected, then abort the load " + "batch. Bounds cheap small-batch rounds; the batch is held at exactly " + "the target running-request count the whole time (no drain tail).", + ) + parser.add_argument( + "--min-steady-seconds", + type=float, + default=DEFAULT_MIN_STEADY_SECONDS, + help="Stop a round only after at least this much steady-state wall time " + "(and --min-steady-steps) has elapsed. Bounds expensive large-batch " + "rounds, where a fixed step count would run many slow steps.", + ) + parser.add_argument( + "--round-timeout", + type=float, + default=DEFAULT_ROUND_TIMEOUT_SECONDS, + help="Per-round wall-clock budget in seconds to collect the target " + "steps before giving up and using what was collected.", + ) + parser.add_argument( + "--ramp-token-slack", + type=int, + default=0, + help="Extra per-request tokens on top of the step budget so requests " + "outlive the whole-batch prefill ramp. Required for long --input-len " + "at high batch sizes, where the ramp exceeds the request lifetime and " + "full-batch alignment becomes unreachable; size it as roughly " + "ramp_seconds / step_time.", + ) + parser.add_argument( + "--repeats", + type=int, + default=1, + help="Times to repeat the whole sweep; per batch_tokens the median " + "steps_per_sec across repeats is taken.", + ) + parser.add_argument( + "--local-tokenizer-path", + type=str, + default=None, + help="Override the tokenizer path (defaults to the one reported by " + "/server_info).", + ) + + +def add_fit_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--max-batch-tokens", + type=int, + default=None, + help="Override the diagonal table's max_batch_tokens metadata (defaults " + "to the largest probed batch_tokens). Ignored for off-diagonal fits.", + ) + parser.add_argument( + "--no-self-check", + dest="self_check", + action="store_false", + help="Skip the read-back + lookup self-check of the written table.", + ) + parser.add_argument( + "--no-plot", + dest="plot", + action="store_false", + help="Skip the .plot.png fit plot (needs plotly + kaleido).", + ) + + +def run_settings(*, args: argparse.Namespace) -> RoundSettings: + return RoundSettings( + input_len=args.input_len, + temperature=args.temperature, + min_steady_steps=args.min_steady_steps, + min_steady_seconds=args.min_steady_seconds, + round_timeout_seconds=args.round_timeout, + ramp_token_slack=args.ramp_token_slack, + ) + + +def run_batch_sizes(*, args: argparse.Namespace) -> list[int]: + if args.batch_size is not None: + return args.batch_size + return build_request_count_sweep(args.max_batch_size) + + +def cli_main() -> None: + parser = argparse.ArgumentParser( + description=( + "Profile a DSpark SPS cost table from an already-running DSpark " + "server. Subcommands: 'run' collects raw per-cell data, 'fit' builds " + "the table (and plot) from that data, 'all' does both." + ) + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + run_parser = subparsers.add_parser( + "run", help="Collect raw sweep data (append one line per cell to jsonl)." + ) + add_out_arg(run_parser) + add_run_args(run_parser) + + fit_parser = subparsers.add_parser( + "fit", help="Fit the table and render the plot from a prior run's jsonl." + ) + add_out_arg(fit_parser) + add_fit_args(fit_parser) + + all_parser = subparsers.add_parser("all", help="Run then fit in one shot.") + add_out_arg(all_parser) + add_run_args(all_parser) + add_fit_args(all_parser) + + args = parser.parse_args() + logging.basicConfig( + level=getattr(logging, args.log_level.upper()), + format="%(message)s", + ) + + if args.command == "run": + run_profile( + base_url=args.base_url, + batch_sizes=run_batch_sizes(args=args), + settings=run_settings(args=args), + out=args.out, + repeats=args.repeats, + local_tokenizer_path=args.local_tokenizer_path, + fracs=args.fracs, + ) + elif args.command == "fit": + fit_profile( + out=args.out, + max_batch_tokens=args.max_batch_tokens, + self_check=args.self_check, + plot=args.plot, + ) + else: + profile_all( + base_url=args.base_url, + batch_sizes=run_batch_sizes(args=args), + settings=run_settings(args=args), + out=args.out, + max_batch_tokens=args.max_batch_tokens, + repeats=args.repeats, + self_check=args.self_check, + local_tokenizer_path=args.local_tokenizer_path, + fracs=args.fracs, + plot=args.plot, + ) + + +if __name__ == "__main__": + cli_main() diff --git a/python/sglang/benchmark/dspark_sts_fit.py b/python/sglang/benchmark/dspark_sts_fit.py new file mode 100644 index 000000000..ec8f09431 --- /dev/null +++ b/python/sglang/benchmark/dspark_sts_fit.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +import argparse +import glob +import logging +import math +from pathlib import Path +from typing import Optional + +import torch + +from sglang.srt.speculative.dspark_components.dspark_sts import ( + DSparkStsCalibration, +) + +logger = logging.getLogger(__name__) + +_EPS_PROB = 1e-8 + + +def default_temperature_grid() -> torch.Tensor: + return torch.logspace(math.log10(0.1), math.log10(10.0), steps=41) + + +def expected_calibration_error( + *, + probs: torch.Tensor, + targets: torch.Tensor, + num_bins: int, +) -> float: + probs = probs.reshape(-1).to(torch.float64).clamp(_EPS_PROB, 1.0 - _EPS_PROB) + targets = targets.reshape(-1).to(torch.float64) + total = probs.numel() + if total == 0: + return float("nan") + bin_index = (probs * num_bins).long().clamp_(0, num_bins - 1) + count = torch.zeros(num_bins, dtype=torch.float64) + pred_sum = torch.zeros(num_bins, dtype=torch.float64) + target_sum = torch.zeros(num_bins, dtype=torch.float64) + count.scatter_add_(0, bin_index, torch.ones_like(probs)) + pred_sum.scatter_add_(0, bin_index, probs) + target_sum.scatter_add_(0, bin_index, targets) + denom = count.clamp_min(1.0) + bin_error = (pred_sum / denom - target_sum / denom).abs() + return float((bin_error * count).sum().item() / total) + + +def fit_sts_temperatures( + *, + logits: torch.Tensor, + prefix_mask: torch.Tensor, + grid: torch.Tensor, + num_bins: int = 15, +) -> dict[str, list[float]]: + logits = logits.to(torch.float64) + prefix_mask = prefix_mask.to(torch.float64) + num_samples, gamma = logits.shape + if num_samples == 0: + raise ValueError("fit_sts_temperatures requires at least one sample.") + grid_values = grid.to(torch.float64).tolist() + + temperatures: list[float] = [] + ece_before: list[float] = [] + ece_after: list[float] = [] + + survival_at_one = torch.ones(num_samples, dtype=torch.float64) + survival_fitted = torch.ones(num_samples, dtype=torch.float64) + for position in range(gamma): + position_logits = logits[:, position] + position_target = prefix_mask[:, position] + + survival_at_one = survival_at_one * torch.sigmoid(position_logits) + ece_before.append( + expected_calibration_error( + probs=survival_at_one, + targets=position_target, + num_bins=num_bins, + ) + ) + + best_temperature = grid_values[0] + best_survival = survival_fitted * torch.sigmoid( + position_logits / best_temperature + ) + best_ece = expected_calibration_error( + probs=best_survival, targets=position_target, num_bins=num_bins + ) + for temperature in grid_values[1:]: + candidate_survival = survival_fitted * torch.sigmoid( + position_logits / temperature + ) + candidate_ece = expected_calibration_error( + probs=candidate_survival, + targets=position_target, + num_bins=num_bins, + ) + if candidate_ece < best_ece: + best_ece = candidate_ece + best_temperature = temperature + best_survival = candidate_survival + + temperatures.append(float(best_temperature)) + ece_after.append(float(best_ece)) + survival_fitted = best_survival + + return { + "temperatures": temperatures, + "ece_before": ece_before, + "ece_after": ece_after, + } + + +def load_collected_shards(*, data_glob: str) -> tuple[torch.Tensor, torch.Tensor]: + shard_paths = sorted(glob.glob(data_glob)) + if not shard_paths: + raise ValueError(f"No STS data shards matched {data_glob!r}.") + + logits_shards: list[torch.Tensor] = [] + prefix_mask_shards: list[torch.Tensor] = [] + shard_gamma: Optional[int] = None + for shard_path in shard_paths: + shard = torch.load(shard_path, map_location="cpu") + shard_logits = shard["logits"] + shard_prefix_mask = shard["prefix_mask"] + if shard_logits.shape != shard_prefix_mask.shape: + raise ValueError( + f"Shard {shard_path!r} logits / prefix_mask shape mismatch: " + f"{tuple(shard_logits.shape)} vs {tuple(shard_prefix_mask.shape)}." + ) + if shard_gamma is None: + shard_gamma = int(shard_logits.shape[1]) + elif int(shard_logits.shape[1]) != shard_gamma: + raise ValueError( + f"Shard {shard_path!r} gamma {int(shard_logits.shape[1])} disagrees " + f"with earlier shards' gamma {shard_gamma}." + ) + logits_shards.append(shard_logits) + prefix_mask_shards.append(shard_prefix_mask) + + return torch.cat(logits_shards, dim=0), torch.cat(prefix_mask_shards, dim=0) + + +def fit( + *, + data_glob: str, + out: Path, + num_bins: int = 15, + gamma: Optional[int] = None, +) -> None: + logits, prefix_mask = load_collected_shards(data_glob=data_glob) + resolved_gamma = int(logits.shape[1]) + if gamma is not None and gamma != resolved_gamma: + raise ValueError( + f"Collected shards have gamma={resolved_gamma} but --gamma={gamma}." + ) + num_samples = int(logits.shape[0]) + + result = fit_sts_temperatures( + logits=logits, + prefix_mask=prefix_mask, + grid=default_temperature_grid(), + num_bins=num_bins, + ) + calibration = DSparkStsCalibration( + temperatures=result["temperatures"], + dataset=data_glob, + num_samples=num_samples, + ece_before=result["ece_before"], + ece_after=result["ece_after"], + ) + out.write_text(calibration.to_json(), encoding="utf-8") + + print( + f"Fit STS temperatures over {num_samples} samples (gamma={resolved_gamma}) " + f"-> {out}" + ) + print("pos temperature ece_before ece_after") + for position in range(resolved_gamma): + print( + f"{position:>3} {result['temperatures'][position]:>11.4f} " + f"{result['ece_before'][position]:>10.4f} " + f"{result['ece_after'][position]:>9.4f}" + ) + + +def main() -> None: + logging.basicConfig(level=logging.INFO) + parser = argparse.ArgumentParser( + description="Fit DSpark Sequential Temperature Scaling (STS) calibration " + "temperatures from collected confidence shards." + ) + parser.add_argument( + "--data-glob", + required=True, + help="Glob of collected .pt shards, each a dict with [n, gamma] " + "'logits' and 'prefix_mask' tensors.", + ) + parser.add_argument( + "--out", + required=True, + type=Path, + help="Output STS calibration JSON path.", + ) + parser.add_argument( + "--num-bins", + type=int, + default=15, + help="Number of equal-width ECE bins.", + ) + parser.add_argument( + "--gamma", + type=int, + default=None, + help="Optional gamma override to validate the shards against.", + ) + args = parser.parse_args() + + fit( + data_glob=args.data_glob, + out=args.out, + num_bins=args.num_bins, + gamma=args.gamma, + ) + + +if __name__ == "__main__": + main() diff --git a/python/sglang/benchmark/serving.py b/python/sglang/benchmark/serving.py index 9bebe908c..cd80d277f 100644 --- a/python/sglang/benchmark/serving.py +++ b/python/sglang/benchmark/serving.py @@ -109,6 +109,9 @@ class RequestFuncOutput: cached_tokens: int = 0 cached_tokens_details: Optional[Dict[str, Any]] = None spec_accept_length: float = 0.0 + spec_cap_length: float = 0.0 + spec_block_accept_length: float = 0.0 + spec_cap_lens_histogram: List[int] = field(default_factory=list) @staticmethod def init_new(request_func_input: RequestFuncInput): @@ -483,6 +486,15 @@ async def async_request_openai_chat_completions( output.spec_accept_length = ( _meta_info.get("spec_accept_length", 0.0) or 0.0 ) + output.spec_cap_length = ( + _meta_info.get("spec_cap_length", 0.0) or 0.0 + ) + output.spec_block_accept_length = ( + _meta_info.get("spec_block_accept_length", 0.0) or 0.0 + ) + output.spec_cap_lens_histogram = ( + _meta_info.get("spec_cap_lens_histogram", []) or [] + ) if getattr(args, "cache_report", False): _extract_cache_from_sglext(response_json, output) else: diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 9188151db..0ede772fb 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py @@ -53,12 +53,14 @@ def apply_deepseek_v4_defaults(server_args: ServerArgs, model_arch: str) -> None ) if server_args.speculative_algorithm is not None: - assert ( - server_args.speculative_algorithm == "EAGLE" - ), f"Only EAGLE speculative algorithm is supported for {model_arch}" - assert ( - server_args.speculative_eagle_topk == 1 - ), f"Only EAGLE speculative algorithm with topk == 1 is supported for {model_arch}" + assert server_args.speculative_algorithm in ( + "EAGLE", + "DSPARK", + ), f"Only EAGLE and DSPARK speculative algorithms are supported for {model_arch}" + if server_args.speculative_algorithm == "EAGLE": + assert ( + server_args.speculative_eagle_topk == 1 + ), f"Only EAGLE speculative algorithm with topk == 1 is supported for {model_arch}" def validate_deepseek_v4_cp(server_args: ServerArgs) -> None: diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index a827fcdc5..cbda12929 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -265,6 +265,171 @@ def _handle_dflash(server_args: ServerArgs) -> None: ) +def _target_checkpoint_bundles_dspark_draft(server_args: ServerArgs) -> bool: + from sglang.srt.speculative.dspark_components.dspark_config import ( + checkpoint_bundles_dspark_draft, + ) + + return checkpoint_bundles_dspark_draft(server_args.get_model_config().hf_config) + + +def _handle_dspark(server_args: ServerArgs) -> None: + if not server_args.device.startswith("cuda"): + raise ValueError("DSpark speculative decoding only supports CUDA device.") + + if server_args.enable_dp_attention: + if not server_args.enable_dp_lm_head: + raise ValueError("DSpark with dp attention requires --enable-dp-lm-head.") + if server_args.moe_a2a_backend != "none": + raise ValueError( + "DSpark with dp attention only supports the built-in TP MoE " + f"(moe_a2a_backend='none'), got {server_args.moe_a2a_backend!r}." + ) + if server_args.attn_cp_size > 1: + raise ValueError( + "DSpark with dp attention does not support context parallel " + f"(attn_cp_size={server_args.attn_cp_size})." + ) + if ( + server_args.speculative_moe_a2a_backend is not None + and server_args.speculative_moe_a2a_backend != server_args.moe_a2a_backend + ): + raise ValueError( + "DSpark ignores --speculative-moe-a2a-backend; with dp attention it " + f"must match the target moe_a2a_backend={server_args.moe_a2a_backend!r} " + f"(got {server_args.speculative_moe_a2a_backend!r})." + ) + + if server_args.pp_size != 1: + raise ValueError( + "Currently DSpark speculative decoding only supports pp_size == 1." + ) + + if server_args.speculative_draft_model_path is None: + if _target_checkpoint_bundles_dspark_draft(server_args): + server_args.speculative_draft_model_path = server_args.model_path + server_args.speculative_draft_model_revision = server_args.revision + logger.info( + "DSpark draft weights are bundled in the target checkpoint; " + "defaulting --speculative-draft-model-path to --model-path (%s).", + server_args.model_path, + ) + else: + raise ValueError( + "DSpark dense speculative decoding requires setting " + "--speculative-draft-model-path." + ) + + if server_args.speculative_num_steps is None: + server_args.speculative_num_steps = 1 + elif int(server_args.speculative_num_steps) != 1: + logger.warning( + "DSpark only supports speculative_num_steps == 1; overriding speculative_num_steps=%s to 1.", + server_args.speculative_num_steps, + ) + server_args.speculative_num_steps = 1 + + if server_args.speculative_eagle_topk is None: + server_args.speculative_eagle_topk = 1 + elif int(server_args.speculative_eagle_topk) != 1: + logger.warning( + "DSpark only supports speculative_eagle_topk == 1; overriding speculative_eagle_topk=%s to 1.", + server_args.speculative_eagle_topk, + ) + server_args.speculative_eagle_topk = 1 + + gamma: Optional[int] = None + if server_args.speculative_dspark_block_size is not None: + if int(server_args.speculative_dspark_block_size) <= 0: + raise ValueError( + "DSpark requires --speculative-dspark-block-size to be positive, " + f"got {server_args.speculative_dspark_block_size}." + ) + gamma = int(server_args.speculative_dspark_block_size) + else: + from sglang.srt.speculative.dspark_components.dspark_config import ( + DEFAULT_DSPARK_GAMMA, + read_draft_checkpoint_gamma, + ) + + try: + gamma = read_draft_checkpoint_gamma(server_args=server_args) + except Exception as e: + logger.warning( + "Failed to read DSpark gamma from draft model config; " + "cannot cross-check --speculative-num-draft-tokens. Error: %s", + e, + ) + if gamma is None and server_args.speculative_num_draft_tokens is None: + gamma = DEFAULT_DSPARK_GAMMA + logger.warning( + "DSpark gamma is not set; defaulting to %d.", + gamma, + ) + + if gamma is not None: + verify_window = int(gamma) + 1 + if ( + server_args.speculative_num_draft_tokens is not None + and int(server_args.speculative_num_draft_tokens) != verify_window + ): + raise ValueError( + "DSpark speculative_num_draft_tokens must equal gamma + 1 " + f"(= {verify_window} for gamma={gamma}), but got " + f"speculative_num_draft_tokens={server_args.speculative_num_draft_tokens}." + ) + server_args.speculative_num_draft_tokens = verify_window + + if server_args.speculative_num_draft_tokens is None: + raise ValueError( + "DSpark could not resolve speculative_num_draft_tokens; set " + "--speculative-dspark-block-size (= gamma)." + ) + if int(server_args.speculative_num_draft_tokens) < 2: + raise ValueError( + "DSpark speculative_num_draft_tokens must be >= 2 (= gamma + 1), " + f"got {server_args.speculative_num_draft_tokens}." + ) + + if server_args.max_running_requests is None: + server_args.max_running_requests = 48 + logger.warning( + "Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests." + ) + + if server_args.enable_mixed_chunk: + server_args.enable_mixed_chunk = False + logger.warning( + "Mixed chunked prefill is disabled because of using dspark speculative decoding." + ) + + from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyMode, + read_ragged_verify_mode, + ) + + ragged_mode = read_ragged_verify_mode() + if ( + server_args.speculative_dspark_align_verify_tokens_to_graph_tier + and ragged_mode is not RaggedVerifyMode.COMPACT + ): + logger.warning( + "--speculative-dspark-align-verify-tokens-to-graph-tier only takes " + "effect with SGLANG_RAGGED_VERIFY_MODE=compact (got %r); it will be " + "a no-op.", + ragged_mode.value, + ) + if ( + server_args.speculative_dspark_sps_table_path + and ragged_mode is RaggedVerifyMode.STATIC + ): + logger.warning( + "--speculative-dspark-sps-table-path feeds the ragged-verify budget " + "scheduler, which is off under SGLANG_RAGGED_VERIFY_MODE=static; it " + "will be a no-op." + ) + + def _resolve_dflash_draft_attention_backend(server_args: ServerArgs) -> None: """Resolve `speculative_draft_attention_backend` to a final, supported value. diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 465553fbb..2cfbd736a 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -122,6 +122,7 @@ def is_deepseek_v4(config) -> bool: return _hf_arch(config) in ( "DeepseekV4ForCausalLM", "DeepseekV4ForCausalLMNextN", + "DeepseekV4ForCausalLMDSpark", ) @@ -247,12 +248,14 @@ class ModelConfig: language_only: bool = False, disable_hybrid_swa_memory: bool = False, model_config_parser: str = "auto", + speculative_algorithm: Optional[str] = None, ) -> None: # Parse args self.model_path = model_path self.revision = revision self.quantization = quantization self.is_draft_model = is_draft_model + self.speculative_algorithm = speculative_algorithm self.model_impl = model_impl self.sampling_defaults = sampling_defaults self.quantize_and_serve = quantize_and_serve @@ -528,6 +531,7 @@ class ModelConfig: is_draft_model=is_draft_model, disable_hybrid_swa_memory=server_args.disable_hybrid_swa_memory, model_config_parser=server_args.model_config_parser, + speculative_algorithm=server_args.speculative_algorithm, **kwargs, ) @@ -547,8 +551,23 @@ class ModelConfig: is_draft_model and self.hf_config.architectures[0] == "DeepseekV4ForCausalLM" ): - self.hf_config.architectures[0] = "DeepseekV4ForCausalLMNextN" - self.hf_config.num_nextn_predict_layers = 1 + from sglang.srt.speculative.dspark_components.dspark_config import ( + checkpoint_bundles_dspark_draft, + ) + + # A dspark-bundled checkpoint may also carry MTP layers; the + # selected algorithm decides which draft arch to load. + if checkpoint_bundles_dspark_draft(self.hf_config) and ( + self.speculative_algorithm in (None, "DSPARK") + ): + self.hf_config.architectures[0] = "DeepseekV4ForCausalLMDSpark" + logger.info( + "Draft checkpoint bundles a DSpark head; loading draft arch " + "DeepseekV4ForCausalLMDSpark." + ) + else: + self.hf_config.architectures[0] = "DeepseekV4ForCausalLMNextN" + self.hf_config.num_nextn_predict_layers = 1 if is_draft_model and self.hf_config.architectures[0] == "Glm4MoeForCausalLM": self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN" @@ -636,7 +655,12 @@ class ModelConfig: logger.info(f"Hybrid swa model: {self.hf_config.architectures=}") self.is_deepseek_v4_arch = any( - arch in ["DeepseekV4ForCausalLM", "DeepseekV4ForCausalLMNextN"] + arch + in [ + "DeepseekV4ForCausalLM", + "DeepseekV4ForCausalLMNextN", + "DeepseekV4ForCausalLMDSpark", + ] for arch in self.hf_config.architectures ) @@ -786,6 +810,7 @@ class ModelConfig: elif ( "DeepseekV4ForCausalLM" in self.hf_config.architectures or "DeepseekV4ForCausalLMNextN" in self.hf_config.architectures + or "DeepseekV4ForCausalLMDSpark" in self.hf_config.architectures ): self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim self.qk_nope_head_dim = self.hf_config.head_dim - self.qk_rope_head_dim @@ -1715,6 +1740,7 @@ multimodal_model_archs = [ piecewise_cuda_graph_disabled_model_archs = [ "DeepseekV4ForCausalLM", "DeepseekV4ForCausalLMNextN", + "DeepseekV4ForCausalLMDSpark", "Qwen3NextForCausalLM", "BailingMoeV2_5ForCausalLM", "LLaDAModelLM", @@ -1844,6 +1870,7 @@ def is_hybrid_swa_model( "Llama4ForConditionalGeneration", "DeepseekV4ForCausalLM", "DeepseekV4ForCausalLMNextN", + "DeepseekV4ForCausalLMDSpark", "GptOssForCausalLM", *MIMO_V2_MODEL_ARCHS, "MiMoV2MTP", diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index acd094a18..40a6c11ca 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -143,6 +143,10 @@ class DecodeReqToTokenPool: ) self.free_slots = list(range(1, self._alloc_size)) + # Slot-reuse generation counter; mirrors ReqToTokenPool. Required even + # here: HybridMambaDecodeReqToTokenPool borrows this __init__ while + # inheriting ReqToTokenPool.alloc, which bumps it. + self.req_generation = torch.zeros(self._alloc_size, dtype=torch.int64) def write(self, indices, values): self.req_to_token[indices] = values @@ -171,6 +175,7 @@ class DecodeReqToTokenPool: for r in reqs: if r.req_pool_idx is None: r.req_pool_idx = select_index[offset] + self.req_generation[r.req_pool_idx] += 1 offset += 1 return [r.req_pool_idx for r in reqs] @@ -181,6 +186,7 @@ class DecodeReqToTokenPool: def clear(self): self.free_slots = list(range(1, self._alloc_size)) + self.req_generation.zero_() class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d9fb89b6d..edf356ef1 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -215,6 +215,7 @@ class Envs: # Logging Options SGLANG_LOG_GC = EnvBool(False) SGLANG_LOG_FORWARD_ITERS = EnvBool(False) + SGLANG_LOG_DECODE_GRAPH_KEY = EnvBool(False) SGLANG_LOG_MS = EnvBool(False) SGLANG_LOG_REQUEST_EXCEEDED_MS = EnvInt(-1) SGLANG_LOG_REQUEST_HEADERS = EnvTuple(tuple()) @@ -262,6 +263,20 @@ class Envs: SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE = EnvBool(False) SGLANG_FORCE_SHUTDOWN = EnvBool(False) SGLANG_DEBUG_MEMORY_POOL = EnvBool(False) + SGLANG_DSPARK_DEBUG_CONFIDENCE_PREFIX_SCHEDULER = EnvBool(False) + SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS = EnvBool(False) + SGLANG_DSPARK_DEBUG_DUMP = EnvTuple(tuple()) + SGLANG_DSPARK_LOG_SPS_PRED_INTERVAL = EnvInt(0) + SGLANG_DSPARK_STS_COLLECT_PATH = EnvStr("") + SGLANG_DSPARK_BLOCK_ACCEPT_ESTIMATE_PATH = EnvStr("") + SGLANG_DSPARK_BLOCK_ACCEPT_ONLINE_INTERVAL = EnvInt(0) + SGLANG_DSPARK_ENABLE_SPS_RECORD = EnvBool(False) + SGLANG_DSPARK_FAST_KERNEL = EnvBool(True) + SGLANG_DSPARK_FP32_LM_HEAD = EnvBool(False) + SGLANG_DSPARK_FAST_SAMPLING = EnvBool(True) + SGLANG_DSPARK_OPT_MARKOV_W2_BF16 = EnvBool(True) + SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD = EnvBool(True) + SGLANG_DSPARK_ENABLE_MULTI_STREAM = EnvBool(True) SGLANG_DEBUG_REVERT_PR = EnvInt(0) SGLANG_PHASE_CHECKER_DEBUG = EnvBool(False) SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False) @@ -705,6 +720,9 @@ class Envs: # Spec Config SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True) + SGLANG_RAGGED_VERIFY_MODE = EnvStr("static") + SGLANG_DSPARK_CONFIDENCE_RELAY_LAG_STEPS = EnvInt(2) + SGLANG_TEST_RAGGED_VERIFY_FORCE_UNIFORM_CAPTURE = EnvBool(False) # Skip draft_extend while adaptive spec is at steps=0 (drafting disabled). # Saves the per-step draft forward, but the draft KV goes stale: an upshift # back to steps>0 starts from a cold draft state (low accept until it recovers). diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index 901200f79..d60b32c6d 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -42,6 +42,8 @@ class AttentionBackend(ABC): prefill_attention_backend_str: Optional[str] = None decode_attention_backend_str: Optional[str] = None + supports_ragged_verify_graph: bool = False + def init_forward_metadata(self, forward_batch: ForwardBatch): """Eager entry point. Default = ``_out_graph(fb) + _in_graph(fb)``. diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 3d09014f3..02956eac0 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -21,6 +21,11 @@ import torch.nn.functional as F from sglang.jit_kernel.dsv4.online_c128_mtp import OnlineC128MTPController from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.dsv4.attn_metadata_kernels import ( + BuildCausalSwaPageIndices, + BuildPageTablePositions, + ExpandPrefillCausally, +) from sglang.srt.layers.attention.dsv4.compressor_v2 import ( CompressorBackendMixin, FusedCompressMetadata, @@ -49,7 +54,20 @@ from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import ( from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_parallel +from sglang.srt.speculative.dspark_components.kernels.dspark_attn_metadata import ( + BuildBlockSeqLensCausal, + BuildDsparkSwaPageIndices, + ComputeDsparkWindowGather, +) from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc +from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyMode, + compute_ragged_extend_lengths, + compute_target_verify_graph_key, + compute_uniform_extend_lengths, + read_ragged_verify_mode, + resolve_ragged_verify_layout, +) from sglang.srt.utils import ceil_align, is_xpu from sglang.srt.utils.common import is_sm120_supported @@ -58,6 +76,7 @@ if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout _is_sm120 = is_sm120_supported() _is_xpu = is_xpu() @@ -412,6 +431,10 @@ class DSV4RawVerifyMetadata: seq_lens_cpu: Optional[List[int]] = None c128_compress_metadata: Optional[FusedCompressMetadata] = None + extend_start_loc: Optional[torch.Tensor] = None + verify_lens: Optional[torch.Tensor] = None + total_verify_tokens: int = 0 + def copy_(self, other: DSV4RawVerifyMetadata): self.req_pool_indices.copy_(other.req_pool_indices) self.seq_lens.copy_(other.seq_lens) @@ -423,6 +446,10 @@ class DSV4RawVerifyMetadata: self.c128_compress_metadata, other.c128_compress_metadata ) + self.extend_start_loc = other.extend_start_loc + self.verify_lens = other.verify_lens + self.total_verify_tokens = other.total_verify_tokens + @dataclass class DSV4RawDecodeMetadata: @@ -456,7 +483,7 @@ class DeepseekV4AttnBackend( AttentionBackend, C4IndexerBackendMixin, CompressorBackendMixin ): use_captured_forward_metadata_for_breakable_cuda_graph: bool = True - + supports_ragged_verify_graph: bool = True needs_cpu_seq_lens: bool = False def __init__( @@ -504,6 +531,21 @@ class DeepseekV4AttnBackend( self.speculative_num_draft_tokens: int = ( model_runner.server_args.speculative_num_draft_tokens ) + if self.speculative_num_draft_tokens is not None: + # Persistent target-verify metadata buffers. Allocated here (not + # lazily) so they are ordinary tensors: the first touch of a lazy + # buffer would inherit the caller's context, and a creation inside + # an inference_mode forward would forbid the in-place updates the + # graph-capture path performs outside inference mode. + num_reqs = self.req_to_token.shape[0] + self.extend_seq_lens_buffer = torch.full( + (num_reqs,), + self.speculative_num_draft_tokens, + **self.cuda_int32_kwargs, + ) + self.extend_start_loc_buffer = torch.zeros( + num_reqs, **self.cuda_int32_kwargs + ) self.speculative_step_id = speculative_step_id self.forward_metadata: Union[ DSV4Metadata, @@ -514,14 +556,56 @@ class DeepseekV4AttnBackend( # Draft-extend and online-c128 verify metadata are host-planned, so # spec runs keep the relay publish (the mirror only exists under # spec-v2; without spec the flag has no consumer either way). - if model_runner.server_args.speculative_algorithm is not None: + # DSPARK is the exception: its draft path carries its own host lens + # (reserved_seq_lens_cpu) and its verify prep is device-side. + spec_alg = model_runner.spec_algorithm + if not spec_alg.is_none() and not spec_alg.is_dspark(): self.needs_cpu_seq_lens = True self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device) + self.is_dspark_draft = model_runner.is_draft_worker and spec_alg.is_dspark() + def _move_to_device(self, x: List[int]) -> torch.Tensor: pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True) return pin_tensor.to(self.device, non_blocking=True) + def _resolve_verify_layout( + self, + forward_batch: ForwardBatch, + bs: int, + ) -> Optional[RaggedVerifyLayout]: + layout = resolve_ragged_verify_layout(forward_batch) + if layout is None: + return None + if read_ragged_verify_mode() is not RaggedVerifyMode.COMPACT: + return None + if get_parallel().attn_cp_size > 1: + raise NotImplementedError( + "DSV4 ragged verify does not support context parallel (CP); " + "set SGLANG_RAGGED_VERIFY_MODE off for CP runs." + ) + if self.online_c128_mtp.enabled(): + raise NotImplementedError( + "DSV4 ragged verify does not support online c128 MTP; " + "set SGLANG_RAGGED_VERIFY_MODE off or disable online compress." + ) + # Layout invariants (verify_lens >= 1, total == sum) are enforced in + # RaggedVerifyLayout.__post_init__; don't re-check the device tensor + # here -- that would D2H-sync the host-free verify prep path. + layout = layout.padded_to_bucket(padded_bs=bs) + return layout + + def _target_verify_graph_key( + self, + bs: int, + ragged_layout: Optional[RaggedVerifyLayout], + ) -> Tuple[int, int]: + return compute_target_verify_graph_key( + bs=bs, + num_draft_tokens=self.speculative_num_draft_tokens, + ragged_layout=ragged_layout, + ) + def _make_target_verify_c128_metadata( self, req_pool_indices: torch.Tensor, @@ -623,6 +707,7 @@ class DeepseekV4AttnBackend( need_compress: bool = True, use_prefill_cuda_graph: bool = False, online_c128_state_slot_offset: int = 0, + dspark_block_size: Optional[int] = None, ) -> DSV4Metadata: seq_lens_casual, req_pool_indices_repeated = self.expand_prefill_casually( num_tokens=num_tokens, @@ -642,6 +727,7 @@ class DeepseekV4AttnBackend( out_loc=out_cache_loc, need_compress=need_compress, is_prefill=True, + dspark_block_size=dspark_block_size, ) indexer_metadata = ( self.init_forward_metadata_indexer( @@ -709,17 +795,29 @@ class DeepseekV4AttnBackend( out_cache_loc: Optional[torch.Tensor] = None, use_prefill_cuda_graph: bool = False, online_c128_state_slot_offset: int = 0, + ragged_layout: Optional[RaggedVerifyLayout] = None, ) -> Union[DSV4Metadata, DSV4RawVerifyMetadata]: if envs.SGLANG_PREP_IN_CUDA_GRAPH.get(): assert out_cache_loc is not None + bs = len(seq_lens) seq_lens_cpu_list = ( seq_lens_cpu.tolist() if seq_lens_cpu is not None else None ) - if not hasattr(self, "extend_seq_lens_buffer"): - self.extend_seq_lens_buffer = torch.tensor( - [self.speculative_num_draft_tokens] * 1025, device=self.device + if ragged_layout is None: + self.extend_seq_lens_buffer[:bs].fill_( + self.speculative_num_draft_tokens ) - extend_seq_lens = self.extend_seq_lens_buffer[: len(seq_lens)] + extend_seq_lens = self.extend_seq_lens_buffer[:bs] + extend_start_loc = None + verify_lens = None + total_verify_tokens = self.speculative_num_draft_tokens * bs + else: + self.extend_seq_lens_buffer[:bs].copy_(ragged_layout.verify_lens) + self.extend_start_loc_buffer[:bs].copy_(ragged_layout.extend_start_loc) + extend_seq_lens = self.extend_seq_lens_buffer[:bs] + extend_start_loc = self.extend_start_loc_buffer[:bs] + verify_lens = self.extend_seq_lens_buffer[:bs] + total_verify_tokens = ragged_layout.graph_num_tokens return DSV4RawVerifyMetadata( req_pool_indices=req_pool_indices, @@ -735,6 +833,9 @@ class DeepseekV4AttnBackend( use_prefill_cuda_graph, online_c128_state_slot_offset, ), + extend_start_loc=extend_start_loc, + verify_lens=verify_lens, + total_verify_tokens=total_verify_tokens, ) else: seq_lens_cpu = seq_lens.tolist() @@ -746,6 +847,7 @@ class DeepseekV4AttnBackend( out_cache_loc=out_cache_loc, use_prefill_cuda_graph=use_prefill_cuda_graph, online_c128_state_slot_offset=online_c128_state_slot_offset, + ragged_layout=ragged_layout, ) def init_forward_metadata_target_verify_old( @@ -757,13 +859,27 @@ class DeepseekV4AttnBackend( out_cache_loc: Optional[torch.Tensor] = None, use_prefill_cuda_graph: bool = False, online_c128_state_slot_offset: int = 0, + ragged_layout: Optional[RaggedVerifyLayout] = None, ) -> DSV4Metadata: - batch_size = len(seq_lens) - seq_lens = seq_lens + self.speculative_num_draft_tokens - seq_lens_cpu = [x + self.speculative_num_draft_tokens for x in seq_lens_cpu] - extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * batch_size - extend_seq_lens = self._move_to_device(extend_seq_lens_cpu) - num_tokens = self.speculative_num_draft_tokens * batch_size + if ragged_layout is None: + lengths = compute_uniform_extend_lengths( + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + extend_len=self.speculative_num_draft_tokens, + ) + extend_seq_lens = self._move_to_device(lengths.extend_seq_lens_cpu) + else: + lengths = compute_ragged_extend_lengths( + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + ragged_layout=ragged_layout, + ) + extend_seq_lens = ragged_layout.verify_lens + seq_lens = lengths.seq_lens_extended + seq_lens_cpu = lengths.seq_lens_cpu_extended + extend_seq_lens_cpu = lengths.extend_seq_lens_cpu + num_tokens = lengths.num_tokens + extend_start_loc = lengths.extend_start_loc if out_cache_loc is None: out_cache_loc = seq_lens.new_zeros(num_tokens) return self.init_forward_metadata_prefill( @@ -775,12 +891,46 @@ class DeepseekV4AttnBackend( num_tokens=num_tokens, extend_seq_lens=extend_seq_lens, extend_seq_lens_cpu=extend_seq_lens_cpu, - extend_start_loc=None, + extend_start_loc=extend_start_loc, need_compress=True, use_prefill_cuda_graph=use_prefill_cuda_graph, online_c128_state_slot_offset=online_c128_state_slot_offset, ) + def init_forward_metadata_dspark_draft_block( + self, + max_seq_len: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + block_size: int, + ) -> DSV4Metadata: + if seq_lens_cpu is None: + seq_lens_cpu_list = seq_lens.tolist() + else: + seq_lens_cpu_list = [int(x) for x in seq_lens_cpu.tolist()] + lengths = compute_uniform_extend_lengths( + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu_list, + extend_len=block_size, + ) + extend_seq_lens = self._move_to_device(lengths.extend_seq_lens_cpu) + return self.init_forward_metadata_prefill( + max_seq_len=max_seq_len, + req_pool_indices=req_pool_indices, + seq_lens=lengths.seq_lens_extended, + seq_lens_cpu=lengths.seq_lens_cpu_extended, + out_cache_loc=out_cache_loc, + num_tokens=lengths.num_tokens, + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=lengths.extend_seq_lens_cpu, + extend_start_loc=lengths.extend_start_loc, + need_compress=False, + use_prefill_cuda_graph=False, + dspark_block_size=block_size, + ) + def make_forward_metadata_from_raw_verify( self, raw_metadata: DSV4RawVerifyMetadata, @@ -791,15 +941,35 @@ class DeepseekV4AttnBackend( out_cache_loc = raw_metadata.out_cache_loc bs, num_draft_tokens = len(seq_lens), self.speculative_num_draft_tokens - seq_lens = seq_lens + self.speculative_num_draft_tokens extend_seq_lens = raw_metadata.extend_seq_lens assert extend_seq_lens is not None - seq_lens_casual, req_pool_indices_repeated = ( - self.expand_extend_with_same_length( - bs, num_draft_tokens, seq_lens, req_pool_indices + is_ragged = raw_metadata.verify_lens is not None + if is_ragged: + seq_lens = seq_lens + extend_seq_lens + num_q_tokens = raw_metadata.total_verify_tokens + assert num_q_tokens > 0, "ragged verify raw metadata is stale/empty" + seq_lens_casual, req_pool_indices_repeated = ( + self._expand_prefill_casually_vectorized( + num_tokens=num_q_tokens, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + extend_start_loc=raw_metadata.extend_start_loc, + req_pool_indices=req_pool_indices, + padded_num_tokens=out_cache_loc.shape[0], + ) + ) + else: + seq_lens = seq_lens + self.speculative_num_draft_tokens + num_q_tokens = num_draft_tokens * bs + seq_lens_casual, req_pool_indices_repeated = ( + self.expand_extend_with_same_length( + bs=bs, + qo_len=num_draft_tokens, + seq_lens=seq_lens, + req_pool_indices=req_pool_indices, + ) ) - ) core_attn_metadata = self.make_core_attn_metadata( req_to_token=self.req_to_token, req_pool_indices_repeated=req_pool_indices_repeated, @@ -820,7 +990,7 @@ class DeepseekV4AttnBackend( seq_lens_cpu=None, extend_lens_cpu=None, use_prefill_cuda_graph=True, - num_q_tokens=num_draft_tokens * bs, + num_q_tokens=num_q_tokens, online_state_slot_offset=online_c128_state_slot_offset, ) c128_compress_metadata = raw_metadata.c128_compress_metadata @@ -941,6 +1111,35 @@ class DeepseekV4AttnBackend( ) ) + if self.is_dspark_draft and forward_batch.forward_mode.is_target_verify(): + block_size = int(forward_batch.spec_info.draft_token_num) + seq_lens_casual = self._dspark_seq_lens_casual( + seq_lens=forward_batch.seq_lens, block_size=block_size + ) + req_pool_indices_repeated = ( + forward_batch.req_pool_indices.repeat_interleave(block_size) + ) + ( + swa_page_indices, + swa_topk_lengths, + ) = self.get_dspark_swa_page_indices( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + out_loc=out_cache_loc, + block_size=block_size, + ) + metadata.core_attn_metadata.swa_page_indices = swa_page_indices + metadata.core_attn_metadata.swa_topk_lengths = swa_topk_lengths + + def _dspark_seq_lens_casual( + self, *, seq_lens: torch.Tensor, block_size: int + ) -> torch.Tensor: + return BuildBlockSeqLensCausal.execute( + seq_lens=seq_lens, + block_size=block_size, + device=self.cuda_int32_kwargs["device"], + ) + def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, @@ -997,6 +1196,7 @@ class DeepseekV4AttnBackend( actual_max_seq_len = seq_lens_cpu.max().item() assert actual_max_seq_len <= chosen_max_seq_len + graph_key = bs if bucket == _GraphBucket.DECODE_OR_IDLE: assert out_cache_loc is not None assert len(out_cache_loc.shape) == 1, f"{out_cache_loc.shape=}" @@ -1017,16 +1217,48 @@ class DeepseekV4AttnBackend( seq_lens=seq_lens, out_cache_loc=out_cache_loc_padded, ) + elif bucket == _GraphBucket.TARGET_VERIFY and self.is_dspark_draft: + block_size = self.speculative_num_draft_tokens - 1 + num_tokens_block = block_size * bs + assert out_cache_loc is not None + out_cache_loc_padded = torch.nn.functional.pad( + out_cache_loc, + pad=(0, num_tokens_block - len(out_cache_loc)), + mode="constant", + value=0, + ) + self.online_c128_mtp.prepare_forward( + actual_forward_mode, + req_pool_indices, + seq_lens, + ) + temp_metadata = self.init_forward_metadata_dspark_draft_block( + max_seq_len=chosen_max_seq_len, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + out_cache_loc=out_cache_loc_padded, + block_size=block_size, + ) elif bucket == _GraphBucket.TARGET_VERIFY: verify_bs = _get_target_verify_bs(forward_batch) + ragged_layout = self._resolve_verify_layout(forward_batch, bs=bs) + graph_key, num_tokens_v = self._target_verify_graph_key( + bs=bs, ragged_layout=ragged_layout + ) if self.online_c128_mtp.enabled() and verify_bs == 0: self.online_c128_mtp.clear() self.forward_metadata = self.cuda_graph_metadata_of_bucket_and_bs[ bucket - ][bs] + ][graph_key] return assert out_cache_loc is not None - num_tokens_v = self.speculative_num_draft_tokens * bs + assert num_tokens_v >= len(out_cache_loc), ( + f"ragged verify token-keyed graph requires the decode cuda-graph " + f"runner to supply out_cache_loc sized to graph_num_tokens " + f"({num_tokens_v}), got {len(out_cache_loc)}; the decode graph " + "runner does not yet route token-keyed ragged captures." + ) out_cache_loc_padded = torch.nn.functional.pad( out_cache_loc, pad=(0, num_tokens_v - len(out_cache_loc)), @@ -1047,6 +1279,7 @@ class DeepseekV4AttnBackend( out_cache_loc=out_cache_loc_padded, use_prefill_cuda_graph=True, online_c128_state_slot_offset=online_c128_state_slot_offset, + ragged_layout=ragged_layout, ) elif bucket == _GraphBucket.DRAFT_EXTEND: self.online_c128_mtp.prepare_forward( @@ -1081,7 +1314,7 @@ class DeepseekV4AttnBackend( raise NotImplementedError self.replay_cuda_graph_metadata_from( - bs=bs, temp_metadata=temp_metadata, bucket=bucket + bs=graph_key, temp_metadata=temp_metadata, bucket=bucket ) if in_capture: @@ -1152,7 +1385,18 @@ class DeepseekV4AttnBackend( seq_lens=seq_lens, out_cache_loc=out_cache_loc, ) + elif self.is_dspark_draft and logical_forward_mode.is_target_verify(): + block_size = int(forward_batch.spec_info.draft_token_num) + metadata = self.init_forward_metadata_dspark_draft_block( + max_seq_len=max_seq_len, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + out_cache_loc=forward_batch.out_cache_loc, + block_size=block_size, + ) elif logical_forward_mode.is_target_verify(): + ragged_layout = self._resolve_verify_layout(forward_batch, bs=len(seq_lens)) metadata = self.init_forward_metadata_target_verify( max_seq_len=max_seq_len, req_pool_indices=req_pool_indices, @@ -1160,6 +1404,7 @@ class DeepseekV4AttnBackend( seq_lens_cpu=seq_lens_cpu, out_cache_loc=forward_batch.out_cache_loc, online_c128_state_slot_offset=online_c128_state_slot_offset, + ragged_layout=ragged_layout, ) elif logical_forward_mode.is_prefill(include_draft_extend_v2=True): extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu @@ -1588,46 +1833,18 @@ class DeepseekV4AttnBackend( extend_seq_lens_tensor: Optional[torch.Tensor] = None, extend_start_loc: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - if ( - seq_lens_tensor is not None - and extend_seq_lens_tensor is not None - and extend_start_loc is not None - ): - return self._expand_prefill_casually_vectorized( - num_tokens=num_tokens, - seq_lens=seq_lens_tensor, - extend_seq_lens=extend_seq_lens_tensor, - extend_start_loc=extend_start_loc, - req_pool_indices=req_pool_indices, - padded_num_tokens=padded_num_tokens, - ) - - seq_lens_casual = torch.empty(num_tokens, **self.cuda_int32_kwargs) - idx_to_req_repeated = torch.empty(num_tokens, **self.cuda_int32_kwargs) - offset = 0 - for i, (kv_len, qo_len) in enumerate(zip(seq_lens, extend_seq_lens)): - out = seq_lens_casual[offset : offset + qo_len] - offset += qo_len - torch.arange(kv_len - qo_len + 1, kv_len + 1, out=out) - idx_to_req_repeated[offset - qo_len : offset].fill_(i) - - assert offset == num_tokens - req_pool_indices_repeated = req_pool_indices[idx_to_req_repeated] - - if padded_num_tokens is not None and padded_num_tokens > num_tokens: - pad_size = padded_num_tokens - num_tokens - seq_lens_casual = torch.nn.functional.pad( - seq_lens_casual, - (0, pad_size), - value=1, - ) - req_pool_indices_repeated = torch.nn.functional.pad( - req_pool_indices_repeated, - (0, pad_size), - value=req_pool_indices_repeated[-1].item(), - ) - - return seq_lens_casual, req_pool_indices_repeated + assert seq_lens_tensor is not None and extend_seq_lens_tensor is not None + result = ExpandPrefillCausally.execute( + req_pool_indices=req_pool_indices, + seq_lens=seq_lens_tensor, + extend_seq_lens=extend_seq_lens_tensor, + extend_start_loc=extend_start_loc, + seq_lens_cpu=seq_lens, + extend_seq_lens_cpu=extend_seq_lens, + num_tokens=num_tokens, + padded_num_tokens=padded_num_tokens, + ) + return result.seq_lens_casual, result.req_pool_indices_repeated def _expand_prefill_casually_vectorized( self, @@ -1638,41 +1855,21 @@ class DeepseekV4AttnBackend( req_pool_indices: torch.Tensor, padded_num_tokens: Optional[int], ) -> Tuple[torch.Tensor, torch.Tensor]: - repeats = extend_seq_lens.to(torch.int64) - req_pool_indices_repeated = torch.repeat_interleave( - req_pool_indices, repeats, output_size=num_tokens + result = ExpandPrefillCausally.execute( + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + extend_start_loc=extend_start_loc, + seq_lens_cpu=None, + extend_seq_lens_cpu=None, + num_tokens=num_tokens, + padded_num_tokens=padded_num_tokens, ) - - start_positions = seq_lens.to(torch.int32) - extend_seq_lens.to(torch.int32) + 1 - start_positions_repeated = torch.repeat_interleave( - start_positions, repeats, output_size=num_tokens - ) - start_locs_repeated = torch.repeat_interleave( - extend_start_loc.to(torch.int32), repeats, output_size=num_tokens - ) - token_offsets = ( - torch.arange(num_tokens, **self.cuda_int32_kwargs) - start_locs_repeated - ) - seq_lens_casual = start_positions_repeated + token_offsets - - if padded_num_tokens is not None and padded_num_tokens > num_tokens: - pad_size = padded_num_tokens - num_tokens - seq_lens_casual = torch.nn.functional.pad( - seq_lens_casual, - (0, pad_size), - value=1, - ) - req_pool_indices_repeated = torch.cat( - ( - req_pool_indices_repeated, - req_pool_indices_repeated[-1:].expand(pad_size), - ) - ) - - return seq_lens_casual, req_pool_indices_repeated + return result.seq_lens_casual, result.req_pool_indices_repeated def expand_extend_with_same_length( self, + *, bs: int, qo_len: int, seq_lens: torch.Tensor, @@ -1697,27 +1894,49 @@ class DeepseekV4AttnBackend( out_loc: torch.Tensor, need_compress: bool = True, is_prefill: bool = False, + dspark_block_size: Optional[int] = None, ) -> DSV4AttnMetadata: assert self.swa_page_size == SWA_WINDOW - seq_lens_casual = seq_lens_casual.to(torch.int32) - - swa_page_indices = self.get_swa_page_indices( - seq_lens_casual=seq_lens_casual, + prep = BuildPageTablePositions.execute( + req_to_token=req_to_token, req_pool_indices_repeated=req_pool_indices_repeated, + seq_lens_casual=seq_lens_casual, + max_seq_len=max_seq_len, + page_size=self.page_size, + swa_window=SWA_WINDOW, ) + seq_lens_casual = prep.seq_lens_casual - swa_page_indices = _pad_last_dim( - swa_page_indices, multiples_of=PAGE_INDEX_ALIGNED_SIZE - ) + raw_positions = prep.positions_casual + if dspark_block_size is not None: + assert ( + self.is_dspark_draft + and dspark_block_size == self.speculative_num_draft_tokens - 1 + ), ( + f"dspark_block_size={dspark_block_size} must equal gamma = " + f"speculative_num_draft_tokens-1={self.speculative_num_draft_tokens - 1} " + f"and is only valid on the DSpark draft backend " + f"(is_dspark_draft={self.is_dspark_draft})." + ) + swa_page_indices, swa_topk_lengths = self.get_dspark_swa_page_indices( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + out_loc=out_loc, + block_size=dspark_block_size, + ) + else: + swa_page_indices = BuildCausalSwaPageIndices.execute( + req_to_token=self.req_to_token, + full_to_swa_mapping=self.token_to_kv_pool.full_to_swa_index_mapping, + req_pool_indices_repeated=req_pool_indices_repeated, + seq_lens_casual=seq_lens_casual, + swa_window=SWA_WINDOW, + page_index_aligned_size=PAGE_INDEX_ALIGNED_SIZE, + ) + swa_topk_lengths = prep.swa_topk_lengths - raw_positions = seq_lens_casual - 1 - swa_topk_lengths = torch.clamp(seq_lens_casual, max=SWA_WINDOW) - - page_table = req_to_token[ - req_pool_indices_repeated, : max_seq_len : self.page_size - ] - page_table = (page_table // self.page_size).to(torch.int32) + page_table = prep.page_table core_attn_metadata = DSV4AttnMetadata( page_size=self.page_size, @@ -1743,24 +1962,34 @@ class DeepseekV4AttnBackend( core_attn_metadata.c128_flashmla_metadata = None return core_attn_metadata - def get_swa_page_indices( + def get_dspark_swa_page_indices( self, + *, seq_lens_casual: torch.Tensor, req_pool_indices_repeated: torch.Tensor, - ) -> torch.Tensor: - pos_causal = seq_lens_casual - 1 - num_qo_tokens = seq_lens_casual.size(0) - offsets = pos_causal.unsqueeze(1) - torch.arange( - SWA_WINDOW, **self.cuda_int32_kwargs - ).unsqueeze(0) - invalid_offset_mask = offsets < 0 - offsets.masked_fill_(invalid_offset_mask, 0) - raw_indices = self.req_to_token[req_pool_indices_repeated[:, None], offsets] - assert raw_indices.shape == (num_qo_tokens, SWA_WINDOW) - raw_indices.masked_fill_(invalid_offset_mask, -1) - swa_indices = self.token_to_kv_pool.translate_loc_from_full_to_swa(raw_indices) - # flash_mla attention requires int32 page indices. - return swa_indices.to(torch.int32) + out_loc: torch.Tensor, + block_size: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + gather = ComputeDsparkWindowGather.execute( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + block_size=block_size, + swa_window=SWA_WINDOW, + ) + + swa_page_indices, swa_topk_lengths = BuildDsparkSwaPageIndices.execute( + req_to_token=self.req_to_token, + full_to_swa_mapping=self.token_to_kv_pool.full_to_swa_index_mapping, + req_pool_indices_per_request=gather.req_pool_indices_per_request, + offsets=gather.offsets, + invalid=gather.invalid, + out_loc=out_loc[: gather.num_q], + context_lens=gather.context_lens, + block_size=block_size, + swa_window=SWA_WINDOW, + page_index_aligned_size=PAGE_INDEX_ALIGNED_SIZE, + ) + return swa_page_indices, swa_topk_lengths class DeepseekV4MultiStepBackend(DeepseekV4AttnBackend): diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index ec225dd17..1755e8212 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -41,6 +41,7 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_parallel from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc +from sglang.srt.speculative.ragged_verify import resolve_ragged_verify_layout from sglang.srt.utils import ceil_align if TYPE_CHECKING: @@ -863,6 +864,12 @@ class DeepseekV4HipRadixBackend( out_cache_loc=out_cache_loc_padded, ) elif bucket == _GraphBucket.TARGET_VERIFY: + if resolve_ragged_verify_layout(forward_batch) is not None: + raise NotImplementedError( + "DSV4 ragged verify is not supported on the HIP backend " + "(DeepseekV4HipRadixBackend) cuda-graph path; disable " + "SGLANG_RAGGED_VERIFY_MODE or use a CUDA device." + ) assert out_cache_loc is not None num_tokens_v = self.speculative_num_draft_tokens * bs out_cache_loc_padded = torch.nn.functional.pad( @@ -950,6 +957,12 @@ class DeepseekV4HipRadixBackend( out_cache_loc=out_cache_loc, ) elif forward_batch.forward_mode.is_target_verify(): + if resolve_ragged_verify_layout(forward_batch) is not None: + raise NotImplementedError( + "DSV4 ragged verify is not supported on the HIP backend " + "(DeepseekV4HipRadixBackend); disable SGLANG_RAGGED_VERIFY_MODE " + "or use a CUDA device." + ) metadata = self.init_forward_metadata_target_verify( max_seq_len=max_seq_len, req_pool_indices=req_pool_indices, diff --git a/python/sglang/srt/layers/attention/dsv4/attn_metadata_kernels.py b/python/sglang/srt/layers/attention/dsv4/attn_metadata_kernels.py new file mode 100644 index 000000000..bbb28952c --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/attn_metadata_kernels.py @@ -0,0 +1,526 @@ +from __future__ import annotations + +from typing import Optional + +import msgspec +import torch +import triton +import triton.language as tl + + +def _inputs_on_cuda(*args, **kwargs) -> bool: + """Route kernel dispatch by input placement: the first tensor argument + decides. CUDA inputs take the fused triton kernel; CPU inputs take the + torch reference implementation (triton is CUDA-only, and CPU-side callers + such as unit tests exercise the reference path).""" + for value in (*args, *kwargs.values()): + if isinstance(value, torch.Tensor): + return value.is_cuda + raise AssertionError("kernel dispatch requires at least one tensor argument") + + +class ExpandPrefillCausallyResult(msgspec.Struct): + seq_lens_casual: torch.Tensor + req_pool_indices_repeated: torch.Tensor + + +class ExpandPrefillCausally: + @classmethod + def execute(cls, *args, **kwargs) -> ExpandPrefillCausallyResult: + if _inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_start_loc: Optional[torch.Tensor], + seq_lens_cpu: Optional[list[int]], + extend_seq_lens_cpu: Optional[list[int]], + num_tokens: int, + padded_num_tokens: Optional[int], + ) -> ExpandPrefillCausallyResult: + return expand_prefill_causally( + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + extend_start_loc=extend_start_loc, + seq_lens_cpu=seq_lens_cpu, + extend_seq_lens_cpu=extend_seq_lens_cpu, + num_tokens=num_tokens, + padded_num_tokens=padded_num_tokens, + ) + + @classmethod + def triton( + cls, + *, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_start_loc: Optional[torch.Tensor], + seq_lens_cpu: Optional[list[int]], + extend_seq_lens_cpu: Optional[list[int]], + num_tokens: int, + padded_num_tokens: Optional[int], + ) -> ExpandPrefillCausallyResult: + return expand_prefill_causally_triton( + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + num_tokens=num_tokens, + padded_num_tokens=padded_num_tokens, + ) + + +def expand_prefill_causally( + *, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_start_loc: Optional[torch.Tensor], + seq_lens_cpu: Optional[list[int]], + extend_seq_lens_cpu: Optional[list[int]], + num_tokens: int, + padded_num_tokens: Optional[int], +) -> ExpandPrefillCausallyResult: + device = req_pool_indices.device + cuda_int32_kwargs = {"dtype": torch.int32, "device": device} + + if extend_start_loc is not None: + repeats = extend_seq_lens.to(torch.int64) + req_pool_indices_repeated = torch.repeat_interleave( + req_pool_indices, repeats, output_size=num_tokens + ) + start_positions = seq_lens.to(torch.int32) - extend_seq_lens.to(torch.int32) + 1 + start_positions_repeated = torch.repeat_interleave( + start_positions, repeats, output_size=num_tokens + ) + start_locs_repeated = torch.repeat_interleave( + extend_start_loc.to(torch.int32), repeats, output_size=num_tokens + ) + token_offsets = ( + torch.arange(num_tokens, **cuda_int32_kwargs) - start_locs_repeated + ) + seq_lens_casual = start_positions_repeated + token_offsets + + if padded_num_tokens is not None and padded_num_tokens > num_tokens: + pad_size = padded_num_tokens - num_tokens + seq_lens_casual = torch.nn.functional.pad( + seq_lens_casual, (0, pad_size), value=1 + ) + req_pool_indices_repeated = torch.cat( + ( + req_pool_indices_repeated, + req_pool_indices_repeated[-1:].expand(pad_size), + ) + ) + return ExpandPrefillCausallyResult( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + ) + + assert seq_lens_cpu is not None and extend_seq_lens_cpu is not None + seq_lens_casual = torch.empty(num_tokens, **cuda_int32_kwargs) + idx_to_req_repeated = torch.empty(num_tokens, **cuda_int32_kwargs) + offset = 0 + for i, (kv_len, qo_len) in enumerate(zip(seq_lens_cpu, extend_seq_lens_cpu)): + out = seq_lens_casual[offset : offset + qo_len] + offset += qo_len + torch.arange(kv_len - qo_len + 1, kv_len + 1, out=out) + idx_to_req_repeated[offset - qo_len : offset].fill_(i) + + assert offset == num_tokens + req_pool_indices_repeated = req_pool_indices[idx_to_req_repeated] + + if padded_num_tokens is not None and padded_num_tokens > num_tokens: + pad_size = padded_num_tokens - num_tokens + seq_lens_casual = torch.nn.functional.pad( + seq_lens_casual, (0, pad_size), value=1 + ) + req_pool_indices_repeated = torch.nn.functional.pad( + req_pool_indices_repeated, + (0, pad_size), + value=req_pool_indices_repeated[-1].item(), + ) + return ExpandPrefillCausallyResult( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + ) + + +@triton.jit +def _expand_prefill_causally_kernel( + req_pool_ptr, + seq_lens_ptr, + extend_seq_lens_ptr, + seq_lens_casual_ptr, + req_pool_repeated_ptr, + bs, + num_tokens, + total_tokens, + BLOCK: tl.constexpr, + BS_P2: tl.constexpr, +): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < total_tokens + + b = tl.arange(0, BS_P2) + bmask = b < bs + extend = tl.load(extend_seq_lens_ptr + b, mask=bmask, other=0).to(tl.int32) + start_locs = tl.cumsum(extend, axis=0) - extend + + is_real = offs < num_tokens + t = tl.where(is_real, offs, 0).to(tl.int32) + started = (start_locs[None, :] <= t[:, None]) & bmask[None, :] + r = tl.sum(started.to(tl.int32), axis=1) - 1 + r = tl.where(is_real, r, bs - 1).to(tl.int64) + + seq_len = tl.load(seq_lens_ptr + r, mask=mask, other=0).to(tl.int32) + ext = tl.load(extend_seq_lens_ptr + r, mask=mask, other=0).to(tl.int32) + start_loc = tl.sum(tl.where(started, extend[None, :], 0).to(tl.int32), axis=1) - ext + causal = (seq_len - ext + 1) + (t - start_loc) + causal = tl.where(is_real, causal, 1) + + rp = tl.load(req_pool_ptr + r, mask=mask, other=0) + tl.store(seq_lens_casual_ptr + offs, causal, mask=mask) + tl.store(req_pool_repeated_ptr + offs, rp, mask=mask) + + +def expand_prefill_causally_triton( + *, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + num_tokens: int, + padded_num_tokens: Optional[int], +) -> ExpandPrefillCausallyResult: + bs = req_pool_indices.shape[0] + device = req_pool_indices.device + total_tokens = ( + padded_num_tokens + if padded_num_tokens is not None and padded_num_tokens > num_tokens + else num_tokens + ) + + seq_lens_casual = torch.empty(total_tokens, dtype=torch.int32, device=device) + req_pool_indices_repeated = torch.empty( + total_tokens, dtype=req_pool_indices.dtype, device=device + ) + BLOCK = 256 + _expand_prefill_causally_kernel[(triton.cdiv(total_tokens, BLOCK),)]( + req_pool_indices, + seq_lens, + extend_seq_lens, + seq_lens_casual, + req_pool_indices_repeated, + bs, + num_tokens, + total_tokens, + BLOCK=BLOCK, + BS_P2=triton.next_power_of_2(max(bs, 1)), + ) + return ExpandPrefillCausallyResult( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + ) + + +class PageTablePositionsResult(msgspec.Struct): + seq_lens_casual: torch.Tensor + positions_casual: torch.Tensor + page_table: torch.Tensor + swa_topk_lengths: torch.Tensor + + +class BuildPageTablePositions: + @classmethod + def execute(cls, *args, **kwargs) -> PageTablePositionsResult: + if _inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + req_to_token: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + seq_lens_casual: torch.Tensor, + max_seq_len: int, + page_size: int, + swa_window: int, + ) -> PageTablePositionsResult: + return build_page_table_positions( + req_to_token=req_to_token, + req_pool_indices_repeated=req_pool_indices_repeated, + seq_lens_casual=seq_lens_casual, + max_seq_len=max_seq_len, + page_size=page_size, + swa_window=swa_window, + ) + + @classmethod + def triton( + cls, + *, + req_to_token: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + seq_lens_casual: torch.Tensor, + max_seq_len: int, + page_size: int, + swa_window: int, + ) -> PageTablePositionsResult: + return build_page_table_positions_triton( + req_to_token=req_to_token, + req_pool_indices_repeated=req_pool_indices_repeated, + seq_lens_casual=seq_lens_casual, + max_seq_len=max_seq_len, + page_size=page_size, + swa_window=swa_window, + ) + + +def build_page_table_positions( + *, + req_to_token: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + seq_lens_casual: torch.Tensor, + max_seq_len: int, + page_size: int, + swa_window: int, +) -> PageTablePositionsResult: + seq_lens_casual = seq_lens_casual.to(torch.int32) + positions_casual = seq_lens_casual - 1 + page_table = req_to_token[ + req_pool_indices_repeated.to(torch.int64), :max_seq_len:page_size + ] + page_table = (page_table // page_size).to(torch.int32) + swa_topk_lengths = torch.clamp(seq_lens_casual, max=swa_window) + return PageTablePositionsResult( + seq_lens_casual=seq_lens_casual, + positions_casual=positions_casual, + page_table=page_table, + swa_topk_lengths=swa_topk_lengths, + ) + + +@triton.jit +def _page_table_positions_kernel( + req_to_token_ptr, + req_pool_ptr, + seq_lens_ptr, + seq_lens_out_ptr, + positions_out_ptr, + page_table_ptr, + topk_out_ptr, + rt_stride, + num_pages, + page_size, + swa_window, + BLOCK_P: tl.constexpr, +): + row = tl.program_id(0) + seq_len = tl.load(seq_lens_ptr + row).to(tl.int32) + tl.store(seq_lens_out_ptr + row, seq_len) + tl.store(positions_out_ptr + row, seq_len - 1) + tl.store(topk_out_ptr + row, tl.minimum(seq_len, swa_window)) + + rp = tl.load(req_pool_ptr + row).to(tl.int64) + base = req_to_token_ptr + rp * rt_stride + out_base = page_table_ptr + row.to(tl.int64) * num_pages + for p0 in range(0, num_pages, BLOCK_P): + p = p0 + tl.arange(0, BLOCK_P) + pmask = p < num_pages + tok = tl.load(base + p.to(tl.int64) * page_size, mask=pmask, other=0).to( + tl.int32 + ) + tl.store(out_base + p, tok // page_size, mask=pmask) + + +def build_page_table_positions_triton( + *, + req_to_token: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + seq_lens_casual: torch.Tensor, + max_seq_len: int, + page_size: int, + swa_window: int, +) -> PageTablePositionsResult: + num_q = seq_lens_casual.shape[0] + num_pages = (max_seq_len + page_size - 1) // page_size + device = seq_lens_casual.device + + seq_lens_out = torch.empty(num_q, dtype=torch.int32, device=device) + positions_out = torch.empty(num_q, dtype=torch.int32, device=device) + page_table = torch.empty((num_q, num_pages), dtype=torch.int32, device=device) + topk_out = torch.empty(num_q, dtype=torch.int32, device=device) + BLOCK_P = 256 + _page_table_positions_kernel[(num_q,)]( + req_to_token, + req_pool_indices_repeated, + seq_lens_casual, + seq_lens_out, + positions_out, + page_table, + topk_out, + req_to_token.stride(0), + num_pages, + page_size, + swa_window, + BLOCK_P=BLOCK_P, + ) + return PageTablePositionsResult( + seq_lens_casual=seq_lens_out, + positions_casual=positions_out, + page_table=page_table, + swa_topk_lengths=topk_out, + ) + + +class BuildCausalSwaPageIndices: + @classmethod + def execute(cls, *args, **kwargs) -> torch.Tensor: + if _inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + req_to_token: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + seq_lens_casual: torch.Tensor, + swa_window: int, + page_index_aligned_size: int, + ) -> torch.Tensor: + return build_causal_swa_page_indices( + req_to_token=req_to_token, + full_to_swa_mapping=full_to_swa_mapping, + req_pool_indices_repeated=req_pool_indices_repeated, + seq_lens_casual=seq_lens_casual, + swa_window=swa_window, + page_index_aligned_size=page_index_aligned_size, + ) + + @classmethod + def triton( + cls, + *, + req_to_token: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + seq_lens_casual: torch.Tensor, + swa_window: int, + page_index_aligned_size: int, + ) -> torch.Tensor: + return build_causal_swa_page_indices_triton( + req_to_token=req_to_token, + full_to_swa_mapping=full_to_swa_mapping, + req_pool_indices_repeated=req_pool_indices_repeated, + seq_lens_casual=seq_lens_casual, + swa_window=swa_window, + page_index_aligned_size=page_index_aligned_size, + ) + + +def build_causal_swa_page_indices( + *, + req_to_token: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + seq_lens_casual: torch.Tensor, + swa_window: int, + page_index_aligned_size: int, +) -> torch.Tensor: + device = seq_lens_casual.device + pos_causal = seq_lens_casual - 1 + num_qo_tokens = seq_lens_casual.size(0) + offsets = pos_causal.unsqueeze(1) - torch.arange( + swa_window, dtype=torch.int32, device=device + ).unsqueeze(0) + invalid_offset_mask = offsets < 0 + offsets.masked_fill_(invalid_offset_mask, 0) + raw_indices = req_to_token[req_pool_indices_repeated[:, None], offsets] + assert raw_indices.shape == (num_qo_tokens, swa_window) + raw_indices.masked_fill_(invalid_offset_mask, -1) + swa_indices = full_to_swa_mapping[raw_indices] + swa_indices = swa_indices.to(torch.int32) + + padded_width = ( + (swa_window + page_index_aligned_size - 1) // page_index_aligned_size + ) * page_index_aligned_size + if padded_width == swa_window: + return swa_indices + return torch.nn.functional.pad( + swa_indices, (0, padded_width - swa_window), value=-1 + ) + + +@triton.jit +def _causal_swa_page_indices_kernel( + req_to_token_ptr, + full_to_swa_ptr, + req_pool_ptr, + seq_lens_ptr, + out_ptr, + rt_stride, + swa_window, + padded_width, + BLOCK_K: tl.constexpr, +): + row = tl.program_id(0) + pos = tl.load(seq_lens_ptr + row).to(tl.int64) - 1 + rp = tl.load(req_pool_ptr + row).to(tl.int64) + base = req_to_token_ptr + rp * rt_stride + out_base = out_ptr + row.to(tl.int64) * padded_width + + for k0 in range(0, padded_width, BLOCK_K): + k = k0 + tl.arange(0, BLOCK_K) + kmask = k < padded_width + off = pos - k.to(tl.int64) + valid = (k < swa_window) & (off >= 0) & kmask + full_loc = tl.load(base + tl.where(valid, off, 0), mask=valid, other=-1).to( + tl.int64 + ) + swa = tl.load(full_to_swa_ptr + full_loc, mask=valid, other=-1).to(tl.int32) + tl.store(out_base + k, tl.where(valid, swa, -1), mask=kmask) + + +def build_causal_swa_page_indices_triton( + *, + req_to_token: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + seq_lens_casual: torch.Tensor, + swa_window: int, + page_index_aligned_size: int, +) -> torch.Tensor: + num_qo_tokens = seq_lens_casual.size(0) + padded_width = ( + (swa_window + page_index_aligned_size - 1) // page_index_aligned_size + ) * page_index_aligned_size + out = torch.empty( + (num_qo_tokens, padded_width), + dtype=torch.int32, + device=seq_lens_casual.device, + ) + BLOCK_K = 256 + _causal_swa_page_indices_kernel[(num_qo_tokens,)]( + req_to_token, + full_to_swa_mapping, + req_pool_indices_repeated, + seq_lens_casual, + out, + req_to_token.stride(0), + swa_window, + padded_width, + BLOCK_K=BLOCK_K, + ) + return out diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 94d307d67..56aaa042b 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -29,6 +29,7 @@ from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_server_args +from sglang.srt.speculative.ragged_verify import build_ragged_target_verify_geometry from sglang.srt.speculative.spec_info import SpecInput, SpeculativeAlgorithm from sglang.srt.utils import get_compiler_backend @@ -216,6 +217,8 @@ class FlashAttentionBackend(AttentionBackend): - For each forward batch, init_replay_cuda_graph will be called first and then replay the graph. """ + supports_ragged_verify_graph: bool = True + def __init__( self, model_runner: ModelRunner, @@ -253,11 +256,12 @@ class FlashAttentionBackend(AttentionBackend): self.max_num_pages = ( self.max_context_len + self.page_size - 1 ) // self.page_size - # Opt out of the seq_lens_cpu D2H only for dflash (the worker adapted to - # the GPU-only relay); EAGLE/MTP/standalone/non-spec keep the CPU mirror. + # Opt out of the seq_lens_cpu D2H only for dflash/dspark (their workers + # adapted to the GPU-only relay); EAGLE/MTP/standalone/non-spec keep the + # CPU mirror. self.needs_cpu_seq_lens = not SpeculativeAlgorithm.from_string( model_runner.server_args.speculative_algorithm - ).is_dflash() + ).is_dflash_family() self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA self.skip_prefill = skip_prefill self.attn_cp_size = model_runner.attn_cp_size @@ -272,6 +276,15 @@ class FlashAttentionBackend(AttentionBackend): self.speculative_num_draft_tokens = ( model_runner.server_args.speculative_num_draft_tokens ) + if ( + self.speculative_num_draft_tokens is not None + and model_runner.is_draft_worker + ): + self.speculative_num_draft_tokens = SpeculativeAlgorithm.from_string( + model_runner.server_args.speculative_algorithm + ).get_num_tokens_per_bs_for_target_verify( + int(self.speculative_num_draft_tokens), is_draft_worker=True + ) self.speculative_step_id = speculative_step_id # Local attention settings @@ -714,26 +727,47 @@ class FlashAttentionBackend(AttentionBackend): self._maybe_init_local_attn_metadata(forward_batch, metadata, device) elif forward_batch.forward_mode.is_target_verify(): if self.topk <= 1: - metadata.cache_seqlens_int32 = ( - forward_batch.seq_lens + self.speculative_num_draft_tokens - ).to(torch.int32) - metadata.max_seq_len_q = self.speculative_num_draft_tokens - metadata.max_seq_len_k = ( - seq_lens_cpu.max().item() + self.speculative_num_draft_tokens - ) - metadata.cu_seqlens_q = torch.arange( - 0, - batch_size * self.speculative_num_draft_tokens + 1, - self.speculative_num_draft_tokens, - dtype=torch.int32, - device=device, - ) - metadata.cu_seqlens_k = torch.nn.functional.pad( - torch.cumsum( - metadata.cache_seqlens_int32, dim=0, dtype=torch.int32 - ), - (1, 0), + ragged_layout = getattr( + forward_batch.spec_info, "ragged_verify_layout", None ) + if ragged_layout is not None: + geometry = build_ragged_target_verify_geometry( + seq_lens=forward_batch.seq_lens, layout=ragged_layout + ) + metadata.cache_seqlens_int32 = geometry.cache_seqlens_int32 + # Device-only layouts carry no host lens; the verify + # window is a valid varlen upper bound. + metadata.max_seq_len_q = ( + geometry.max_seq_len_q + if geometry.max_seq_len_q is not None + else self.speculative_num_draft_tokens + ) + metadata.max_seq_len_k = int( + metadata.cache_seqlens_int32.max().item() + ) + metadata.cu_seqlens_q = geometry.cu_seqlens_q + metadata.cu_seqlens_k = geometry.cu_seqlens_k + else: + metadata.cache_seqlens_int32 = ( + forward_batch.seq_lens + self.speculative_num_draft_tokens + ).to(torch.int32) + metadata.max_seq_len_q = self.speculative_num_draft_tokens + metadata.max_seq_len_k = ( + seq_lens_cpu.max().item() + self.speculative_num_draft_tokens + ) + metadata.cu_seqlens_q = torch.arange( + 0, + batch_size * self.speculative_num_draft_tokens + 1, + self.speculative_num_draft_tokens, + dtype=torch.int32, + device=device, + ) + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum( + metadata.cache_seqlens_int32, dim=0, dtype=torch.int32 + ), + (1, 0), + ) metadata.page_table = self.req_to_token_pool.req_to_token[ forward_batch.req_pool_indices, : metadata.max_seq_len_k ] @@ -2563,9 +2597,18 @@ class FlashAttentionBackend(AttentionBackend): elif forward_mode.is_target_verify(): if self.topk <= 1: metadata = self.target_verify_metadata[bs] - metadata.cache_seqlens_int32.copy_( - (seq_lens + self.speculative_num_draft_tokens) - ) + ragged_layout = getattr(spec_info, "ragged_verify_layout", None) + if ragged_layout is not None: + padded = ragged_layout.padded_to_bucket(padded_bs=bs) + geometry = build_ragged_target_verify_geometry( + seq_lens=seq_lens, layout=padded + ) + metadata.cache_seqlens_int32.copy_(geometry.cache_seqlens_int32) + metadata.cu_seqlens_q.copy_(geometry.cu_seqlens_q) + else: + metadata.cache_seqlens_int32.copy_( + (seq_lens + self.speculative_num_draft_tokens) + ) # Page table built on-device (self-guards on cache_seqlens); # max_seq_len_k left unset -- unread here (scheduler_metadata is diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 8e537259a..a55f7f5e7 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -657,6 +657,16 @@ class FlashInferAttnBackend(AttentionBackend): forward_mode = forward_batch.forward_mode spec_info = forward_batch.spec_info + if ( + spec_info is not None + and spec_info.ragged_verify_layout is not None + and forward_mode.is_target_verify() + ): + raise NotImplementedError( + "FlashInfer does not support ragged verify in cuda graph; " + "disable SGLANG_RAGGED_VERIFY_MODE for this configuration." + ) + if in_capture: num_tokens = forward_batch.positions.numel() self._prepare_cuda_graph_metadata(bs, num_tokens, forward_mode, spec_info) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 006911a48..ddcd1c345 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -32,6 +32,10 @@ from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_buffer +from sglang.srt.speculative.ragged_verify import ( + build_ragged_target_verify_geometry, + resolve_ragged_verify_layout, +) from sglang.srt.utils import is_flashinfer_available from sglang.srt.utils.common import is_sm90_supported, is_sm120_supported @@ -43,6 +47,7 @@ if is_flashinfer_available(): if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout from sglang.srt.speculative.spec_info import SpecInput # Constants @@ -69,6 +74,7 @@ class TRTLLMMHAMetadata: swa_page_table: torch.Tensor = None # full->SWA translated out_cache_loc (SWA KV-store write target) swa_out_cache_loc: torch.Tensor = None + is_ragged_verify: bool = False class TRTLLMHAAttnBackend(FlashInferAttnBackend): @@ -79,6 +85,8 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): # seq_lens_cpu D2H sync; opt out of it, matching trtllm_mla / triton. needs_cpu_seq_lens: bool = False + supports_ragged_verify_graph: bool = True + def __init__( self, model_runner: ModelRunner, @@ -353,6 +361,10 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): "cache_seqlens": torch.zeros( max_bs, dtype=torch.int32, device=self.device ), + # Static uniform preset (Q_MODE_NONE: the fused kernel never + # rewrites it). Ragged verify overwrites the [:bs+1] slice + # eagerly on every capture/replay-prep, and the ragged-verify + # mode is fixed for the whole server run, so the two never mix. "cu_seqlens_q": torch.arange( 0, max_bs * self.speculative_num_draft_tokens + 1, @@ -449,7 +461,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): self.decode_cuda_graph_metadata[bs] = metadata elif forward_mode.is_target_verify(): # Target Verify (topk = 1) - tokens_per_req = num_tokens // bs metadata.cache_seqlens_int32 = self.target_verify_metadata["cache_seqlens"][ :bs ] @@ -459,7 +470,14 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): metadata.cu_seqlens_k = self.target_verify_metadata["cu_seqlens_k"][ : bs + 1 ] - metadata.max_seq_len_q = tokens_per_req + metadata.is_ragged_verify = ( + spec_info is not None and spec_info.ragged_verify_layout is not None + ) + metadata.max_seq_len_q = ( + self.speculative_num_draft_tokens + if metadata.is_ragged_verify + else num_tokens // bs + ) metadata.page_table = self.target_verify_metadata["page_table"][:bs, :] self._bind_swa_page_table( metadata, @@ -536,6 +554,13 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): elif forward_mode.is_target_verify(): # Here we only support topk = 1 for now. metadata = self.target_verify_metadata[bs] + if spec_info is not None and spec_info.ragged_verify_layout is not None: + # Ragged verify: the per-request k-extension is not a + # uniform scalar seqlen_offset, so the fused kernel cannot + # rebuild this metadata. It is written eagerly on every + # capture/replay-prep in init_forward_metadata_out_graph; + # record nothing here. + return seqlen_offset = metadata.max_seq_len_q elif forward_mode.is_draft_extend_v2(): metadata = self.draft_extend_metadata[bs] @@ -628,6 +653,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): forward_mode = forward_batch.forward_mode spec_info = forward_batch.spec_info + if ( + forward_mode.is_target_verify() + and resolve_ragged_verify_layout(forward_batch) is not None + ): + self._assert_ragged_verify_supported() + if in_capture: num_tokens = forward_batch.positions.numel() self._build_cuda_graph_metadata( @@ -638,6 +669,11 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): self.forward_metadata = self.decode_cuda_graph_metadata[bs] elif forward_mode.is_target_verify(): self.forward_metadata = self.target_verify_metadata[bs] + ragged_layout = resolve_ragged_verify_layout(forward_batch) + if ragged_layout is not None: + self._write_ragged_verify_graph_metadata( + self.forward_metadata, forward_batch, ragged_layout, bs + ) elif forward_mode.is_draft_extend_v2(): self.forward_metadata = self.draft_extend_metadata[bs] else: @@ -645,6 +681,53 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): f"Invalid forward mode: {forward_mode=} for CUDA Graph replay." ) + def _assert_ragged_verify_supported(self) -> None: + if self.is_xqa_impl: + raise NotImplementedError( + "Compact ragged verify (variable-length cum_seq_lens_q) " + "requires the trtllm-gen decode kernel; the xqa impl (sm90 / sm120) " + "rejects it. Disable SGLANG_RAGGED_VERIFY_MODE for this configuration." + ) + + def _write_ragged_verify_graph_metadata( + self, + metadata: TRTLLMMHAMetadata, + forward_batch: ForwardBatch, + ragged_layout: RaggedVerifyLayout, + bs: int, + ) -> None: + """Eagerly rebuild the target-verify graph metadata for ragged verify. + + The per-request verify lengths make the k-extension non-uniform, which + the fused in-graph kernel cannot express (scalar ``seqlen_offset`` + only), so this runs out-of-graph on every capture/replay-prep and + ``_apply_cuda_graph_metadata`` records nothing for ragged batches. + """ + seq_lens = forward_batch.seq_lens[:bs] + req_pool_indices = forward_batch.req_pool_indices[:bs] + padded_layout = ragged_layout.padded_to_bucket(padded_bs=bs) + geometry = build_ragged_target_verify_geometry( + seq_lens=seq_lens, layout=padded_layout + ) + metadata.cache_seqlens_int32.copy_(geometry.cache_seqlens_int32) + metadata.cu_seqlens_q.copy_(geometry.cu_seqlens_q) + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32) + ) + self._fill_page_table_device( + metadata, req_pool_indices, metadata.cache_seqlens_int32 + ) + # The fused in-graph kernel also skips ragged batches, so refill the + # SWA write-target buffer here (out_cache_loc -> SWA locs). + if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None: + n = forward_batch.out_cache_loc.shape[0] + self.cuda_graph_swa_out_cache_loc[n:].zero_() + self.cuda_graph_swa_out_cache_loc[:n].copy_( + self.token_to_kv_pool.translate_loc_from_full_to_swa( + forward_batch.out_cache_loc + ) + ) + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): self._apply_cuda_graph_metadata( bs=forward_batch.batch_size, @@ -689,23 +772,42 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0) ) elif forward_batch.forward_mode.is_target_verify(): - # Only support topk = 1 for now. - tokens_per_req = forward_batch.input_ids.shape[0] // batch_size - metadata.cache_seqlens_int32 = (forward_batch.seq_lens + tokens_per_req).to( - torch.int32 - ) - metadata.max_seq_len_q = tokens_per_req - metadata.cu_seqlens_q = torch.arange( - 0, - batch_size * tokens_per_req + 1, - tokens_per_req, - dtype=torch.int32, - device=device, - ) - metadata.cu_seqlens_k = torch.nn.functional.pad( - torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32), - (1, 0), - ) + ragged_layout = resolve_ragged_verify_layout(forward_batch) + if ragged_layout is not None: + self._assert_ragged_verify_supported() + geometry = build_ragged_target_verify_geometry( + seq_lens=seqlens_in_batch, layout=ragged_layout + ) + metadata.cache_seqlens_int32 = geometry.cache_seqlens_int32 + # Device-only layouts carry no host lens; the verify window + # is a valid varlen upper bound. + metadata.max_seq_len_q = ( + geometry.max_seq_len_q + if geometry.max_seq_len_q is not None + else self.speculative_num_draft_tokens + ) + metadata.cu_seqlens_q = geometry.cu_seqlens_q + metadata.cu_seqlens_k = geometry.cu_seqlens_k + metadata.is_ragged_verify = True + else: + tokens_per_req = forward_batch.input_ids.shape[0] // batch_size + metadata.cache_seqlens_int32 = ( + forward_batch.seq_lens + tokens_per_req + ).to(torch.int32) + metadata.max_seq_len_q = tokens_per_req + metadata.cu_seqlens_q = torch.arange( + 0, + batch_size * tokens_per_req + 1, + tokens_per_req, + dtype=torch.int32, + device=device, + ) + metadata.cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum( + metadata.cache_seqlens_int32, dim=0, dtype=torch.int32 + ), + (1, 0), + ) else: metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32) @@ -906,21 +1008,40 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend_v2() ): - o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( - query=q, - kv_cache=kv_cache, - workspace_buffer=self.workspace_buffer, - block_tables=page_table, - seq_lens=self.forward_metadata.cache_seqlens_int32, - max_seq_len=self.max_context_len, - bmm1_scale=bmm1_scale, - bmm2_scale=bmm2_scale, - window_left=layer.sliding_window_size, - sinks=attention_sink, - skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), - out_dtype=self.q_data_type, # model_runner.dtype - q_len_per_req=self.forward_metadata.max_seq_len_q, - ) + if self.forward_metadata.is_ragged_verify: + o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( + query=q, + kv_cache=kv_cache, + workspace_buffer=self.workspace_buffer, + block_tables=page_table, + seq_lens=self.forward_metadata.cache_seqlens_int32, + max_seq_len=self.max_context_len, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + window_left=layer.sliding_window_size, + sinks=attention_sink, + skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), + out_dtype=self.q_data_type, + q_len_per_req=None, + max_q_len=self.forward_metadata.max_seq_len_q, + cum_seq_lens_q=self.forward_metadata.cu_seqlens_q, + ) + else: + o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( + query=q, + kv_cache=kv_cache, + workspace_buffer=self.workspace_buffer, + block_tables=page_table, + seq_lens=self.forward_metadata.cache_seqlens_int32, + max_seq_len=self.max_context_len, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + window_left=layer.sliding_window_size, + sinks=attention_sink, + skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), + out_dtype=self.q_data_type, + q_len_per_req=self.forward_metadata.max_seq_len_q, + ) else: o = flashinfer.prefill.trtllm_batch_context_with_kv_cache( query=q, diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index fa9a1d262..d2f515434 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -428,7 +428,10 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): video_tokens=recv_obj.video_tokens, spec_verify_ct=recv_obj.spec_verify_ct, spec_num_correct_drafts=recv_obj.spec_num_correct_drafts, + spec_num_block_accept_tokens=recv_obj.spec_num_block_accept_tokens, + spec_num_cap_tokens=recv_obj.spec_num_cap_tokens, spec_correct_drafts_histogram=recv_obj.spec_correct_drafts_histogram, + spec_cap_lens_histogram=recv_obj.spec_cap_lens_histogram, input_token_logprobs_val=recv_obj.input_token_logprobs_val, input_token_logprobs_idx=recv_obj.input_token_logprobs_idx, output_token_logprobs_val=recv_obj.output_token_logprobs_val, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 60985e55e..bf773a2d8 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1274,8 +1274,11 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True): spec_verify_ct: Optional[List[int]] = None # Accepted drafts spec_num_correct_drafts: Optional[List[int]] = None + spec_num_block_accept_tokens: Optional[List[int]] = None + spec_num_cap_tokens: Optional[List[int]] = None # Acceptance histogram spec_correct_drafts_histogram: Optional[List[List[int]]] = None + spec_cap_lens_histogram: Optional[List[List[int]]] = None class BatchStrOutput(BaseBatchReq, kw_only=True): @@ -1349,8 +1352,11 @@ class BatchStrOutput(BaseBatchReq, kw_only=True): spec_verify_ct: Optional[List[int]] = None # Accepted drafts spec_num_correct_drafts: Optional[List[int]] = None + spec_num_block_accept_tokens: Optional[List[int]] = None + spec_num_cap_tokens: Optional[List[int]] = None # Acceptance histogram spec_correct_drafts_histogram: Optional[List[List[int]]] = None + spec_cap_lens_histogram: Optional[List[List[int]]] = None class BatchEmbeddingOutput(BaseBatchReq, kw_only=True): diff --git a/python/sglang/srt/managers/min_free_slots_delayer.py b/python/sglang/srt/managers/min_free_slots_delayer.py index 9f65d9ec6..9c0020939 100644 --- a/python/sglang/srt/managers/min_free_slots_delayer.py +++ b/python/sglang/srt/managers/min_free_slots_delayer.py @@ -4,7 +4,7 @@ from typing import Optional def resolve_min_free_slots( user_value: Optional[int], max_running_requests: int, - is_dflash: bool = False, + is_dflash_family: bool = False, ) -> Optional[int]: """Resolve the min-free-slots threshold (None = disabled). @@ -16,7 +16,7 @@ def resolve_min_free_slots( max_running_requests = max(0, int(max_running_requests)) formula = min(4, max(2, (max_running_requests + 5) // 6)) if user_value is None: - user_value = formula if is_dflash else None + user_value = formula if is_dflash_family else None if user_value is None or user_value <= 1: return None diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 99e809a40..e3675ad4c 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -161,6 +161,15 @@ def _handle_output_by_index(output, i): spec_correct_drafts_histogram=_extract_field_by_index( output, "spec_correct_drafts_histogram", i ), + spec_num_block_accept_tokens=_extract_field_by_index( + output, "spec_num_block_accept_tokens", i + ), + spec_num_cap_tokens=_extract_field_by_index( + output, "spec_num_cap_tokens", i + ), + spec_cap_lens_histogram=_extract_field_by_index( + output, "spec_cap_lens_histogram", i + ), time_stats=_extract_field_by_index(output, "time_stats", i), finished_reasons=_extract_field_by_index(output, "finished_reasons", i), decoded_texts=_extract_field_by_index(output, "decoded_texts", i), @@ -263,6 +272,15 @@ def _handle_output_by_index(output, i): spec_correct_drafts_histogram=_extract_field_by_index( output, "spec_correct_drafts_histogram", i ), + spec_num_block_accept_tokens=_extract_field_by_index( + output, "spec_num_block_accept_tokens", i + ), + spec_num_cap_tokens=_extract_field_by_index( + output, "spec_num_cap_tokens", i + ), + spec_cap_lens_histogram=_extract_field_by_index( + output, "spec_cap_lens_histogram", i + ), time_stats=_extract_field_by_index(output, "time_stats", i), finished_reasons=_extract_field_by_index(output, "finished_reasons", i), output_strs=_extract_field_by_index(output, "output_strs", i), diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index d27c89ad6..5777ac983 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -1,8 +1,9 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Optional, Sequence +from typing import TYPE_CHECKING, Any, Optional, Sequence +import msgspec import torch from sglang.kernels.ops.speculative.gather_spec_extras import gather_spec_extras @@ -35,7 +36,8 @@ def decide_needs_cpu_seq_lens( if server_args.enable_two_batch_overlap: # FIXME: support TBO without seq lens cpu value return True - if SpeculativeAlgorithm.from_string(server_args.speculative_algorithm).is_ngram(): + algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm) + if algo.is_ngram(): # ngram's USE_FULL_MASK verify path reads seq_lens_cpu per req to size # the tree mask, regardless of the attn backend (e.g. Triton opts out). return True @@ -46,6 +48,19 @@ def decide_needs_cpu_seq_lens( ) +def decide_needs_confidence_relay(server_args: ServerArgs) -> bool: + from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyMode, + read_ragged_verify_mode, + ) + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm) + if not algo.is_dspark(): + return False + return read_ragged_verify_mode() is not RaggedVerifyMode.STATIC + + _is_cuda = is_cuda() _is_hip = is_hip() _is_npu = is_npu() @@ -100,6 +115,16 @@ def resolve_forward_inputs(batch: ScheduleBatch, future_map: FutureMap) -> None: future_map._resolve_spec_extras(batch) +CONFIDENCE_RELAY_RING_LAG: int = 2 +CONFIDENCE_RELAY_RING_DEPTH: int = CONFIDENCE_RELAY_RING_LAG + 1 + + +class ResolvedConfidence(msgspec.Struct): + + confidence: torch.Tensor + generation: torch.Tensor + + @dataclass class RelayPayload: """Per-iteration stash payload for the FutureMap bufs. Non-spec fills only @@ -125,6 +150,85 @@ class RelayPayload: ) +class ConfidenceRelay(msgspec.Struct): + + device: torch.device + req_pool_size: int + pool: Any + confidence_buf: Optional[torch.Tensor] = None + conf_ring: Optional[torch.Tensor] = None + gen_ring: Optional[torch.Tensor] = None + copy_done: Optional[list] = None + ring_pos: int = 0 + initialized: bool = False + + def _lazy_init(self, confidence: torch.Tensor) -> None: + self.initialized = True + gamma = confidence.shape[-1] + self.confidence_buf = torch.empty( + (self.req_pool_size, gamma), dtype=torch.float32, device=self.device + ) + if _is_cuda: + depth = CONFIDENCE_RELAY_RING_DEPTH + self.conf_ring = torch.empty( + (depth, self.req_pool_size, gamma), + dtype=torch.float32, + pin_memory=True, + ) + self.gen_ring = torch.zeros((depth, self.req_pool_size), dtype=torch.int64) + self.copy_done = [ + torch.get_device_module(self.device).Event() for _ in range(depth) + ] + + def scatter(self, indices: torch.Tensor, confidence: torch.Tensor) -> None: + if not self.initialized: + self._lazy_init(confidence) + self.confidence_buf[indices] = confidence.to(self.confidence_buf.dtype) + + def issue_ring_copy(self, *, stream, publish_ready) -> None: + if not self.initialized or stream is None or publish_ready is None: + return + slot = self.ring_pos % CONFIDENCE_RELAY_RING_DEPTH + stream.wait_event(publish_ready) + with torch.get_device_module(self.device).stream(stream): + self.conf_ring[slot].copy_(self.confidence_buf, non_blocking=True) + self.copy_done[slot].record() + self.gen_ring[slot].copy_(self.pool.req_generation) + self.ring_pos += 1 + + def resolve( + self, batch: ScheduleBatch, *, stream, publish_ready + ) -> Optional[ResolvedConfidence]: + if not self.initialized: + return None + draft_input = batch.spec_info + if draft_input is None: + return None + fi = draft_input.future_indices + if fi is None or fi.shape[0] == 0: + return None + + if stream is None or publish_ready is None: + idx = batch.req_pool_indices + idx_cpu = batch.req_pool_indices_cpu + return ResolvedConfidence( + confidence=self.confidence_buf[idx].cpu(), + generation=self.pool.req_generation[idx_cpu].clone(), + ) + + if self.ring_pos < CONFIDENCE_RELAY_RING_LAG: + return None + slot = (self.ring_pos - CONFIDENCE_RELAY_RING_LAG) % CONFIDENCE_RELAY_RING_DEPTH + if not self.copy_done[slot].query(): + return None + + idx_cpu = batch.req_pool_indices_cpu + return ResolvedConfidence( + confidence=self.conf_ring[slot][idx_cpu], + generation=self.gen_ring[slot][idx_cpu], + ) + + class FutureMap: """Always-on pool-indexed relay for cross-iter values. Forward writes via publish/stash; next iter reads via resolve_forward_inputs / resolve_seq_lens_cpu. @@ -136,6 +240,7 @@ class FutureMap: spec_algo: SpeculativeAlgorithm, req_to_token_pool: ReqToTokenPool, needs_cpu_seq_lens: bool = True, + needs_confidence_relay: bool = False, ): # Bufs indexed by req_pool_idx; slot 0 mirrors KV padding row so # CUDA-graph padded batches (req_pool_idx == 0) are harmless. @@ -144,6 +249,7 @@ class FutureMap: # Computed by decide_needs_cpu_seq_lens(); see that helper for the # full decision (per-backend flag + TBO / piecewise CG overrides). self.needs_cpu_seq_lens = needs_cpu_seq_lens + self.needs_confidence_relay = needs_confidence_relay self.req_pool_size = req_to_token_pool.req_to_token.shape[0] if _DEBUG_ASSERT: @@ -181,6 +287,12 @@ class FutureMap: # resolve; arm/consume strictly alternate across all batch interleavings. self._publish_fresh = False + self.confidence_relay = ConfidenceRelay( + device=self.device, + req_pool_size=self.req_pool_size, + pool=req_to_token_pool, + ) + def _lazy_init_forward_buf(self, payload: RelayPayload): # Local import (see decide_needs_cpu_seq_lens): keep module-level deps leaf. from sglang.srt.speculative.spec_utils import spec_need_hidden_states @@ -235,6 +347,17 @@ class FutureMap: device=self.device, ) + def resolve_confidence_cpu( + self, batch: ScheduleBatch + ) -> Optional[ResolvedConfidence]: + if not self.needs_confidence_relay: + return None + return self.confidence_relay.resolve( + batch, + stream=self.fwd_prepare_d2h_stream, + publish_ready=self.publish_ready, + ) + def _resolve_spec_extras(self, batch: ScheduleBatch) -> None: if self.spec_algo.is_ngram(): # FIXME: remove once precomputed draft is supported. @@ -339,11 +462,19 @@ class FutureMap: # mirror is not poisoned. _assert_nonneg_and_invalidate(batch.seq_lens, self.new_seq_lens_buf, fi) - def publish(self, future_indices: torch.Tensor, new_seq_lens: torch.Tensor) -> None: + def publish( + self, + future_indices: torch.Tensor, + new_seq_lens: torch.Tensor, + confidence: Optional[torch.Tensor] = None, + ) -> None: indices = future_indices if indices.shape[0] == 0: return # DP idle self.new_seq_lens_buf[indices] = new_seq_lens.to(self.new_seq_lens_buf.dtype) + publish_confidence = self.needs_confidence_relay and confidence is not None + if publish_confidence: + self.confidence_relay.scatter(indices, confidence) # Only spec_v2 needs the event; it gates the seq_lens D2H on the private stream. if self.spec_algo.is_some(): device_module = torch.get_device_module(self.device) @@ -356,6 +487,11 @@ class FutureMap: device_module.current_stream().wait_event(self.publish_ready) self.publish_ready.record() self._publish_fresh = True + if publish_confidence: + self.confidence_relay.issue_ring_copy( + stream=self.fwd_prepare_d2h_stream, + publish_ready=self.publish_ready, + ) def stash(self, future_indices: torch.Tensor, payload: RelayPayload) -> None: if self.spec_algo.is_ngram(): diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 13a1bdf7c..4c9be3c9c 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -965,11 +965,17 @@ class Req(ReqDllmMixin): # Per-request count of accepted draft tokens (excludes the bonus token). self.spec_num_correct_drafts = 0 + self.spec_num_block_accept_tokens = 0 + + self.spec_num_cap_tokens = 0 + # Acceptance histogram for speculative decoding. # List index = number of accepted tokens in a step, List value = count of steps with that many accepted tokens. # Example: histogram[0] = 5 means 5 steps with 0 accepted tokens, histogram[3] = 10 means 10 steps with 3 accepted tokens. self.spec_correct_drafts_histogram: List[int] = [] + self.spec_cap_lens_histogram: List[int] = [] + # The number of times this request has been retracted / preempted. self.retraction_count = 0 self.retraction_mb_id = None @@ -1094,6 +1100,14 @@ class Req(ReqDllmMixin): ) self.spec_correct_drafts_histogram[num_correct_drafts] += 1 + def update_spec_cap_lens_histogram(self, cap_len: int): + cap_len = int(cap_len) + if len(self.spec_cap_lens_histogram) <= cap_len: + self.spec_cap_lens_histogram.extend( + [0] * (cap_len - len(self.spec_cap_lens_histogram) + 1) + ) + self.spec_cap_lens_histogram[cap_len] += 1 + def extend_image_inputs(self, image_inputs): if self.multimodal_inputs is None: self.multimodal_inputs = image_inputs @@ -1865,6 +1879,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): can_run_dp_cuda_graph: bool = False can_run_dp_breakable_cuda_graph: bool = False tbo_split_seq_index: Optional[int] = None + spec_verify_tier_num_tokens: int = -1 # For processing logprobs return_logprob: bool = False @@ -1910,6 +1925,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # For DP attention global_num_tokens: Optional[List[int]] = None global_num_tokens_for_logprob: Optional[List[int]] = None + global_spec_verify_tier_num_tokens: Optional[List[int]] = None # === Compound crossing to ForwardBatch (carry their own device tensors) === # Sampling info diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ea657dfa1..41e21015a 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -154,6 +154,7 @@ from sglang.srt.managers.min_free_slots_delayer import ( from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors from sglang.srt.managers.overlap_utils import ( RelayPayload, + decide_needs_confidence_relay, decide_needs_cpu_seq_lens, resolve_forward_inputs, ) @@ -884,7 +885,7 @@ class Scheduler( min_free_slots = resolve_min_free_slots( self.server_args.min_free_slots_delay, self.max_running_requests, - is_dflash=self.spec_algorithm.is_dflash(), + is_dflash_family=self.spec_algorithm.is_dflash_family(), ) if min_free_slots is not None: self.min_free_slots_delayer = MinFreeSlotsDelayer( @@ -1252,12 +1253,24 @@ class Scheduler( else: attn_backends = (self.tp_worker.model_runner.attn_backend,) needs_cpu_seq_lens = decide_needs_cpu_seq_lens(self.server_args, attn_backends) + needs_confidence_relay = decide_needs_confidence_relay(self.server_args) self.future_map = self.spec_algorithm.create_future_map( self.device, self.req_to_token_pool, needs_cpu_seq_lens=needs_cpu_seq_lens, + needs_confidence_relay=needs_confidence_relay, ) + self._confidence_budget_prepare = None + if ( + needs_confidence_relay + and self.enable_overlap + and self.draft_worker is not None + ): + self._confidence_budget_prepare = ( + self.draft_worker.get_confidence_budget_prepare() + ) + if use_mlx(): # MLX uses its own overlap loop and does not create CUDA streams, # but the normal non-overlap scheduler path still relays decode @@ -2173,7 +2186,7 @@ class Scheduler( self._add_request_to_queue(req) return - if self.spec_algorithm.is_dflash(): + if self.spec_algorithm.is_dflash_family(): error_msg = validate_dflash_request(req, self.enable_overlap) if error_msg is not None: req.set_finish_with_abort(error_msg) @@ -3242,6 +3255,8 @@ class Scheduler( # Self-gates on batch.spec_info.future_indices; non-spec_v2 # no-ops (ForwardBatch.init_new lazily computes the sum). self.future_map.resolve_seq_lens_cpu(batch) + if self._confidence_budget_prepare is not None: + self._confidence_budget_prepare(batch, self.future_map) with self.forward_stream_ctx: self.forward_stream.wait_stream(self.schedule_stream) @@ -3805,6 +3820,11 @@ class Scheduler( if RECORD_STEP_TIME: ret["step_time_dict"] = self.metrics_reporter.step_time_dict + if self.spec_algorithm.is_dspark() and self.draft_worker is not None: + info_record = self.draft_worker.dump_info_records() + if info_record is not None: + ret["dspark_info_record"] = info_record + # This field is not serializable. ret.pop("model_config", None) @@ -3817,6 +3837,8 @@ class Scheduler( "pp_max_micro_batch_size", "speculative_accept_threshold_single", "speculative_accept_threshold_acc", + "dspark_force_budget_frac", + "dspark_clear_info_records", ] ) @@ -3834,6 +3856,30 @@ class Scheduler( ) if_success = False break + elif k == "dspark_force_budget_frac": + if not self.spec_algorithm.is_dspark() or not hasattr( + self.draft_worker, "set_dspark_forced_budget_frac" + ): + logging.warning( + "dspark_force_budget_frac requires a DSpark draft worker." + ) + if_success = False + break + if v is not None and not (0.0 < float(v) <= 1.0): + logging.warning( + f"dspark_force_budget_frac must be in (0, 1] or null, got {v}." + ) + if_success = False + break + elif k == "dspark_clear_info_records": + if not self.spec_algorithm.is_dspark() or not hasattr( + self.draft_worker, "clear_info_records" + ): + logging.warning( + "dspark_clear_info_records requires a DSpark draft worker." + ) + if_success = False + break if if_success: if ( @@ -3848,7 +3894,18 @@ class Scheduler( self.metrics_reporter.spec_total_num_accept_tokens = ( self.metrics_reporter.spec_total_num_forward_ct ) = 0 - get_server_args().override(source="update_server_args", **server_args_dict) + # DSpark control keys are worker commands, not server args; route + # them to the draft worker and keep them out of the override. + remaining = dict(server_args_dict) + frac = remaining.pop("dspark_force_budget_frac", None) + if "dspark_force_budget_frac" in server_args_dict: + self.draft_worker.set_dspark_forced_budget_frac( + None if frac is None else float(frac) + ) + if remaining.pop("dspark_clear_info_records", None): + self.draft_worker.clear_info_records() + if remaining: + get_server_args().override(source="update_server_args", **remaining) logger.info(f"Global server args updated! {get_server_args()=}") server_args = dict(vars(get_server_args())) 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 8d0888064..8207f4b2c 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -18,6 +18,7 @@ from sglang.srt.environ import envs from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.schedule_batch import ( FINISH_ABORT, + FINISH_MATCHED_TOKEN, Req, ScheduleBatch, ) @@ -26,6 +27,7 @@ from sglang.srt.mem_cache.common import ( release_kv_cache, ) from sglang.srt.runtime_context import get_server_args +from sglang.srt.speculative.base_spec_worker import BaseSpecWorker from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer @@ -542,6 +544,17 @@ class SchedulerBatchResultProcessor: result.num_correct_drafts = sum(accept_lens) - len(batch.reqs) result.num_correct_drafts_per_req_cpu = [x - 1 for x in accept_lens] + block_accept_lens = ( + result.block_accept_lens.tolist() + if result.block_accept_lens is not None + else None + ) + result.num_block_accept_tokens = ( + sum(block_accept_lens) if block_accept_lens else 0 + ) + cap_lens = result.cap_lens.tolist() if result.cap_lens is not None else None + result.num_cap_tokens = sum(cap_lens) if cap_lens else 0 + # Feed the adaptive controller now that accept_lens is on CPU, # instead of doing a synchronous GPU→CPU copy in the worker hot path. # BaseSpecWorker provides a no-op default for non-adaptive workers. @@ -579,6 +592,12 @@ class SchedulerBatchResultProcessor: req.spec_num_correct_drafts += num_correct_drafts req.update_spec_correct_drafts_histogram(num_correct_drafts) + if block_accept_lens is not None: + req.spec_num_block_accept_tokens += block_accept_lens[i] + if cap_lens is not None: + req.spec_num_cap_tokens += cap_lens[i] + req.update_spec_cap_lens_histogram(cap_lens[i]) + predict_tokens.append(accept_tokens) return predict_tokens @@ -656,7 +675,10 @@ class SchedulerBatchResultProcessor: self.metrics_reporter.num_generated_tokens += len(batch.reqs) if not batch.spec_algorithm.is_none(): self.metrics_reporter.update_spec_metrics( - batch.batch_size(), result.num_correct_drafts + batch.batch_size(), + result.num_correct_drafts, + num_block_accept_tokens=result.num_block_accept_tokens, + num_cap_tokens=result.num_cap_tokens, ) if self.server_args.enable_metrics: self.metrics_collector.increment_decode_cuda_graph_pass( @@ -828,6 +850,14 @@ class SchedulerBatchResultProcessor: self.decode_offload_manager.offload_kv_cache(req) if req.finished(): + # isinstance narrowing: create_worker may also return plain + # TpModelWorker-based drafts, which carry no spec-worker hooks. + if isinstance(self.draft_worker, BaseSpecWorker): + self.draft_worker.note_request_finished( + rid=req.rid, + natural_stop=isinstance(req.finished_reason, FINISH_MATCHED_TOKEN), + ) + # delete feature to save memory if req.multimodal_inputs is not None and req.session is None: req.multimodal_inputs.release_features() diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index f82face7c..7d6d767ef 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -139,6 +139,8 @@ class SchedulerMetricsReporter: self.spec_num_forward_ct = 0 self.spec_total_num_accept_tokens = 0 # lifetime self.spec_total_num_forward_ct = 0 + self.spec_num_block_accept_tokens = 0 + self.spec_num_cap_tokens = 0 # For PD disaggregation self.kv_transfer_speed_gb_s: float = 0.0 @@ -348,9 +350,17 @@ class SchedulerMetricsReporter: "num_draft_tokens": num_draft_tokens or 0, } - def update_spec_metrics(self, bs: int, num_correct_drafts: int): + def update_spec_metrics( + self, + bs: int, + num_correct_drafts: int, + num_block_accept_tokens: int = 0, + num_cap_tokens: int = 0, + ): self.spec_num_accept_tokens += num_correct_drafts + bs 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 @@ -510,6 +520,8 @@ class SchedulerMetricsReporter: self.spec_num_forward_ct = 0 self.spec_total_num_accept_tokens = 0 self.spec_total_num_forward_ct = 0 + self.spec_num_block_accept_tokens = 0 + self.spec_num_cap_tokens = 0 def report_prefill_stats( self, @@ -733,6 +745,8 @@ class SchedulerMetricsReporter: if self.scheduler.spec_algorithm.is_none(): spec_accept_length = 0 spec_accept_rate = 0 + spec_cap_length = 0 + 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 @@ -746,10 +760,38 @@ class SchedulerMetricsReporter: spec_accept_rate = ( num_correct_drafts / total_draft_tokens if total_draft_tokens > 0 else 0 ) + spec_cap_length = ( + self.spec_num_cap_tokens / self.spec_num_forward_ct + if self.spec_num_forward_ct > 0 + else 0 + ) + from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyMode, + read_ragged_verify_mode, + ) + + spec_block_accept_length = ( + self.spec_num_block_accept_tokens / self.spec_num_forward_ct + if self.spec_num_forward_ct > 0 + and read_ragged_verify_mode() is RaggedVerifyMode.CAP_ACCEPT + else 0 + ) 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_block_accept_tokens = 0 + self.spec_num_cap_tokens = 0 msg += f"accept len: {spec_accept_length:.2f}, accept rate: {spec_accept_rate:.2f}, " + if spec_cap_length > 0: + msg += f"cap len: {spec_cap_length:.2f}, " + if spec_block_accept_length > 0: + msg += f"block accept len: {spec_block_accept_length:.2f}, " + if self.scheduler.spec_algorithm.is_dspark(): + draft_worker = self.scheduler.draft_worker + if draft_worker is not None: + estimate_suffix = draft_worker.block_accept_estimate_log_suffix() + if estimate_suffix: + msg += f"{estimate_suffix}, " if self.current_scheduler_metrics_enabled: spec_snapshot = self._active_spec_config_snapshot() @@ -825,6 +867,8 @@ class SchedulerMetricsReporter: # Speculative decoding self.stats.spec_accept_length = spec_accept_length self.stats.spec_accept_rate = spec_accept_rate + self.stats.spec_cap_length = spec_cap_length + self.stats.spec_block_accept_length = spec_block_accept_length self.stats.spec_num_steps = spec_num_steps self.stats.spec_num_draft_tokens = spec_num_draft_tokens diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 6b691d84e..f33118a15 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -278,7 +278,10 @@ class _GenerationStreamAccumulator: video_tokens: list = field(default_factory=list) spec_verify_ct: list = field(default_factory=list) spec_num_correct_drafts: list = field(default_factory=list) + spec_num_block_accept_tokens: list = field(default_factory=list) + spec_num_cap_tokens: list = field(default_factory=list) spec_correct_drafts_histogram: list = field(default_factory=list) + spec_cap_lens_histogram: list = field(default_factory=list) retraction_counts: list = field(default_factory=list) output_hidden_states: Optional[list] = None routed_experts: Optional[list] = None @@ -406,7 +409,10 @@ class _GenerationStreamAccumulator: if not self.spec_algorithm.is_none(): self.spec_verify_ct.append(req.spec_verify_ct) self.spec_num_correct_drafts.append(req.spec_num_correct_drafts) + self.spec_num_block_accept_tokens.append(req.spec_num_block_accept_tokens) + self.spec_num_cap_tokens.append(req.spec_num_cap_tokens) self.spec_correct_drafts_histogram.append(req.spec_correct_drafts_histogram) + self.spec_cap_lens_histogram.append(req.spec_cap_lens_histogram) if self.return_logprob: if ( @@ -528,7 +534,10 @@ class _GenerationStreamAccumulator: http_worker_ipcs=self.http_worker_ipcs, spec_verify_ct=self.spec_verify_ct, spec_num_correct_drafts=self.spec_num_correct_drafts, + spec_num_block_accept_tokens=self.spec_num_block_accept_tokens, + spec_num_cap_tokens=self.spec_num_cap_tokens, spec_correct_drafts_histogram=self.spec_correct_drafts_histogram, + spec_cap_lens_histogram=self.spec_cap_lens_histogram, time_stats=wrap_as_pickle(self.time_stats), finished_reasons=self.finished_reasons, decoded_texts=self.decoded_texts, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 04d40a474..b22c2f6f7 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -32,6 +32,7 @@ from collections import deque from contextlib import nullcontext from datetime import datetime from enum import Enum +from functools import lru_cache from http import HTTPStatus from typing import Any, Awaitable, Dict, Iterable, List, Optional, Tuple, Union @@ -142,6 +143,19 @@ _REQUEST_STATE_WAIT_TIMEOUT = envs.SGLANG_REQUEST_STATE_WAIT_TIMEOUT.get() logger = logging.getLogger(__name__) + +@lru_cache(maxsize=1) +def _ragged_verify_cap_accept() -> bool: + # The mode env is fixed at server launch; cache to keep it off the + # per-request metrics path. + from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyMode, + read_ragged_verify_mode, + ) + + return read_ragged_verify_mode() is RaggedVerifyMode.CAP_ACCEPT + + _INCREMENTAL_STREAMING_META_INFO_KEYS = ( "output_token_logprobs", "output_top_logprobs", @@ -2366,6 +2380,25 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): meta_info["spec_num_proposed_drafts"] = num_proposed_drafts meta_info["spec_verify_ct"] = recv_obj.spec_verify_ct[i] + if ( + getattr(recv_obj, "spec_num_cap_tokens", None) is not None + and len(recv_obj.spec_num_cap_tokens) > i + and recv_obj.spec_num_cap_tokens[i] > 0 + ): + meta_info["spec_cap_length"] = ( + recv_obj.spec_num_cap_tokens[i] / recv_obj.spec_verify_ct[i] + ) + if ( + _ragged_verify_cap_accept() + and getattr(recv_obj, "spec_num_block_accept_tokens", None) + is not None + and len(recv_obj.spec_num_block_accept_tokens) > i + ): + meta_info["spec_block_accept_length"] = ( + recv_obj.spec_num_block_accept_tokens[i] + / recv_obj.spec_verify_ct[i] + ) + # FIXME: backward-compat aliases, remove in next release. meta_info["spec_accepted_drafts"] = num_correct_drafts meta_info["spec_proposed_drafts"] = num_proposed_drafts @@ -2383,6 +2416,14 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): meta_info["spec_accept_histogram"] = ( recv_obj.spec_correct_drafts_histogram[i] ) + if ( + getattr(recv_obj, "spec_cap_lens_histogram", None) + and len(recv_obj.spec_cap_lens_histogram) > i + and recv_obj.spec_cap_lens_histogram[i] + ): + meta_info["spec_cap_lens_histogram"] = recv_obj.spec_cap_lens_histogram[ + i + ] def _request_has_grammar(self, obj: GenerateReqInput) -> bool: return ( diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index c932de486..40ad38891 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -44,6 +44,8 @@ class GenerationBatchResult: ] = None num_correct_drafts: int = 0 # no bonus included num_correct_drafts_per_req_cpu: Optional[List[int]] = None + num_block_accept_tokens: int = 0 + num_cap_tokens: int = 0 # FDFO dLLM batching: per-request accepted block length and carried algo state. accept_length_per_req_cpu: Optional[List[int]] = None dllm_algo_state: Optional[List[Any]] = None @@ -68,6 +70,10 @@ class GenerationBatchResult: # sync path: forward stream -> output processor accept_lens: Optional[torch.Tensor] = None + block_accept_lens: Optional[torch.Tensor] = None + + cap_lens: Optional[torch.Tensor] = None + # Next-iter seq_lens; published via on_publish. new_seq_lens: Optional[torch.Tensor] = None @@ -135,6 +141,12 @@ class GenerationBatchResult: if self.accept_lens is not None: self.accept_lens = _async_d2h(self.accept_lens) + if self.block_accept_lens is not None: + self.block_accept_lens = _async_d2h(self.block_accept_lens) + + if self.cap_lens is not None: + self.cap_lens = _async_d2h(self.cap_lens) + # Sub-objects only declare their device fields; the single copy+safety # primitive (_async_d2h: pinned D2H + record_stream) is injected here so # all device->host copying and lifetime safety lives in one place. diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 3da49c6a4..7f4d5810e 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -264,6 +264,7 @@ class ReqToTokenPool: (self._alloc_size, max_context_len), dtype=torch.int32, device=device ) self.free_slots = list(range(1, self._alloc_size)) + self.req_generation = torch.zeros(self._alloc_size, dtype=torch.int64) def write(self, indices, values): self.req_to_token[indices] = values @@ -295,6 +296,7 @@ class ReqToTokenPool: for r in reqs: if r.req_pool_idx is None: r.req_pool_idx = select_index[offset] + self.req_generation[r.req_pool_idx] += 1 offset += 1 return [r.req_pool_idx for r in reqs] @@ -305,6 +307,7 @@ class ReqToTokenPool: def clear(self): self.free_slots = list(range(1, self._alloc_size)) + self.req_generation.zero_() class MambaPool: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 64be70903..1ffad3397 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -423,9 +423,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): # auxiliary hidden capture mode. TODO: expose this to server args? self.eagle_use_aux_hidden_state = False self.eagle_draft_num_layers = None - self.dflash_use_aux_hidden_state = False - self.dflash_target_layer_ids = None - self.dflash_draft_num_layers = None + self.dflash_family_use_aux_hidden_state = False + self.dflash_family_target_layer_ids = None + self.dflash_family_draft_num_layers = None if ( (self.spec_algorithm.is_eagle() or self.spec_algorithm.is_standalone()) and not self.is_draft_worker @@ -465,10 +465,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): # if there is no aux layer, set to None self.eagle_aux_hidden_state_layer_ids = None - if self.spec_algorithm.is_dflash() and not self.is_draft_worker: + if self.spec_algorithm.is_dflash_family() and not self.is_draft_worker: from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config - # Select target layers to capture for building DFlash context features. + # Select target layers to capture for building draft context features. draft_model_config = self._build_model_config( server_args, model_path=(server_args.speculative_draft_model_path), @@ -486,8 +486,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) if target_num_layers is None: raise ValueError( - "DFLASH requires target num_hidden_layers in config. " - f"Got target={target_num_layers}." + "Block-draft-with-target-kv spec requires target num_hidden_layers " + f"in config. Got target={target_num_layers}." ) target_num_layers = int(target_num_layers) @@ -496,19 +496,37 @@ class ModelRunner(ModelRunnerKVCacheMixin): and trained_target_layers != target_num_layers ): logger.warning( - "DFLASH draft config num_target_layers=%s differs from runtime target num_hidden_layers=%s; " + "Draft config num_target_layers=%s differs from runtime target num_hidden_layers=%s; " "selecting capture layers based on the runtime target model.", trained_target_layers, target_num_layers, ) - self.dflash_use_aux_hidden_state = True - self.dflash_draft_num_layers = int(draft_num_layers) - self.dflash_target_layer_ids = dflash_draft_config.resolve_target_layer_ids( + target_layer_ids = dflash_draft_config.resolve_target_layer_ids( target_num_layers=int(target_num_layers), draft_num_layers=int(draft_num_layers), ) + if self.spec_algorithm.is_dspark(): + from sglang.srt.speculative.dspark_components.dspark_config import ( + parse_dspark_draft_config, + ) + + dspark_draft_config = parse_dspark_draft_config( + draft_hf_config=draft_model_config.hf_config + ) + if not dspark_draft_config.require_markov(): + raise ValueError( + "DSPARK requires markov_rank > 0 in the draft config, " + f"got markov_rank={dspark_draft_config.markov_rank}." + ) + if dspark_draft_config.target_layer_ids is not None: + target_layer_ids = list(dspark_draft_config.target_layer_ids) + + self.dflash_family_use_aux_hidden_state = True + self.dflash_family_draft_num_layers = int(draft_num_layers) + self.dflash_family_target_layer_ids = target_layer_ids + # Apply the rank zero filter to logger if server_args.show_time_cost: enable_show_time_cost() @@ -1057,13 +1075,23 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.model.set_eagle3_layers_to_capture( self.eagle_aux_hidden_state_layer_ids ) - if self.dflash_use_aux_hidden_state: - if not hasattr(self.model, "set_dflash_layers_to_capture"): + if self.dflash_family_use_aux_hidden_state: + if self.spec_algorithm.is_dspark() and hasattr( + self.model, "set_dspark_layers_to_capture" + ): + self.model.set_dspark_layers_to_capture( + self.dflash_family_target_layer_ids + ) + elif hasattr(self.model, "set_dflash_layers_to_capture"): + self.model.set_dflash_layers_to_capture( + self.dflash_family_target_layer_ids + ) + else: raise ValueError( - f"Model {self.model.__class__.__name__} does not implement " - "set_dflash_layers_to_capture, which is required for DFLASH." + f"Model {self.model.__class__.__name__} implements neither " + "set_dspark_layers_to_capture nor set_dflash_layers_to_capture, " + "one of which is required for DFLASH/DSPARK." ) - self.model.set_dflash_layers_to_capture(self.dflash_target_layer_ids) def remote_instance_init_transfer_engine(self): try: diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 4bc650ba5..18a8cb272 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -154,13 +154,13 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): * (1 + int(eagle_draft_num_layers) / int(num_layers)) ) - # DFLASH: scale cell_size to account for draft model KV cache - if mr.spec_algorithm.is_dflash() and not mr.is_draft_worker: + # DFLASH/DSPARK: scale cell_size to account for draft model KV cache + if mr.spec_algorithm.is_dflash_family() and not mr.is_draft_worker: from sglang.srt.speculative.dflash_utils import ( scale_kv_cell_size_per_token_for_dflash, ) - draft_num_layers = mr.dflash_draft_num_layers + draft_num_layers = mr.dflash_family_draft_num_layers if ( draft_num_layers is not None and int(draft_num_layers) > 0 diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 16c473715..8e8a1b6fa 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -42,6 +42,7 @@ from sglang.srt.distributed.parallel_state import ( set_pdmux_status, ) from sglang.srt.dllm.config import DllmConfig +from sglang.srt.environ import envs from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp from sglang.srt.layers.dp_attention import ( DpPaddingMode, @@ -91,12 +92,11 @@ from sglang.srt.model_executor.runner_utils.deepep_adapter import ( ) from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_stream_groups from sglang.srt.runtime_context import get_flags, get_parallel +from sglang.srt.speculative.ragged_verify import resolve_ragged_verify_layout from sglang.srt.utils import ( empty_context, get_available_gpu_memory, require_attn_tp_gather, - require_gathered_buffer, - require_mlp_sync, require_mlp_tp_gather, ) from sglang.srt.utils.profile_utils import export_cuda_graph_capture_trace @@ -112,6 +112,15 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + +def ragged_verify_compact_graphs_enabled(spec_algorithm: SpeculativeAlgorithm) -> bool: + if not spec_algorithm.supports_ragged_verify(): + return False + from sglang.srt.speculative.ragged_verify import ragged_verify_compact_enabled + + return ragged_verify_compact_enabled() def build_replay_fb_view( @@ -188,10 +197,19 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): self.enable_torch_compile = get_flags().capture.enable_torch_compile self.disable_padding = model_runner.server_args.disable_cuda_graph_padding self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder - self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args) - self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args) - self.require_mlp_sync = require_mlp_sync(model_runner.server_args) + self.require_mlp_tp_gather = require_mlp_tp_gather( + model_runner.server_args + ) and not self._forward_is_dp_local(model_runner) self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args) + # Composite predicates derive from the instance values so the dp-local + # draft exemption above stays consistent (require_gathered_buffer == + # mlp_tp_gather or attn_tp_gather; require_mlp_sync adds dp attention). + self.require_gathered_buffer = ( + self.require_mlp_tp_gather or self.require_attn_tp_gather + ) + self.require_mlp_sync = ( + model_runner.server_args.enable_dp_attention or self.require_gathered_buffer + ) self.enable_two_batch_overlap = ( model_runner.server_args.enable_two_batch_overlap ) @@ -257,6 +275,29 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): if KTRANSFORMERS_AVAILABLE: KTMoEWrapper.set_capture_batch_sizes(self.capture_bs) + self.ragged_verify_mode = ( + ragged_verify_compact_graphs_enabled(self.model_runner.spec_algorithm) + and (self.capture_forward_mode == ForwardMode.TARGET_VERIFY) + and not self.model_runner.is_draft_worker + ) + self.capture_num_tokens: Optional[list[int]] = ( + self._build_ragged_verify_token_buckets() + if self.ragged_verify_mode + else None + ) + self._ragged_graph_size = 0 + if self.ragged_verify_mode and ( + self.enable_two_batch_overlap + or model_runner.server_args.enable_lora + or self.disable_padding + ): + raise ValueError( + "Compact ragged verify does not support two-batch-overlap, " + "LoRA, or disable-cuda-graph-padding (bs pads to the captured " + "tier); disable SGLANG_RAGGED_VERIFY_MODE or the conflicting " + "feature." + ) + # If returning hidden states is enabled, set initial capture hidden mode to full to avoid double-capture on startup if self.enable_return_hidden_states: self.capture_hidden_mode = CaptureHiddenMode.FULL @@ -362,6 +403,11 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): f"Capture cuda graph failed: {e}\n" f"{CUDA_GRAPH_CAPTURE_FAILED_MSG}" ) + def _build_ragged_verify_token_buckets(self) -> list[int]: + buckets = sorted({bs * self.num_tokens_per_bs for bs in self.capture_bs}) + assert buckets and buckets[0] > 0, f"{buckets=}" + return buckets + def _autotune_buffers(self): """Reuse these static decode buffers (sized to max_bs) for the warmup flashinfer-autotune dummy forward instead of allocating a throwaway set @@ -385,13 +431,16 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): def _cache_loc_dtype(self): return torch.int64 - def _make_graph_key(self, bs, stream_idx=None, variant_label=None): + def _make_graph_key(self, size, stream_idx=None, variant_label=None): return ShapeKey( - size=bs, + size=size, stream_idx=stream_idx, variant_label=variant_label, ) + def _capture_graph_size(self, *, bs: int, num_tokens: int) -> int: + return num_tokens if self.ragged_verify_mode else bs + def _resolve_lora_variant(self, forward_batch: ForwardBatch): if not getattr(self, "record_nolora_graph", False): return None @@ -401,16 +450,69 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): return "lora" return "nolora" + @staticmethod + def _forward_is_dp_local(model_runner) -> bool: + """The DSpark dense draft runs attn-TP-local (draft_tp_context): each + DP rank drafts independently with no cross-DP collective, so its + hand-built batches carry no dp-global metadata and must key graphs by + local batch size. Everything else keeps the dp-global padding path.""" + if not model_runner.is_draft_worker: + return False + if not model_runner.spec_algorithm.is_dspark(): + return False + from sglang.srt.speculative.dspark_components.dspark_config import ( + draft_is_deepseek_v4, + ) + + return not draft_is_deepseek_v4(server_args=model_runner.server_args) + + def _ragged_capture_slots(self, num_tokens: int) -> int: + if envs.SGLANG_TEST_RAGGED_VERIFY_FORCE_UNIFORM_CAPTURE.get(): + return num_tokens // self.num_tokens_per_bs + return min(num_tokens, self.max_bs) + + def _capture_ragged_verify_layout(self, num_tokens: int): + if not self.ragged_verify_mode: + return None + if envs.SGLANG_TEST_RAGGED_VERIFY_FORCE_UNIFORM_CAPTURE.get(): + return None + from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyLayout, + build_capture_verify_lens, + ) + + verify_lens_cpu = build_capture_verify_lens( + num_tokens=num_tokens, + num_slots=self._ragged_capture_slots(num_tokens), + num_draft_tokens=self.num_tokens_per_bs, + ) + return RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=verify_lens_cpu, + device=self.device, + grid=self.capture_num_tokens, + ) + def can_run_graph(self, forward_batch: ForwardBatch): # Disable for token embedding overrides (dynamic per-request) if forward_batch.replace_embeds is not None: return False + + ragged_layout = ( + resolve_ragged_verify_layout(forward_batch) + if self.ragged_verify_mode + else None + ) + if ragged_layout is not None: + return self._can_run_ragged_verify_graph(forward_batch, ragged_layout) + if self.ragged_verify_mode and forward_batch.forward_mode.is_target_verify(): + return False + if self.require_mlp_tp_gather: cuda_graph_bs = ( max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs if self.model_runner.spec_algorithm.is_eagle() or self.model_runner.spec_algorithm.is_standalone() - or self.model_runner.spec_algorithm.is_dflash() + or self.model_runner.spec_algorithm.is_dflash_family() else max(forward_batch.global_num_tokens_cpu) ) else: @@ -472,6 +574,46 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): and is_ngram_supported ) + def _can_run_ragged_verify_graph(self, forward_batch: ForwardBatch, ragged_layout): + if not self.attn_backend.supports_ragged_verify_graph: + return False + + admission_tokens = ragged_layout.graph_num_tokens + is_tokens_supported = admission_tokens <= self.capture_num_tokens[ + -1 + ] and forward_batch.batch_size <= self._ragged_capture_slots(admission_tokens) + + is_dp_supported = ( + forward_batch.can_run_dp_cuda_graph if self.require_mlp_sync else True + ) + + is_encoder_lens_supported = ( + torch.all(forward_batch.encoder_lens > 0) + if self.is_encoder_decoder + else True + ) + + requested_capture_hidden_mode = max( + forward_batch.capture_hidden_mode, + ( + forward_batch.spec_info.capture_hidden_mode + if getattr(forward_batch.spec_info, "capture_hidden_mode", None) + is not None + else CaptureHiddenMode.NULL + ), + ) + capture_hidden_mode_matches = ( + requested_capture_hidden_mode == CaptureHiddenMode.NULL + or requested_capture_hidden_mode == self.capture_hidden_mode + ) + + return ( + is_tokens_supported + and is_dp_supported + and is_encoder_lens_supported + and capture_hidden_mode_matches + ) + def _init_profile_context_and_memory_record(self): profile_context = profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], @@ -510,17 +652,22 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): self, size: int, stream_idx: Optional[int] = None, + num_tokens: Optional[int] = None, ): """Build the dummy decode ForwardBatch for capture at size (=bs), populate static input buffers, choose the active attn backend, and optionally build pp_proxy_tensors. + num_tokens defaults to the uniform bs * num_tokens_per_bs; ragged + verify capture passes the decoupled (slots, tier tokens) pair. + Returns (forward_batch, attn_backend, pp_proxy_tensors); pp_proxy_tensors is None unless pp_size > 1. """ bs = size buffers: DecodeInputBuffers = self.buffers - num_tokens = bs * self.num_tokens_per_bs + if num_tokens is None: + num_tokens = bs * self.num_tokens_per_bs # Registry-owned FB-shared slots come through the registry (which # shares physical storage with self.buffers via source=...); the rest @@ -754,8 +901,8 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): stream_idx: Optional[int] = None, variant_label: Optional[str] = None, ): - bs = size - num_tokens = bs * self.num_tokens_per_bs + num_tokens = size * self.num_tokens_per_bs + bs = self._ragged_capture_slots(num_tokens) if self.ragged_verify_mode else size # Sanity-check: --debug-cuda-graph requires breakable backend. if self.model_runner.server_args.debug_cuda_graph: @@ -764,7 +911,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): ), "Breakable CUDA graph is required for --debug-cuda-graph" forward_batch, attn_backend, pp_proxy_tensors = self.capture_prepare( - size, stream_idx=stream_idx + bs, stream_idx=stream_idx, num_tokens=num_tokens ) # All setup hooks below read get_attn_backend() (TboForwardBatchPreparer, @@ -808,9 +955,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): {k: v.clone() for k, v in pp_proxy_tensors.tensors.items()} ) if ( - self.model_runner.spec_algorithm.is_dflash() + self.model_runner.spec_algorithm.is_dflash_family() and self.model_runner.is_draft_worker and "input_embeds" in inspect.signature(forward).parameters + and not hasattr(self.model_runner.model, "forward_embed") ): kwargs["input_embeds"] = self.buffers.input_embeds[:num_tokens] @@ -835,7 +983,11 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): # wires no buffer here. (SWA write loc rides the `swa_out_cache_loc` rail.) with canary_ctx: - shape_key = self._make_graph_key(bs, stream_idx, variant_label) + shape_key = self._make_graph_key( + self._capture_graph_size(bs=bs, num_tokens=num_tokens), + stream_idx, + variant_label, + ) post_warmup_hook = getattr( self.model_runner.attn_backend, "on_after_cuda_graph_warmup", @@ -892,15 +1044,35 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): forward_batch: ForwardBatch, pp_proxy_tensors: Optional[PPProxyTensors] = None, ): + ragged_layout = ( + resolve_ragged_verify_layout(forward_batch) + if self.ragged_verify_mode + else None + ) + is_ragged = ragged_layout is not None + self.deepep_adapter.replay() if not forward_batch.needs_forward_metadata_init(): # Pre-planned (plan-stream load_batch already ran). # In speculative decoding, these two fields are still needed. + graph_size_key = ( + self._ragged_graph_size + if is_ragged + else self._capture_graph_size( + bs=self.bs, num_tokens=self.bs * self.num_tokens_per_bs + ) + ) + if is_ragged: + assert self.raw_num_token == ragged_layout.graph_num_tokens, ( + f"stale ragged raw_num_token {self.raw_num_token} != " + f"{ragged_layout.graph_num_tokens}" + ) self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) if ( - self.model_runner.spec_algorithm.is_dflash() + not is_ragged + and self.model_runner.spec_algorithm.is_dflash_family() and self.model_runner.is_draft_worker and forward_batch.input_embeds is not None ): @@ -910,7 +1082,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): variant_label = self._resolve_lora_variant(forward_batch) stream_idx = get_current_stream_idx() if self.enable_pdmux else None self._replay_graph_key = self._make_graph_key( - self.bs, stream_idx, variant_label + graph_size_key, stream_idx, variant_label ) return @@ -918,37 +1090,57 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): self.recapture_if_needed(forward_batch) raw_bs = forward_batch.batch_size - raw_num_token = raw_bs * self.num_tokens_per_bs - if self.require_mlp_tp_gather: - max_num_tokens = max(forward_batch.global_num_tokens_cpu) - max_batch_size = ( - max_num_tokens / self.num_tokens_per_bs - if self.model_runner.spec_algorithm.is_eagle() - or self.model_runner.spec_algorithm.is_standalone() - or self.model_runner.spec_algorithm.is_dflash() - else max_num_tokens + if is_ragged: + raw_num_token = ragged_layout.graph_num_tokens + graph_size_key = self._ragged_graph_num_tokens(raw_num_token) + assert graph_size_key == ragged_layout.graph_num_tokens, ( + f"ragged verify tier mismatch: runner tier {graph_size_key} != " + f"layout graph_num_tokens {ragged_layout.graph_num_tokens}" ) - bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs) + bs = self._ragged_capture_slots(graph_size_key) + assert bs >= raw_bs, ( + f"ragged capture slots {bs} (tier {graph_size_key}) < raw_bs " + f"{raw_bs}; the planner must reject this batch before replay" + ) + padded_num_tokens = graph_size_key else: - bs = self._pad_to_bucket(raw_bs, self.capture_bs) + raw_num_token = raw_bs * self.num_tokens_per_bs + if self.require_mlp_tp_gather: + max_num_tokens = max(forward_batch.global_num_tokens_cpu) + max_batch_size = ( + max_num_tokens / self.num_tokens_per_bs + if self.model_runner.spec_algorithm.is_eagle() + or self.model_runner.spec_algorithm.is_standalone() + or self.model_runner.spec_algorithm.is_dflash_family() + else max_num_tokens + ) + bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs) + else: + bs = self._pad_to_bucket(raw_bs, self.capture_bs) + padded_num_tokens = bs * self.num_tokens_per_bs + graph_size_key = self._capture_graph_size( + bs=bs, num_tokens=padded_num_tokens + ) self.buffer_registry.fill_from( forward_batch, raw_bs=raw_bs, padded_bs=bs, raw_num_tokens=raw_num_token, - padded_num_tokens=bs * self.num_tokens_per_bs, + padded_num_tokens=padded_num_tokens, pp_proxy_tensors=pp_proxy_tensors, ) if ( - self.model_runner.spec_algorithm.is_dflash() + not is_ragged + and self.model_runner.spec_algorithm.is_dflash_family() and self.model_runner.is_draft_worker and forward_batch.input_embeds is not None ): buffers.input_embeds[:raw_num_token].copy_(forward_batch.input_embeds) - # Padded tokens aren't read, so skip zeroing them. + # Padded tokens aren't read, so skip zeroing. Ragged input_ids arrive + # from the planner already padded to the tier, invalid slots zeroed. if self.enable_two_batch_overlap: self.tbo_plugin.replay_prepare( forward_mode=self.capture_forward_mode, @@ -956,7 +1148,11 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): num_token_non_padded=len(forward_batch.input_ids), spec_info=forward_batch.spec_info, ) - if forward_batch.forward_mode.is_idle() and forward_batch.spec_info is not None: + if ( + not is_ragged + and forward_batch.forward_mode.is_idle() + and forward_batch.spec_info is not None + ): forward_batch.spec_info.custom_mask = buffers.custom_mask if self.enable_pdmux: stream_idx = get_current_stream_idx() @@ -968,17 +1164,18 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): buffers=buffers, bs=bs, raw_bs=raw_bs, - num_tokens=bs * self.num_tokens_per_bs, + num_tokens=padded_num_tokens, seq_len_fill_value=self.seq_len_fill_value, capture_forward_mode=self.capture_forward_mode, is_encoder_decoder=self.is_encoder_decoder, ) attn_backend.init_forward_metadata_out_graph(fb_view) - # Store fields self.raw_bs = raw_bs self.raw_num_token = raw_num_token self.bs = bs + if is_ragged: + self._ragged_graph_size = graph_size_key if self.model_runner.hisparse_coordinator is not None: self.model_runner.hisparse_coordinator.num_real_reqs.fill_(raw_bs) @@ -986,9 +1183,14 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): variant_label = self._resolve_lora_variant(forward_batch) stream_idx = get_current_stream_idx() if self.enable_pdmux else None self._replay_graph_key = self._make_graph_key( - self.bs, stream_idx, variant_label + graph_size_key, stream_idx, variant_label ) + def _ragged_graph_num_tokens(self, total_verify_tokens: int) -> int: + from sglang.srt.speculative.ragged_verify import round_up_grid + + return round_up_grid(total_verify_tokens, self.capture_num_tokens) + def execute( self, forward_batch: ForwardBatch, @@ -1001,19 +1203,46 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): if self.model_runner.device_timer else contextlib.nullcontext() ) + # Publish a read-done event for the WAR barrier: a cuda-graph forward + # finishes its shared req_to_token / SWA reads at this pre-replay + # snapshot, so plain DECODE and block-draft TARGET_VERIFY qualify. + publish_read_done = forward_batch.forward_mode.is_decode() or ( + forward_batch.forward_mode.is_target_verify() + and self.model_runner.spec_algorithm.is_dflash_family() + ) + # Exception: breakable-graph verify replays (captured forward metadata) + # re-read req_to_token *during* replay, so the pre-replay snapshot is + # too early -- record the event after replay instead. + read_done_post_replay = ( + publish_read_done + and forward_batch.forward_mode.is_target_verify() + and self.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph + ) with timer_ctx, self.backend.replay_session(): self.load_batch(forward_batch, pp_proxy_tensors) - # Publish a read-done event for the WAR barrier: a cuda-graph forward - # finishes its shared req_to_token / SWA reads at this pre-replay - # snapshot, so plain DECODE and DFLASH TARGET_VERIFY both qualify. - if forward_batch.forward_mode.is_decode() or ( - forward_batch.forward_mode.is_target_verify() - and self.model_runner.spec_algorithm.is_dflash() - ): + if envs.SGLANG_LOG_DECODE_GRAPH_KEY.get(): + logger.info( + "Decode graph replay: worker=%s key_size=%s (%s) mode=%s raw_bs=%d%s", + "draft" if self.model_runner.is_draft_worker else "target", + self._replay_graph_key.size, + "num_tokens" if self.ragged_verify_mode else "bs", + forward_batch.forward_mode.name, + forward_batch.batch_size, + ( + f" slots={self._ragged_capture_slots(self._replay_graph_key.size)}" + if self.ragged_verify_mode + else "" + ), + ) + if publish_read_done and not read_done_post_replay: read_done = self.device_module.Event() read_done.record() self.model_runner.war_fastpath_read_done_event = read_done output = self.backend.replay(self._replay_graph_key, forward_batch) + if read_done_post_replay: + read_done = self.device_module.Event() + read_done.record() + self.model_runner.war_fastpath_read_done_event = read_done if isinstance(output, LogitsProcessorOutput): if self.is_dllm: @@ -1083,7 +1312,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): dtype=self.model_runner.dtype, device=self.model_runner.device, ) - elif self.model_runner.spec_algorithm.is_dflash(): + elif self.model_runner.spec_algorithm.is_dflash_family(): from sglang.srt.speculative.dflash_info import DFlashVerifyInput from sglang.srt.speculative.dflash_utils import ( resolve_dflash_verify_mask_policy, @@ -1097,7 +1326,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): spec_info = DFlashVerifyInput( draft_token=None, positions=None, - draft_token_num=self.model_runner.server_args.speculative_num_draft_tokens, + draft_token_num=self.num_tokens_per_bs, custom_mask=( None if (self.model_runner.is_draft_worker or not build_custom_mask) @@ -1108,6 +1337,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): if self.model_runner.is_draft_worker else CaptureHiddenMode.FULL ), + ragged_verify_layout=self._capture_ragged_verify_layout(num_tokens), ) elif self.model_runner.spec_algorithm.is_ngram(): diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 94970e151..d658b5c59 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -182,7 +182,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # Ported from main #27468. if ( model_runner.server_args.enable_return_hidden_states - or model_runner.spec_algorithm.is_dflash() + or model_runner.spec_algorithm.is_dflash_family() ): self.capture_hidden_mode = CaptureHiddenMode.FULL # EAGLE captures FULL hidden states for the target and LAST for the diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index ae841d56e..5773395a0 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -2035,6 +2035,8 @@ class DeepseekV4Model(nn.Module): if self.dsa_enable_prefill_cp: self.cp_size = get_parallel().attn_cp_size + self.dspark_layers_to_capture: Optional[List[int]] = None + def get_input_embeddings(self) -> nn.Module: return self.embed_tokens @@ -2220,7 +2222,18 @@ class DeepseekV4Model(nn.Module): if hasattr(forward_batch, _attr): delattr(forward_batch, _attr) - if self._can_run_tbo(forward_batch): + capture_dspark = self.dspark_layers_to_capture is not None + if capture_dspark and dsa_use_prefill_cp(forward_batch): + raise NotImplementedError( + "DSpark aux hidden-state capture is not supported together with " + "DeepSeek-V4 prefill context parallelism (attn_cp_size > 1). Disable one " + "of them: DSpark static-verify is CP-off for v1." + ) + dspark_aux_hidden_states: List[torch.Tensor] = [] + # DSpark aux capture needs the per-layer eager loop (TBO's overlapped + # execution cannot expose per-layer completed hidden states), so skip + # TBO when capturing -- a perf-only downgrade, not a correctness one. + if self._can_run_tbo(forward_batch) and not capture_dspark: # Two-batch-overlap prefill (EP / mori). Cross-layer mHC fusion is # disabled here (each layer self-contained), so no trailing hc_post. hidden_states = self._forward_layers_tbo( @@ -2251,6 +2264,14 @@ class DeepseekV4Model(nn.Module): prev_post=prev_post, prev_comb=prev_comb, ) + if capture_dspark and i in self.dspark_layers_to_capture: + if use_fused: + completed = layer.hc_post( + hidden_states, prev_residual, prev_post, prev_comb + ) + else: + completed = hidden_states + dspark_aux_hidden_states.append(completed.mean(dim=1)) if use_fused and last_layer is not None: hidden_states = last_layer.hc_post( hidden_states, prev_residual, prev_post, prev_comb @@ -2276,6 +2297,9 @@ class DeepseekV4Model(nn.Module): ) hidden_states = self.norm(hidden_states) + if capture_dspark: + return (hidden_states, pre_hc_head), dspark_aux_hidden_states + return hidden_states, pre_hc_head @@ -2352,6 +2376,16 @@ class DeepseekV4ForCausalLM(nn.Module): def get_input_embeddings(self) -> nn.Module: return self.model.get_input_embeddings() + def set_dspark_layers_to_capture(self, layer_ids: List[int]) -> None: + if not self.pp_group.is_last_rank: + return + if layer_ids is None: + raise ValueError( + "DSPARK requires explicit layer_ids for aux hidden capture." + ) + self.capture_aux_hidden_states = True + self.model.dspark_layers_to_capture = list(layer_ids) + def determine_num_fused_shared_experts(self): self.num_fused_shared_experts = 0 if get_server_args().disable_shared_experts_fusion: @@ -2429,7 +2463,9 @@ class DeepseekV4ForCausalLM(nn.Module): self.lm_head, forward_batch, aux_hidden_states, - hidden_states_before_norm=pre_hc_head, + hidden_states_before_norm=( + None if aux_hidden_states is not None else pre_hc_head + ), ) def _setup_fp8_wo_a_scales(self, is_nextn: bool) -> None: diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py new file mode 100644 index 000000000..98a411afd --- /dev/null +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -0,0 +1,892 @@ +from __future__ import annotations + +import logging +from typing import Iterable, List, Optional, Tuple + +import msgspec +import torch +import torch.nn.functional as F +from torch import nn + +from sglang.jit_kernel.dsv4 import fused_q_norm_rope, fused_rope_inplace +from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config +from sglang.srt.environ import envs +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.forward_context import get_token_to_kv_pool +from sglang.srt.model_executor.runner import get_is_capture_mode +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.dbrx import ReplicatedLinear +from sglang.srt.models.deepseek_v4 import ( + DEEPSEEK_V4_STACKED_PARAMS_MAPPING, + DeepseekV4DecoderLayer, + MqaAttentionBase, + _dequant_fp8_wo_a, + hc_head_torch, + make_hc_head_params, +) +from sglang.srt.models.dspark import ( + DSparkConfidenceHead, + StepSampler, + gather_and_crop_vocab, + run_markov_block, +) +from sglang.srt.runtime_context import get_parallel +from sglang.srt.speculative.dspark_components.dspark_config import ( + parse_dspark_draft_config, +) +from sglang.srt.speculative.dspark_components.kernels.dspark_draft_model import ( + BuildStepLocal, + CommitKvProj, +) +from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyMode, + read_ragged_verify_mode, +) +from sglang.srt.utils import add_prefix, is_blackwell_supported +from sglang.srt.utils.async_probe import maybe_detect_in_closed_range + +logger = logging.getLogger(__name__) + +_PAD_NUM_HEADS = 64 + + +def apply_rotary_emb( + x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + y = x + x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if x.ndim == 3: + freqs_cis = freqs_cis.view(x.size(0), 1, x.size(-1)) + else: + freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) + x = torch.view_as_real(x * freqs_cis).flatten(-2) + y.copy_(x) + return y + + +class DSparkAttention(MqaAttentionBase): + + def __init__( + self, + config: DeepSeekV4Config, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_streams: Optional[List[torch.cuda.Stream]] = None, + ) -> None: + super().__init__( + config, + layer_id, + quant_config, + prefix, + attn_tp_rank=get_parallel().attn_tp_rank, + attn_tp_size=get_parallel().attn_tp_size, + compress_ratio=0, + fuse_wqa_wkv=False, + wo_a_fp8=False, + wo_a_keeps_quant_config=False, + wo_b_reduce_results=True, + rope_original_seq_len=0, + ) + assert ( + self.compress_ratio == 0 + ), "DSpark draft attention requires compress_ratio == 0." + self.window_size = int( + getattr(config, "sliding_window", None) or config.window_size + ) + + self.attn = RadixAttention( + self.n_local_heads, + self.head_dim, + self.softmax_scale, + num_kv_heads=1, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + self._use_fast_kernel = envs.SGLANG_DSPARK_FAST_KERNEL.get() + self.alt_streams = alt_streams + self._multi_stream_bs_limit = 128 if is_blackwell_supported() else 64 + + def kv_proj_only(self, x: torch.Tensor) -> torch.Tensor: + kv, _ = self.wkv(x) + return kv + + def _local_attn_sink(self) -> torch.Tensor: + if self.attn_tp_size == 1: + return self.attn_sink + if self._attn_sink_local is None: + rank = self.attn_tp_rank + num_heads = self.n_local_heads + sink = self.attn_sink.new_zeros(max(num_heads, _PAD_NUM_HEADS)) + sink[:num_heads] = self.attn_sink[rank * num_heads : (rank + 1) * num_heads] + self._attn_sink_local = sink + return self._attn_sink_local + + def _store_block_kv( + self, + *, + kv: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + attn_backend, + pool: DeepSeekV4TokenToKVPool, + ) -> None: + pool.set_swa_key_buffer_radix_fused_norm_rope( + layer_id=self.layer_id, + swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch), + kv=kv, + kv_weight=self.kv_norm.weight.data, + eps=self.eps, + freqs_cis=self.freqs_cis, + positions=positions, + ) + + def _compute_q( + self, + x: torch.Tensor, + positions: torch.Tensor, + q_out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + q, _ = self.wq_a(x) + q = self.q_norm(q) + q, _ = self.wq_b(q) + q = q.view(-1, self.n_local_heads, self.head_dim) + if self._use_fast_kernel: + if q_out is None: + q_out = torch.empty_like(q) + fused_q_norm_rope(q, q_out, self.eps, self.freqs_cis, positions) + return q_out + else: + q = q * torch.rsqrt( + q.float().square().mean(-1, keepdim=True) + self.eps + ).to(q.dtype) + apply_rotary_emb(q[..., -self.rope_head_dim :], self.freqs_cis[positions]) + if q_out is not None: + q_out.copy_(q) + return q_out + return q + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + from sglang.srt.model_executor.forward_context import get_attn_backend + + pool = _resolve_dspark_pool() + attn_backend = get_attn_backend() + rd = self.rope_head_dim + + enable_multi_stream = ( + self.alt_streams is not None + and get_is_capture_mode() + and hidden_states.shape[0] <= self._multi_stream_bs_limit + ) + + q_padded: Optional[torch.Tensor] = None + q_out: Optional[torch.Tensor] = None + if self.n_local_heads < _PAD_NUM_HEADS: + q_padded = hidden_states.new_empty( + hidden_states.shape[0], _PAD_NUM_HEADS, self.head_dim + ) + q_out = q_padded[:, : self.n_local_heads, :] + + if enable_multi_stream: + current_stream = torch.cuda.current_stream() + stream_kv = self.alt_streams[0] + stream_kv.wait_stream(current_stream) + with torch.cuda.stream(stream_kv): + kv = self.kv_proj_only(hidden_states) + self._store_block_kv( + kv=kv, + positions=positions, + forward_batch=forward_batch, + attn_backend=attn_backend, + pool=pool, + ) + q = self._compute_q(hidden_states, positions, q_out=q_out) + current_stream.wait_stream(stream_kv) + else: + kv = self.kv_proj_only(hidden_states) + self._store_block_kv( + kv=kv, + positions=positions, + forward_batch=forward_batch, + attn_backend=attn_backend, + pool=pool, + ) + q = self._compute_q(hidden_states, positions, q_out=q_out) + + if q_padded is not None: + q = q_padded + attn_sink = self._local_attn_sink() + + o = attn_backend.forward( + q=q, + k=kv, + v=kv, + layer=self.attn, + forward_batch=forward_batch, + compress_ratio=0, + attn_sink=attn_sink, + save_kv_cache=False, + ) + if o.shape[1] != self.n_local_heads: + o = o[:, : self.n_local_heads, :] + + if self._use_fast_kernel: + fused_rope_inplace( + o[..., -rd:], None, self.freqs_cis, positions=positions, inverse=True + ) + else: + apply_rotary_emb(o[..., -rd:], self.freqs_cis[positions], inverse=True) + + o = o.view( + o.shape[0], + self.n_local_groups, + o.shape[1] * o.shape[2] // self.n_local_groups, + ) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + if self._use_fast_kernel: + o = torch.einsum("bgd,grd->bgr", o, wo_a) + else: + o = torch.einsum("bgd,grd->bgr", o.float(), wo_a.float()).to(q.dtype) + out, _ = self.wo_b(o.reshape(o.shape[0], o.shape[1] * o.shape[2])) + return out + + +def _resolve_dspark_pool() -> DeepSeekV4TokenToKVPool: + pool = get_token_to_kv_pool() + assert isinstance(pool, DeepSeekV4TokenToKVPool), ( + "DSpark draft attention requires a DeepSeekV4TokenToKVPool, " + f"got {type(pool).__name__}." + ) + return pool + + +class MarkovW2ShardGeometry(msgspec.Struct, frozen=True): + + tp_size: int + org_vocab_start: int + org_vocab_end: int + num_embeddings_per_partition: int + num_embeddings_padded: int + + +class DSparkV4MarkovHead(nn.Module): + + markov_head_type = "vanilla" + + def __init__(self, *, vocab_size: int, markov_rank: int) -> None: + super().__init__() + self.vocab_size = int(vocab_size) + self.markov_rank = int(markov_rank) + if self.markov_rank <= 0: + raise ValueError( + f"DSparkV4MarkovHead requires markov_rank > 0, got {self.markov_rank}." + ) + self.markov_w1 = VocabParallelEmbedding( + self.vocab_size, self.markov_rank, enable_tp=False + ) + self._opt_markov_w2_bf16 = envs.SGLANG_DSPARK_OPT_MARKOV_W2_BF16.get() + self._opt_markov_w2_tp_shard = envs.SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD.get() + markov_w2_dtype = torch.bfloat16 if self._opt_markov_w2_bf16 else torch.float32 + self.markov_w2 = nn.Linear( + self.markov_rank, self.vocab_size, bias=False, dtype=markov_w2_dtype + ) + self._tp_shard: Optional[MarkovW2ShardGeometry] = None + + def configure_tp_shard(self, *, lm_head: nn.Module) -> None: + if not self._opt_markov_w2_tp_shard: + return + if int(lm_head.org_vocab_size) != self.vocab_size: + raise ValueError( + "DSpark markov_w2 TP-shard requires lm_head.org_vocab_size == " + f"markov vocab_size, got {int(lm_head.org_vocab_size)} vs " + f"{self.vocab_size}." + ) + tp_size = int(lm_head.tp_size) + per_partition = int(lm_head.num_embeddings_per_partition) + num_padded = int(lm_head.num_embeddings_padded) + if per_partition * tp_size != num_padded: + raise ValueError( + "DSpark markov_w2 TP-shard could not align to the lm_head partition: " + f"num_embeddings_per_partition({per_partition}) * tp_size({tp_size}) != " + f"num_embeddings_padded({num_padded})." + ) + attn_tp_size = get_parallel().attn_tp_group.world_size + if attn_tp_size != tp_size: + raise ValueError( + "DSpark markov_w2 TP-shard needs the attn-TP group (used for the per-step " + f"all-gather) to equal the lm_head shard group, got attn_tp_size=" + f"{attn_tp_size} vs lm_head tp_size={tp_size}. This config (e.g. DP " + "attention without --enable-dp-lm-head, where lm_head shards over the " + "global TP group) is unsupported; disable " + "SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD." + ) + self._tp_shard = MarkovW2ShardGeometry( + tp_size=tp_size, + org_vocab_start=int(lm_head.shard_indices.org_vocab_start_index), + org_vocab_end=int(lm_head.shard_indices.org_vocab_end_index), + num_embeddings_per_partition=per_partition, + num_embeddings_padded=num_padded, + ) + + def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.markov_w1(token_ids.long()) + + def project_bias( + self, latent_states: torch.Tensor, *, weight: Optional[torch.Tensor] = None + ) -> torch.Tensor: + weight = self.markov_w2.weight if weight is None else weight + if self._opt_markov_w2_bf16: + return F.linear(latent_states.to(weight.dtype), weight).float() + return F.linear(latent_states.float(), weight) + + def compute_step_bias( + self, token_ids: torch.Tensor, hidden_states: Optional[torch.Tensor] + ) -> torch.Tensor: + del hidden_states + return self.project_bias(self.get_prev_embeddings(token_ids)) + + def apply_step_logits( + self, + logits: torch.Tensor, + *, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + ) -> torch.Tensor: + if self._tp_shard is not None: + return self._apply_step_logits_sharded( + base_local=logits, token_ids=token_ids + ) + return logits + self.compute_step_bias(token_ids, hidden_states) + + def _apply_step_logits_sharded( + self, *, base_local: torch.Tensor, token_ids: torch.Tensor + ) -> torch.Tensor: + shard = self._tp_shard + latent = self.get_prev_embeddings(token_ids) + weight_local = self.markov_w2.weight[ + shard.org_vocab_start : shard.org_vocab_end + ] + if self._opt_markov_w2_bf16: + bias = F.linear(latent.to(weight_local.dtype), weight_local) + else: + bias = F.linear(latent.float(), weight_local) + step_local = BuildStepLocal.execute(bias=bias, base_local=base_local) + if shard.tp_size > 1: + full = get_parallel().attn_tp_group.all_gather(step_local, dim=-1) + else: + full = step_local + return full[..., : self.vocab_size] + + def forward(self, token_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + embed = self.get_prev_embeddings(token_ids) + logits = self.project_bias(embed) + return logits, embed + + def sample_block( + self, + base_logits: torch.Tensor, + *, + first_prev_tokens: torch.Tensor, + hidden_states: Optional[torch.Tensor], + sampler: StepSampler, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return run_markov_block( + self, + base_logits, + first_prev_tokens=first_prev_tokens, + hidden_states=hidden_states, + sampler=sampler, + ) + + +def build_dspark_v4_confidence_head( + *, config: DeepSeekV4Config, markov_rank: int +) -> Optional[DSparkConfidenceHead]: + if read_ragged_verify_mode() is RaggedVerifyMode.STATIC: + return None + if not hasattr(config, "enable_confidence_head"): + logger.warning( + "DSpark draft config has no enable_confidence_head field; treating the " + "confidence head as enabled." + ) + with_markov_cfg = getattr(config, "confidence_head_with_markov", None) + with_markov = ( + (markov_rank > 0) if with_markov_cfg is None else bool(with_markov_cfg) + ) + if with_markov and markov_rank <= 0: + raise ValueError( + "DSpark V4 confidence_head_with_markov requires markov_rank > 0, " + f"got markov_rank={markov_rank}." + ) + return DSparkConfidenceHead( + hidden_size=int(config.hidden_size), + markov_rank=int(markov_rank), + with_markov=with_markov, + bias=False, + ) + + +class DSparkV4Stage(DeepseekV4DecoderLayer): + + def __init__( + self, + config: DeepSeekV4Config, + layer_id: int, + stage_id: int, + num_stages: int, + num_target_layers: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_streams: Optional[List[torch.cuda.Stream]] = None, + ) -> None: + super().__init__( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=prefix, + is_nextn=True, + alt_streams=alt_streams, + ) + self.stage_id = stage_id + self.dim = config.hidden_size + + if stage_id == 0: + if num_target_layers <= 0: + raise ValueError( + "DSpark needs target layers for the target-hidden projection." + ) + self.main_proj = ReplicatedLinear( + config.hidden_size * num_target_layers, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("main_proj", prefix), + ) + self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + if stage_id == num_stages - 1: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + ( + self.hc_head_fn, + self.hc_head_base, + self.hc_head_scale, + ) = make_hc_head_params(config.hc_mult, config.hidden_size) + + def _build_self_attn( + self, + *, + config: DeepSeekV4Config, + layer_id: int, + quant_config: Optional[QuantizationConfig], + prefix: str, + alt_streams: Optional[List[torch.cuda.Stream]], + compress_ratio_override: Optional[int], + ) -> nn.Module: + del compress_ratio_override + return DSparkAttention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=prefix, + alt_streams=alt_streams, + ) + + def _hc_pre_block( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + y, post, comb, _ = self.hc_pre(x, hc_fn, hc_scale, hc_base) + return y, post, comb + + def _hc_post_block( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ) -> torch.Tensor: + return self.hc_post(x, residual, post, comb) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + residual = hidden_states + x, post, comb = self._hc_pre_block( + hidden_states, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base + ) + x = self.input_layernorm(x) + x = self.self_attn(positions, x, forward_batch) + x = self._hc_post_block(x, residual, post, comb) + + residual = x + x, post, comb = self._hc_pre_block( + x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base + ) + x = self.post_attention_layernorm(x) + x = self._run_ffn(x, forward_batch) + x = self._hc_post_block(x, residual, post, comb) + return x + + def _run_ffn(self, x: torch.Tensor, forward_batch: ForwardBatch) -> torch.Tensor: + shape = x.shape + x = x.reshape(-1, self.dim) + y = self._run_moe_ffn_dp_sync( + x, forward_batch, input_ids=None, input_ids_global=None + ) + return y.view(shape) + + +class DeepseekV4ForCausalLMDSpark(nn.Module): + + def __init__( + self, + config: DeepSeekV4Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + + dspark_config = parse_dspark_draft_config(draft_hf_config=config) + if not dspark_config.require_markov(): + raise ValueError( + "DSpark V4 draft requires markov_rank > 0, " + f"got markov_rank={dspark_config.markov_rank}." + ) + self.gamma = int( + dspark_config.resolve_gamma(default=int(config.num_hidden_layers)) + ) + self.block_size = self.gamma + if dspark_config.target_layer_ids is not None: + self.num_stages = len(dspark_config.target_layer_ids) + else: + self.num_stages = int(getattr(config, "num_nextn_predict_layers", 1) or 1) + + target_num_layers = ( + int(dspark_config.num_target_layers) + if dspark_config.num_target_layers is not None + else int(getattr(config, "num_hidden_layers", 1)) + ) + if dspark_config.target_layer_ids is not None: + self.num_target_features = len(dspark_config.target_layer_ids) + else: + self.num_target_features = target_num_layers + + self.start_layer = 0 + self.end_layer = self.num_stages + use_multi_stream = ( + envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get() + and envs.SGLANG_DSPARK_ENABLE_MULTI_STREAM.get() + and torch.cuda.is_available() + ) + self.alt_streams: Optional[List[torch.cuda.Stream]] = ( + [torch.cuda.Stream()] if use_multi_stream else None + ) + self.stages = nn.ModuleList( + [ + DSparkV4Stage( + config=config, + layer_id=stage_id, + stage_id=stage_id, + num_stages=self.num_stages, + num_target_layers=self.num_target_features, + quant_config=quant_config, + prefix=add_prefix(f"stages.{stage_id}", prefix), + alt_streams=self.alt_streams, + ) + for stage_id in range(self.num_stages) + ] + ) + self.markov_head = DSparkV4MarkovHead( + vocab_size=int(config.vocab_size), + markov_rank=int(dspark_config.markov_rank), + ) + self.confidence_head = build_dspark_v4_confidence_head( + config=config, markov_rank=int(dspark_config.markov_rank) + ) + self.hc_mult = int(config.hc_mult) + self.norm_eps = float(config.rms_norm_eps) + self.hc_eps = float(config.hc_eps) + + self.embed_tokens: Optional[nn.Module] = None + self.lm_head: Optional[nn.Module] = None + self._use_fp32_lm_head = envs.SGLANG_DSPARK_FP32_LM_HEAD.get() + self._opt_markov_w2_tp_shard = envs.SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD.get() + + @property + def enable_confidence_head(self) -> bool: + return self.confidence_head is not None + + def attach_shared_modules( + self, *, embed_tokens: nn.Module, lm_head: nn.Module + ) -> None: + self.embed_tokens = embed_tokens + self.lm_head = lm_head + self.markov_head.configure_tp_shard(lm_head=lm_head) + + def project_target_hidden(self, main_hidden: torch.Tensor) -> torch.Tensor: + stage0 = self.stages[0] + projected, _ = stage0.main_proj(main_hidden) + return stage0.main_norm(projected) + + def write_target_hidden_kv( + self, + *, + main_hidden: torch.Tensor, + swa_loc: torch.Tensor, + positions: torch.Tensor, + pool: DeepSeekV4TokenToKVPool, + ) -> None: + main_x = self.project_target_hidden(main_hidden) + swa_loc = swa_loc.to(torch.int32) + kvs = CommitKvProj.execute( + main_x=main_x, + wkv_linears=[stage.self_attn.wkv for stage in self.stages], + ) + for stage, kv in zip(self.stages, kvs): + attn = stage.self_attn + pool.set_swa_key_buffer_radix_fused_norm_rope( + layer_id=attn.layer_id, + swa_loc=swa_loc, + kv=kv, + kv_weight=attn.kv_norm.weight.data, + eps=attn.eps, + freqs_cis=attn.freqs_cis, + positions=positions, + ) + + def forward_embed(self, input_ids: torch.Tensor) -> torch.Tensor: + if self.embed_tokens is None: + raise ValueError( + "DeepseekV4ForCausalLMDSpark requires the target embed_tokens " + "(call attach_shared_modules first)." + ) + x = self.embed_tokens(input_ids) + x = x.unsqueeze(1).repeat(1, self.hc_mult, 1) + return x + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + get_embedding: bool = False, + pp_proxy_tensors=None, + ) -> LogitsProcessorOutput: + del get_embedding, pp_proxy_tensors + if input_embeds is None: + input_embeds = self.forward_embed(input_ids) + x = input_embeds + for stage in self.stages: + x = stage(positions, x, forward_batch) + + return LogitsProcessorOutput(next_token_logits=None, hidden_states=x) + + def collapse_hc_head(self, x: torch.Tensor) -> torch.Tensor: + last = self.stages[-1] + return hc_head_torch( + x, + last.hc_head_fn, + last.hc_head_scale, + last.hc_head_base, + norm_eps=self.norm_eps, + hc_eps=self.hc_eps, + ) + + def compute_base_logits(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + + x_post_hc = self.collapse_hc_head(x) + return self._logits_from_x_post_hc(x_post_hc), x_post_hc + + def _logits_from_x_post_hc(self, x_post_hc: torch.Tensor) -> torch.Tensor: + if self.lm_head is None: + raise ValueError( + "DeepseekV4ForCausalLMDSpark requires the target lm_head " + "(call attach_shared_modules first)." + ) + last = self.stages[-1] + x = last.norm(x_post_hc) + weight = self.lm_head.weight + if self._use_fp32_lm_head: + local_logits = F.linear(x.float(), weight.float()) + else: + local_logits = torch.matmul(x.to(weight.dtype), weight.T) + if self._opt_markov_w2_tp_shard: + return local_logits + return gather_and_crop_vocab(local_logits, self.lm_head) + + def compute_confidence( + self, + *, + anchor_tokens: torch.Tensor, + sampled_tokens: torch.Tensor, + x_post_hc: torch.Tensor, + ) -> Optional[torch.Tensor]: + confidence_head = self.confidence_head + if confidence_head is None: + return None + bs = int(anchor_tokens.shape[0]) + x_post_hc = x_post_hc.view(bs, self.gamma, -1) + if confidence_head.with_markov: + prev_seq = torch.cat( + [anchor_tokens.view(-1, 1), sampled_tokens[:, : self.gamma - 1]], dim=1 + ) + markov_embed_stack = self.markov_head.get_prev_embeddings(prev_seq) + else: + markov_embed_stack = None + confidence_raw = confidence_head(x_post_hc, markov_embed_stack) + confidence = confidence_head.apply_sts(confidence_raw) + maybe_detect_in_closed_range( + confidence, 0.0, 1.0, "DSpark confidence must lie in [0, 1]." + ) + return confidence + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None: + params_dict = dict(self.named_parameters()) + loaded_params = set() + + weights = list(weights) + if any(name.endswith(".wo_a.scale") for name, _ in weights): + weights = list(_dequant_fp8_wo_a(weights)) + + stacked_params_mapping = DEEPSEEK_V4_STACKED_PARAMS_MAPPING + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts, + ) + + for name, loaded_weight in weights: + mapped = self._remap_dspark_weight_name(name) + if mapped is None: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in mapped: + continue + candidate = mapped.replace(weight_name, param_name) + if candidate not in params_dict: + continue + param = params_dict[candidate] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(candidate) + break + else: + for ( + param_name, + weight_name, + expert_id, + shard_id, + ) in expert_params_mapping: + if weight_name not in mapped: + continue + candidate = mapped.replace(weight_name, param_name) + if candidate not in params_dict: + continue + param = params_dict[candidate] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + candidate, + shard_id=shard_id, + expert_id=expert_id, + ) + loaded_params.add(candidate) + break + else: + if mapped not in params_dict: + logger.warning( + "DSpark V4 draft: unexpected weight %r -> %r", name, mapped + ) + continue + param = params_dict[mapped] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(mapped) + + self._assert_confidence_head_loaded( + params_dict=params_dict, loaded_params=loaded_params + ) + + def _assert_confidence_head_loaded( + self, *, params_dict: dict, loaded_params: set + ) -> None: + if self.confidence_head is None: + return + confidence_param_names = { + name for name in params_dict if name.startswith("confidence_head.") + } + missing = confidence_param_names - loaded_params + if missing: + raise ValueError( + f"DSpark V4 confidence head is enabled but the checkpoint is missing " + f"{sorted(missing)}. Provide a checkpoint with trained confidence weights, " + f"or disable the confidence head (enable_confidence_head=False)." + ) + + def _remap_dspark_weight_name(self, name: str) -> Optional[str]: + if name.startswith(("embed.", "embed_tokens.", "head.", "lm_head.")): + return None + if "rotary_emb.inv_freq" in name: + return None + + if not name.startswith("mtp."): + return None + parts = name.split(".", 2) + if len(parts) < 3: + return None + stage_id, rest = parts[1], parts[2] + + if rest.startswith("markov_head."): + return f"markov_head.{rest[len('markov_head.'):]}" + + if rest.startswith("confidence_head."): + if self.confidence_head is None: + return None + return f"confidence_head.{rest[len('confidence_head.'):]}" + + mapped_rest = rest + mapped_rest = mapped_rest.replace("attn.", "self_attn.", 1) + mapped_rest = mapped_rest.replace("ffn.", "mlp.", 1) + mapped_rest = mapped_rest.replace("attn_norm.", "input_layernorm.", 1) + mapped_rest = mapped_rest.replace("ffn_norm.", "post_attention_layernorm.", 1) + mapped_rest = mapped_rest.replace(".w1.", ".gate_proj.") + mapped_rest = mapped_rest.replace(".w2.", ".down_proj.") + mapped_rest = mapped_rest.replace(".w3.", ".up_proj.") + mapped_rest = mapped_rest.replace(".gate.tid2eid", ".topk.tid2eid") + mapped_rest = mapped_rest.replace(".gate.bias", ".gate.e_score_correction_bias") + mapped_rest = mapped_rest.replace(".scale", ".weight_scale_inv") + return f"stages.{stage_id}.{mapped_rest}" + + +EntryClass = [DeepseekV4ForCausalLMDSpark] diff --git a/python/sglang/srt/models/dspark.py b/python/sglang/srt/models/dspark.py new file mode 100644 index 000000000..ea8bcb0ad --- /dev/null +++ b/python/sglang/srt/models/dspark.py @@ -0,0 +1,509 @@ +from __future__ import annotations + +import logging +from typing import Callable, Iterable, Optional, Tuple + +import torch +from torch import nn + +from sglang.srt.distributed.communication_op import tensor_model_parallel_all_gather +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.dflash import DFlashDraftModel +from sglang.srt.speculative.dspark_components.dspark_config import ( + parse_dspark_draft_config, +) +from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyMode, + read_ragged_verify_mode, +) + +logger = logging.getLogger(__name__) + +StepSampler = Callable[[torch.Tensor, int], torch.Tensor] + + +def gather_and_crop_vocab( + local_logits: torch.Tensor, lm_head: nn.Module +) -> torch.Tensor: + full_logits = tensor_model_parallel_all_gather(local_logits, dim=-1) + return full_logits[..., : int(lm_head.org_vocab_size)] + + +def run_markov_block( + head: nn.Module, + base_logits: torch.Tensor, + *, + first_prev_tokens: torch.Tensor, + hidden_states: Optional[torch.Tensor], + sampler: StepSampler, +) -> Tuple[torch.Tensor, torch.Tensor]: + batch_size, proposal_len = base_logits.shape[:2] + if proposal_len == 0: + empty = torch.empty(batch_size, 0, dtype=torch.long, device=base_logits.device) + return empty, base_logits + + sampled_tokens = [] + corrected_logits = [] + prev_tokens = first_prev_tokens.long() + for step_idx in range(proposal_len): + step_hidden = None if hidden_states is None else hidden_states[:, step_idx, ...] + step_logits = head.apply_step_logits( + base_logits[:, step_idx, :], + token_ids=prev_tokens, + hidden_states=step_hidden, + ) + next_tokens = sampler(step_logits, step_idx) + sampled_tokens.append(next_tokens) + corrected_logits.append(step_logits.unsqueeze(1)) + prev_tokens = next_tokens + return ( + torch.stack(sampled_tokens, dim=1), + torch.cat(corrected_logits, dim=1), + ) + + +class VanillaMarkov(nn.Module): + + markov_head_type = "vanilla" + + def __init__(self, *, vocab_size: int, markov_rank: int) -> None: + super().__init__() + self.vocab_size = int(vocab_size) + self.markov_rank = int(markov_rank) + if self.markov_rank <= 0: + raise ValueError( + f"VanillaMarkov requires markov_rank > 0, got {self.markov_rank}." + ) + self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank) + self.markov_w2 = nn.Linear(self.markov_rank, self.vocab_size, bias=False) + + def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.markov_w1(token_ids.long()) + + def project_bias(self, latent_states: torch.Tensor) -> torch.Tensor: + return self.markov_w2(latent_states) + + def compute_step_bias( + self, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + ) -> torch.Tensor: + del hidden_states + return self.project_bias(self.get_prev_embeddings(token_ids)) + + def apply_step_logits( + self, + logits: torch.Tensor, + *, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + ) -> torch.Tensor: + return logits + self.compute_step_bias(token_ids, hidden_states) + + def apply_block_logits( + self, + base_logits: torch.Tensor, + *, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + ) -> torch.Tensor: + if base_logits.size(-2) == 0: + return base_logits + return base_logits + self.compute_step_bias(token_ids, hidden_states) + + def sample_block( + self, + base_logits: torch.Tensor, + *, + first_prev_tokens: torch.Tensor, + hidden_states: Optional[torch.Tensor], + sampler: StepSampler, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return run_markov_block( + self, + base_logits, + first_prev_tokens=first_prev_tokens, + hidden_states=hidden_states, + sampler=sampler, + ) + + +class GatedMarkovHead(VanillaMarkov): + + markov_head_type = "gated" + + def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int) -> None: + super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) + self.gate_proj = nn.Linear(int(hidden_size) + markov_rank, markov_rank) + + def compute_gate( + self, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + ) -> torch.Tensor: + if hidden_states is None: + raise ValueError("GatedMarkovHead requires hidden_states.") + prev_embeddings = self.get_prev_embeddings(token_ids) + gate_inputs = torch.cat([hidden_states, prev_embeddings], dim=-1) + return torch.sigmoid(self.gate_proj(gate_inputs)) + + def compute_step_bias( + self, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + ) -> torch.Tensor: + prev_embeddings = self.get_prev_embeddings(token_ids) + gate = self.compute_gate(token_ids, hidden_states).to( + dtype=prev_embeddings.dtype + ) + return self.project_bias(gate * prev_embeddings) + + +class RNNHead(VanillaMarkov): + + markov_head_type = "rnn" + + def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int) -> None: + super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) + self.hidden_size = int(hidden_size) + self.state_size = markov_rank + self.joint_proj = nn.Linear(2 * markov_rank + self.hidden_size, 3 * markov_rank) + + def _rnn_step( + self, + state: torch.Tensor, + prev_embeddings: torch.Tensor, + hidden_states: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + z = torch.cat([state, prev_embeddings, hidden_states], dim=-1) + gate_raw, candidate_raw, output_raw = self.joint_proj(z).chunk(3, dim=-1) + gate = torch.sigmoid(gate_raw) + candidate = torch.tanh(candidate_raw) + new_state = gate * state + (1.0 - gate) * candidate + bias = self.project_bias(torch.tanh(output_raw)) + return new_state, bias + + def compute_step_bias( + self, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + ) -> torch.Tensor: + if hidden_states is None: + raise ValueError("RNNHead requires hidden_states.") + prev_embeddings = self.get_prev_embeddings(token_ids) + state = torch.zeros_like(prev_embeddings) + _, bias = self._rnn_step(state, prev_embeddings, hidden_states) + return bias + + def apply_block_logits( + self, + base_logits: torch.Tensor, + *, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + ) -> torch.Tensor: + if hidden_states is None: + raise ValueError("RNNHead requires hidden_states.") + block_size = base_logits.size(-2) + if block_size == 0: + return base_logits + leading_shape = base_logits.shape[:-2] + state = torch.zeros( + *leading_shape, + self.markov_rank, + device=base_logits.device, + dtype=hidden_states.dtype, + ) + output_logits = [] + for k in range(block_size): + prev_emb = self.get_prev_embeddings(token_ids[..., k]) + state, bias = self._rnn_step(state, prev_emb, hidden_states[..., k, :]) + output_logits.append(base_logits[..., k, :] + bias) + return torch.stack(output_logits, dim=-2) + + def sample_block( + self, + base_logits: torch.Tensor, + *, + first_prev_tokens: torch.Tensor, + hidden_states: Optional[torch.Tensor], + sampler: StepSampler, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if hidden_states is None: + raise ValueError("RNNHead requires hidden_states.") + batch_size, proposal_len = base_logits.shape[:2] + if proposal_len == 0: + empty = torch.empty( + batch_size, 0, dtype=torch.long, device=base_logits.device + ) + return empty, base_logits + + state = torch.zeros( + batch_size, + self.markov_rank, + device=base_logits.device, + dtype=hidden_states.dtype, + ) + sampled_tokens = [] + corrected_logits = [] + prev_tokens = first_prev_tokens.long() + for step_idx in range(proposal_len): + prev_emb = self.get_prev_embeddings(prev_tokens) + state, bias = self._rnn_step(state, prev_emb, hidden_states[:, step_idx, :]) + step_logits = base_logits[:, step_idx, :] + bias + next_tokens = sampler(step_logits, step_idx) + sampled_tokens.append(next_tokens) + corrected_logits.append(step_logits.unsqueeze(1)) + prev_tokens = next_tokens + return ( + torch.stack(sampled_tokens, dim=1), + torch.cat(corrected_logits, dim=1), + ) + + +def build_markov_head(config) -> Optional[nn.Module]: + markov_rank = int(getattr(config, "markov_rank", 0)) + if markov_rank <= 0: + raise ValueError( + "DSpark requires markov_rank > 0 (the Markov head is the core of the " + f"semi-AR draft); got markov_rank={markov_rank}." + ) + markov_head_type = str(getattr(config, "markov_head_type", "vanilla")).lower() + vocab_size = int(config.vocab_size) + hidden_size = int(config.hidden_size) + if markov_head_type == "vanilla": + return VanillaMarkov(vocab_size=vocab_size, markov_rank=markov_rank) + if markov_head_type == "gated": + return GatedMarkovHead( + vocab_size=vocab_size, markov_rank=markov_rank, hidden_size=hidden_size + ) + if markov_head_type == "rnn": + return RNNHead( + vocab_size=vocab_size, markov_rank=markov_rank, hidden_size=hidden_size + ) + raise ValueError(f"Unsupported DSpark markov_head_type={markov_head_type!r}.") + + +class DSparkConfidenceHead(nn.Module): + + def __init__( + self, + *, + hidden_size: int, + markov_rank: int, + with_markov: bool = True, + bias: bool = True, + dtype: torch.dtype = torch.float32, + ) -> None: + super().__init__() + self.with_markov = bool(with_markov) + input_dim = int(hidden_size) + (int(markov_rank) if self.with_markov else 0) + self.proj = nn.Linear(input_dim, 1, bias=bias, dtype=dtype) + self.register_buffer( + "sts_temperatures", torch.ones((), dtype=torch.float32), persistent=False + ) + self._last_confidence_raw: Optional[torch.Tensor] = None + + def forward( + self, + hidden_states: torch.Tensor, + markov_embed_stack: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if self.with_markov: + if markov_embed_stack is None: + raise ValueError( + "DSparkConfidenceHead(with_markov=True) requires markov_embed_stack." + ) + features = torch.cat( + [hidden_states, markov_embed_stack.to(dtype=hidden_states.dtype)], + dim=-1, + ) + else: + features = hidden_states + features = features.to(dtype=self.proj.weight.dtype) + return self.proj(features).squeeze(-1) + + def apply_sts(self, confidence_raw: torch.Tensor) -> torch.Tensor: + self._last_confidence_raw = confidence_raw + return torch.sigmoid(confidence_raw.float() / self.sts_temperatures) + + +def build_confidence_head(config) -> Optional[nn.Module]: + if read_ragged_verify_mode() is RaggedVerifyMode.STATIC: + return None + if not hasattr(config, "enable_confidence_head"): + logger.warning( + "DSpark draft config has no enable_confidence_head field; treating the " + "confidence head as enabled." + ) + hidden_size = int(config.hidden_size) + markov_rank = int(getattr(config, "markov_rank", 0)) + with_markov = bool(getattr(config, "confidence_head_with_markov", markov_rank > 0)) + if with_markov and markov_rank <= 0: + raise ValueError( + "DSpark confidence_head_with_markov requires markov_rank > 0, " + f"got markov_rank={markov_rank}." + ) + return DSparkConfidenceHead( + hidden_size=hidden_size, + markov_rank=markov_rank, + with_markov=with_markov, + ) + + +_DSPARK_SKIPPED_WEIGHT_PREFIXES = ( + "embed_tokens.", + "lm_head.", + "rotary_emb.", +) + + +class DSparkDraftMixin: + + def __init__(self, config, quant_config=None, prefix: str = "") -> None: + super().__init__(config=config, quant_config=quant_config, prefix=prefix) + dspark_config = parse_dspark_draft_config(draft_hf_config=config) + if not dspark_config.require_markov(): + raise ValueError( + "DSpark draft requires markov_rank > 0, " + f"got markov_rank={dspark_config.markov_rank}." + ) + self.gamma = int(dspark_config.resolve_gamma(default=self.block_size)) + self.markov_head = build_markov_head(config) + self.confidence_head = build_confidence_head(config) + self.lm_head: Optional[nn.Module] = None + + def attach_shared_modules( + self, *, embed_tokens: nn.Module, lm_head: nn.Module + ) -> None: + del embed_tokens + self.lm_head = lm_head + + def compute_base_logits( + self, hidden: torch.Tensor + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + if self.lm_head is None: + raise ValueError( + "DSpark dense draft requires the target lm_head " + "(call attach_shared_modules first)." + ) + weight = self.lm_head.weight + if hidden.dtype != weight.dtype: + hidden = hidden.to(weight.dtype) + local_logits = torch.matmul(hidden, weight.T) + base_logits = gather_and_crop_vocab(local_logits, self.lm_head) + return base_logits, None + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + markov_weights = [] + confidence_weights = [] + backbone_weights = [] + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if any(name.startswith(p) for p in _DSPARK_SKIPPED_WEIGHT_PREFIXES): + continue + if name.startswith("confidence_head."): + if self.confidence_head is None: + continue + confidence_weights.append((name, loaded_weight)) + elif name.startswith("markov_head."): + markov_weights.append((name, loaded_weight)) + else: + backbone_weights.append((name, loaded_weight)) + + super().load_weights(backbone_weights) + + for name, loaded_weight in markov_weights: + if name not in params_dict: + raise ValueError( + f"DSpark unexpected markov weight {name!r} not found in model " + f"parameters (known markov params require a {type(self.markov_head).__name__} head)." + ) + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + self._load_confidence_weights( + confidence_weights=confidence_weights, params_dict=params_dict + ) + + def _load_confidence_weights( + self, + *, + confidence_weights: list, + params_dict: dict, + ) -> None: + if self.confidence_head is None: + return + loaded_names = set() + for name, loaded_weight in confidence_weights: + if name not in params_dict: + raise ValueError( + f"DSpark unexpected confidence weight {name!r} not found in " + "model parameters." + ) + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_names.add(name) + + confidence_param_names = { + name for name in params_dict if name.startswith("confidence_head.") + } + missing = confidence_param_names - loaded_names + if missing: + raise ValueError( + f"DSpark confidence head is enabled but the checkpoint is missing " + f"{sorted(missing)}. Provide a checkpoint with trained confidence weights, " + f"or disable the confidence head (enable_confidence_head=False)." + ) + + def write_target_hidden_kv( + self, + *, + target_hidden: torch.Tensor, + pool, + positions: torch.Tensor, + cache_loc: torch.Tensor, + cache_loc_2d: Optional[torch.Tensor] = None, + commit_lens: Optional[torch.Tensor] = None, + ) -> None: + ctx_hidden = self.project_target_hidden(target_hidden) + for layer in self.layers: + attn = layer.self_attn + k, v = attn.kv_proj_only(ctx_hidden) + k = attn.apply_k_norm(k) + k = attn.apply_k_rope(positions, k) + k = k.view(-1, attn.num_kv_heads, attn.head_dim) + v = v.view(-1, attn.num_kv_heads, attn.head_dim) + if cache_loc_2d is not None and commit_lens is not None: + pool.set_kv_buffer_prefix_valid( + attn.attn, + cache_loc_2d, + commit_lens, + k, + v, + attn.attn.k_scale, + attn.attn.v_scale, + ) + else: + pool.set_kv_buffer( + attn.attn, + cache_loc, + k, + v, + attn.attn.k_scale, + attn.attn.v_scale, + ) + + +class DSparkDraftModel(DSparkDraftMixin, DFlashDraftModel): + + pass + + +class Qwen3DSparkModel(DSparkDraftModel): + pass + + +EntryClass = [Qwen3DSparkModel] diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index cc4d2d6a1..ffb917037 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -109,6 +109,8 @@ class SchedulerStats: # Speculative decoding spec_accept_length: float = 0.0 spec_accept_rate: float = 0.0 + spec_cap_length: float = 0.0 + spec_block_accept_length: float = 0.0 # Adaptive speculative decoding (currently active tier). spec_num_steps: int = 0 spec_num_draft_tokens: int = 0 @@ -423,6 +425,18 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): labelnames=labels.keys(), multiprocess_mode="mostrecent", ) + self.spec_cap_length = Gauge( + name="sglang:spec_cap_length", + documentation="Mean DSpark confidence-scheduled verify window per verify step, incl the bonus slot (0 when no cap is scheduled).", + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) + self.spec_block_accept_length = Gauge( + name="sglang:spec_block_accept_length", + documentation="Mean uncapped full-block accept length per verify step (accept + cap-trimmed drafts; exact only in DSpark cap-accept mode).", + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) self.spec_num_steps = Gauge( name="sglang:spec_num_steps", documentation="Currently active speculative_num_steps.", @@ -1287,6 +1301,8 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): # Speculative decoding self._log_gauge(self.spec_accept_length, stats.spec_accept_length) self._log_gauge(self.spec_accept_rate, stats.spec_accept_rate) + self._log_gauge(self.spec_cap_length, stats.spec_cap_length) + self._log_gauge(self.spec_block_accept_length, stats.spec_block_accept_length) self._log_gauge(self.spec_num_steps, stats.spec_num_steps) self._log_gauge(self.spec_num_draft_tokens, stats.spec_num_draft_tokens) diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index 680bac645..5d4c3772c 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py @@ -31,6 +31,8 @@ class SamplingBatchInfo: # Whether all requests use greedy sampling is_all_greedy: bool + is_any_greedy: bool + # Whether any requests use top_p sampling need_top_p_sampling: bool @@ -188,6 +190,7 @@ class SamplingBatchInfo: min_ps=min_ps, sampling_seed=sampling_seed, is_all_greedy=all(r.sampling_params.top_k <= 1 for r in reqs), + is_any_greedy=any(r.sampling_params.top_k <= 1 for r in reqs), need_top_p_sampling=any(r.sampling_params.top_p != 1.0 for r in reqs), need_top_k_sampling=any(r.sampling_params.top_k != TOP_K_ALL for r in reqs), need_min_p_sampling=any(r.sampling_params.min_p > 0 for r in reqs), @@ -409,6 +412,7 @@ class SamplingBatchInfo: setattr(self, item, torch.cat([self_val, other_val])) self.is_all_greedy &= other.is_all_greedy + self.is_any_greedy |= other.is_any_greedy self.need_top_p_sampling |= other.need_top_p_sampling self.need_top_k_sampling |= other.need_top_k_sampling self.need_min_p_sampling |= other.need_min_p_sampling diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 89b03f53a..d27e50d0e 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1655,7 +1655,7 @@ class ServerArgs: # ------------------------------------------------------------------------- speculative_algorithm: A[ Optional[str], - "Speculative algorithm. Builtins: EAGLE, EAGLE3, NEXTN, STANDALONE, NGRAM, DFLASH. Or any name registered via `SpeculativeAlgorithm.register`.", + "Speculative algorithm. Builtins: EAGLE, EAGLE3, NEXTN, STANDALONE, NGRAM, DFLASH, DSPARK. Or any name registered via `SpeculativeAlgorithm.register`.", ] = None speculative_draft_model_path: A[ Optional[str], @@ -1691,6 +1691,37 @@ class ServerArgs: Optional[int], "DFLASH only. Block size (verify window length). Alias of --speculative-num-draft-tokens for DFLASH.", ] = None + speculative_dspark_block_size: A[ + Optional[int], + "DSPARK only. Draft block size gamma (number of proposed draft tokens). The verify window is gamma + 1, so this sets --speculative-num-draft-tokens = gamma + 1. Omit to auto-infer gamma from the draft checkpoint block_size.", + ] = None + speculative_dspark_sps_table_path: A[ + Optional[str], + "DSPARK only. Path to a pre-profiled SPS cost table (JSON) built offline with " + "sglang.benchmark.dspark_sps_profiler, consumed by the ragged-verify " + "scheduler (cap-accept / compact). Omit for an uninitialized flat " + "constant-SPS table: the budget degenerates to verify-all (zero throughput " + "gain by itself).", + ] = None + speculative_dspark_confidence_sts_path: A[ + Optional[str], + "DSPARK only. Optional path to a per-position STS (sequential temperature " + "scaling) calibration JSON, fit offline with sglang.benchmark.dspark_sts_fit. " + "Calibrates the confidence-head survival probabilities the ragged-verify " + "scheduler consumes. Omit to use identity (no calibration); losslessness is " + "unaffected either way.", + ] = None + speculative_dspark_align_verify_tokens_to_graph_tier: A[ + bool, + "DSPARK compact ragged-verify only. Fill the per-request verify lengths so " + "the total verify-token count reaches the cuda-graph tier the forward is " + "already padded to: round the dp-max scheduled total up to the captured " + "token bucket and let the top-k allocator admit that many real draft tokens " + "(confidence-ordered). This recovers the padding the forward pays for anyway " + "-- both the cuda-graph bucket round-up and the dp cross-rank max -- turning " + "it into extra real verification at the same step time. Off by default; when " + "off the schedule is byte-for-byte unchanged.", + ] = False speculative_accept_threshold_single: A[ float, "Accept a draft token if its probability in the target model is greater than this threshold.", diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py index c871798be..95d28026b 100644 --- a/python/sglang/srt/speculative/base_spec_worker.py +++ b/python/sglang/srt/speculative/base_spec_worker.py @@ -338,6 +338,14 @@ class BaseSpecWorker(ABC): """ pass + def note_request_finished(self, *, rid: str, natural_stop: bool) -> None: + """Hook called by the batch-result processor when a request finishes. + + Default no-op. DSpark overrides this to settle / censor its + block-accept estimator state for the finished request. + """ + pass + def activate_step_by_batch(self, batch_size: int) -> None: """Activate the optimal adaptive step for the current batch size. diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index f0f3dc6fb..ef2b8c136 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/python/sglang/srt/speculative/dflash_info.py @@ -16,6 +16,7 @@ from sglang.srt.speculative.spec_info import SpecInput, SpecInputType if TYPE_CHECKING: from sglang.srt.managers.tp_worker import TpModelWorker + from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout @dataclass @@ -41,6 +42,8 @@ class DFlashVerifyInput(SpecInput): # Shape info for padding (e.g., DP attention / CUDA graph). num_tokens_per_req: int = -1 + ragged_verify_layout: Optional[RaggedVerifyLayout] = None + def __post_init__(self): super().__init__(spec_input_type=SpecInputType.DFLASH_VERIFY) if self.num_tokens_per_req == -1: @@ -99,20 +102,29 @@ class DFlashVerifyInput(SpecInput): device = req_pool_indices.device bs = len(req_pool_indices) - qo_indptr = torch.arange( - 0, - (bs + 1) * self.draft_token_num, - step=self.draft_token_num, - dtype=torch.int32, - device=device, - ) + layout = self.ragged_verify_layout + + if layout is None: + qo_indptr = torch.arange( + 0, + (bs + 1) * self.draft_token_num, + step=self.draft_token_num, + dtype=torch.int32, + device=device, + ) + verify_lens = self.draft_token_num + kv_indices_extra = self.draft_token_num * bs + else: + qo_indptr = layout.qo_indptr_device + verify_lens = layout.verify_lens + kv_indices_extra = layout.total_verify_tokens cum_kv_seq_len = torch.zeros((bs + 1,), dtype=torch.int32, device=device) - paged_kernel_lens = paged_kernel_lens + self.draft_token_num + paged_kernel_lens = paged_kernel_lens + verify_lens cum_kv_seq_len[1:] = torch.cumsum(paged_kernel_lens, dim=0) kv_indices = torch.empty( - paged_kernel_lens_sum + self.draft_token_num * bs, + paged_kernel_lens_sum + kv_indices_extra, dtype=torch.int32, device=device, ) diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index 866ea093b..2a9185987 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/python/sglang/srt/speculative/dflash_info_v2.py @@ -59,6 +59,8 @@ class DFlashDraftInputV2(SpecInput): # Filled by scheduler after dispatch. future_indices: Optional[torch.Tensor] = None + verify_token_budget: Optional[int] = None + def __post_init__(self): super().__init__(spec_input_type=SpecInputType.DFLASH_DRAFT) diff --git a/python/sglang/srt/speculative/dspark_components/__init__.py b/python/sglang/srt/speculative/dspark_components/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/speculative/dspark_components/dspark_block_accept_estimator.py b/python/sglang/srt/speculative/dspark_components/dspark_block_accept_estimator.py new file mode 100644 index 000000000..0a77646d7 --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/dspark_block_accept_estimator.py @@ -0,0 +1,835 @@ +from __future__ import annotations + +import json +import logging +import math +from collections import deque +from pathlib import Path +from typing import Any, List, Optional, Tuple, Union + +import msgspec +import torch + +from sglang.srt.environ import envs +from sglang.srt.kv_canary.runner.future_tensor import DelayedDeviceHostHandler +from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout + +logger = logging.getLogger(__name__) + +_GATHER_ROW_CHUNK = 512 +_STATE_SWEEP_INTERVAL = 1024 +_STATE_EXPIRE_STEPS = 4096 +_FLUSH_EVERY_STEPS = 16 +_PENDING_BUCKET_MIN = 16 +_DEFAULT_ONLINE_WINDOW_STEPS = 256 + + +SKIP_STEP_WARNING = ( + "skipping step: {} (pending blocks of affected requests " + "are dropped by the seq-len continuity check)" +) + + +def block_accept_skip_reason( + *, + logits_adjustments_are_noop: bool, + corrected_logits: Optional[Any], +) -> Optional[str]: + if not logits_adjustments_are_noop: + return ( + "non-noop logits adjustments (penalizer/logit_bias/grammar) " + "in batch; cross-step conditioning of the gathered target " + "probabilities would be state-dependent" + ) + if corrected_logits is None: + return "corrected_logits unavailable (folded draft path)" + return None + + +def warn_once(warned_reasons: set, *, reason: str) -> None: + if reason not in warned_reasons: + warned_reasons.add(reason) + logger.warning( + "DSPARK block accept estimate recorder: %s (warned once)", reason + ) + + +def gather_chunked_token_logprobs( + *, + logits, + row_indices, + token_indices, + per_row_temps, + chunk_size: int, +): + """Chunked per-row token logprob gather: logprob of token_indices[i] under + logits[row_indices[i]] / per_row_temps[i], computed chunk_size rows at a + time to bound the fp32 softmax workspace.""" + results = [] + for start in range(0, row_indices.shape[0], chunk_size): + end = start + chunk_size + rows = logits[row_indices[start:end]].to(torch.float32) + rows = rows / per_row_temps[start:end, None] + log_norm = torch.logsumexp(rows, dim=-1) + token_logits = rows.gather(dim=1, index=token_indices[start:end, None]).squeeze( + 1 + ) + results.append(token_logits - log_norm) + return torch.cat(results) + + +def _pending_bucket(count: int) -> int: + if count == 0: + return 0 + bucket = _PENDING_BUCKET_MIN + while bucket < count: + bucket *= 2 + return bucket + + +class _CeilingSnapshot(msgspec.Struct): + window_lo: float + window_hi: float + window_blocks: int + window_horizon: int + cumulative_lo: float + cumulative_hi: float + cumulative_blocks: int + + +class _OnlineCeiling: + def __init__(self, *, log_interval: int, window_steps: int) -> None: + self._log_interval = log_interval + self._window_steps = window_steps + self._steps: deque[Tuple[int, float, float, int]] = deque() + self._win_lo = 0.0 + self._win_hi = 0.0 + self._win_count = 0 + self._cum_lo = 0.0 + self._cum_hi = 0.0 + self._cum_count = 0 + self._max_forward_ct = 0 + + def add(self, *, forward_ct: int, lo: float, hi: float) -> None: + self._max_forward_ct = max(self._max_forward_ct, forward_ct) + if self._steps and self._steps[-1][0] == forward_ct: + fct, slo, shi, c = self._steps[-1] + self._steps[-1] = (fct, slo + lo, shi + hi, c + 1) + else: + self._steps.append((forward_ct, lo, hi, 1)) + self._win_lo += lo + self._win_hi += hi + self._win_count += 1 + self._cum_lo += lo + self._cum_hi += hi + self._cum_count += 1 + self._evict(forward_ct=self._max_forward_ct) + + def _evict(self, *, forward_ct: int) -> None: + cutoff = forward_ct - self._window_steps + while self._steps and self._steps[0][0] <= cutoff: + _, slo, shi, c = self._steps.popleft() + self._win_lo -= slo + self._win_hi -= shi + self._win_count -= c + + def estimate(self) -> Optional[_CeilingSnapshot]: + if self._cum_count == 0: + return None + return _CeilingSnapshot( + window_lo=self._win_lo / self._win_count, + window_hi=self._win_hi / self._win_count, + window_blocks=self._win_count, + window_horizon=min(self._window_steps, self._max_forward_ct), + cumulative_lo=self._cum_lo / self._cum_count, + cumulative_hi=self._cum_hi / self._cum_count, + cumulative_blocks=self._cum_count, + ) + + def maybe_log(self, *, forward_ct: int) -> None: + if self._log_interval <= 0 or forward_ct % self._log_interval != 0: + return + snap = self.estimate() + if snap is None: + return + logger.info( + "DSpark uncapped-acc-len estimate (forward_ct=%d): " + "last %d passes ~%.3f [%.3f, %.3f] w=%.3f (%d blocks) | " + "cumulative ~%.3f [%.3f, %.3f] w=%.3f (%d blocks)", + forward_ct, + snap.window_horizon, + 0.5 * (snap.window_lo + snap.window_hi), + snap.window_lo, + snap.window_hi, + snap.window_hi - snap.window_lo, + snap.window_blocks, + 0.5 * (snap.cumulative_lo + snap.cumulative_hi), + snap.cumulative_lo, + snap.cumulative_hi, + snap.cumulative_hi - snap.cumulative_lo, + snap.cumulative_blocks, + ) + + +class _PendingBlock(msgspec.Struct): + forward_ct: int + anchor_pos: int + window: int + trimmed_tokens: List[int] + next_offset: int + q_lps: List[float] = [] + est_prod: float = 1.0 + est_lo_extra: float = 0.0 + + +class _RequestState(msgspec.Struct): + expected_seq_len: int = -1 + last_seen_ct: int = 0 + pending: List[_PendingBlock] = [] + + +class _PendingPlan(msgspec.Struct): + rows: List[int] + tokens: List[int] + slot_lookup: dict[tuple[int, int, int], int] + + +class _SettleBatch(msgspec.Struct): + forward_ct: int + rids: List[str] + row_meta: List[List[int]] + drafts: List[List[int]] + q_all: List[List[float]] + target_diag: List[List[float]] + pending_logprobs: List[float] + slot_lookup: dict[tuple[int, int, int], int] + + @classmethod + def from_bundle(cls, bundle: dict[str, Any]) -> _SettleBatch: + return cls( + forward_ct=bundle["forward_ct"], + rids=bundle["rids"], + row_meta=bundle["row_meta"].tolist(), + drafts=bundle["draft_tokens"].tolist(), + q_all=bundle["q_all"].tolist(), + target_diag=bundle["target_diag_logprobs"].tolist(), + pending_logprobs=bundle["pending_logprobs"].tolist(), + slot_lookup=bundle["pending_slot_lookup"], + ) + + +class BlockAcceptEstimateRecorder: + def __init__( + self, + *, + path: str, + gamma: int, + device: Union[str, torch.device], + online_log_interval: int = 0, + online_window_steps: int = 0, + ) -> None: + self._gamma = gamma + self._last_forward_ct = 0 + if path: + self._path: Optional[Path] = Path(path) + self._path.parent.mkdir(parents=True, exist_ok=True) + self._file = self._path.open("w") + else: + self._path = None + self._file = None + self._device = torch.device(device) + self._states: dict[str, _RequestState] = {} + self._steps_since_flush = 0 + self._observed_step_ct = 0 + self._discontinuity_drop_ct = 0 + self._skipped_step_ct = 0 + self._warned_skip_reasons: set[str] = set() + self._finish_intents: dict[str, bool] = {} + + self._online = _OnlineCeiling( + log_interval=online_log_interval, + window_steps=( + online_window_steps + if online_window_steps > 0 + else ( + online_log_interval + if online_log_interval > 0 + else _DEFAULT_ONLINE_WINDOW_STEPS + ) + ), + ) + + self._retained_h2d: List[torch.Tensor] = [] + self._delayed: Optional[DelayedDeviceHostHandler] = None + if self._device.type == "cuda": + self._delayed = DelayedDeviceHostHandler( + d2h_stream=torch.cuda.Stream(device=self._device) + ) + + logger.info( + "DSPARK block accept estimate recorder enabled: path=%s gamma=%d " + "async=%s online_log_interval=%d", + path, + gamma, + self._delayed is not None, + online_log_interval, + ) + + def observe_verify_step( + self, + *, + forward_ct: int, + rids: List[str], + draft_tokens: torch.Tensor, + corrected_logits: Optional[torch.Tensor], + draft_temperatures: torch.Tensor, + greedy_mask: torch.Tensor, + target_logits: torch.Tensor, + target_temperatures: torch.Tensor, + truncated_sampling_mask: Optional[torch.Tensor], + logits_adjustments_are_noop: bool, + correct_len: torch.Tensor, + cap_trim_lens: torch.Tensor, + bonus: torch.Tensor, + prefix_lens: torch.Tensor, + layout: Optional[RaggedVerifyLayout], + ) -> None: + if ( + self._delayed is not None + and torch.cuda.is_available() + and torch.cuda.is_current_stream_capturing() + ): + return + + skip_reason = self._skip_reason( + logits_adjustments_are_noop=logits_adjustments_are_noop, + corrected_logits=corrected_logits, + ) + if skip_reason is not None: + self._skip_step(reason=skip_reason) + + def compute_on_device() -> Optional[dict[str, Any]]: + if skip_reason is not None: + return None + return self._build_device_bundle( + forward_ct=forward_ct, + rids=rids, + draft_tokens=draft_tokens, + corrected_logits=corrected_logits, + draft_temperatures=draft_temperatures, + greedy_mask=greedy_mask, + target_logits=target_logits, + target_temperatures=target_temperatures, + truncated_sampling_mask=truncated_sampling_mask, + correct_len=correct_len, + cap_trim_lens=cap_trim_lens, + bonus=bonus, + prefix_lens=prefix_lens, + layout=layout, + ) + + if self._delayed is not None: + self._delayed.step( + compute_on_device=compute_on_device, + postprocess_on_host=self._settle_and_write, + ) + else: + bundle = compute_on_device() + if bundle is not None: + self._settle_and_write(bundle) + + def flush(self) -> None: + if self._delayed is not None: + self._delayed.step( + compute_on_device=lambda: None, + postprocess_on_host=self._settle_and_write, + ) + self._apply_all_finish_intents() + if self._file is not None: + self._file.flush() + self._steps_since_flush = 0 + + def note_request_finished(self, *, rid: str, natural_stop: bool) -> None: + if self._delayed is None: + self._finalize_request( + rid=rid, natural_stop=natural_stop, forward_ct=self._last_forward_ct + ) + else: + self._finish_intents[rid] = natural_stop + + def _apply_all_finish_intents(self) -> None: + for rid in list(self._finish_intents): + self._finalize_request( + rid=rid, + natural_stop=self._finish_intents.pop(rid), + forward_ct=self._last_forward_ct, + ) + + def _finalize_request( + self, *, rid: str, natural_stop: bool, forward_ct: int + ) -> None: + state = self._states.pop(rid, None) + if state is None: + return + for block in state.pending: + if natural_stop: + self._finalize_eos_online(block, forward_ct=forward_ct) + else: + self._finalize_at_end_online(block, forward_ct=forward_ct) + if natural_stop and state.pending: + self._write_eos_marker(rid=rid, blocks=state.pending) + + def _finalize_eos_online(self, block: _PendingBlock, *, forward_ct: int) -> None: + lo = block.window + 1.0 + block.est_lo_extra + self._online.add(forward_ct=forward_ct, lo=lo, hi=lo) + + def _write_eos_marker(self, *, rid: str, blocks: List[_PendingBlock]) -> None: + if self._file is None: + return + marker = {"rid": rid, "eos_end": [block.forward_ct for block in blocks]} + self._file.write(json.dumps(marker) + "\n") + + def online_estimate(self) -> Optional[_CeilingSnapshot]: + return self._online.estimate() + + def estimate_log_suffix(self) -> Optional[str]: + snap = self.online_estimate() + if snap is None: + return None + mid = 0.5 * (snap.cumulative_lo + snap.cumulative_hi) + return ( + f"est uncap acc len: {mid:.2f} " + f"[{snap.cumulative_lo:.2f}, {snap.cumulative_hi:.2f}]" + ) + + def drain_pending_online(self) -> None: + for state in self._states.values(): + for block in state.pending: + self._finalize_at_end_online(block, forward_ct=self._last_forward_ct) + state.pending = [] + + def _finalize_walk_online( + self, block: _PendingBlock, *, diverged: bool, forward_ct: int + ) -> None: + base = block.window + 1.0 + lo = base + block.est_lo_extra + if diverged: + offset = block.next_offset - 1 + tail = ( + block.est_prod * (self._gamma - offset) if offset < self._gamma else 0.0 + ) + else: + tail = 0.0 + self._online.add(forward_ct=forward_ct, lo=lo, hi=lo + tail) + + def _finalize_at_end_online(self, block: _PendingBlock, *, forward_ct: int) -> None: + base = block.window + 1.0 + lo = base + block.est_lo_extra + tail = block.est_prod * (self._gamma - block.next_offset + 1) + self._online.add(forward_ct=forward_ct, lo=lo, hi=lo + tail) + + def _build_device_bundle( + self, + *, + forward_ct: int, + rids: List[str], + draft_tokens: torch.Tensor, + corrected_logits: torch.Tensor, + draft_temperatures: torch.Tensor, + greedy_mask: torch.Tensor, + target_logits: torch.Tensor, + target_temperatures: torch.Tensor, + truncated_sampling_mask: Optional[torch.Tensor], + correct_len: torch.Tensor, + cap_trim_lens: torch.Tensor, + bonus: torch.Tensor, + prefix_lens: torch.Tensor, + layout: Optional[RaggedVerifyLayout], + ) -> dict[str, Any]: + gamma = self._gamma + rows_per_request = gamma + 1 + bs = len(rids) + device = target_logits.device + assert draft_tokens.shape == (bs, gamma) + assert corrected_logits.shape[0] == bs and corrected_logits.shape[1] == gamma + assert target_logits.shape[0] == bs * rows_per_request + + if truncated_sampling_mask is not None: + truncated_mask = truncated_sampling_mask + else: + truncated_mask = torch.zeros(bs, dtype=torch.bool, device=device) + if layout is not None: + verify_lens = layout.verify_lens + else: + verify_lens = torch.full( + (bs,), rows_per_request, dtype=torch.int32, device=device + ) + + draft_temps_full = ( + draft_temperatures.reshape(bs).to(torch.float32).repeat_interleave(gamma) + ) + target_temps_full = ( + target_temperatures.reshape(bs) + .to(torch.float32) + .repeat_interleave(rows_per_request) + ) + draft_flat = draft_tokens.reshape(-1) + + q_all = self._gather_logprobs( + logits=corrected_logits.reshape(bs * gamma, -1), + row_indices=torch.arange(bs * gamma, device=device), + token_indices=draft_flat, + temps=draft_temps_full, + ).reshape(bs, gamma) + target_diag = self._gather_logprobs( + logits=target_logits, + row_indices=self._diag_rows(bs=bs, rows_per_request=rows_per_request), + token_indices=draft_flat, + temps=target_temps_full, + ).reshape(bs, gamma) + + self._retained_h2d = [] + plan = self._plan_pending(bs=bs, rows_per_request=rows_per_request, rids=rids) + pending_logprobs = self._gather_pending( + plan=plan, + target_logits=target_logits, + target_temps_full=target_temps_full, + device=device, + ) + + return { + "forward_ct": int(forward_ct), + "rids": list(rids), + "row_meta": self._pack_row_meta( + correct_len=correct_len, + cap_trim_lens=cap_trim_lens, + bonus=bonus, + prefix_lens=prefix_lens, + greedy_mask=greedy_mask, + truncated_mask=truncated_mask, + verify_lens=verify_lens, + ), + "draft_tokens": draft_tokens, + "q_all": q_all, + "target_diag_logprobs": target_diag, + "pending_logprobs": pending_logprobs, + "pending_slot_lookup": plan.slot_lookup, + } + + def _diag_rows(self, *, bs: int, rows_per_request: int) -> torch.Tensor: + device = self._device + return ( + (torch.arange(bs, device=device) * rows_per_request)[:, None] + + torch.arange(self._gamma, device=device)[None, :] + ).reshape(-1) + + def _plan_pending( + self, *, bs: int, rows_per_request: int, rids: List[str] + ) -> _PendingPlan: + gamma = self._gamma + rows: List[int] = [] + tokens: List[int] = [] + slot_lookup: dict[tuple[int, int, int], int] = {} + for b in range(bs): + state = self._states.get(rids[b]) + if state is None or not state.pending or state.expected_seq_len < 0: + continue + expected_seq_len = state.expected_seq_len + for block_idx, block in enumerate(state.pending): + offset = block.next_offset + while offset <= gamma: + row = block.anchor_pos + offset - expected_seq_len + if row < 0 or row >= rows_per_request: + break + slot_lookup[(b, block_idx, offset)] = len(rows) + rows.append(b * rows_per_request + row) + tokens.append(block.trimmed_tokens[offset - block.window - 1]) + offset += 1 + return _PendingPlan(rows=rows, tokens=tokens, slot_lookup=slot_lookup) + + def _gather_pending( + self, + *, + plan: _PendingPlan, + target_logits: torch.Tensor, + target_temps_full: torch.Tensor, + device: torch.device, + ) -> torch.Tensor: + bucket = _pending_bucket(len(plan.rows)) + rows = plan.rows + [0] * (bucket - len(plan.rows)) + tokens = plan.tokens + [0] * (bucket - len(plan.tokens)) + return self._gather_logprobs( + logits=target_logits, + row_indices=self._host_to_device_async(rows, device=device), + token_indices=self._host_to_device_async(tokens, device=device), + temps=target_temps_full, + ) + + def _pack_row_meta( + self, + *, + correct_len: torch.Tensor, + cap_trim_lens: torch.Tensor, + bonus: torch.Tensor, + prefix_lens: torch.Tensor, + greedy_mask: torch.Tensor, + truncated_mask: torch.Tensor, + verify_lens: torch.Tensor, + ) -> torch.Tensor: + return torch.stack( + [ + correct_len.to(torch.int64), + cap_trim_lens.to(torch.int64), + bonus.to(torch.int64), + prefix_lens.to(torch.int64), + greedy_mask.to(torch.int64), + truncated_mask.to(torch.int64), + verify_lens.to(torch.int64), + ], + dim=1, + ) + + def _settle_and_write(self, bundle: dict[str, Any]) -> None: + batch = _SettleBatch.from_bundle(bundle) + self._last_forward_ct = batch.forward_ct + for b in range(len(batch.rids)): + self._settle_row(b=b, batch=batch) + self._finish_step(forward_ct=batch.forward_ct) + self._apply_all_finish_intents() + + def _settle_row(self, *, b: int, batch: _SettleBatch) -> None: + forward_ct = batch.forward_ct + rid = batch.rids[b] + state = self._states.setdefault(rid, _RequestState()) + state.last_seen_ct = forward_ct + + cl, cap_trim, bonus_token, seq_len, is_greedy, is_truncated, verify_len = ( + batch.row_meta[b] + ) + window = verify_len - 1 + assert 0 <= cl <= window <= self._gamma + + self._drop_pending_on_discontinuity( + state, seq_len=seq_len, forward_ct=forward_ct + ) + state.expected_seq_len = seq_len + cl + 1 + + if is_greedy or is_truncated: + if is_truncated and not is_greedy: + self._warn_once( + reason="requests with top-k/top-p/min-p sampling are " + "excluded per-row; the estimator only supports " + "pure-temperature sampling (processed target distribution " + "would differ from plain softmax(logits/T))" + ) + state.pending = [] + return + + record: dict[str, Any] = { + "rid": rid, + "fct": forward_ct, + "w": window, + "cl": cl, + "ct": cap_trim, + } + num_old_pending = len(state.pending) + if cl == window and window < self._gamma: + self._open_block( + state, + record, + drafts_row=batch.drafts[b], + q_all_row=batch.q_all[b], + window=window, + seq_len=seq_len, + forward_ct=forward_ct, + ) + else: + self._online.add(forward_ct=forward_ct, lo=cl + 1.0, hi=cl + 1.0) + + pending_gathers = self._settle_pending( + b=b, + batch=batch, + state=state, + realized=batch.drafts[b][:cl] + [bonus_token], + cl=cl, + seq_len=seq_len, + num_old_pending=num_old_pending, + ) + if pending_gathers: + record["pg"] = pending_gathers + if self._file is not None: + self._file.write(json.dumps(record) + "\n") + + def _open_block( + self, + state: _RequestState, + record: dict[str, Any], + *, + drafts_row: List[int], + q_all_row: List[float], + window: int, + seq_len: int, + forward_ct: int, + ) -> None: + trimmed_tokens = drafts_row[window : self._gamma] + q_lps = q_all_row[window : self._gamma] + state.pending.append( + _PendingBlock( + forward_ct=forward_ct, + anchor_pos=seq_len - 1, + window=window, + trimmed_tokens=trimmed_tokens, + next_offset=window + 1, + q_lps=q_lps, + ) + ) + record["trimmed_tokens"] = trimmed_tokens + record["q_lp"] = q_lps + + def _settle_pending( + self, + *, + b: int, + batch: _SettleBatch, + state: _RequestState, + realized: List[int], + cl: int, + seq_len: int, + num_old_pending: int, + ) -> List[list]: + gamma = self._gamma + pending_gathers: List[list] = [] + kept_pending: List[_PendingBlock] = [] + for block_idx, block in enumerate(state.pending): + diverged = False + while block.next_offset <= gamma: + row = block.anchor_pos + block.next_offset - seq_len + assert row >= 0 + if row > cl: + break + token = block.trimmed_tokens[block.next_offset - block.window - 1] + if block_idx < num_old_pending: + p_lp = batch.pending_logprobs[ + batch.slot_lookup[(b, block_idx, block.next_offset)] + ] + else: + p_lp = batch.target_diag[b][row] + pending_gathers.append( + [block.forward_ct, block.next_offset, p_lp, token, realized[row]] + ) + self._accumulate_online(block, p_lp=p_lp) + block.next_offset += 1 + if realized[row] != token: + diverged = True + break + if not diverged and block.next_offset <= gamma: + kept_pending.append(block) + else: + self._finalize_walk_online( + block, diverged=diverged, forward_ct=batch.forward_ct + ) + state.pending = kept_pending + return pending_gathers + + def _accumulate_online(self, block: _PendingBlock, *, p_lp: float) -> None: + a = min(1.0, math.exp(p_lp - block.q_lps[block.next_offset - block.window - 1])) + block.est_prod *= a + block.est_lo_extra += block.est_prod + + def _drop_pending_on_discontinuity( + self, state: _RequestState, *, seq_len: int, forward_ct: int + ) -> None: + if state.expected_seq_len < 0 or seq_len == state.expected_seq_len: + return + if not state.pending: + return + self._discontinuity_drop_ct += len(state.pending) + for block in state.pending: + self._finalize_at_end_online(block, forward_ct=forward_ct) + state.pending = [] + + def _finish_step(self, *, forward_ct: int) -> None: + self._observed_step_ct += 1 + if self._file is not None: + self._steps_since_flush += 1 + if self._steps_since_flush >= _FLUSH_EVERY_STEPS: + self._file.flush() + self._steps_since_flush = 0 + if self._observed_step_ct % _STATE_SWEEP_INTERVAL == 0: + self._sweep_states(forward_ct=forward_ct) + self._online.maybe_log(forward_ct=forward_ct) + + def _host_to_device_async( + self, values: List[int], *, device: torch.device + ) -> torch.Tensor: + host = torch.tensor(values, dtype=torch.long, pin_memory=device.type == "cuda") + self._retained_h2d.append(host) + return host.to(device=device, non_blocking=True) + + def _gather_logprobs( + self, + *, + logits: torch.Tensor, + row_indices: torch.Tensor, + token_indices: torch.Tensor, + temps: torch.Tensor, + ) -> torch.Tensor: + if row_indices.numel() == 0: + return torch.zeros(0, dtype=torch.float32, device=logits.device) + per_row_temps = temps[row_indices].clamp_min(1e-5) + return gather_chunked_token_logprobs( + logits=logits, + row_indices=row_indices, + token_indices=token_indices, + per_row_temps=per_row_temps, + chunk_size=_GATHER_ROW_CHUNK, + ) + + def _sweep_states(self, *, forward_ct: int) -> None: + expired = [ + rid + for rid, state in self._states.items() + if forward_ct - state.last_seen_ct > _STATE_EXPIRE_STEPS + ] + for rid in expired: + for block in self._states[rid].pending: + self._finalize_at_end_online(block, forward_ct=forward_ct) + del self._states[rid] + self._finish_intents.pop(rid, None) + + def _skip_reason( + self, + *, + logits_adjustments_are_noop: bool, + corrected_logits: Optional[torch.Tensor], + ) -> Optional[str]: + return block_accept_skip_reason( + logits_adjustments_are_noop=logits_adjustments_are_noop, + corrected_logits=corrected_logits, + ) + + def _skip_step(self, *, reason: str) -> None: + self._skipped_step_ct += 1 + self._warn_once(reason=SKIP_STEP_WARNING.format(reason)) + + def _warn_once(self, *, reason: str) -> None: + warn_once(self._warned_skip_reasons, reason=reason) + + +def create_block_accept_estimate_recorder( + *, gamma: int, device: Union[str, torch.device], tp_rank: int +) -> Optional[BlockAcceptEstimateRecorder]: + if tp_rank != 0: + return None + + path = envs.SGLANG_DSPARK_BLOCK_ACCEPT_ESTIMATE_PATH.get() + online_log_interval = envs.SGLANG_DSPARK_BLOCK_ACCEPT_ONLINE_INTERVAL.get() + if not path and online_log_interval <= 0: + return None + + return BlockAcceptEstimateRecorder( + path=path, + gamma=gamma, + device=device, + online_log_interval=online_log_interval, + ) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_config.py b/python/sglang/srt/speculative/dspark_components/dspark_config.py new file mode 100644 index 000000000..0cac50977 --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/dspark_config.py @@ -0,0 +1,295 @@ +from __future__ import annotations + +import json +import logging +from typing import TYPE_CHECKING, Any, List, Optional + +import msgspec + +from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + +DEFAULT_DSPARK_GAMMA = 7 +SUPPORTED_DSPARK_MARKOV_HEAD_TYPES = ("vanilla", "gated", "rnn") + +# The dsv4 self-drafting checkpoint runs its draft attention on the dedicated +# DeepSeek-V4 backend instead of the generic draft-backend fallback. +DSV4_DRAFT_ATTENTION_BACKEND = "dsv4" + + +def draft_is_deepseek_v4(*, server_args: ServerArgs) -> bool: + from sglang.srt.configs.model_config import is_deepseek_v4 + from sglang.srt.utils.hf_transformers_utils import get_config + + draft_model_path = server_args.speculative_draft_model_path + if not draft_model_path: + return False + draft_hf_config = get_config( + draft_model_path, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.speculative_draft_model_revision, + model_override_args=json.loads(server_args.json_model_override_args), + model_config_parser=server_args.model_config_parser, + ) + return draft_hf_config is not None and is_deepseek_v4(draft_hf_config) + + +def dspark_gamma_from_num_draft_tokens(num_draft_tokens: int) -> int: + gamma = int(num_draft_tokens) - 1 + if gamma < 1: + raise ValueError( + "DSpark speculative_num_draft_tokens must be >= 2 (= gamma + 1), " + f"got {num_draft_tokens}." + ) + return gamma + + +class DSparkDraftConfig(msgspec.Struct, frozen=True): + num_hidden_layers: Optional[int] + num_target_layers: Optional[int] + gamma: Optional[int] + target_layer_ids: Optional[List[int]] + mask_token: str + mask_token_id: Optional[int] + markov_rank: int + markov_head_type: Optional[str] + + def resolve_gamma(self, *, default: Optional[int] = None) -> Optional[int]: + return self.gamma if self.gamma is not None else default + + def require_markov(self) -> bool: + return int(self.markov_rank) > 0 + + +class DSparkRuntimeConfig(msgspec.Struct, frozen=True): + gamma: int + verify_num_draft_tokens: int + mask_token_id: int + + +def resolve_runtime_config( + *, + draft_hf_config: Any, + speculative_num_draft_tokens: Optional[int], + target_vocab_size: int, +) -> DSparkRuntimeConfig: + """Resolve and validate the worker-facing DSpark runtime knobs (gamma, + verify window, mask token) from the draft checkpoint config, with + server_args.speculative_num_draft_tokens taking precedence for gamma.""" + draft_config = parse_dspark_draft_config(draft_hf_config=draft_hf_config) + if not draft_config.require_markov(): + raise ValueError( + "DSpark draft requires markov_rank > 0; got " + f"markov_rank={draft_config.markov_rank}." + ) + + if speculative_num_draft_tokens is None: + gamma = int(draft_config.resolve_gamma(default=None) or 0) + if gamma < 1: + raise ValueError( + "DSpark could not resolve gamma from the draft config and " + "speculative_num_draft_tokens is unset." + ) + else: + gamma = dspark_gamma_from_num_draft_tokens(int(speculative_num_draft_tokens)) + config_gamma = draft_config.resolve_gamma(default=None) + if config_gamma is not None and int(config_gamma) != gamma: + logger.warning( + "DSpark gamma mismatch: using gamma=%s (from " + "speculative_num_draft_tokens=%s) but draft config block_size=%s.", + gamma, + speculative_num_draft_tokens, + config_gamma, + ) + + if draft_config.mask_token_id is None: + raise ValueError( + "DSpark requires mask_token_id to be set in the draft model config." + ) + mask_token_id = int(draft_config.mask_token_id) + if mask_token_id >= target_vocab_size: + raise ValueError( + f"DSpark mask_token_id={mask_token_id} is outside the target " + f"vocab size {target_vocab_size}." + ) + + return DSparkRuntimeConfig( + gamma=gamma, + verify_num_draft_tokens=gamma + 1, + mask_token_id=mask_token_id, + ) + + +def read_draft_checkpoint_gamma(*, server_args: ServerArgs) -> Optional[int]: + """Load the draft checkpoint's hf config and read its DSpark gamma + (block_size). Raises on config-load failure; callers pick the fallback.""" + from sglang.srt.utils.hf_transformers_utils import get_config + + draft_hf_config = get_config( + server_args.speculative_draft_model_path, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.speculative_draft_model_revision, + model_override_args=json.loads(server_args.json_model_override_args), + ) + return parse_dspark_draft_config(draft_hf_config=draft_hf_config).resolve_gamma( + default=None + ) + + +def checkpoint_bundles_dspark_draft(hf_config: Any) -> bool: + """The checkpoint carries a bundled DSpark draft head, marked by the + prefixed dspark_* keys on the target hf config. Single source of truth + for the bundling convention (draft-path defaulting, draft-arch remap).""" + return any( + _cfg_get(hf_config, key, None) is not None + for key in ( + "dspark_block_size", + "dspark_markov_rank", + "dspark_noise_token_id", + "dspark_target_layer_ids", + ) + ) + + +def _cfg_get(config: Any, key: str, default: Any = None) -> Any: + if isinstance(config, dict): + return config.get(key, default) + return getattr(config, key, default) + + +def _get_text_config(config: Any) -> Any: + if config is None: + return None + if isinstance(config, dict): + return config.get("text_config", config) + text_config = getattr(config, "text_config", None) + if text_config is not None: + return text_config + return config + + +def _get_dspark_config(config: Any) -> dict: + cfg = _cfg_get(config, "dspark_config", None) + if cfg is None: + return {} + if isinstance(cfg, dict): + return cfg + try: + return dict(cfg) + except Exception: + return {} + + +def parse_dspark_draft_config(*, draft_hf_config: Any) -> DSparkDraftConfig: + base = parse_dflash_draft_config(draft_hf_config=draft_hf_config) + + dspark_cfg = _get_dspark_config(draft_hf_config) + text_config = _get_text_config(draft_hf_config) + + prefixed_block_size = _cfg_get(draft_hf_config, "dspark_block_size", None) + prefixed_markov_rank = _cfg_get(draft_hf_config, "dspark_markov_rank", None) + prefixed_markov_head_type = _cfg_get( + draft_hf_config, "dspark_markov_head_type", None + ) + prefixed_noise_token_id = _cfg_get(draft_hf_config, "dspark_noise_token_id", None) + prefixed_target_layer_ids = _cfg_get( + draft_hf_config, "dspark_target_layer_ids", None + ) + uses_prefixed = any( + value is not None + for value in ( + prefixed_block_size, + prefixed_markov_rank, + prefixed_noise_token_id, + prefixed_target_layer_ids, + ) + ) + + raw_markov_rank = ( + prefixed_markov_rank + if prefixed_markov_rank is not None + else dspark_cfg.get( + "markov_rank", + _cfg_get( + text_config, "markov_rank", _cfg_get(draft_hf_config, "markov_rank", 0) + ), + ) + ) + markov_rank = int(raw_markov_rank) if raw_markov_rank is not None else 0 + if markov_rank < 0: + raise ValueError(f"DSpark markov_rank must be >= 0, got {markov_rank}.") + + markov_head_type = ( + prefixed_markov_head_type + if prefixed_markov_head_type is not None + else dspark_cfg.get( + "markov_head_type", + _cfg_get( + text_config, + "markov_head_type", + _cfg_get(draft_hf_config, "markov_head_type", None), + ), + ) + ) + if markov_rank > 0 and markov_head_type is None and not uses_prefixed: + raise ValueError( + "DSpark requires markov_head_type when markov_rank > 0, got None." + ) + if markov_head_type is not None: + markov_head_type = str(markov_head_type).lower() + if markov_head_type not in SUPPORTED_DSPARK_MARKOV_HEAD_TYPES: + raise ValueError( + f"Unsupported DSpark markov_head_type={markov_head_type!r}. " + f"Supported: {SUPPORTED_DSPARK_MARKOV_HEAD_TYPES}." + ) + + raw_mask_token_id = ( + prefixed_noise_token_id + if prefixed_noise_token_id is not None + else dspark_cfg.get( + "mask_token_id", + _cfg_get( + text_config, + "mask_token_id", + _cfg_get(draft_hf_config, "mask_token_id", base.mask_token_id), + ), + ) + ) + mask_token_id = int(raw_mask_token_id) if raw_mask_token_id is not None else None + if mask_token_id is not None and mask_token_id < 0: + raise ValueError( + f"DSpark mask_token_id must be non-negative, got {mask_token_id}." + ) + + gamma = ( + int(prefixed_block_size) if prefixed_block_size is not None else base.block_size + ) + + if prefixed_target_layer_ids is not None: + if not isinstance(prefixed_target_layer_ids, (list, tuple)) or not len( + prefixed_target_layer_ids + ): + raise ValueError( + "DSpark dspark_target_layer_ids must be a non-empty list of ints, " + f"got {prefixed_target_layer_ids!r}." + ) + target_layer_ids: Optional[List[int]] = [ + int(x) for x in prefixed_target_layer_ids + ] + else: + target_layer_ids = base.target_layer_ids + + return DSparkDraftConfig( + num_hidden_layers=base.num_hidden_layers, + num_target_layers=base.num_target_layers, + gamma=gamma, + target_layer_ids=target_layer_ids, + mask_token=base.mask_token, + mask_token_id=mask_token_id, + markov_rank=markov_rank, + markov_head_type=markov_head_type, + ) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py new file mode 100644 index 000000000..7eb9a13f3 --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -0,0 +1,421 @@ +from __future__ import annotations + +import logging +from contextlib import nullcontext +from typing import Optional + +import msgspec +import torch + +from sglang.srt.environ import envs +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + ForwardBatch, + ForwardMode, +) +from sglang.srt.runtime_context import get_parallel +from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 +from sglang.srt.speculative.draft_worker_common import make_draft_input_v2 +from sglang.srt.speculative.dspark_components.dspark_planner import VerifyWindow +from sglang.srt.speculative.dspark_components.kernels.dspark_draft_model import ( + SampleStepTokens, +) +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.speculative.spec_utils import draft_tp_context + +logger = logging.getLogger(__name__) + + +class DraftBlockResult(msgspec.Struct, frozen=True): + draft_tokens: torch.Tensor + corrected_logits: Optional[torch.Tensor] + greedy_mask: torch.Tensor + temperatures: torch.Tensor + + +class DraftForwardResult(msgspec.Struct, frozen=True): + draft_block_ids: torch.Tensor + raw_hidden: torch.Tensor + draft_hidden_3d: torch.Tensor + can_run_graph: bool + + +class DraftProposal(msgspec.Struct, frozen=True): + draft_block_ids: torch.Tensor + draft_block: DraftBlockResult + draft_hidden: Optional[torch.Tensor] + confidence: Optional[torch.Tensor] = None + confidence_tap: Optional[torch.Tensor] = None + folded: bool = False + + +def greedy_step_sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor: + del step_idx + return torch.argmax(step_logits, dim=-1) + + +class DsparkDraftSampler: + + def __init__(self, *, model, gamma, max_bs, device, confidence_fn=None, out=None): + self.model = model + self.markov_head = model.markov_head + self.gamma = int(gamma) + if out is not None: + assert out.shape == (int(max_bs) * self.gamma,) and out.dtype == torch.int64 + self.out = out + else: + self.out = torch.empty( + (int(max_bs) * self.gamma,), dtype=torch.int64, device=device + ) + self.confidence_fn = confidence_fn + self.confidence_out = ( + torch.empty((int(max_bs), self.gamma), dtype=torch.float32, device=device) + if confidence_fn is not None + else None + ) + + def __call__(self, hidden_states, input_ids): + bs = hidden_states.shape[0] // self.gamma + base_logits, confidence_tap = self.model.compute_base_logits(hidden_states) + base_logits = base_logits.view(bs, self.gamma, -1) + anchor = input_ids.view(bs, self.gamma)[:, 0] + draft_tokens, _ = self.markov_head.sample_block( + base_logits, + first_prev_tokens=anchor, + hidden_states=hidden_states.view(bs, self.gamma, -1), + sampler=greedy_step_sampler, + ) + self.out[: draft_tokens.numel()].copy_(draft_tokens.reshape(-1)) + if self.confidence_out is not None: + confidence = self.confidence_fn( + draft_hidden=hidden_states.view(bs, self.gamma, -1), + anchor_tokens=anchor, + draft_tokens=draft_tokens, + confidence_tap=confidence_tap, + ) + self.confidence_out[:bs].copy_(confidence) + + +def maybe_build_draft_sampler( + *, + draft_model, + gamma: int, + max_bs: int, + device, + tp_rank: int, + confidence_fn=None, + out=None, +) -> Optional[DsparkDraftSampler]: + """Build the graph-folded greedy draft sampler, or return None (with the + reason logged) when the draft model cannot support folding and the + proposal must stay eager.""" + + def _eager(reason): + if tp_rank == 0: + logger.info("DSpark draft greedy proposal kept eager (reason=%s).", reason) + return None + + if gamma <= 0: + return _eager("gamma<=0") + if not hasattr(draft_model, "compute_base_logits"): + return _eager("no compute_base_logits") + if getattr(draft_model, "markov_head", None) is None: + return _eager("no markov head") + if tp_rank == 0: + logger.info("DSpark draft greedy proposal folded into the draft cuda graph.") + return DsparkDraftSampler( + model=draft_model, + gamma=gamma, + max_bs=max_bs, + device=device, + confidence_fn=confidence_fn, + out=out, + ) + + +def make_next_draft_input( + *, + bonus_tokens: torch.Tensor, + new_seq_lens: torch.Tensor, +) -> DFlashDraftInputV2: + return make_draft_input_v2(bonus_tokens=bonus_tokens, new_seq_lens=new_seq_lens) + + +def resolve_greedy_mask( + *, + bs: int, + sampling_info, + device: torch.device, +) -> torch.Tensor: + if sampling_info is None: + return torch.ones(bs, dtype=torch.bool, device=device) + return (sampling_info.top_ks <= 1).view(-1) + + +def sample_draft_block( + *, + base_logits: torch.Tensor, + anchor_tokens: torch.Tensor, + draft_hidden: torch.Tensor, + sampling_info, + markov_head, + device: torch.device, +) -> DraftBlockResult: + bs = base_logits.shape[0] + greedy_mask = resolve_greedy_mask(bs=bs, sampling_info=sampling_info, device=device) + any_sampling = sampling_info is not None and not sampling_info.is_all_greedy + fast_sampling = envs.SGLANG_DSPARK_FAST_SAMPLING.get() + + if sampling_info is None: + temperatures = torch.ones(bs, dtype=torch.float32, device=device) + else: + temperatures = ( + sampling_info.temperatures.view(-1).to(torch.float32).clamp_min(1e-5) + ) + + if not any_sampling: + + def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor: + return torch.argmax(step_logits, dim=-1) + + else: + + def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor: + if fast_sampling: + exp_noise = torch.empty( + step_logits.shape, dtype=torch.float32, device=step_logits.device + ).exponential_(1) + return SampleStepTokens.execute( + step_logits=step_logits, + temperatures=temperatures, + greedy_mask=greedy_mask, + exp_noise=exp_noise, + ) + else: + probs = torch.softmax( + step_logits.float() / temperatures[:, None], dim=-1 + ) + argmax_tokens = torch.argmax(step_logits, dim=-1) + sampled_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1) + return torch.where(greedy_mask, argmax_tokens, sampled_tokens) + + draft_tokens, corrected_logits = markov_head.sample_block( + base_logits, + first_prev_tokens=anchor_tokens, + hidden_states=draft_hidden, + sampler=sampler, + ) + return DraftBlockResult( + draft_tokens=draft_tokens, + corrected_logits=corrected_logits, + greedy_mask=greedy_mask, + temperatures=temperatures, + ) + + +class DraftBlockProposer: + def __init__( + self, + *, + draft_model, + draft_model_runner, + gamma: int, + mask_token_id: int, + draft_block_spec_info, + dp_moe_sync: bool = False, + ) -> None: + self.draft_model = draft_model + self.draft_model_runner = draft_model_runner + self.gamma = gamma + self._mask_token_id = mask_token_id + self._draft_block_spec_info = draft_block_spec_info + self._draft_sampler = None + self._dp_moe_sync = dp_moe_sync + + def attach_draft_sampler(self, draft_sampler) -> None: + self._draft_sampler = draft_sampler + + def _base_logits_context(self): + if self._dp_moe_sync: + return draft_tp_context(get_parallel().attn_tp_group) + return nullcontext() + + def propose( + self, + *, + batch: ScheduleBatch, + draft_input: DFlashDraftInputV2, + verify_window: VerifyWindow, + bs: int, + device: str, + target_model, + sampling_info, + ) -> DraftProposal: + embed_module = target_model.get_input_embeddings() + fwd = self._run_forward( + batch=batch, + draft_input=draft_input, + verify_window=verify_window, + bs=bs, + device=device, + embed_module=embed_module, + ) + draft_block_ids = fwd.draft_block_ids + + draft_sampler = self._draft_sampler + all_greedy = sampling_info is None or sampling_info.is_all_greedy + folded_confidence = None + confidence_tap = None + folded = False + if draft_sampler is not None and fwd.can_run_graph and all_greedy: + folded = True + if sampling_info is None: + temperatures = torch.ones(bs, dtype=torch.float32, device=device) + else: + temperatures = ( + sampling_info.temperatures.view(-1) + .to(torch.float32) + .clamp_min(1e-5) + ) + draft_block = DraftBlockResult( + draft_tokens=draft_sampler.out[: bs * self.gamma].view(bs, self.gamma), + corrected_logits=None, + greedy_mask=resolve_greedy_mask( + bs=bs, sampling_info=sampling_info, device=device + ), + temperatures=temperatures, + ) + if draft_sampler.confidence_out is not None: + folded_confidence = draft_sampler.confidence_out[:bs] + else: + with self._base_logits_context(): + base_logits, confidence_tap = self.draft_model.compute_base_logits( + fwd.raw_hidden + ) + base_logits = base_logits.view(bs, self.gamma, -1) + draft_block = sample_draft_block( + base_logits=base_logits, + anchor_tokens=draft_block_ids[:, 0], + draft_hidden=fwd.draft_hidden_3d, + sampling_info=sampling_info, + markov_head=self.draft_model.markov_head, + device=device, + ) + return DraftProposal( + draft_block_ids=draft_block_ids, + draft_block=draft_block, + draft_hidden=fwd.draft_hidden_3d, + confidence=folded_confidence, + confidence_tap=confidence_tap, + folded=folded, + ) + + def run_idle_participation(self, batch: ScheduleBatch) -> None: + if not self._dp_moe_sync or batch.global_num_tokens is None: + return + device = self.draft_model_runner.device + empty_long = torch.empty((0,), dtype=torch.int64, device=device) + idle_batch = ForwardBatch( + forward_mode=ForwardMode.IDLE, + batch_size=0, + input_ids=empty_long, + req_pool_indices=empty_long, + seq_lens=empty_long, + out_cache_loc=empty_long, + seq_lens_sum=0, + seq_lens_cpu=torch.empty((0,), dtype=torch.int64), + positions=empty_long, + spec_algorithm=SpeculativeAlgorithm.DSPARK, + spec_info=self._draft_block_spec_info, + capture_hidden_mode=CaptureHiddenMode.NULL, + ) + self._fill_dp_moe_sync_metadata(idle_batch, batch) + with torch.inference_mode(): + self.draft_model_runner.forward(idle_batch) + + def _run_forward( + self, + *, + batch: ScheduleBatch, + draft_input: DFlashDraftInputV2, + verify_window: VerifyWindow, + bs: int, + device: str, + embed_module, + ) -> DraftForwardResult: + gamma = self.gamma + prefix_lens = batch.seq_lens + positions_2d = verify_window.positions_2d + verify_cache_loc_2d = verify_window.verify_cache_loc_2d + + draft_block_ids = torch.full( + (bs, gamma), int(self._mask_token_id), dtype=torch.long, device=device + ) + draft_block_ids[:, 0].copy_(draft_input.bonus_tokens.view(-1)) + draft_positions = positions_2d[:, :gamma].reshape(-1) + draft_cache_loc = verify_cache_loc_2d[:, :gamma].reshape(-1) + + draft_owns_embed = hasattr(self.draft_model, "forward_embed") + draft_input_embeds: Optional[torch.Tensor] = None + if not draft_owns_embed: + noise_embedding = embed_module(draft_block_ids) + draft_input_embeds = noise_embedding.view(-1, noise_embedding.shape[-1]) + + if batch.seq_lens_cpu is not None: + draft_seq_lens_cpu = batch.seq_lens_cpu + gamma + draft_seq_lens_sum = int(draft_seq_lens_cpu.sum()) + elif draft_input.reserved_seq_lens_cpu is not None: + draft_seq_lens_cpu = draft_input.reserved_seq_lens_cpu + draft_seq_lens_sum = int(draft_input.reserved_seq_lens_sum) + else: + raise RuntimeError("DSpark decode expected batch.seq_lens_cpu, got None") + + draft_forward_batch = ForwardBatch( + forward_mode=ForwardMode.TARGET_VERIFY, + batch_size=bs, + input_ids=draft_block_ids.flatten(), + req_pool_indices=batch.req_pool_indices, + seq_lens=prefix_lens, + out_cache_loc=draft_cache_loc, + seq_lens_sum=draft_seq_lens_sum, + seq_lens_cpu=draft_seq_lens_cpu, + positions=draft_positions, + input_embeds=draft_input_embeds, + spec_algorithm=SpeculativeAlgorithm.DSPARK, + spec_info=self._draft_block_spec_info, + capture_hidden_mode=CaptureHiddenMode.NULL, + ) + self._fill_dp_moe_sync_metadata(draft_forward_batch, batch) + with torch.inference_mode(): + draft_out = self.draft_model_runner.forward(draft_forward_batch) + logits_output = draft_out.logits_output + raw_hidden = logits_output.hidden_states + if raw_hidden is None: + raise RuntimeError("DSpark draft model returned no hidden states.") + draft_hidden_3d = raw_hidden.view(bs, gamma, -1) + return DraftForwardResult( + draft_block_ids=draft_block_ids, + raw_hidden=raw_hidden, + draft_hidden_3d=draft_hidden_3d, + can_run_graph=draft_out.can_run_graph, + ) + + def _fill_dp_moe_sync_metadata( + self, forward_batch: ForwardBatch, batch: ScheduleBatch + ) -> None: + if not self._dp_moe_sync or batch.global_num_tokens is None: + return + gnt, gnt_logprob = ( + self._draft_block_spec_info.get_spec_adjusted_global_num_tokens(batch) + ) + device = self.draft_model_runner.device + forward_batch.global_num_tokens_cpu = gnt + forward_batch.global_num_tokens_for_logprob_cpu = gnt_logprob + forward_batch.global_num_tokens_gpu = torch.tensor(gnt, dtype=torch.int64).to( + device, non_blocking=True + ) + forward_batch.global_num_tokens_for_logprob_gpu = torch.tensor( + gnt_logprob, dtype=torch.int64 + ).to(device, non_blocking=True) + forward_batch.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph diff --git a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py new file mode 100644 index 000000000..329c3f455 --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py @@ -0,0 +1,157 @@ +from typing import Optional + +import torch + +from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.speculative.dspark_components.kernels.dspark_verify_window import ( + BuildCommitInjectLayout, +) +from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout + + +class TargetHiddenKvInjector: + def __init__( + self, + *, + draft_model, + draft_model_runner, + model_runner, + device, + verify_num_draft_tokens: int, + block_pos_offsets: torch.Tensor, + ) -> None: + self.draft_model = draft_model + self.draft_model_runner = draft_model_runner + self.model_runner = model_runner + self.device = device + self.verify_num_draft_tokens = verify_num_draft_tokens + self._block_pos_offsets = block_pos_offsets + + def inject_target_hidden( + self, + *, + target_hidden: torch.Tensor, + cache_loc: torch.Tensor, + positions: torch.Tensor, + cache_loc_2d: Optional[torch.Tensor] = None, + commit_lens: Optional[torch.Tensor] = None, + ) -> None: + if target_hidden is None or target_hidden.numel() == 0: + return + device = self.model_runner.device + cache_loc = cache_loc.to(device=device, dtype=torch.int64, non_blocking=True) + positions = positions.to(device=device, dtype=torch.int64, non_blocking=True) + target_hidden = target_hidden.to(device=device, non_blocking=True) + n_real = positions.shape[0] + if target_hidden.shape[0] > n_real: + target_hidden = target_hidden[:n_real] + if cache_loc_2d is not None: + cache_loc_2d = cache_loc_2d.to( + device=device, dtype=torch.int64, non_blocking=True + ) + if commit_lens is not None: + commit_lens = commit_lens.to( + device=device, dtype=torch.int32, non_blocking=True + ) + + pool = self.draft_model_runner.token_to_kv_pool + if hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope"): + self._inject_mla( + pool=pool, + target_hidden=target_hidden, + cache_loc=cache_loc, + positions=positions, + cache_loc_2d=cache_loc_2d, + commit_lens=commit_lens, + ) + return + + with torch.inference_mode(): + self.draft_model.write_target_hidden_kv( + target_hidden=target_hidden, + pool=pool, + positions=positions, + cache_loc=cache_loc, + cache_loc_2d=cache_loc_2d, + commit_lens=commit_lens, + ) + + def _inject_mla( + self, + *, + pool, + target_hidden: torch.Tensor, + cache_loc: torch.Tensor, + positions: torch.Tensor, + cache_loc_2d: Optional[torch.Tensor], + commit_lens: Optional[torch.Tensor], + ) -> None: + swa_loc = pool.translate_loc_from_full_to_swa(cache_loc).to(torch.int32) + if commit_lens is not None and cache_loc_2d is not None: + bs, verify_len = cache_loc_2d.shape + col = torch.arange(verify_len, device=cache_loc.device).view(1, -1) + committed_mask = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1) + swa_loc = torch.where(committed_mask, swa_loc, torch.full_like(swa_loc, -1)) + + with torch.inference_mode(): + self.draft_model.write_target_hidden_kv( + main_hidden=target_hidden, + swa_loc=swa_loc, + positions=positions, + pool=pool, + ) + + def inject_ragged( + self, + *, + batch: ScheduleBatch, + layout: RaggedVerifyLayout, + hidden_strided: torch.Tensor, + commit_lens: torch.Tensor, + bs: int, + ) -> None: + stride = self.verify_num_draft_tokens + prefix_lens = batch.seq_lens + hidden = hidden_strided.view(bs, stride, -1) + + pool = self.draft_model_runner.token_to_kv_pool + if hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope"): + if hidden_strided.numel() == 0: + return + inject_layout = BuildCommitInjectLayout.execute( + req_pool_indices=batch.req_pool_indices, + req_to_token=self.model_runner.req_to_token_pool.req_to_token, + prefix_lens=prefix_lens, + block_pos_offsets=self._block_pos_offsets[:stride], + full_to_swa_mapping=pool.full_to_swa_index_mapping, + commit_lens=commit_lens, + stride=stride, + ) + with torch.inference_mode(): + self.draft_model.write_target_hidden_kv( + main_hidden=hidden.reshape(-1, hidden.shape[-1]), + swa_loc=inject_layout.swa_loc, + positions=inject_layout.positions, + pool=pool, + ) + return + + positions_2d = prefix_lens.unsqueeze(1) + self._block_pos_offsets + verify_cache_loc = assign_extend_cache_locs_func( + req_pool_indices=batch.req_pool_indices, + req_to_token=self.model_runner.req_to_token_pool.req_to_token, + start_offset=prefix_lens, + end_offset=prefix_lens + stride, + batch_size=bs, + draft_token_num=stride, + device=self.device, + ) + verify_cache_loc_2d = verify_cache_loc.view(bs, stride) + self.inject_target_hidden( + target_hidden=hidden.reshape(-1, hidden.shape[-1]), + cache_loc=verify_cache_loc, + cache_loc_2d=verify_cache_loc_2d, + positions=positions_2d.reshape(-1), + commit_lens=commit_lens, + ) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_observability.py b/python/sglang/srt/speculative/dspark_components/dspark_observability.py new file mode 100644 index 000000000..e08d11d38 --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/dspark_observability.py @@ -0,0 +1,961 @@ +from __future__ import annotations + +import logging +import math +import statistics +import time +from collections import deque +from contextlib import contextmanager, nullcontext +from enum import Enum +from typing import Callable, ContextManager, Iterator, Optional, Union + +import msgspec +import torch + +from sglang.srt.environ import envs +from sglang.srt.kv_canary.runner.future_tensor import FutureTensors +from sglang.srt.runtime_context import get_parallel +from sglang.srt.sampling.sampling_params import TOP_K_ALL +from sglang.srt.speculative.dflash_utils import compute_dflash_correct_drafts_and_bonus +from sglang.srt.speculative.dspark_components.dspark_block_accept_estimator import ( + create_block_accept_estimate_recorder, +) +from sglang.srt.speculative.dspark_components.dspark_sts import StsDataRecorder +from sglang.srt.speculative.dspark_components.dspark_verify import ( + verify_logits_adjustments_are_noop, +) + +logger = logging.getLogger(__name__) + +_NULL_SEGMENT = nullcontext() + +ALL_COMPONENTS_TOKEN = "all" + + +class InfoComponent(str, Enum): + CORE = "core" + STEP_CPU_TIME = "step_cpu_time" + STEP_GPU_TIME = "step_gpu_time" + DRAFT_GPU_TIME = "draft_gpu_time" + TARGET_VERIFY_GPU_TIME = "target_verify_gpu_time" + REQS = "reqs" + + +class InfoSegment(str, Enum): + STEP = "step" + DRAFT = "draft" + TARGET_VERIFY = "target_verify" + + +INFO_DUMP_MAX_RECORDS = 200_000 +INFO_DUMP_MAX_STEP_CPU_SECONDS = 1.0 + + +def resolve_enabled_components() -> set[InfoComponent]: + """Components enabled via env: SGLANG_DSPARK_DEBUG_DUMP tokens, plus the + published SPS-profiling switch SGLANG_DSPARK_ENABLE_SPS_RECORD=1, which is + an alias for the core,step_cpu_time components the SPS table fit needs.""" + components = resolve_components(envs.SGLANG_DSPARK_DEBUG_DUMP.get()) + if envs.SGLANG_DSPARK_ENABLE_SPS_RECORD.get(): + components |= {InfoComponent.CORE, InfoComponent.STEP_CPU_TIME} + return components + + +def resolve_components(raw: tuple[str, ...]) -> set[InfoComponent]: + tokens = {token.strip() for token in raw if token.strip()} + if not tokens: + return set() + if ALL_COMPONENTS_TOKEN in tokens: + return set(InfoComponent) + try: + return {InfoComponent(token) for token in tokens} + except ValueError as exc: + valid = [component.value for component in InfoComponent] + raise ValueError( + f"Invalid SGLANG_DSPARK_DEBUG_DUMP token in {sorted(tokens)}; " + f"valid: {valid} or '{ALL_COMPONENTS_TOKEN}'." + ) from exc + + +class ReqDetail(msgspec.Struct, omit_defaults=True): + req_pool_index: int + prefix_len: int + verify_len: int + acc_len: int + correct_drafts: int + cap_trim: int + bonus_token: int + draft_tokens: list[int] + rid: Optional[str] = None + confidence: Optional[list[float]] = None + survival: Optional[list[float]] = None + + +class DecodeStepRecord(msgspec.Struct, omit_defaults=True): + forward_ct: int + bs: int = -1 + mode: str = "" + budget: Optional[int] = None + lag_steps: Optional[int] = None + num_running_reqs: int = -1 + num_verify_tokens: int = -1 + verify_tokens_local: int = -1 + verify_tokens_dp_synced: int = -1 + verify_tokens_graph_key: int = -1 + predicted_step_ms: Optional[float] = None + predicted_theta: Optional[float] = None + step_cpu_ms: Optional[float] = None + step_gpu_ms: Optional[float] = None + draft_gpu_ms: Optional[float] = None + target_verify_gpu_ms: Optional[float] = None + reqs: Optional[list[ReqDetail]] = None + + +class DecodeStepObservation(msgspec.Struct): + forward_ct: int + bs: int + mode: str + budget: Optional[int] + lag_steps: Optional[int] + num_verify_tokens: int + verify_tokens_local: int + verify_tokens_dp_synced: int + verify_tokens_graph_key: int + predicted_step_ms: Optional[float] + predicted_theta: Optional[float] + verify_lens: Optional[torch.Tensor] + confidence: Optional[torch.Tensor] + req_pool_indices: torch.Tensor + prefix_lens: torch.Tensor + draft_tokens: torch.Tensor + bonus_tokens: torch.Tensor + correct_len: torch.Tensor + cap_trim_lens: torch.Tensor + commit_lens: torch.Tensor + rids: Optional[list[str]] + + +class _PendingStep(msgspec.Struct): + forward_ct: int + bs: int + mode: str + budget: Optional[int] + lag_steps: Optional[int] + num_verify_tokens: int + verify_tokens_local: int + verify_tokens_dp_synced: int + verify_tokens_graph_key: int + predicted_step_ms: Optional[float] + predicted_theta: Optional[float] + step_cpu_ms: Optional[float] + rids: Optional[list[str]] + future: Optional[FutureTensors] + segment_events: dict[InfoSegment, tuple[torch.cuda.Event, torch.cuda.Event]] + + +class DsparkInfoDumper: + def __init__( + self, + *, + components: set[Union[InfoComponent, str]], + gamma: int, + verify_num_draft_tokens: int, + attn_tp_rank: int, + device: torch.device, + mode_value: str, + sps_report_interval: int = 0, + max_records: int = INFO_DUMP_MAX_RECORDS, + max_step_cpu_seconds: float = INFO_DUMP_MAX_STEP_CPU_SECONDS, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self.gamma = int(gamma) + self.verify_num_draft_tokens = int(verify_num_draft_tokens) + self.attn_tp_rank = int(attn_tp_rank) + self.device = device + self.mode_value = mode_value + self._clock = clock + self._max_step_cpu_seconds = max_step_cpu_seconds + + self._components: set[InfoComponent] = { + InfoComponent(component) for component in components + } + self._sps_report_interval = int(sps_report_interval) + if self._sps_report_interval > 0: + self._components.add(InfoComponent.STEP_GPU_TIME) + # Dedup within an attention-TP group only: records describe the + # DP-rank-local batch, so under dp-attention every DP rank must keep + # dumping (the SPS profiler reads one payload per DP rank). + self.enabled = bool(self._components) and self.attn_tp_rank == 0 + self._sps_window: list[tuple[float, float]] = [] + self._sps_mismatched = 0 + + self._records: deque[DecodeStepRecord] = deque(maxlen=max_records) + self._pending: Optional[_PendingStep] = None + self._prev_stamp: Optional[float] = None + + self._d2h_stream: Optional[torch.cuda.Stream] = None + if self.enabled and InfoComponent.REQS in self._components: + self._d2h_stream = torch.cuda.Stream(device=device) + + self._current_segments: dict[ + InfoSegment, tuple[torch.cuda.Event, torch.cuda.Event] + ] = {} + self._open_segments: dict[InfoSegment, torch.cuda.Event] = {} + + def begin_step(self) -> None: + if not self.enabled: + return + self._current_segments = {} + self._open_segments = {} + if InfoComponent.STEP_GPU_TIME in self._components: + self._open_segment(InfoSegment.STEP) + + def segment(self, name: Union[InfoSegment, str]) -> ContextManager[None]: + if not self.enabled: + return _NULL_SEGMENT + segment = InfoSegment(name) + if not self._segment_enabled(segment): + return _NULL_SEGMENT + return self._active_segment(segment) + + @contextmanager + def _active_segment(self, segment: InfoSegment) -> Iterator[None]: + self._open_segment(segment) + try: + yield + finally: + self._close_segment(segment) + + def observe_decode_step(self, obs: DecodeStepObservation) -> None: + if not self.enabled: + return + if InfoComponent.STEP_GPU_TIME in self._components: + self._close_segment(InfoSegment.STEP) + + now = self._clock() + step_cpu_ms = self._step_cpu_ms(now=now) + self._drain_pending() + + future = ( + self._stage_reqs(obs) if InfoComponent.REQS in self._components else None + ) + self._pending = _PendingStep( + forward_ct=int(obs.forward_ct), + bs=int(obs.bs), + mode=obs.mode, + budget=None if obs.budget is None else int(obs.budget), + lag_steps=None if obs.lag_steps is None else int(obs.lag_steps), + num_verify_tokens=int(obs.num_verify_tokens), + verify_tokens_local=int(obs.verify_tokens_local), + verify_tokens_dp_synced=int(obs.verify_tokens_dp_synced), + verify_tokens_graph_key=int(obs.verify_tokens_graph_key), + predicted_step_ms=obs.predicted_step_ms, + predicted_theta=obs.predicted_theta, + step_cpu_ms=step_cpu_ms, + rids=obs.rids, + future=future, + segment_events=self._current_segments, + ) + self._current_segments = {} + self._prev_stamp = now + + def note_non_decode_step(self) -> None: + if not self.enabled: + return + self._drain_pending() + self._prev_stamp = None + self._current_segments = {} + self._open_segments = {} + + def flush(self) -> None: + if not self.enabled: + return + self._drain_pending() + + def clear(self) -> None: + self._records.clear() + self._pending = None + self._prev_stamp = None + self._current_segments = {} + self._open_segments = {} + self._sps_window = [] + self._sps_mismatched = 0 + + def dump(self) -> Optional[dict]: + if not self.enabled: + return None + self.flush() + return { + "mode": self.mode_value, + "gamma": self.gamma, + "verify_num_draft_tokens": self.verify_num_draft_tokens, + "components": sorted(component.value for component in self._components), + "records": [msgspec.to_builtins(record) for record in self._records], + } + + def _segment_enabled(self, segment: InfoSegment) -> bool: + if segment is InfoSegment.STEP: + return InfoComponent.STEP_GPU_TIME in self._components + if segment is InfoSegment.DRAFT: + return InfoComponent.DRAFT_GPU_TIME in self._components + if segment is InfoSegment.TARGET_VERIFY: + return InfoComponent.TARGET_VERIFY_GPU_TIME in self._components + return False + + def _open_segment(self, segment: InfoSegment) -> None: + start = torch.cuda.Event(enable_timing=True) + start.record() + self._open_segments[segment] = start + + def _close_segment(self, segment: InfoSegment) -> None: + start = self._open_segments.pop(segment, None) + if start is None: + return + end = torch.cuda.Event(enable_timing=True) + end.record() + self._current_segments[segment] = (start, end) + + def _stage_reqs(self, obs: DecodeStepObservation) -> Optional[FutureTensors]: + tensors: dict[str, torch.Tensor] = { + "req_pool_indices": obs.req_pool_indices, + "prefix_lens": obs.prefix_lens, + "draft_tokens": obs.draft_tokens, + "bonus_tokens": obs.bonus_tokens, + "correct_len": obs.correct_len, + "cap_trim_lens": obs.cap_trim_lens, + "commit_lens": obs.commit_lens, + } + if obs.verify_lens is not None: + tensors["verify_lens"] = obs.verify_lens + if obs.confidence is not None: + tensors["confidence"] = obs.confidence + return FutureTensors.device_to_host(tensors, d2h_stream=self._d2h_stream) + + def _drain_pending(self) -> None: + pending = self._pending + self._pending = None + if pending is None: + return + + record = DecodeStepRecord(forward_ct=pending.forward_ct) + if InfoComponent.CORE in self._components: + record.bs = pending.bs + record.mode = pending.mode + record.budget = pending.budget + record.lag_steps = pending.lag_steps + record.num_running_reqs = pending.bs + record.num_verify_tokens = pending.num_verify_tokens + record.verify_tokens_local = pending.verify_tokens_local + record.verify_tokens_dp_synced = pending.verify_tokens_dp_synced + record.verify_tokens_graph_key = pending.verify_tokens_graph_key + record.predicted_step_ms = pending.predicted_step_ms + record.predicted_theta = pending.predicted_theta + if InfoComponent.STEP_CPU_TIME in self._components: + record.step_cpu_ms = pending.step_cpu_ms + if InfoComponent.STEP_GPU_TIME in self._components: + record.step_gpu_ms = self._segment_ms(pending, InfoSegment.STEP) + if InfoComponent.DRAFT_GPU_TIME in self._components: + record.draft_gpu_ms = self._segment_ms(pending, InfoSegment.DRAFT) + if InfoComponent.TARGET_VERIFY_GPU_TIME in self._components: + record.target_verify_gpu_ms = self._segment_ms( + pending, InfoSegment.TARGET_VERIFY + ) + if InfoComponent.REQS in self._components and pending.future is not None: + record.reqs = self._build_reqs( + host=pending.future.wait(), bs=pending.bs, rids=pending.rids + ) + elif pending.future is not None: + pending.future.wait() + + self._records.append(record) + if self._sps_report_interval > 0: + self._report_sps_prediction(pending=pending, step_gpu_ms=record.step_gpu_ms) + + def _report_sps_prediction( + self, *, pending: _PendingStep, step_gpu_ms: Optional[float] + ) -> None: + predicted = pending.predicted_step_ms + if predicted is None or step_gpu_ms is None: + return + matched = ( + pending.budget is not None + and pending.bs + pending.budget == pending.num_verify_tokens + ) + if not matched: + self._sps_mismatched += 1 + return + self._sps_window.append((predicted, step_gpu_ms)) + if len(self._sps_window) < self._sps_report_interval: + return + + predictions = [p for p, _ in self._sps_window] + actuals = [a for _, a in self._sps_window] + abs_err = [abs(p - a) for p, a in self._sps_window] + rel_err = [abs(p - a) / a * 100 for p, a in self._sps_window if a > 0] + total = len(self._sps_window) + self._sps_mismatched + logger.info( + "DSpark SPS prediction: n=%d mean predicted=%.3fms mean actual=%.3fms " + "MAE=%.3fms median rel-err=%.1f%% mean bias(pred-actual)=%+.3fms " + "M_mismatch_rate=%.1f%% (%d/%d)", + len(self._sps_window), + statistics.fmean(predictions), + statistics.fmean(actuals), + statistics.fmean(abs_err), + statistics.median(rel_err) if rel_err else float("nan"), + statistics.fmean([p - a for p, a in self._sps_window]), + self._sps_mismatched / total * 100 if total else 0.0, + self._sps_mismatched, + total, + ) + self._sps_window = [] + self._sps_mismatched = 0 + + def _step_cpu_ms(self, *, now: float) -> Optional[float]: + prev = self._prev_stamp + if prev is None: + return None + step_cpu = now - prev + if not (0.0 < step_cpu <= self._max_step_cpu_seconds): + return None + return round(step_cpu * 1000.0, 4) + + def _segment_ms( + self, pending: _PendingStep, segment: InfoSegment + ) -> Optional[float]: + events = pending.segment_events.get(segment) + if events is None: + return None + start, end = events + end.synchronize() + elapsed_ms = start.elapsed_time(end) + if elapsed_ms > self._max_step_cpu_seconds * 1000.0: + return None + return round(elapsed_ms, 4) + + def _build_reqs( + self, *, host: dict, bs: int, rids: Optional[list[str]] + ) -> list[ReqDetail]: + req_ids = host["req_pool_indices"].tolist() + prefixes = host["prefix_lens"].tolist() + draft_rows = host["draft_tokens"].tolist() + bonus = host["bonus_tokens"].tolist() + correct = host["correct_len"].tolist() + cap_trim = host["cap_trim_lens"].tolist() + commit = host["commit_lens"].tolist() + verify_lens = host["verify_lens"].tolist() if "verify_lens" in host else None + if "confidence" in host: + conf_host = host["confidence"].float() + conf_rows = conf_host.tolist() + survival_rows = torch.cumprod(conf_host, dim=1).tolist() + else: + conf_rows = None + survival_rows = None + + reqs: list[ReqDetail] = [] + for row in range(bs): + verify_len = ( + self.verify_num_draft_tokens + if verify_lens is None + else int(verify_lens[row]) + ) + reqs.append( + ReqDetail( + rid=None if rids is None else rids[row], + req_pool_index=int(req_ids[row]), + prefix_len=int(prefixes[row]), + verify_len=verify_len, + acc_len=int(commit[row]), + correct_drafts=int(correct[row]), + cap_trim=int(cap_trim[row]), + bonus_token=int(bonus[row]), + draft_tokens=[int(t) for t in draft_rows[row]], + confidence=( + None + if conf_rows is None + else [round(float(p), 4) for p in conf_rows[row]] + ), + survival=( + None + if survival_rows is None + else [round(float(p), 4) for p in survival_rows[row]] + ), + ) + ) + return reqs + + +EPS_PROB = 1e-8 + + +def _format_float(value: float, digits: int = 4) -> str: + value = float(value) + if math.isnan(value): + return "nan" + return f"{value:.{digits}f}" + + +class PerPositionConfidenceMetrics: + + def __init__( + self, + *, + gamma: int, + device: torch.device, + num_coarse_bins: int = 15, + num_fine_bins: int = 1024, + ) -> None: + self.gamma = int(gamma) + self.num_coarse_bins = int(num_coarse_bins) + self.num_fine_bins = int(num_fine_bins) + self.coarse_count = torch.zeros( + (self.gamma, self.num_coarse_bins), dtype=torch.float64, device=device + ) + self.coarse_pred = torch.zeros_like(self.coarse_count) + self.coarse_target = torch.zeros_like(self.coarse_count) + self.fine_pos = torch.zeros( + (self.gamma, self.num_fine_bins), dtype=torch.float64, device=device + ) + self.fine_neg = torch.zeros_like(self.fine_pos) + self.brier_num = torch.zeros(self.gamma, dtype=torch.float64, device=device) + + def update(self, *, survival: torch.Tensor, prefix_mask: torch.Tensor) -> None: + assert survival.shape == prefix_mask.shape + assert survival.dim() == 2 and survival.shape[1] == self.gamma + + probs = survival.to(torch.float64).clamp(EPS_PROB, 1.0 - EPS_PROB) + targets = prefix_mask.to(torch.float64) + bs = probs.shape[0] + + probs_flat = probs.reshape(-1) + targets_flat = targets.reshape(-1) + weights = torch.ones_like(probs_flat) + pos_idx = ( + torch.arange(self.gamma, device=probs.device) + .view(1, -1) + .expand(bs, self.gamma) + .reshape(-1) + ) + + coarse_idx = ( + (probs_flat * self.num_coarse_bins) + .long() + .clamp_(0, self.num_coarse_bins - 1) + ) + flat_coarse = pos_idx * self.num_coarse_bins + coarse_idx + self.coarse_count.view(-1).scatter_add_(0, flat_coarse, weights) + self.coarse_pred.view(-1).scatter_add_(0, flat_coarse, probs_flat) + self.coarse_target.view(-1).scatter_add_(0, flat_coarse, targets_flat) + + fine_idx = ( + (probs_flat * self.num_fine_bins).long().clamp_(0, self.num_fine_bins - 1) + ) + flat_fine = pos_idx * self.num_fine_bins + fine_idx + self.fine_pos.view(-1).scatter_add_(0, flat_fine, targets_flat) + self.fine_neg.view(-1).scatter_add_(0, flat_fine, 1.0 - targets_flat) + + self.brier_num.add_((probs - targets).pow(2).sum(dim=0)) + + @staticmethod + def _auroc_from_hist(pos_hist: torch.Tensor, neg_hist: torch.Tensor) -> float: + total_pos = float(pos_hist.sum()) + total_neg = float(neg_hist.sum()) + if total_pos <= 0.0 or total_neg <= 0.0: + return float("nan") + cum_neg = torch.cumsum(neg_hist, dim=0) + cum_neg_before = cum_neg - neg_hist + pair = (pos_hist * cum_neg_before).sum() + 0.5 * (pos_hist * neg_hist).sum() + return float(pair) / (total_pos * total_neg) + + def compute(self) -> list[dict]: + coarse_count = self.coarse_count.cpu() + coarse_pred = self.coarse_pred.cpu() + coarse_target = self.coarse_target.cpu() + fine_pos = self.fine_pos.cpu() + fine_neg = self.fine_neg.cpu() + brier_num = self.brier_num.cpu() + + out: list[dict] = [] + for pos in range(self.gamma): + weights = coarse_count[pos] + total = float(weights.sum()) + if total <= 1e-12: + out.append( + { + "position": pos, + "total_weight": 0.0, + "ece": float("nan"), + "auc": float("nan"), + "brier": float("nan"), + "pred_mean": float("nan"), + "target_mean": float("nan"), + "reliability": [], + } + ) + continue + + denom = weights.clamp_min(1e-12) + avg_pred = coarse_pred[pos] / denom + avg_target = coarse_target[pos] / denom + bin_err = (avg_pred - avg_target).abs() + ece = float((bin_err * weights).sum()) / total + auc = self._auroc_from_hist(fine_pos[pos], fine_neg[pos]) + brier = float(brier_num[pos]) / total + reliability = [] + for bin_idx in range(self.num_coarse_bins): + weight = float(weights[bin_idx]) + if weight <= 0.0: + continue + reliability.append( + { + "bin": bin_idx, + "range": [ + bin_idx / self.num_coarse_bins, + (bin_idx + 1) / self.num_coarse_bins, + ], + "avg_pred": float(avg_pred[bin_idx]), + "avg_target": float(avg_target[bin_idx]), + "weight": weight, + } + ) + out.append( + { + "position": pos, + "total_weight": total, + "ece": ece, + "auc": auc, + "brier": brier, + "pred_mean": float(coarse_pred[pos].sum()) / total, + "target_mean": float(coarse_target[pos].sum()) / total, + "reliability": reliability, + } + ) + return out + + def format_table(self) -> str: + rows = self.compute() + header = ( + f"{'pos':>3} {'count':>12} {'pred':>8} {'target':>8} " + f"{'ece':>8} {'auc':>8} {'brier':>8}" + ) + lines = [ + "DSpark confidence-head per-position calibration " + "(cumprod survival vs leading-correct-prefix)", + header, + ] + for row in rows: + lines.append( + f"{row['position']:>3} {row['total_weight']:>12.0f} " + f"{_format_float(row['pred_mean']):>8} " + f"{_format_float(row['target_mean']):>8} " + f"{_format_float(row['ece']):>8} " + f"{_format_float(row['auc']):>8} " + f"{_format_float(row['brier']):>8}" + ) + return "\n".join(lines) + + +class ConfidenceMetricsProbe: + + def __init__( + self, + *, + gamma: int, + verify_num_draft_tokens: int, + tp_rank: int, + print_every: int = 256, + ) -> None: + self.gamma = int(gamma) + self.verify_num_draft_tokens = int(verify_num_draft_tokens) + self.tp_rank = int(tp_rank) + self.print_every = int(print_every) + self._metrics: Optional[PerPositionConfidenceMetrics] = None + self._step_ct: int = 0 + self._compact_warned: bool = False + + def maybe_observe( + self, + *, + carries_confidence: bool, + is_compact_mode: bool, + confidence_raw: Optional[torch.Tensor], + verify_ids_2d: torch.Tensor, + target_logits: torch.Tensor, + bs: int, + ) -> None: + if not envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS.get(): + return + if self.tp_rank != 0: + return + if not carries_confidence: + return + if is_compact_mode: + if not self._compact_warned: + logger.warning( + "SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS is ignored under " + "SGLANG_RAGGED_VERIFY_MODE=compact (padded verify rows corrupt the " + "per-position prefix label); run cap-accept or static to measure it." + ) + self._compact_warned = True + return + if confidence_raw is None: + return + + target_predict = torch.argmax(target_logits, dim=-1).view( + bs, self.verify_num_draft_tokens + ) + num_correct_drafts, _ = compute_dflash_correct_drafts_and_bonus( + candidates=verify_ids_2d, + target_predict=target_predict, + ) + positions = torch.arange(self.gamma, device=confidence_raw.device).view(1, -1) + prefix_mask = (positions < num_correct_drafts.view(-1, 1)).to(torch.float32) + survival = torch.cumprod(torch.sigmoid(confidence_raw.float()), dim=1) + + if self._metrics is None: + self._metrics = PerPositionConfidenceMetrics( + gamma=self.gamma, device=confidence_raw.device + ) + self._metrics.update(survival=survival, prefix_mask=prefix_mask) + self._step_ct += 1 + if self._step_ct % self.print_every == 0: + logger.info("%s", self._metrics.format_table()) + + +_STS_COLLECT_FLUSH_EVERY: int = 256 + + +class DsparkStepObservers: + """Facade over the per-step observability sinks (info dumper, confidence + probe, STS collection, block-accept estimator). The worker's decode path + makes one call per step; all sink gating and field derivation live here + so the hot path stays free of observer plumbing.""" + + def __init__( + self, + *, + planner, + gamma: int, + verify_num_draft_tokens: int, + tp_rank: int, + device, + simulate_acc_len: float, + ) -> None: + self._planner = planner + self._gamma = int(gamma) + self._verify_num_draft_tokens = int(verify_num_draft_tokens) + self._simulate_acc_len = float(simulate_acc_len) + + self._confidence_probe = ConfidenceMetricsProbe( + gamma=gamma, + verify_num_draft_tokens=verify_num_draft_tokens, + tp_rank=tp_rank, + ) + self._info_dumper = DsparkInfoDumper( + components=resolve_enabled_components(), + gamma=gamma, + verify_num_draft_tokens=verify_num_draft_tokens, + attn_tp_rank=get_parallel().attn_tp_rank, + device=device, + mode_value=planner.mode_value, + sps_report_interval=envs.SGLANG_DSPARK_LOG_SPS_PRED_INTERVAL.get(), + ) + self._block_accept_recorder = create_block_accept_estimate_recorder( + gamma=gamma, device=device, tp_rank=tp_rank + ) + if self._simulate_acc_len > 0 and self._block_accept_recorder is not None: + raise ValueError( + "SGLANG_DSPARK_BLOCK_ACCEPT_ESTIMATE_PATH cannot be combined with " + "SGLANG_SIMULATE_ACC_LEN (simulated correct_len breaks the " + "accept-probability bookkeeping of the estimator)." + ) + self._sts_collect_path = envs.SGLANG_DSPARK_STS_COLLECT_PATH.get() + self._sts_recorder: Optional[StsDataRecorder] = None + + # --- step lifecycle ------------------------------------------------- + + def begin_step(self) -> None: + self._info_dumper.begin_step() + + def segment(self, name: Union[InfoSegment, str]) -> ContextManager[None]: + return self._info_dumper.segment(name) + + def note_prefill_step(self) -> None: + self._info_dumper.note_non_decode_step() + if self._block_accept_recorder is not None: + self._block_accept_recorder.flush() + + def note_idle_decode_step(self) -> None: + self._info_dumper.note_non_decode_step() + + # --- scheduler-facing hooks ------------------------------------------ + + def dump_info_records(self) -> Optional[dict]: + dumped = self._info_dumper.dump() + if dumped is None: + return None + dumped["simulate_acc_len"] = ( + self._simulate_acc_len if self._simulate_acc_len > 0 else None + ) + return dumped + + def clear_info_records(self) -> None: + self._info_dumper.clear() + + def block_accept_estimate_log_suffix(self) -> Optional[str]: + if self._block_accept_recorder is None: + return None + return self._block_accept_recorder.estimate_log_suffix() + + def note_request_finished(self, *, rid: str, natural_stop: bool) -> None: + if self._block_accept_recorder is None: + return + self._block_accept_recorder.note_request_finished( + rid=rid, natural_stop=natural_stop + ) + + # --- per-step observation -------------------------------------------- + + def observe_verify_step( + self, + *, + forward_ct: int, + reqs, + bs: int, + proposal_folded: bool, + verify_ids_2d: torch.Tensor, + target_logits: Optional[torch.Tensor], + layout, + confidence: Optional[torch.Tensor], + prefix_lens: torch.Tensor, + draft_tokens: torch.Tensor, + draft_block, + sampling_info, + correct_len: torch.Tensor, + cap_trim_lens: torch.Tensor, + bonus: torch.Tensor, + commit_lens: torch.Tensor, + verify_token_budget: Optional[int], + req_pool_indices: torch.Tensor, + verify_tier_num_tokens: int, + dp_tier_num_tokens: Optional[int], + ) -> None: + planner = self._planner + if not proposal_folded: + self._maybe_record_sts_collect( + verify_ids_2d=verify_ids_2d, + target_logits=target_logits, + bs=bs, + ) + self._confidence_probe.maybe_observe( + carries_confidence=planner.carries_confidence, + is_compact_mode=planner.is_compact_mode, + confidence_raw=planner.last_confidence_raw, + verify_ids_2d=verify_ids_2d, + target_logits=target_logits, + bs=bs, + ) + if self._block_accept_recorder is not None and not proposal_folded: + self._block_accept_recorder.observe_verify_step( + forward_ct=forward_ct, + rids=[req.rid for req in reqs], + draft_tokens=draft_tokens, + corrected_logits=draft_block.corrected_logits, + draft_temperatures=draft_block.temperatures, + greedy_mask=draft_block.greedy_mask, + target_logits=target_logits, + target_temperatures=( + sampling_info.temperatures + if sampling_info is not None + else draft_block.temperatures + ), + truncated_sampling_mask=( + (sampling_info.top_ks != TOP_K_ALL) + | (sampling_info.top_ps != 1.0) + | (sampling_info.min_ps > 0) + if sampling_info is not None + else None + ), + logits_adjustments_are_noop=verify_logits_adjustments_are_noop( + sampling_info + ), + correct_len=correct_len, + cap_trim_lens=cap_trim_lens, + bonus=bonus, + prefix_lens=prefix_lens, + layout=layout, + ) + if self._info_dumper.enabled: + budget_decision = planner.take_budget_decision() + predicted_step_ms = ( + None + if budget_decision is None + or budget_decision.predicted_step_seconds is None + else budget_decision.predicted_step_seconds * 1e3 + ) + predicted_theta = ( + None if budget_decision is None else budget_decision.predicted_theta + ) + num_verify_tokens = ( + layout.graph_num_tokens + if layout is not None + else int(verify_ids_2d.numel()) + ) + self._info_dumper.observe_decode_step( + DecodeStepObservation( + forward_ct=forward_ct, + bs=bs, + mode=planner.mode_value, + budget=verify_token_budget, + lag_steps=planner.lag_steps, + num_verify_tokens=num_verify_tokens, + verify_tokens_local=verify_tier_num_tokens, + verify_tokens_dp_synced=( + -1 if dp_tier_num_tokens is None else int(dp_tier_num_tokens) + ), + verify_tokens_graph_key=num_verify_tokens, + predicted_step_ms=predicted_step_ms, + predicted_theta=predicted_theta, + verify_lens=layout.verify_lens if layout is not None else None, + confidence=confidence, + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + draft_tokens=draft_tokens, + bonus_tokens=bonus, + correct_len=correct_len, + cap_trim_lens=cap_trim_lens, + commit_lens=commit_lens, + rids=[req.rid for req in reqs], + ) + ) + + def _maybe_record_sts_collect( + self, + *, + verify_ids_2d: torch.Tensor, + target_logits: Optional[torch.Tensor], + bs: int, + ) -> None: + if not self._sts_collect_path: + return + if not self._planner.carries_confidence: + return + confidence_raw = self._planner.last_confidence_raw + if confidence_raw is None: + return + if self._sts_recorder is None: + self._sts_recorder = StsDataRecorder( + path_stem=self._sts_collect_path, + gamma=self._gamma, + flush_every=_STS_COLLECT_FLUSH_EVERY, + ) + target_predict = torch.argmax(target_logits, dim=-1).view( + bs, self._verify_num_draft_tokens + ) + num_correct_drafts, _ = compute_dflash_correct_drafts_and_bonus( + candidates=verify_ids_2d, + target_predict=target_predict, + ) + self._sts_recorder.record( + confidence_raw=confidence_raw, + num_correct_drafts=num_correct_drafts, + ) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_planner.py b/python/sglang/srt/speculative/dspark_components/dspark_planner.py new file mode 100644 index 000000000..d71f8b224 --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/dspark_planner.py @@ -0,0 +1,1117 @@ +from __future__ import annotations + +import logging +from typing import Optional, Union + +import msgspec +import torch + +from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func +from sglang.srt.distributed import get_tp_group +from sglang.srt.environ import envs +from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.managers.overlap_utils import ( + CONFIDENCE_RELAY_RING_LAG, + FutureMap, + ResolvedConfidence, +) +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.runtime_context import get_parallel +from sglang.srt.server_args import ServerArgs +from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 +from sglang.srt.speculative.dflash_utils import apply_dflash_verify_logits_adjustments +from sglang.srt.speculative.dspark_components.dspark_sps import ( + SpsAdditiveCostTable, + SpsCostTable, + _interp_clamped, + build_uninitialized_sps_table, + is_uninitialized_sps_table, + load_sps_table_from_path, +) +from sglang.srt.speculative.dspark_components.dspark_sts import ( + load_sts_calibration_from_path, +) +from sglang.srt.speculative.dspark_components.kernels.dspark_schedule import ( + ScheduleVerifyLensTopk, + compute_sort_survival, +) +from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyLayout, + RaggedVerifyMode, + read_ragged_verify_mode, + round_up_grid, +) +from sglang.srt.utils.async_probe import ( + maybe_assert_async, + maybe_detect_in_closed_range, +) +from sglang.srt.utils.common import require_mlp_tp_gather + +logger = logging.getLogger(__name__) + + +class VerifyWindow(msgspec.Struct, frozen=True): + positions_2d: torch.Tensor + verify_cache_loc: torch.Tensor + verify_cache_loc_2d: torch.Tensor + + +class DSparkVerifyPlanner: + def __init__( + self, + *, + draft_model, + gamma: int, + model_runner, + device, + tp_rank: int, + server_args: ServerArgs, + verify_num_draft_tokens: int, + ) -> None: + self.draft_model = draft_model + self.gamma = gamma + self.model_runner = model_runner + self.device = device + self.server_args = server_args + self.verify_num_draft_tokens = verify_num_draft_tokens + self._align_verify_tokens_to_graph_tier = ( + server_args.speculative_dspark_align_verify_tokens_to_graph_tier + ) + + self._confidence_head = getattr(self.draft_model, "confidence_head", None) + + sts_path = server_args.speculative_dspark_confidence_sts_path + if sts_path and self._confidence_head is not None: + calibration = load_sts_calibration_from_path(sts_path) + sts_temperatures = torch.tensor( + calibration.temperatures, dtype=torch.float32, device=device + ) + if envs.SGLANG_DSPARK_STS_COLLECT_PATH.get() and not bool( + torch.all(sts_temperatures == 1.0) + ): + raise ValueError( + "DSpark STS data collection (SGLANG_DSPARK_STS_COLLECT_PATH) " + "requires identity temperatures, but a non-identity calibration " + f"was loaded from {sts_path}. Collect pre-calibration logits with " + "no table (omit --speculative-dspark-confidence-sts-path)." + ) + if sts_temperatures.numel() != self.gamma: + raise ValueError( + "DSpark STS calibration was fit for gamma=" + f"{sts_temperatures.numel()} but the runtime gamma is " + f"{self.gamma}; refit the table for gamma={self.gamma} or omit " + "--speculative-dspark-confidence-sts-path." + ) + self._confidence_head.sts_temperatures = sts_temperatures + if tp_rank == 0: + logger.info( + "DSpark STS calibration loaded from %s (gamma=%d); per-position " + "temperatures applied to confidence-head survival.", + sts_path, + self.gamma, + ) + elif sts_path and self._confidence_head is None: + if tp_rank == 0: + logger.warning( + "DSpark STS calibration path given but no confidence head present " + "(static mode / head-less checkpoint); ignoring %s.", + sts_path, + ) + + self._ragged_verify_mode = read_ragged_verify_mode() + self._schedule_cfg = DSparkScheduleConfig(gamma=self.gamma) + self._budget_planner: Optional[HostConfidenceBudgetPlanner] = None + self._dynamic_graph_tier = False + self._dp_tier_gather_enabled = False + self._is_verify_all = True + if self._ragged_verify_mode is not RaggedVerifyMode.STATIC: + if self._confidence_head is None: + raise ValueError( + f"DSpark ragged-verify mode {self._ragged_verify_mode.value!r} " + f"schedules per-request verify lengths from the draft confidence " + f"head, but this DSpark draft checkpoint has no confidence head -- " + f"the checkpoint is wrong/incomplete (it ships no " + f"enable_confidence_head + trained confidence_head weights). Use a " + f"draft checkpoint that includes the confidence head, or run " + f"SGLANG_RAGGED_VERIFY_MODE=static." + ) + self._require_prep_in_cuda_graph() + sps_table = build_sps_cost_table( + server_args=self.server_args, + verify_num_draft_tokens=self.verify_num_draft_tokens, + ) + self._is_verify_all = ( + self._ragged_verify_mode is RaggedVerifyMode.COMPACT + and is_uninitialized_sps_table(sps_table) + ) + relay_lag_steps = ( + 0 + if self.server_args.disable_overlap_schedule + else CONFIDENCE_RELAY_RING_LAG + ) + self._budget_planner = HostConfidenceBudgetPlanner( + sps_table=sps_table, + cfg=self._schedule_cfg, + model_runner=self.model_runner, + relay_lag_steps=relay_lag_steps, + ) + self._dynamic_graph_tier = not is_dp_attention_enabled() + self._dp_tier_gather_enabled = ( + self._ragged_verify_mode is RaggedVerifyMode.COMPACT + and is_dp_attention_enabled() + and get_parallel().attn_tp_size == 1 + and get_parallel().attn_cp_size == 1 + and require_mlp_tp_gather(self.server_args) + and not self.server_args.disable_overlap_schedule + and not self.server_args.speculative_skip_dp_mlp_sync + and self.server_args.disaggregation_mode == "null" + and self.server_args.pp_size == 1 + and not envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get() + ) + if tp_rank == 0: + sps_table_source = ( + self.server_args.speculative_dspark_sps_table_path + or "uninitialized" + ) + logger.info( + "DSpark ragged-verify scheduler enabled (mode=%s, lag=%d, " + "relay_lag=%d, sps_table=%s, graph_tier=%s).", + self._ragged_verify_mode.value, + self._budget_planner.lag_steps, + relay_lag_steps, + sps_table_source, + ( + "dynamic" + if self._dynamic_graph_tier + else ( + "dp-gathered" if self._dp_tier_gather_enabled else "pinned" + ) + ), + ) + if isinstance(sps_table, SpsCostTable) and is_uninitialized_sps_table( + sps_table + ): + logger.warning( + "DSpark SPS table is uninitialized (flat): the verify " + "budget degenerates to verify-all (zero scheduling gain). " + "Pass a profiled --speculative-dspark-sps-table-path." + ) + + def _require_prep_in_cuda_graph(self) -> None: + if not envs.SGLANG_PREP_IN_CUDA_GRAPH.get(): + raise ValueError( + f"DSpark ragged-verify mode {self._ragged_verify_mode.value!r} " + f"requires SGLANG_PREP_IN_CUDA_GRAPH=1 (the captured-graph prepare " + f"path). It is currently disabled, which would put per-step " + f"verify_lens_cpu host reads on the critical path. Set " + f"SGLANG_PREP_IN_CUDA_GRAPH=1 or run SGLANG_RAGGED_VERIFY_MODE=static." + ) + + @property + def carries_confidence(self) -> bool: + return self._confidence_head is not None + + @property + def last_confidence_raw(self) -> Optional[torch.Tensor]: + if self._confidence_head is None: + return None + return self._confidence_head._last_confidence_raw + + @property + def schedules_verify_budget(self) -> bool: + return self._budget_planner is not None + + @property + def is_compact_mode(self) -> bool: + return self._ragged_verify_mode is RaggedVerifyMode.COMPACT + + @property + def is_verify_all(self) -> bool: + return self._is_verify_all + + @property + def mode_value(self) -> str: + return self._ragged_verify_mode.value + + @property + def lag_steps(self) -> Optional[int]: + if self._budget_planner is None: + return None + return self._budget_planner.lag_steps + + def take_budget_decision(self) -> Optional[VerifyBudgetDecision]: + if self._budget_planner is None: + return None + return self._budget_planner.take_last_decision() + + def should_run_compact(self, *, layout: Optional[RaggedVerifyLayout]) -> bool: + return ( + self._ragged_verify_mode is RaggedVerifyMode.COMPACT and layout is not None + ) + + def compute_confidence_tensor( + self, + *, + draft_hidden: Optional[torch.Tensor], + anchor_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + confidence_tap: Optional[torch.Tensor] = None, + ) -> Optional[torch.Tensor]: + if self._confidence_head is None: + return None + compute_confidence_hook = getattr(self.draft_model, "compute_confidence", None) + if compute_confidence_hook is not None: + assert ( + confidence_tap is not None + ), "dsv4 compute_confidence needs the compute_base_logits tap" + with torch.inference_mode(): + return compute_confidence_hook( + anchor_tokens=anchor_tokens, + sampled_tokens=draft_tokens, + x_post_hc=confidence_tap, + ) + assert draft_hidden is not None + return compute_confidence( + draft_hidden=draft_hidden, + anchor_tokens=anchor_tokens, + draft_tokens=draft_tokens, + confidence_head=self._confidence_head, + markov_head=self.draft_model.markov_head, + gamma=self.gamma, + ) + + def prepare_verify_budget( + self, batch: ScheduleBatch, future_map: FutureMap + ) -> None: + draft_input = batch.spec_info + if self._budget_planner is None: + return + if draft_input is None: + local_tier_num_tokens = 0 if batch.batch_size() == 0 else -1 + self._maybe_gather_dp_verify_tier( + batch=batch, local_tier_num_tokens=local_tier_num_tokens + ) + return + if batch.forward_mode.is_extend() or batch.is_extend_in_batch: + self._budget_planner.note_non_decode_step() + self._maybe_gather_dp_verify_tier(batch=batch, local_tier_num_tokens=0) + return + resolved = future_map.resolve_confidence_cpu(batch) + draft_input.verify_token_budget = self._budget_from_resolved( + resolved=resolved, req_pool_indices_cpu=batch.req_pool_indices_cpu + ) + batch.spec_verify_tier_num_tokens = local_verify_tier_num_tokens( + bs=batch.batch_size(), + verify_token_budget=draft_input.verify_token_budget, + verify_num_draft_tokens=self.verify_num_draft_tokens, + min_verify_len=self._schedule_cfg.min_verify_len, + ) + self._maybe_gather_dp_verify_tier( + batch=batch, local_tier_num_tokens=batch.spec_verify_tier_num_tokens + ) + + def _maybe_gather_dp_verify_tier( + self, *, batch: ScheduleBatch, local_tier_num_tokens: int + ) -> None: + if not self._dp_tier_gather_enabled: + return + if batch.is_extend_in_batch: + batch.global_spec_verify_tier_num_tokens = None + return + cpu_group = get_tp_group().cpu_group + local_tensor = torch.tensor([local_tier_num_tokens], dtype=torch.int64) + gathered = torch.empty( + (torch.distributed.get_world_size(group=cpu_group),), dtype=torch.int64 + ) + torch.distributed.all_gather_into_tensor( + gathered, local_tensor, group=cpu_group + ) + batch.global_spec_verify_tier_num_tokens = gathered.tolist() + + def note_non_decode_step(self) -> None: + if self._budget_planner is not None: + self._budget_planner.note_non_decode_step() + + def set_forced_budget_frac(self, frac) -> None: + if self._budget_planner is not None: + self._budget_planner.forced_budget_frac = frac + + def compute_budget_sync( + self, + *, + confidence: torch.Tensor, + prefix_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + ) -> Optional[int]: + del prefix_lens + if self._budget_planner is None: + return None + req_pool_indices_cpu = req_pool_indices.to("cpu").to(torch.int64) + generation = self.model_runner.req_to_token_pool.req_generation[ + req_pool_indices_cpu + ].clone() + resolved = ResolvedConfidence( + confidence=confidence.to("cpu"), + generation=generation, + ) + return self._budget_from_resolved( + resolved=resolved, req_pool_indices_cpu=req_pool_indices_cpu + ) + + def resolve_verify_token_budget( + self, + *, + draft_input: DFlashDraftInputV2, + confidence: Optional[torch.Tensor], + prefix_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + ) -> Optional[int]: + """Per-step verify-token budget: under overlap it was precomputed into + the draft input by prepare_verify_budget; otherwise compute it now.""" + if not self.schedules_verify_budget or confidence is None: + return None + if not self.server_args.disable_overlap_schedule: + return draft_input.verify_token_budget + return self.compute_budget_sync( + confidence=confidence, + prefix_lens=prefix_lens, + req_pool_indices=req_pool_indices, + ) + + def confidence_budget_prepare(self): + if not self.schedules_verify_budget: + return None + return self.prepare_verify_budget + + def _budget_from_resolved( + self, + *, + resolved: Optional[ResolvedConfidence], + req_pool_indices_cpu: torch.Tensor, + ) -> Optional[int]: + if resolved is None: + self._budget_planner.note_non_decode_step() + return None + current_generation = self.model_runner.req_to_token_pool.req_generation[ + req_pool_indices_cpu.to(torch.int64) + ] + return int( + self._budget_planner.compute_budget( + confidence=resolved.confidence, + generation=resolved.generation, + current_generation=current_generation, + req_pool_indices_cpu=req_pool_indices_cpu, + ) + ) + + def schedule_layout( + self, + *, + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + device: torch.device, + confidence: Optional[torch.Tensor], + budget: Optional[int], + global_num_reqs: Optional[int] = None, + dp_tier_num_tokens: Optional[int] = None, + ) -> Optional[RaggedVerifyLayout]: + if self._ragged_verify_mode is RaggedVerifyMode.STATIC: + return None + verify_lens = self._schedule_verify_lens( + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + device=device, + confidence=confidence, + budget=self._budget_aligned_to_graph_tier( + req_pool_indices=req_pool_indices, + budget=budget, + global_num_reqs=global_num_reqs, + dp_tier_num_tokens=dp_tier_num_tokens, + ), + ) + if verify_lens is None: + assert dp_tier_num_tokens is None, ( + "dp tier agreement present but local verify lens are None; " + "the gathered hint and the local budget diverged" + ) + if self._ragged_verify_mode is RaggedVerifyMode.COMPACT: + return uniform_ragged_layout( + bs=len(req_pool_indices), + device=device, + verify_num_draft_tokens=self.verify_num_draft_tokens, + ragged_verify_mode=self._ragged_verify_mode, + model_runner=self.model_runner, + tier_num_reqs=global_num_reqs, + ) + return None + bs = int(verify_lens.shape[0]) + tier_num_reqs = bs if global_num_reqs is None else global_num_reqs + if dp_tier_num_tokens is not None: + assert global_num_reqs is not None, ( + "dp tier agreement requires the dp-global request count; " + "keying the tier off the local bs diverges across ranks" + ) + tier_num_tokens = dp_tier_num_tokens + elif self._dynamic_graph_tier and budget is not None: + tier_num_tokens = local_verify_tier_num_tokens( + bs=tier_num_reqs, + verify_token_budget=budget, + verify_num_draft_tokens=self.verify_num_draft_tokens, + min_verify_len=self._schedule_cfg.min_verify_len, + ) + else: + tier_num_tokens = None + if ragged_layout_exceeds_captured_grid( + num_reqs=tier_num_reqs, + verify_num_draft_tokens=self.verify_num_draft_tokens, + model_runner=self.model_runner, + tier_tokens_hint=tier_num_tokens, + ): + return None + graph_num_tokens_floor = verify_layout_graph_num_tokens_floor( + num_reqs=tier_num_reqs, + ragged_verify_mode=self._ragged_verify_mode, + verify_num_draft_tokens=self.verify_num_draft_tokens, + model_runner=self.model_runner, + tier_num_tokens=tier_num_tokens, + ) + capture_num_tokens = ragged_capture_num_tokens(model_runner=self.model_runner) + if graph_num_tokens_floor > 0 and capture_num_tokens is not None: + graph_num_tokens = round_up_grid(graph_num_tokens_floor, capture_num_tokens) + return RaggedVerifyLayout.from_verify_lens_device( + verify_lens=verify_lens, graph_num_tokens=graph_num_tokens + ) + verify_lens_cpu = verify_lens.to("cpu").tolist() + grid = verify_layout_grid( + verify_lens_cpu=verify_lens_cpu, + ragged_verify_mode=self._ragged_verify_mode, + model_runner=self.model_runner, + ) + return RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=verify_lens_cpu, + device=device, + grid=grid, + graph_num_tokens_floor=graph_num_tokens_floor, + ) + + def _budget_aligned_to_graph_tier( + self, + *, + req_pool_indices: torch.Tensor, + budget: Optional[int], + global_num_reqs: Optional[int], + dp_tier_num_tokens: Optional[int], + ) -> Optional[int]: + # Flag off (default): returns budget unchanged, so the schedule below is + # byte-for-byte the original. On: ceils role 1's verify-token total up to the + # padded graph tier graph_num_tokens = round_up(dp-max tier, captured token + # bucket), which folds in the cuda-graph bucket round-up (H1) and the dp + # cross-rank max (H2); role 2 (the single top-k) then admits that many real + # draft tokens. graph_num_tokens is derived from the same (request count, + # gathered dp tier, original budget) inputs the layout below uses, so the two + # agree by construction -- this only feeds the larger budget into the top-k, + # it does not touch the layout's own tier computation. + if not self._align_verify_tokens_to_graph_tier or budget is None: + return budget + tier_num_reqs = ( + int(req_pool_indices.shape[0]) + if global_num_reqs is None + else global_num_reqs + ) + if dp_tier_num_tokens is not None: + tier_num_tokens = dp_tier_num_tokens + elif self._dynamic_graph_tier: + tier_num_tokens = local_verify_tier_num_tokens( + bs=tier_num_reqs, + verify_token_budget=budget, + verify_num_draft_tokens=self.verify_num_draft_tokens, + min_verify_len=self._schedule_cfg.min_verify_len, + ) + else: + tier_num_tokens = None + graph_num_tokens_floor = verify_layout_graph_num_tokens_floor( + num_reqs=tier_num_reqs, + ragged_verify_mode=self._ragged_verify_mode, + verify_num_draft_tokens=self.verify_num_draft_tokens, + model_runner=self.model_runner, + tier_num_tokens=tier_num_tokens, + ) + capture_num_tokens = ragged_capture_num_tokens(model_runner=self.model_runner) + if graph_num_tokens_floor <= 0 or capture_num_tokens is None: + return budget + graph_num_tokens = round_up_grid(graph_num_tokens_floor, capture_num_tokens) + return graph_tier_fill_budget( + graph_num_tokens=graph_num_tokens, + bs=int(req_pool_indices.shape[0]), + verify_num_draft_tokens=self.verify_num_draft_tokens, + min_verify_len=self._schedule_cfg.min_verify_len, + ) + + def _schedule_verify_lens( + self, + *, + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + device: torch.device, + confidence: Optional[torch.Tensor], + budget: Optional[int], + ) -> Optional[torch.Tensor]: + if self._budget_planner is None or confidence is None or budget is None: + return None + verify_lens = ScheduleVerifyLensTopk.execute( + confidence=confidence, + budget=budget, + cfg=self._schedule_cfg, + ).to(device=device, dtype=torch.int32) + + if envs.SGLANG_ENABLE_ASYNC_ASSERT.get(): + verify_lens_64 = verify_lens.to(torch.int64) + effective_floor = max(self._schedule_cfg.min_verify_len, 1) + maybe_assert_async( + (verify_lens_64 - effective_floor).sum() <= budget, + f"DSpark verify-len budget violated (budget={budget})", + ) + + if envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_PREFIX_SCHEDULER.get(): + self._log_verify_lens_decision( + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + budget=budget, + sort_survival=compute_sort_survival(confidence), + verify_lens=verify_lens, + ) + + broadcast_group, group_size = verify_lens_broadcast_group( + tp_size=self.server_args.tp_size + ) + if group_size > 1: + broadcast_group.broadcast(verify_lens, src=0) + + return verify_lens + + def _log_verify_lens_decision( + self, + *, + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + budget: int, + sort_survival: torch.Tensor, + verify_lens: torch.Tensor, + ) -> None: + cfg = self._schedule_cfg + max_len = cfg.resolved_max_verify_len() + req_ids = req_pool_indices.tolist() + prefixes = prefix_lens.tolist() + lens = verify_lens.tolist() + sort_rows = sort_survival.to(torch.float32).tolist() + logger.info( + "[DSPARK-CPS] num_reqs=%d budget=%d gamma=%d verify_len_range=[%d,%d]", + len(req_ids), + budget, + cfg.gamma, + cfg.min_verify_len, + max_len, + ) + for row in range(len(req_ids)): + survival_str = "[" + ", ".join(f"{p:.3f}" for p in sort_rows[row]) + "]" + logger.info( + "[DSPARK-CPS] req=%d prefix=%d verify_len=%d sort_survival=%s", + int(req_ids[row]), + int(prefixes[row]), + int(lens[row]), + survival_str, + ) + + +def local_verify_tier_num_tokens( + *, + bs: int, + verify_token_budget: Optional[int], + verify_num_draft_tokens: int, + min_verify_len: int, +) -> int: + if verify_token_budget is None: + return -1 + floor_tokens = bs * max(min_verify_len, 1) + return min(floor_tokens + verify_token_budget, bs * verify_num_draft_tokens) + + +def graph_tier_fill_budget( + *, + graph_num_tokens: int, + bs: int, + verify_num_draft_tokens: int, + min_verify_len: int, +) -> int: + # top-k budget (tokens above the per-request floor) that makes the scheduled + # total reach the padded graph tier, capped at bs * verify_num_draft_tokens + # since a request cannot verify more than its proposed drafts. Inverse of + # local_verify_tier_num_tokens: total = floor_tokens + budget. + fill_total = min(graph_num_tokens, bs * verify_num_draft_tokens) + floor_tokens = bs * max(min_verify_len, 1) + return max(0, fill_total - floor_tokens) + + +def dp_global_verify_tier_num_tokens( + *, + global_tier_num_tokens: Optional[list[int]], +) -> Optional[int]: + if global_tier_num_tokens is None: + return None + if any(tier_num_tokens < 0 for tier_num_tokens in global_tier_num_tokens): + return None + max_tier_num_tokens = max(global_tier_num_tokens, default=0) + return max_tier_num_tokens if max_tier_num_tokens > 0 else None + + +def idle_ragged_layout( + *, + tier_num_reqs: int, + dp_tier_num_tokens: Optional[int], + device: torch.device, + verify_num_draft_tokens: int, + model_runner, +) -> Optional[RaggedVerifyLayout]: + if ragged_capture_num_tokens(model_runner=model_runner) is None: + dp_tier_num_tokens = None + if dp_tier_num_tokens is None: + return uniform_ragged_layout( + bs=tier_num_reqs, + device=device, + verify_num_draft_tokens=verify_num_draft_tokens, + ragged_verify_mode=RaggedVerifyMode.COMPACT, + model_runner=model_runner, + ) + if ragged_layout_exceeds_captured_grid( + num_reqs=tier_num_reqs, + verify_num_draft_tokens=verify_num_draft_tokens, + model_runner=model_runner, + tier_tokens_hint=dp_tier_num_tokens, + ): + return None + verify_lens_cpu = [1] * tier_num_reqs + grid = verify_layout_grid( + verify_lens_cpu=verify_lens_cpu, + ragged_verify_mode=RaggedVerifyMode.COMPACT, + model_runner=model_runner, + ) + return RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=verify_lens_cpu, + device=device, + grid=grid, + graph_num_tokens_floor=dp_tier_num_tokens, + ) + + +def uniform_ragged_layout( + *, + bs: int, + device: torch.device, + verify_num_draft_tokens: int, + ragged_verify_mode: RaggedVerifyMode, + model_runner, + tier_num_reqs: Optional[int] = None, +) -> Optional[RaggedVerifyLayout]: + tier_num_reqs = bs if tier_num_reqs is None else tier_num_reqs + if ragged_layout_exceeds_captured_grid( + num_reqs=tier_num_reqs, + verify_num_draft_tokens=verify_num_draft_tokens, + model_runner=model_runner, + ): + return None + verify_lens_cpu = [verify_num_draft_tokens] * bs + grid = verify_layout_grid( + verify_lens_cpu=verify_lens_cpu, + ragged_verify_mode=ragged_verify_mode, + model_runner=model_runner, + ) + graph_num_tokens_floor = verify_layout_graph_num_tokens_floor( + num_reqs=tier_num_reqs, + ragged_verify_mode=ragged_verify_mode, + verify_num_draft_tokens=verify_num_draft_tokens, + model_runner=model_runner, + ) + return RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=verify_lens_cpu, + device=device, + grid=grid, + graph_num_tokens_floor=graph_num_tokens_floor, + ) + + +def verify_lens_broadcast_group(*, tp_size: int) -> tuple: + if is_dp_attention_enabled(): + return get_parallel().attn_tp_group, get_parallel().attn_tp_size + return get_tp_group(), tp_size + + +def verify_layout_grid( + *, + verify_lens_cpu: list[int], + ragged_verify_mode: RaggedVerifyMode, + model_runner, +) -> list[int]: + total = sum(verify_lens_cpu) + if ragged_verify_mode is not RaggedVerifyMode.COMPACT: + return [total] + capture_num_tokens = ragged_capture_num_tokens(model_runner=model_runner) + if capture_num_tokens is None: + return [total] + return capture_num_tokens + + +def verify_layout_graph_num_tokens_floor( + *, + num_reqs: int, + ragged_verify_mode: RaggedVerifyMode, + verify_num_draft_tokens: int, + model_runner, + tier_num_tokens: Optional[int] = None, +) -> int: + if ( + ragged_verify_mode is not RaggedVerifyMode.COMPACT + or ragged_capture_num_tokens(model_runner=model_runner) is None + ): + return 0 + if tier_num_tokens is not None: + return min(tier_num_tokens, num_reqs * verify_num_draft_tokens) + return num_reqs * verify_num_draft_tokens + + +def ragged_capture_num_tokens(*, model_runner) -> Optional[list[int]]: + runner = model_runner.decode_cuda_graph_runner + if runner is None or not runner.ragged_verify_mode: + return None + return runner.capture_num_tokens + + +def ragged_capture_max_slots(*, model_runner) -> Optional[int]: + runner = model_runner.decode_cuda_graph_runner + if runner is None or not runner.ragged_verify_mode: + return None + return runner.max_bs + + +def ragged_layout_exceeds_captured_grid( + *, + num_reqs: int, + verify_num_draft_tokens: int, + model_runner, + tier_tokens_hint: Optional[int] = None, +) -> bool: + capture_num_tokens = ragged_capture_num_tokens(model_runner=model_runner) + if capture_num_tokens is None: + return False + max_slots = ragged_capture_max_slots(model_runner=model_runner) + if max_slots is not None and num_reqs > max_slots: + return True + tier_tokens = ( + tier_tokens_hint + if tier_tokens_hint is not None + else num_reqs * verify_num_draft_tokens + ) + return tier_tokens > capture_num_tokens[-1] + + +def alloc_verify_window( + *, + batch: ScheduleBatch, + bs: int, + device: str, + verify_num_draft_tokens: int, + block_pos_offsets: torch.Tensor, + model_runner, +) -> VerifyWindow: + prefix_lens = batch.seq_lens + verify_w = verify_num_draft_tokens + positions_2d = prefix_lens.unsqueeze(1) + block_pos_offsets + verify_cache_loc = assign_extend_cache_locs_func( + req_pool_indices=batch.req_pool_indices, + req_to_token=model_runner.req_to_token_pool.req_to_token, + start_offset=prefix_lens, + end_offset=prefix_lens + verify_w, + batch_size=bs, + draft_token_num=verify_w, + device=device, + ) + verify_cache_loc_2d = verify_cache_loc.view(bs, verify_w) + return VerifyWindow( + positions_2d=positions_2d, + verify_cache_loc=verify_cache_loc, + verify_cache_loc_2d=verify_cache_loc_2d, + ) + + +def apply_logits_adjustments_strided( + *, + next_token_logits: torch.Tensor, + sampling_info, + verify_num_draft_tokens: int, +) -> None: + if sampling_info is None: + return + apply_dflash_verify_logits_adjustments( + next_token_logits=next_token_logits, + sampling_info=sampling_info, + draft_token_num=verify_num_draft_tokens, + ) + + +def build_markov_embed_stack( + *, + anchor_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + markov_head, + gamma: int, +) -> torch.Tensor: + prev_seq = torch.cat( + [anchor_tokens.view(-1, 1), draft_tokens[:, : gamma - 1]], dim=1 + ) + return markov_head.get_prev_embeddings(prev_seq) + + +def compute_confidence( + *, + draft_hidden: torch.Tensor, + anchor_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + confidence_head, + markov_head, + gamma: int, +) -> torch.Tensor: + assert confidence_head is not None + if confidence_head.with_markov: + markov_embed_stack = build_markov_embed_stack( + anchor_tokens=anchor_tokens, + draft_tokens=draft_tokens, + markov_head=markov_head, + gamma=gamma, + ) + else: + markov_embed_stack = None + confidence_raw = confidence_head(draft_hidden, markov_embed_stack) + confidence = confidence_head.apply_sts(confidence_raw) + maybe_detect_in_closed_range(confidence, 0.0, 1.0, "DSpark confidence") + return confidence + + +class DSparkScheduleConfig(msgspec.Struct): + gamma: int + min_verify_len: int = 1 + max_verify_len: int = 0 + survival_eps: float = 1e-6 + + def resolved_max_verify_len(self) -> int: + return self.max_verify_len or (self.gamma + 1) + + def validate(self) -> None: + max_len = self.resolved_max_verify_len() + if self.gamma < 1: + raise ValueError(f"DSpark gamma must be >= 1, got {self.gamma}.") + if not (0 <= self.min_verify_len <= max_len <= self.gamma + 1): + raise ValueError( + "DSpark verify-len config must satisfy 0 <= min <= max <= gamma+1, " + f"got min={self.min_verify_len}, max={max_len}, gamma={self.gamma}." + ) + if self.survival_eps < 0: + raise ValueError(f"survival_eps must be >= 0, got {self.survival_eps}.") + + +class VerifyBudgetDecision(msgspec.Struct): + budget: int + predicted_step_seconds: Optional[float] = None + predicted_theta: Optional[float] = None + + +def compute_verify_token_budget( + *, + history_survival_probs: torch.Tensor, + sps_table: Union[SpsCostTable, SpsAdditiveCostTable], + cfg: DSparkScheduleConfig, +) -> VerifyBudgetDecision: + num_requests = history_survival_probs.shape[0] + max_len = cfg.resolved_max_verify_len() + + candidates = history_survival_probs[:, :max_len].flatten() + candidates = candidates[candidates >= cfg.survival_eps].to(torch.float64) + candidates_sorted = torch.sort(candidates, descending=True).values + prefix_sum = torch.cumsum(candidates_sorted, dim=0) + + tau_star = num_requests + torch.cat( + [torch.zeros(1, dtype=torch.float64), prefix_sum] + ) + if isinstance(sps_table, SpsAdditiveCostTable): + step_time = _additive_step_time_tensor( + table=sps_table, + num_requests=int(num_requests), + num_budgets=int(tau_star.numel()), + ) + theta = tau_star / step_time + idx = int(torch.argmax(theta)) + predicted_step_seconds = float(step_time[idx]) + else: + batch_tokens = num_requests + torch.arange(tau_star.numel(), dtype=torch.int64) + sps = _lookup_sps_tensor(sps_table=sps_table, batch_tokens=batch_tokens) + theta = tau_star * sps + idx = int(torch.argmax(theta)) + sps_at_idx = float(sps[idx]) + predicted_step_seconds = 1.0 / sps_at_idx if sps_at_idx > 0 else None + return VerifyBudgetDecision( + budget=idx, + predicted_step_seconds=predicted_step_seconds, + predicted_theta=float(theta[idx]), + ) + + +def _lookup_sps_tensor( + *, sps_table: SpsCostTable, batch_tokens: torch.Tensor +) -> torch.Tensor: + probes = torch.tensor(sps_table.sample_batch_tokens, dtype=torch.int64) + sps = torch.tensor(sps_table.sample_steps_per_sec, dtype=torch.float64) + idx = torch.bucketize(batch_tokens, probes, right=True) - 1 + idx = idx.clamp_(0, probes.numel() - 1) + return sps[idx] + + +def _additive_step_time_tensor( + *, table: SpsAdditiveCostTable, num_requests: int, num_budgets: int +) -> torch.Tensor: + floor = table.bias_seconds + _interp_clamped( + table.bs_probes, table.alpha_seconds, float(num_requests) + ) + m_probes = torch.tensor(table.m_probes, dtype=torch.float64) + theta_vals = torch.tensor(table.theta_seconds, dtype=torch.float64) + m = (num_requests + torch.arange(num_budgets, dtype=torch.float64)).clamp_( + min=float(table.m_probes[0]), max=float(table.m_probes[-1]) + ) + hi = torch.bucketize(m, m_probes, right=True).clamp_(1, m_probes.numel() - 1) + lo = hi - 1 + span = (m_probes[hi] - m_probes[lo]).clamp_(min=1e-9) + frac = (m - m_probes[lo]) / span + theta_at_m = theta_vals[lo] + frac * (theta_vals[hi] - theta_vals[lo]) + return floor + theta_at_m + + +class HostConfidenceBudgetPlanner: + + def __init__( + self, + *, + sps_table: SpsCostTable, + cfg: DSparkScheduleConfig, + model_runner, + relay_lag_steps: int = 1, + ) -> None: + cfg.validate() + self.sps_table = sps_table + self.cfg = cfg + self._model_runner = model_runner + self.forced_budget_frac: Optional[float] = None + self.last_decision: Optional[VerifyBudgetDecision] = None + self.lag_steps = max( + int(envs.SGLANG_DSPARK_CONFIDENCE_RELAY_LAG_STEPS.get()), 1 + ) + self.carry_steps = max(self.lag_steps - int(relay_lag_steps), 0) + self._carry_confidence: Optional[torch.Tensor] = None + self._carry_generation: Optional[torch.Tensor] = None + self._carry_pos = 0 + + def compute_budget( + self, + *, + confidence: torch.Tensor, + generation: torch.Tensor, + current_generation: torch.Tensor, + req_pool_indices_cpu: torch.Tensor, + ) -> int: + lagged_confidence, lagged_generation = self._shift_to_lag( + confidence=confidence, + generation=generation, + req_pool_indices_cpu=req_pool_indices_cpu, + ) + survival = self._two_steps_prior_survival( + lagged_confidence=lagged_confidence, + lagged_generation=lagged_generation, + current_generation=current_generation, + ) + forced_frac = self.forced_budget_frac + if forced_frac is not None: + full_budget = int(survival[:, : self.cfg.resolved_max_verify_len()].numel()) + forced_budget = max(0, int(float(forced_frac) * full_budget)) + self.last_decision = VerifyBudgetDecision(budget=forced_budget) + return forced_budget + decision = compute_verify_token_budget( + history_survival_probs=survival, + sps_table=self.sps_table, + cfg=self.cfg, + ) + self.last_decision = decision + return decision.budget + + def take_last_decision(self) -> Optional[VerifyBudgetDecision]: + decision = self.last_decision + self.last_decision = None + return decision + + def note_non_decode_step(self) -> None: + self.last_decision = None + + def _shift_to_lag( + self, + *, + confidence: torch.Tensor, + generation: torch.Tensor, + req_pool_indices_cpu: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.carry_steps == 0: + return confidence, generation + self._ensure_carry(gamma=confidence.shape[-1]) + slot = self._carry_pos % self.carry_steps + rows = req_pool_indices_cpu.to(torch.int64) + lagged_confidence = self._carry_confidence[slot, rows].clone() + lagged_generation = self._carry_generation[slot, rows].clone() + self._carry_confidence[slot, rows] = confidence.to(torch.float32) + self._carry_generation[slot, rows] = generation.to(torch.int64) + self._carry_pos += 1 + return lagged_confidence, lagged_generation + + def _two_steps_prior_survival( + self, + *, + lagged_confidence: torch.Tensor, + lagged_generation: torch.Tensor, + current_generation: torch.Tensor, + ) -> torch.Tensor: + k_survival = torch.cumprod(lagged_confidence.to(torch.float32), dim=1) + current_gen = current_generation.to(torch.int64) + fresh = ( + (current_gen >= 1) & (lagged_generation.to(torch.int64) == current_gen) + ).view(-1, 1) + return torch.where(fresh, k_survival, torch.ones_like(k_survival)) + + def _ensure_carry(self, *, gamma: int) -> None: + if self._carry_confidence is not None: + return + req_pool_size = int(self._model_runner.req_to_token_pool.req_to_token.shape[0]) + self._carry_confidence = torch.zeros( + (self.carry_steps, req_pool_size, gamma), dtype=torch.float32 + ) + self._carry_generation = torch.zeros( + (self.carry_steps, req_pool_size), + dtype=torch.int64, + ) + + +def build_sps_cost_table( + *, + server_args: ServerArgs, + verify_num_draft_tokens: int, +) -> Union[SpsCostTable, SpsAdditiveCostTable]: + sps_table_path = server_args.speculative_dspark_sps_table_path + if sps_table_path: + return load_sps_table_from_path(sps_table_path) + max_batch_tokens = max( + 1, + int(server_args.max_running_requests or 1) * verify_num_draft_tokens, + ) + return build_uninitialized_sps_table(max_batch_tokens=max_batch_tokens) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_sps.py b/python/sglang/srt/speculative/dspark_components/dspark_sps.py new file mode 100644 index 000000000..f19483f2d --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/dspark_sps.py @@ -0,0 +1,164 @@ +from __future__ import annotations + +import bisect +from typing import Optional + +import msgspec + + +def floor_probe_index(edges: list[int], batch_tokens: int) -> int: + idx = bisect.bisect_right(edges, batch_tokens) - 1 + return max(0, min(idx, len(edges) - 1)) + + +class SpsCostTable(msgspec.Struct, frozen=True): + sample_batch_tokens: list[int] + sample_steps_per_sec: list[float] + max_batch_tokens: int + + def __post_init__(self) -> None: + if not self.sample_batch_tokens: + raise ValueError("SpsCostTable requires at least one probe.") + if self.sample_batch_tokens != sorted(set(self.sample_batch_tokens)): + raise ValueError( + "sample_batch_tokens must be strictly increasing (monotone-sorted " + f"invariant), got {self.sample_batch_tokens}." + ) + if len(self.sample_batch_tokens) != len(self.sample_steps_per_sec): + raise ValueError( + "sample_batch_tokens and sample_steps_per_sec must have equal length, " + f"got {len(self.sample_batch_tokens)} vs {len(self.sample_steps_per_sec)}." + ) + if self.max_batch_tokens < self.sample_batch_tokens[-1]: + raise ValueError( + "max_batch_tokens must be >= the largest probe, got " + f"{self.max_batch_tokens} < {self.sample_batch_tokens[-1]}." + ) + + def lookup(self, batch_tokens: int) -> float: + return self.sample_steps_per_sec[ + floor_probe_index(self.sample_batch_tokens, batch_tokens) + ] + + def to_json(self) -> str: + return msgspec.json.encode(self).decode("utf-8") + + @classmethod + def from_json(cls, data: str) -> SpsCostTable: + return msgspec.json.decode(data.encode("utf-8"), type=cls) + + +def _interp_clamped(xs: list[int], ys: list[float], x: float) -> float: + if x <= xs[0]: + return ys[0] + if x >= xs[-1]: + return ys[-1] + hi = bisect.bisect_right(xs, x) + lo = hi - 1 + frac = (x - xs[lo]) / (xs[hi] - xs[lo]) + return ys[lo] + frac * (ys[hi] - ys[lo]) + + +class SpsAdditiveCostTable(msgspec.Struct, frozen=True): + + bias_seconds: float + bs_probes: list[int] + alpha_seconds: list[float] + m_probes: list[int] + theta_seconds: list[float] + + def __post_init__(self) -> None: + for name, probes, values in ( + ("bs", self.bs_probes, self.alpha_seconds), + ("m", self.m_probes, self.theta_seconds), + ): + if not probes: + raise ValueError(f"SpsAdditiveCostTable requires {name}_probes.") + if probes != sorted(set(probes)): + raise ValueError( + f"{name}_probes must be strictly increasing, got {probes}." + ) + if len(probes) != len(values): + raise ValueError( + f"{name}_probes and its values must have equal length, got " + f"{len(probes)} vs {len(values)}." + ) + if self.bias_seconds <= 0: + raise ValueError(f"bias_seconds must be > 0, got {self.bias_seconds}.") + + def step_time(self, *, num_reqs: int, budget: int) -> float: + return ( + self.bias_seconds + + _interp_clamped(self.bs_probes, self.alpha_seconds, float(num_reqs)) + + _interp_clamped( + self.m_probes, self.theta_seconds, float(num_reqs + budget) + ) + ) + + def to_json(self) -> str: + return msgspec.json.encode(self).decode("utf-8") + + @classmethod + def from_json(cls, data: str) -> SpsAdditiveCostTable: + return msgspec.json.decode(data.encode("utf-8"), type=cls) + + +def profile_sps_table( + *, + probes: list[tuple[int, float]], + max_batch_tokens: Optional[int] = None, +) -> SpsCostTable: + if not probes: + raise ValueError("profile_sps_table requires at least one probe.") + + sorted_probes = sorted(probes, key=lambda probe: probe[0]) + + sample_batch_tokens: list[int] = [] + sample_steps_per_sec: list[float] = [] + for batch_tokens, steps_per_sec in sorted_probes: + batch_tokens = int(batch_tokens) + if batch_tokens < 1: + raise ValueError( + f"profile_sps_table requires batch_tokens >= 1, got {batch_tokens}." + ) + if sample_batch_tokens and batch_tokens == sample_batch_tokens[-1]: + raise ValueError( + "profile_sps_table requires unique batch_tokens per probe; " + f"batch_tokens={batch_tokens} appears more than once. Median the " + "repeated samples per batch_tokens before calling the assembler." + ) + sample_batch_tokens.append(batch_tokens) + sample_steps_per_sec.append(float(steps_per_sec)) + + resolved_max = ( + int(max_batch_tokens) + if max_batch_tokens is not None + else sample_batch_tokens[-1] + ) + return SpsCostTable( + sample_batch_tokens=sample_batch_tokens, + sample_steps_per_sec=sample_steps_per_sec, + max_batch_tokens=resolved_max, + ) + + +def load_sps_table_from_path(path: str): + with open(path, "r", encoding="utf-8") as f: + data = f.read() + if '"bias_seconds"' in data: + return SpsAdditiveCostTable.from_json(data) + return SpsCostTable.from_json(data) + + +def build_uninitialized_sps_table(*, max_batch_tokens: int) -> SpsCostTable: + return SpsCostTable( + sample_batch_tokens=[1], + sample_steps_per_sec=[1.0], + max_batch_tokens=max_batch_tokens, + ) + + +def is_uninitialized_sps_table(table: SpsCostTable | SpsAdditiveCostTable) -> bool: + if isinstance(table, SpsAdditiveCostTable): + return False + return len(table.sample_batch_tokens) <= 1 diff --git a/python/sglang/srt/speculative/dspark_components/dspark_sts.py b/python/sglang/srt/speculative/dspark_components/dspark_sts.py new file mode 100644 index 000000000..01d1a8842 --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/dspark_sts.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from pathlib import Path + +import msgspec +import torch + + +class DSparkStsCalibration(msgspec.Struct, frozen=True, omit_defaults=True): + temperatures: list[float] + dataset: str = "" + num_samples: int = 0 + ece_before: list[float] = [] + ece_after: list[float] = [] + + def __post_init__(self) -> None: + if not self.temperatures: + raise ValueError("DSparkStsCalibration requires at least one temperature.") + for temperature in self.temperatures: + if temperature <= 0: + raise ValueError( + "DSparkStsCalibration temperatures must all be > 0, got " + f"{self.temperatures}." + ) + + def to_json(self) -> str: + return msgspec.json.encode(self).decode("utf-8") + + @classmethod + def from_json(cls, data: str) -> DSparkStsCalibration: + return msgspec.json.decode(data.encode("utf-8"), type=cls) + + +def load_sts_calibration_from_path(path: str) -> DSparkStsCalibration: + with open(path, "r", encoding="utf-8") as f: + return DSparkStsCalibration.from_json(f.read()) + + +class StsDataRecorder: + def __init__(self, *, path_stem: str, gamma: int, flush_every: int) -> None: + self.path_stem = path_stem + self.gamma = int(gamma) + self.flush_every = int(flush_every) + self._logits_buffer: list[torch.Tensor] = [] + self._prefix_mask_buffer: list[torch.Tensor] = [] + self._shard_ct = 0 + + def record( + self, *, confidence_raw: torch.Tensor, num_correct_drafts: torch.Tensor + ) -> None: + logits = confidence_raw.detach().to(device="cpu", dtype=torch.float32) + positions = torch.arange(self.gamma).view(1, -1) + counts = ( + num_correct_drafts.detach().to(device="cpu", dtype=torch.int64).view(-1, 1) + ) + prefix_mask = (positions < counts).to(torch.float32) + self._logits_buffer.append(logits) + self._prefix_mask_buffer.append(prefix_mask) + if len(self._logits_buffer) >= self.flush_every: + self.flush() + + def flush(self) -> None: + if not self._logits_buffer: + return + shard_path = Path(f"{self.path_stem}.{self._shard_ct}.pt") + shard_path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "logits": torch.cat(self._logits_buffer, dim=0), + "prefix_mask": torch.cat(self._prefix_mask_buffer, dim=0), + }, + shard_path, + ) + self._logits_buffer.clear() + self._prefix_mask_buffer.clear() + self._shard_ct += 1 diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py new file mode 100644 index 000000000..43099c9fe --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -0,0 +1,716 @@ +from __future__ import annotations + +from typing import Optional + +import msgspec +import torch + +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode +from sglang.srt.speculative.dflash_info import DFlashVerifyInput +from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 +from sglang.srt.speculative.dflash_utils import apply_dflash_verify_logits_adjustments +from sglang.srt.speculative.dspark_components.dspark_draft import DraftBlockResult +from sglang.srt.speculative.dspark_components.dspark_kv_inject import ( + TargetHiddenKvInjector, +) +from sglang.srt.speculative.dspark_components.dspark_planner import ( + VerifyWindow, + apply_logits_adjustments_strided, +) +from sglang.srt.speculative.dspark_components.kernels.dspark_accept import ( + AcceptGreedy, + AcceptSampling, + FinalizeAcceptLens, + SelectMixedAccept, + SoftmaxTemp, + accept_greedy_triton, + finalize_accept_lens_triton, +) +from sglang.srt.speculative.dspark_components.kernels.dspark_verify_window import ( + BuildCommitInjectLayout, + BuildOutTokens, + BuildRaggedVerifyWindow, + RaggedVerifyWindow, + ScatterCompactToStrided, + scatter_compact_to_strided_into, +) +from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout + + +def verify_logits_adjustments_are_noop(sampling_info) -> bool: + if sampling_info is None: + return True + if sampling_info.has_custom_logit_processor: + return False + if getattr(sampling_info, "acc_linear_penalties", None) is not None: + return False + penalizer = getattr(sampling_info, "penalizer_orchestrator", None) + if penalizer is not None and penalizer.is_required: + return False + if getattr(sampling_info, "vocab_mask", None) is not None: + return False + if getattr(sampling_info, "logit_bias", None) is not None: + return False + return True + + +class TargetVerifyResult(msgspec.Struct, frozen=True): + logits_output: object + can_run_cuda_graph: bool + + +class TargetVerifyExecutor: + def __init__( + self, + *, + target_worker, + gamma: int, + verify_num_draft_tokens: int, + model_runner, + kv_injector: TargetHiddenKvInjector, + verify_epilogue=None, + simulate_acc_len: float = 0.0, + ) -> None: + self.target_worker = target_worker + self.gamma = int(gamma) + self.verify_num_draft_tokens = verify_num_draft_tokens + self.model_runner = model_runner + self.kv_injector = kv_injector + self.verify_epilogue = verify_epilogue + self._verify_backend_self_adds_seq_lens_cache: Optional[bool] = None + self._simulate_acc_len = float(simulate_acc_len) + self._simulated_correct_drafts_buf: Optional[torch.Tensor] = None + + def accept_and_finalize( + self, + *, + folded_accept: bool, + bs: int, + verify_ids_2d: torch.Tensor, + target_logits: Optional[torch.Tensor], + draft_block: DraftBlockResult, + sampling_info, + draft_input: DFlashDraftInputV2, + layout: Optional[RaggedVerifyLayout], + prefix_lens: torch.Tensor, + draft_tokens: torch.Tensor, + ) -> AcceptOuts: + """Produce the per-request accept outcome after target verify. + + Folded path: the accept/finalize/out-token kernels already ran inside + the target-verify cuda graph (DsparkVerifyEpilogue); read its buffers. + Eager path: run them here, including the SGLANG_SIMULATE_ACC_LEN + override. + """ + if folded_accept: + return self.verify_epilogue.read_accept(bs) + + correct_len, bonus, cap_trim_lens = accept_draft_tokens( + candidates=verify_ids_2d, + target_logits=target_logits, + draft_block=draft_block, + sampling_info=sampling_info, + draft_input=draft_input, + gamma=self.gamma, + verify_num_draft_tokens=self.verify_num_draft_tokens, + cutoff_layout=layout, + ) + if self._simulate_acc_len > 0: + correct_len = self._simulated_correct_len( + bs=bs, dtype=correct_len.dtype, device=correct_len.device + ) + + finalized = FinalizeAcceptLens.execute( + correct_len=correct_len, + cap_trim_lens=cap_trim_lens, + prefix_lens=prefix_lens, + ) + out_tokens = BuildOutTokens.execute( + draft_tokens=draft_tokens, + correct_len=correct_len, + bonus=bonus, + verify_num_draft_tokens=self.verify_num_draft_tokens, + gamma=self.gamma, + ) + return AcceptOuts( + correct_len=correct_len, + bonus=bonus, + cap_trim_lens=finalized.cap_trim_lens, + commit_lens=finalized.commit_lens, + new_seq_lens=finalized.new_seq_lens, + out_tokens=out_tokens, + ) + + def _simulated_correct_len( + self, *, bs: int, dtype: torch.dtype, device: torch.device + ) -> torch.Tensor: + buf = self._simulated_correct_drafts_buf + if buf is None or buf.numel() < bs or buf.dtype != dtype: + correct_target = int( + round(min(max(self._simulate_acc_len - 1.0, 0.0), float(self.gamma))) + ) + buf = torch.full( + (max(bs, 512),), correct_target, dtype=dtype, device=device + ) + self._simulated_correct_drafts_buf = buf + return buf[:bs] + + def run_idle_participation( + self, + *, + batch: ScheduleBatch, + idle_layout: Optional[RaggedVerifyLayout], + ) -> None: + """Run a dummy target-verify forward so an idle DP rank joins the + token-keyed collective ops of the busy ranks' verify step.""" + device = self.model_runner.device + if self.verify_epilogue is not None: + self.verify_epilogue.begin_step(None, armed=False) + num_dummy_tokens = ( + idle_layout.graph_num_tokens if idle_layout is not None else 0 + ) + verify_input = DFlashVerifyInput( + draft_token=torch.zeros( + (num_dummy_tokens,), dtype=torch.int64, device=device + ), + positions=torch.zeros( + (num_dummy_tokens,), dtype=torch.int64, device=device + ), + draft_token_num=self.verify_num_draft_tokens, + custom_mask=None, + capture_hidden_mode=CaptureHiddenMode.FULL, + ragged_verify_layout=idle_layout, + ) + batch.out_cache_loc = torch.zeros( + (num_dummy_tokens,), dtype=torch.int64, device=device + ) + if idle_layout is not None: + num_dummy_slots = int(idle_layout.verify_lens.numel()) + batch.seq_lens = torch.ones( + (num_dummy_slots,), dtype=torch.int64, device=device + ) + batch.req_pool_indices = torch.zeros( + (num_dummy_slots,), dtype=torch.int64, device=device + ) + batch.seq_lens_cpu = torch.ones((num_dummy_slots,), dtype=torch.int64) + batch.seq_lens_sum = num_dummy_slots + batch.forward_mode = ForwardMode.TARGET_VERIFY + verify_forward_batch, _ = verify_input.prepare_for_verify( + batch, self.target_worker + ) + self.target_worker.forward_batch_generation( + batch=None, + forward_batch=verify_forward_batch, + is_verify=True, + skip_attn_backend_init=True, + ) + + def run_non_compact( + self, + *, + batch: ScheduleBatch, + draft_input: DFlashDraftInputV2, + verify_ids_2d: torch.Tensor, + verify_window: VerifyWindow, + sampling_info, + ) -> TargetVerifyResult: + verify_w = self.verify_num_draft_tokens + positions_2d = verify_window.positions_2d + verify_cache_loc = verify_window.verify_cache_loc + + verify_input = DFlashVerifyInput( + draft_token=verify_ids_2d.reshape(-1), + positions=positions_2d.reshape(-1), + draft_token_num=verify_w, + custom_mask=None, + capture_hidden_mode=CaptureHiddenMode.FULL, + ) + batch.out_cache_loc = verify_cache_loc + seq_lens_cpu_backup = batch.seq_lens_cpu + seq_lens_sum_backup = batch.seq_lens_sum + if not self._verify_backend_self_adds_seq_lens(): + if seq_lens_cpu_backup is not None: + batch.seq_lens_cpu = seq_lens_cpu_backup + verify_w + batch.seq_lens_sum = int(batch.seq_lens_cpu.sum()) + elif draft_input.reserved_seq_lens_cpu is not None: + batch.seq_lens_cpu = draft_input.reserved_seq_lens_cpu + batch.seq_lens_sum = int(draft_input.reserved_seq_lens_sum) + + result = self._forward_prepared_verify( + batch=batch, + verify_input=verify_input, + seq_lens_cpu_backup=seq_lens_cpu_backup, + seq_lens_sum_backup=seq_lens_sum_backup, + ) + + if sampling_info is not None: + apply_dflash_verify_logits_adjustments( + next_token_logits=result.logits_output.next_token_logits, + sampling_info=sampling_info, + draft_token_num=verify_w, + ) + + return result + + def _forward_prepared_verify( + self, + *, + batch: ScheduleBatch, + verify_input: DFlashVerifyInput, + seq_lens_cpu_backup, + seq_lens_sum_backup, + ) -> TargetVerifyResult: + verify_forward_batch, _ = verify_input.prepare_for_verify( + batch, self.target_worker + ) + batch.seq_lens_cpu = seq_lens_cpu_backup + batch.seq_lens_sum = seq_lens_sum_backup + + target_out = self.target_worker.forward_batch_generation( + batch=None, + forward_batch=verify_forward_batch, + is_verify=True, + skip_attn_backend_init=True, + ) + return TargetVerifyResult( + logits_output=target_out.logits_output, + can_run_cuda_graph=target_out.can_run_cuda_graph, + ) + + def commit_hidden( + self, + *, + batch: ScheduleBatch, + layout: Optional[RaggedVerifyLayout], + hidden_strided: Optional[torch.Tensor], + verify_window: VerifyWindow, + logits_output, + commit_lens: torch.Tensor, + bs: int, + run_compact: bool, + ) -> None: + if run_compact: + self.kv_injector.inject_ragged( + batch=batch, + layout=layout, + hidden_strided=hidden_strided, + commit_lens=commit_lens, + bs=bs, + ) + return + hidden = logits_output.hidden_states + if hidden is None: + raise RuntimeError("DSpark verify requires target hidden states, got None.") + hidden = hidden.view(bs, self.verify_num_draft_tokens, -1) + self.kv_injector.inject_target_hidden( + target_hidden=hidden.reshape(-1, hidden.shape[-1]), + cache_loc=verify_window.verify_cache_loc, + cache_loc_2d=verify_window.verify_cache_loc_2d, + positions=verify_window.positions_2d.reshape(-1), + commit_lens=commit_lens, + ) + + def _run_ragged( + self, + *, + batch: ScheduleBatch, + layout: RaggedVerifyLayout, + ragged_window: RaggedVerifyWindow, + sampling_info, + ) -> TargetVerifyResult: + verify_input = DFlashVerifyInput( + draft_token=ragged_window.verify_ids, + positions=ragged_window.positions, + draft_token_num=self.verify_num_draft_tokens, + custom_mask=None, + capture_hidden_mode=CaptureHiddenMode.FULL, + ragged_verify_layout=layout, + ) + batch.out_cache_loc = ragged_window.verify_cache_loc + seq_lens_cpu_backup = batch.seq_lens_cpu + seq_lens_sum_backup = batch.seq_lens_sum + if seq_lens_cpu_backup is not None: + verify_lens_cpu = ( + layout.verify_lens_cpu + if layout.verify_lens_cpu is not None + else layout.verify_lens.cpu().tolist() + ) + batch.seq_lens_cpu = seq_lens_cpu_backup + torch.tensor( + verify_lens_cpu, dtype=seq_lens_cpu_backup.dtype + ) + batch.seq_lens_sum = int(batch.seq_lens_cpu.sum()) + + return self._forward_prepared_verify( + batch=batch, + verify_input=verify_input, + seq_lens_cpu_backup=seq_lens_cpu_backup, + seq_lens_sum_backup=seq_lens_sum_backup, + ) + + def run_compact( + self, + *, + batch: ScheduleBatch, + layout: RaggedVerifyLayout, + draft_block_ids: torch.Tensor, + draft_tokens: torch.Tensor, + bs: int, + device: str, + sampling_info, + inject_gate: bool = False, + ) -> tuple[TargetVerifyResult, torch.Tensor]: + ragged_window = BuildRaggedVerifyWindow.execute( + batch=batch, + layout=layout, + draft_block_ids=draft_block_ids, + draft_tokens=draft_tokens, + bs=bs, + device=device, + verify_num_draft_tokens=self.verify_num_draft_tokens, + model_runner=self.model_runner, + ) + if self.verify_epilogue is not None: + self.verify_epilogue.begin_step(layout.verify_lens, armed=inject_gate) + target_verify = self._run_ragged( + batch=batch, + layout=layout, + ragged_window=ragged_window, + sampling_info=sampling_info, + ) + logits_output = target_verify.logits_output + + stride = self.verify_num_draft_tokens + if self.verify_epilogue is not None and target_verify.can_run_cuda_graph: + strided_logits = self.verify_epilogue.strided_logits + hidden_strided = self.verify_epilogue.strided_hidden + assert strided_logits is not None and hidden_strided is not None, ( + "verify epilogue buffers unwritten after a graph replay -- the " + "replayed graph was captured without the epilogue" + ) + strided_logits = strided_logits[: bs * stride] + hidden_strided = hidden_strided[: bs * stride] + else: + compact_logits = logits_output.next_token_logits + strided_logits = ScatterCompactToStrided.execute( + compact=compact_logits, + layout=layout, + fill_value=0.0, + verify_num_draft_tokens=stride, + ) + compact_hidden = logits_output.hidden_states + if compact_hidden is None: + raise RuntimeError( + "DSpark verify requires target hidden states, got None." + ) + hidden_strided = ScatterCompactToStrided.execute( + compact=compact_hidden, + layout=layout, + fill_value=0.0, + verify_num_draft_tokens=stride, + ) + apply_logits_adjustments_strided( + next_token_logits=strided_logits, + sampling_info=sampling_info, + verify_num_draft_tokens=stride, + ) + logits_output.next_token_logits = strided_logits + logits_output.hidden_states = hidden_strided + return target_verify, hidden_strided + + def _verify_backend_self_adds_seq_lens(self) -> bool: + if self._verify_backend_self_adds_seq_lens_cache is None: + backend = self.target_worker.model_runner.attn_backend + self._verify_backend_self_adds_seq_lens_cache = hasattr( + backend, "make_forward_metadata_from_raw_verify" + ) + return self._verify_backend_self_adds_seq_lens_cache + + +class CommitInjectCtx(msgspec.Struct): + + draft_model: object + block_pos_offsets: torch.Tensor + resolve_pool: object + resolve_req_to_token: object + + +class AcceptOuts(msgspec.Struct): + correct_len: torch.Tensor + bonus: torch.Tensor + cap_trim_lens: torch.Tensor + commit_lens: torch.Tensor + new_seq_lens: torch.Tensor + out_tokens: torch.Tensor + + +class DsparkVerifyEpilogue: + + def __init__( + self, + *, + max_bs: int, + verify_num_draft_tokens: int, + device, + commit_ctx: Optional[CommitInjectCtx] = None, + ) -> None: + self.max_bs = int(max_bs) + self.stride = int(verify_num_draft_tokens) + self.gamma = self.stride - 1 + self.commit_ctx = commit_ctx + self.inject_gate_buf = torch.zeros((1,), dtype=torch.int32, device=device) + self.verify_lens_buf = torch.zeros( + (self.max_bs,), dtype=torch.int64, device=device + ) + self.draft_tokens_buf = torch.zeros( + (self.max_bs * self.gamma,), dtype=torch.int64, device=device + ) + self.correct_len_buf = torch.zeros( + (self.max_bs,), dtype=torch.int64, device=device + ) + self.bonus_buf = torch.zeros((self.max_bs,), dtype=torch.int64, device=device) + self.cap_trim_lens_buf = torch.zeros( + (self.max_bs,), dtype=torch.int32, device=device + ) + self.commit_lens_buf = torch.zeros( + (self.max_bs,), dtype=torch.int32, device=device + ) + self.new_seq_lens_buf = torch.zeros( + (self.max_bs,), dtype=torch.int64, device=device + ) + self.out_tokens_buf = torch.zeros( + (self.max_bs, self.stride), dtype=torch.int64, device=device + ) + self.strided_logits: Optional[torch.Tensor] = None + self.strided_hidden: Optional[torch.Tensor] = None + + def capture_hook(self, runner, out, forward_batch, num_tokens) -> None: + if runner.model_runner.is_draft_worker or not runner.ragged_verify_mode: + return + if ( + not isinstance(out, LogitsProcessorOutput) + or out.next_token_logits is None + or out.hidden_states is None + ): + return + self( + compact_logits=out.next_token_logits, + compact_hidden=out.hidden_states, + input_ids=forward_batch.input_ids, + seq_lens=forward_batch.seq_lens, + req_pool_indices=forward_batch.req_pool_indices, + bs=forward_batch.batch_size, + ) + + def begin_step(self, verify_lens, armed: bool) -> None: + if verify_lens is None: + self.verify_lens_buf.zero_() + else: + bs = verify_lens.shape[0] + self.verify_lens_buf[:bs].copy_(verify_lens) + if bs < self.max_bs: + self.verify_lens_buf[bs:].zero_() + self.inject_gate_buf.fill_(1 if armed else 0) + + def read_accept(self, bs: int) -> AcceptOuts: + return AcceptOuts( + correct_len=self.correct_len_buf[:bs], + bonus=self.bonus_buf[:bs], + cap_trim_lens=self.cap_trim_lens_buf[:bs], + commit_lens=self.commit_lens_buf[:bs], + new_seq_lens=self.new_seq_lens_buf[:bs], + out_tokens=self.out_tokens_buf[:bs], + ) + + @property + def folds_commit(self) -> bool: + if self.commit_ctx is None: + return False + pool = self.commit_ctx.resolve_pool() + return hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope") + + def _ensure_out( + self, buf: Optional[torch.Tensor], compact: torch.Tensor + ) -> torch.Tensor: + if ( + buf is not None + and buf.dtype == compact.dtype + and buf.shape[1] == compact.shape[1] + ): + return buf + assert not torch.cuda.is_current_stream_capturing(), ( + "DsparkVerifyEpilogue output buffers must be allocated during " + "warmup, not inside graph capture (pool memory is unreadable " + "post-replay)." + ) + return torch.empty( + (self.max_bs * self.stride, compact.shape[1]), + dtype=compact.dtype, + device=compact.device, + ) + + def __call__( + self, + *, + compact_logits: torch.Tensor, + compact_hidden: torch.Tensor, + input_ids: torch.Tensor, + seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + bs: int, + ) -> None: + self.strided_logits = self._ensure_out(self.strided_logits, compact_logits) + self.strided_hidden = self._ensure_out(self.strided_hidden, compact_hidden) + verify_lens = self.verify_lens_buf[:bs] + self._scatter(compact_logits, compact_hidden, verify_lens, bs) + commit_lens = self._accept(input_ids, seq_lens, verify_lens, bs) + if self.folds_commit: + self._commit_inject( + commit_lens, verify_lens, seq_lens, req_pool_indices, bs + ) + + def _scatter(self, compact_logits, compact_hidden, verify_lens, bs: int) -> None: + scatter_compact_to_strided_into( + compact=compact_logits, + verify_lens=verify_lens, + out=self.strided_logits[: bs * self.stride], + stride=self.stride, + fill_value=0.0, + ) + scatter_compact_to_strided_into( + compact=compact_hidden, + verify_lens=verify_lens, + out=self.strided_hidden[: bs * self.stride], + stride=self.stride, + fill_value=0.0, + ) + + def _accept(self, input_ids, seq_lens, verify_lens, bs: int) -> torch.Tensor: + candidates = torch.zeros( + (bs * self.stride, 1), dtype=input_ids.dtype, device=input_ids.device + ) + scatter_compact_to_strided_into( + compact=input_ids.view(-1, 1), + verify_lens=verify_lens, + out=candidates, + stride=self.stride, + fill_value=0, + ) + correct_len, bonus, cap_trim_lens = accept_greedy_triton( + candidates=candidates.view(bs, self.stride), + target_logits=self.strided_logits[: bs * self.stride], + verify_num_draft_tokens=self.stride, + cutoff_verify_lens=verify_lens, + ) + finalized = finalize_accept_lens_triton( + correct_len=correct_len, + cap_trim_lens=cap_trim_lens, + prefix_lens=seq_lens[:bs], + ) + out_tokens = BuildOutTokens.execute( + draft_tokens=self.draft_tokens_buf[: bs * self.gamma].view(bs, self.gamma), + correct_len=correct_len, + bonus=bonus, + verify_num_draft_tokens=self.stride, + gamma=self.gamma, + ) + self.correct_len_buf[:bs].copy_(correct_len) + self.bonus_buf[:bs].copy_(bonus) + self.cap_trim_lens_buf[:bs].copy_(cap_trim_lens.to(torch.int32)) + self.commit_lens_buf[:bs].copy_(finalized.commit_lens) + self.new_seq_lens_buf[:bs].copy_(finalized.new_seq_lens) + self.out_tokens_buf[:bs].copy_(out_tokens.view(bs, self.stride)) + return finalized.commit_lens + + def _commit_inject( + self, commit_lens, verify_lens, seq_lens, req_pool_indices, bs: int + ) -> None: + ctx = self.commit_ctx + pool = ctx.resolve_pool() + gated_commit_lens = ( + torch.minimum(commit_lens, verify_lens.to(torch.int32)) + * self.inject_gate_buf + ) + inject_layout = BuildCommitInjectLayout.execute( + req_pool_indices=req_pool_indices, + req_to_token=ctx.resolve_req_to_token(), + prefix_lens=seq_lens[:bs], + block_pos_offsets=ctx.block_pos_offsets[: self.stride], + full_to_swa_mapping=pool.full_to_swa_index_mapping, + commit_lens=gated_commit_lens, + stride=self.stride, + ) + with torch.inference_mode(): + ctx.draft_model.write_target_hidden_kv( + main_hidden=self.strided_hidden[: bs * self.stride], + swa_loc=inject_layout.swa_loc, + positions=inject_layout.positions, + pool=pool, + ) + + +def accept_draft_tokens( + *, + candidates: torch.Tensor, + target_logits: torch.Tensor, + draft_block: DraftBlockResult, + sampling_info, + draft_input: DFlashDraftInputV2, + gamma: int, + verify_num_draft_tokens: int, + cutoff_layout: Optional[RaggedVerifyLayout] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + greedy_mask = draft_block.greedy_mask + cutoff_verify_lens = None if cutoff_layout is None else cutoff_layout.verify_lens + all_greedy = sampling_info is None or sampling_info.is_all_greedy + if all_greedy: + return AcceptGreedy.execute( + candidates=candidates, + target_logits=target_logits, + verify_num_draft_tokens=verify_num_draft_tokens, + cutoff_verify_lens=cutoff_verify_lens, + ) + bs, gamma_rows, vocab = draft_block.corrected_logits.shape + draft_probs = SoftmaxTemp.execute( + logits=draft_block.corrected_logits.reshape(bs * gamma_rows, vocab), + temperatures=draft_block.temperatures, + rows_per_request=gamma_rows, + ).view(bs, gamma_rows, vocab) + if not sampling_info.is_any_greedy: + return AcceptSampling.execute( + candidates=candidates, + target_logits=target_logits, + draft_probs=draft_probs, + sampling_info=sampling_info, + draft_input=draft_input, + gamma=gamma, + verify_num_draft_tokens=verify_num_draft_tokens, + cutoff_verify_lens=cutoff_verify_lens, + ) + greedy_len, greedy_bonus, greedy_trim = AcceptGreedy.execute( + candidates=candidates, + target_logits=target_logits, + verify_num_draft_tokens=verify_num_draft_tokens, + cutoff_verify_lens=cutoff_verify_lens, + ) + sampling_len, sampling_bonus, sampling_trim = AcceptSampling.execute( + candidates=candidates, + target_logits=target_logits, + draft_probs=draft_probs, + sampling_info=sampling_info, + draft_input=draft_input, + gamma=gamma, + verify_num_draft_tokens=verify_num_draft_tokens, + cutoff_verify_lens=cutoff_verify_lens, + ) + selected = SelectMixedAccept.execute( + greedy_mask=greedy_mask, + greedy_len=greedy_len, + greedy_bonus=greedy_bonus, + greedy_trim=greedy_trim, + sampling_len=sampling_len, + sampling_bonus=sampling_bonus, + sampling_trim=sampling_trim, + ) + return selected.correct_len, selected.bonus, selected.cap_trim_lens diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py new file mode 100644 index 000000000..ee525f02a --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -0,0 +1,693 @@ +import logging +from contextlib import nullcontext +from typing import Optional + +import torch + +from sglang.srt.environ import envs +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.managers.scheduler import GenerationBatchResult +from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + compute_position, +) +from sglang.srt.runtime_context import get_parallel +from sglang.srt.server_args import ServerArgs +from sglang.srt.speculative.base_spec_worker import BaseSpecWorker +from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 +from sglang.srt.speculative.draft_worker_common import ( + build_block_pos_offsets, + build_draft_tp_worker, + make_draft_block_spec_info, + make_draft_sampler_capture_hook, +) +from sglang.srt.speculative.dspark_components.dspark_config import ( + DSV4_DRAFT_ATTENTION_BACKEND, + draft_is_deepseek_v4, + resolve_runtime_config, +) +from sglang.srt.speculative.dspark_components.dspark_draft import ( + DraftBlockProposer, + make_next_draft_input, + maybe_build_draft_sampler, +) +from sglang.srt.speculative.dspark_components.dspark_kv_inject import ( + TargetHiddenKvInjector, +) +from sglang.srt.speculative.dspark_components.dspark_observability import ( + DsparkStepObservers, + InfoSegment, +) +from sglang.srt.speculative.dspark_components.dspark_planner import ( + DSparkVerifyPlanner, + alloc_verify_window, + dp_global_verify_tier_num_tokens, + idle_ragged_layout, +) +from sglang.srt.speculative.dspark_components.dspark_verify import ( + CommitInjectCtx, + DsparkVerifyEpilogue, + TargetVerifyExecutor, + verify_logits_adjustments_are_noop, +) +from sglang.srt.speculative.spec_utils import draft_tp_context +from sglang.srt.utils import get_available_gpu_memory, is_cuda + +logger = logging.getLogger(__name__) + + +class DSparkWorkerV2(BaseSpecWorker): + + def __init__( + self, + server_args: ServerArgs, + gpu_id: int, + tp_rank: int, + dp_rank: Optional[int], + moe_ep_rank: int, + attn_cp_rank: int, + moe_dp_rank: int, + nccl_port: int, + target_worker: TpModelWorker, + ): + self.server_args = server_args + self.gpu_id = gpu_id + self.tp_rank = tp_rank + self.dp_rank = dp_rank + self.moe_ep_rank = moe_ep_rank + self.attn_cp_rank = attn_cp_rank + self.moe_dp_rank = moe_dp_rank + self.nccl_port = nccl_port + self._target_worker = target_worker + self.model_runner = target_worker.model_runner + self.page_size = server_args.page_size + self.device = target_worker.device + + self._draft_is_moe = draft_is_deepseek_v4(server_args=server_args) + self._draft_dp_context_enabled = ( + server_args.enable_dp_attention and not self._draft_is_moe + ) + attn_tp_size = server_args.tp_size // max(server_args.dp_size, 1) + if server_args.enable_dp_attention and self._draft_is_moe and attn_tp_size > 1: + raise ValueError( + "DSpark + dp attention with a DeepSeek-V4 (MoE) draft requires " + "attn_tp == 1 (set --dp-size == --tp). attn_tp > 1 corrupts the " + "MoE-under-DP all-reduce." + ) + + with self._draft_context(): + bundle = build_draft_tp_worker( + server_args=server_args, + gpu_id=gpu_id, + tp_rank=tp_rank, + dp_rank=dp_rank, + moe_ep_rank=moe_ep_rank, + attn_cp_rank=attn_cp_rank, + moe_dp_rank=moe_dp_rank, + nccl_port=nccl_port, + target_model_config=target_worker.model_runner.model_config, + algo_label="DSPARK", + attention_backend_override=( + DSV4_DRAFT_ATTENTION_BACKEND if self._draft_is_moe else None + ), + ) + self._draft_worker = bundle.draft_worker + self.draft_model_runner = bundle.draft_model_runner + self.draft_model = bundle.draft_model + self._draft_sampler = None + + runtime_config = resolve_runtime_config( + draft_hf_config=self.draft_model_runner.model_config.hf_config, + speculative_num_draft_tokens=server_args.speculative_num_draft_tokens, + target_vocab_size=int( + self.target_worker.model_runner.model_config.vocab_size + ), + ) + self.gamma = runtime_config.gamma + self.verify_num_draft_tokens = runtime_config.verify_num_draft_tokens + self.speculative_num_draft_tokens = self.verify_num_draft_tokens + self._mask_token_id = runtime_config.mask_token_id + + if self.tp_rank == 0: + logger.info( + "Initialized DSpark draft runner. attention_backend=%s, model=%s, " + "gamma=%s, verify_num_draft_tokens=%s, mask_token_id=%s, " + "markov_head=%s", + bundle.resolved_attention_backend, + self.draft_model.__class__.__name__, + self.gamma, + self.verify_num_draft_tokens, + self._mask_token_id, + type(self.draft_model.markov_head).__name__, + ) + + self._block_pos_offsets = build_block_pos_offsets( + length=self.verify_num_draft_tokens, device=self.device + ) + self._draft_block_spec_info = make_draft_block_spec_info( + draft_token_num=int(self.gamma), device=self.device + ) + + target_model = self.target_worker.model_runner.model + lm_head = getattr(target_model, "lm_head", None) + if lm_head is None or not hasattr(lm_head, "weight"): + raise RuntimeError( + "DSpark requires the target model to expose `lm_head` with `weight`." + ) + self.draft_model.attach_shared_modules( + embed_tokens=self._resolve_target_embed_tokens(target_model), + lm_head=lm_head, + ) + + self._verify_planner = DSparkVerifyPlanner( + draft_model=self.draft_model, + gamma=self.gamma, + model_runner=self.model_runner, + device=self.device, + tp_rank=self.tp_rank, + server_args=self.server_args, + verify_num_draft_tokens=self.verify_num_draft_tokens, + ) + if ( + server_args.enable_dp_attention + and not self._draft_is_moe + and self._verify_planner.is_compact_mode + and not server_args.disable_cuda_graph + ): + raise ValueError( + "DSpark dense-draft compact verify under --enable-dp-attention does not " + "yet support cuda graph (idle DP groups cannot join the token-keyed " + "compact graph). Re-run with --disable-cuda-graph (eager is lossless), " + "or use SGLANG_RAGGED_VERIFY_MODE=static. The dsv4 (MoE) draft supports " + "cuda graph under DP." + ) + self._kv_injector = TargetHiddenKvInjector( + draft_model=self.draft_model, + draft_model_runner=self.draft_model_runner, + model_runner=self.model_runner, + device=self.device, + verify_num_draft_tokens=self.verify_num_draft_tokens, + block_pos_offsets=self._block_pos_offsets, + ) + self._proposer = DraftBlockProposer( + draft_model=self.draft_model, + draft_model_runner=self.draft_model_runner, + gamma=self.gamma, + mask_token_id=self._mask_token_id, + draft_block_spec_info=self._draft_block_spec_info, + dp_moe_sync=self._draft_is_moe and server_args.enable_dp_attention, + ) + self._verify_epilogue = None + if ( + self._verify_planner.is_compact_mode + and not server_args.disable_cuda_graph + and is_cuda() + ): + self._verify_epilogue = DsparkVerifyEpilogue( + max_bs=max(server_args.cuda_graph_config.decode.bs), + verify_num_draft_tokens=self.verify_num_draft_tokens, + device=self.device, + commit_ctx=CommitInjectCtx( + draft_model=self.draft_model, + block_pos_offsets=self._block_pos_offsets, + resolve_pool=lambda: self.draft_model_runner.token_to_kv_pool, + resolve_req_to_token=lambda: ( + self.model_runner.req_to_token_pool.req_to_token + ), + ), + ) + self.model_runner.capture_tail_hooks.append( + self._verify_epilogue.capture_hook + ) + + self._simulate_acc_len = float(envs.SGLANG_SIMULATE_ACC_LEN.get()) + if ( + self._simulate_acc_len > 0 + and self._simulate_acc_len != 1.0 + and not self._verify_planner.is_verify_all + ): + raise ValueError( + "SGLANG_SIMULATE_ACC_LEN>1.0 with DSpark requires a verify-all " + "schedule (SGLANG_RAGGED_VERIFY_MODE=static, or =compact with the " + "uninitialized/flat SPS table): a constant simulated correct_len>0 " + "can exceed a trimmed request's verify budget (cap-accept, or " + "compact with a profiled SPS table) and break the cutoff/cap " + "accounting. SGLANG_SIMULATE_ACC_LEN=1.0 yields correct_len=0 " + "(commit is the bonus token only), which stays within every verify " + "budget and is safe in any mode. Got mode=" + f"{self._verify_planner.mode_value!r}, simulate_acc_len=" + f"{self._simulate_acc_len}." + ) + + self._verify_executor = TargetVerifyExecutor( + target_worker=self.target_worker, + gamma=self.gamma, + verify_num_draft_tokens=self.verify_num_draft_tokens, + model_runner=self.model_runner, + kv_injector=self._kv_injector, + verify_epilogue=self._verify_epilogue, + simulate_acc_len=self._simulate_acc_len, + ) + + self._forced_budget_frac: Optional[float] = None + + self._observers = DsparkStepObservers( + planner=self._verify_planner, + gamma=self.gamma, + verify_num_draft_tokens=self.verify_num_draft_tokens, + tp_rank=self.tp_rank, + device=self.device, + simulate_acc_len=self._simulate_acc_len, + ) + + def _resolve_target_embed_tokens(self, target_model): + if hasattr(target_model, "get_input_embeddings"): + return target_model.get_input_embeddings() + return target_model.model.get_input_embeddings() + + @property + def carries_confidence(self) -> bool: + return self._verify_planner.carries_confidence + + @property + def target_worker(self) -> TpModelWorker: + return self._target_worker + + @property + def draft_worker(self): + return self._draft_worker + + @property + def spec_v2_attn_backends(self) -> tuple: + return ( + self._target_worker.model_runner.attn_backend, + self.draft_model_runner.attn_backend, + ) + + def __getattr__(self, name): + if name == "_target_worker": + raise AttributeError(name) + return getattr(self.target_worker, name) + + def _draft_context(self): + if self._draft_dp_context_enabled: + return draft_tp_context(get_parallel().attn_tp_group) + return nullcontext() + + def alloc_memory_pool( + self, + memory_pool_config=None, + req_to_token_pool=None, + token_to_kv_pool_allocator=None, + ): + self._draft_worker.alloc_memory_pool( + memory_pool_config=memory_pool_config, + req_to_token_pool=req_to_token_pool, + token_to_kv_pool_allocator=token_to_kv_pool_allocator, + ) + + def init_attention_backends(self): + with self._draft_context(): + self._draft_worker.init_attention_backends() + + def init_cuda_graphs(self): + capture_decode_cuda_graph = not self.server_args.disable_cuda_graph + if is_cuda() and capture_decode_cuda_graph: + available_mem = get_available_gpu_memory(self.device, self.gpu_id) + if available_mem < 1.0: + capture_decode_cuda_graph = False + logger.warning( + "Disable DSpark draft cuda graph because only %.2f GB GPU " + "memory is available after target backend initialization.", + available_mem, + ) + with self._draft_context(): + if capture_decode_cuda_graph: + self._draft_sampler = self._maybe_build_draft_sampler() + if self._draft_sampler is not None: + self.draft_model_runner.capture_tail_hooks.append( + make_draft_sampler_capture_hook(self._draft_sampler) + ) + self._proposer.attach_draft_sampler(self._draft_sampler) + self._draft_worker.init_cuda_graphs( + capture_decode_cuda_graph=capture_decode_cuda_graph + ) + + def _maybe_build_draft_sampler(self): + return maybe_build_draft_sampler( + draft_model=self.draft_model, + gamma=self.gamma, + max_bs=max(self.server_args.cuda_graph_config.decode.bs), + device=self.device, + tp_rank=self.tp_rank, + confidence_fn=( + self._verify_planner.compute_confidence_tensor + if self._verify_planner.carries_confidence + else None + ), + out=( + self._verify_epilogue.draft_tokens_buf + if self._verify_epilogue is not None + else None + ), + ) + + def clear_cache_pool(self): + pass + + def set_dspark_forced_budget_frac(self, frac: Optional[float]) -> None: + self._forced_budget_frac = frac + self._verify_planner.set_forced_budget_frac(frac) + + def dump_info_records(self) -> Optional[dict]: + return self._observers.dump_info_records() + + def clear_info_records(self) -> None: + self._observers.clear_info_records() + + def block_accept_estimate_log_suffix(self) -> Optional[str]: + return self._observers.block_accept_estimate_log_suffix() + + def note_request_finished(self, *, rid: str, natural_stop: bool) -> None: + self._observers.note_request_finished(rid=rid, natural_stop=natural_stop) + + def forward_batch_generation( + self, + batch: ScheduleBatch, + on_publish=None, + ) -> GenerationBatchResult: + if getattr(batch, "return_logprob", False): + raise ValueError( + "DSpark speculative decoding does not support return_logprob yet." + ) + + if batch.forward_mode.is_extend() or batch.is_extend_in_batch: + self._verify_planner.note_non_decode_step() + self._observers.note_prefill_step() + return self._forward_prefill(batch, on_publish) + + return self._forward_decode(batch, on_publish) + + def _forward_prefill( + self, batch: ScheduleBatch, on_publish + ) -> GenerationBatchResult: + if batch.forward_mode.is_idle(): + if self.server_args.enable_dp_attention: + batch.capture_hidden_mode = CaptureHiddenMode.FULL + self.target_worker.forward_batch_generation(batch) + return self._decode_idle_result(on_publish=on_publish) + + batch.capture_hidden_mode = CaptureHiddenMode.FULL + batch_output = self.target_worker.forward_batch_generation(batch) + logits_output = batch_output.logits_output + next_token_ids = batch_output.next_token_ids + batch_output.new_seq_lens = batch.seq_lens + if on_publish is not None: + on_publish(batch_output.new_seq_lens) + + if logits_output.hidden_states is None: + raise RuntimeError( + "DSpark requires target aux hidden capture for prefill, but got None. " + "Make sure the target model has DFlash layers-to-capture configured." + ) + if batch.extend_lens is None or batch.prefix_lens is None: + raise RuntimeError( + "DSpark expected extend_lens / prefix_lens in extend mode, got None." + ) + if batch.out_cache_loc is None: + raise RuntimeError("DSpark prefill expected out_cache_loc, but got None.") + + device = next_token_ids.device + ctx_lens = torch.tensor(batch.extend_lens, dtype=torch.int32, device=device) + draft_seq_lens = torch.tensor( + batch.prefix_lens, dtype=torch.int32, device=device + ) + positions, _ = compute_position( + self.model_runner.server_args.attention_backend, + draft_seq_lens, + ctx_lens, + int(sum(batch.extend_lens)), + ) + self._kv_injector.inject_target_hidden( + target_hidden=logits_output.hidden_states, + cache_loc=batch.out_cache_loc, + positions=positions, + ) + logits_output.hidden_states = None + + batch_output.next_draft_input = make_next_draft_input( + bonus_tokens=next_token_ids, + new_seq_lens=batch.seq_lens, + ) + return batch_output + + def _idle_verify_ragged_layout(self, batch: ScheduleBatch): + if batch.global_num_tokens is None or not self._verify_planner.is_compact_mode: + return None + global_bs = max(batch.global_num_tokens) + if global_bs <= 0: + return None + return idle_ragged_layout( + tier_num_reqs=global_bs, + dp_tier_num_tokens=self._dp_verify_tier_num_tokens(batch), + device=self.device, + verify_num_draft_tokens=self.verify_num_draft_tokens, + model_runner=self.model_runner, + ) + + def _dp_verify_tier_num_tokens(self, batch: ScheduleBatch) -> Optional[int]: + if not ( + self._draft_is_moe + and self.server_args.enable_dp_attention + and batch.global_num_tokens is not None + and self._verify_planner.is_compact_mode + ): + return None + return dp_global_verify_tier_num_tokens( + global_tier_num_tokens=batch.global_spec_verify_tier_num_tokens + ) + + def _decode_idle_result( + self, + *, + on_publish, + ) -> GenerationBatchResult: + next_draft_input = make_next_draft_input( + bonus_tokens=torch.empty((0,), device=self.device, dtype=torch.int64), + new_seq_lens=torch.empty((0,), device=self.device, dtype=torch.int64), + ) + if on_publish is not None: + on_publish(next_draft_input.new_seq_lens) + return GenerationBatchResult( + logits_output=None, + next_token_ids=torch.empty((0,), dtype=torch.int64, device=self.device), + accept_lens=torch.empty((0,), dtype=torch.int32, device=self.device), + block_accept_lens=torch.empty((0,), dtype=torch.int32, device=self.device), + next_draft_input=next_draft_input, + can_run_cuda_graph=False, + speculative_num_draft_tokens=int(self.verify_num_draft_tokens), + new_seq_lens=next_draft_input.new_seq_lens, + ) + + def _forward_decode( + self, batch: ScheduleBatch, on_publish + ) -> GenerationBatchResult: + if batch.spec_info is None: + batch.spec_info = DFlashDraftInputV2.create_idle_input(device=self.device) + draft_input = batch.spec_info + if not isinstance(draft_input, DFlashDraftInputV2): + raise RuntimeError( + "DSpark spec-v2 expected DFlashDraftInputV2 state on the running batch." + ) + + if batch.forward_mode.is_idle(): + self._observers.note_idle_decode_step() + if self.server_args.enable_dp_attention: + if self._draft_is_moe: + self._proposer.run_idle_participation(batch) + self._verify_executor.run_idle_participation( + batch=batch, idle_layout=self._idle_verify_ragged_layout(batch) + ) + return self._decode_idle_result(on_publish=on_publish) + + batch.seq_lens.record_stream( + torch.get_device_module(self.device).current_stream() + ) + bs = len(batch.seq_lens) + device = self.device + prefix_lens = batch.seq_lens + + self._observers.begin_step() + + target_model = self.target_worker.model_runner.model + + verify_window = alloc_verify_window( + batch=batch, + bs=bs, + device=device, + verify_num_draft_tokens=self.verify_num_draft_tokens, + block_pos_offsets=self._block_pos_offsets, + model_runner=self.model_runner, + ) + + sampling_info = batch.sampling_info + with self._draft_context(), self._observers.segment(InfoSegment.DRAFT): + proposal = self._proposer.propose( + batch=batch, + draft_input=draft_input, + verify_window=verify_window, + bs=bs, + device=device, + target_model=target_model, + sampling_info=sampling_info, + ) + draft_block_ids = proposal.draft_block_ids + draft_block = proposal.draft_block + draft_tokens = draft_block.draft_tokens + + confidence = proposal.confidence + if confidence is None: + confidence = self._verify_planner.compute_confidence_tensor( + draft_hidden=proposal.draft_hidden, + anchor_tokens=draft_block_ids[:, 0], + draft_tokens=draft_tokens, + confidence_tap=proposal.confidence_tap, + ) + + verify_token_budget = self._verify_planner.resolve_verify_token_budget( + draft_input=draft_input, + confidence=confidence, + prefix_lens=prefix_lens, + req_pool_indices=batch.req_pool_indices, + ) + + global_num_reqs = ( + max(batch.global_num_tokens) + if self._draft_is_moe + and self.server_args.enable_dp_attention + and batch.global_num_tokens is not None + else None + ) + layout = self._verify_planner.schedule_layout( + req_pool_indices=batch.req_pool_indices, + prefix_lens=prefix_lens, + device=device, + confidence=confidence, + budget=verify_token_budget, + global_num_reqs=global_num_reqs, + dp_tier_num_tokens=self._dp_verify_tier_num_tokens(batch), + ) + run_compact = self._verify_planner.should_run_compact(layout=layout) + + verify_ids_2d = torch.cat( + [draft_block_ids[:, :1], draft_tokens], dim=1 + ).contiguous() + + fold_eligible = ( + self._verify_executor.verify_epilogue is not None + and proposal.folded + and verify_logits_adjustments_are_noop(sampling_info) + and self._simulate_acc_len <= 0 + ) + with self._observers.segment(InfoSegment.TARGET_VERIFY): + if run_compact: + target_verify, hidden_strided = self._verify_executor.run_compact( + batch=batch, + layout=layout, + draft_block_ids=draft_block_ids, + draft_tokens=draft_tokens, + bs=bs, + device=device, + sampling_info=sampling_info, + inject_gate=fold_eligible, + ) + else: + target_verify = self._verify_executor.run_non_compact( + batch=batch, + draft_input=draft_input, + verify_ids_2d=verify_ids_2d, + verify_window=verify_window, + sampling_info=sampling_info, + ) + hidden_strided = None + logits_output = target_verify.logits_output + can_run_cuda_graph = target_verify.can_run_cuda_graph + + epilogue = self._verify_executor.verify_epilogue + folded_accept = fold_eligible and run_compact and can_run_cuda_graph + accept = self._verify_executor.accept_and_finalize( + folded_accept=folded_accept, + bs=bs, + verify_ids_2d=verify_ids_2d, + target_logits=logits_output.next_token_logits, + draft_block=draft_block, + sampling_info=sampling_info, + draft_input=draft_input, + layout=layout, + prefix_lens=prefix_lens, + draft_tokens=draft_tokens, + ) + if on_publish is not None: + if confidence is not None: + on_publish(accept.new_seq_lens, confidence=confidence) + else: + on_publish(accept.new_seq_lens) + + folded_commit = folded_accept and epilogue.folds_commit + if not folded_commit: + self._verify_executor.commit_hidden( + batch=batch, + layout=layout, + hidden_strided=hidden_strided, + verify_window=verify_window, + logits_output=logits_output, + commit_lens=accept.commit_lens, + bs=bs, + run_compact=run_compact, + ) + logits_output.hidden_states = None + + self._observers.observe_verify_step( + forward_ct=int(batch.forward_iter), + reqs=batch.reqs, + bs=bs, + proposal_folded=proposal.folded, + verify_ids_2d=verify_ids_2d, + target_logits=logits_output.next_token_logits, + layout=layout, + confidence=confidence, + prefix_lens=prefix_lens, + draft_tokens=draft_tokens, + draft_block=draft_block, + sampling_info=sampling_info, + correct_len=accept.correct_len, + cap_trim_lens=accept.cap_trim_lens, + bonus=accept.bonus, + commit_lens=accept.commit_lens, + verify_token_budget=verify_token_budget, + req_pool_indices=batch.req_pool_indices, + verify_tier_num_tokens=int(batch.spec_verify_tier_num_tokens), + dp_tier_num_tokens=self._dp_verify_tier_num_tokens(batch), + ) + + next_draft_input = make_next_draft_input( + bonus_tokens=accept.bonus, + new_seq_lens=accept.new_seq_lens, + ) + return GenerationBatchResult( + logits_output=logits_output, + next_token_ids=accept.out_tokens.reshape(-1), + accept_lens=accept.commit_lens, + block_accept_lens=accept.commit_lens + accept.cap_trim_lens, + cap_lens=( + layout.verify_lens.to(torch.int32) if layout is not None else None + ), + can_run_cuda_graph=can_run_cuda_graph, + next_draft_input=next_draft_input, + speculative_num_draft_tokens=int(self.verify_num_draft_tokens), + new_seq_lens=accept.new_seq_lens, + ) + + def get_confidence_budget_prepare(self): + return self._verify_planner.confidence_budget_prepare() diff --git a/python/sglang/srt/speculative/dspark_components/kernels/__init__.py b/python/sglang/srt/speculative/dspark_components/kernels/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dispatch.py b/python/sglang/srt/speculative/dspark_components/kernels/dispatch.py new file mode 100644 index 000000000..032b99b9f --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/kernels/dispatch.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +import torch + + +def inputs_on_cuda(*args, **kwargs) -> bool: + """Route kernel dispatch by input placement: the first tensor argument + decides. CUDA inputs take the fused triton kernel; CPU inputs take the + torch reference implementation (triton is CUDA-only, and CPU-side callers + such as unit tests exercise the reference path).""" + for value in (*args, *kwargs.values()): + if isinstance(value, torch.Tensor): + return value.is_cuda + raise AssertionError("kernel dispatch requires at least one tensor argument") diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py b/python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py new file mode 100644 index 000000000..560c17077 --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py @@ -0,0 +1,862 @@ +from __future__ import annotations + +from typing import Optional + +import msgspec +import torch +import triton +import triton.language as tl + +from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 +from sglang.srt.speculative.dflash_utils import ( + _get_or_create_chain_verify_buffers, + build_dflash_verify_target_probs, + compute_dflash_correct_drafts_and_bonus, +) +from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda +from sglang.srt.speculative.reject_sampling import chain_speculative_sampling_triton + + +class AcceptSampling: + @classmethod + def execute( + cls, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + candidates: torch.Tensor, + target_logits: torch.Tensor, + draft_probs: torch.Tensor, + sampling_info, + draft_input: DFlashDraftInputV2, + gamma: int, + verify_num_draft_tokens: int, + cutoff_verify_lens: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return accept_sampling( + candidates=candidates, + target_logits=target_logits, + draft_probs=draft_probs, + sampling_info=sampling_info, + draft_input=draft_input, + gamma=gamma, + verify_num_draft_tokens=verify_num_draft_tokens, + cutoff_verify_lens=cutoff_verify_lens, + ) + + @classmethod + def triton( + cls, + *, + candidates: torch.Tensor, + target_logits: torch.Tensor, + draft_probs: torch.Tensor, + sampling_info, + draft_input: DFlashDraftInputV2, + gamma: int, + verify_num_draft_tokens: int, + cutoff_verify_lens: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return accept_sampling_triton( + candidates=candidates, + target_logits=target_logits, + draft_probs=draft_probs, + sampling_info=sampling_info, + draft_input=draft_input, + gamma=gamma, + verify_num_draft_tokens=verify_num_draft_tokens, + cutoff_verify_lens=cutoff_verify_lens, + ) + + +def _accept_sampling_core( + *, + candidates: torch.Tensor, + target_logits: torch.Tensor, + draft_probs: torch.Tensor, + sampling_info, + draft_input: DFlashDraftInputV2, + gamma: int, + verify_num_draft_tokens: int, + cutoff_verify_lens: Optional[torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + bs = candidates.shape[0] + device = candidates.device + if not sampling_info.need_top_k_sampling and not sampling_info.need_top_p_sampling: + target_probs = SoftmaxTemp.execute( + logits=target_logits, + temperatures=sampling_info.temperatures, + rows_per_request=verify_num_draft_tokens, + ).view(bs, verify_num_draft_tokens, -1) + else: + target_probs = build_dflash_verify_target_probs( + next_token_logits=target_logits, + sampling_info=sampling_info, + draft_token_num=verify_num_draft_tokens, + bs=bs, + max_top_k=draft_input.max_top_k, + uniform_top_k_value=draft_input.uniform_top_k_value, + ) + ( + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + predicts, + accept_index, + accept_token_num, + ) = _get_or_create_chain_verify_buffers( + bs=bs, + draft_token_num=verify_num_draft_tokens, + device=device, + ) + uniform_samples = torch.rand((bs, gamma), dtype=torch.float32, device=device) + uniform_samples_final = torch.rand((bs,), dtype=torch.float32, device=device) + chain_speculative_sampling_triton( + predicts=predicts, + accept_index=accept_index, + accept_token_num=accept_token_num, + candidates=candidates, + retrive_index=retrieve_index, + retrive_next_token=retrieve_next_token, + retrive_next_sibling=retrieve_next_sibling, + uniform_samples=uniform_samples, + uniform_samples_for_final_sampling=uniform_samples_final, + target_probs=target_probs, + draft_probs=draft_probs, + threshold_single=1.0, + threshold_acc=1.0, + deterministic=True, + ) + correct_len = accept_token_num + if cutoff_verify_lens is not None: + correct_len, cap_trim_lens = CapCorrectLen.execute( + correct_len=correct_len, verify_lens=cutoff_verify_lens + ) + else: + cap_trim_lens = torch.zeros_like(correct_len) + return correct_len, cap_trim_lens, accept_index, predicts + + +def accept_sampling( + *, + candidates: torch.Tensor, + target_logits: torch.Tensor, + draft_probs: torch.Tensor, + sampling_info, + draft_input: DFlashDraftInputV2, + gamma: int, + verify_num_draft_tokens: int, + cutoff_verify_lens: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + bs = candidates.shape[0] + device = candidates.device + correct_len, cap_trim_lens, accept_index, predicts = _accept_sampling_core( + candidates=candidates, + target_logits=target_logits, + draft_probs=draft_probs, + sampling_info=sampling_info, + draft_input=draft_input, + gamma=gamma, + verify_num_draft_tokens=verify_num_draft_tokens, + cutoff_verify_lens=cutoff_verify_lens, + ) + row_ids = torch.arange(bs, dtype=torch.long, device=device) + accept_pos = accept_index[row_ids, correct_len.to(torch.long)].to(torch.long) + bonus = predicts[accept_pos].to(torch.int64) + return correct_len, bonus, cap_trim_lens + + +@triton.jit +def _gather_two_level_bonus_kernel( + accept_index_ptr, + predicts_ptr, + correct_len_ptr, + out_ptr, + cols, + n, + BLOCK: tl.constexpr, +): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + cl = tl.load(correct_len_ptr + offs, mask=mask, other=0).to(tl.int64) + accept_pos = tl.load(accept_index_ptr + offs * cols + cl, mask=mask, other=0).to( + tl.int64 + ) + bonus = tl.load(predicts_ptr + accept_pos, mask=mask, other=0) + tl.store(out_ptr + offs, bonus.to(tl.int64), mask=mask) + + +def gather_two_level_bonus_triton( + *, + accept_index: torch.Tensor, + predicts: torch.Tensor, + correct_len: torch.Tensor, +) -> torch.Tensor: + bs, cols = accept_index.shape + accept_index = accept_index.contiguous() + predicts = predicts.contiguous() + correct_len = correct_len.contiguous() + out = torch.empty(bs, dtype=torch.int64, device=accept_index.device) + BLOCK = 256 + grid = (triton.cdiv(bs, BLOCK),) + _gather_two_level_bonus_kernel[grid]( + accept_index, predicts, correct_len, out, cols, bs, BLOCK=BLOCK + ) + return out + + +def accept_sampling_triton( + *, + candidates: torch.Tensor, + target_logits: torch.Tensor, + draft_probs: torch.Tensor, + sampling_info, + draft_input: DFlashDraftInputV2, + gamma: int, + verify_num_draft_tokens: int, + cutoff_verify_lens: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + correct_len, cap_trim_lens, accept_index, predicts = _accept_sampling_core( + candidates=candidates, + target_logits=target_logits, + draft_probs=draft_probs, + sampling_info=sampling_info, + draft_input=draft_input, + gamma=gamma, + verify_num_draft_tokens=verify_num_draft_tokens, + cutoff_verify_lens=cutoff_verify_lens, + ) + bonus = gather_two_level_bonus_triton( + accept_index=accept_index, predicts=predicts, correct_len=correct_len + ) + return correct_len, bonus, cap_trim_lens + + +try: + from flashinfer.sampling import softmax as _flashinfer_softmax +except ImportError: + _flashinfer_softmax = None + + +class SoftmaxTemp: + @classmethod + def execute(cls, *args, **kwargs) -> torch.Tensor: + if not inputs_on_cuda(*args, **kwargs): + return cls.torch(*args, **kwargs) + if _flashinfer_softmax is not None: + return cls.flashinfer(*args, **kwargs) + return cls.triton(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + logits: torch.Tensor, + temperatures: torch.Tensor, + rows_per_request: int, + ) -> torch.Tensor: + return softmax_temp( + logits=logits, + temperatures=temperatures, + rows_per_request=rows_per_request, + ) + + @classmethod + def triton( + cls, + *, + logits: torch.Tensor, + temperatures: torch.Tensor, + rows_per_request: int, + ) -> torch.Tensor: + return softmax_temp_triton( + logits=logits, + temperatures=temperatures, + rows_per_request=rows_per_request, + ) + + @classmethod + def flashinfer( + cls, + *, + logits: torch.Tensor, + temperatures: torch.Tensor, + rows_per_request: int, + ) -> torch.Tensor: + return softmax_temp_flashinfer( + logits=logits, + temperatures=temperatures, + rows_per_request=rows_per_request, + ) + + +def softmax_temp( + *, + logits: torch.Tensor, + temperatures: torch.Tensor, + rows_per_request: int, +) -> torch.Tensor: + num_rows = logits.shape[0] + bs = num_rows // rows_per_request + assert ( + bs * rows_per_request == num_rows + ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + temp_per_row = torch.repeat_interleave( + temperatures.reshape(bs).to(torch.float32), rows_per_request, dim=0 + ) + scaled = logits.to(torch.float32) / temp_per_row[:, None] + return torch.softmax(scaled, dim=-1) + + +@triton.jit +def _softmax_temp_kernel( + logits_ptr, + temp_ptr, + out_ptr, + vocab, + rows_per_request, + logits_row_stride, + BLOCK_V: tl.constexpr, +): + row = tl.program_id(0) + temp = tl.load(temp_ptr + row // rows_per_request).to(tl.float32) + base = logits_ptr + row.to(tl.int64) * logits_row_stride + out_base = out_ptr + row.to(tl.int64) * vocab + + row_max = -float("inf") + for v0 in range(0, vocab, BLOCK_V): + offs = v0 + tl.arange(0, BLOCK_V) + vmask = offs < vocab + x = tl.load(base + offs, mask=vmask, other=-float("inf")).to(tl.float32) + x = x / temp + row_max = tl.maximum(row_max, tl.max(x, axis=0)) + + sum_exp = 0.0 + for v0 in range(0, vocab, BLOCK_V): + offs = v0 + tl.arange(0, BLOCK_V) + vmask = offs < vocab + x = tl.load(base + offs, mask=vmask, other=-float("inf")).to(tl.float32) + x = x / temp + e = tl.exp(x - row_max) + e = tl.where(vmask, e, 0.0) + sum_exp += tl.sum(e, axis=0) + + for v0 in range(0, vocab, BLOCK_V): + offs = v0 + tl.arange(0, BLOCK_V) + vmask = offs < vocab + x = tl.load(base + offs, mask=vmask, other=-float("inf")).to(tl.float32) + x = x / temp + e = tl.exp(x - row_max) + tl.store(out_base + offs, e / sum_exp, mask=vmask) + + +def softmax_temp_triton( + *, + logits: torch.Tensor, + temperatures: torch.Tensor, + rows_per_request: int, +) -> torch.Tensor: + num_rows, vocab = logits.shape[0], logits.shape[-1] + bs = num_rows // rows_per_request + assert ( + bs * rows_per_request == num_rows + ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + temperatures = temperatures.reshape(bs).to(torch.float32).contiguous() + out = torch.empty((num_rows, vocab), dtype=torch.float32, device=logits.device) + BLOCK_V = 4096 + _softmax_temp_kernel[(num_rows,)]( + logits, + temperatures, + out, + vocab, + rows_per_request, + logits.stride(0), + BLOCK_V=BLOCK_V, + ) + return out + + +def softmax_temp_flashinfer( + *, + logits: torch.Tensor, + temperatures: torch.Tensor, + rows_per_request: int, +) -> torch.Tensor: + if _flashinfer_softmax is None: + raise RuntimeError( + "softmax_temp_flashinfer requires flashinfer.sampling.softmax, " + "which is unavailable in this environment" + ) + num_rows, vocab = logits.shape[0], logits.shape[-1] + bs = num_rows // rows_per_request + assert ( + bs * rows_per_request == num_rows + ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + temp_per_row = torch.repeat_interleave( + temperatures.reshape(bs).to(torch.float32), rows_per_request, dim=0 + ).contiguous() + logits_2d = logits.to(torch.float32).contiguous() + return _flashinfer_softmax(logits=logits_2d, temperature=temp_per_row) + + +class MixedAcceptSelectResult(msgspec.Struct): + correct_len: torch.Tensor + bonus: torch.Tensor + cap_trim_lens: torch.Tensor + + +class SelectMixedAccept: + @classmethod + def execute(cls, *args, **kwargs) -> MixedAcceptSelectResult: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + greedy_mask: torch.Tensor, + greedy_len: torch.Tensor, + greedy_bonus: torch.Tensor, + greedy_trim: torch.Tensor, + sampling_len: torch.Tensor, + sampling_bonus: torch.Tensor, + sampling_trim: torch.Tensor, + ) -> MixedAcceptSelectResult: + return select_mixed_accept( + greedy_mask=greedy_mask, + greedy_len=greedy_len, + greedy_bonus=greedy_bonus, + greedy_trim=greedy_trim, + sampling_len=sampling_len, + sampling_bonus=sampling_bonus, + sampling_trim=sampling_trim, + ) + + @classmethod + def triton( + cls, + *, + greedy_mask: torch.Tensor, + greedy_len: torch.Tensor, + greedy_bonus: torch.Tensor, + greedy_trim: torch.Tensor, + sampling_len: torch.Tensor, + sampling_bonus: torch.Tensor, + sampling_trim: torch.Tensor, + ) -> MixedAcceptSelectResult: + return select_mixed_accept_triton( + greedy_mask=greedy_mask, + greedy_len=greedy_len, + greedy_bonus=greedy_bonus, + greedy_trim=greedy_trim, + sampling_len=sampling_len, + sampling_bonus=sampling_bonus, + sampling_trim=sampling_trim, + ) + + +def select_mixed_accept( + *, + greedy_mask: torch.Tensor, + greedy_len: torch.Tensor, + greedy_bonus: torch.Tensor, + greedy_trim: torch.Tensor, + sampling_len: torch.Tensor, + sampling_bonus: torch.Tensor, + sampling_trim: torch.Tensor, +) -> MixedAcceptSelectResult: + correct_len = torch.where( + greedy_mask, greedy_len.to(sampling_len.dtype), sampling_len + ) + bonus = torch.where(greedy_mask, greedy_bonus, sampling_bonus) + cap_trim_lens = torch.where( + greedy_mask, greedy_trim.to(sampling_trim.dtype), sampling_trim + ) + return MixedAcceptSelectResult( + correct_len=correct_len, bonus=bonus, cap_trim_lens=cap_trim_lens + ) + + +@triton.jit +def _mixed_accept_select_kernel( + greedy_mask_ptr, + greedy_len_ptr, + greedy_bonus_ptr, + greedy_trim_ptr, + sampling_len_ptr, + sampling_bonus_ptr, + sampling_trim_ptr, + correct_len_ptr, + bonus_ptr, + cap_trim_ptr, + bs, + BLOCK: tl.constexpr, +): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < bs + is_greedy = tl.load(greedy_mask_ptr + offs, mask=mask, other=0) != 0 + + g_len = tl.load(greedy_len_ptr + offs, mask=mask, other=0) + s_len = tl.load(sampling_len_ptr + offs, mask=mask, other=0) + tl.store(correct_len_ptr + offs, tl.where(is_greedy, g_len, s_len), mask=mask) + + g_bonus = tl.load(greedy_bonus_ptr + offs, mask=mask, other=0) + s_bonus = tl.load(sampling_bonus_ptr + offs, mask=mask, other=0) + tl.store(bonus_ptr + offs, tl.where(is_greedy, g_bonus, s_bonus), mask=mask) + + g_trim = tl.load(greedy_trim_ptr + offs, mask=mask, other=0) + s_trim = tl.load(sampling_trim_ptr + offs, mask=mask, other=0) + tl.store(cap_trim_ptr + offs, tl.where(is_greedy, g_trim, s_trim), mask=mask) + + +def select_mixed_accept_triton( + *, + greedy_mask: torch.Tensor, + greedy_len: torch.Tensor, + greedy_bonus: torch.Tensor, + greedy_trim: torch.Tensor, + sampling_len: torch.Tensor, + sampling_bonus: torch.Tensor, + sampling_trim: torch.Tensor, +) -> MixedAcceptSelectResult: + bs = greedy_mask.shape[0] + device = greedy_mask.device + + correct_len = torch.empty(bs, dtype=sampling_len.dtype, device=device) + bonus = torch.empty(bs, dtype=sampling_bonus.dtype, device=device) + cap_trim_lens = torch.empty(bs, dtype=sampling_trim.dtype, device=device) + BLOCK = 256 + _mixed_accept_select_kernel[(triton.cdiv(bs, BLOCK),)]( + greedy_mask, + greedy_len, + greedy_bonus, + greedy_trim, + sampling_len, + sampling_bonus, + sampling_trim, + correct_len, + bonus, + cap_trim_lens, + bs, + BLOCK=BLOCK, + ) + return MixedAcceptSelectResult( + correct_len=correct_len, bonus=bonus, cap_trim_lens=cap_trim_lens + ) + + +class AcceptGreedy: + @classmethod + def execute( + cls, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + candidates: torch.Tensor, + target_logits: torch.Tensor, + verify_num_draft_tokens: int, + cutoff_verify_lens: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return accept_greedy( + candidates=candidates, + target_logits=target_logits, + verify_num_draft_tokens=verify_num_draft_tokens, + cutoff_verify_lens=cutoff_verify_lens, + ) + + @classmethod + def triton( + cls, + *, + candidates: torch.Tensor, + target_logits: torch.Tensor, + verify_num_draft_tokens: int, + cutoff_verify_lens: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return accept_greedy_triton( + candidates=candidates, + target_logits=target_logits, + verify_num_draft_tokens=verify_num_draft_tokens, + cutoff_verify_lens=cutoff_verify_lens, + ) + + +def accept_greedy( + *, + candidates: torch.Tensor, + target_logits: torch.Tensor, + verify_num_draft_tokens: int, + cutoff_verify_lens: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + bs = candidates.shape[0] + target_predict = torch.argmax(target_logits, dim=-1).view( + bs, verify_num_draft_tokens + ) + correct_len, bonus = compute_dflash_correct_drafts_and_bonus( + candidates=candidates, + target_predict=target_predict, + ) + cap_trim_lens = torch.zeros_like(correct_len) + if cutoff_verify_lens is not None: + correct_len, cap_trim_lens = CapCorrectLen.execute( + correct_len=correct_len, verify_lens=cutoff_verify_lens + ) + row_ids = torch.arange(bs, device=target_predict.device) + bonus = target_predict[row_ids, correct_len.to(torch.long)].to(torch.int64) + return correct_len, bonus, cap_trim_lens + + +@triton.jit +def _gather_row_bonus_kernel( + table_ptr, + idx_ptr, + out_ptr, + cols, + n, + BLOCK: tl.constexpr, +): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + idx = tl.load(idx_ptr + offs, mask=mask, other=0).to(tl.int64) + val = tl.load(table_ptr + offs * cols + idx, mask=mask, other=0) + tl.store(out_ptr + offs, val.to(tl.int64), mask=mask) + + +def gather_row_bonus_triton(*, table: torch.Tensor, idx: torch.Tensor) -> torch.Tensor: + bs, cols = table.shape + table = table.contiguous() + idx = idx.contiguous() + out = torch.empty(bs, dtype=torch.int64, device=table.device) + BLOCK = 256 + grid = (triton.cdiv(bs, BLOCK),) + _gather_row_bonus_kernel[grid](table, idx, out, cols, bs, BLOCK=BLOCK) + return out + + +def accept_greedy_triton( + *, + candidates: torch.Tensor, + target_logits: torch.Tensor, + verify_num_draft_tokens: int, + cutoff_verify_lens: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + bs = candidates.shape[0] + target_predict = torch.argmax(target_logits, dim=-1).view( + bs, verify_num_draft_tokens + ) + correct_len, bonus = compute_dflash_correct_drafts_and_bonus( + candidates=candidates, + target_predict=target_predict, + ) + cap_trim_lens = torch.zeros_like(correct_len) + if cutoff_verify_lens is not None: + correct_len, cap_trim_lens = CapCorrectLen.execute( + correct_len=correct_len, verify_lens=cutoff_verify_lens + ) + bonus = gather_row_bonus_triton(table=target_predict, idx=correct_len) + return correct_len, bonus, cap_trim_lens + + +class FinalizeAcceptLensResult(msgspec.Struct): + commit_lens: torch.Tensor + new_seq_lens: torch.Tensor + cap_trim_lens: torch.Tensor + + +class FinalizeAcceptLens: + @classmethod + def execute(cls, *args, **kwargs) -> FinalizeAcceptLensResult: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + correct_len: torch.Tensor, + cap_trim_lens: torch.Tensor, + prefix_lens: torch.Tensor, + ) -> FinalizeAcceptLensResult: + return finalize_accept_lens( + correct_len=correct_len, + cap_trim_lens=cap_trim_lens, + prefix_lens=prefix_lens, + ) + + @classmethod + def triton( + cls, + *, + correct_len: torch.Tensor, + cap_trim_lens: torch.Tensor, + prefix_lens: torch.Tensor, + ) -> FinalizeAcceptLensResult: + return finalize_accept_lens_triton( + correct_len=correct_len, + cap_trim_lens=cap_trim_lens, + prefix_lens=prefix_lens, + ) + + +def finalize_accept_lens( + *, + correct_len: torch.Tensor, + cap_trim_lens: torch.Tensor, + prefix_lens: torch.Tensor, +) -> FinalizeAcceptLensResult: + commit_lens = correct_len.to(torch.int32) + 1 + new_seq_lens = prefix_lens + commit_lens.to(prefix_lens.dtype) + return FinalizeAcceptLensResult( + commit_lens=commit_lens, + new_seq_lens=new_seq_lens, + cap_trim_lens=cap_trim_lens.to(torch.int32), + ) + + +@triton.jit +def _finalize_accept_lens_kernel( + correct_len_ptr, + cap_trim_ptr, + prefix_lens_ptr, + commit_lens_ptr, + new_seq_lens_ptr, + cap_trim_out_ptr, + bs, + BLOCK: tl.constexpr, +): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < bs + commit = tl.load(correct_len_ptr + offs, mask=mask, other=0).to(tl.int32) + 1 + prefix = tl.load(prefix_lens_ptr + offs, mask=mask, other=0) + trim = tl.load(cap_trim_ptr + offs, mask=mask, other=0).to(tl.int32) + tl.store(commit_lens_ptr + offs, commit, mask=mask) + tl.store(new_seq_lens_ptr + offs, prefix + commit, mask=mask) + tl.store(cap_trim_out_ptr + offs, trim, mask=mask) + + +def finalize_accept_lens_triton( + *, + correct_len: torch.Tensor, + cap_trim_lens: torch.Tensor, + prefix_lens: torch.Tensor, +) -> FinalizeAcceptLensResult: + bs = correct_len.shape[0] + device = correct_len.device + + commit_lens = torch.empty(bs, dtype=torch.int32, device=device) + new_seq_lens = torch.empty(bs, dtype=prefix_lens.dtype, device=device) + cap_trim_out = torch.empty(bs, dtype=torch.int32, device=device) + BLOCK = 256 + _finalize_accept_lens_kernel[(triton.cdiv(bs, BLOCK),)]( + correct_len, + cap_trim_lens, + prefix_lens, + commit_lens, + new_seq_lens, + cap_trim_out, + bs, + BLOCK=BLOCK, + ) + return FinalizeAcceptLensResult( + commit_lens=commit_lens, + new_seq_lens=new_seq_lens, + cap_trim_lens=cap_trim_out, + ) + + +class CapCorrectLen: + @classmethod + def execute(cls, *args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + correct_len: torch.Tensor, + verify_lens: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return cap_correct_len( + correct_len=correct_len, + verify_lens=verify_lens, + ) + + @classmethod + def triton( + cls, + *, + correct_len: torch.Tensor, + verify_lens: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return cap_correct_len_triton( + correct_len=correct_len, + verify_lens=verify_lens, + ) + + +def cap_correct_len( + *, + correct_len: torch.Tensor, + verify_lens: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + ell_r = (verify_lens.to(device=correct_len.device) - 1).to(correct_len.dtype) + capped = torch.minimum(correct_len, ell_r) + cap_trim_lens = correct_len - capped + return capped, cap_trim_lens + + +@triton.jit +def _cap_correct_len_kernel( + correct_len_ptr, + verify_lens_ptr, + capped_ptr, + trim_ptr, + n, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + cl = tl.load(correct_len_ptr + offs, mask=mask, other=0).to(tl.int64) + vl = tl.load(verify_lens_ptr + offs, mask=mask, other=0).to(tl.int64) + ell = vl - 1 + capped = tl.minimum(cl, ell) + trim = cl - capped + tl.store(capped_ptr + offs, capped, mask=mask) + tl.store(trim_ptr + offs, trim, mask=mask) + + +def cap_correct_len_triton( + *, + correct_len: torch.Tensor, + verify_lens: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + device = correct_len.device + correct_len = correct_len.contiguous() + verify_lens = verify_lens.to(device=device).contiguous() + n = correct_len.shape[0] + capped = torch.empty_like(correct_len) + trim = torch.empty_like(correct_len) + BLOCK = 1024 + grid = (triton.cdiv(n, BLOCK),) + _cap_correct_len_kernel[grid]( + correct_len, verify_lens, capped, trim, n, BLOCK=BLOCK + ) + return capped, trim diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_attn_metadata.py b/python/sglang/srt/speculative/dspark_components/kernels/dspark_attn_metadata.py new file mode 100644 index 000000000..a18076a92 --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/kernels/dspark_attn_metadata.py @@ -0,0 +1,491 @@ +from __future__ import annotations + +from typing import Tuple + +import msgspec +import torch +import triton +import triton.language as tl + +from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda +from sglang.srt.utils import ceil_align + + +class DsparkWindowGather(msgspec.Struct, frozen=True): + num_q: int + bs: int + context_lens: torch.Tensor + req_pool_indices_per_request: torch.Tensor + offsets: torch.Tensor + invalid: torch.Tensor + + +class ComputeDsparkWindowGather: + @classmethod + def execute(cls, *args, **kwargs) -> DsparkWindowGather: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + seq_lens_casual: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + block_size: int, + swa_window: int, + ) -> DsparkWindowGather: + return compute_dspark_window_gather( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + block_size=block_size, + swa_window=swa_window, + ) + + @classmethod + def triton( + cls, + *, + seq_lens_casual: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + block_size: int, + swa_window: int, + ) -> DsparkWindowGather: + return compute_dspark_window_gather_triton( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + block_size=block_size, + swa_window=swa_window, + ) + + +class BuildDsparkSwaPageIndices: + @classmethod + def execute(cls, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + req_to_token: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + req_pool_indices_per_request: torch.Tensor, + offsets: torch.Tensor, + invalid: torch.Tensor, + out_loc: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + swa_window: int, + page_index_aligned_size: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return build_dspark_swa_page_indices( + req_to_token=req_to_token, + full_to_swa_mapping=full_to_swa_mapping, + req_pool_indices_per_request=req_pool_indices_per_request, + offsets=offsets, + invalid=invalid, + out_loc=out_loc, + context_lens=context_lens, + block_size=block_size, + swa_window=swa_window, + page_index_aligned_size=page_index_aligned_size, + ) + + @classmethod + def triton( + cls, + *, + req_to_token: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + req_pool_indices_per_request: torch.Tensor, + offsets: torch.Tensor, + invalid: torch.Tensor, + out_loc: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + swa_window: int, + page_index_aligned_size: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return build_dspark_swa_page_indices_triton( + req_to_token=req_to_token, + full_to_swa_mapping=full_to_swa_mapping, + req_pool_indices_per_request=req_pool_indices_per_request, + offsets=offsets, + out_loc=out_loc, + context_lens=context_lens, + block_size=block_size, + swa_window=swa_window, + page_index_aligned_size=page_index_aligned_size, + ) + + +def compute_dspark_window_gather( + *, + seq_lens_casual: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + block_size: int, + swa_window: int, +) -> DsparkWindowGather: + seq_lens_casual = seq_lens_casual.to(torch.int32) + num_q = seq_lens_casual.size(0) + assert num_q % block_size == 0, ( + f"DSpark draft block forward must be uniform-gamma: num_q={num_q} not " + f"divisible by block_size={block_size}." + ) + bs = num_q // block_size + device = seq_lens_casual.device + + first_token = torch.arange(bs, device=device, dtype=torch.int64) * block_size + prefix_lens = (seq_lens_casual[first_token] - 1).to(torch.int32) + context_lens = torch.clamp(prefix_lens, max=swa_window).to(torch.int32) + req_pool_indices_per_request = req_pool_indices_repeated[first_token] + + offsets = ( + prefix_lens.to(torch.int64).unsqueeze(1) + - swa_window + + torch.arange(swa_window, device=device, dtype=torch.int64).unsqueeze(0) + ) + invalid = offsets < 0 + offsets = offsets.clamp(min=0) + + return DsparkWindowGather( + num_q=num_q, + bs=bs, + context_lens=context_lens, + req_pool_indices_per_request=req_pool_indices_per_request, + offsets=offsets, + invalid=invalid, + ) + + +@triton.jit +def _window_gather_kernel( + seq_lens_casual_ptr, + req_pool_rep_ptr, + context_lens_ptr, + req_pool_out_ptr, + offsets_ptr, + invalid_ptr, + block_size, + swa_window, + W_BLOCK: tl.constexpr, +): + i = tl.program_id(0) + ft = i * block_size + prefix = tl.load(seq_lens_casual_ptr + ft).to(tl.int64) - 1 + tl.store(context_lens_ptr + i, tl.minimum(prefix, swa_window).to(tl.int32)) + tl.store(req_pool_out_ptr + i, tl.load(req_pool_rep_ptr + ft)) + col = tl.arange(0, W_BLOCK) + cmask = col < swa_window + off = prefix - swa_window + col + tl.store(invalid_ptr + i * swa_window + col, off < 0, mask=cmask) + tl.store(offsets_ptr + i * swa_window + col, tl.maximum(off, 0), mask=cmask) + + +def compute_dspark_window_gather_triton( + *, + seq_lens_casual: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + block_size: int, + swa_window: int, +) -> DsparkWindowGather: + seq_lens_casual = seq_lens_casual.to(torch.int32).contiguous() + num_q = seq_lens_casual.size(0) + assert num_q % block_size == 0, ( + f"DSpark draft block forward must be uniform-gamma: num_q={num_q} not " + f"divisible by block_size={block_size}." + ) + bs = num_q // block_size + device = seq_lens_casual.device + req_pool_indices_repeated = req_pool_indices_repeated.to(device=device).contiguous() + context_lens = torch.empty(bs, dtype=torch.int32, device=device) + req_pool_out = torch.empty(bs, dtype=req_pool_indices_repeated.dtype, device=device) + offsets = torch.empty((bs, swa_window), dtype=torch.int64, device=device) + invalid = torch.empty((bs, swa_window), dtype=torch.bool, device=device) + W_BLOCK = triton.next_power_of_2(swa_window) + _window_gather_kernel[(bs,)]( + seq_lens_casual, + req_pool_indices_repeated, + context_lens, + req_pool_out, + offsets, + invalid, + block_size, + swa_window, + W_BLOCK=W_BLOCK, + ) + return DsparkWindowGather( + num_q=num_q, + bs=bs, + context_lens=context_lens, + req_pool_indices_per_request=req_pool_out, + offsets=offsets, + invalid=invalid, + ) + + +def build_dspark_swa_page_indices( + *, + req_to_token: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + req_pool_indices_per_request: torch.Tensor, + offsets: torch.Tensor, + invalid: torch.Tensor, + out_loc: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + swa_window: int, + page_index_aligned_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + if offsets.ndim != 2 or offsets.shape[1] != swa_window: + raise ValueError( + "offsets must be [bs, swa_window]; " + f"got shape={tuple(offsets.shape)} (swa_window={swa_window})." + ) + bs = offsets.shape[0] + device = offsets.device + context_lens = context_lens.to(device=device, dtype=torch.int32) + + window_full_locs = req_to_token[ + req_pool_indices_per_request[:, None].to(torch.int64), offsets + ] + window_full_locs = window_full_locs.masked_fill(invalid, 0) + window_swa_locs = full_to_swa_mapping[window_full_locs].to(torch.int32) + window_swa_locs = window_swa_locs.masked_fill(invalid, -1) + + block_full_locs = out_loc[: bs * block_size].view(bs, block_size) + block_swa_locs = full_to_swa_mapping[block_full_locs].to(torch.int32) + + target_width = ceil_align(swa_window + block_size, page_index_aligned_size) + + swa_page_indices = _compact_dspark_window_then_block( + window_swa_locs=window_swa_locs, + block_swa_locs=block_swa_locs, + context_lens=context_lens, + target_width=target_width, + block_size=block_size, + swa_window=swa_window, + ) + + swa_page_indices = ( + swa_page_indices.view(bs, 1, target_width) + .expand(bs, block_size, target_width) + .reshape(bs * block_size, target_width) + .contiguous() + ) + swa_topk_lengths = ( + (context_lens + block_size) + .view(bs, 1) + .expand(bs, block_size) + .reshape(bs * block_size) + .contiguous() + .to(torch.int32) + ) + return swa_page_indices, swa_topk_lengths + + +def _compact_dspark_window_then_block( + *, + window_swa_locs: torch.Tensor, + block_swa_locs: torch.Tensor, + context_lens: torch.Tensor, + target_width: int, + block_size: int, + swa_window: int, +) -> torch.Tensor: + bs = window_swa_locs.shape[0] + device = window_swa_locs.device + out = torch.full((bs, target_width), -1, dtype=torch.int32, device=device) + + j = torch.arange(swa_window, device=device, dtype=torch.int32).view(1, -1) + shift = (swa_window - context_lens.view(-1, 1)).to(torch.int32) + src_col = (shift + j).clamp_(min=0, max=swa_window - 1).to(torch.int64) + gathered = torch.gather(window_swa_locs, dim=1, index=src_col) + valid = j < context_lens.view(-1, 1) + out[:, :swa_window] = torch.where(valid, gathered, -1) + + block_col = context_lens.view(-1, 1) + torch.arange( + block_size, device=device, dtype=torch.int32 + ).view(1, -1) + block_rows = torch.arange(bs, device=device).view(-1, 1).expand(-1, block_size) + out[block_rows, block_col] = block_swa_locs + return out + + +@triton.jit +def _swa_page_indices_kernel( + req_to_token_ptr, + full_to_swa_ptr, + req_pool_ptr, + offsets_ptr, + out_loc_ptr, + context_lens_ptr, + out_ptr, + topk_ptr, + rt_stride, + swa_window, + block_size, + target_width, + TW_BLOCK: tl.constexpr, +): + q = tl.program_id(0) + i = q // block_size + cl = tl.load(context_lens_ptr + i) + rp = tl.load(req_pool_ptr + i).to(tl.int64) + k = tl.arange(0, TW_BLOCK) + kmask = k < target_width + in_window = k < cl + src_col = tl.minimum(tl.maximum((swa_window - cl) + k, 0), swa_window - 1) + wmask = kmask & in_window + off = tl.load(offsets_ptr + i * swa_window + src_col, mask=wmask, other=0).to( + tl.int64 + ) + win_full = tl.load(req_to_token_ptr + rp * rt_stride + off, mask=wmask, other=0).to( + tl.int64 + ) + win_swa = tl.load(full_to_swa_ptr + win_full, mask=wmask, other=-1).to(tl.int32) + + in_block = (k >= cl) & (k < cl + block_size) + bmask = kmask & in_block + bcol = tl.maximum(k - cl, 0) + blk_full = tl.load(out_loc_ptr + i * block_size + bcol, mask=bmask, other=0).to( + tl.int64 + ) + blk_swa = tl.load(full_to_swa_ptr + blk_full, mask=bmask, other=-1).to(tl.int32) + + val = tl.where(in_window, win_swa, tl.where(in_block, blk_swa, -1)) + tl.store(out_ptr + q * target_width + k, val.to(tl.int32), mask=kmask) + tl.store(topk_ptr + q, (cl + block_size).to(tl.int32)) + + +def build_dspark_swa_page_indices_triton( + *, + req_to_token: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + req_pool_indices_per_request: torch.Tensor, + offsets: torch.Tensor, + out_loc: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + swa_window: int, + page_index_aligned_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + if offsets.ndim != 2 or offsets.shape[1] != swa_window: + raise ValueError( + "offsets must be [bs, swa_window]; " + f"got shape={tuple(offsets.shape)} (swa_window={swa_window})." + ) + bs = offsets.shape[0] + device = offsets.device + req_pool = req_pool_indices_per_request.to(device=device).contiguous() + offsets = offsets.to(torch.int64).contiguous() + out_loc = out_loc[: bs * block_size].contiguous() + context_lens = context_lens.to(device=device, dtype=torch.int32).contiguous() + rt_stride = req_to_token.stride(0) + target_width = ceil_align(swa_window + block_size, page_index_aligned_size) + n_q = bs * block_size + swa_page_indices = torch.empty( + (n_q, target_width), dtype=torch.int32, device=device + ) + swa_topk_lengths = torch.empty(n_q, dtype=torch.int32, device=device) + TW_BLOCK = triton.next_power_of_2(target_width) + _swa_page_indices_kernel[(n_q,)]( + req_to_token, + full_to_swa_mapping, + req_pool, + offsets, + out_loc, + context_lens, + swa_page_indices, + swa_topk_lengths, + rt_stride, + swa_window, + block_size, + target_width, + TW_BLOCK=TW_BLOCK, + ) + return swa_page_indices, swa_topk_lengths + + +class BuildBlockSeqLensCausal: + @classmethod + def execute(cls, *args, **kwargs) -> torch.Tensor: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + seq_lens: torch.Tensor, + block_size: int, + device: torch.device, + ) -> torch.Tensor: + return build_block_seq_lens_causal( + seq_lens=seq_lens, + block_size=block_size, + device=device, + ) + + @classmethod + def triton( + cls, + *, + seq_lens: torch.Tensor, + block_size: int, + device: torch.device, + ) -> torch.Tensor: + return build_block_seq_lens_causal_triton( + seq_lens=seq_lens, + block_size=block_size, + device=device, + ) + + +def build_block_seq_lens_causal( + *, + seq_lens: torch.Tensor, + block_size: int, + device: torch.device, +) -> torch.Tensor: + prefix = seq_lens.to(torch.int32) + steps = torch.arange(1, block_size + 1, device=device, dtype=torch.int32) + return (prefix[:, None] + steps[None, :]).reshape(-1) + + +@triton.jit +def _block_seq_lens_casual_kernel( + seq_lens_ptr, + out_ptr, + block_size, + n_out, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_out + row = offs // block_size + col = offs % block_size + prefix = tl.load(seq_lens_ptr + row, mask=mask, other=0) + tl.store(out_ptr + offs, (prefix + col + 1).to(tl.int32), mask=mask) + + +def build_block_seq_lens_causal_triton( + *, + seq_lens: torch.Tensor, + block_size: int, + device: torch.device, +) -> torch.Tensor: + seq_lens = seq_lens.to(device=device, dtype=torch.int64).contiguous() + n_rows = seq_lens.shape[0] + n_out = n_rows * block_size + out = torch.empty(n_out, dtype=torch.int32, device=device) + BLOCK = 256 + grid = (triton.cdiv(n_out, BLOCK),) + _block_seq_lens_casual_kernel[grid](seq_lens, out, block_size, n_out, BLOCK=BLOCK) + return out diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_draft_model.py b/python/sglang/srt/speculative/dspark_components/kernels/dspark_draft_model.py new file mode 100644 index 000000000..d09c717e1 --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/kernels/dspark_draft_model.py @@ -0,0 +1,443 @@ +from __future__ import annotations + +from typing import Optional + +import msgspec +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda + +_BLOCK_V = 1024 +_IDX_SENTINEL = tl.constexpr(2147483647) + + +class SampleStepTokens: + @classmethod + def execute( + cls, + *, + step_logits: torch.Tensor, + temperatures: torch.Tensor, + greedy_mask: torch.Tensor, + exp_noise: torch.Tensor, + ) -> torch.Tensor: + if step_logits.is_cuda: + return cls.triton( + step_logits=step_logits, + temperatures=temperatures, + greedy_mask=greedy_mask, + exp_noise=exp_noise, + ) + return cls.torch( + step_logits=step_logits, + temperatures=temperatures, + greedy_mask=greedy_mask, + exp_noise=exp_noise, + ) + + @classmethod + def torch( + cls, + *, + step_logits: torch.Tensor, + temperatures: torch.Tensor, + greedy_mask: torch.Tensor, + exp_noise: torch.Tensor, + ) -> torch.Tensor: + return sample_step_tokens( + step_logits=step_logits, + temperatures=temperatures, + greedy_mask=greedy_mask, + exp_noise=exp_noise, + ) + + @classmethod + def triton( + cls, + *, + step_logits: torch.Tensor, + temperatures: torch.Tensor, + greedy_mask: torch.Tensor, + exp_noise: torch.Tensor, + ) -> torch.Tensor: + return sample_step_tokens_triton( + step_logits=step_logits, + temperatures=temperatures, + greedy_mask=greedy_mask, + exp_noise=exp_noise, + ) + + +def sample_step_tokens( + *, + step_logits: torch.Tensor, + temperatures: torch.Tensor, + greedy_mask: torch.Tensor, + exp_noise: torch.Tensor, +) -> torch.Tensor: + probs = torch.softmax(step_logits.float() / temperatures[:, None], dim=-1) + noise = torch.where(greedy_mask[:, None], 1.0, exp_noise) + return probs.div_(noise).argmax(dim=-1) + + +@triton.jit +def _online_partial_kernel( + logits_ptr, + temperatures_ptr, + greedy_mask_ptr, + exp_noise_ptr, + tile_max_ptr, + partial_key_ptr, + partial_idx_ptr, + V, + stride_row, + n_tiles, + BLOCK_V: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + offs = tile * BLOCK_V + tl.arange(0, BLOCK_V) + mask = offs < V + logits = tl.load( + logits_ptr + row * stride_row + offs, mask=mask, other=float("-inf") + ).to(tl.float32) + temperature = tl.load(temperatures_ptr + row) + s = logits / temperature + tile_max = tl.max(s, axis=0) + greedy = tl.load(greedy_mask_ptr + row) != 0 + noise = tl.load(exp_noise_ptr + row * V + offs, mask=mask, other=1.0) + denom = tl.where(greedy, 1.0, noise) + key = tl.exp(s - tile_max) / denom + key = tl.where(mask, key, -1.0) + tile_best = tl.max(key, axis=0) + idx = tl.where(key == tile_best, offs, _IDX_SENTINEL) + tl.store(tile_max_ptr + row * n_tiles + tile, tile_max) + tl.store(partial_key_ptr + row * n_tiles + tile, tile_best) + tl.store(partial_idx_ptr + row * n_tiles + tile, tl.min(idx, axis=0)) + + +@triton.jit +def _online_combine_kernel( + tile_max_ptr, + partial_key_ptr, + partial_idx_ptr, + next_tokens_ptr, + n_tiles, + BLOCK_TILES: tl.constexpr, +): + row = tl.program_id(0) + offs = tl.arange(0, BLOCK_TILES) + mask = offs < n_tiles + tile_max = tl.load( + tile_max_ptr + row * n_tiles + offs, mask=mask, other=float("-inf") + ) + keys = tl.load(partial_key_ptr + row * n_tiles + offs, mask=mask, other=-1.0) + idxs = tl.load( + partial_idx_ptr + row * n_tiles + offs, mask=mask, other=_IDX_SENTINEL + ) + global_max = tl.max(tile_max, axis=0) + rescaled = keys * tl.exp(tile_max - global_max) + rescaled = tl.where(mask, rescaled, -1.0) + best = tl.max(rescaled, axis=0) + cand = tl.where(rescaled == best, idxs, _IDX_SENTINEL) + tl.store(next_tokens_ptr + row, tl.min(cand, axis=0).to(tl.int64)) + + +def sample_step_tokens_triton( + *, + step_logits: torch.Tensor, + temperatures: torch.Tensor, + greedy_mask: torch.Tensor, + exp_noise: torch.Tensor, +) -> torch.Tensor: + bs, V = step_logits.shape + device = step_logits.device + assert step_logits.stride(1) == 1, "step_logits rows must be contiguous" + stride_row = step_logits.stride(0) + temperatures = temperatures.to(torch.float32).contiguous() + greedy_mask = greedy_mask.to(torch.int32).contiguous() + exp_noise = exp_noise.to(torch.float32).contiguous() + + n_tiles = triton.cdiv(V, _BLOCK_V) + block_tiles = triton.next_power_of_2(n_tiles) + + tile_max = torch.empty((bs, n_tiles), dtype=torch.float32, device=device) + partial_key = torch.empty((bs, n_tiles), dtype=torch.float32, device=device) + partial_idx = torch.empty((bs, n_tiles), dtype=torch.int32, device=device) + next_tokens = torch.empty((bs,), dtype=torch.int64, device=device) + + tile_grid = (bs, n_tiles) + row_grid = (bs,) + + _online_partial_kernel[tile_grid]( + step_logits, + temperatures, + greedy_mask, + exp_noise, + tile_max, + partial_key, + partial_idx, + V, + stride_row, + n_tiles, + BLOCK_V=_BLOCK_V, + ) + _online_combine_kernel[row_grid]( + tile_max, + partial_key, + partial_idx, + next_tokens, + n_tiles, + BLOCK_TILES=block_tiles, + ) + return next_tokens + + +_STACKED_WEIGHT_CACHE: dict[int, _StackedWkvWeight] = {} + + +class CommitKvProj: + @classmethod + def execute( + cls, + *, + main_x: torch.Tensor, + wkv_linears: list[torch.nn.Module], + ) -> list[torch.Tensor]: + if main_x.is_cuda and _fused_commit_kv_proj_supported(wkv_linears=wkv_linears): + return cls.triton(main_x=main_x, wkv_linears=wkv_linears) + return cls.torch(main_x=main_x, wkv_linears=wkv_linears) + + @classmethod + def torch( + cls, + *, + main_x: torch.Tensor, + wkv_linears: list[torch.nn.Module], + ) -> list[torch.Tensor]: + return commit_kv_proj(main_x=main_x, wkv_linears=wkv_linears) + + @classmethod + def triton( + cls, + *, + main_x: torch.Tensor, + wkv_linears: list[torch.nn.Module], + ) -> list[torch.Tensor]: + return commit_kv_proj_fused(main_x=main_x, wkv_linears=wkv_linears) + + +def commit_kv_proj( + *, + main_x: torch.Tensor, + wkv_linears: list[torch.nn.Module], +) -> list[torch.Tensor]: + return [linear(main_x)[0] for linear in wkv_linears] + + +def commit_kv_proj_fused( + *, + main_x: torch.Tensor, + wkv_linears: list[torch.nn.Module], +) -> list[torch.Tensor]: + num_stages = len(wkv_linears) + stacked = _stacked_wkv_weight(wkv_linears=wkv_linears) + + if stacked.fp8_scale is not None: + quant_method = wkv_linears[0].quant_method + kv_all = quant_method.w8a8_block_fp8_linear( + input=main_x, + weight=stacked.weight, + block_size=quant_method.quant_config.weight_block_size, + weight_scale=stacked.fp8_scale, + input_scale=None, + bias=None, + ) + else: + kv_all = torch.nn.functional.linear(main_x, stacked.weight) + + head_dim = kv_all.shape[-1] // num_stages + return [ + kv_all[:, i * head_dim : (i + 1) * head_dim].contiguous() + for i in range(num_stages) + ] + + +class _StackedWkvWeight(msgspec.Struct): + weight: torch.Tensor + fp8_scale: Optional[torch.Tensor] + + +def _stacked_wkv_weight(*, wkv_linears: list[torch.nn.Module]) -> _StackedWkvWeight: + key = id(wkv_linears[0]) + cached = _STACKED_WEIGHT_CACHE.get(key) + if cached is None: + cached = _build_stacked_wkv_weight(wkv_linears=wkv_linears) + _STACKED_WEIGHT_CACHE[key] = cached + return cached + + +def _block_quant_stack_applies(*, wkv_linears: list[torch.nn.Module]) -> bool: + quant_method = wkv_linears[0].quant_method + block_quant = hasattr(quant_method, "block_quant") and quant_method.block_quant + if not (block_quant and hasattr(quant_method, "w8a8_block_fp8_linear")): + return False + block_out = quant_method.quant_config.weight_block_size[0] + return all( + linear.weight.dtype == torch.float8_e4m3fn + and linear.weight.shape[0] % block_out == 0 + for linear in wkv_linears + ) + + +def _dequant_supported(linear: torch.nn.Module) -> bool: + """Mirrors the preconditions asserted in _dequant_linear_weight.""" + weight = linear.weight + if weight.dtype in (torch.bfloat16, torch.float16, torch.float32): + return True + if weight.dtype != torch.float8_e4m3fn: + return False + block = 128 + out_dim, in_dim = weight.shape + expected_scale_shape = ( + (out_dim + block - 1) // block, + (in_dim + block - 1) // block, + ) + return tuple(linear.weight_scale_inv.shape) == expected_scale_shape + + +def _fused_commit_kv_proj_supported(*, wkv_linears: list[torch.nn.Module]) -> bool: + """Whether _build_stacked_wkv_weight can handle these weights; unsupported + quant schemes fall back to the per-linear torch path in execute().""" + if _block_quant_stack_applies(wkv_linears=wkv_linears): + return True + return all(_dequant_supported(linear) for linear in wkv_linears) + + +def _build_stacked_wkv_weight( + *, wkv_linears: list[torch.nn.Module] +) -> _StackedWkvWeight: + if _block_quant_stack_applies(wkv_linears=wkv_linears): + weight = torch.cat([linear.weight for linear in wkv_linears], dim=0) + if wkv_linears[0].weight_scale_inv.dtype == torch.int32: + from sglang.srt.layers.quantization.fp8_utils import ( + inverse_transform_scale_ue8m0, + transform_scale_ue8m0, + ) + + sf_fp32 = torch.cat( + [ + inverse_transform_scale_ue8m0( + linear.weight_scale_inv, mn=linear.weight.shape[0] + ) + for linear in wkv_linears + ], + dim=0, + ) + scale = transform_scale_ue8m0(sf_fp32, mn=weight.shape[0]) + return _StackedWkvWeight(weight=weight, fp8_scale=scale) + scale = torch.cat([linear.weight_scale_inv for linear in wkv_linears], dim=0) + if scale.dim() >= 2 and scale.stride(-2) != 1: + scale = scale.transpose(-2, -1).contiguous().transpose(-2, -1) + return _StackedWkvWeight(weight=weight, fp8_scale=scale) + weight = torch.cat( + [_dequant_linear_weight(linear) for linear in wkv_linears], dim=0 + ) + return _StackedWkvWeight(weight=weight, fp8_scale=None) + + +def _dequant_linear_weight(linear: torch.nn.Module) -> torch.Tensor: + weight = linear.weight + if weight.dtype in (torch.bfloat16, torch.float16, torch.float32): + return weight.to(torch.bfloat16) + assert weight.dtype == torch.float8_e4m3fn, ( + f"unsupported wkv weight dtype {weight.dtype} for the fused commit kv proj; " + f"execute() should have routed this to the torch path " + f"(_fused_commit_kv_proj_supported)" + ) + block = 128 + scale = linear.weight_scale_inv + out_dim, in_dim = weight.shape + expected_scale_shape = ( + (out_dim + block - 1) // block, + (in_dim + block - 1) // block, + ) + assert tuple(scale.shape) == expected_scale_shape, ( + f"wkv weight_scale_inv shape {tuple(scale.shape)} does not match the " + f"128x128 block grid {expected_scale_shape} for weight {tuple(weight.shape)}; " + f"execute() should have routed this to the torch path " + f"(_fused_commit_kv_proj_supported)" + ) + scale_full = scale.repeat_interleave(block, dim=0)[:out_dim] + scale_full = scale_full.repeat_interleave(block, dim=1)[:, :in_dim] + return (weight.to(torch.float32) * scale_full.to(torch.float32)).to(torch.bfloat16) + + +_BLOCK = 1024 + + +class BuildStepLocal: + @classmethod + def execute(cls, *args, **kwargs) -> torch.Tensor: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch(cls, *, bias: torch.Tensor, base_local: torch.Tensor) -> torch.Tensor: + return build_step_local(bias=bias, base_local=base_local) + + @classmethod + def triton(cls, *, bias: torch.Tensor, base_local: torch.Tensor) -> torch.Tensor: + return build_step_local_triton(bias=bias, base_local=base_local) + + +def build_step_local(*, bias: torch.Tensor, base_local: torch.Tensor) -> torch.Tensor: + per_partition = base_local.shape[-1] + pad = per_partition - bias.shape[-1] + padded = ( + F.pad(bias.to(torch.float32), (0, pad)) if pad > 0 else bias.to(torch.float32) + ) + return base_local + padded + + +@triton.jit +def _build_step_local_kernel( + bias_ptr, + base_ptr, + out_ptr, + org_width, + per_partition, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + offs = tile * BLOCK + tl.arange(0, BLOCK) + mask = offs < per_partition + base = tl.load(base_ptr + row * per_partition + offs, mask=mask, other=0.0).to( + tl.float32 + ) + bias = tl.load( + bias_ptr + row * org_width + offs, mask=offs < org_width, other=0.0 + ).to(tl.float32) + tl.store(out_ptr + row * per_partition + offs, base + bias, mask=mask) + + +def build_step_local_triton( + *, bias: torch.Tensor, base_local: torch.Tensor +) -> torch.Tensor: + bs, per_partition = base_local.shape + org_width = bias.shape[-1] + base_local = base_local.contiguous() + bias = bias.contiguous() + out = torch.empty( + (bs, per_partition), dtype=torch.float32, device=base_local.device + ) + grid = (bs, triton.cdiv(per_partition, _BLOCK)) + _build_step_local_kernel[grid]( + bias, base_local, out, org_width, per_partition, BLOCK=_BLOCK + ) + return out diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_schedule.py b/python/sglang/srt/speculative/dspark_components/kernels/dspark_schedule.py new file mode 100644 index 000000000..52ffc74bb --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/kernels/dspark_schedule.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import triton +import triton.language as tl + +from sglang.srt.speculative.dspark_components.kernels.dispatch import ( + inputs_on_cuda, +) + +if TYPE_CHECKING: + from sglang.srt.speculative.dspark_components.dspark_planner import ( + DSparkScheduleConfig, + ) + + +class ScheduleVerifyLensTopk: + @classmethod + def execute(cls, *args, **kwargs) -> torch.Tensor: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + confidence: torch.Tensor, + budget: int, + cfg: DSparkScheduleConfig, + ) -> torch.Tensor: + return schedule_verify_lens_topk(confidence=confidence, budget=budget, cfg=cfg) + + @classmethod + def triton( + cls, + *, + confidence: torch.Tensor, + budget: int, + cfg: DSparkScheduleConfig, + ) -> torch.Tensor: + return schedule_verify_lens_topk_triton( + confidence=confidence, budget=budget, cfg=cfg + ) + + +def compute_sort_survival(confidence: torch.Tensor) -> torch.Tensor: + return torch.cumprod(confidence.to(torch.float32), dim=1) + + +def schedule_verify_lens_topk( + *, + confidence: torch.Tensor, + budget: int, + cfg: DSparkScheduleConfig, +) -> torch.Tensor: + return schedule_verify_lens_topk_from_survival( + survival_probs=compute_sort_survival(confidence), budget=budget, cfg=cfg + ) + + +def schedule_verify_lens_topk_from_survival( + *, + survival_probs: torch.Tensor, + budget: int, + cfg: DSparkScheduleConfig, +) -> torch.Tensor: + num_requests, _gamma = survival_probs.shape + max_len = cfg.resolved_max_verify_len() + device = survival_probs.device + + selected_extra = torch.zeros(num_requests, dtype=torch.int64, device=device) + if budget > 0: + candidate_window = survival_probs[:, :max_len] + num_candidates = candidate_window.numel() + if num_candidates > 0: + request_index = ( + torch.arange(num_requests, device=device) + .view(num_requests, 1) + .expand_as(candidate_window) + ) + position_index = ( + torch.arange(candidate_window.shape[1], device=device) + .view(1, candidate_window.shape[1]) + .expand_as(candidate_window) + ) + valid = candidate_window >= cfg.survival_eps + + flat_prob = candidate_window.reshape(-1).to(torch.float64) + flat_request = request_index.reshape(-1) + flat_position = position_index.reshape(-1) + flat_valid = valid.reshape(-1) + + order = _value_independent_descending_order( + probs=flat_prob, + positions=flat_position, + requests=flat_request, + valid=flat_valid, + ) + + take = min(int(budget), num_candidates) + chosen = order[:take] + chosen_requests = flat_request[chosen] + chosen_valid = flat_valid[chosen].to(torch.int64) + selected_extra.scatter_add_(0, chosen_requests, chosen_valid) + + min_len = torch.full( + (num_requests,), cfg.min_verify_len, dtype=torch.int64, device=device + ) + verify_lens = min_len + selected_extra + lower_bound = max(cfg.min_verify_len, 1) + verify_lens = torch.clamp(verify_lens, min=lower_bound, max=max_len) + return verify_lens.to(torch.int32) + + +def _value_independent_descending_order( + *, + probs: torch.Tensor, + positions: torch.Tensor, + requests: torch.Tensor, + valid: torch.Tensor, +) -> torch.Tensor: + masked_prob = torch.where(valid, probs, torch.full_like(probs, float("-inf"))) + num_candidates = masked_prob.numel() + order = torch.arange(num_candidates, device=probs.device) + order = order[torch.argsort(requests[order], stable=True)] + order = order[torch.argsort(positions[order], stable=True)] + order = order[torch.argsort(-masked_prob[order], stable=True)] + return order + + +@triton.jit +def _schedule_topk_prep_kernel( + confidence_ptr, + survival_ptr, + selected_extra_ptr, + gamma, + cols, + G_P2: tl.constexpr, +): + row = tl.program_id(0) + g = tl.arange(0, G_P2) + conf = tl.load( + confidence_ptr + row.to(tl.int64) * gamma + g, mask=g < gamma, other=1.0 + ).to(tl.float32) + surv = tl.cumprod(conf, axis=0) + tl.store(survival_ptr + row.to(tl.int64) * cols + g, surv, mask=g < cols) + tl.store(selected_extra_ptr + row, 0) + + +@triton.jit +def _schedule_topk_finalize_kernel( + selected_extra_ptr, + out_ptr, + min_verify_len, + lower_bound, + max_len, + bs, + BLOCK: tl.constexpr, +): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < bs + extra = tl.load(selected_extra_ptr + offs, mask=mask, other=0).to(tl.int32) + lens = min_verify_len + extra + lens = tl.maximum(lens, lower_bound) + lens = tl.minimum(lens, max_len) + tl.store(out_ptr + offs, lens, mask=mask) + + +@triton.jit +def _schedule_topk_selected_extra_kernel( + survival_ptr, + selected_extra_ptr, + budget, + cols, + n, + survival_eps, + BLOCK_C: tl.constexpr, + BLOCK_CP: tl.constexpr, +): + pid = tl.program_id(0) + c = pid * BLOCK_C + tl.arange(0, BLOCK_C) + cmask = c < n + r = c // cols + p = c % cols + sp = tl.load(survival_ptr + c, mask=cmask, other=0.0) + valid_c = sp >= survival_eps + mp = tl.where(valid_c, sp, float("-inf")) + rank = tl.zeros([BLOCK_C], dtype=tl.int32) + for cp0 in range(0, n, BLOCK_CP): + cp = cp0 + tl.arange(0, BLOCK_CP) + cpmask = cp < n + rp = cp // cols + pp = cp % cols + spp = tl.load(survival_ptr + cp, mask=cpmask, other=0.0) + validp = spp >= survival_eps + mpp = tl.where(validp, spp, float("-inf")) + gt = mpp[None, :] > mp[:, None] + eq = mpp[None, :] == mp[:, None] + pos_lt = pp[None, :] < p[:, None] + pos_eq = pp[None, :] == p[:, None] + req_lt = rp[None, :] < r[:, None] + before = gt | (eq & (pos_lt | (pos_eq & req_lt))) + before = before & cpmask[None, :] + rank += tl.sum(before.to(tl.int32), axis=1) + selected = valid_c & (rank < budget) + tl.atomic_add(selected_extra_ptr + r, selected.to(tl.int32), mask=cmask) + + +def schedule_verify_lens_topk_triton( + *, + confidence: torch.Tensor, + budget: int, + cfg: DSparkScheduleConfig, +) -> torch.Tensor: + num_requests, gamma = confidence.shape + max_len = cfg.resolved_max_verify_len() + device = confidence.device + cols = min(max_len, gamma) + n = num_requests * cols + + selected_extra = torch.empty(num_requests, dtype=torch.int32, device=device) + survival = torch.empty((num_requests, cols), dtype=torch.float32, device=device) + _schedule_topk_prep_kernel[(num_requests,)]( + confidence.contiguous(), + survival, + selected_extra, + gamma, + cols, + G_P2=triton.next_power_of_2(max(gamma, 1)), + ) + if budget > 0 and n > 0: + BLOCK_C = 64 + BLOCK_CP = 256 + grid = (triton.cdiv(n, BLOCK_C),) + _schedule_topk_selected_extra_kernel[grid]( + survival, + selected_extra, + int(budget), + cols, + n, + float(cfg.survival_eps), + BLOCK_C=BLOCK_C, + BLOCK_CP=BLOCK_CP, + ) + + verify_lens = torch.empty(num_requests, dtype=torch.int32, device=device) + BLOCK = 256 + _schedule_topk_finalize_kernel[(triton.cdiv(num_requests, BLOCK),)]( + selected_extra, + verify_lens, + int(cfg.min_verify_len), + max(cfg.min_verify_len, 1), + int(max_len), + num_requests, + BLOCK=BLOCK, + ) + return verify_lens diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py b/python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py new file mode 100644 index 000000000..62f396d1a --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py @@ -0,0 +1,871 @@ +from __future__ import annotations + +import msgspec +import torch +import triton +import triton.language as tl + +from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda +from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout + + +class RaggedVerifyWindow(msgspec.Struct, frozen=True): + positions: torch.Tensor + verify_cache_loc: torch.Tensor + verify_ids: torch.Tensor + + +class BuildRaggedVerifyWindow: + @classmethod + def execute(cls, *args, **kwargs) -> RaggedVerifyWindow: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + batch: ScheduleBatch, + layout: RaggedVerifyLayout, + draft_block_ids: torch.Tensor, + draft_tokens: torch.Tensor, + bs: int, + device: str, + verify_num_draft_tokens: int, + model_runner, + ) -> RaggedVerifyWindow: + return build_ragged_verify_window( + batch=batch, + layout=layout, + draft_block_ids=draft_block_ids, + draft_tokens=draft_tokens, + bs=bs, + device=device, + verify_num_draft_tokens=verify_num_draft_tokens, + model_runner=model_runner, + ) + + @classmethod + def triton( + cls, + *, + batch: ScheduleBatch, + layout: RaggedVerifyLayout, + draft_block_ids: torch.Tensor, + draft_tokens: torch.Tensor, + bs: int, + device: str, + verify_num_draft_tokens: int, + model_runner, + ) -> RaggedVerifyWindow: + return build_ragged_verify_window_triton( + batch=batch, + layout=layout, + draft_block_ids=draft_block_ids, + draft_tokens=draft_tokens, + bs=bs, + device=device, + verify_num_draft_tokens=verify_num_draft_tokens, + model_runner=model_runner, + ) + + +def build_ragged_verify_window( + *, + batch: ScheduleBatch, + layout: RaggedVerifyLayout, + draft_block_ids: torch.Tensor, + draft_tokens: torch.Tensor, + bs: int, + device: str, + verify_num_draft_tokens: int, + model_runner, +) -> RaggedVerifyWindow: + prefix_lens = batch.seq_lens + verify_lens = layout.verify_lens.to(device=device, dtype=torch.int32) + padded_total = layout.graph_num_tokens + + req_id, within, valid = compact_row_index( + verify_lens=verify_lens, padded_total=padded_total, device=device + ) + safe_req = req_id.clamp(max=bs - 1) + positions = torch.where( + valid, + prefix_lens.to(torch.int64)[safe_req] + within, + torch.zeros_like(within), + ) + real_cache_loc = assign_extend_cache_locs_func( + req_pool_indices=batch.req_pool_indices, + req_to_token=model_runner.req_to_token_pool.req_to_token, + start_offset=prefix_lens, + end_offset=prefix_lens + verify_lens.to(prefix_lens.dtype), + batch_size=bs, + draft_token_num=verify_num_draft_tokens, + device=device, + ) + verify_cache_loc = torch.nn.functional.pad( + real_cache_loc, (0, padded_total - real_cache_loc.shape[0]) + ) + verify_cache_loc = torch.where( + valid, verify_cache_loc, torch.zeros_like(verify_cache_loc) + ) + + verify_ids = compact_verify_ids( + draft_block_ids=draft_block_ids, + draft_tokens=draft_tokens, + layout=layout, + device=device, + ) + + return RaggedVerifyWindow( + positions=positions, + verify_cache_loc=verify_cache_loc, + verify_ids=verify_ids, + ) + + +@triton.jit +def _ragged_finalize_kernel( + req_ptr, + within_ptr, + prefix_ptr, + cache_ptr, + pos_out_ptr, + cache_out_ptr, + bs, + n, + real_len, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + req = tl.load(req_ptr + offs, mask=mask, other=0) + within = tl.load(within_ptr + offs, mask=mask, other=0) + valid = req < bs + safe_req = tl.minimum(req, bs - 1) + prefix = tl.load(prefix_ptr + safe_req, mask=mask, other=0) + pos = tl.where(valid, prefix + within, 0) + lmask = mask & (offs < real_len) + cl = tl.load(cache_ptr + offs, mask=lmask, other=0) + cl = tl.where(valid, cl, 0) + tl.store(pos_out_ptr + offs, pos, mask=mask) + tl.store(cache_out_ptr + offs, cl, mask=mask) + + +def build_ragged_verify_window_triton( + *, + batch: ScheduleBatch, + layout: RaggedVerifyLayout, + draft_block_ids: torch.Tensor, + draft_tokens: torch.Tensor, + bs: int, + device: str, + verify_num_draft_tokens: int, + model_runner, +) -> RaggedVerifyWindow: + prefix_lens = batch.seq_lens + verify_lens = layout.verify_lens.to(device=device, dtype=torch.int32) + padded_total = layout.graph_num_tokens + + req_id, within, _valid = compact_row_index_triton( + verify_lens=verify_lens, padded_total=padded_total, device=device + ) + real_cache_loc = assign_extend_cache_locs_func( + req_pool_indices=batch.req_pool_indices, + req_to_token=model_runner.req_to_token_pool.req_to_token, + start_offset=prefix_lens, + end_offset=prefix_lens + verify_lens.to(prefix_lens.dtype), + batch_size=bs, + draft_token_num=verify_num_draft_tokens, + device=device, + ) + prefix_i64 = prefix_lens.to(device=device, dtype=torch.int64).contiguous() + positions = torch.empty(padded_total, dtype=torch.int64, device=device) + verify_cache_loc = torch.empty( + padded_total, dtype=real_cache_loc.dtype, device=device + ) + BLOCK = 256 + grid = (triton.cdiv(padded_total, BLOCK),) + _ragged_finalize_kernel[grid]( + req_id, + within, + prefix_i64, + real_cache_loc, + positions, + verify_cache_loc, + bs, + padded_total, + real_cache_loc.shape[0], + BLOCK=BLOCK, + ) + + verify_ids = compact_verify_ids_triton( + draft_block_ids=draft_block_ids, + draft_tokens=draft_tokens, + layout=layout, + device=device, + ) + return RaggedVerifyWindow( + positions=positions, + verify_cache_loc=verify_cache_loc, + verify_ids=verify_ids, + ) + + +_SEARCH_NBITS = 11 + + +class CompactRowIndex: + @classmethod + def execute( + cls, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + verify_lens: torch.Tensor, + padded_total: int, + device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return compact_row_index( + verify_lens=verify_lens, + padded_total=padded_total, + device=device, + ) + + @classmethod + def triton( + cls, + *, + verify_lens: torch.Tensor, + padded_total: int, + device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return compact_row_index_triton( + verify_lens=verify_lens, + padded_total=padded_total, + device=device, + ) + + +class CompactVerifyIds: + @classmethod + def execute(cls, *args, **kwargs) -> torch.Tensor: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + draft_block_ids: torch.Tensor, + draft_tokens: torch.Tensor, + layout: RaggedVerifyLayout, + device: str, + ) -> torch.Tensor: + return compact_verify_ids( + draft_block_ids=draft_block_ids, + draft_tokens=draft_tokens, + layout=layout, + device=device, + ) + + @classmethod + def triton( + cls, + *, + draft_block_ids: torch.Tensor, + draft_tokens: torch.Tensor, + layout: RaggedVerifyLayout, + device: str, + ) -> torch.Tensor: + return compact_verify_ids_triton( + draft_block_ids=draft_block_ids, + draft_tokens=draft_tokens, + layout=layout, + device=device, + ) + + +def compact_verify_ids( + *, + draft_block_ids: torch.Tensor, + draft_tokens: torch.Tensor, + layout: RaggedVerifyLayout, + device: str, +) -> torch.Tensor: + req_id, within, valid = compact_row_index( + verify_lens=layout.verify_lens, + padded_total=layout.graph_num_tokens, + device=device, + ) + bs = layout.verify_lens.shape[0] + safe_req = req_id.clamp(max=bs - 1) + anchors = draft_block_ids[:, 0] + drafts = draft_tokens[safe_req, (within - 1).clamp_min(0)] + verify_ids = torch.where(within == 0, anchors[safe_req], drafts) + verify_ids = torch.where(valid, verify_ids, torch.zeros_like(verify_ids)) + return verify_ids.to(torch.int64) + + +def compact_row_index( + *, + verify_lens: torch.Tensor, + padded_total: int, + device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + verify_lens = verify_lens.to(device=device, dtype=torch.int64) + bs = int(verify_lens.numel()) + incl = torch.cumsum(verify_lens, dim=0) + start = incl - verify_lens + real_total = incl[-1] + row = torch.arange(padded_total, device=device, dtype=torch.int64) + valid = row < real_total + req_id = torch.searchsorted(incl, row, right=True) + req_id = torch.where(valid, req_id, torch.full_like(req_id, bs)) + within = torch.where( + valid, row - start[req_id.clamp(max=bs - 1)], torch.zeros_like(row) + ) + return req_id, within, valid + + +@triton.jit +def _compact_row_index_kernel( + incl_ptr, + req_out_ptr, + within_out_ptr, + valid_out_ptr, + bs, + n, + BLOCK: tl.constexpr, + NBITS: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + row = offs.to(tl.int64) + real_total = tl.load(incl_ptr + (bs - 1)) + lo = tl.zeros([BLOCK], dtype=tl.int32) + hi = tl.full([BLOCK], bs, dtype=tl.int32) + for _ in range(NBITS): + mid = (lo + hi) // 2 + active = lo < hi + val = tl.load(incl_ptr + tl.minimum(mid, bs - 1), mask=mask, other=0) + go_right = val <= row + lo = tl.where(active & go_right, mid + 1, lo) + hi = tl.where(active & (~go_right), mid, hi) + req = lo + gidx = tl.maximum(req - 1, 0) + start = tl.load(incl_ptr + gidx, mask=mask, other=0) + start = tl.where(req > 0, start, 0) + valid = row < real_total + within = tl.where(valid, row - start, 0) + req_final = tl.where(valid, req.to(tl.int64), bs) + tl.store(req_out_ptr + offs, req_final, mask=mask) + tl.store(within_out_ptr + offs, within, mask=mask) + tl.store(valid_out_ptr + offs, valid, mask=mask) + + +def compact_row_index_triton( + *, + verify_lens: torch.Tensor, + padded_total: int, + device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + verify_lens = verify_lens.to(device=device, dtype=torch.int64).contiguous() + bs = verify_lens.shape[0] + incl = torch.cumsum(verify_lens, dim=0).contiguous() + req = torch.empty(padded_total, dtype=torch.int64, device=device) + within = torch.empty(padded_total, dtype=torch.int64, device=device) + valid = torch.empty(padded_total, dtype=torch.bool, device=device) + BLOCK = 256 + grid = (triton.cdiv(padded_total, BLOCK),) + _compact_row_index_kernel[grid]( + incl, req, within, valid, bs, padded_total, BLOCK=BLOCK, NBITS=_SEARCH_NBITS + ) + return req, within, valid + + +@triton.jit +def _compact_verify_ids_gather_kernel( + req_ptr, + within_ptr, + draft_block_ids_ptr, + draft_tokens_ptr, + out_ptr, + bs, + gamma, + n, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + req = tl.load(req_ptr + offs, mask=mask, other=0) + within = tl.load(within_ptr + offs, mask=mask, other=0) + valid = req < bs + safe_req = tl.minimum(req, bs - 1) + anchor = tl.load(draft_block_ids_ptr + safe_req * gamma, mask=mask, other=0) + wcol = tl.maximum(within - 1, 0) + draft = tl.load(draft_tokens_ptr + safe_req * gamma + wcol, mask=mask, other=0) + v = tl.where(within == 0, anchor, draft) + v = tl.where(valid, v, 0) + tl.store(out_ptr + offs, v.to(tl.int64), mask=mask) + + +def compact_verify_ids_triton( + *, + draft_block_ids: torch.Tensor, + draft_tokens: torch.Tensor, + layout: RaggedVerifyLayout, + device: str, +) -> torch.Tensor: + req, within, _valid = compact_row_index_triton( + verify_lens=layout.verify_lens, + padded_total=layout.graph_num_tokens, + device=device, + ) + bs = layout.verify_lens.shape[0] + gamma = draft_tokens.shape[1] + draft_block_ids = draft_block_ids.to(device=device, dtype=torch.int64).contiguous() + draft_tokens = draft_tokens.to(device=device, dtype=torch.int64).contiguous() + n = layout.graph_num_tokens + out = torch.empty(n, dtype=torch.int64, device=device) + BLOCK = 256 + grid = (triton.cdiv(n, BLOCK),) + _compact_verify_ids_gather_kernel[grid]( + req, within, draft_block_ids, draft_tokens, out, bs, gamma, n, BLOCK=BLOCK + ) + return out + + +class ScatterCompactToStrided: + @classmethod + def execute(cls, *args, **kwargs) -> torch.Tensor: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + compact: torch.Tensor, + layout: RaggedVerifyLayout, + fill_value: float, + verify_num_draft_tokens: int, + ) -> torch.Tensor: + return scatter_compact_to_strided( + compact=compact, + layout=layout, + fill_value=fill_value, + verify_num_draft_tokens=verify_num_draft_tokens, + ) + + @classmethod + def triton( + cls, + *, + compact: torch.Tensor, + layout: RaggedVerifyLayout, + fill_value: float, + verify_num_draft_tokens: int, + ) -> torch.Tensor: + return scatter_compact_to_strided_triton( + compact=compact, + layout=layout, + fill_value=fill_value, + verify_num_draft_tokens=verify_num_draft_tokens, + ) + + +def scatter_compact_to_strided( + *, + compact: torch.Tensor, + layout: RaggedVerifyLayout, + fill_value: float, + verify_num_draft_tokens: int, +) -> torch.Tensor: + stride = verify_num_draft_tokens + bs = layout.verify_lens.shape[0] + dim = compact.shape[1] + device = compact.device + compact = compact[: layout.graph_num_tokens] + strided = torch.full( + (bs * stride + 1, dim), fill_value, dtype=compact.dtype, device=device + ) + req_id, within, valid = compact_row_index( + verify_lens=layout.verify_lens, + padded_total=layout.graph_num_tokens, + device=device, + ) + sink = bs * stride + strided_pos = torch.where( + valid, + req_id.clamp(max=bs - 1) * stride + within, + torch.full_like(within, sink), + ) + strided.index_copy_(0, strided_pos, compact) + return strided[: bs * stride] + + +@triton.jit +def _scatter_compact_to_strided_kernel( + compact_ptr, + verify_lens_ptr, + start_ptr, + out_ptr, + stride, + dim, + fill_value, + BLOCK_D: tl.constexpr, +): + o = tl.program_id(0).to(tl.int64) + dblk = tl.program_id(1) + i = o // stride + w = o % stride + vl_i = tl.load(verify_lens_ptr + i) + start_i = tl.load(start_ptr + i) + d = dblk * BLOCK_D + tl.arange(0, BLOCK_D) + dmask = d < dim + in_range = w < vl_i + src = tl.where(in_range, start_i + w, 0) + val = tl.load(compact_ptr + src * dim + d, mask=dmask & in_range, other=0) + val = tl.where(in_range, val, fill_value) + tl.store(out_ptr + o * dim + d, val, mask=dmask) + + +def scatter_compact_to_strided_into( + *, + compact: torch.Tensor, + verify_lens: torch.Tensor, + out: torch.Tensor, + stride: int, + fill_value: float, +) -> torch.Tensor: + dim = compact.shape[1] + fill_value = float(fill_value) if out.dtype.is_floating_point else int(fill_value) + compact = compact.contiguous() + verify_lens = verify_lens.to(dtype=torch.int64).contiguous() + start = (torch.cumsum(verify_lens, dim=0) - verify_lens).contiguous() + n_out = out.shape[0] + BLOCK_D = 1024 + grid = (n_out, triton.cdiv(dim, BLOCK_D)) + _scatter_compact_to_strided_kernel[grid]( + compact, + verify_lens, + start, + out, + stride, + dim, + fill_value, + BLOCK_D=BLOCK_D, + ) + return out + + +def scatter_compact_to_strided_triton( + *, + compact: torch.Tensor, + layout: RaggedVerifyLayout, + fill_value: float, + verify_num_draft_tokens: int, +) -> torch.Tensor: + stride = verify_num_draft_tokens + bs = layout.verify_lens.shape[0] + dim = compact.shape[1] + device = compact.device + out = torch.empty((bs * stride, dim), dtype=compact.dtype, device=device) + return scatter_compact_to_strided_into( + compact=compact, + verify_lens=layout.verify_lens.to(device=device), + out=out, + stride=stride, + fill_value=fill_value, + ) + + +class CommitInjectLayoutResult(msgspec.Struct): + swa_loc: torch.Tensor + positions: torch.Tensor + + +class BuildCommitInjectLayout: + @classmethod + def execute(cls, *args, **kwargs) -> CommitInjectLayoutResult: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + prefix_lens: torch.Tensor, + block_pos_offsets: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + commit_lens: torch.Tensor, + stride: int, + ) -> CommitInjectLayoutResult: + return build_commit_inject_layout( + req_pool_indices=req_pool_indices, + req_to_token=req_to_token, + prefix_lens=prefix_lens, + block_pos_offsets=block_pos_offsets, + full_to_swa_mapping=full_to_swa_mapping, + commit_lens=commit_lens, + stride=stride, + ) + + @classmethod + def triton( + cls, + *, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + prefix_lens: torch.Tensor, + block_pos_offsets: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + commit_lens: torch.Tensor, + stride: int, + ) -> CommitInjectLayoutResult: + return build_commit_inject_layout_triton( + req_pool_indices=req_pool_indices, + req_to_token=req_to_token, + prefix_lens=prefix_lens, + block_pos_offsets=block_pos_offsets, + full_to_swa_mapping=full_to_swa_mapping, + commit_lens=commit_lens, + stride=stride, + ) + + +def build_commit_inject_layout( + *, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + prefix_lens: torch.Tensor, + block_pos_offsets: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + commit_lens: torch.Tensor, + stride: int, +) -> CommitInjectLayoutResult: + from sglang.kernels.ops.speculative.cache_locs import ( + assign_extend_cache_locs_func, + ) + + bs = req_pool_indices.shape[0] + device = req_pool_indices.device + + positions_2d = prefix_lens.unsqueeze(1) + block_pos_offsets[:stride] + positions = positions_2d.reshape(-1).to(dtype=torch.int64) + + cache_loc = assign_extend_cache_locs_func( + req_pool_indices=req_pool_indices, + req_to_token=req_to_token, + start_offset=prefix_lens, + end_offset=prefix_lens + stride, + batch_size=bs, + draft_token_num=stride, + device=device, + ).to(dtype=torch.int64) + swa_loc = full_to_swa_mapping[cache_loc].to(torch.int32) + + col = torch.arange(stride, device=device).view(1, -1) + committed = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1) + swa_loc = torch.where(committed, swa_loc, torch.full_like(swa_loc, -1)) + + return CommitInjectLayoutResult(swa_loc=swa_loc, positions=positions) + + +@triton.jit +def _commit_inject_layout_kernel( + req_pool_ptr, + req_to_token_ptr, + prefix_lens_ptr, + block_pos_offsets_ptr, + full_to_swa_ptr, + commit_lens_ptr, + swa_loc_ptr, + positions_ptr, + rt_stride, + stride, + n, + BLOCK: tl.constexpr, +): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + r = offs // stride + c = offs % stride + + prefix = tl.load(prefix_lens_ptr + r, mask=mask, other=0).to(tl.int64) + pos_off = tl.load(block_pos_offsets_ptr + c, mask=mask, other=0).to(tl.int64) + rp = tl.load(req_pool_ptr + r, mask=mask, other=0).to(tl.int64) + full_loc = tl.load( + req_to_token_ptr + rp * rt_stride + prefix + pos_off, mask=mask, other=0 + ).to(tl.int64) + swa = tl.load(full_to_swa_ptr + full_loc, mask=mask, other=-1).to(tl.int32) + + commit_len = tl.load(commit_lens_ptr + r, mask=mask, other=0).to(tl.int64) + swa = tl.where(c.to(tl.int64) < commit_len, swa, -1) + + tl.store(swa_loc_ptr + offs, swa, mask=mask) + tl.store(positions_ptr + offs, prefix + pos_off, mask=mask) + + +def build_commit_inject_layout_triton( + *, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + prefix_lens: torch.Tensor, + block_pos_offsets: torch.Tensor, + full_to_swa_mapping: torch.Tensor, + commit_lens: torch.Tensor, + stride: int, +) -> CommitInjectLayoutResult: + bs = req_pool_indices.shape[0] + n = bs * stride + device = req_pool_indices.device + + swa_loc = torch.empty(n, dtype=torch.int32, device=device) + positions = torch.empty(n, dtype=torch.int64, device=device) + BLOCK = 256 + _commit_inject_layout_kernel[(triton.cdiv(n, BLOCK),)]( + req_pool_indices, + req_to_token, + prefix_lens, + block_pos_offsets, + full_to_swa_mapping, + commit_lens, + swa_loc, + positions, + req_to_token.stride(0), + stride, + n, + BLOCK=BLOCK, + ) + return CommitInjectLayoutResult(swa_loc=swa_loc, positions=positions) + + +class BuildOutTokens: + @classmethod + def execute(cls, *args, **kwargs) -> torch.Tensor: + if inputs_on_cuda(*args, **kwargs): + return cls.triton(*args, **kwargs) + return cls.torch(*args, **kwargs) + + @classmethod + def torch( + cls, + *, + draft_tokens: torch.Tensor, + correct_len: torch.Tensor, + bonus: torch.Tensor, + verify_num_draft_tokens: int, + gamma: int, + ) -> torch.Tensor: + return build_out_tokens( + draft_tokens=draft_tokens, + correct_len=correct_len, + bonus=bonus, + verify_num_draft_tokens=verify_num_draft_tokens, + gamma=gamma, + ) + + @classmethod + def triton( + cls, + *, + draft_tokens: torch.Tensor, + correct_len: torch.Tensor, + bonus: torch.Tensor, + verify_num_draft_tokens: int, + gamma: int, + ) -> torch.Tensor: + return build_out_tokens_triton( + draft_tokens=draft_tokens, + correct_len=correct_len, + bonus=bonus, + verify_num_draft_tokens=verify_num_draft_tokens, + gamma=gamma, + ) + + +def build_out_tokens( + *, + draft_tokens: torch.Tensor, + correct_len: torch.Tensor, + bonus: torch.Tensor, + verify_num_draft_tokens: int, + gamma: int, +) -> torch.Tensor: + bs = draft_tokens.shape[0] + out_tokens = torch.empty( + (bs, verify_num_draft_tokens), + dtype=torch.int64, + device=draft_tokens.device, + ) + out_tokens[:, :gamma].copy_(draft_tokens) + out_tokens[:, gamma].fill_(0) + out_tokens.scatter_(1, correct_len.to(torch.int64)[:, None], bonus[:, None]) + return out_tokens + + +@triton.jit +def _build_out_tokens_kernel( + draft_tokens_ptr, + correct_len_ptr, + bonus_ptr, + out_ptr, + gamma, + T, + n_out, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_out + b = offs // T + k = offs % T + cl = tl.load(correct_len_ptr + b, mask=mask, other=0).to(tl.int32) + bonus = tl.load(bonus_ptr + b, mask=mask, other=0) + draft_mask = mask & (k < gamma) + draft = tl.load(draft_tokens_ptr + b * gamma + k, mask=draft_mask, other=0) + val = tl.where(k == cl, bonus, tl.where(k < gamma, draft, 0)) + tl.store(out_ptr + offs, val.to(tl.int64), mask=mask) + + +def build_out_tokens_triton( + *, + draft_tokens: torch.Tensor, + correct_len: torch.Tensor, + bonus: torch.Tensor, + verify_num_draft_tokens: int, + gamma: int, +) -> torch.Tensor: + bs = draft_tokens.shape[0] + T = verify_num_draft_tokens + device = draft_tokens.device + draft_tokens = draft_tokens.to(torch.int64).contiguous() + correct_len_i = correct_len.to(torch.int64).contiguous() + bonus_i = bonus.to(torch.int64).contiguous() + out = torch.empty((bs, T), dtype=torch.int64, device=device) + n_out = bs * T + BLOCK = 256 + grid = (triton.cdiv(n_out, BLOCK),) + _build_out_tokens_kernel[grid]( + draft_tokens, correct_len_i, bonus_i, out, gamma, T, n_out, BLOCK=BLOCK + ) + return out diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index f43e2cfc7..4e193bb97 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -393,6 +393,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): top_ks=torch.full((num_seqs,), -1, dtype=torch.int32), min_ps=torch.zeros((num_seqs,), dtype=torch.float), is_all_greedy=False, + is_any_greedy=False, need_top_p_sampling=False, need_top_k_sampling=False, need_min_p_sampling=False, diff --git a/python/sglang/srt/speculative/ragged_verify.py b/python/sglang/srt/speculative/ragged_verify.py new file mode 100644 index 000000000..1dcb618ec --- /dev/null +++ b/python/sglang/srt/speculative/ragged_verify.py @@ -0,0 +1,304 @@ +from __future__ import annotations + +import bisect +from enum import Enum +from typing import List, Optional, Sequence, Tuple + +import msgspec +import torch + +from sglang.srt.environ import envs + + +class RaggedVerifyMode(str, Enum): + STATIC = "static" + CAP_ACCEPT = "cap-accept" + COMPACT = "compact" + + +def read_ragged_verify_mode() -> RaggedVerifyMode: + value = envs.SGLANG_RAGGED_VERIFY_MODE.get() + for mode in RaggedVerifyMode: + if value == mode.value: + return mode + raise ValueError( + f"invalid SGLANG_RAGGED_VERIFY_MODE={value!r}; expected one of " + f"{', '.join(repr(m.value) for m in RaggedVerifyMode)}" + ) + + +def ragged_verify_compact_enabled() -> bool: + return read_ragged_verify_mode() == RaggedVerifyMode.COMPACT + + +def round_up_grid(total: int, grid: Sequence[int]) -> int: + if not grid: + raise ValueError("round_up_grid requires a non-empty grid") + if total > grid[-1]: + raise ValueError( + f"total {total} exceeds max grid tier {grid[-1]}; " + "the caller must reject this batch before selecting a graph tier" + ) + index = bisect.bisect_left(grid, total) + return grid[index] + + +class RaggedVerifyLayout(msgspec.Struct, frozen=True): + verify_lens: torch.Tensor + graph_num_tokens: int + extend_start_loc: torch.Tensor + qo_indptr_device: torch.Tensor + verify_lens_cpu: Optional[list[int]] = None + total_verify_tokens: Optional[int] = None + qo_indptr_host: Optional[torch.Tensor] = None + kv_indptr_host: Optional[torch.Tensor] = None + kv_lens_host: Optional[torch.Tensor] = None + max_q_len: Optional[int] = None + max_kv_len: Optional[int] = None + + def __post_init__(self) -> None: + if self.verify_lens_cpu is None: + return + if not self.verify_lens_cpu: + raise ValueError("RaggedVerifyLayout requires at least one request") + if min(self.verify_lens_cpu) < 1: + raise ValueError( + f"every request must verify the anchor (verify_len >= 1), got " + f"{self.verify_lens_cpu}" + ) + if self.total_verify_tokens != sum(self.verify_lens_cpu): + raise ValueError( + f"total_verify_tokens {self.total_verify_tokens} != " + f"sum(verify_lens_cpu) {sum(self.verify_lens_cpu)}" + ) + if not (self.total_verify_tokens <= self.graph_num_tokens): + raise ValueError( + f"total_verify_tokens {self.total_verify_tokens} exceeds " + f"graph_num_tokens {self.graph_num_tokens}" + ) + + @property + def bs(self) -> int: + return int(self.verify_lens.shape[0]) + + @classmethod + def _assemble_device( + cls, + *, + verify_lens: torch.Tensor, + graph_num_tokens: int, + verify_lens_cpu: Optional[list[int]] = None, + total_verify_tokens: Optional[int] = None, + ) -> RaggedVerifyLayout: + from sglang.srt.speculative.ragged_verify_kernels import ( + BuildQoIndptr, + ) + + verify_lens = verify_lens.to(torch.int32) + indptr = BuildQoIndptr.execute(verify_lens=verify_lens) + return cls( + verify_lens=verify_lens, + graph_num_tokens=graph_num_tokens, + extend_start_loc=indptr.extend_start_loc, + qo_indptr_device=indptr.qo_indptr, + verify_lens_cpu=verify_lens_cpu, + total_verify_tokens=total_verify_tokens, + ) + + @classmethod + def _assemble( + cls, + *, + verify_lens_cpu: list[int], + total_verify_tokens: int, + graph_num_tokens: int, + device: torch.device, + ) -> RaggedVerifyLayout: + verify_lens = torch.tensor(verify_lens_cpu, dtype=torch.int32, device=device) + return cls._assemble_device( + verify_lens=verify_lens, + graph_num_tokens=graph_num_tokens, + verify_lens_cpu=verify_lens_cpu, + total_verify_tokens=total_verify_tokens, + ) + + @classmethod + def from_verify_lens_device( + cls, + *, + verify_lens: torch.Tensor, + graph_num_tokens: int, + ) -> RaggedVerifyLayout: + return cls._assemble_device( + verify_lens=verify_lens, graph_num_tokens=graph_num_tokens + ) + + @classmethod + def from_verify_lens( + cls, + *, + verify_lens_cpu: Sequence[int], + device: torch.device, + grid: Sequence[int], + graph_num_tokens_floor: int = 0, + ) -> RaggedVerifyLayout: + verify_lens_list = [int(v) for v in verify_lens_cpu] + total_verify_tokens = sum(verify_lens_list) + bucket_input = max(total_verify_tokens, graph_num_tokens_floor) + graph_num_tokens = round_up_grid(total=bucket_input, grid=grid) + + return cls._assemble( + verify_lens_cpu=verify_lens_list, + total_verify_tokens=total_verify_tokens, + graph_num_tokens=graph_num_tokens, + device=device, + ) + + def padded_to_bucket(self, *, padded_bs: int) -> RaggedVerifyLayout: + from sglang.srt.speculative.ragged_verify_kernels import ( + PaddedToBucket, + ) + + padded = PaddedToBucket.execute( + verify_lens=self.verify_lens, + graph_num_tokens=self.graph_num_tokens, + bs=self.bs, + padded_bs=padded_bs, + ) + + return RaggedVerifyLayout._assemble_device( + verify_lens=padded, + graph_num_tokens=self.graph_num_tokens, + total_verify_tokens=self.graph_num_tokens, + ) + + +def build_capture_verify_lens( + *, + num_tokens: int, + num_slots: int, + num_draft_tokens: int, +) -> list[int]: + if num_slots < 1 or num_tokens < num_slots: + raise ValueError( + f"capture layout needs 1 <= num_slots <= num_tokens, got " + f"num_slots={num_slots}, num_tokens={num_tokens}" + ) + if num_tokens > num_slots * num_draft_tokens: + raise ValueError( + f"capture layout cannot pack num_tokens={num_tokens} into " + f"{num_slots} rows of at most {num_draft_tokens} tokens" + ) + base = num_tokens // num_slots + rem = num_tokens - base * num_slots + return [base + 1] * rem + [base] * (num_slots - rem) + + +def resolve_ragged_verify_layout(forward_batch) -> Optional[RaggedVerifyLayout]: + """Layout riding the batch's spec input, or None. Tolerates the runner's + ad-hoc replay batch views, which may not carry spec_info at all.""" + spec_info = getattr(forward_batch, "spec_info", None) + if spec_info is None: + return None + return spec_info.ragged_verify_layout + + +class RaggedTargetVerifyGeometry(msgspec.Struct): + cache_seqlens_int32: torch.Tensor + cu_seqlens_q: torch.Tensor + cu_seqlens_k: torch.Tensor + max_seq_len_q: Optional[int] + + +def build_ragged_target_verify_geometry( + *, + seq_lens: torch.Tensor, + layout: RaggedVerifyLayout, +) -> RaggedTargetVerifyGeometry: + cache_seqlens_int32 = (seq_lens + layout.verify_lens).to(torch.int32) + cu_seqlens_q = layout.qo_indptr_device.to(torch.int32) + cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(cache_seqlens_int32, dim=0, dtype=torch.int32), (1, 0) + ) + max_seq_len_q = ( + max(layout.verify_lens_cpu) if layout.verify_lens_cpu is not None else None + ) + return RaggedTargetVerifyGeometry( + cache_seqlens_int32=cache_seqlens_int32, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seq_len_q=max_seq_len_q, + ) + + +def compute_target_verify_graph_key( + *, + bs: int, + num_draft_tokens: int, + ragged_layout: Optional[RaggedVerifyLayout], +) -> Tuple[int, int]: + num_tokens_full_block = num_draft_tokens * bs + if ragged_layout is None: + return bs, num_tokens_full_block + graph_num_tokens = ragged_layout.graph_num_tokens + assert graph_num_tokens <= num_tokens_full_block, ( + f"ragged verify graph_num_tokens={graph_num_tokens} exceeds full block " + f"num_draft*bs={num_tokens_full_block}" + ) + total_verify_tokens = ragged_layout.total_verify_tokens + if total_verify_tokens is not None: + assert total_verify_tokens <= graph_num_tokens, ( + f"ragged verify total_verify_tokens={total_verify_tokens} exceeds the " + f"round-up bucket graph_num_tokens={graph_num_tokens}" + ) + return graph_num_tokens, graph_num_tokens + + +class VerifyExtendLengths(msgspec.Struct, frozen=True): + seq_lens_extended: torch.Tensor + seq_lens_cpu_extended: List[int] + extend_seq_lens_cpu: List[int] + num_tokens: int + extend_start_loc: Optional[torch.Tensor] + + +def compute_uniform_extend_lengths( + *, + seq_lens: torch.Tensor, + seq_lens_cpu: List[int], + extend_len: int, +) -> VerifyExtendLengths: + batch_size = len(seq_lens_cpu) + seq_lens_extended = seq_lens + extend_len + seq_lens_cpu_extended = [x + extend_len for x in seq_lens_cpu] + extend_seq_lens_cpu = [extend_len] * batch_size + num_tokens = extend_len * batch_size + return VerifyExtendLengths( + seq_lens_extended=seq_lens_extended, + seq_lens_cpu_extended=seq_lens_cpu_extended, + extend_seq_lens_cpu=extend_seq_lens_cpu, + num_tokens=num_tokens, + extend_start_loc=None, + ) + + +def compute_ragged_extend_lengths( + *, + seq_lens: torch.Tensor, + seq_lens_cpu: List[int], + ragged_layout: RaggedVerifyLayout, +) -> VerifyExtendLengths: + extend_seq_lens_cpu = list(ragged_layout.verify_lens_cpu) + seq_lens_extended = seq_lens + ragged_layout.verify_lens + seq_lens_cpu_extended = [ + raw + length for raw, length in zip(seq_lens_cpu, extend_seq_lens_cpu) + ] + num_tokens = ragged_layout.total_verify_tokens + extend_start_loc = ragged_layout.extend_start_loc + return VerifyExtendLengths( + seq_lens_extended=seq_lens_extended, + seq_lens_cpu_extended=seq_lens_cpu_extended, + extend_seq_lens_cpu=extend_seq_lens_cpu, + num_tokens=num_tokens, + extend_start_loc=extend_start_loc, + ) diff --git a/python/sglang/srt/speculative/ragged_verify_kernels.py b/python/sglang/srt/speculative/ragged_verify_kernels.py new file mode 100644 index 000000000..68f8ad898 --- /dev/null +++ b/python/sglang/srt/speculative/ragged_verify_kernels.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +import msgspec +import torch +import triton +import triton.language as tl + + +class PaddedToBucket: + @classmethod + def execute( + cls, + *, + verify_lens: torch.Tensor, + graph_num_tokens: int, + bs: int, + padded_bs: int, + ) -> torch.Tensor: + impl = cls.triton if verify_lens.is_cuda else cls.torch + return impl( + verify_lens=verify_lens, + graph_num_tokens=graph_num_tokens, + bs=bs, + padded_bs=padded_bs, + ) + + @classmethod + def torch( + cls, + *, + verify_lens: torch.Tensor, + graph_num_tokens: int, + bs: int, + padded_bs: int, + ) -> torch.Tensor: + return pad_verify_lens_to_bucket( + verify_lens=verify_lens, + graph_num_tokens=graph_num_tokens, + bs=bs, + padded_bs=padded_bs, + ) + + @classmethod + def triton( + cls, + *, + verify_lens: torch.Tensor, + graph_num_tokens: int, + bs: int, + padded_bs: int, + ) -> torch.Tensor: + return pad_verify_lens_to_bucket_triton( + verify_lens=verify_lens, + graph_num_tokens=graph_num_tokens, + bs=bs, + padded_bs=padded_bs, + ) + + +def pad_verify_lens_to_bucket( + *, + verify_lens: torch.Tensor, + graph_num_tokens: int, + bs: int, + padded_bs: int, +) -> torch.Tensor: + assert padded_bs >= bs, ( + f"padded_bs {padded_bs} < bs {bs}: the captured tier cannot hold this " + "batch's requests" + ) + device = verify_lens.device + num_pad_reqs = padded_bs - bs + padded = verify_lens.to(torch.int32) + leftover = graph_num_tokens - padded.to(torch.int64).sum() + if num_pad_reqs > 0: + base = leftover // num_pad_reqs + rem = leftover - base * num_pad_reqs + pad_block = base + ( + torch.arange(num_pad_reqs, device=device, dtype=torch.int64) < rem + ) + padded = torch.cat([padded, pad_block.to(torch.int32)]) + else: + padded = padded.clone() + padded[-1] = (padded[-1].to(torch.int64) + leftover).to(torch.int32) + return padded + + +@triton.jit +def _padded_to_bucket_kernel( + verify_lens_ptr, + out_ptr, + bs, + padded_bs, + graph_num_tokens, + BLOCK: tl.constexpr, +): + idx = tl.arange(0, BLOCK) + valid = idx < padded_bs + is_real = idx < bs + vl = tl.load(verify_lens_ptr + idx, mask=is_real, other=0).to(tl.int64) + leftover = graph_num_tokens - tl.sum(vl) + num_pad = padded_bs - bs + num_pad_safe = tl.maximum(num_pad, 1) + base = leftover // num_pad_safe + rem = leftover - base * num_pad_safe + pad_len = base + tl.where((idx - bs) < rem, 1, 0) + final = tl.where(is_real, vl, pad_len) + final = final + tl.where((num_pad == 0) & (idx == bs - 1), leftover, 0) + tl.store(out_ptr + idx, final.to(tl.int32), mask=valid) + + +def pad_verify_lens_to_bucket_triton( + *, + verify_lens: torch.Tensor, + graph_num_tokens: int, + bs: int, + padded_bs: int, +) -> torch.Tensor: + assert padded_bs >= bs, ( + f"padded_bs {padded_bs} < bs {bs}: the captured tier cannot hold this " + "batch's requests" + ) + device = verify_lens.device + verify_lens = verify_lens.to(torch.int32).contiguous() + out = torch.empty(padded_bs, dtype=torch.int32, device=device) + BLOCK = triton.next_power_of_2(max(padded_bs, 1)) + _padded_to_bucket_kernel[(1,)]( + verify_lens, + out, + bs, + padded_bs, + graph_num_tokens, + BLOCK=BLOCK, + ) + return out + + +class QoIndptrResult(msgspec.Struct): + qo_indptr: torch.Tensor + extend_start_loc: torch.Tensor + + +class BuildQoIndptr: + @classmethod + def execute(cls, *, verify_lens: torch.Tensor) -> QoIndptrResult: + impl = cls.triton if verify_lens.is_cuda else cls.torch + return impl(verify_lens=verify_lens) + + @classmethod + def torch(cls, *, verify_lens: torch.Tensor) -> QoIndptrResult: + return build_qo_indptr(verify_lens=verify_lens) + + @classmethod + def triton(cls, *, verify_lens: torch.Tensor) -> QoIndptrResult: + return build_qo_indptr_triton(verify_lens=verify_lens) + + +def build_qo_indptr(*, verify_lens: torch.Tensor) -> QoIndptrResult: + verify_lens = verify_lens.to(torch.int32) + cumsum = torch.cumsum(verify_lens, dim=0).to(torch.int32) + zero = torch.zeros(1, dtype=torch.int32, device=verify_lens.device) + qo_indptr = torch.cat([zero, cumsum]) + extend_start_loc = qo_indptr[:-1].clone() + return QoIndptrResult(qo_indptr=qo_indptr, extend_start_loc=extend_start_loc) + + +@triton.jit +def _qo_indptr_kernel( + verify_lens_ptr, + qo_indptr_ptr, + extend_start_loc_ptr, + bs, + BLOCK: tl.constexpr, +): + idx = tl.arange(0, BLOCK) + valid = idx < bs + vl = tl.load(verify_lens_ptr + idx, mask=valid, other=0).to(tl.int32) + incl = tl.cumsum(vl, axis=0) + excl = incl - vl + tl.store(qo_indptr_ptr, 0) + tl.store(qo_indptr_ptr + 1 + idx, incl, mask=valid) + tl.store(extend_start_loc_ptr + idx, excl, mask=valid) + + +def build_qo_indptr_triton(*, verify_lens: torch.Tensor) -> QoIndptrResult: + bs = verify_lens.shape[0] + device = verify_lens.device + verify_lens = verify_lens.contiguous() + qo_indptr = torch.empty(bs + 1, dtype=torch.int32, device=device) + extend_start_loc = torch.empty(bs, dtype=torch.int32, device=device) + BLOCK = triton.next_power_of_2(max(bs, 1)) + _qo_indptr_kernel[(1,)]( + verify_lens, + qo_indptr, + extend_start_loc, + bs, + BLOCK=BLOCK, + ) + return QoIndptrResult(qo_indptr=qo_indptr, extend_start_loc=extend_start_loc) diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index f891c15fa..a9867dc74 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -23,6 +23,7 @@ if TYPE_CHECKING: from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.base_spec_worker import BaseSpecWorker from sglang.srt.speculative.ngram_worker import NGRAMWorker + from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout class SpeculativeAlgorithm(Enum): @@ -33,6 +34,7 @@ class SpeculativeAlgorithm(Enum): """ DFLASH = auto() + DSPARK = auto() EAGLE = auto() EAGLE3 = auto() FROZEN_KV_MTP = auto() @@ -109,6 +111,12 @@ class SpeculativeAlgorithm(Enum): def is_dflash(self) -> bool: return self == SpeculativeAlgorithm.DFLASH + def is_dspark(self) -> bool: + return self == SpeculativeAlgorithm.DSPARK + + def is_dflash_family(self) -> bool: + return self.is_dflash() or self.is_dspark() + def is_standalone(self) -> bool: return self == SpeculativeAlgorithm.STANDALONE @@ -116,7 +124,13 @@ class SpeculativeAlgorithm(Enum): return self == SpeculativeAlgorithm.NGRAM def supports_target_verify_for_draft(self) -> bool: - return self.is_dflash() + return self.is_dflash_family() + + def supports_ragged_verify(self) -> bool: + """Whether this algorithm's verify step may carry a RaggedVerifyLayout + (per-request verify lengths); gates the token-bucket-keyed verify + graphs in the decode cuda graph runner.""" + return self.is_dspark() def has_draft_kv(self) -> bool: """Whether the draft phase writes KV chains. NGRAM does not (its tree @@ -134,10 +148,17 @@ class SpeculativeAlgorithm(Enum): device: torch.device, req_to_token_pool, needs_cpu_seq_lens: bool = True, + needs_confidence_relay: bool = False, ) -> FutureMap: from sglang.srt.managers.overlap_utils import FutureMap - return FutureMap(device, self, req_to_token_pool, needs_cpu_seq_lens) + return FutureMap( + device, + self, + req_to_token_pool, + needs_cpu_seq_lens, + needs_confidence_relay, + ) def build_disagg_draft_input( self, @@ -166,13 +187,22 @@ class SpeculativeAlgorithm(Enum): """ from sglang.srt.arg_groups.speculative_hook import ( _handle_dflash, + _handle_dspark, _handle_eagle_family, _handle_frozen_kv_mtp, _handle_ngram, ) + # Validate for every algorithm at startup: the metrics paths read the + # ragged-verify mode env and must not be where a typo'd value raises. + from sglang.srt.speculative.ragged_verify import read_ragged_verify_mode + + read_ragged_verify_mode() + if self.is_dflash(): _handle_dflash(server_args) + elif self.is_dspark(): + _handle_dspark(server_args) elif self.is_frozen_kv_mtp(): _handle_frozen_kv_mtp(server_args) elif self.is_eagle() or self.is_standalone(): @@ -188,6 +218,8 @@ class SpeculativeAlgorithm(Enum): # graph support. We can use it for target verify, or we can use it for # other cases which is not target verify but fixed length prefill. # Here, we expose this interface to allow the other use cases. + if self.is_dspark() and is_draft_worker: + return num_draft_tokens - 1 return num_draft_tokens def create_worker( @@ -204,6 +236,13 @@ class SpeculativeAlgorithm(Enum): return DFlashWorkerV2 + if self.is_dspark(): + from sglang.srt.speculative.dspark_components.dspark_worker_v2 import ( + DSparkWorkerV2, + ) + + return DSparkWorkerV2 + if self.is_frozen_kv_mtp(): # V2 worker drives both overlap and non-overlap (scheduler runs it # synchronously when overlap is disabled), same as EAGLE. @@ -254,6 +293,14 @@ class SpecInputType(IntEnum): class SpecInput(ABC): + # Per-request verify lengths for the ragged-verify graphs (see + # sglang.srt.speculative.ragged_verify); verify inputs of algorithms with + # supports_ragged_verify() override it per step. Must stay a class-level + # default, not an __init__ assignment: dataclass subclasses declare it as + # a field and run __post_init__ -> super().__init__ *after* field + # assignment, so an init-time default would clobber the passed layout. + ragged_verify_layout: Optional[RaggedVerifyLayout] = None + def __init__(self, spec_input_type: SpecInputType): self.spec_input_type = spec_input_type @@ -324,7 +371,7 @@ def create_dummy_verify_input( seq_lens_sum=None, seq_lens_cpu=None, ) - elif spec_algorithm.is_dflash(): + elif spec_algorithm.is_dflash_family(): from sglang.srt.speculative.dflash_info import DFlashVerifyInput # Dummy warmup only needs shape metadata; avoid forcing custom-mask mode. diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py index ce3f3178f..cd3b1665b 100644 --- a/python/sglang/srt/speculative/spec_registry.py +++ b/python/sglang/srt/speculative/spec_registry.py @@ -76,6 +76,12 @@ class CustomSpecAlgo: def is_dflash(self) -> bool: return False + def is_dspark(self) -> bool: + return False + + def is_dflash_family(self) -> bool: + return False + def is_standalone(self) -> bool: return False @@ -85,6 +91,9 @@ class CustomSpecAlgo: def supports_target_verify_for_draft(self) -> bool: return False + def supports_ragged_verify(self) -> bool: + return False + def has_draft_kv(self) -> bool: # Conservative default: the larger KV reserve. return True diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 437ea74b3..70a47954d 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -207,9 +207,9 @@ def spec_need_hidden_states(server_args: Optional[ServerArgs] = None) -> bool: server_args = get_server_args() # STANDALONE drafts don't consume `spec_info.hidden_states` (vanilla LLM). - # multi_layer_eagle and DFLASH don't relay hidden_states through FutureMap. + # multi_layer_eagle, DFLASH, and DSPARK don't relay hidden_states through FutureMap. # TODO(lsyin): also skip when step == 1. - if server_args.speculative_algorithm in ("STANDALONE", "DFLASH"): + if server_args.speculative_algorithm in ("STANDALONE", "DFLASH", "DSPARK"): return False return not server_args.enable_multi_layer_eagle @@ -711,7 +711,7 @@ def spec_prepare_for_decode(batch: ScheduleBatch) -> None: """eagle/ngram share a stateless free function; dflash keeps stateful prep on its draft input -- the dispatcher routes. """ - if batch.spec_algorithm.is_dflash(): + if batch.spec_algorithm.is_dflash_family(): 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/utils/async_probe.py b/python/sglang/srt/utils/async_probe.py index b511476d5..92d98f87b 100644 --- a/python/sglang/srt/utils/async_probe.py +++ b/python/sglang/srt/utils/async_probe.py @@ -76,6 +76,12 @@ def sanitize_nan_logits(logits: torch.Tensor, msg: str = ""): torch.nan_to_num_(logits, nan=-1e30, posinf=1e30, neginf=-1e30) +def maybe_assert_async(cond: torch.Tensor, msg: str = ""): + if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get(): + return + torch._assert_async(cond, msg) + + def maybe_detect_nan(tensor: Optional[torch.Tensor], msg: str = ""): """Async NaN check — no GPU-CPU sync, error surfaces at next sync point.""" if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get(): @@ -96,6 +102,19 @@ def maybe_detect_inf(tensor: Optional[torch.Tensor], msg: str = ""): torch._assert_async(~torch.any(torch.isinf(tensor)), f"Inf detected! {msg}") +def maybe_detect_in_closed_range( + tensor: Optional[torch.Tensor], low: float, high: float, msg: str = "" +): + if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get(): + return + if tensor is None or tensor.numel() == 0: + return + torch._assert_async( + ((tensor >= low) & (tensor <= high)).all(), + f"value outside [{low}, {high}]: {msg}", + ) + + def maybe_detect_oob(indices: Optional[torch.Tensor], low: int, high: int, msg: str): """Async OOB check — no GPU-CPU sync, error surfaces at next sync point. diff --git a/test/registered/core/test_basic_sanity_dspark.py b/test/registered/core/test_basic_sanity_dspark.py new file mode 100644 index 000000000..c27d9056d --- /dev/null +++ b/test/registered/core/test_basic_sanity_dspark.py @@ -0,0 +1,94 @@ +import unittest + +from sglang.srt.utils import is_sm100_supported, kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.basic_api_contract_kit import BasicAPIContractMixin +from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin +from sglang.test.kits.basic_scheduler_stress_kit import BasicSchedulerStressMixin +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.kits.fwd_occupancy_kit import FwdOccupancyMixin +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large") + +TARGET_MODEL = "Qwen/Qwen3-14B" +DRAFT_MODEL = "deepseek-ai/dspark_qwen3_14b_block7" + +# trtllm_mha prefill requires SM100 (Blackwell); use the Hopper-native pair elsewhere. +if is_sm100_supported(): + ATTENTION_BACKEND = "trtllm_mha" + DRAFT_ATTENTION_BACKEND = "fa4" +else: + ATTENTION_BACKEND = "fa3" + DRAFT_ATTENTION_BACKEND = "fa3" + + +class TestBasicSanityDSpark( + BasicAPIContractMixin, + BasicDecodeCorrectnessMixin, + BasicSchedulerStressMixin, + FwdOccupancyMixin, + GSM8KMixin, + CustomTestCase, +): + served_model_name = TARGET_MODEL + model = TARGET_MODEL + + fwd_occupancy_threshold = 60 + fwd_occupancy_max_new_tokens = 4096 + fwd_occupancy_acc_length_threshold: float = 2.0 + + gsm8k_num_questions = 200 + gsm8k_accuracy_thres = 0.80 + gsm8k_accept_length_thres = 2.0 + + attention_backend = ATTENTION_BACKEND + draft_attention_backend = DRAFT_ATTENTION_BACKEND + + process = None + + @classmethod + def setUpClass(cls): + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + TARGET_MODEL, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--attention-backend", + cls.attention_backend, + "--speculative-draft-attention-backend", + cls.draft_attention_backend, + "--speculative-algorithm", + "DSPARK", + "--speculative-draft-model-path", + DRAFT_MODEL, + "--cuda-graph-max-bs-decode", + "4", + "--mem-fraction-static", + "0.7", + "--page-size", + "1", + "--enable-metrics", + "--disable-piecewise-cuda-graph", + ], + env={ + "SGLANG_ENABLE_METRICS_DEVICE_TIMER": "1", + "SGLANG_RAGGED_VERIFY_MODE": "compact", + }, + ) + + @classmethod + def tearDownClass(cls): + if cls.process is not None: + kill_process_tree(cls.process.pid) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/ops/test_aiter_greedy_sample_amd.py b/test/registered/ops/test_aiter_greedy_sample_amd.py index 473c3b4e9..e605d1d28 100644 --- a/test/registered/ops/test_aiter_greedy_sample_amd.py +++ b/test/registered/ops/test_aiter_greedy_sample_amd.py @@ -52,6 +52,7 @@ def _make_sampling_info(batch_size, vocab_size, device="cuda"): top_ks=torch.zeros(batch_size, device=device, dtype=torch.int32), min_ps=torch.zeros(batch_size, device=device), is_all_greedy=True, + is_any_greedy=True, need_top_p_sampling=False, need_top_k_sampling=False, need_min_p_sampling=False, diff --git a/test/registered/scheduler/test_min_free_slots_delayer.py b/test/registered/scheduler/test_min_free_slots_delayer.py index 4dee7bf3b..0fdcb6153 100644 --- a/test/registered/scheduler/test_min_free_slots_delayer.py +++ b/test/registered/scheduler/test_min_free_slots_delayer.py @@ -14,17 +14,17 @@ class TestResolveMinFreeSlots(unittest.TestCase): def test_unset_non_dflash_disables(self): # Unset + not DFlash -> trigger stays disabled. - self.assertIsNone(resolve_min_free_slots(None, 512, is_dflash=False)) + self.assertIsNone(resolve_min_free_slots(None, 512, is_dflash_family=False)) def test_unset_dflash_auto_enables(self): # Unset + DFlash -> falls back to the legacy formula (full mapping). - self.assertEqual(resolve_min_free_slots(None, 512, is_dflash=True), 4) - self.assertEqual(resolve_min_free_slots(None, 8, is_dflash=True), 2) + self.assertEqual(resolve_min_free_slots(None, 512, is_dflash_family=True), 4) + self.assertEqual(resolve_min_free_slots(None, 8, is_dflash_family=True), 2) def test_unset_dflash_small_cluster_disables(self): # DFlash auto-default still respects the < 8 guard. - self.assertIsNone(resolve_min_free_slots(None, 7, is_dflash=True)) - self.assertIsNone(resolve_min_free_slots(None, 0, is_dflash=True)) + self.assertIsNone(resolve_min_free_slots(None, 7, is_dflash_family=True)) + self.assertIsNone(resolve_min_free_slots(None, 0, is_dflash_family=True)) def test_le_one_disables(self): # <= 1 can never batch, so it is a no-op. @@ -47,7 +47,7 @@ class TestResolveMinFreeSlots(unittest.TestCase): def test_user_value_overrides_dflash_default(self): # An explicit user value wins over the DFlash auto-default. - self.assertEqual(resolve_min_free_slots(3, 512, is_dflash=True), 3) + self.assertEqual(resolve_min_free_slots(3, 512, is_dflash_family=True), 3) class TestMinFreeSlotsDelayer(unittest.TestCase): diff --git a/test/registered/spec/dspark/test_dspark_block_accept_estimator.py b/test/registered/spec/dspark/test_dspark_block_accept_estimator.py new file mode 100644 index 000000000..021c7191d --- /dev/null +++ b/test/registered/spec/dspark/test_dspark_block_accept_estimator.py @@ -0,0 +1,579 @@ +import json +import math +import tempfile +import unittest +from pathlib import Path + +import torch + +from sglang.srt.speculative.dspark_components.dspark_block_accept_estimator import ( + BlockAcceptEstimateRecorder, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +_GAMMA = 3 +_VOCAB = 8 + + +class _FakeLayout: + def __init__(self, verify_lens: torch.Tensor): + self.verify_lens = verify_lens + + +def _reference_logprob( + logits_row: torch.Tensor, token: int, temperature: float +) -> float: + scaled = logits_row.to(torch.float32) / temperature + return float(scaled[token] - torch.logsumexp(scaled, dim=-1)) + + +def _make_recorder(tmp_dir: str) -> tuple[BlockAcceptEstimateRecorder, Path]: + path = Path(tmp_dir) / "estimate.jsonl" + recorder = BlockAcceptEstimateRecorder(path=str(path), gamma=_GAMMA, device="cpu") + return recorder, path + + +class _FakeDelayed: + + def __init__(self): + self._pending = None + + def step(self, *, compute_on_device, postprocess_on_host): + if self._pending is not None: + result, post = self._pending + self._pending = None + if result is not None: + post(result) + result = compute_on_device() + self._pending = (result, postprocess_on_host) if result is not None else None + + +def _observe( + recorder: BlockAcceptEstimateRecorder, + *, + forward_ct: int, + rid: str, + drafts: list[int], + corrected_logits: torch.Tensor, + target_logits: torch.Tensor, + verify_len: int, + correct_len: int, + bonus: int, + seq_len: int, + cap_trim: int = 0, + temperature: float = 1.0, +) -> None: + recorder.observe_verify_step( + forward_ct=forward_ct, + rids=[rid], + draft_tokens=torch.tensor([drafts], dtype=torch.int64), + corrected_logits=corrected_logits.unsqueeze(0), + draft_temperatures=torch.tensor([temperature], dtype=torch.float32), + greedy_mask=torch.tensor([False]), + target_logits=target_logits, + target_temperatures=torch.tensor([[temperature]], dtype=torch.float32), + truncated_sampling_mask=None, + logits_adjustments_are_noop=True, + correct_len=torch.tensor([correct_len], dtype=torch.int32), + cap_trim_lens=torch.tensor([cap_trim], dtype=torch.int32), + bonus=torch.tensor([bonus], dtype=torch.int64), + prefix_lens=torch.tensor([seq_len], dtype=torch.int64), + layout=_FakeLayout(torch.tensor([verify_len], dtype=torch.int32)), + ) + + +def _read_records(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text().splitlines()] + + +class TestBlockAcceptEstimateRecorder(CustomTestCase): + def test_exact_block_when_rejected_inside_window(self): + with tempfile.TemporaryDirectory() as tmp: + recorder, path = _make_recorder(tmp) + corrected = torch.randn(_GAMMA, _VOCAB) + target = torch.randn((_GAMMA + 1), _VOCAB) + _observe( + recorder, + forward_ct=1, + rid="r0", + drafts=[1, 2, 3], + corrected_logits=corrected, + target_logits=target, + verify_len=3, + correct_len=1, + bonus=5, + seq_len=10, + ) + recorder._file.flush() + + records = _read_records(path) + self.assertEqual(len(records), 1) + self.assertEqual(records[0]["w"], 2) + self.assertEqual(records[0]["cl"], 1) + self.assertNotIn("q_lp", records[0]) + self.assertNotIn("pg", records[0]) + + def test_censored_block_gathers_q_and_same_step_bonus_row_p(self): + with tempfile.TemporaryDirectory() as tmp: + recorder, path = _make_recorder(tmp) + corrected = torch.randn(_GAMMA, _VOCAB) + target = torch.randn((_GAMMA + 1), _VOCAB) + drafts = [1, 2, 3] + _observe( + recorder, + forward_ct=1, + rid="r0", + drafts=drafts, + corrected_logits=corrected, + target_logits=target, + verify_len=2, + correct_len=1, + bonus=2, + seq_len=10, + ) + recorder._file.flush() + + records = _read_records(path) + self.assertEqual(len(records), 1) + record = records[0] + self.assertEqual(record["w"], 1) + self.assertEqual(record["cl"], 1) + self.assertEqual(record["trimmed_tokens"], [2, 3]) + self.assertEqual(len(record["q_lp"]), 2) + self.assertAlmostEqual( + record["q_lp"][0], + _reference_logprob(corrected[1], 2, 1.0), + places=4, + ) + self.assertAlmostEqual( + record["q_lp"][1], + _reference_logprob(corrected[2], 3, 1.0), + places=4, + ) + self.assertEqual(len(record["pg"]), 1) + src_fct, offset, p_lp, draft_token, realized_token = record["pg"][0] + self.assertEqual((src_fct, offset), (1, 2)) + self.assertEqual(draft_token, 2) + self.assertEqual(realized_token, 2) + self.assertAlmostEqual( + p_lp, _reference_logprob(target[1], 2, 1.0), places=4 + ) + + def test_pending_block_resolves_in_later_step(self): + with tempfile.TemporaryDirectory() as tmp: + recorder, path = _make_recorder(tmp) + corrected1 = torch.randn(_GAMMA, _VOCAB) + target1 = torch.randn((_GAMMA + 1), _VOCAB) + _observe( + recorder, + forward_ct=1, + rid="r0", + drafts=[1, 2, 3], + corrected_logits=corrected1, + target_logits=target1, + verify_len=2, + correct_len=1, + bonus=2, + seq_len=10, + ) + + corrected2 = torch.randn(_GAMMA, _VOCAB) + target2 = torch.randn((_GAMMA + 1), _VOCAB) + _observe( + recorder, + forward_ct=2, + rid="r0", + drafts=[3, 6, 7], + corrected_logits=corrected2, + target_logits=target2, + verify_len=4, + correct_len=3, + bonus=4, + seq_len=12, + ) + recorder._file.flush() + + records = _read_records(path) + self.assertEqual(len(records), 2) + step2 = records[1] + self.assertEqual(len(step2["pg"]), 1) + src_fct, offset, p_lp, draft_token, realized_token = step2["pg"][0] + self.assertEqual((src_fct, offset), (1, 3)) + self.assertEqual(draft_token, 3) + self.assertEqual(realized_token, 3) + self.assertAlmostEqual( + p_lp, _reference_logprob(target2[0], 3, 1.0), places=4 + ) + self.assertEqual(recorder._states["r0"].pending, []) + + def test_divergence_drops_block_after_final_gather(self): + with tempfile.TemporaryDirectory() as tmp: + recorder, path = _make_recorder(tmp) + corrected1 = torch.randn(_GAMMA, _VOCAB) + target1 = torch.randn((_GAMMA + 1), _VOCAB) + _observe( + recorder, + forward_ct=1, + rid="r0", + drafts=[1, 2, 3], + corrected_logits=corrected1, + target_logits=target1, + verify_len=2, + correct_len=1, + bonus=6, + seq_len=10, + ) + recorder._file.flush() + + records = _read_records(path) + src_fct, offset, p_lp, draft_token, realized_token = records[0]["pg"][0] + self.assertEqual(draft_token, 2) + self.assertEqual(realized_token, 6) + self.assertEqual(recorder._states["r0"].pending, []) + + def test_temperature_scales_logprobs(self): + with tempfile.TemporaryDirectory() as tmp: + recorder, path = _make_recorder(tmp) + corrected = torch.randn(_GAMMA, _VOCAB) + target = torch.randn((_GAMMA + 1), _VOCAB) + _observe( + recorder, + forward_ct=1, + rid="r0", + drafts=[1, 2, 3], + corrected_logits=corrected, + target_logits=target, + verify_len=2, + correct_len=1, + bonus=2, + seq_len=10, + temperature=0.7, + ) + recorder._file.flush() + + record = _read_records(path)[0] + self.assertAlmostEqual( + record["q_lp"][0], + _reference_logprob(corrected[1], 2, 0.7), + places=4, + ) + self.assertAlmostEqual( + record["pg"][0][2], + _reference_logprob(target[1], 2, 0.7), + places=4, + ) + + def test_greedy_row_is_skipped_but_seq_len_bookkeeping_advances(self): + with tempfile.TemporaryDirectory() as tmp: + recorder, path = _make_recorder(tmp) + corrected = torch.randn(_GAMMA, _VOCAB) + target = torch.randn((_GAMMA + 1), _VOCAB) + recorder.observe_verify_step( + forward_ct=1, + rids=["r0"], + draft_tokens=torch.tensor([[1, 2, 3]], dtype=torch.int64), + corrected_logits=corrected.unsqueeze(0), + draft_temperatures=torch.tensor([1.0]), + greedy_mask=torch.tensor([True]), + target_logits=target, + target_temperatures=torch.tensor([[1.0]]), + truncated_sampling_mask=None, + logits_adjustments_are_noop=True, + correct_len=torch.tensor([2], dtype=torch.int32), + cap_trim_lens=torch.tensor([0], dtype=torch.int32), + bonus=torch.tensor([5], dtype=torch.int64), + prefix_lens=torch.tensor([10], dtype=torch.int64), + layout=_FakeLayout(torch.tensor([4], dtype=torch.int32)), + ) + recorder._file.flush() + + self.assertEqual(path.read_text(), "") + self.assertEqual(recorder._states["r0"].expected_seq_len, 13) + + def test_seq_len_discontinuity_drops_pending_blocks(self): + with tempfile.TemporaryDirectory() as tmp: + recorder, path = _make_recorder(tmp) + corrected1 = torch.randn(_GAMMA, _VOCAB) + target1 = torch.randn((_GAMMA + 1), _VOCAB) + _observe( + recorder, + forward_ct=1, + rid="r0", + drafts=[1, 2, 3], + corrected_logits=corrected1, + target_logits=target1, + verify_len=2, + correct_len=1, + bonus=2, + seq_len=10, + ) + self.assertEqual(len(recorder._states["r0"].pending), 1) + + corrected2 = torch.randn(_GAMMA, _VOCAB) + target2 = torch.randn((_GAMMA + 1), _VOCAB) + _observe( + recorder, + forward_ct=2, + rid="r0", + drafts=[3, 6, 7], + corrected_logits=corrected2, + target_logits=target2, + verify_len=4, + correct_len=3, + bonus=4, + seq_len=11, + ) + recorder._file.flush() + + records = _read_records(path) + self.assertNotIn("pg", records[1]) + self.assertEqual(recorder._states["r0"].pending, []) + self.assertEqual(recorder._discontinuity_drop_ct, 1) + + def test_truncated_sampling_row_is_excluded_while_clean_row_records(self): + with tempfile.TemporaryDirectory() as tmp: + recorder, path = _make_recorder(tmp) + corrected = torch.randn(2, _GAMMA, _VOCAB) + target = torch.randn(2 * (_GAMMA + 1), _VOCAB) + recorder.observe_verify_step( + forward_ct=1, + rids=["r_clean", "r_top_p"], + draft_tokens=torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int64), + corrected_logits=corrected, + draft_temperatures=torch.tensor([1.0, 1.0]), + greedy_mask=torch.tensor([False, False]), + target_logits=target, + target_temperatures=torch.tensor([[1.0], [1.0]]), + truncated_sampling_mask=torch.tensor([False, True]), + logits_adjustments_are_noop=True, + correct_len=torch.tensor([1, 1], dtype=torch.int32), + cap_trim_lens=torch.tensor([0, 0], dtype=torch.int32), + bonus=torch.tensor([2, 7], dtype=torch.int64), + prefix_lens=torch.tensor([10, 20], dtype=torch.int64), + layout=_FakeLayout(torch.tensor([2, 2], dtype=torch.int32)), + ) + recorder._file.flush() + + records = _read_records(path) + self.assertEqual([r["rid"] for r in records], ["r_clean"]) + self.assertEqual(recorder._skipped_step_ct, 0) + self.assertEqual(recorder._states["r_top_p"].pending, []) + self.assertEqual(recorder._states["r_top_p"].expected_seq_len, 22) + + +def _offline_estimate(path: Path, gamma: int) -> tuple[float, float, int]: + from collections import defaultdict + + blocks: list[dict] = [] + gathers: dict[tuple, list] = defaultdict(list) + for line in path.read_text().splitlines(): + rec = json.loads(line) + blocks.append(rec) + for src_fct, offset, p_lp, draft_token, realized_token in rec.get("pg", []): + gathers[(rec["rid"], src_fct)].append( + [offset, p_lp, draft_token, realized_token] + ) + los: list[float] = [] + his: list[float] = [] + for rec in blocks: + cl, w = rec["cl"], rec["w"] + if "q_lp" not in rec: + los.append(cl + 1.0) + his.append(cl + 1.0) + continue + q_lps = rec["q_lp"] + entries = {e[0]: e for e in gathers.get((rec["rid"], rec["fct"]), [])} + base, prod, lo_extra, tail = w + 1.0, 1.0, 0.0, 0.0 + for offset in range(w + 1, gamma + 1): + entry = entries.get(offset) + if entry is None: + tail = prod * (gamma - offset + 1) + break + _, p_lp, draft_token, realized_token = entry + a = min(1.0, math.exp(p_lp - q_lps[offset - w - 1])) + prod *= a + lo_extra += prod + if draft_token != realized_token: + if offset < gamma: + tail = prod * (gamma - offset) + break + los.append(base + lo_extra) + his.append(base + lo_extra + tail) + n = len(los) + return sum(los) / n, sum(his) / n, n + + +class TestOnlineCeilingEstimate(CustomTestCase): + def test_online_estimate_matches_offline_aggregation(self): + bs, steps = 4, 14 + gen = torch.Generator().manual_seed(11) + seq = [50 + 3 * b for b in range(bs)] + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "est.jsonl" + recorder = BlockAcceptEstimateRecorder( + path=str(path), + gamma=_GAMMA, + device="cpu", + online_log_interval=3, + online_window_steps=2, + ) + for t in range(steps): + verify_lens, correct_lens, drafts, bonus, prefix = [], [], [], [], [] + for b in range(bs): + vl = int(torch.randint(1, _GAMMA + 2, (1,), generator=gen)) + window = vl - 1 + cl = int(torch.randint(0, window + 1, (1,), generator=gen)) + row = torch.randint(0, _VOCAB, (_GAMMA,), generator=gen).tolist() + if cl < _GAMMA and int(torch.randint(0, 2, (1,), generator=gen)): + bt = row[cl] + else: + bt = int(torch.randint(0, _VOCAB, (1,), generator=gen)) + verify_lens.append(vl) + correct_lens.append(cl) + drafts.append(row) + bonus.append(bt) + prefix.append(seq[b]) + seq[b] += cl + 1 + recorder.observe_verify_step( + forward_ct=t + 1, + rids=[f"r{b}" for b in range(bs)], + draft_tokens=torch.tensor(drafts, dtype=torch.int64), + corrected_logits=torch.randn(bs, _GAMMA, _VOCAB, generator=gen), + draft_temperatures=torch.ones(bs), + greedy_mask=torch.zeros(bs, dtype=torch.bool), + target_logits=torch.randn(bs * (_GAMMA + 1), _VOCAB, generator=gen), + target_temperatures=torch.ones(bs), + truncated_sampling_mask=None, + logits_adjustments_are_noop=True, + correct_len=torch.tensor(correct_lens, dtype=torch.int32), + cap_trim_lens=torch.tensor( + [_GAMMA - (v - 1) for v in verify_lens], dtype=torch.int32 + ), + bonus=torch.tensor(bonus, dtype=torch.int64), + prefix_lens=torch.tensor(prefix, dtype=torch.int64), + layout=_FakeLayout(torch.tensor(verify_lens, dtype=torch.int32)), + ) + recorder.drain_pending_online() + recorder._file.flush() + + off_lo, off_hi, off_n = _offline_estimate(path, _GAMMA) + snap = recorder.online_estimate() + self.assertIsNotNone(snap) + self.assertEqual(snap.cumulative_blocks, off_n) + self.assertAlmostEqual(snap.cumulative_lo, off_lo, places=6) + self.assertAlmostEqual(snap.cumulative_hi, off_hi, places=6) + self.assertGreater(off_n, bs) + self.assertLessEqual(snap.window_blocks, snap.cumulative_blocks) + + def test_online_window_evicts_forward_passes_outside_horizon(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "est.jsonl" + recorder = BlockAcceptEstimateRecorder( + path=str(path), + gamma=_GAMMA, + device="cpu", + online_log_interval=1, + online_window_steps=3, + ) + target = torch.randn((_GAMMA + 1), _VOCAB) + seq = 10 + for t in range(10): + cl = t % 3 + _observe( + recorder, + forward_ct=t + 1, + rid="r0", + drafts=[1, 2, 3], + corrected_logits=torch.randn(_GAMMA, _VOCAB), + target_logits=target, + verify_len=_GAMMA + 1, + correct_len=cl, + bonus=5, + seq_len=seq, + ) + seq += cl + 1 + snap = recorder.online_estimate() + self.assertIsNotNone(snap) + self.assertLessEqual(snap.window_horizon, 3) + self.assertLessEqual(snap.window_blocks, 3) + self.assertEqual(snap.cumulative_blocks, 10) + + +class TestNaturalStopEosTail(CustomTestCase): + def _finalize_kept_block(self, *, natural_stop: bool): + with tempfile.TemporaryDirectory() as tmp: + recorder, _ = _make_recorder(tmp) + _observe( + recorder, + forward_ct=1, + rid="r0", + drafts=[1, 2, 3], + corrected_logits=torch.randn(_GAMMA, _VOCAB), + target_logits=torch.randn((_GAMMA + 1), _VOCAB), + verify_len=2, + correct_len=1, + bonus=2, + seq_len=10, + ) + self.assertEqual(len(recorder._states["r0"].pending), 1) + recorder.note_request_finished(rid="r0", natural_stop=natural_stop) + self.assertNotIn("r0", recorder._states) + return recorder.online_estimate() + + def test_natural_eos_caps_tail_to_zero(self): + snap = self._finalize_kept_block(natural_stop=True) + self.assertEqual(snap.cumulative_blocks, 1) + self.assertAlmostEqual(snap.cumulative_lo, snap.cumulative_hi, places=6) + + def test_external_finish_keeps_optimistic_tail(self): + snap = self._finalize_kept_block(natural_stop=False) + self.assertEqual(snap.cumulative_blocks, 1) + self.assertGreater(snap.cumulative_hi, snap.cumulative_lo) + + +class TestAsyncFinishIntent(CustomTestCase): + def test_intent_buffered_then_applied_at_next_drain(self): + with tempfile.TemporaryDirectory() as tmp: + recorder, _ = _make_recorder(tmp) + recorder._delayed = _FakeDelayed() + + _observe( + recorder, + forward_ct=1, + rid="r0", + drafts=[1, 2, 3], + corrected_logits=torch.randn(_GAMMA, _VOCAB), + target_logits=torch.randn((_GAMMA + 1), _VOCAB), + verify_len=2, + correct_len=1, + bonus=2, + seq_len=10, + ) + recorder.note_request_finished(rid="r0", natural_stop=True) + self.assertIn("r0", recorder._finish_intents) + self.assertIsNone(recorder.online_estimate()) + + _observe( + recorder, + forward_ct=2, + rid="r1", + drafts=[4, 5, 6], + corrected_logits=torch.randn(_GAMMA, _VOCAB), + target_logits=torch.randn((_GAMMA + 1), _VOCAB), + verify_len=_GAMMA + 1, + correct_len=1, + bonus=7, + seq_len=20, + ) + self.assertNotIn("r0", recorder._finish_intents) + self.assertNotIn("r0", recorder._states) + snap = recorder.online_estimate() + self.assertEqual(snap.cumulative_blocks, 1) + self.assertAlmostEqual(snap.cumulative_lo, snap.cumulative_hi, places=6) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_dspark_confidence_metrics.py b/test/registered/spec/dspark/test_dspark_confidence_metrics.py new file mode 100644 index 000000000..a91eb777b --- /dev/null +++ b/test/registered/spec/dspark/test_dspark_confidence_metrics.py @@ -0,0 +1,159 @@ +import unittest + +import torch + +from sglang.srt.environ import envs +from sglang.srt.speculative.dspark_components.dspark_observability import ( + ConfidenceMetricsProbe, + PerPositionConfidenceMetrics, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") + + +def _cpu_metrics(gamma: int) -> PerPositionConfidenceMetrics: + return PerPositionConfidenceMetrics(gamma=gamma, device=torch.device("cpu")) + + +class TestPerPositionConfidenceMetrics(CustomTestCase): + def test_perfectly_calibrated_has_low_ece(self): + torch.manual_seed(0) + n = 40000 + survival = torch.full((n, 1), 0.3, dtype=torch.float64) + prefix_mask = (torch.rand(n, 1) < 0.3).to(torch.float64) + metrics = _cpu_metrics(gamma=1) + metrics.update(survival=survival, prefix_mask=prefix_mask) + row = metrics.compute()[0] + self.assertLess(row["ece"], 0.03) + self.assertAlmostEqual(row["pred_mean"], 0.3, places=4) + + def test_overconfident_has_high_ece_and_pred_above_target(self): + torch.manual_seed(0) + n = 40000 + survival = torch.full((n, 1), 0.9, dtype=torch.float64) + prefix_mask = (torch.rand(n, 1) < 0.3).to(torch.float64) + metrics = _cpu_metrics(gamma=1) + metrics.update(survival=survival, prefix_mask=prefix_mask) + row = metrics.compute()[0] + self.assertGreater(row["ece"], 0.4) + self.assertGreater(row["pred_mean"], row["target_mean"]) + + def test_separable_scores_give_auc_near_one(self): + torch.manual_seed(0) + n = 20000 + pos = torch.rand(n, 1) * 0.3 + 0.7 + neg = torch.rand(n, 1) * 0.3 + survival = torch.cat([pos, neg], dim=0) + prefix_mask = torch.cat([torch.ones(n, 1), torch.zeros(n, 1)], dim=0) + metrics = _cpu_metrics(gamma=1) + metrics.update(survival=survival, prefix_mask=prefix_mask) + self.assertGreater(metrics.compute()[0]["auc"], 0.99) + + def test_random_scores_give_auc_near_half(self): + torch.manual_seed(0) + n = 40000 + survival = torch.rand(n, 1) + prefix_mask = (torch.rand(n, 1) < 0.5).to(torch.float64) + metrics = _cpu_metrics(gamma=1) + metrics.update(survival=survival, prefix_mask=prefix_mask) + auc = metrics.compute()[0]["auc"] + self.assertGreater(auc, 0.45) + self.assertLess(auc, 0.55) + + def test_batched_update_matches_per_sample_update(self): + torch.manual_seed(0) + bs, gamma = 32, 5 + survival = torch.rand(bs, gamma) + prefix_mask = (torch.rand(bs, gamma) < 0.5).to(torch.float64) + + batched = _cpu_metrics(gamma=gamma) + batched.update(survival=survival, prefix_mask=prefix_mask) + + per_sample = _cpu_metrics(gamma=gamma) + for row_idx in range(bs): + per_sample.update( + survival=survival[row_idx : row_idx + 1], + prefix_mask=prefix_mask[row_idx : row_idx + 1], + ) + + for name in ( + "coarse_count", + "coarse_pred", + "coarse_target", + "fine_pos", + "fine_neg", + "brier_num", + ): + self.assertTrue( + torch.allclose( + getattr(batched, name), getattr(per_sample, name), atol=1e-9 + ), + msg=name, + ) + + +def _probe_inputs(bs: int, gamma: int): + torch.manual_seed(0) + verify_num_draft_tokens = gamma + 1 + vocab = 16 + verify_ids_2d = torch.randint(0, vocab, (bs, verify_num_draft_tokens)) + target_logits = torch.randn(bs * verify_num_draft_tokens, vocab) + confidence_raw = torch.randn(bs, gamma) + return verify_ids_2d, target_logits, confidence_raw + + +class TestConfidenceMetricsProbe(CustomTestCase): + def _observe(self, probe, *, carries_confidence=True, is_compact_mode=False): + verify_ids_2d, target_logits, confidence_raw = _probe_inputs(bs=3, gamma=4) + probe.maybe_observe( + carries_confidence=carries_confidence, + is_compact_mode=is_compact_mode, + confidence_raw=confidence_raw, + verify_ids_2d=verify_ids_2d, + target_logits=target_logits, + bs=3, + ) + + def test_disabled_env_is_noop(self): + probe = ConfidenceMetricsProbe(gamma=4, verify_num_draft_tokens=5, tp_rank=0) + self._observe(probe) + self.assertIsNone(probe._metrics) + self.assertEqual(probe._step_ct, 0) + + def test_non_rank0_is_noop(self): + probe = ConfidenceMetricsProbe(gamma=4, verify_num_draft_tokens=5, tp_rank=1) + with envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS.override(True): + self._observe(probe) + self.assertIsNone(probe._metrics) + + def test_missing_confidence_head_is_noop(self): + probe = ConfidenceMetricsProbe(gamma=4, verify_num_draft_tokens=5, tp_rank=0) + with envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS.override(True): + self._observe(probe, carries_confidence=False) + self.assertIsNone(probe._metrics) + + def test_compact_mode_warns_once_and_skips(self): + probe = ConfidenceMetricsProbe(gamma=4, verify_num_draft_tokens=5, tp_rank=0) + with envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS.override(True): + self._observe(probe, is_compact_mode=True) + self.assertTrue(probe._compact_warned) + self._observe(probe, is_compact_mode=True) + self.assertIsNone(probe._metrics) + self.assertEqual(probe._step_ct, 0) + + def test_enabled_path_accumulates(self): + probe = ConfidenceMetricsProbe( + gamma=4, verify_num_draft_tokens=5, tp_rank=0, print_every=2 + ) + with envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS.override(True): + self._observe(probe) + self.assertIsInstance(probe._metrics, PerPositionConfidenceMetrics) + self.assertEqual(probe._step_ct, 1) + self._observe(probe) + self.assertEqual(probe._step_ct, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_dspark_dp_tier.py b/test/registered/spec/dspark/test_dspark_dp_tier.py new file mode 100644 index 000000000..f0bcb4c7b --- /dev/null +++ b/test/registered/spec/dspark/test_dspark_dp_tier.py @@ -0,0 +1,99 @@ +import random +import unittest + +from sglang.srt.speculative.dspark_components.dspark_planner import ( + dp_global_verify_tier_num_tokens, + local_verify_tier_num_tokens, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class TestLocalVerifyTierNumTokens(CustomTestCase): + def test_no_budget_returns_sentinel(self): + self.assertEqual( + local_verify_tier_num_tokens( + bs=8, + verify_token_budget=None, + verify_num_draft_tokens=6, + min_verify_len=1, + ), + -1, + ) + + def test_budget_adds_to_anchor_floor(self): + self.assertEqual( + local_verify_tier_num_tokens( + bs=8, + verify_token_budget=10, + verify_num_draft_tokens=6, + min_verify_len=1, + ), + 18, + ) + + # Clamp/floor variants (verify-all clamp, min_verify_len floor, min=0) are + # covered by the TestBusyIdleGraphKeyIdentity sweep bounds. + + +class TestDpGlobalVerifyTierNumTokens(CustomTestCase): + def test_any_sentinel_pins_everyone(self): + # The sweep never emits a -1 contribution, so this is the only guard + # on "any rank without a budget pins everyone"; losing it forks graph + # keys across DP ranks. + self.assertIsNone( + dp_global_verify_tier_num_tokens(global_tier_num_tokens=[100, -1, 50, 0]) + ) + + +class TestBusyIdleGraphKeyIdentity(CustomTestCase): + def test_busy_and_idle_floors_agree_on_random_topologies(self): + rng = random.Random(20260703) + for _ in range(2000): + verify_num_draft_tokens = rng.randint(2, 8) + min_verify_len = rng.randint(0, verify_num_draft_tokens - 1) + effective_min = max(min_verify_len, 1) + num_ranks = rng.randint(1, 8) + contributions = [] + num_reqs_per_rank = [] + for _ in range(num_ranks): + if rng.random() < 0.3: + num_reqs_per_rank.append(0) + contributions.append(0) + continue + bs = rng.randint(1, 512) + budget = rng.randint(0, bs * verify_num_draft_tokens) + num_reqs_per_rank.append(bs) + contributions.append( + local_verify_tier_num_tokens( + bs=bs, + verify_token_budget=budget, + verify_num_draft_tokens=verify_num_draft_tokens, + min_verify_len=min_verify_len, + ) + ) + tier_num_tokens = dp_global_verify_tier_num_tokens( + global_tier_num_tokens=contributions + ) + global_num_reqs = max(num_reqs_per_rank) + if tier_num_tokens is None: + self.assertEqual(global_num_reqs, 0) + continue + + self.assertGreaterEqual(tier_num_tokens, global_num_reqs * effective_min) + self.assertLessEqual( + tier_num_tokens, global_num_reqs * verify_num_draft_tokens + ) + + busy_floor = min(tier_num_tokens, global_num_reqs * verify_num_draft_tokens) + self.assertEqual(busy_floor, tier_num_tokens) + + idle_lens_total = global_num_reqs + idle_bucket_input = max(idle_lens_total, tier_num_tokens) + self.assertEqual(idle_bucket_input, tier_num_tokens) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_dspark_draft_path_default.py b/test/registered/spec/dspark/test_dspark_draft_path_default.py new file mode 100644 index 000000000..b09ae0651 --- /dev/null +++ b/test/registered/spec/dspark/test_dspark_draft_path_default.py @@ -0,0 +1,88 @@ +import unittest +from types import SimpleNamespace + +from sglang.srt.arg_groups.speculative_hook import ( + _handle_dspark, + _target_checkpoint_bundles_dspark_draft, +) +from sglang.srt.server_args import ServerArgs +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +_BUNDLED_MODEL_PATH = "deepseek-ai/DeepSeek-V4-Flash-DSpark" +_PLAIN_MODEL_PATH = "deepseek-ai/DeepSeek-V4-Flash" + + +def _bundled_hf_config() -> SimpleNamespace: + return SimpleNamespace( + architectures=["DeepseekV4ForCausalLM"], + dspark_block_size=5, + dspark_markov_rank=256, + dspark_target_layer_ids=[40, 41, 42], + dspark_noise_token_id=128799, + ) + + +def _plain_hf_config() -> SimpleNamespace: + return SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]) + + +def _make_dspark_server_args( + *, model_path: str, hf_config: SimpleNamespace +) -> ServerArgs: + server_args = ServerArgs(model_path="dummy") + server_args.model_path = model_path + server_args.device = "cuda" + server_args.speculative_algorithm = "DSPARK" + server_args.speculative_draft_model_path = None + server_args.speculative_dspark_block_size = 5 + server_args.model_config = SimpleNamespace(hf_config=hf_config) + return server_args + + +class TestTargetCheckpointBundlesDsparkDraft(CustomTestCase): + def test_bundled_dsv4_config_is_detected(self): + server_args = _make_dspark_server_args( + model_path=_BUNDLED_MODEL_PATH, hf_config=_bundled_hf_config() + ) + self.assertTrue(_target_checkpoint_bundles_dspark_draft(server_args)) + + def test_plain_target_config_is_not_detected(self): + server_args = _make_dspark_server_args( + model_path=_PLAIN_MODEL_PATH, hf_config=_plain_hf_config() + ) + self.assertFalse(_target_checkpoint_bundles_dspark_draft(server_args)) + + +class TestDsparkDraftPathDefaulting(CustomTestCase): + def test_bundled_checkpoint_defaults_draft_path_to_model_path(self): + server_args = _make_dspark_server_args( + model_path=_BUNDLED_MODEL_PATH, hf_config=_bundled_hf_config() + ) + _handle_dspark(server_args) + self.assertEqual(server_args.speculative_draft_model_path, _BUNDLED_MODEL_PATH) + self.assertEqual(server_args.speculative_num_draft_tokens, 6) + + def test_plain_target_without_draft_path_raises(self): + server_args = _make_dspark_server_args( + model_path=_PLAIN_MODEL_PATH, hf_config=_plain_hf_config() + ) + with self.assertRaises(ValueError): + _handle_dspark(server_args) + + def test_explicit_draft_path_is_not_overwritten(self): + server_args = _make_dspark_server_args( + model_path=_BUNDLED_MODEL_PATH, hf_config=_bundled_hf_config() + ) + server_args.speculative_draft_model_path = "deepseek-ai/some-other-dspark-draft" + _handle_dspark(server_args) + self.assertEqual( + server_args.speculative_draft_model_path, + "deepseek-ai/some-other-dspark-draft", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_dspark_info_dumper.py b/test/registered/spec/dspark/test_dspark_info_dumper.py new file mode 100644 index 000000000..0e948a673 --- /dev/null +++ b/test/registered/spec/dspark/test_dspark_info_dumper.py @@ -0,0 +1,398 @@ +import unittest + +import torch + +from sglang.srt.environ import envs +from sglang.srt.speculative.dspark_components.dspark_observability import ( + DecodeStepObservation, + DsparkInfoDumper, + InfoComponent, + _PendingStep, + logger, + resolve_components, + resolve_enabled_components, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class FakeClock: + def __init__(self) -> None: + self.now = 100.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +def make_dumper(components, **kwargs): + clock = FakeClock() + dumper = DsparkInfoDumper( + components=set(components), + gamma=5, + verify_num_draft_tokens=6, + attn_tp_rank=0, + device=torch.device("cpu"), + mode_value="static", + clock=clock, + **kwargs, + ) + return dumper, clock + + +def make_obs( + *, + forward_ct, + bs=4, + num_verify_tokens=24, + predicted_step_ms=None, + predicted_theta=None, +): + return DecodeStepObservation( + forward_ct=forward_ct, + bs=bs, + mode="static", + budget=100, + lag_steps=0, + num_verify_tokens=num_verify_tokens, + verify_tokens_local=num_verify_tokens, + verify_tokens_dp_synced=num_verify_tokens, + verify_tokens_graph_key=num_verify_tokens, + predicted_step_ms=predicted_step_ms, + predicted_theta=predicted_theta, + verify_lens=torch.full((bs,), 6, dtype=torch.int32), + confidence=torch.full((bs, 5), 0.9), + req_pool_indices=torch.arange(bs, dtype=torch.int64), + prefix_lens=torch.full((bs,), 128, dtype=torch.int64), + draft_tokens=torch.zeros((bs, 5), dtype=torch.int64), + bonus_tokens=torch.zeros((bs,), dtype=torch.int64), + correct_len=torch.full((bs,), 3, dtype=torch.int32), + cap_trim_lens=torch.zeros((bs,), dtype=torch.int32), + commit_lens=torch.full((bs,), 4, dtype=torch.int32), + rids=[f"r{i}" for i in range(bs)], + ) + + +class TestResolveComponents(CustomTestCase): + def test_empty_disables(self): + self.assertEqual(resolve_components(()), set()) + + def test_all_expands_to_every_component(self): + self.assertEqual(resolve_components(("all",)), set(InfoComponent)) + + def test_subset_and_whitespace_are_kept(self): + self.assertEqual( + resolve_components((" core ", "reqs")), + {InfoComponent.CORE, InfoComponent.REQS}, + ) + + def test_unknown_component_raises(self): + with self.assertRaises(ValueError): + resolve_components(("core", "bogus")) + + def test_sps_record_env_enables_core_and_cpu_timing(self): + """SGLANG_DSPARK_ENABLE_SPS_RECORD=1 is the published SPS-profiling + switch; it must keep enabling the components the table fit reads.""" + with envs.SGLANG_DSPARK_ENABLE_SPS_RECORD.override(True): + self.assertEqual( + resolve_enabled_components(), + {InfoComponent.CORE, InfoComponent.STEP_CPU_TIME}, + ) + + def test_sps_record_env_unions_with_debug_dump(self): + with envs.SGLANG_DSPARK_ENABLE_SPS_RECORD.override(True): + with envs.SGLANG_DSPARK_DEBUG_DUMP.override("reqs"): + self.assertEqual( + resolve_enabled_components(), + { + InfoComponent.CORE, + InfoComponent.STEP_CPU_TIME, + InfoComponent.REQS, + }, + ) + + +class TestCoreAndCpuTiming(CustomTestCase): + def test_disabled_dumper_records_nothing(self): + dumper, clock = make_dumper(set()) + dumper.begin_step() + dumper.observe_decode_step(make_obs(forward_ct=1)) + self.assertIsNone(dumper.dump()) + + def test_non_root_rank_is_disabled(self): + clock = FakeClock() + dumper = DsparkInfoDumper( + components={"core"}, + gamma=5, + verify_num_draft_tokens=6, + attn_tp_rank=1, + device=torch.device("cpu"), + mode_value="static", + clock=clock, + ) + self.assertFalse(dumper.enabled) + dumper.observe_decode_step(make_obs(forward_ct=1)) + self.assertIsNone(dumper.dump()) + + def test_one_record_per_step_including_the_last(self): + dumper, clock = make_dumper({"core", "step_cpu_time"}) + for forward_ct in range(1, 4): + dumper.observe_decode_step(make_obs(forward_ct=forward_ct)) + clock.advance(0.01) + records = dumper.dump()["records"] + self.assertEqual([r["forward_ct"] for r in records], [1, 2, 3]) + + def test_step_cpu_ms_is_attributed_to_the_step_it_measures(self): + dumper, clock = make_dumper({"core", "step_cpu_time"}) + dumper.observe_decode_step(make_obs(forward_ct=1)) + clock.advance(0.02) + dumper.observe_decode_step(make_obs(forward_ct=2)) + records = dumper.dump()["records"] + first = next(r for r in records if r["forward_ct"] == 1) + second = next(r for r in records if r["forward_ct"] == 2) + self.assertNotIn("step_cpu_ms", first) + self.assertAlmostEqual(second["step_cpu_ms"], 20.0, places=3) + + def test_core_fields_present(self): + dumper, _ = make_dumper({"core"}) + dumper.observe_decode_step(make_obs(forward_ct=7, bs=3, num_verify_tokens=18)) + record = dumper.dump()["records"][0] + self.assertEqual(record["bs"], 3) + self.assertEqual(record["num_running_reqs"], 3) + self.assertEqual(record["num_verify_tokens"], 18) + self.assertEqual(record["mode"], "static") + + def test_core_only_omits_timing_fields(self): + dumper, clock = make_dumper({"core"}) + dumper.observe_decode_step(make_obs(forward_ct=1)) + clock.advance(0.01) + dumper.observe_decode_step(make_obs(forward_ct=2)) + for record in dumper.dump()["records"]: + self.assertNotIn("step_cpu_ms", record) + + def test_non_decode_step_resets_cpu_pairing(self): + dumper, clock = make_dumper({"core", "step_cpu_time"}) + dumper.observe_decode_step(make_obs(forward_ct=1)) + clock.advance(0.02) + dumper.note_non_decode_step() + clock.advance(0.02) + dumper.observe_decode_step(make_obs(forward_ct=3)) + records = dumper.dump()["records"] + self.assertEqual([r["forward_ct"] for r in records], [1, 3]) + for record in records: + self.assertNotIn("step_cpu_ms", record) + + def test_oversized_gap_nulls_cpu_ms_but_keeps_record(self): + dumper, clock = make_dumper({"core", "step_cpu_time"}, max_step_cpu_seconds=0.5) + dumper.observe_decode_step(make_obs(forward_ct=1)) + clock.advance(0.6) + dumper.observe_decode_step(make_obs(forward_ct=2)) + records = dumper.dump()["records"] + self.assertEqual([r["forward_ct"] for r in records], [1, 2]) + second = next(r for r in records if r["forward_ct"] == 2) + self.assertNotIn("step_cpu_ms", second) + + def test_ring_buffer_evicts_oldest(self): + dumper, clock = make_dumper({"core"}, max_records=3) + for forward_ct in range(1, 8): + dumper.observe_decode_step(make_obs(forward_ct=forward_ct)) + clock.advance(0.01) + records = dumper.dump()["records"] + self.assertEqual([r["forward_ct"] for r in records], [5, 6, 7]) + + def test_dump_is_repeatable(self): + dumper, clock = make_dumper({"core"}) + dumper.observe_decode_step(make_obs(forward_ct=1)) + clock.advance(0.01) + dumper.observe_decode_step(make_obs(forward_ct=2)) + self.assertEqual(dumper.dump(), dumper.dump()) + + def test_clear_drops_all_records_and_pending(self): + dumper, clock = make_dumper({"core"}) + dumper.observe_decode_step(make_obs(forward_ct=1)) + clock.advance(0.01) + dumper.observe_decode_step(make_obs(forward_ct=2)) + dumper.clear() + self.assertEqual(dumper.dump()["records"], []) + dumper.observe_decode_step(make_obs(forward_ct=9)) + clock.advance(0.01) + dumper.observe_decode_step(make_obs(forward_ct=10)) + self.assertEqual([r["forward_ct"] for r in dumper.dump()["records"]], [9, 10]) + + +class TestPredictedStepFields(CustomTestCase): + def test_predicted_fields_recorded_under_core(self): + dumper, clock = make_dumper({"core"}) + dumper.observe_decode_step( + make_obs(forward_ct=1, predicted_step_ms=1.5, predicted_theta=200.0) + ) + clock.advance(0.01) + dumper.observe_decode_step(make_obs(forward_ct=2)) + record = next(r for r in dumper.dump()["records"] if r["forward_ct"] == 1) + self.assertAlmostEqual(record["predicted_step_ms"], 1.5) + self.assertAlmostEqual(record["predicted_theta"], 200.0) + + def test_predicted_fields_omitted_when_none(self): + dumper, clock = make_dumper({"core"}) + dumper.observe_decode_step(make_obs(forward_ct=1)) + clock.advance(0.01) + dumper.observe_decode_step(make_obs(forward_ct=2)) + record = next(r for r in dumper.dump()["records"] if r["forward_ct"] == 1) + self.assertNotIn("predicted_step_ms", record) + self.assertNotIn("predicted_theta", record) + + +def _pending(*, bs, budget, num_verify_tokens, predicted_step_ms): + return _PendingStep( + forward_ct=1, + bs=bs, + mode="compact", + budget=budget, + lag_steps=1, + num_verify_tokens=num_verify_tokens, + verify_tokens_local=num_verify_tokens, + verify_tokens_dp_synced=num_verify_tokens, + verify_tokens_graph_key=num_verify_tokens, + predicted_step_ms=predicted_step_ms, + predicted_theta=1.0, + step_cpu_ms=None, + rids=None, + future=None, + segment_events={}, + ) + + +class TestOnlineSpsReporter(CustomTestCase): + def test_report_interval_enables_dumper_and_gpu_timing(self): + dumper, _ = make_dumper(set(), sps_report_interval=2) + self.assertTrue(dumper.enabled) + self.assertIn(InfoComponent.STEP_GPU_TIME, dumper._components) + + def test_report_interval_zero_leaves_dumper_disabled(self): + dumper, _ = make_dumper(set(), sps_report_interval=0) + self.assertFalse(dumper.enabled) + + def test_reporter_logs_summary_every_interval_matched_steps(self): + dumper, _ = make_dumper(set(), sps_report_interval=2) + matched = dict(bs=4, budget=20, num_verify_tokens=24) + with self.assertLogs(logger, level="INFO") as cm: + dumper._report_sps_prediction( + pending=_pending(**matched, predicted_step_ms=10.0), step_gpu_ms=12.0 + ) + dumper._report_sps_prediction( + pending=_pending(**matched, predicted_step_ms=8.0), step_gpu_ms=9.0 + ) + self.assertEqual(sum("SPS prediction" in m for m in cm.output), 1) + self.assertEqual(dumper._sps_window, []) + + def test_reporter_counts_mismatch_and_excludes_it_from_means(self): + dumper, _ = make_dumper(set(), sps_report_interval=1) + with self.assertLogs(logger, level="INFO") as cm: + dumper._report_sps_prediction( + pending=_pending( + bs=4, budget=99, num_verify_tokens=24, predicted_step_ms=10.0 + ), + step_gpu_ms=12.0, + ) + dumper._report_sps_prediction( + pending=_pending( + bs=4, budget=20, num_verify_tokens=24, predicted_step_ms=10.0 + ), + step_gpu_ms=12.0, + ) + self.assertTrue(any("M_mismatch_rate=50.0%" in m for m in cm.output)) + + def test_reporter_skips_steps_missing_prediction_or_actual(self): + dumper, _ = make_dumper(set(), sps_report_interval=2) + dumper._report_sps_prediction( + pending=_pending( + bs=4, budget=20, num_verify_tokens=24, predicted_step_ms=None + ), + step_gpu_ms=12.0, + ) + dumper._report_sps_prediction( + pending=_pending( + bs=4, budget=20, num_verify_tokens=24, predicted_step_ms=10.0 + ), + step_gpu_ms=None, + ) + self.assertEqual(dumper._sps_window, []) + self.assertEqual(dumper._sps_mismatched, 0) + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA for d2h staging") +class TestReqsAndGpuTiming(CustomTestCase): + def _cuda_obs(self, *, forward_ct, bs=4): + obs = make_obs(forward_ct=forward_ct, bs=bs) + return DecodeStepObservation( + forward_ct=obs.forward_ct, + bs=obs.bs, + mode=obs.mode, + budget=obs.budget, + lag_steps=obs.lag_steps, + num_verify_tokens=obs.num_verify_tokens, + verify_tokens_local=obs.verify_tokens_local, + verify_tokens_dp_synced=obs.verify_tokens_dp_synced, + verify_tokens_graph_key=obs.verify_tokens_graph_key, + predicted_step_ms=obs.predicted_step_ms, + predicted_theta=obs.predicted_theta, + verify_lens=obs.verify_lens.cuda(), + confidence=obs.confidence.cuda(), + req_pool_indices=obs.req_pool_indices.cuda(), + prefix_lens=obs.prefix_lens.cuda(), + draft_tokens=obs.draft_tokens.cuda(), + bonus_tokens=obs.bonus_tokens.cuda(), + correct_len=obs.correct_len.cuda(), + cap_trim_lens=obs.cap_trim_lens.cuda(), + commit_lens=obs.commit_lens.cuda(), + rids=obs.rids, + ) + + def _make(self, components): + return DsparkInfoDumper( + components=set(components), + gamma=5, + verify_num_draft_tokens=6, + attn_tp_rank=0, + device=torch.device("cuda"), + mode_value="static", + ) + + def test_reqs_component_stages_per_request_detail(self): + dumper = self._make({"core", "reqs"}) + dumper.observe_decode_step(self._cuda_obs(forward_ct=1, bs=3)) + dumper.observe_decode_step(self._cuda_obs(forward_ct=2, bs=3)) + record = next(r for r in dumper.dump()["records"] if r["forward_ct"] == 1) + self.assertEqual(len(record["reqs"]), 3) + req = record["reqs"][0] + self.assertEqual(req["rid"], "r0") + self.assertEqual(req["verify_len"], 6) + self.assertEqual(req["acc_len"], 4) + self.assertEqual(req["correct_drafts"], 3) + self.assertEqual(len(req["survival"]), 5) + + def test_gpu_timing_populates_segment_fields(self): + dumper = self._make( + {"step_gpu_time", "draft_gpu_time", "target_verify_gpu_time"} + ) + for forward_ct in (1, 2): + dumper.begin_step() + with dumper.segment("draft"): + torch.zeros(1024, device="cuda").sum() + with dumper.segment("target_verify"): + torch.zeros(1024, device="cuda").sum() + dumper.observe_decode_step(self._cuda_obs(forward_ct=forward_ct)) + record = next(r for r in dumper.dump()["records"] if r["forward_ct"] == 1) + # The segments launch real kernels, so resolved event pairs must + # measure strictly positive time; 0.0 would mean the events never ran. + self.assertGreater(record["step_gpu_ms"], 0.0) + self.assertGreater(record["draft_gpu_ms"], 0.0) + self.assertGreater(record["target_verify_gpu_ms"], 0.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_dspark_kernel_parity.py b/test/registered/spec/dspark/test_dspark_kernel_parity.py new file mode 100644 index 000000000..9faa2f7a6 --- /dev/null +++ b/test/registered/spec/dspark/test_dspark_kernel_parity.py @@ -0,0 +1,562 @@ +"""Seeded triton-vs-torch parity sweep for the DSpark kernels. + +Guards against toolchain drift (triton/torch upgrades) silently diverging +the triton implementations from their torch references. Every method calls +the production kernel pair directly via `Cls.torch(...)` / `Cls.triton(...)` +(no env-var dispatch) on a small set of adversarial inputs and compares +exactly, or with the tolerance the kernel is specified to meet. +""" + +import types +import unittest + +import torch + +from sglang.srt.layers.attention.dsv4 import attn_metadata_kernels +from sglang.srt.speculative import ragged_verify_kernels +from sglang.srt.speculative.dspark_components.dspark_planner import ( + DSparkScheduleConfig, +) +from sglang.srt.speculative.dspark_components.kernels import ( + dspark_accept, + dspark_attn_metadata, + dspark_draft_model, + dspark_schedule, + dspark_verify_window, +) +from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small") + +DEVICE = torch.device("cuda") +VOCAB = 129280 + + +def _ri(lo, hi, shape, dtype=torch.int64, g=None): + return torch.randint(lo, hi, shape, device=DEVICE, dtype=dtype, generator=g) + + +def _layout(verify_lens, graph_num_tokens): + return RaggedVerifyLayout.from_verify_lens_device( + verify_lens=verify_lens, graph_num_tokens=graph_num_tokens + ) + + +class _Bf16Linear(torch.nn.Module): + quant_method = None + + def __init__(self, weight): + super().__init__() + self.weight = weight + + def forward(self, x): + return torch.nn.functional.linear(x, self.weight), None + + +def _case_accept_greedy(tc): + torch.manual_seed(0) + bs, t = 8, 6 + candidates = _ri(0, 200, (bs, t)) + target_logits = torch.randn(bs * t, 200, device=DEVICE) + for cutoff in (None, _ri(1, t + 1, (bs,), torch.int32)): + tc._parity( + dspark_accept.AcceptGreedy, + candidates=candidates, + target_logits=target_logits, + verify_num_draft_tokens=t, + cutoff_verify_lens=cutoff, + ) + # gather_row_bonus: bonus token at a per-row column index. + table, idx = _ri(0, VOCAB, (64, t)), _ri(0, t, (64,), torch.int32) + ref = table[torch.arange(64, device=DEVICE), idx.long()] + tc._eq(dspark_accept.gather_row_bonus_triton(table=table, idx=idx), ref) + + +def _case_accept_sampling(tc): + torch.manual_seed(1) + bs, t = 64, 6 + accept_index = _ri(0, bs * t, (bs, t)) + predicts = _ri(0, VOCAB, (bs * t,)) + correct_len = _ri(0, t, (bs,), torch.int32) + rows = torch.arange(bs, device=DEVICE) + ref = predicts[accept_index[rows, correct_len.long()].long()] + got = dspark_accept.gather_two_level_bonus_triton( + accept_index=accept_index, predicts=predicts, correct_len=correct_len + ) + tc._eq(got, ref) + + +def _case_build_block_seq_lens_causal(tc): + torch.manual_seed(2) + seq_lens = _ri(1, 100000, (128,)) + for block_size in (1, 5, 7): + tc._parity( + dspark_attn_metadata.BuildBlockSeqLensCausal, + seq_lens=seq_lens, + block_size=block_size, + device=DEVICE, + ) + + +def _case_build_out_tokens(tc): + torch.manual_seed(3) + bs, gamma = 64, 5 + for cl_dtype in (torch.int32, torch.int64): + # Bonus insertion swept through every position 0..gamma. + cl = (torch.arange(bs, device=DEVICE) % (gamma + 1)).to(cl_dtype) + tc._parity( + dspark_verify_window.BuildOutTokens, + draft_tokens=_ri(0, VOCAB, (bs, gamma)), + correct_len=cl, + bonus=_ri(0, VOCAB, (bs,)), + verify_num_draft_tokens=gamma + 1, + gamma=gamma, + ) + + +def _case_build_ragged_verify_window(tc): + torch.manual_seed(4) + gamma, t, bs = 5, 6, 8 + verify_lens = _ri(1, t + 1, (bs,), torch.int32) + batch = types.SimpleNamespace( + seq_lens=_ri(1, 20, (bs,)), + req_pool_indices=torch.randperm(bs + 3, device=DEVICE)[:bs], + ) + model_runner = types.SimpleNamespace( + req_to_token_pool=types.SimpleNamespace( + req_to_token=_ri(0, 1_000_000, (bs + 3, 64), torch.int32) + ) + ) + for graph_num_tokens in (bs * t, (bs + 3) * t): # tight and bucket padding + tc._parity( + dspark_verify_window.BuildRaggedVerifyWindow, + batch=batch, + layout=_layout(verify_lens, graph_num_tokens), + draft_block_ids=_ri(0, VOCAB, (bs, gamma)), + draft_tokens=_ri(0, VOCAB, (bs, gamma)), + bs=bs, + device=DEVICE, + verify_num_draft_tokens=t, + model_runner=model_runner, + ) + + +def _case_build_step_local(tc): + torch.manual_seed(5) + for org_width, per_partition, bias_dtype in ( + (32320, 32384, torch.bfloat16), + (5000, 8192, torch.float32), + ): + bias = (torch.randn(3, org_width, device=DEVICE) * 3.0).to(bias_dtype) + base = torch.randn(3, per_partition, device=DEVICE) + got, _ = tc._parity( + dspark_draft_model.BuildStepLocal, bias=bias, base_local=base + ) + # Padding columns beyond org_width must stay pure base. + tc.assertTrue(torch.equal(got[:, org_width:], base[:, org_width:])) + + +def _case_cap_correct_len(tc): + torch.manual_seed(6) + bs, nd = 64, 6 + verify_lens = _ri(1, nd + 1, (bs,), torch.int32) + for cl_dtype in (torch.int32, torch.int64): + cl = (torch.arange(bs, device=DEVICE) % (nd + 1)).to(cl_dtype) + tc._parity(dspark_accept.CapCorrectLen, correct_len=cl, verify_lens=verify_lens) + + +def _case_causal_swa_page_indices(tc): + swa, num_pool, pool_len, num_q = 128, 64, 600, 40 + g = torch.Generator(device=DEVICE).manual_seed(7) + kw = dict( + req_to_token=_ri(0, 40000, (num_pool, pool_len), torch.int32, g), + full_to_swa_mapping=_ri(0, 1 << 20, (40000,), torch.int64, g), + req_pool_indices_repeated=_ri(0, num_pool, (num_q,), torch.int32, g), + swa_window=swa, + page_index_aligned_size=96, + ) + # Lens short of / straddling / beyond the SWA window boundary. + for lo, hi in ((1, swa), (swa - 4, swa + 4), (swa + 1, pool_len)): + lens = _ri(lo, hi, (num_q,), torch.int32, g) + cls = attn_metadata_kernels.BuildCausalSwaPageIndices + ref = cls.torch(seq_lens_casual=lens, **kw) + got = cls.triton(seq_lens_casual=lens, **kw) + tc.assertEqual(got.shape, ref.shape) + tc.assertEqual(got.dtype, ref.dtype) + # Parity holds on the attended region; padding slots must be -1. + col = torch.arange(ref.shape[1], device=DEVICE).view(1, -1) + attended = col < torch.clamp(lens, max=swa).view(-1, 1) + tc.assertTrue(torch.equal(got[attended], ref[attended])) + tc.assertTrue(bool((got[~attended] == -1).all())) + + +def _case_commit_inject_layout(tc): + stride, num_pool, pool_len, n_full, bs = 7, 300, 400, 50000, 64 + g = torch.Generator(device=DEVICE).manual_seed(8) + pool_perm = torch.randperm(num_pool, device=DEVICE, generator=g) + kw = dict( + req_pool_indices=pool_perm[:bs], + req_to_token=_ri(0, n_full, (num_pool, pool_len), torch.int64, g), + prefix_lens=_ri(1, pool_len - stride, (bs,), torch.int64, g), + block_pos_offsets=torch.arange(stride, device=DEVICE), + full_to_swa_mapping=_ri(0, 1 << 20, (n_full,), torch.int64, g), + commit_lens=_ri(0, stride + 1, (bs,), torch.int32, g), + stride=stride, + ) + tc._parity(dspark_verify_window.BuildCommitInjectLayout, **kw) + # commit_len edges: 0 masks the whole row to -1, stride keeps it all. + kw.update( + req_pool_indices=kw["req_pool_indices"][:2], + prefix_lens=kw["prefix_lens"][:2], + commit_lens=torch.tensor([0, stride], device=DEVICE, dtype=torch.int32), + ) + edge = dspark_verify_window.BuildCommitInjectLayout.triton(**kw) + swa_2d = edge.swa_loc.view(2, stride) + tc.assertTrue(bool((swa_2d[0] == -1).all())) + tc.assertTrue(bool((swa_2d[1] >= 0).all())) + + +def _case_commit_kv_proj(tc): + hidden, head_dim, num_stages = 1024, 576, 3 + g = torch.Generator(device=DEVICE).manual_seed(9) + linears = [ + _Bf16Linear( + (torch.randn(head_dim, hidden, device=DEVICE, generator=g) * 0.02).to( + torch.bfloat16 + ) + ) + for _ in range(num_stages) + ] + main_x = (torch.randn(56, hidden, device=DEVICE, generator=g) * 0.5).to( + torch.bfloat16 + ) + cls = dspark_draft_model.CommitKvProj + ref = cls.torch(main_x=main_x, wkv_linears=linears) + got = cls.triton(main_x=main_x, wkv_linears=linears) + tc.assertEqual(len(got), num_stages) + for kv_got, kv_ref in zip(got, ref): + tc.assertEqual(kv_got.shape, kv_ref.shape) + tc.assertTrue(kv_got.is_contiguous()) + torch.testing.assert_close(kv_got.float(), kv_ref.float(), rtol=2e-2, atol=2e-3) + # fp8 blockwise weight dequant path (2x3 grid of 128x128 blocks). + out_dim, in_dim, block = 192, 384, 128 + w8 = torch.randn(out_dim, in_dim, device=DEVICE, generator=g).to( + torch.float8_e4m3fn + ) + scale = torch.rand(2, 3, device=DEVICE, generator=g) + 0.5 + sf = scale.repeat_interleave(block, 0)[:out_dim] + sf = sf.repeat_interleave(block, 1)[:, :in_dim] + expected = (w8.to(torch.float32) * sf).to(torch.bfloat16) + stub = types.SimpleNamespace(weight=w8, weight_scale_inv=scale) + tc._eq(dspark_draft_model._dequant_linear_weight(stub), expected) + + +def _case_compact_layout(tc): + torch.manual_seed(10) + gamma, t, bs = 5, 6, 64 + verify_lens = _ri(1, t + 1, (bs,), torch.int32) + total = int(verify_lens.sum().item()) + for padded_total in (total, bs * t): # exact and bucket padding + tc._parity( + dspark_verify_window.CompactRowIndex, + verify_lens=verify_lens, + padded_total=padded_total, + device=DEVICE, + ) + tc._parity( + dspark_verify_window.CompactVerifyIds, + draft_block_ids=_ri(0, VOCAB, (bs, gamma)), + draft_tokens=_ri(0, VOCAB, (bs, gamma)), + layout=_layout(verify_lens, padded_total), + device=DEVICE, + ) + + +def _case_swa_page_indices(tc): + torch.manual_seed(11) + block_size, num_q, max_reqs, n_full = 5, 320, 300, 50000 + _, gather = tc._parity( + dspark_attn_metadata.ComputeDsparkWindowGather, + seq_lens_casual=_ri(1, 300, (num_q,), torch.int32), + req_pool_indices_repeated=_ri(0, max_reqs, (num_q,)), + block_size=block_size, + swa_window=128, + ) + tc._parity( + dspark_attn_metadata.BuildDsparkSwaPageIndices, + req_to_token=_ri(0, n_full, (max_reqs, 400), torch.int32), + full_to_swa_mapping=_ri(0, 20000, (n_full,), torch.int32), + req_pool_indices_per_request=gather.req_pool_indices_per_request, + offsets=gather.offsets, + invalid=gather.invalid, + out_loc=_ri(0, n_full, (num_q,)), + context_lens=gather.context_lens, + block_size=block_size, + swa_window=128, + page_index_aligned_size=64, + ) + + +def _case_expand_prefill_causally(tc): + torch.manual_seed(12) + # Vectorized branch: ragged extends with padded token count. + bs = 64 + extend = _ri(1, 8, (bs,)) + num_tokens = int(extend.sum()) + req_pool_indices = torch.randperm(512, device=DEVICE)[:bs] + seq_lens = _ri(8, 500, (bs,)) + tc._parity( + attn_metadata_kernels.ExpandPrefillCausally, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + extend_seq_lens=extend, + extend_start_loc=torch.cumsum(extend, dim=0) - extend, + seq_lens_cpu=None, + extend_seq_lens_cpu=None, + num_tokens=num_tokens, + padded_num_tokens=num_tokens + 5, + ) + # Loop branch: uniform extend with CPU lens and no padding. + bs2, block = 8, 6 + tc._parity( + attn_metadata_kernels.ExpandPrefillCausally, + req_pool_indices=req_pool_indices[:bs2], + seq_lens=seq_lens[:bs2], + extend_seq_lens=torch.full((bs2,), block, device=DEVICE), + extend_start_loc=None, + seq_lens_cpu=[int(x) for x in seq_lens[:bs2].tolist()], + extend_seq_lens_cpu=[block] * bs2, + num_tokens=bs2 * block, + padded_num_tokens=None, + ) + + +def _case_finalize_accept_lens(tc): + torch.manual_seed(13) + bs = 64 + for prefix_dtype in (torch.int32, torch.int64): + tc._parity( + dspark_accept.FinalizeAcceptLens, + correct_len=_ri(0, 7, (bs,), torch.int32), + cap_trim_lens=_ri(0, 4, (bs,)), + prefix_lens=_ri(1, 4000, (bs,), prefix_dtype), + ) + + +def _case_mixed_accept_select(tc): + torch.manual_seed(14) + bs = 64 + # Mixed dtypes between the greedy and sampling lanes. + tc._parity( + dspark_accept.SelectMixedAccept, + greedy_mask=torch.rand(bs, device=DEVICE) < 0.5, + greedy_len=_ri(0, 7, (bs,)), + greedy_bonus=_ri(0, 100000, (bs,)), + greedy_trim=_ri(0, 4, (bs,)), + sampling_len=_ri(0, 7, (bs,), torch.int32), + sampling_bonus=_ri(0, 100000, (bs,)), + sampling_trim=_ri(0, 4, (bs,), torch.int32), + ) + + +def _case_padded_to_bucket(tc): + torch.manual_seed(15) + for bs, padded_bs, graph_num_tokens in ((3, 6, 16), (2, 8, 16), (8, 128, 768)): + verify_lens = _ri(1, 7, (bs,), torch.int32) + if int(verify_lens.sum()) > graph_num_tokens: + verify_lens = torch.ones(bs, dtype=torch.int32, device=DEVICE) + got, _ = tc._parity( + ragged_verify_kernels.PaddedToBucket, + verify_lens=verify_lens, + graph_num_tokens=graph_num_tokens, + bs=bs, + padded_bs=padded_bs, + ) + # Padding rows must absorb exactly the leftover budget. + tc.assertEqual(int(got.to(torch.int64).sum()), graph_num_tokens) + if padded_bs > bs: + tc.assertTrue(torch.equal(got[:bs], verify_lens)) + + +def _case_page_table_positions(tc): + num_pool, pool_len = 128, 4096 + g = torch.Generator(device=DEVICE).manual_seed(16) + req_to_token = _ri(0, 1 << 20, (num_pool, pool_len), torch.int32, g) + # Large page + non-pool-aligned max_seq_len, then page_size 1. + for num_q, page_size, max_seq_len in ((300, 64, 4000), (56, 1, 4096)): + tc._parity( + attn_metadata_kernels.BuildPageTablePositions, + req_to_token=req_to_token, + req_pool_indices_repeated=_ri(0, num_pool, (num_q,), torch.int32, g), + seq_lens_casual=_ri(1, pool_len, (num_q,), torch.int64, g), + max_seq_len=max_seq_len, + page_size=page_size, + swa_window=128, + ) + + +def _case_qo_indptr(tc): + torch.manual_seed(17) + cls = ragged_verify_kernels.BuildQoIndptr + for dtype in (torch.int32, torch.int64): + verify_lens = _ri(1, 8, (129,), dtype) # straddles the 128 block + ref = cls.torch(verify_lens=verify_lens) + got = cls.triton(verify_lens=verify_lens.to(torch.int32)) + tc._eq(got, ref) + # Aliasing regression: the two outputs must not share storage. + vl = torch.tensor([3, 1, 5], device=DEVICE, dtype=torch.int32) + got = cls.triton(verify_lens=vl) + got.extend_start_loc.fill_(-7) + tc.assertEqual(got.qo_indptr[:2].tolist(), [0, 3]) + + +def _case_sample_step_tokens(tc): + torch.manual_seed(18) + cls = dspark_draft_model.SampleStepTokens + # Injected noise makes stochastic sampling exactly comparable. + for vocab, dtype in ((130000, torch.bfloat16), (5003, torch.float32)): + bs = 3 + tc._parity( + cls, + step_logits=(torch.randn(bs, vocab, device=DEVICE) * 4.0).to(dtype), + temperatures=torch.rand(bs, device=DEVICE) + 0.5, + greedy_mask=(torch.arange(bs, device=DEVICE) % 2) == 0, + exp_noise=torch.empty(bs, vocab, device=DEVICE).exponential_(1), + ) + # Greedy tie straddling a triton block boundary picks the smaller index. + logits = torch.zeros(1, 2050, device=DEVICE) + logits[0, 1000] = logits[0, 1100] = 5.0 + tokens = cls.triton( + step_logits=logits, + temperatures=torch.tensor([1.0], device=DEVICE), + greedy_mask=torch.tensor([True], device=DEVICE), + exp_noise=torch.ones(1, 2050, device=DEVICE), + ) + tc.assertEqual(tokens.item(), 1000) + # Non-contiguous strided cropped view must match its contiguous copy. + view = (torch.randn(2, 129536, device=DEVICE) * 4.0)[:, :VOCAB] + tc.assertFalse(view.is_contiguous()) + kw = dict( + temperatures=torch.rand(2, device=DEVICE) + 0.5, + greedy_mask=torch.tensor([True, False], device=DEVICE), + exp_noise=torch.empty(2, VOCAB, device=DEVICE).exponential_(1), + ) + tc._eq( + cls.triton(step_logits=view, **kw), + cls.triton(step_logits=view.contiguous(), **kw), + ) + + +def _case_scatter_compact_to_strided(tc): + torch.manual_seed(19) + t, bs, dim = 6, 8, 4096 + verify_lens = _ri(1, t + 1, (bs,), torch.int32) + total = int(verify_lens.sum().item()) + for graph_num_tokens in (total, bs * t): # exact and bucket padding + compact = torch.randn( + graph_num_tokens, dim, dtype=torch.bfloat16, device=DEVICE + ) + tc._parity( + dspark_verify_window.ScatterCompactToStrided, + compact=compact, + layout=_layout(verify_lens, graph_num_tokens), + fill_value=0.0, + verify_num_draft_tokens=t, + ) + + +def _case_schedule_verify_lens_topk(tc): + torch.manual_seed(20) + gamma, bs = 5, 64 + cfg = DSparkScheduleConfig(gamma=gamma) + cls = dspark_schedule.ScheduleVerifyLensTopk + base = torch.rand(bs, gamma, device=DEVICE) + confidences = ( + torch.full((bs, gamma), 0.5, device=DEVICE), # all-ties + (base * 4).floor() / 4, # coarse quantization + torch.where(base < 0.3, torch.zeros_like(base), base), # invalid zeros + ) + for confidence in confidences: + for budget in (0, 1, 3, 7, 1000): + tc._parity(cls, confidence=confidence, budget=budget, cfg=cfg) + + +def _case_softmax_temp(tc): + g = torch.Generator(device=DEVICE).manual_seed(21) + cls = dspark_accept.SoftmaxTemp + # bf16 logits, non-power-of-two rows_per_request, full vocab. + logits = (torch.randn(56, VOCAB, device=DEVICE, generator=g) * 8.0).to( + torch.bfloat16 + ) + temps = (torch.rand(8, device=DEVICE, generator=g) * 1.5 + 0.05).float() + ref = cls.torch(logits=logits, temperatures=temps, rows_per_request=7) + got = cls.triton(logits=logits, temperatures=temps, rows_per_request=7) + tc.assertEqual(got.dtype, torch.float32) + torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-6) + torch.testing.assert_close( + got.sum(dim=-1), torch.ones_like(got.sum(dim=-1)), rtol=1e-5, atol=1e-5 + ) + # Column-shaped (bs, 1) temperatures. + logits2 = torch.randn(6, 512, device=DEVICE, generator=g).to(torch.bfloat16) + temps2 = (torch.rand(2, 1, device=DEVICE, generator=g) + 0.3).float() + ref2 = cls.torch(logits=logits2, temperatures=temps2, rows_per_request=3) + got2 = cls.triton(logits=logits2, temperatures=temps2, rows_per_request=3) + torch.testing.assert_close(got2, ref2, rtol=1e-5, atol=1e-7) + + +_CASES = [ + ("accept_greedy", _case_accept_greedy), + ("accept_sampling", _case_accept_sampling), + ("build_block_seq_lens_causal", _case_build_block_seq_lens_causal), + ("build_out_tokens", _case_build_out_tokens), + ("build_ragged_verify_window", _case_build_ragged_verify_window), + ("build_step_local", _case_build_step_local), + ("cap_correct_len", _case_cap_correct_len), + ("causal_swa_page_indices", _case_causal_swa_page_indices), + ("commit_inject_layout", _case_commit_inject_layout), + ("commit_kv_proj", _case_commit_kv_proj), + ("compact_layout", _case_compact_layout), + ("swa_page_indices", _case_swa_page_indices), + ("expand_prefill_causally", _case_expand_prefill_causally), + ("finalize_accept_lens", _case_finalize_accept_lens), + ("mixed_accept_select", _case_mixed_accept_select), + ("padded_to_bucket", _case_padded_to_bucket), + ("page_table_positions", _case_page_table_positions), + ("qo_indptr", _case_qo_indptr), + ("sample_step_tokens", _case_sample_step_tokens), + ("scatter_compact_to_strided", _case_scatter_compact_to_strided), + ("schedule_verify_lens_topk", _case_schedule_verify_lens_topk), + ("softmax_temp", _case_softmax_temp), +] + + +class TestDsparkKernelParity(CustomTestCase): + def _eq(self, got, ref): + """Exact comparison of tensors, tuples, and msgspec result structs.""" + if isinstance(ref, tuple): + for g, r in zip(got, ref): + self._eq(g, r) + elif hasattr(ref, "__struct_fields__"): + for name in ref.__struct_fields__: + self._eq(getattr(got, name), getattr(ref, name)) + elif isinstance(ref, torch.Tensor): + self.assertEqual(got.dtype, ref.dtype) + self.assertTrue(torch.equal(got, ref)) + else: + self.assertEqual(got, ref) + + def _parity(self, cls, **kw): + got, ref = cls.triton(**kw), cls.torch(**kw) + self._eq(got, ref) + return got, ref + + def test_all_kernels_triton_matches_torch(self): + for name, case in _CASES: + with self.subTest(kernel=name): + case(self) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_dspark_scheduler.py b/test/registered/spec/dspark/test_dspark_scheduler.py new file mode 100644 index 000000000..1bb77b532 --- /dev/null +++ b/test/registered/spec/dspark/test_dspark_scheduler.py @@ -0,0 +1,558 @@ +import functools +import types +import unittest + +import torch + +from sglang.srt.speculative.dspark_components.dspark_planner import ( + DSparkScheduleConfig, + HostConfidenceBudgetPlanner, + VerifyBudgetDecision, + compute_verify_token_budget, + graph_tier_fill_budget, +) +from sglang.srt.speculative.dspark_components.dspark_sps import ( + SpsAdditiveCostTable, + SpsCostTable, +) +from sglang.srt.speculative.dspark_components.kernels.dspark_schedule import ( + schedule_verify_lens_topk_from_survival, +) +from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") + + +def _flat_table( + steps_per_sec: float = 1.0, max_batch_tokens: int = 4096 +) -> SpsCostTable: + return SpsCostTable( + sample_batch_tokens=[1], + sample_steps_per_sec=[steps_per_sec], + max_batch_tokens=max_batch_tokens, + ) + + +def _cliff_table() -> SpsCostTable: + return SpsCostTable( + sample_batch_tokens=[1, 2, 3, 4, 5, 6, 7, 8], + sample_steps_per_sec=[1.0, 1.0, 1.0, 0.5, 0.45, 0.44, 0.43, 0.42], + max_batch_tokens=64, + ) + + +def _additive_table() -> SpsAdditiveCostTable: + return SpsAdditiveCostTable( + bias_seconds=0.01, + bs_probes=[1, 100], + alpha_seconds=[0.0, 0.05], + m_probes=[1, 200], + theta_seconds=[0.0, 0.02], + ) + + +def _survival_from_confidence(confidence: torch.Tensor) -> torch.Tensor: + return torch.cumprod(confidence, dim=1) + + +def _bruteforce_budget( + *, + history_survival_probs: torch.Tensor, + sps_table: SpsCostTable, + cfg: DSparkScheduleConfig, +) -> int: + num_requests = history_survival_probs.shape[0] + max_len = cfg.resolved_max_verify_len() + candidates = history_survival_probs[:, :max_len].flatten() + candidates = [float(x) for x in candidates.tolist() if float(x) >= cfg.survival_eps] + candidates.sort(reverse=True) + + best_extra, best_theta = 0, float("-inf") + for extra in range(len(candidates) + 1): + tau_star = num_requests + sum(candidates[:extra]) + theta = tau_star * sps_table.lookup(num_requests + extra) + if theta > best_theta: + best_theta, best_extra = theta, extra + return best_extra + + +def schedule_verify_lens_topk_vanilla( + *, + survival_probs: torch.Tensor, + budget: int, + cfg: DSparkScheduleConfig, +) -> torch.Tensor: + cfg.validate() + num_requests, _gamma = survival_probs.shape + max_len = cfg.resolved_max_verify_len() + device = survival_probs.device + + valid_rows = (survival_probs >= cfg.survival_eps).tolist() + survival_rows = survival_probs.to(torch.float64).tolist() + + candidates: list[tuple[float, int, int]] = [] + for request in range(num_requests): + for position in range(min(max_len, len(valid_rows[request]))): + if valid_rows[request][position]: + candidates.append((survival_rows[request][position], position, request)) + + candidates.sort(key=lambda candidate: (-candidate[0], candidate[1], candidate[2])) + + selected_extra = [0] * num_requests + for _survival, _position, request in candidates[: max(int(budget), 0)]: + selected_extra[request] += 1 + + lower_bound = max(cfg.min_verify_len, 1) + verify_lens = [ + min(max(cfg.min_verify_len + extra, lower_bound), max_len) + for extra in selected_extra + ] + return torch.tensor(verify_lens, dtype=torch.int32, device=device) + + +_TOPK_IMPLS = ( + schedule_verify_lens_topk_from_survival, + schedule_verify_lens_topk_vanilla, +) + + +def _for_each_impl(test_method): + @functools.wraps(test_method) + def wrapper(self): + for impl in _TOPK_IMPLS: + with self.subTest(impl=impl.__name__): + test_method(self, impl) + + return wrapper + + +class TestComputeVerifyTokenBudget(CustomTestCase): + def test_budget_argmax_matches_bruteforce_scan_across_sps_cliffs(self): + torch.manual_seed(1) + cfg = DSparkScheduleConfig(gamma=7) + table = _cliff_table() + for _ in range(20): + confidence = torch.rand(3, 7, dtype=torch.float32) * 0.5 + 0.45 + survival = _survival_from_confidence(confidence) + expected = _bruteforce_budget( + history_survival_probs=survival, sps_table=table, cfg=cfg + ) + actual = compute_verify_token_budget( + history_survival_probs=survival, sps_table=table, cfg=cfg + ).budget + self.assertEqual(actual, expected) + + def test_budget_pool_is_independent_of_min_verify_len(self): + survival = torch.tensor([[0.9, 0.8, 0.7, 0.6]], dtype=torch.float32) + cfg_no_min = DSparkScheduleConfig(gamma=4, min_verify_len=0) + cfg_min2 = DSparkScheduleConfig(gamma=4, min_verify_len=2) + table = _flat_table() + budget_no_min = compute_verify_token_budget( + history_survival_probs=survival, sps_table=table, cfg=cfg_no_min + ).budget + budget_min2 = compute_verify_token_budget( + history_survival_probs=survival, sps_table=table, cfg=cfg_min2 + ).budget + # The budget candidate pool spans all positions; the per-request floor + # is enforced later at verify-lens scheduling, not at budget time. + self.assertEqual(budget_no_min, 4) + self.assertEqual(budget_min2, 4) + + def test_budget_drops_candidates_below_survival_eps(self): + survival = torch.tensor([[0.9, 1e-9, 1e-12]], dtype=torch.float32) + cfg = DSparkScheduleConfig(gamma=3, survival_eps=1e-6) + table = _flat_table() + budget = compute_verify_token_budget( + history_survival_probs=survival, sps_table=table, cfg=cfg + ).budget + self.assertLessEqual(budget, 1) + + def test_decision_predicted_step_matches_additive_table_at_budget(self): + survival = torch.tensor([[0.9, 0.8, 0.7, 0.6]], dtype=torch.float32) + cfg = DSparkScheduleConfig(gamma=4) + table = _additive_table() + decision = compute_verify_token_budget( + history_survival_probs=survival, sps_table=table, cfg=cfg + ) + # Reference via the independent scalar interpolation path (step_time), + # not the tensor helper the implementation itself uses. + self.assertAlmostEqual( + decision.predicted_step_seconds, + table.step_time(num_reqs=1, budget=int(decision.budget)), + places=5, + ) + self.assertGreater(decision.predicted_theta, 0.0) + + def test_decision_predicted_step_is_inverse_sps_for_diagonal_table(self): + survival = torch.tensor([[0.9, 0.8, 0.7, 0.6]], dtype=torch.float32) + cfg = DSparkScheduleConfig(gamma=4) + table = _cliff_table() + decision = compute_verify_token_budget( + history_survival_probs=survival, sps_table=table, cfg=cfg + ) + expected_sps = table.lookup(1 + decision.budget) + self.assertIsNotNone(decision.predicted_step_seconds) + self.assertAlmostEqual( + decision.predicted_step_seconds, 1.0 / expected_sps, places=9 + ) + self.assertGreater(decision.predicted_theta, 0.0) + + +def _make_budget_planner() -> HostConfidenceBudgetPlanner: + return HostConfidenceBudgetPlanner( + sps_table=_flat_table(), + cfg=DSparkScheduleConfig(gamma=4), + model_runner=None, + ) + + +class TestBudgetDecisionLifecycle(CustomTestCase): + def test_take_last_decision_is_consume_once(self): + planner = _make_budget_planner() + planner.last_decision = VerifyBudgetDecision( + budget=3, predicted_step_seconds=0.01, predicted_theta=100.0 + ) + first = planner.take_last_decision() + self.assertEqual(first.budget, 3) + self.assertIsNone(planner.take_last_decision()) + + def test_note_non_decode_step_clears_decision(self): + planner = _make_budget_planner() + planner.last_decision = VerifyBudgetDecision(budget=1) + planner.note_non_decode_step() + self.assertIsNone(planner.take_last_decision()) + + +class TestScheduleVerifyLensTopk(CustomTestCase): + @_for_each_impl + def test_topk_does_not_exceed_budget(self, impl): + torch.manual_seed(2) + survival = _survival_from_confidence(torch.rand(5, 7) * 0.4 + 0.55) + cfg = DSparkScheduleConfig(gamma=7) + floor = max(cfg.min_verify_len, 1) + for budget in (0, 1, 5, 12, 100): + verify_lens = impl(survival_probs=survival, budget=budget, cfg=cfg) + total_extra = int((verify_lens.to(torch.int64) - floor).sum().item()) + self.assertLessEqual(total_extra, budget) + self.assertGreaterEqual(int(verify_lens.min().item()), 1) + + @_for_each_impl + def test_total_equals_anchors_plus_lens(self, impl): + survival = torch.tensor( + [[0.90, 0.80, 0.30, 0.20], [0.85, 0.70, 0.25, 0.15]], + dtype=torch.float32, + ) + num_requests, max_len, budget = 2, 4, 2 + cfg = DSparkScheduleConfig(gamma=max_len, min_verify_len=1) + verify_lens = impl(survival_probs=survival, budget=budget, cfg=cfg) + actual_total = num_requests + int(verify_lens.to(torch.int64).sum().item()) + admitted = budget + expected_total = num_requests + num_requests * cfg.min_verify_len + admitted + self.assertEqual(actual_total, expected_total) + + @_for_each_impl + def test_higher_confidence_admitted_first(self, impl): + survival = torch.tensor( + [[0.99, 0.98, 0.97, 0.96], [0.40, 0.30, 0.20, 0.10]], + dtype=torch.float32, + ) + cfg = DSparkScheduleConfig(gamma=4) + verify_lens = impl(survival_probs=survival, budget=2, cfg=cfg) + extra = verify_lens.to(torch.int64) - cfg.min_verify_len + self.assertEqual(int(extra[0].item()), 2) + self.assertEqual(int(extra[1].item()), 0) + + @_for_each_impl + def test_min_and_max_enter_the_budget(self, impl): + survival = torch.tensor([[0.99, 0.99, 0.99, 0.99, 0.99]], dtype=torch.float32) + cfg = DSparkScheduleConfig(gamma=5, min_verify_len=1, max_verify_len=3) + verify_lens = impl(survival_probs=survival, budget=100, cfg=cfg) + self.assertGreaterEqual(int(verify_lens.min().item()), 1) + self.assertLessEqual(int(verify_lens.max().item()), 3) + + @_for_each_impl + def test_large_budget_selects_all_candidates(self, impl): + survival = torch.tensor([[0.9, 0.8, 0.7]], dtype=torch.float32) + cfg = DSparkScheduleConfig(gamma=3) + verify_lens = impl(survival_probs=survival, budget=1000, cfg=cfg) + # anchor (min_verify_len=1) + all 3 admitted drafts + self.assertEqual(int(verify_lens[0].item()), 4) + + @_for_each_impl + def test_tie_break_is_value_independent(self, impl): + survival = torch.tensor([[0.8, 0.8, 0.8], [0.8, 0.8, 0.8]], dtype=torch.float32) + cfg = DSparkScheduleConfig(gamma=3) + floor = max(cfg.min_verify_len, 1) + verify_lens = impl(survival_probs=survival, budget=3, cfg=cfg) + total_extra = int((verify_lens.to(torch.int64) - floor).sum().item()) + self.assertEqual(total_extra, 3) + + +class TestVerifyLenAnchorContract(CustomTestCase): + + @_for_each_impl + def test_explicit_zero_min_still_clamped_to_anchor(self, impl): + survival = _survival_from_confidence( + torch.tensor([[0.9, 0.8, 0.7], [0.6, 0.5, 0.4]], dtype=torch.float32) + ) + cfg = DSparkScheduleConfig(gamma=3, min_verify_len=0) + verify_lens = impl(survival_probs=survival, budget=0, cfg=cfg) + self.assertGreaterEqual(int(verify_lens.min().item()), 1) + self.assertTrue( + torch.equal(verify_lens, torch.tensor([1, 1], dtype=torch.int32)) + ) + + def test_non_flat_table_small_budget_feeds_ragged_layout(self): + table = SpsCostTable( + sample_batch_tokens=[2, 3], + sample_steps_per_sec=[1.0, 0.1], + max_batch_tokens=64, + ) + cfg = DSparkScheduleConfig(gamma=3) + survival = _survival_from_confidence( + torch.tensor([[0.90, 0.80, 0.70], [0.85, 0.60, 0.40]], dtype=torch.float32) + ) + budget = compute_verify_token_budget( + history_survival_probs=survival, sps_table=table, cfg=cfg + ).budget + self.assertEqual(budget, 0) + verify_lens = schedule_verify_lens_topk_from_survival( + survival_probs=survival, budget=budget, cfg=cfg + ) + self.assertGreaterEqual(int(verify_lens.min().item()), 1) + verify_lens_cpu = verify_lens.to(torch.int64).tolist() + layout = RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=verify_lens_cpu, + device=torch.device("cpu"), + grid=[sum(verify_lens_cpu)], + ) + self.assertEqual(layout.verify_lens_cpu, verify_lens_cpu) + + +class TestNonAnticipating(CustomTestCase): + @_for_each_impl + def test_lens_topk_non_anticipating_under_future_perturbation(self, impl): + base = torch.tensor( + [ + [0.95, 0.90, 0.80, 0.40], + [0.92, 0.70, 0.30, 0.10], + [0.99, 0.98, 0.50, 0.05], + ], + dtype=torch.float32, + ) + cfg = DSparkScheduleConfig(gamma=4) + budget = 5 + baseline = impl(survival_probs=base, budget=budget, cfg=cfg) + + request, cut = 1, 2 + for delta in (-0.05, -0.2, 0.05, 0.0): + perturbed = base.clone() + future = perturbed[request, cut:] + perturbed[request, cut:] = torch.clamp( + torch.minimum(future + delta, base[request, cut - 1]), min=0.0 + ) + verify_lens = impl(survival_probs=perturbed, budget=budget, cfg=cfg) + admitted_prefix_unchanged = min(int(baseline[request].item()), cut) == min( + int(verify_lens[request].item()), cut + ) + self.assertTrue( + admitted_prefix_unchanged, + msg=f"prefix admission changed under future perturbation delta={delta}", + ) + + @_for_each_impl + def test_other_requests_unaffected_by_one_request_future(self, impl): + base = torch.tensor( + [[0.95, 0.90, 0.20], [0.93, 0.88, 0.15]], dtype=torch.float32 + ) + cfg = DSparkScheduleConfig(gamma=3) + budget = 2 + baseline = impl(survival_probs=base, budget=budget, cfg=cfg) + perturbed = base.clone() + perturbed[0, 2] = 0.01 + verify_lens = impl(survival_probs=perturbed, budget=budget, cfg=cfg) + self.assertEqual(int(baseline[1].item()), int(verify_lens[1].item())) + + +class TestVanillaMatchesReference(CustomTestCase): + def test_random_inputs_match_reference(self): + torch.manual_seed(20260630) + num_trials = 4000 + for trial in range(num_trials): + num_requests = int(torch.randint(1, 6, ()).item()) + gamma = int(torch.randint(1, 9, ()).item()) + + dtype = torch.float32 if trial % 2 == 0 else torch.float64 + confidence = torch.rand(num_requests, gamma, dtype=dtype) + if trial % 3 == 0: + confidence = (confidence * 4).round() / 4 + if trial % 7 == 0: + confidence = torch.ones(num_requests, gamma, dtype=dtype) + survival = torch.cumprod(confidence, dim=1) + + min_verify_len = int(torch.randint(0, gamma + 1, ()).item()) + if torch.rand(()).item() < 0.5: + max_verify_len = 0 + else: + max_verify_len = int( + torch.randint(min_verify_len, gamma + 1, ()).item() + ) + survival_eps = float( + [1e-6, 1e-3, 0.1, 0.5][int(torch.randint(0, 4, ()).item())] + ) + budget = int(torch.randint(0, num_requests * gamma + 3, ()).item()) + + cfg = DSparkScheduleConfig( + gamma=gamma, + min_verify_len=min_verify_len, + max_verify_len=max_verify_len, + survival_eps=survival_eps, + ) + reference = schedule_verify_lens_topk_from_survival( + survival_probs=survival, budget=budget, cfg=cfg + ) + vanilla = schedule_verify_lens_topk_vanilla( + survival_probs=survival, budget=budget, cfg=cfg + ) + self.assertTrue( + torch.equal(reference, vanilla), + msg=( + f"mismatch on trial {trial}: budget={budget} " + f"min={min_verify_len} max={max_verify_len} eps={survival_eps} " + f"survival={survival.tolist()} " + f"reference={reference.tolist()} vanilla={vanilla.tolist()}" + ), + ) + + +class TestDSparkScheduleConfig(CustomTestCase): + def test_validate_rejects_min_greater_than_max(self): + with self.assertRaises(ValueError): + DSparkScheduleConfig(gamma=4, min_verify_len=3, max_verify_len=2).validate() + + def test_validate_rejects_max_greater_than_gamma_plus_one(self): + with self.assertRaises(ValueError): + DSparkScheduleConfig(gamma=4, max_verify_len=6).validate() + + def test_zero_max_resolves_to_gamma_plus_one(self): + cfg = DSparkScheduleConfig(gamma=7) + self.assertEqual(cfg.resolved_max_verify_len(), 8) + + +class TestGraphTierFillBudget(CustomTestCase): + def test_floor_scales_with_min_verify_len(self): + """The subtracted floor is bs * max(min_verify_len, 1).""" + self.assertEqual( + graph_tier_fill_budget( + graph_num_tokens=60, bs=10, verify_num_draft_tokens=6, min_verify_len=2 + ), + 60 - 20, + ) + + def test_feeding_budget_fills_topk_total_to_tier(self): + """Feeding the fill budget to the top-k lifts the total to min(tier, bs*K).""" + bs = 6 + cfg = DSparkScheduleConfig(gamma=6, min_verify_len=1) + cap = cfg.resolved_max_verify_len() + survival = torch.full((bs, cap), 0.99, dtype=torch.float32) + for graph_num_tokens in (bs, bs * cap // 2, bs * cap, bs * cap + 12): + budget = graph_tier_fill_budget( + graph_num_tokens=graph_num_tokens, + bs=bs, + verify_num_draft_tokens=cap, + min_verify_len=cfg.min_verify_len, + ) + verify_lens = schedule_verify_lens_topk_from_survival( + survival_probs=survival, budget=budget, cfg=cfg + ) + total = int(verify_lens.to(torch.int64).sum().item()) + self.assertEqual(total, min(graph_num_tokens, bs * cap)) + + +class _FakeRaggedRunner(types.SimpleNamespace): + pass + + +def _fake_model_runner(capture_num_tokens, max_bs): + runner = _FakeRaggedRunner( + ragged_verify_mode=True, + capture_num_tokens=capture_num_tokens, + max_bs=max_bs, + ) + return types.SimpleNamespace(decode_cuda_graph_runner=runner) + + +class TestBudgetTierSelection(CustomTestCase): + def test_floor_uses_tier_hint_capped_at_uniform_window(self): + from sglang.srt.speculative.dspark_components.dspark_planner import ( + verify_layout_graph_num_tokens_floor, + ) + from sglang.srt.speculative.ragged_verify import RaggedVerifyMode + + model_runner = _fake_model_runner([8, 16, 1024], max_bs=128) + floor = verify_layout_graph_num_tokens_floor( + num_reqs=100, + ragged_verify_mode=RaggedVerifyMode.COMPACT, + verify_num_draft_tokens=8, + model_runner=model_runner, + tier_num_tokens=150, + ) + self.assertEqual(floor, 150) + capped = verify_layout_graph_num_tokens_floor( + num_reqs=10, + ragged_verify_mode=RaggedVerifyMode.COMPACT, + verify_num_draft_tokens=8, + model_runner=model_runner, + tier_num_tokens=150, + ) + self.assertEqual(capped, 80) + pinned = verify_layout_graph_num_tokens_floor( + num_reqs=100, + ragged_verify_mode=RaggedVerifyMode.COMPACT, + verify_num_draft_tokens=8, + model_runner=model_runner, + ) + self.assertEqual(pinned, 800) + + def test_exceeds_gate_checks_slots_and_tier(self): + from sglang.srt.speculative.dspark_components.dspark_planner import ( + ragged_layout_exceeds_captured_grid, + ) + + model_runner = _fake_model_runner([8, 16, 1024], max_bs=128) + self.assertTrue( + ragged_layout_exceeds_captured_grid( + num_reqs=129, + verify_num_draft_tokens=8, + model_runner=model_runner, + tier_tokens_hint=200, + ) + ) + self.assertFalse( + ragged_layout_exceeds_captured_grid( + num_reqs=128, + verify_num_draft_tokens=8, + model_runner=model_runner, + tier_tokens_hint=512, + ) + ) + self.assertFalse( + ragged_layout_exceeds_captured_grid( + num_reqs=128, + verify_num_draft_tokens=8, + model_runner=model_runner, + ) + ) + self.assertTrue( + ragged_layout_exceeds_captured_grid( + num_reqs=128, + verify_num_draft_tokens=9, + model_runner=model_runner, + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_dspark_sps_profiler.py b/test/registered/spec/dspark/test_dspark_sps_profiler.py new file mode 100644 index 000000000..be191e8a4 --- /dev/null +++ b/test/registered/spec/dspark/test_dspark_sps_profiler.py @@ -0,0 +1,278 @@ +import unittest + +from sglang.benchmark.dspark_sps_profiler import ( + LoadInfo, + ServerContext, + SpsRow, + build_request_count_sweep, + build_table_from_summaries, + count_aligned_steps, + postprocess_round, + resolve_cuda_graph_max_bs, + round_summary_dict, + validate_sweep_against_server, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") + + +def make_load_info() -> LoadInfo: + return LoadInfo( + num_requests=4, max_new_tokens=1200, wall_seconds=1.0, reached_target=True + ) + + +def make_rows( + *, + num_rows: int = 30, + num_running_reqs: int = 4, + num_verify_tokens: int = 32, + step_time: float = 0.01, + first_forward_ct: int = 0, +) -> list[SpsRow]: + return [ + SpsRow( + forward_ct=first_forward_ct + index, + num_running_reqs=num_running_reqs, + num_verify_tokens=num_verify_tokens, + step_time=step_time, + ) + for index in range(num_rows) + ] + + +def make_context(**overrides) -> ServerContext: + values = dict( + base_url="http://localhost:30000", + tokenizer_path="dummy", + tp_size=4, + dp_size=1, + verify_num_draft_tokens=8, + simulate_acc_len=1.0, + cuda_graph_max_bs=128, + skip_max_running_requests_threshold=float("inf"), + skip_token_capacity_threshold=float("inf"), + ) + values.update(overrides) + return ServerContext(**values) + + +class TestPostprocessRound(CustomTestCase): + def test_single_rank_round_builds_probe_from_median_step_time(self): + outcome = postprocess_round( + rank_rows=[make_rows(step_time=0.01)], + batch_size_per_rank=4, + dp_size=1, + verify_num_draft_tokens=8, + min_steady_steps=16, + load_info=make_load_info(), + ) + self.assertEqual(outcome.batch_tokens, 32) + self.assertAlmostEqual(outcome.steps_per_sec, 100.0) + self.assertEqual(outcome.match_fraction, 1.0) + + def test_round_warmup_steps_are_dropped_from_timing(self): + slow_head = make_rows(num_rows=8, step_time=0.5, first_forward_ct=0) + steady_tail = make_rows(num_rows=20, step_time=0.01, first_forward_ct=8) + outcome = postprocess_round( + rank_rows=[slow_head + steady_tail], + batch_size_per_rank=4, + dp_size=1, + verify_num_draft_tokens=8, + min_steady_steps=16, + load_info=make_load_info(), + ) + self.assertAlmostEqual(outcome.steps_per_sec, 100.0) + + def test_off_target_batch_rows_are_filtered_out(self): + ramp = make_rows(num_rows=10, num_running_reqs=2, num_verify_tokens=16) + steady = make_rows(num_rows=30, first_forward_ct=10, step_time=0.02) + outcome = postprocess_round( + rank_rows=[ramp + steady], + batch_size_per_rank=4, + dp_size=1, + verify_num_draft_tokens=8, + min_steady_steps=16, + load_info=make_load_info(), + ) + self.assertAlmostEqual(outcome.steps_per_sec, 50.0) + self.assertAlmostEqual(outcome.match_fraction, 1.0) + + def test_mid_round_instability_raises(self): + head = make_rows(num_rows=15) + gap = make_rows( + num_rows=40, num_running_reqs=3, num_verify_tokens=24, first_forward_ct=15 + ) + tail = make_rows(num_rows=15, first_forward_ct=55) + with self.assertRaisesRegex(RuntimeError, "unstable mid-round"): + postprocess_round( + rank_rows=[head + gap + tail], + batch_size_per_rank=4, + dp_size=1, + verify_num_draft_tokens=8, + min_steady_steps=16, + load_info=make_load_info(), + ) + + def test_round_that_never_stabilizes_raises(self): + rows = make_rows(num_rows=50, num_running_reqs=3, num_verify_tokens=24) + rows += make_rows(num_rows=2, first_forward_ct=50) + with self.assertRaisesRegex(RuntimeError, "never stabilized"): + postprocess_round( + rank_rows=[rows], + batch_size_per_rank=4, + dp_size=1, + verify_num_draft_tokens=8, + min_steady_steps=16, + load_info=make_load_info(), + ) + + +class TestPostprocessRoundCrossRank(CustomTestCase): + def test_two_uniform_ranks_average_their_step_times(self): + outcome = postprocess_round( + rank_rows=[make_rows(step_time=0.01), make_rows(step_time=0.03)], + batch_size_per_rank=4, + dp_size=2, + verify_num_draft_tokens=8, + min_steady_steps=16, + load_info=make_load_info(), + ) + self.assertEqual(outcome.batch_size_per_rank, 4) + self.assertEqual(outcome.batch_tokens, 32) + self.assertAlmostEqual(outcome.steps_per_sec, 50.0) + self.assertEqual(len(outcome.per_rank_median_step_time), 2) + self.assertAlmostEqual(outcome.per_rank_median_step_time[0], 0.01) + self.assertAlmostEqual(outcome.per_rank_median_step_time[1], 0.03) + + def test_rank_with_no_new_records_raises(self): + with self.assertRaisesRegex(RuntimeError, "no new decode-step records"): + postprocess_round( + rank_rows=[make_rows(), []], + batch_size_per_rank=4, + dp_size=2, + verify_num_draft_tokens=8, + min_steady_steps=16, + load_info=make_load_info(), + ) + + def test_disjoint_forward_ct_ranges_raise(self): + with self.assertRaisesRegex(RuntimeError, "no common forward_ct"): + postprocess_round( + rank_rows=[ + make_rows(first_forward_ct=0), + make_rows(first_forward_ct=1000), + ], + batch_size_per_rank=4, + dp_size=2, + verify_num_draft_tokens=8, + min_steady_steps=16, + load_info=make_load_info(), + ) + + def test_rank_below_expected_verify_tokens_raises(self): + # A rank reporting fewer verify tokens than bs_per_rank * K is not + # running the uniform static verify; above-expected counts are + # tolerated (the recorded count is the replayed graph tier). + with self.assertRaisesRegex(RuntimeError, "num_verify_tokens"): + postprocess_round( + rank_rows=[make_rows(), make_rows(num_verify_tokens=24)], + batch_size_per_rank=4, + dp_size=2, + verify_num_draft_tokens=8, + min_steady_steps=16, + load_info=make_load_info(), + ) + + def test_rank_count_mismatch_raises(self): + with self.assertRaisesRegex(RuntimeError, "DP ranks"): + postprocess_round( + rank_rows=[make_rows()], + batch_size_per_rank=4, + dp_size=2, + verify_num_draft_tokens=8, + min_steady_steps=16, + load_info=make_load_info(), + ) + + +class TestTableAssembly(CustomTestCase): + def test_repeats_take_the_median_per_batch_tokens(self): + rounds = [ + postprocess_round( + rank_rows=[make_rows(step_time=step_time)], + batch_size_per_rank=4, + dp_size=1, + verify_num_draft_tokens=8, + min_steady_steps=16, + load_info=make_load_info(), + ) + for step_time in (0.01, 0.02, 0.04) + ] + table = build_table_from_summaries( + summaries=[ + round_summary_dict(outcome=outcome, repeat=repeat) + for repeat, outcome in enumerate(rounds) + ], + max_batch_tokens=None, + offdiag=False, + ) + self.assertEqual(table.sample_batch_tokens, [32]) + self.assertAlmostEqual(table.sample_steps_per_sec[0], 50.0) + + +class TestSweepHelpers(CustomTestCase): + def test_request_count_sweep_tapers_and_hits_the_max(self): + sweep = build_request_count_sweep(100) + self.assertEqual(sweep[:4], [1, 2, 4, 8]) + self.assertEqual(sweep[-1], 100) + self.assertIn(64, sweep) + + def test_sweep_beyond_captured_cuda_graphs_raises(self): + with self.assertRaisesRegex(ValueError, "cuda graphs"): + validate_sweep_against_server( + context=make_context(cuda_graph_max_bs=64), + batch_sizes=[8, 128], + ) + + def test_sweep_within_captured_cuda_graphs_passes(self): + validate_sweep_against_server( + context=make_context(cuda_graph_max_bs=64, dp_size=2), + batch_sizes=[8, 64], + ) + + def test_resolve_cuda_graph_max_bs_prefers_captured_list(self): + internal_state = { + "cuda_graph_config": {"decode": {"bs": [1, 2, 160], "max_bs": 128}} + } + self.assertEqual(resolve_cuda_graph_max_bs(internal_state=internal_state), 160) + + def test_resolve_cuda_graph_max_bs_handles_missing_config(self): + self.assertIsNone(resolve_cuda_graph_max_bs(internal_state={})) + + +class TestCountAlignedSteps(CustomTestCase): + def test_off_target_steps_are_not_counted(self): + rows = make_rows(num_rows=10, num_running_reqs=3) + self.assertEqual( + count_aligned_steps(rank_rows=[rows], batch_size_per_rank=4), 0 + ) + + +class TestMinSteadySteps(CustomTestCase): + def test_min_steady_steps_rejects_thin_probes(self): + with self.assertRaisesRegex(RuntimeError, "never stabilized"): + postprocess_round( + rank_rows=[make_rows(num_rows=20)], + batch_size_per_rank=4, + dp_size=1, + verify_num_draft_tokens=8, + min_steady_steps=32, + load_info=make_load_info(), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_dspark_sps_table.py b/test/registered/spec/dspark/test_dspark_sps_table.py new file mode 100644 index 000000000..ac13c3dc2 --- /dev/null +++ b/test/registered/spec/dspark/test_dspark_sps_table.py @@ -0,0 +1,200 @@ +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +from sglang.srt.speculative.dspark_components.dspark_sps import ( + SpsAdditiveCostTable, + SpsCostTable, + build_uninitialized_sps_table, + is_uninitialized_sps_table, + load_sps_table_from_path, + profile_sps_table, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +def _make_table() -> SpsCostTable: + return SpsCostTable( + sample_batch_tokens=[8, 16, 32, 64], + sample_steps_per_sec=[1000.0, 950.0, 500.0, 480.0], + max_batch_tokens=128, + ) + + +class TestSpsCostTableInvariants(CustomTestCase): + def test_rejects_non_increasing_batch_tokens(self): + with self.assertRaises(ValueError): + SpsCostTable( + sample_batch_tokens=[8, 8, 16], + sample_steps_per_sec=[1.0, 2.0, 3.0], + max_batch_tokens=16, + ) + + def test_rejects_unsorted_batch_tokens(self): + with self.assertRaises(ValueError): + SpsCostTable( + sample_batch_tokens=[16, 8], + sample_steps_per_sec=[1.0, 2.0], + max_batch_tokens=16, + ) + + def test_rejects_length_mismatch(self): + with self.assertRaises(ValueError): + SpsCostTable( + sample_batch_tokens=[8, 16], + sample_steps_per_sec=[1.0], + max_batch_tokens=16, + ) + + def test_rejects_empty_table(self): + with self.assertRaises(ValueError): + SpsCostTable( + sample_batch_tokens=[], + sample_steps_per_sec=[], + max_batch_tokens=0, + ) + + def test_rejects_max_below_largest_probe(self): + with self.assertRaises(ValueError): + SpsCostTable( + sample_batch_tokens=[8, 16], + sample_steps_per_sec=[1.0, 2.0], + max_batch_tokens=15, + ) + + +class TestSpsCostTableLookup(CustomTestCase): + def test_lookup_exact_probe_returns_that_sps(self): + table = _make_table() + self.assertEqual(table.lookup(8), 1000.0) + self.assertEqual(table.lookup(16), 950.0) + self.assertEqual(table.lookup(32), 500.0) + self.assertEqual(table.lookup(64), 480.0) + + def test_lookup_floors_to_lower_captured_probe(self): + table = _make_table() + self.assertEqual(table.lookup(31), 950.0) + self.assertEqual(table.lookup(63), 500.0) + + def test_lookup_below_first_probe_clamps_to_first(self): + table = _make_table() + self.assertEqual(table.lookup(1), 1000.0) + self.assertEqual(table.lookup(7), 1000.0) + + def test_lookup_above_last_probe_clamps_to_last(self): + table = _make_table() + self.assertEqual(table.lookup(65), 480.0) + self.assertEqual(table.lookup(10_000), 480.0) + + +class TestLoadSpsTableFromPath(CustomTestCase): + def test_load_from_path_round_trips_table_and_lookup(self): + table = _make_table() + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "sps.json" + path.write_text(table.to_json(), encoding="utf-8") + loaded = load_sps_table_from_path(str(path)) + self.assertEqual(loaded.sample_batch_tokens, table.sample_batch_tokens) + self.assertEqual(loaded.sample_steps_per_sec, table.sample_steps_per_sec) + self.assertEqual(loaded.max_batch_tokens, table.max_batch_tokens) + for batch_tokens in (1, 8, 31, 64, 200): + self.assertEqual(loaded.lookup(batch_tokens), table.lookup(batch_tokens)) + + +class TestFlatTableLookupIsConstant(CustomTestCase): + def test_flat_table_lookup_is_one_for_any_batch(self): + flat = SpsCostTable( + sample_batch_tokens=[1], + sample_steps_per_sec=[1.0], + max_batch_tokens=4096, + ) + for batch_tokens in (0, 1, 2, 17, 256, 100_000): + self.assertEqual(flat.lookup(batch_tokens), 1.0) + + +class TestProfileSpsTable(CustomTestCase): + def test_profile_sorts_out_of_order_probes(self): + table = profile_sps_table( + probes=[(32, 500.0), (8, 1000.0), (16, 950.0)], + ) + self.assertEqual(table.sample_batch_tokens, [8, 16, 32]) + self.assertEqual(table.sample_steps_per_sec, [1000.0, 950.0, 500.0]) + + def test_profile_rejects_duplicate_batch_tokens(self): + with self.assertRaises(ValueError): + profile_sps_table(probes=[(8, 1000.0), (8, 900.0)]) + + def test_profile_rejects_empty_probes(self): + with self.assertRaises(ValueError): + profile_sps_table(probes=[]) + + def test_profile_max_batch_tokens_defaults_to_largest_probe(self): + table = profile_sps_table(probes=[(8, 1000.0), (64, 480.0), (16, 950.0)]) + self.assertEqual(table.max_batch_tokens, 64) + + def test_profile_honors_explicit_max_batch_tokens(self): + table = profile_sps_table( + probes=[(8, 1000.0), (16, 950.0)], max_batch_tokens=256 + ) + self.assertEqual(table.max_batch_tokens, 256) + + +def _build_sps_cost_table_for(*, sps_table_path): + from sglang.srt.speculative.dspark_components.dspark_planner import ( + build_sps_cost_table, + ) + + server_args = SimpleNamespace( + speculative_dspark_sps_table_path=sps_table_path, + max_running_requests=4, + ) + return build_sps_cost_table(server_args=server_args, verify_num_draft_tokens=5) + + +class TestBuildSpsCostTableContract(CustomTestCase): + def test_unset_table_path_returns_flat_table(self): + for sps_table_path in (None, ""): + table = _build_sps_cost_table_for(sps_table_path=sps_table_path) + self.assertEqual(table.sample_batch_tokens, [1]) + self.assertEqual(table.sample_steps_per_sec, [1.0]) + self.assertEqual(table.max_batch_tokens, 20) + + def test_real_path_loads_table(self): + table = _make_table() + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "sps.json" + path.write_text(table.to_json(), encoding="utf-8") + loaded = _build_sps_cost_table_for(sps_table_path=str(path)) + self.assertEqual(loaded.sample_batch_tokens, table.sample_batch_tokens) + self.assertEqual(loaded.sample_steps_per_sec, table.sample_steps_per_sec) + self.assertEqual(loaded.max_batch_tokens, table.max_batch_tokens) + + +class TestIsUninitializedSpsTable(CustomTestCase): + def test_additive_table_is_never_uninitialized(self): + table = SpsAdditiveCostTable( + bias_seconds=0.1, + bs_probes=[128, 192, 256], + alpha_seconds=[0.0, 0.008, 0.016], + m_probes=[384, 512, 1024], + theta_seconds=[0.0, 0.02, 0.1], + ) + self.assertFalse(is_uninitialized_sps_table(table)) + + def test_placeholder_diagonal_table_is_uninitialized(self): + self.assertTrue( + is_uninitialized_sps_table( + build_uninitialized_sps_table(max_batch_tokens=128) + ) + ) + + def test_real_diagonal_table_is_initialized(self): + self.assertFalse(is_uninitialized_sps_table(_make_table())) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_dspark_sts.py b/test/registered/spec/dspark/test_dspark_sts.py new file mode 100644 index 000000000..2d32a9bef --- /dev/null +++ b/test/registered/spec/dspark/test_dspark_sts.py @@ -0,0 +1,142 @@ +import tempfile +import unittest +from pathlib import Path + +import torch + +from sglang.benchmark.dspark_sts_fit import ( + default_temperature_grid, + expected_calibration_error, + fit_sts_temperatures, +) +from sglang.srt.models.dspark import DSparkConfidenceHead +from sglang.srt.speculative.dspark_components.dspark_sts import ( + DSparkStsCalibration, + StsDataRecorder, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class TestApplySts(CustomTestCase): + def test_default_buffer_is_identity_sigmoid(self): + head = DSparkConfidenceHead(hidden_size=8, markov_rank=4, with_markov=False) + confidence_raw = torch.randn(3, 5) * 9.0 + out = head.apply_sts(confidence_raw) + self.assertTrue(torch.equal(out, torch.sigmoid(confidence_raw.float()))) + + def test_per_position_temperature_scales_each_column(self): + head = DSparkConfidenceHead(hidden_size=8, markov_rank=4, with_markov=False) + head.sts_temperatures = torch.tensor([0.5, 1.0, 2.0]) + confidence_raw = torch.full((2, 3), 2.0) + out = head.apply_sts(confidence_raw) + # Hand-computed sigmoid(2.0 / T) per column: identical raw logits with + # distinct per-column T catch wrong-axis broadcasts, and dividing (not + # multiplying) by T is what separates 0.982 from 0.731 in column 0. + expected_row = [0.98201379, 0.88079708, 0.73105858] + for row in out.tolist(): + for got, want in zip(row, expected_row): + self.assertAlmostEqual(got, want, places=6) + + def test_apply_sts_stashes_raw_logit(self): + head = DSparkConfidenceHead(hidden_size=8, markov_rank=4, with_markov=False) + confidence_raw = torch.randn(2, 5) + head.apply_sts(confidence_raw) + self.assertIs(head._last_confidence_raw, confidence_raw) + + +class TestDSparkStsCalibration(CustomTestCase): + def test_json_round_trip_preserves_fields(self): + calibration = DSparkStsCalibration( + temperatures=[1.5, 2.0, 0.5], + dataset="shards.*.pt", + num_samples=1234, + ece_before=[0.3, 0.2, 0.1], + ece_after=[0.02, 0.01, 0.03], + ) + restored = DSparkStsCalibration.from_json(calibration.to_json()) + self.assertEqual(restored.temperatures, calibration.temperatures) + self.assertEqual(restored.dataset, calibration.dataset) + self.assertEqual(restored.num_samples, calibration.num_samples) + self.assertEqual(restored.ece_before, calibration.ece_before) + self.assertEqual(restored.ece_after, calibration.ece_after) + + def test_rejects_empty_temperatures(self): + with self.assertRaises(ValueError): + DSparkStsCalibration(temperatures=[]) + + def test_rejects_non_positive_temperature(self): + with self.assertRaises(ValueError): + DSparkStsCalibration(temperatures=[1.0, 0.0, 2.0]) + with self.assertRaises(ValueError): + DSparkStsCalibration(temperatures=[1.0, -0.5]) + + +class TestExpectedCalibrationError(CustomTestCase): + def test_perfectly_calibrated_probs_have_low_ece(self): + torch.manual_seed(0) + probs = torch.full((20000,), 0.3, dtype=torch.float64) + targets = (torch.rand(20000) < 0.3).to(torch.float64) + ece = expected_calibration_error(probs=probs, targets=targets, num_bins=15) + self.assertLess(ece, 0.02) + + def test_overconfident_probs_have_high_ece(self): + probs = torch.full((20000,), 0.95, dtype=torch.float64) + targets = torch.full((20000,), 0.3, dtype=torch.float64) + ece = expected_calibration_error(probs=probs, targets=targets, num_bins=15) + self.assertGreater(ece, 0.5) + + +class TestFitStsTemperatures(CustomTestCase): + def test_recovers_scale_and_reduces_ece(self): + torch.manual_seed(0) + num_samples, gamma, scale = 60000, 4, 2.5 + base_logit = torch.tensor([2.0, 1.2, 0.8, 0.4]) + true_logit = base_logit[None, :] + torch.randn(num_samples, gamma) * 0.5 + true_prob = torch.sigmoid(true_logit) + accept = (torch.rand(num_samples, gamma) < true_prob).to(torch.float64) + prefix_mask = torch.cumprod(accept, dim=1) + overconfident_logits = true_logit * scale + + result = fit_sts_temperatures( + logits=overconfident_logits, + prefix_mask=prefix_mask, + grid=default_temperature_grid(), + num_bins=15, + ) + + self.assertEqual(len(result["temperatures"]), gamma) + for temperature in result["temperatures"]: + self.assertGreater(temperature, scale / 1.5) + self.assertLess(temperature, scale * 1.5) + mean_before = sum(result["ece_before"]) / gamma + mean_after = sum(result["ece_after"]) / gamma + self.assertLess(mean_after, 0.25 * mean_before) + + +class TestStsDataRecorder(CustomTestCase): + def test_builds_prefix_mask_and_writes_shard(self): + gamma = 4 + confidence_raw = torch.randn(4, gamma) + num_correct_drafts = torch.tensor([0, 2, 4, 1], dtype=torch.int32) + expected_prefix_mask = torch.tensor( + [[0, 0, 0, 0], [1, 1, 0, 0], [1, 1, 1, 1], [1, 0, 0, 0]], + dtype=torch.float32, + ) + with tempfile.TemporaryDirectory() as tmp: + stem = str(Path(tmp) / "shard") + recorder = StsDataRecorder(path_stem=stem, gamma=gamma, flush_every=10) + recorder.record( + confidence_raw=confidence_raw, + num_correct_drafts=num_correct_drafts, + ) + recorder.flush() + shard = torch.load(f"{stem}.0.pt") + self.assertTrue(torch.equal(shard["prefix_mask"], expected_prefix_mask)) + self.assertTrue(torch.equal(shard["logits"], confidence_raw.to(torch.float32))) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_ragged_verify.py b/test/registered/spec/dspark/test_ragged_verify.py new file mode 100644 index 000000000..95a713b05 --- /dev/null +++ b/test/registered/spec/dspark/test_ragged_verify.py @@ -0,0 +1,132 @@ +import unittest + +import torch + +from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyLayout, + build_ragged_target_verify_geometry, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +_DEVICE = torch.device("cpu") +_GRID = [8, 16, 24, 32, 64] + +# The backend capability checks (supports_ragged_verify_graph) live in +# test_ragged_verify_backend_capability.py: importing the backend modules +# pulls GPU-only wheels, which fail to import on the CPU runners. + + +class TestRaggedTargetVerifyGeometry(CustomTestCase): + def test_mixed_verify_lens_geometry(self): + layout = RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=[8, 1, 3], device=_DEVICE, grid=_GRID + ) + seq_lens = torch.tensor([10, 20, 30], dtype=torch.int32) + geometry = build_ragged_target_verify_geometry(seq_lens=seq_lens, layout=layout) + self.assertEqual(geometry.cache_seqlens_int32.tolist(), [18, 21, 33]) + self.assertEqual(geometry.cu_seqlens_q.tolist(), [0, 8, 9, 12]) + self.assertEqual(geometry.cu_seqlens_k.tolist(), [0, 18, 39, 72]) + self.assertEqual(geometry.max_seq_len_q, 8) + + def test_geometry_dtypes_are_int32(self): + layout = RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=[8, 1, 3], device=_DEVICE, grid=_GRID + ) + seq_lens = torch.tensor([10, 20, 30], dtype=torch.int64) + geometry = build_ragged_target_verify_geometry(seq_lens=seq_lens, layout=layout) + self.assertEqual(geometry.cache_seqlens_int32.dtype, torch.int32) + self.assertEqual(geometry.cu_seqlens_q.dtype, torch.int32) + self.assertEqual(geometry.cu_seqlens_k.dtype, torch.int32) + + +class TestPaddedRaggedVerifyGeometry(CustomTestCase): + def test_padded_layout_grows_bs_and_fills_bucket(self): + raw = RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=[8, 1, 3], + device=_DEVICE, + grid=[8, 16, 32, 64], + graph_num_tokens_floor=24, + ) + self.assertEqual(raw.graph_num_tokens, 32) + padded = raw.padded_to_bucket(padded_bs=4) + self.assertEqual(padded.bs, 4) + self.assertEqual(padded.verify_lens.tolist(), [8, 1, 3, 20]) + self.assertEqual(padded.qo_indptr_device.tolist(), [0, 8, 9, 12, 32]) + seq_lens = torch.tensor([10, 20, 30, 1], dtype=torch.int32) + geometry = build_ragged_target_verify_geometry(seq_lens=seq_lens, layout=padded) + self.assertEqual(geometry.cu_seqlens_q.tolist(), [0, 8, 9, 12, 32]) + self.assertEqual(geometry.cache_seqlens_int32.tolist(), [18, 21, 33, 21]) + self.assertEqual(int(geometry.cu_seqlens_k[-1]), 18 + 21 + 33 + 21) + + def test_padded_layout_decoupled_slots_spread_slack(self): + raw = RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=[8, 1, 3], + device=_DEVICE, + grid=[8, 16, 32, 64], + graph_num_tokens_floor=24, + ) + padded = raw.padded_to_bucket(padded_bs=6) + self.assertEqual(padded.bs, 6) + self.assertEqual(padded.verify_lens.tolist(), [8, 1, 3, 7, 7, 6]) + self.assertEqual(int(padded.qo_indptr_device[-1]), 32) + + def test_padded_layout_budget_tier_below_uniform(self): + raw = RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=[8, 1, 3], + device=_DEVICE, + grid=[8, 16, 32, 64], + ) + self.assertEqual(raw.graph_num_tokens, 16) + padded = raw.padded_to_bucket(padded_bs=3) + self.assertEqual(padded.verify_lens.tolist(), [8, 1, 7]) + self.assertEqual(int(padded.qo_indptr_device[-1]), 16) + + def test_padded_layout_zero_len_pad_rows(self): + raw = RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=[8, 8], + device=_DEVICE, + grid=[8, 16, 32, 64], + ) + self.assertEqual(raw.graph_num_tokens, 16) + padded = raw.padded_to_bucket(padded_bs=8) + self.assertEqual(padded.verify_lens.tolist(), [8, 8, 0, 0, 0, 0, 0, 0]) + self.assertEqual(int(padded.qo_indptr_device[-1]), 16) + + +class TestCaptureVerifyLens(CustomTestCase): + def test_small_tier_one_token_rows(self): + from sglang.srt.speculative.ragged_verify import build_capture_verify_lens + + lens = build_capture_verify_lens(num_tokens=8, num_slots=8, num_draft_tokens=8) + self.assertEqual(lens, [1] * 8) + + def test_large_tier_spreads_within_window(self): + from sglang.srt.speculative.ragged_verify import build_capture_verify_lens + + lens = build_capture_verify_lens( + num_tokens=1024, num_slots=128, num_draft_tokens=8 + ) + self.assertEqual(sum(lens), 1024) + self.assertEqual(lens, [8] * 128) + + def test_uneven_tier_rows_stay_legal(self): + from sglang.srt.speculative.ragged_verify import build_capture_verify_lens + + lens = build_capture_verify_lens(num_tokens=24, num_slots=5, num_draft_tokens=8) + self.assertEqual(sum(lens), 24) + self.assertTrue(all(1 <= v <= 8 for v in lens)) + + def test_rejects_overpacked_tier(self): + from sglang.srt.speculative.ragged_verify import build_capture_verify_lens + + with self.assertRaises(ValueError): + build_capture_verify_lens(num_tokens=64, num_slots=4, num_draft_tokens=8) + with self.assertRaises(ValueError): + build_capture_verify_lens(num_tokens=4, num_slots=8, num_draft_tokens=8) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_ragged_verify_backend_capability.py b/test/registered/spec/dspark/test_ragged_verify_backend_capability.py new file mode 100644 index 000000000..b0257ae28 --- /dev/null +++ b/test/registered/spec/dspark/test_ragged_verify_backend_capability.py @@ -0,0 +1,43 @@ +"""Backend opt-in flags for the ragged-verify graphs. + +Runs in the GPU suite because importing the backend modules pulls GPU-only +wheels (sgl_kernel) at module scope, which fail to import on CPU runners. +""" + +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") + + +class TestRaggedVerifyGraphCapability(CustomTestCase): + def test_base_backend_defaults_false(self): + from sglang.srt.layers.attention.base_attn_backend import AttentionBackend + + self.assertFalse(AttentionBackend.supports_ragged_verify_graph) + + def test_ragged_implementing_backends_declare_the_flag(self): + """Every backend with a ragged-verify metadata path must opt in; a + dropped flag silently disables ragged graphs for that backend (the + runner falls back to eager with no other test going red).""" + from sglang.srt.layers.attention.deepseek_v4_backend import ( + DeepseekV4AttnBackend, + ) + from sglang.srt.layers.attention.flashattention_backend import ( + FlashAttentionBackend, + ) + from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend + + for backend in ( + TRTLLMHAAttnBackend, + DeepseekV4AttnBackend, + FlashAttentionBackend, + ): + with self.subTest(backend=backend.__name__): + self.assertTrue(backend.supports_ragged_verify_graph) + + +if __name__ == "__main__": + unittest.main() 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 b13cd3519..0507b350e 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 @@ -93,6 +93,8 @@ def _make_result(num_draft_tokens, accept_lens, flat_tokens): 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, ) diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index da8faa4f9..12088beb6 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -121,6 +121,7 @@ def _make_model_runner( spec.is_eagle.return_value = False spec.is_standalone.return_value = False spec.is_dflash.return_value = False + spec.is_dflash_family.return_value = False spec.is_none.return_value = True mr.spec_algorithm = spec diff --git a/test/registered/unit/sampling/test_sampling_batch_info.py b/test/registered/unit/sampling/test_sampling_batch_info.py index 46ed180b9..0ff0efbd8 100644 --- a/test/registered/unit/sampling/test_sampling_batch_info.py +++ b/test/registered/unit/sampling/test_sampling_batch_info.py @@ -30,6 +30,7 @@ def _make_info(batch_size=2, **overrides): top_ks=torch.full((batch_size,), TOP_K_ALL, dtype=torch.int32), min_ps=torch.zeros(batch_size), is_all_greedy=False, + is_any_greedy=False, need_top_p_sampling=False, need_top_k_sampling=False, need_min_p_sampling=False,