diff --git a/benchmarks/analyze_dflash2_selector_alignment.py b/benchmarks/analyze_dflash2_selector_alignment.py new file mode 100644 index 0000000000..b332106835 --- /dev/null +++ b/benchmarks/analyze_dflash2_selector_alignment.py @@ -0,0 +1,517 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Audit and sweep exact DFlash2 selector proposals from alignment dumps. + +This tool consumes the B1 records emitted by the GPU-runner compact rejection +path when ``VLLM_SPEC_DUMP_ALIGNMENT=1``. It never changes target sampling. +Counterfactual results are one-step overlap proxies on the recorded prefixes; +only an end-to-end run can establish a new acceptance-length result. +""" + +from __future__ import annotations + +import argparse +import glob +import json +import math +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import torch + + +@dataclass(frozen=True) +class ProposalConfig: + name: str + proposal_temperature_scale: float = 1.0 + proposal_top_p: float = 1.0 + unary_scale: float = 1.0 + edge_scale: float = 1.0 + future_beta: float = 0.0 + greedy_mix: float = 0.0 + use_cached_logits: bool = False + + +@dataclass +class AlignmentRecord: + path: Path + step: int + temperature: float + top_p: float + target_topk_ids: torch.Tensor + target_topk_logits: torch.Tensor + candidate_ids: torch.Tensor + realized_logits: torch.Tensor + unary_logits: torch.Tensor + lattice_scores: torch.Tensor + draft_sampled: torch.Tensor + num_sampled: int + + +def _compact_probs(logits: torch.Tensor, top_p: float) -> torch.Tensor: + """Apply the compact top-p contract used by the sparse rejection kernel.""" + logits = logits.to(torch.float64) + order = torch.argsort(logits, descending=True, stable=True) + sorted_probs = torch.softmax(logits[order], dim=-1) + cumulative_before = torch.cumsum(sorted_probs, dim=-1) - sorted_probs + keep_sorted = (top_p >= 1.0) | (cumulative_before < top_p) + kept = torch.zeros_like(sorted_probs) + kept[keep_sorted] = sorted_probs[keep_sorted] + kept /= kept.sum() + probs = torch.zeros_like(kept) + probs[order] = kept + return probs + + +def _distribution_overlap( + target_ids: torch.Tensor, + target_probs: torch.Tensor, + proposal_ids: torch.Tensor, + proposal_probs: torch.Tensor, +) -> float: + matches = target_ids[:, None] == proposal_ids[None, :] + proposal_on_target = torch.where( + matches, + proposal_probs[None, :], + torch.zeros((), dtype=proposal_probs.dtype), + ).sum(dim=1) + return float(torch.minimum(target_probs, proposal_on_target).sum()) + + +def _support_mass( + target_ids: torch.Tensor, + target_probs: torch.Tensor, + proposal_ids: torch.Tensor, +) -> float: + supported = (target_ids[:, None] == proposal_ids[None, :]).any(dim=1) + return float(target_probs[supported].sum()) + + +def _transformed_lattice( + record: AlignmentRecord, config: ProposalConfig +) -> torch.Tensor: + unary = record.unary_logits[:, None, :].to(torch.float64) + lattice = record.lattice_scores.to(torch.float64) + edge = lattice - unary + denominator = record.temperature * config.proposal_temperature_scale + return (config.unary_scale * unary + config.edge_scale * edge) / denominator + + +def _backward_messages(lattice: torch.Tensor) -> torch.Tensor: + """Compute normalized log-sum-exp future messages for a KxK chain.""" + num_steps, top_k, _ = lattice.shape + future = torch.zeros((num_steps, top_k), dtype=lattice.dtype) + for step in range(num_steps - 2, -1, -1): + values = lattice[step + 1] + future[step + 1][None, :] + future[step] = torch.logsumexp(values, dim=-1) + future[step] -= future[step].max() + return future + + +def _path_predecessors(record: AlignmentRecord) -> list[int]: + num_steps = record.candidate_ids.shape[0] + if record.draft_sampled.numel() < num_steps + 1: + raise ValueError(f"{record.path}: draft_sampled is missing the anchor row") + predecessors = [0] + for step in range(num_steps - 1): + proposed = int(record.draft_sampled[step + 1]) + matches = torch.nonzero(record.candidate_ids[step] == proposed).flatten() + if matches.numel() != 1: + raise ValueError( + f"{record.path}: proposed token {proposed} is not unique at step {step}" + ) + predecessors.append(int(matches[0])) + return predecessors + + +def _proposal_probs( + record: AlignmentRecord, + config: ProposalConfig, +) -> list[torch.Tensor]: + if config.use_cached_logits: + rows = [ + _compact_probs(row, config.proposal_top_p) for row in record.realized_logits + ] + else: + lattice = _transformed_lattice(record, config) + future = _backward_messages(lattice) + predecessors = _path_predecessors(record) + rows = [ + _compact_probs( + lattice[step, predecessor] + config.future_beta * future[step], + config.proposal_top_p, + ) + for step, predecessor in enumerate(predecessors) + ] + if config.greedy_mix == 0.0: + return rows + if not 0.0 <= config.greedy_mix <= 1.0: + raise ValueError("greedy_mix must be in [0, 1]") + mixed_rows = [] + for probs in rows: + mixed = probs * (1.0 - config.greedy_mix) + mixed[torch.argmax(probs)] += config.greedy_mix + mixed_rows.append(mixed) + return mixed_rows + + +def _target_rows(record: AlignmentRecord) -> list[torch.Tensor]: + rows = [] + for logits in record.target_topk_logits[: record.candidate_ids.shape[0]]: + rows.append(_compact_probs(logits / record.temperature, record.top_p)) + return rows + + +def _load_record(path: Path) -> AlignmentRecord: + payload = torch.load(path, map_location="cpu", weights_only=True) + if payload.get("format") != "dflash2_selector_alignment_v1": + raise ValueError(f"{path}: unsupported alignment format") + candidate_ids = payload["selector_candidate_ids"].to(torch.int64) + cached_ids = payload["draft_candidate_ids"].to(torch.int64) + if not torch.equal(candidate_ids, cached_ids): + raise ValueError(f"{path}: packed selector IDs do not match request-slot cache") + record = AlignmentRecord( + path=path, + step=int(payload["step"]), + temperature=float(payload["temperature"]), + top_p=float(payload["top_p"]), + target_topk_ids=payload["target_topk_ids"].to(torch.int64), + target_topk_logits=payload["target_topk_logits"].to(torch.float64), + candidate_ids=candidate_ids, + realized_logits=payload["draft_realized_logits"].to(torch.float64), + unary_logits=payload["selector_unary_logits"].to(torch.float64), + lattice_scores=payload["selector_lattice_scores"].to(torch.float64), + draft_sampled=payload["draft_sampled"].to(torch.int64).flatten(), + num_sampled=int(payload["num_sampled"].flatten()[0]), + ) + return record + + +def _mean(values: list[float]) -> float: + return float(sum(values) / len(values)) if values else math.nan + + +def _completion_proxy(per_position_overlap: list[float]) -> float: + survival = 1.0 + result = 1.0 + for overlap in per_position_overlap: + survival *= overlap + result += survival + return result + + +def _evaluate(records: list[AlignmentRecord], config: ProposalConfig) -> dict[str, Any]: + num_steps = records[0].candidate_ids.shape[0] + overlaps = [[] for _ in range(num_steps)] + for record in records: + target_rows = _target_rows(record) + proposal_rows = _proposal_probs(record, config) + for step, (target_probs, proposal_probs) in enumerate( + zip(target_rows, proposal_rows) + ): + overlaps[step].append( + _distribution_overlap( + record.target_topk_ids[step], + target_probs, + record.candidate_ids[step], + proposal_probs, + ) + ) + means = [_mean(values) for values in overlaps] + return { + "config": asdict(config), + "num_records": len(records), + "mean_overlap_by_position": means, + "completion_length_proxy": _completion_proxy(means), + } + + +def _sweep_configs() -> list[ProposalConfig]: + configs = [ProposalConfig(name="current", use_cached_logits=True)] + seen: set[tuple[float, float, float, float, float, float]] = set() + + def add( + family: str, + temperature: float = 1.0, + top_p: float = 1.0, + unary: float = 1.0, + edge: float = 1.0, + beta: float = 0.0, + greedy_mix: float = 0.0, + ) -> None: + key = (temperature, top_p, unary, edge, beta, greedy_mix) + if key in seen: + return + seen.add(key) + configs.append( + ProposalConfig( + name=( + f"{family}:t={temperature:g},p={top_p:g},u={unary:g}," + f"e={edge:g},b={beta:g},g={greedy_mix:g}" + ), + proposal_temperature_scale=temperature, + proposal_top_p=top_p, + unary_scale=unary, + edge_scale=edge, + future_beta=beta, + greedy_mix=greedy_mix, + ) + ) + + for temperature in (0.7, 0.8, 0.9, 0.95, 1.0, 1.05, 1.1, 1.2, 1.3): + add("temperature", temperature=temperature) + for top_p in (0.9, 0.95, 0.98, 0.99): + for temperature in (0.8, 0.9, 1.0): + add("nucleus", temperature=temperature, top_p=top_p) + for greedy_mix in (0.1, 0.2, 0.3, 0.4, 0.5): + for temperature in (0.8, 0.9, 1.0, 1.1): + add( + "greedy-mixture", + temperature=temperature, + greedy_mix=greedy_mix, + ) + for unary in (0.75, 1.0, 1.25): + for edge in (0.5, 0.75, 1.0, 1.25, 1.5): + for temperature in (0.85, 1.0, 1.15): + add( + "edge-calibration", + temperature=temperature, + unary=unary, + edge=edge, + ) + for beta in (0.25, 0.5, 0.75, 1.0): + for temperature in (0.8, 0.9, 1.0, 1.1): + for edge in (0.75, 1.0, 1.25): + add( + "future-message", + temperature=temperature, + edge=edge, + beta=beta, + ) + return configs + + +def _baseline_diagnostics(records: list[AlignmentRecord]) -> dict[str, Any]: + num_steps = records[0].candidate_ids.shape[0] + support = [[] for _ in range(num_steps)] + cached_lattice_max_abs: list[float] = [] + for record in records: + target_rows = _target_rows(record) + predecessors = _path_predecessors(record) + for step, target_probs in enumerate(target_rows): + support[step].append( + _support_mass( + record.target_topk_ids[step], + target_probs, + record.candidate_ids[step], + ) + ) + expected = ( + record.lattice_scores[step, predecessors[step]] / record.temperature + ) + cached_lattice_max_abs.append( + float((expected - record.realized_logits[step]).abs().max()) + ) + support_means = [_mean(values) for values in support] + return { + "observed_mean_completion_tokens_per_round": _mean( + [float(record.num_sampled) for record in records] + ), + "candidate_support_mass_by_position": support_means, + "candidate_support_completion_upper_proxy": _completion_proxy(support_means), + "cached_vs_lattice_max_abs": max(cached_lattice_max_abs, default=math.nan), + } + + +def summarize(pattern: str, top_n: int) -> dict[str, Any]: + paths = [Path(path) for path in sorted(glob.glob(pattern))] + if not paths: + raise ValueError(f"no selector alignment dumps matched {pattern!r}") + records = [_load_record(path) for path in paths] + tune = [record for record in records if record.step % 2 == 0] + holdout = [record for record in records if record.step % 2 == 1] + if not tune or not holdout: + tune = records + holdout = records + + configs = _sweep_configs() + tune_results = [_evaluate(tune, config) for config in configs] + holdout_by_name = { + result["config"]["name"]: result + for result in (_evaluate(holdout, config) for config in configs) + } + current_tune = tune_results[0]["completion_length_proxy"] + current_holdout = holdout_by_name["current"]["completion_length_proxy"] + position_policy = [] + position_tune_overlaps = [] + position_holdout_overlaps = [] + for position in range(records[0].candidate_ids.shape[0]): + best = max( + tune_results, + key=lambda result: result["mean_overlap_by_position"][position], + ) + holdout_best = holdout_by_name[best["config"]["name"]] + position_policy.append( + { + "position": position, + "config": best["config"], + "tune_overlap": best["mean_overlap_by_position"][position], + "holdout_overlap": holdout_best["mean_overlap_by_position"][position], + } + ) + position_tune_overlaps.append(best["mean_overlap_by_position"][position]) + position_holdout_overlaps.append( + holdout_best["mean_overlap_by_position"][position] + ) + ranked = sorted( + tune_results[1:], + key=lambda result: result["completion_length_proxy"], + reverse=True, + ) + leaders = [] + for result in ranked[:top_n]: + holdout_result = holdout_by_name[result["config"]["name"]] + leaders.append( + { + "config": result["config"], + "tune_completion_length_proxy": result["completion_length_proxy"], + "tune_delta": result["completion_length_proxy"] - current_tune, + "holdout_completion_length_proxy": holdout_result[ + "completion_length_proxy" + ], + "holdout_delta": ( + holdout_result["completion_length_proxy"] - current_holdout + ), + "tune_mean_overlap_by_position": result["mean_overlap_by_position"], + "holdout_mean_overlap_by_position": holdout_result[ + "mean_overlap_by_position" + ], + } + ) + + return { + "contract": { + "glob": pattern, + "num_records": len(records), + "num_tune_records": len(tune), + "num_holdout_records": len(holdout), + "split": "even/odd sampler step", + "counterfactual_scope": "one-step recorded-prefix overlap proxy", + "end_to_end_claim": False, + }, + "baseline": _baseline_diagnostics(records), + "current": { + "tune": tune_results[0], + "holdout": holdout_by_name["current"], + }, + "position_policy": { + "selection": "best tune overlap independently at each depth", + "rows": position_policy, + "tune_completion_length_proxy": _completion_proxy(position_tune_overlaps), + "tune_delta": (_completion_proxy(position_tune_overlaps) - current_tune), + "holdout_completion_length_proxy": _completion_proxy( + position_holdout_overlaps + ), + "holdout_delta": ( + _completion_proxy(position_holdout_overlaps) - current_holdout + ), + }, + "leaders": leaders, + } + + +def render_markdown(summary: dict[str, Any]) -> str: + contract = summary["contract"] + baseline = summary["baseline"] + current = summary["current"] + position_policy = summary["position_policy"] + lines = [ + "# DFlash2 Selector Alignment Sweep", + "", + ( + f"- Records: `{contract['num_records']}` " + f"(tune `{contract['num_tune_records']}`, " + f"holdout `{contract['num_holdout_records']}`)" + ), + ( + "- Counterfactual scope: one-step overlap on recorded prefixes; " + "this is not an end-to-end acceptance claim." + ), + ( + "- Observed completion tokens/round: " + f"`{baseline['observed_mean_completion_tokens_per_round']:.4f}`" + ), + ( + "- Current overlap proxy (tune/holdout): " + f"`{current['tune']['completion_length_proxy']:.4f}` / " + f"`{current['holdout']['completion_length_proxy']:.4f}`" + ), + ( + "- Fixed-top16 support upper proxy: " + f"`{baseline['candidate_support_completion_upper_proxy']:.4f}`" + ), + ( + "- Cached/lattice max absolute mismatch: " + f"`{baseline['cached_vs_lattice_max_abs']:.3e}`" + ), + ( + "- Position-wise held-out proxy/delta: " + f"`{position_policy['holdout_completion_length_proxy']:.4f}` / " + f"`{position_policy['holdout_delta']:+.4f}`" + ), + "", + "| candidate | tune proxy | tune delta | holdout proxy | holdout delta |", + "| --- | ---: | ---: | ---: | ---: |", + ] + for row in summary["leaders"]: + lines.append( + f"| `{row['config']['name']}` | " + f"{row['tune_completion_length_proxy']:.4f} | " + f"{row['tune_delta']:+.4f} | " + f"{row['holdout_completion_length_proxy']:.4f} | " + f"{row['holdout_delta']:+.4f} |" + ) + lines.extend( + [ + "", + "## Position-wise calibration policy", + "", + "| position | tune overlap | holdout overlap | configuration |", + "| ---: | ---: | ---: | --- |", + ] + ) + for row in position_policy["rows"]: + lines.append( + f"| {row['position']} | {row['tune_overlap']:.4f} | " + f"{row['holdout_overlap']:.4f} | `{row['config']['name']}` |" + ) + lines.append("") + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--glob", + default="/tmp/spec_alignment_dflash2_selector_*.pt", + help="Glob for GPU-runner DFlash2 selector alignment records.", + ) + parser.add_argument("--top", type=int, default=20) + parser.add_argument("--out-json", type=Path) + parser.add_argument("--out-md", type=Path) + args = parser.parse_args() + + summary = summarize(args.glob, args.top) + markdown = render_markdown(summary) + if args.out_json is not None: + args.out_json.parent.mkdir(parents=True, exist_ok=True) + args.out_json.write_text(json.dumps(summary, indent=2), encoding="utf-8") + if args.out_md is not None: + args.out_md.parent.mkdir(parents=True, exist_ok=True) + args.out_md.write_text(markdown, encoding="utf-8") + print(markdown) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_dsv4_gsm8k_api.py b/benchmarks/benchmark_dsv4_gsm8k_api.py new file mode 100644 index 0000000000..49f19a22b7 --- /dev/null +++ b/benchmarks/benchmark_dsv4_gsm8k_api.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Run a deterministic, sequential GSM8K gate against a vLLM API server.""" + +from __future__ import annotations + +import argparse +import hashlib +import http.client +import json +import time +from decimal import Decimal, InvalidOperation +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import regex as re + +_NUMBER_RE = re.compile(r"(?"] + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _load_jsonl(path: Path) -> list[dict[str, Any]]: + return [ + json.loads(line) + for line in path.read_text(encoding="utf-8").splitlines() + if line and not line.startswith("#") + ] + + +def _last_integer(text: str) -> int | None: + numbers = _NUMBER_RE.findall(text) + if not numbers: + return None + try: + value = Decimal(numbers[-1].replace(",", "")) + except InvalidOperation: + return None + if not value.is_finite() or value != value.to_integral_value(): + return None + return int(value) + + +def _post_completion( + *, + host: str, + port: int, + timeout: int, + payload: dict[str, Any], +) -> tuple[int, dict[str, Any], float]: + connection = http.client.HTTPConnection(host, port, timeout=timeout) + started = time.perf_counter() + try: + connection.request( + "POST", + "/v1/completions", + body=json.dumps(payload).encode("utf-8"), + headers={"Content-Type": "application/json"}, + ) + response = connection.getresponse() + body = response.read() + finally: + connection.close() + elapsed = time.perf_counter() - started + try: + decoded = json.loads(body) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"completion returned non-JSON HTTP {response.status}: " + f"{body.decode('utf-8', errors='replace')}" + ) from exc + return response.status, decoded, elapsed + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:8000") + parser.add_argument("--model", required=True) + parser.add_argument("--train", type=Path, required=True) + parser.add_argument("--test", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--limit", type=int, default=64) + parser.add_argument("--few-shot", type=int, default=5) + parser.add_argument("--max-tokens", type=int, default=256) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--timeout", type=int, default=1200) + parser.add_argument("--min-correct", type=int, default=0) + parser.add_argument("--max-invalid", type=int, default=0) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if args.limit <= 0 or args.few_shot < 0: + raise ValueError("--limit must be positive and --few-shot non-negative") + if args.min_correct < 0 or args.max_invalid < 0: + raise ValueError("quality thresholds must be non-negative") + + parsed = urlparse(args.base_url) + if parsed.scheme != "http" or not parsed.hostname: + raise ValueError("--base-url must be an http URL") + host = parsed.hostname + port = parsed.port or 80 + + train = _load_jsonl(args.train) + test = _load_jsonl(args.test)[: args.limit] + if len(train) < args.few_shot: + raise RuntimeError( + f"training set has {len(train)} rows, fewer than {args.few_shot} shots" + ) + if len(test) < args.limit: + raise RuntimeError(f"test set has only {len(test)} rows, need {args.limit}") + + few_shot = "".join( + f"Question: {row['question']}\nAnswer: {row['answer']}\n\n" + for row in train[: args.few_shot] + ) + rows: list[dict[str, Any]] = [] + started = time.perf_counter() + for index, row in enumerate(test): + prompt = few_shot + f"Question: {row['question']}\nAnswer:" + request = { + "model": args.model, + "prompt": prompt, + "temperature": 0.0, + "top_p": 1.0, + "max_tokens": args.max_tokens, + "stop": _STOP_SEQUENCES, + "seed": args.seed, + } + status, response, elapsed = _post_completion( + host=host, + port=port, + timeout=args.timeout, + payload=request, + ) + choices = response.get("choices") or [] + choice = choices[0] if choices else {} + output = choice.get("text") or "" + expected = _last_integer(row["answer"]) + if expected is None: + raise ValueError(f"GSM8K item {index} has no integral reference answer") + predicted = _last_integer(output) + rows.append( + { + "index": index, + "question": row["question"], + "expected": expected, + "predicted": predicted, + "correct": predicted is not None and predicted == expected, + "invalid": predicted is None, + "status": status, + "elapsed_seconds": elapsed, + "finish_reason": choice.get("finish_reason"), + "usage": response.get("usage"), + "output": output, + } + ) + if status != 200: + raise RuntimeError(f"GSM8K item {index} returned HTTP {status}: {response}") + + wall_seconds = time.perf_counter() - started + correct = sum(bool(row["correct"]) for row in rows) + invalid = sum(bool(row["invalid"]) for row in rows) + completion_tokens = sum( + int((row.get("usage") or {}).get("completion_tokens") or 0) for row in rows + ) + passed = correct >= args.min_correct and invalid <= args.max_invalid + result = { + "contract": { + "base_url": args.base_url, + "model": args.model, + "questions": args.limit, + "few_shot": args.few_shot, + "temperature": 0.0, + "top_p": 1.0, + "seed": args.seed, + "max_tokens": args.max_tokens, + "strictly_sequential": True, + "train_selection": "first_n", + "test_selection": "first_n", + "prompt_format": "gsm8k_question_answer_v1", + "answer_normalization": "last_signed_integral_decimal_v1", + "stop_sequences": _STOP_SEQUENCES, + "min_correct": args.min_correct, + "max_invalid": args.max_invalid, + }, + "input_manifest": { + "train": {"path": str(args.train), "sha256": _sha256(args.train)}, + "test": {"path": str(args.test), "sha256": _sha256(args.test)}, + }, + "correct": correct, + "accuracy": correct / len(rows), + "invalid": invalid, + "wall_seconds": wall_seconds, + "completion_tokens": completion_tokens, + "aggregate_output_tokens_per_second": ( + completion_tokens / wall_seconds if wall_seconds else None + ), + "passed": passed, + "rows": rows, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text( + json.dumps(result, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + print( + json.dumps( + { + "correct": correct, + "samples": len(rows), + "invalid": invalid, + "passed": passed, + } + ), + flush=True, + ) + return 0 if passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_dsv4_quality_api.py b/benchmarks/benchmark_dsv4_quality_api.py index 90f2982443..5eba6b30ea 100644 --- a/benchmarks/benchmark_dsv4_quality_api.py +++ b/benchmarks/benchmark_dsv4_quality_api.py @@ -582,6 +582,30 @@ def main() -> int: key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items() }, + "evaluation_contract": { + "version": "dsv4_api_quality_v1", + "strictly_sequential": True, + "temperature": 0.0, + "top_p": 1.0, + "seed": 20260824, + "human_eval": { + "selection": "first_n", + "limit": args.humaneval_limit, + "endpoint": "chat_completions", + "enable_thinking": False, + "max_tokens": 768, + "execution": "landlock_seccomp_python_isolated_v1", + }, + "longbench": { + "datasets": datasets, + "selection": "length_ge_8000_evenly_spaced_v1", + "limit_per_dataset": args.longbench_limit, + "max_input_tokens": args.longbench_max_input_tokens, + "truncation": "middle_first_half_last_half_v1", + "endpoint_policy": "official_no_chat_dataset_set_v1", + "max_output_tokens": "dataset2maxlen_sha256_manifest", + }, + }, "input_manifest": input_manifest, "human_eval": human_eval, "longbench": longbench, diff --git a/benchmarks/benchmark_sm70_decode.py b/benchmarks/benchmark_sm70_decode.py index 1dfee36b0f..19f27d6dea 100644 --- a/benchmarks/benchmark_sm70_decode.py +++ b/benchmarks/benchmark_sm70_decode.py @@ -30,7 +30,13 @@ def _module_file(module_name: str) -> str | None: module = sys.modules.get(module_name) if module is not None: return getattr(module, "__file__", None) - spec = importlib.util.find_spec(module_name) + try: + spec = importlib.util.find_spec(module_name) + except (AttributeError, ImportError, ValueError): + # Optional extension parents can raise while find_spec imports them. + # Diagnostics must report an unavailable extension, not discard an + # otherwise successful benchmark at result-serialization time. + return None return spec.origin if spec is not None else None diff --git a/benchmarks/benchmark_sm70_dflash2_gsm8k.py b/benchmarks/benchmark_sm70_dflash2_gsm8k.py new file mode 100644 index 0000000000..ac92519225 --- /dev/null +++ b/benchmarks/benchmark_sm70_dflash2_gsm8k.py @@ -0,0 +1,484 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Measure target-only or DFlash2 on a fixed local quality subset.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +import math +import random +import subprocess +import time +from collections import Counter +from pathlib import Path +from typing import Any + +import regex as re +from benchmark_sm70_decode import ( + _diff_spec_metrics, + _hash_ids, + _json_safe, + _module_file, + _module_realpath, + _request_metrics_dict, + _spec_metrics_snapshot, + _tracked_env, +) + +INVALID_ANSWER = -9_999_999 +GSM8K_PROMPT_SUFFIX = ( + "\nPlease reason step by step, and put your final answer within \\boxed{}." +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--draft-model", type=Path) + parser.add_argument("--dataset", type=Path, required=True) + parser.add_argument( + "--dataset-format", + choices=("gsm8k", "turns"), + default="gsm8k", + help=( + "Input schema. 'gsm8k' scores question/answer rows; 'turns' uses " + "the first preformatted turn and leaves task scoring to a separate " + "evaluator." + ), + ) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--mode", choices=("target-only", "dflash"), required=True) + parser.add_argument("--num-questions", type=int, default=64) + parser.add_argument("--start-index", type=int, default=0) + parser.add_argument( + "--dataset-order", + choices=("sequential", "zlab-shuffle42"), + default="sequential", + ) + parser.add_argument("--max-tokens", type=int, default=256) + parser.add_argument("--warmup-tokens", type=int, default=32) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--tensor-parallel-size", type=int, default=4) + parser.add_argument("--max-model-len", type=int, default=2048) + parser.add_argument("--max-num-batched-tokens", type=int, default=512) + parser.add_argument("--max-num-seqs", type=int, default=4) + parser.add_argument("--sequential", action="store_true") + parser.add_argument("--gpu-memory-utilization", type=float, default=0.8) + parser.add_argument( + "--target-kv-cache-dtype", + choices=("auto", "fp8_e5m2"), + default="fp8_e5m2", + ) + parser.add_argument("--enforce-eager", action="store_true") + parser.add_argument( + "--draft-sample-method", + choices=("greedy", "probabilistic"), + default="greedy", + ) + parser.add_argument( + "--draft-attention-backend", + choices=("FLASH_ATTN_V100", "TRITON_ATTN"), + default="FLASH_ATTN_V100", + help="Draft-only attention backend; the target remains on FLASH_ATTN_V100.", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--request-seed", + type=int, + default=0, + help="Sampling seed for every request; use -1 for server-style random seeds.", + ) + parser.add_argument( + "--cuda-profiler-capture", + action="store_true", + help="Wrap the measured generation in cudaProfilerStart/Stop for nsys.", + ) + return parser.parse_args() + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _model_weight_files(model: Path) -> dict[str, str]: + index = model / "model.safetensors.index.json" + if not index.is_file(): + weight = model / "model.safetensors" + return {weight.name: str(weight.resolve())} + weight_map = json.loads(index.read_text())["weight_map"] + return { + name: str((model / name).resolve()) for name in sorted(set(weight_map.values())) + } + + +def _answer_value(text: str) -> int: + numbers = re.findall(r"\d+", text.replace(",", "")) + if not numbers: + return INVALID_ANSWER + try: + return int(ast.literal_eval(numbers[-1])) + except (SyntaxError, ValueError): + return INVALID_ANSWER + + +def _percentile(values: list[float], quantile: float) -> float | None: + if not values: + return None + ordered = sorted(values) + position = (len(ordered) - 1) * quantile + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + weight = position - lower + return ordered[lower] + (ordered[upper] - ordered[lower]) * weight + + +def _distribution(values: list[float]) -> dict[str, float | int | None]: + return { + "count": len(values), + "mean": sum(values) / len(values) if values else None, + "p50": _percentile(values, 0.50), + "p90": _percentile(values, 0.90), + "p99": _percentile(values, 0.99), + "min": min(values) if values else None, + "max": max(values) if values else None, + } + + +def _load_rows(args: argparse.Namespace) -> list[tuple[int, dict[str, Any]]]: + rows = [ + json.loads(line) + for line in args.dataset.read_text(encoding="utf-8").splitlines() + if line + ] + indices = list(range(len(rows))) + if args.dataset_order == "zlab-shuffle42": + random.Random(42).shuffle(indices) + selected_indices = indices[args.start_index : args.start_index + args.num_questions] + if len(selected_indices) != args.num_questions: + raise ValueError( + f"Requested {args.num_questions} rows at index {args.start_index}, " + f"but the dataset provided {len(selected_indices)}." + ) + return [(index, rows[index]) for index in selected_indices] + + +def _prompt_content(row: dict[str, Any], dataset_format: str) -> str: + if dataset_format == "gsm8k": + return str(row["question"]) + GSM8K_PROMPT_SUFFIX + turns = row.get("turns") + if not isinstance(turns, list) or not turns or not isinstance(turns[0], str): + raise ValueError("turns dataset rows must contain a non-empty string list") + return turns[0] + + +def _summarize_requests(cases: list[dict[str, Any]]) -> dict[str, Any]: + metric_names = ( + "queued_time", + "first_token_latency", + "prefill_time", + "decode_time", + "steady_decode_tps", + "tpot_seconds", + ) + summary = {} + for metric_name in metric_names: + values = [ + float(case["request_metrics"][metric_name]) + for case in cases + if case["request_metrics"] is not None + and case["request_metrics"].get(metric_name) is not None + ] + summary[metric_name] = _distribution(values) + + prefill_tps = [ + case["prompt_tokens"] / case["request_metrics"]["prefill_time"] + for case in cases + if case["request_metrics"] is not None + and case["request_metrics"].get("prefill_time") + ] + summary["prefill_tokens_per_second"] = _distribution(prefill_tps) + summary["prompt_tokens"] = _distribution( + [float(case["prompt_tokens"]) for case in cases] + ) + summary["output_tokens"] = _distribution( + [float(case["output_tokens"]) for case in cases] + ) + return summary + + +def main() -> int: + args = _parse_args() + if args.num_questions <= 0: + raise ValueError("--num-questions must be positive") + if args.request_seed < -1: + raise ValueError("--request-seed must be -1 or a non-negative integer") + if args.mode == "dflash" and args.draft_model is None: + raise ValueError("--draft-model is required for --mode dflash") + + import torch + import vllm._C as vllm_c + import vllm._C_stable_libtorch as vllm_c_stable + from transformers import AutoTokenizer + + import vllm + from vllm import LLM, SamplingParams + + rows = _load_rows(args) + tokenizer = AutoTokenizer.from_pretrained(str(args.model)) + prompt_contents = [_prompt_content(row, args.dataset_format) for _, row in rows] + prompts = [ + tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": prompt_content, + } + ], + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + reasoning_effort="xhigh", + ) + for prompt_content in prompt_contents + ] + prompt_token_counts = [len(tokenizer.encode(prompt)) for prompt in prompts] + + speculative_config = None + if args.mode == "dflash": + speculative_config = { + "method": "dflash", + "model": str(args.draft_model), + "revision": "dedf8df68adfb1afeaf7b7480c0a0243108177b4", + "num_speculative_tokens": 7, + "kv_cache_dtype": "auto", + "attention_backend": args.draft_attention_backend, + "draft_sample_method": args.draft_sample_method, + "enforce_eager": args.enforce_eager, + } + + engine_kwargs = { + "model": str(args.model), + "tensor_parallel_size": args.tensor_parallel_size, + "dtype": "half", + "kv_cache_dtype": args.target_kv_cache_dtype, + "attention_backend": "FLASH_ATTN_V100", + "max_model_len": args.max_model_len, + "max_num_batched_tokens": args.max_num_batched_tokens, + "max_num_seqs": args.max_num_seqs, + "gpu_memory_utilization": args.gpu_memory_utilization, + "enable_prefix_caching": False, + "disable_log_stats": False, + "enforce_eager": args.enforce_eager, + "seed": args.seed, + "speculative_config": speculative_config, + } + if args.cuda_profiler_capture: + engine_kwargs["profiler_config"] = {"profiler": "cuda"} + + load_started = time.perf_counter() + llm = LLM(**engine_kwargs) + load_seconds = time.perf_counter() - load_started + + request_seed = None if args.request_seed == -1 else args.request_seed + warmup_sampling = SamplingParams( + temperature=args.temperature, + top_p=0.95, + top_k=20, + max_tokens=args.warmup_tokens, + seed=request_seed, + skip_special_tokens=False, + ) + llm.generate([prompts[0]], warmup_sampling, use_tqdm=False) + + sampling = SamplingParams( + temperature=args.temperature, + top_p=0.95, + top_k=20, + max_tokens=args.max_tokens, + seed=request_seed, + skip_special_tokens=False, + ) + spec_before = _spec_metrics_snapshot(llm) + if args.cuda_profiler_capture: + llm.start_profile() + started = time.perf_counter() + if args.sequential: + outputs = [] + request_spec_metrics = [] + for prompt in prompts: + request_spec_before = _spec_metrics_snapshot(llm) + outputs.append(llm.generate([prompt], sampling, use_tqdm=False)[0]) + request_spec_after = _spec_metrics_snapshot(llm) + request_spec_metrics.append( + _diff_spec_metrics(request_spec_before, request_spec_after) + ) + else: + outputs = llm.generate(prompts, sampling, use_tqdm=False) + request_spec_metrics = [None] * len(outputs) + elapsed_seconds = time.perf_counter() - started + if args.cuda_profiler_capture: + llm.stop_profile() + spec_after = _spec_metrics_snapshot(llm) + + cases = [] + for ( + dataset_index, + row, + ), prompt_content, prompt_tokens, output, request_spec in zip( + rows, prompt_contents, prompt_token_counts, outputs, request_spec_metrics + ): + result = output.outputs[0] + token_ids = list(result.token_ids) + if args.dataset_format == "gsm8k": + prediction = _answer_value(result.text) + expected = _answer_value(row["answer"]) + correct = prediction == expected + else: + prediction = None + expected = None + correct = None + cases.append( + { + "dataset_index": dataset_index, + "question": row.get("question"), + "prompt_content": prompt_content, + "expected_answer": expected, + "predicted_answer": prediction, + "correct": correct, + "prompt_tokens": prompt_tokens, + "output_tokens": len(token_ids), + "finish_reason": result.finish_reason, + "stop_reason": result.stop_reason, + "text": result.text, + "token_ids": token_ids, + "token_hash": _hash_ids(token_ids), + "spec_decode_metrics": request_spec, + "request_metrics": _request_metrics_dict( + output.metrics, + len(token_ids), + ), + } + ) + + total_output_tokens = sum(case["output_tokens"] for case in cases) + scored_cases = [case for case in cases if case["correct"] is not None] + correct = sum(bool(case["correct"]) for case in scored_cases) + invalid = sum(case["predicted_answer"] == INVALID_ANSWER for case in scored_cases) + per_request_acceptance_lengths = [ + float(case["spec_decode_metrics"]["acceptance_length"]) + for case in cases + if case["spec_decode_metrics"] is not None + and case["spec_decode_metrics"].get("acceptance_length") is not None + ] + per_request_completion_tokens_per_verification_step = [ + case["output_tokens"] / case["spec_decode_metrics"]["num_drafts"] + for case in cases + if case["spec_decode_metrics"] is not None + and case["spec_decode_metrics"].get("num_drafts", 0) > 0 + ] + c_extension = Path(vllm_c.__file__).resolve() + c_stable_extension = Path(vllm_c_stable.__file__).resolve() + payload = { + "contract": { + "source_sha": subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).resolve().parents[1], + text=True, + ).strip(), + "mode": args.mode, + "dataset": str(args.dataset), + "dataset_sha256": _sha256_file(args.dataset), + "dataset_format": args.dataset_format, + "start_index": args.start_index, + "num_questions": args.num_questions, + "dataset_order": args.dataset_order, + "model": str(args.model), + "model_config_sha256": _sha256_file(args.model / "config.json"), + "model_index_sha256": ( + _sha256_file(args.model / "model.safetensors.index.json") + if (args.model / "model.safetensors.index.json").is_file() + else None + ), + "model_weight_files": _model_weight_files(args.model), + "model_weights_realpath": str((args.model / "model.safetensors").resolve()), + "draft_model": str(args.draft_model) if args.draft_model else None, + "graph": not args.enforce_eager, + "sequential": args.sequential, + "sampling": { + "temperature": args.temperature, + "top_p": 0.95, + "top_k": 20, + "max_tokens": args.max_tokens, + "seed": request_seed, + "ignore_eos": False, + "thinking": True, + "reasoning_effort": "xhigh", + "prompt_suffix": ( + GSM8K_PROMPT_SUFFIX if args.dataset_format == "gsm8k" else None + ), + }, + "engine_kwargs": engine_kwargs, + }, + "runtime": { + "vllm_version": getattr(vllm, "__version__", None), + "vllm_file": getattr(vllm, "__file__", None), + "torch_version": torch.__version__, + "torch_cuda": torch.version.cuda, + "cuda_device_count": torch.accelerator.device_count(), + "device_capabilities": [ + list(torch.cuda.get_device_capability(index)) + for index in range(torch.accelerator.device_count()) + ], + "c_extension": str(c_extension), + "c_extension_sha256": _sha256_file(c_extension), + "c_stable_extension": str(c_stable_extension), + "c_stable_extension_sha256": _sha256_file(c_stable_extension), + "flash_attn_v100_python": _module_file("flash_attn_v100"), + "flash_attn_v100_cuda": _module_realpath("flash_attn_v100_cuda"), + "tracked_env": _tracked_env(), + "load_seconds": load_seconds, + }, + "results": { + "elapsed_seconds": elapsed_seconds, + "total_output_tokens": total_output_tokens, + "aggregate_output_tokens_per_second": ( + total_output_tokens / elapsed_seconds + ), + "questions_per_second": len(cases) / elapsed_seconds, + "accuracy": correct / len(scored_cases) if scored_cases else None, + "invalid_answer_rate": ( + invalid / len(scored_cases) if scored_cases else None + ), + "finish_reasons": dict( + Counter(str(case["finish_reason"]) for case in cases) + ), + "request_metrics": _summarize_requests(cases), + "spec_decode_metrics": _diff_spec_metrics(spec_before, spec_after), + "per_request_acceptance_length": _distribution( + per_request_acceptance_lengths + ), + "per_request_completion_tokens_per_verification_step": _distribution( + per_request_completion_tokens_per_verification_step + ), + }, + "cases": cases, + } + payload = _json_safe(payload) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + printable = {"contract": payload["contract"], "results": payload["results"]} + print(json.dumps(printable, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/compare_dsv4_quality_results.py b/benchmarks/compare_dsv4_quality_results.py index 84a1b26342..073d8276de 100644 --- a/benchmarks/compare_dsv4_quality_results.py +++ b/benchmarks/compare_dsv4_quality_results.py @@ -66,9 +66,10 @@ def _compare_humaneval( "regressions": regressions, "improvements": improvements, "missing_or_extra_tasks": missing, - "passed": not regressions - and not missing - and candidate_passed >= reference_passed, + # Individual flips remain visible for diagnosis, but are not token- + # identity gates. Acceptance is based on matched inputs and aggregate + # task quality. + "passed": not missing and candidate_passed >= reference_passed, } @@ -121,8 +122,7 @@ def _compare_longbench( "regressions": regressions, "missing_or_extra_rows": row_mismatch, "input_mismatches": input_mismatch, - "passed": not regressions - and not row_mismatch + "passed": not row_mismatch and not input_mismatch and candidate_score + tolerance >= reference_score, } @@ -148,25 +148,89 @@ def _compare_api_quality( reference_hashes = _manifest_hashes(reference.get("input_manifest", {})) candidate_hashes = _manifest_hashes(candidate.get("input_manifest", {})) manifest_equal = bool(reference_hashes) and reference_hashes == candidate_hashes + reference_contract = reference.get("evaluation_contract") + candidate_contract = candidate.get("evaluation_contract") + contract_equal = ( + bool(reference_contract) and reference_contract == candidate_contract + ) humaneval = _compare_humaneval(reference, candidate) longbench = _compare_longbench(reference, candidate, tolerance) return { "input_hashes_equal": manifest_equal, "reference_input_hashes": reference_hashes, "candidate_input_hashes": candidate_hashes, + "evaluation_contract_equal": contract_equal, + "reference_evaluation_contract": reference_contract, + "candidate_evaluation_contract": candidate_contract, "humaneval": humaneval, "longbench": longbench, - "passed": manifest_equal and humaneval["passed"] and longbench["passed"], + "passed": manifest_equal + and contract_equal + and humaneval["passed"] + and longbench["passed"], } def _compare_gsm8k( reference: dict[str, Any], candidate: dict[str, Any] ) -> dict[str, Any]: + reference_hashes = _manifest_hashes(reference.get("input_manifest", {})) + candidate_hashes = _manifest_hashes(candidate.get("input_manifest", {})) + manifest_equal = bool(reference_hashes) and reference_hashes == candidate_hashes + contract_keys = ( + "questions", + "few_shot", + "temperature", + "top_p", + "seed", + "max_tokens", + "strictly_sequential", + "train_selection", + "test_selection", + "prompt_format", + "answer_normalization", + "stop_sequences", + ) + reference_contract = { + key: reference.get("contract", {}).get(key) for key in contract_keys + } + candidate_contract = { + key: candidate.get("contract", {}).get(key) for key in contract_keys + } + contract_equal = reference_contract == candidate_contract reference_rows = {int(row["index"]): row for row in reference["rows"]} candidate_rows = {int(row["index"]): row for row in candidate["rows"]} common = sorted(reference_rows.keys() & candidate_rows.keys()) missing = sorted(reference_rows.keys() ^ candidate_rows.keys()) + expected_samples = reference_contract["questions"] + row_count_matches_contract = ( + isinstance(expected_samples, int) + and len(reference["rows"]) == expected_samples + and len(candidate["rows"]) == expected_samples + and len(reference_rows) == expected_samples + and len(candidate_rows) == expected_samples + ) + + def row_is_consistent(row: dict[str, Any]) -> bool: + expected = row.get("expected") + predicted = row.get("predicted") + return bool( + expected is not None + and bool(row.get("correct")) + == (predicted is not None and predicted == expected) + and bool(row.get("invalid")) == (predicted is None) + ) + + rows_self_consistent = all( + row_is_consistent(row) for row in [*reference["rows"], *candidate["rows"]] + ) + summaries_self_consistent = all( + int(result["correct"]) + == sum(bool(row.get("correct")) for row in result["rows"]) + and int(result["invalid"]) + == sum(bool(row.get("invalid")) for row in result["rows"]) + for result in (reference, candidate) + ) input_mismatch = [ index for index in common @@ -184,21 +248,43 @@ def _compare_gsm8k( for index in common if reference_rows[index]["correct"] and not candidate_rows[index]["correct"] ] + improvements = [ + index + for index in common + if not reference_rows[index]["correct"] and candidate_rows[index]["correct"] + ] exact_predictions = sum( reference_rows[index].get("predicted") == candidate_rows[index].get("predicted") for index in common ) return { + "input_hashes_equal": manifest_equal, + "reference_input_hashes": reference_hashes, + "candidate_input_hashes": candidate_hashes, + "evaluation_contract_equal": contract_equal, + "reference_evaluation_contract": reference_contract, + "candidate_evaluation_contract": candidate_contract, + "row_count_matches_contract": row_count_matches_contract, + "rows_self_consistent": rows_self_consistent, + "summaries_self_consistent": summaries_self_consistent, "reference_correct": int(reference["correct"]), "candidate_correct": int(candidate["correct"]), + "correct_delta": int(candidate["correct"]) - int(reference["correct"]), "reference_invalid": int(reference["invalid"]), "candidate_invalid": int(candidate["invalid"]), "samples": len(common), "exact_prediction_matches": exact_predictions, "regressions": regressions, + "improvements": improvements, "missing_or_extra_rows": missing, "input_mismatches": input_mismatch, - "passed": not regressions + # Report directional flips without requiring greedy/token identity. + # Matched aggregate correctness and validity are the quality gate. + "passed": manifest_equal + and contract_equal + and row_count_matches_contract + and rows_self_consistent + and summaries_self_consistent and not missing and not input_mismatch and int(candidate["correct"]) >= int(reference["correct"]) @@ -256,9 +342,7 @@ def _compare_needle( "anywhere_hit_regressions": anywhere_regressions, "missing_or_extra_rows": missing, "input_mismatches": input_mismatch, - "passed": not final_regressions - and not anywhere_regressions - and not missing + "passed": not missing and not input_mismatch and candidate_hits >= reference_hits and candidate_anywhere >= reference_anywhere, @@ -278,8 +362,8 @@ def _paired_paths( def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--reference-api", type=Path, required=True) - parser.add_argument("--candidate-api", type=Path, required=True) + parser.add_argument("--reference-api", type=Path) + parser.add_argument("--candidate-api", type=Path) parser.add_argument("--reference-gsm8k", type=Path) parser.add_argument("--candidate-gsm8k", type=Path) parser.add_argument("--reference-needle", type=Path) @@ -290,25 +374,31 @@ def main() -> int: if args.longbench_score_tolerance < 0: parser.error("--longbench-score-tolerance must be non-negative") + has_api = _paired_paths(parser, args.reference_api, args.candidate_api, "api") has_gsm8k = _paired_paths( parser, args.reference_gsm8k, args.candidate_gsm8k, "gsm8k" ) has_needle = _paired_paths( parser, args.reference_needle, args.candidate_needle, "needle" ) - api = _compare_api_quality( - _load_json(args.reference_api), - _load_json(args.candidate_api), - args.longbench_score_tolerance, - ) + if not (has_api or has_gsm8k or has_needle): + parser.error("at least one paired quality artifact must be provided") + result: dict[str, Any] = { "contract": { key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items() }, - "api_quality": api, } - gates = [api["passed"]] + gates: list[bool] = [] + if has_api: + assert args.reference_api is not None and args.candidate_api is not None + result["api_quality"] = _compare_api_quality( + _load_json(args.reference_api), + _load_json(args.candidate_api), + args.longbench_score_tolerance, + ) + gates.append(result["api_quality"]["passed"]) if has_gsm8k: assert args.reference_gsm8k is not None and args.candidate_gsm8k is not None result["gsm8k"] = _compare_gsm8k( diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index e21be2eb8c..8918d9e7b0 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -53,6 +53,9 @@ elseif(VLLM_FLASH_ATTN_SM70) COMMAND ${PATCH_EXECUTABLE} --batch --forward -p1 -l -i ${CMAKE_CURRENT_LIST_DIR}/../patches/sm70_flash_attn_d256_k_pingpong.patch + COMMAND + ${PATCH_EXECUTABLE} --batch --forward -p1 -l + -i ${CMAKE_CURRENT_LIST_DIR}/../patches/sm70_flash_attn_d256_gqa_arch.patch BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn ) else() diff --git a/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch b/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch new file mode 100644 index 0000000000..d6d65439a5 --- /dev/null +++ b/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch @@ -0,0 +1,1673 @@ +diff --git a/CMakeLists.txt b/CMakeLists.txt +index acf2219..f5e2732 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -84,6 +84,7 @@ string(REPLACE "-O2" "-O3" CMAKE_CUDA_FLAGS_RELWITHDEBINFO "${CMAKE_CUDA_FLAGS_R + + if (FA2_ENABLED) + set(FA2_GEN_SRCS ++ "csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu" + "csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu" + "csrc/flash_attn/src/flash_fwd_hdim256_causal_prefill_sm70.cu" + "csrc/flash_attn/src/flash_fwd_split_hdim256_causal_prefill_sm70.cu") +@@ -114,7 +115,8 @@ if (FA2_ENABLED) + csrc/flash_attn + csrc/flash_attn/src + csrc/common +- csrc/cutlass/include) ++ csrc/cutlass/include ++ csrc/cutlass/examples/35_gemm_softmax) + + # custom definitions + target_compile_definitions(_vllm_fa2_C PRIVATE +@@ -125,5 +127,16 @@ if (FA2_ENABLED) + # FLASHATTENTION_DISABLE_UNEVEN_K + # FLASHATTENTION_DISABLE_LOCAL + FLASHATTENTION_DISABLE_PYBIND ++ PREFIX_TORCH_EXTENSION ++ PREFIX_QK_FULL_STATS ++ PREFIX_QK_SKIP_APPLY ++ QK_TB_M=128 ++ QK_TB_N=512 ++ QK_WARP_M=64 ++ QK_WARP_N=128 ++ PV_TB_M=64 ++ PV_TB_N=256 ++ PV_WARP_M=32 ++ PV_WARP_N=64 + ) + endif () +diff --git a/csrc/flash_attn/flash_api_torch_lib.cpp b/csrc/flash_attn/flash_api_torch_lib.cpp +index 4b7601a..f5b71a2 100644 +--- a/csrc/flash_attn/flash_api_torch_lib.cpp ++++ b/csrc/flash_attn/flash_api_torch_lib.cpp +@@ -33,6 +33,24 @@ at::Tensor sm70_d256_splitd_dense_splitkv3_fwd( + double softmax_scale, + bool causal); + ++at::Tensor sm70_d256_splitd_dense_state_fwd( ++ const at::Tensor &q, ++ const at::Tensor &k, ++ const at::Tensor &v, ++ at::Tensor &state_max, ++ at::Tensor &state_sum, ++ at::Tensor &out, ++ double softmax_scale, ++ bool causal); ++ ++at::Tensor sm70_d256_gqa_architecture_fwd( ++ const at::Tensor &q, ++ const at::Tensor &k, ++ const at::Tensor &v, ++ at::Tensor &out, ++ double softmax_scale, ++ bool causal); ++ + at::Tensor sm70_d256_splitd_paged_fwd( + const at::Tensor &q, + const at::Tensor &k, +@@ -98,6 +116,17 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { + ops.impl("sm70_d256_splitd_n32_dense_fwd", torch::kCUDA, + make_pytorch_shim(&sm70_d256_splitd_dense_fwd)); + ++ ops.def("sm70_d256_splitd_n32_dense_state_fwd(Tensor q, Tensor k, " ++ "Tensor v, Tensor(a!) state_max, Tensor(b!) state_sum, " ++ "Tensor(c!) out, float softmax_scale, bool causal) -> Tensor(c!)"); ++ ops.impl("sm70_d256_splitd_n32_dense_state_fwd", torch::kCUDA, ++ make_pytorch_shim(&sm70_d256_splitd_dense_state_fwd)); ++ ++ ops.def("sm70_d256_gqa_architecture_fwd(Tensor q, Tensor k, Tensor v, " ++ "Tensor(a!) out, float softmax_scale, bool causal) -> Tensor(a!)"); ++ ops.impl("sm70_d256_gqa_architecture_fwd", torch::kCUDA, ++ make_pytorch_shim(&sm70_d256_gqa_architecture_fwd)); ++ + ops.def("sm70_d256_splitd_n32_dense_splitkv3_fwd(Tensor q, Tensor k, " + "Tensor v, Tensor(a!) partial_out, Tensor(b!) partial_max, " + "Tensor(c!) partial_sum, Tensor(d!) out, float softmax_scale, " +diff --git a/csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu b/csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu +new file mode 100644 +index 0000000..197453f +--- /dev/null ++++ b/csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu +@@ -0,0 +1,1359 @@ ++// SPDX-License-Identifier: BSD-3-Clause ++ ++/*************************************************************************************************** ++ * End-to-end prefix architecture screen for V100/SM70. ++ * ++ * QK GEMM writes FP16 logits and row statistics. PV normalizes logits in its ++ * A-operand transform and writes one reusable FP16 block partial. A float ++ * online accumulator folds that partial into prefix softmax state before the ++ * next block, avoiding resident storage for every block partial. ++ **************************************************************************************************/ ++ ++#include ++#include ++#include ++ ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++ ++#if defined(PREFIX_TORCH_EXTENSION) ++#include ++#include ++#include ++#include ++#include "namespace_config.h" ++#endif ++ ++#include "cutlass/cutlass.h" ++#include "cutlass/device_kernel.h" ++#include "cutlass/epilogue/thread/linear_combination.h" ++#include "cutlass/gemm/kernel/default_gemm.h" ++#include "cutlass/gemm/kernel/gemm.h" ++#include "cutlass/gemm/threadblock/mma_pipelined.h" ++#include "cutlass/layout/matrix.h" ++#include "cutlass/numeric_conversion.h" ++#include "gemm_with_softmax.h" ++ ++namespace { ++ ++using Element = cutlass::half_t; ++ ++#ifndef QK_TB_M ++#define QK_TB_M 128 ++#endif ++#ifndef QK_TB_N ++#define QK_TB_N 128 ++#endif ++#ifndef QK_WARP_M ++#define QK_WARP_M 32 ++#endif ++#ifndef QK_WARP_N ++#define QK_WARP_N 64 ++#endif ++ ++using QKLayoutA = cutlass::layout::RowMajor; ++using QKLayoutB = cutlass::layout::ColumnMajor; ++using QKThreadblockShape = cutlass::gemm::GemmShape; ++using QKWarpShape = cutlass::gemm::GemmShape; ++using QKInstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; ++using QKOutputOp = cutlass::epilogue::thread::LinearCombination< ++ Element, 8, Element, Element>; ++using QKGemm = cutlass::GemmSoftmax< ++ Element, ++ QKLayoutA, ++ Element, ++ QKLayoutB, ++ Element, ++ Element, ++ cutlass::arch::OpClassTensorOp, ++ cutlass::arch::Sm70, ++ QKThreadblockShape, ++ QKWarpShape, ++ QKInstructionShape, ++ QKOutputOp, ++ 2, ++ cutlass::MatrixShape<1, 1024>>; ++ ++#ifndef PV_TB_M ++#define PV_TB_M 64 ++#endif ++#ifndef PV_TB_N ++#define PV_TB_N 128 ++#endif ++#ifndef PV_WARP_M ++#define PV_WARP_M 32 ++#endif ++#ifndef PV_WARP_N ++#define PV_WARP_N 64 ++#endif ++ ++using PVLayout = cutlass::layout::RowMajor; ++using PVAccumulator = Element; ++using PVThreadblockShape = cutlass::gemm::GemmShape; ++using PVWarpShape = cutlass::gemm::GemmShape; ++using PVInstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; ++using PVSwizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; ++using PVOutputOp = cutlass::epilogue::thread::LinearCombination< ++ Element, 8, PVAccumulator, float>; ++ ++constexpr int kPVAlignment = 8; ++ ++using PVDefaultKernel = typename cutlass::gemm::kernel::DefaultGemm< ++ Element, ++ PVLayout, ++ kPVAlignment, ++ Element, ++ PVLayout, ++ kPVAlignment, ++ Element, ++ PVLayout, ++ PVAccumulator, ++ cutlass::arch::OpClassTensorOp, ++ cutlass::arch::Sm70, ++ PVThreadblockShape, ++ PVWarpShape, ++ PVInstructionShape, ++ PVOutputOp, ++ PVSwizzle, ++ 2, ++ false, ++ cutlass::arch::OpMultiplyAdd>::GemmKernel; ++ ++using PVDefaultMma = typename PVDefaultKernel::Mma; ++using PVIteratorA = typename PVDefaultMma::IteratorA; ++using PVIteratorB = typename PVDefaultMma::IteratorB; ++using PVSmemIteratorA = typename PVDefaultMma::SmemIteratorA; ++using PVSmemIteratorB = typename PVDefaultMma::SmemIteratorB; ++ ++__device__ float const* g_row_max = nullptr; ++__device__ float const* g_row_inv_sum = nullptr; ++__device__ float* g_row_sum_out = nullptr; ++__device__ int g_rows = 0; ++ ++#if defined(PREFIX_QK_FULL_STATS) ++struct ExpRowSumTransformA { ++ using InputFragment = typename PVIteratorA::Fragment; ++ using OutputFragment = cutlass::Array< ++ typename PVSmemIteratorA::Element, InputFragment::kElements>; ++ using ThreadMap = typename PVIteratorA::ThreadMap; ++ static constexpr int kAccessesPerVector = ++ PVIteratorA::UnderlyingIterator::kAccessesPerVector; ++ static constexpr int kContiguousIterations = ++ ThreadMap::Iterations::kContiguous; ++ static constexpr int kStridedIterations = ThreadMap::Iterations::kStrided; ++ ++ float row_max[kStridedIterations]; ++ float row_inv_sum[kStridedIterations]; ++ ++ CUTLASS_DEVICE ++ ExpRowSumTransformA() { ++ auto thread_offset = ThreadMap::initial_offset(threadIdx.x); ++#pragma unroll ++ for (int s = 0; s < kStridedIterations; ++s) { ++ int row = blockIdx.x * PVThreadblockShape::kM ++ + thread_offset.strided() + s * ThreadMap::Delta::kStrided; ++ row_max[s] = row < g_rows ? g_row_max[row] : 0.0f; ++ row_inv_sum[s] = row < g_rows ? g_row_inv_sum[row] : 0.0f; ++ } ++ } ++ ++ CUTLASS_DEVICE ++ OutputFragment operator()(InputFragment const& input) { ++ OutputFragment output; ++ constexpr float kLog2E = 1.4426950408889634f; ++ constexpr int kElementsPerAccess = PVIteratorA::AccessType::kElements; ++ auto const* input_access = ++ reinterpret_cast(&input); ++ auto* output_access = ++ reinterpret_cast(&output); ++#pragma unroll ++ for (int s = 0; s < kStridedIterations; ++s) { ++#pragma unroll ++ for (int c = 0; c < kContiguousIterations; ++c) { ++#pragma unroll ++ for (int v = 0; v < kAccessesPerVector; ++v) { ++ int index = v + kAccessesPerVector * (c + s * kContiguousIterations); ++ typename PVIteratorA::AccessType transformed; ++#pragma unroll ++ for (int e = 0; e < kElementsPerAccess; ++e) { ++ float value = static_cast(input_access[index][e]); ++ transformed[e] = Element( ++ exp2f((value - row_max[s]) * kLog2E) * row_inv_sum[s]); ++ } ++ output_access[index] = transformed; ++ } ++ } ++ } ++ return output; ++ } ++}; ++#else ++struct ExpRowSumTransformA { ++ using InputFragment = typename PVIteratorA::Fragment; ++ using OutputFragment = cutlass::Array< ++ typename PVSmemIteratorA::Element, InputFragment::kElements>; ++ using ThreadMap = typename PVIteratorA::ThreadMap; ++ ++ static constexpr int kAccessesPerVector = ++ PVIteratorA::UnderlyingIterator::kAccessesPerVector; ++ static constexpr int kContiguousIterations = ++ ThreadMap::Iterations::kContiguous; ++ static constexpr int kStridedIterations = ThreadMap::Iterations::kStrided; ++ ++ float row_max[kStridedIterations]; ++ float row_sum[kStridedIterations]; ++ bool valid = true; ++ ++ CUTLASS_DEVICE ++ void set_valid(bool is_valid) { ++ valid = is_valid; ++ } ++ ++ CUTLASS_DEVICE ++ ExpRowSumTransformA() { ++ auto thread_offset = ThreadMap::initial_offset(threadIdx.x); ++#pragma unroll ++ for (int s = 0; s < kStridedIterations; ++s) { ++ int row = blockIdx.x * PVThreadblockShape::kM ++ + thread_offset.strided() + s * ThreadMap::Delta::kStrided; ++ row_max[s] = row < g_rows ? g_row_max[row] : 0.0f; ++ row_sum[s] = 0.0f; ++ } ++ } ++ ++ CUTLASS_DEVICE ++ OutputFragment operator()(InputFragment const& input) { ++ if (!valid) { ++ OutputFragment output; ++ output.clear(); ++ return output; ++ } ++ OutputFragment output; ++ constexpr float kLog2E = 1.4426950408889634f; ++ constexpr int kElementsPerAccess = PVIteratorA::AccessType::kElements; ++ auto const* input_access = ++ reinterpret_cast(&input); ++ auto* output_access = ++ reinterpret_cast(&output); ++#pragma unroll ++ for (int s = 0; s < kStridedIterations; ++s) { ++#pragma unroll ++ for (int c = 0; c < kContiguousIterations; ++c) { ++#pragma unroll ++ for (int v = 0; v < kAccessesPerVector; ++v) { ++ int index = v + kAccessesPerVector * (c + s * kContiguousIterations); ++ typename PVIteratorA::AccessType transformed; ++#pragma unroll ++ for (int e = 0; e < kElementsPerAccess; ++e) { ++ float value = static_cast(input_access[index][e]); ++ float weight = exp2f((value - row_max[s]) * kLog2E); ++ if (blockIdx.y == 0) { ++ row_sum[s] += weight; ++ } ++ transformed[e] = Element(weight); ++ } ++ output_access[index] = transformed; ++ } ++ } ++ } ++ return output; ++ } ++ ++ CUTLASS_DEVICE ++ void finalize() { ++ if (blockIdx.y != 0) { ++ return; ++ } ++ constexpr int kThreadsPerRow = ++ ThreadMap::Detail::WarpThreadArrangement::kContiguous; ++ auto thread_offset = ThreadMap::initial_offset(threadIdx.x); ++ int lane = threadIdx.x & 31; ++#pragma unroll ++ for (int s = 0; s < kStridedIterations; ++s) { ++ float sum = row_sum[s]; ++#pragma unroll ++ for (int offset = 1; offset < kThreadsPerRow; offset <<= 1) { ++ sum += __shfl_xor_sync(0xffffffffu, sum, offset); ++ } ++ int row = blockIdx.x * PVThreadblockShape::kM ++ + thread_offset.strided() + s * ThreadMap::Delta::kStrided; ++ if ((lane % kThreadsPerRow) == 0 && row < g_rows) { ++ g_row_sum_out[row] = sum; ++ } ++ } ++ } ++}; ++#endif ++ ++using PVTransformB = cutlass::NumericArrayConverter< ++ typename PVSmemIteratorB::Element, ++ typename PVIteratorB::Element, ++ PVIteratorB::Fragment::kElements>; ++using PVMma = cutlass::gemm::threadblock::MmaPipelined< ++ typename PVDefaultMma::Shape, ++ PVIteratorA, ++ PVSmemIteratorA, ++ PVIteratorB, ++ PVSmemIteratorB, ++ PVAccumulator, ++ PVLayout, ++ typename PVDefaultMma::Policy, ++ ExpRowSumTransformA, ++ PVTransformB>; ++using PVKernel = cutlass::gemm::kernel::Gemm< ++ PVMma, typename PVDefaultKernel::Epilogue, PVSwizzle, false>; ++ ++void check(cudaError_t result, char const* operation) { ++ if (result != cudaSuccess) { ++ std::cerr << operation << ": " << cudaGetErrorString(result) << "\n"; ++ std::exit(EXIT_FAILURE); ++ } ++} ++ ++void check(cutlass::Status status, char const* operation) { ++ if (status != cutlass::Status::kSuccess) { ++ std::cerr << operation << ": CUTLASS status " << int(status) << "\n"; ++ std::exit(EXIT_FAILURE); ++ } ++} ++ ++struct PVLauncher { ++ typename PVKernel::Params params; ++ dim3 grid; ++ dim3 block; ++ int smem_bytes; ++ ++ PVLauncher(Element* scores, Element* value, Element* output, int rows, int k) { ++ cutlass::gemm::GemmCoord problem(rows, 256, k); ++ PVSwizzle swizzle; ++ auto tiled_shape = swizzle.get_tiled_shape( ++ problem, ++ {PVThreadblockShape::kM, PVThreadblockShape::kN, ++ PVThreadblockShape::kK}, ++ 1); ++ params = typename PVKernel::Params( ++ problem, ++ tiled_shape, ++ {scores, PVLayout(k)}, ++ {value, PVLayout(256)}, ++ {output, PVLayout(256)}, ++ {output, PVLayout(256)}, ++ typename PVOutputOp::Params(1.0f, 0.0f), ++ nullptr); ++ grid = swizzle.get_grid_shape(tiled_shape); ++ block = dim3(PVKernel::kThreadCount, 1, 1); ++ smem_bytes = int(sizeof(typename PVKernel::SharedStorage)); ++ if (smem_bytes >= 48 * 1024) { ++ check(cudaFuncSetAttribute( ++ cutlass::Kernel, ++ cudaFuncAttributeMaxDynamicSharedMemorySize, ++ smem_bytes), ++ "set PV dynamic shared memory"); ++ } ++ } ++ ++ void launch(cudaStream_t stream) const { ++ cutlass::Kernel<<>>(params); ++ } ++}; ++ ++__global__ void prepare_weights( ++ float const* block_max, ++ float const* block_sum, ++ float* weights, ++ float* state_max, ++ float* state_sum, ++ int blocks, ++ int rows) { ++ int row = blockIdx.x * blockDim.x + threadIdx.x; ++ if (row >= rows) { ++ return; ++ } ++ float global_max = -CUDART_INF_F; ++ for (int block = 0; block < blocks; ++block) { ++ global_max = fmaxf(global_max, block_max[block * rows + row]); ++ } ++ float denominator = 0.0f; ++ for (int block = 0; block < blocks; ++block) { ++ float scale = exp2f( ++ (block_max[block * rows + row] - global_max) * ++ 1.4426950408889634f); ++#if defined(PREFIX_QK_FULL_STATS) ++ float scaled_sum = scale / block_sum[block * rows + row]; ++ weights[block * rows + row] = scaled_sum; ++ denominator += scaled_sum; ++#else ++ weights[block * rows + row] = scale; ++ denominator += scale * block_sum[block * rows + row]; ++#endif ++ } ++ float inverse = 1.0f / denominator; ++ for (int block = 0; block < blocks; ++block) { ++ weights[block * rows + row] *= inverse; ++ } ++ state_max[row] = global_max; ++ state_sum[row] = denominator; ++} ++ ++__global__ void merge_partials( ++ __half const* partials, ++ float const* weights, ++ __half* output, ++ int blocks, ++ int rows) { ++ int pair = blockIdx.x * blockDim.x + threadIdx.x; ++ constexpr int kPairsPerRow = 128; ++ int total_pairs = rows * kPairsPerRow; ++ if (pair >= total_pairs) { ++ return; ++ } ++ int row = pair / kPairsPerRow; ++ int column_pair = pair - row * kPairsPerRow; ++ float2 accumulator = make_float2(0.0f, 0.0f); ++ auto const* partials2 = reinterpret_cast<__half2 const*>(partials); ++ auto* output2 = reinterpret_cast<__half2*>(output); ++ for (int block = 0; block < blocks; ++block) { ++ int index = (block * rows + row) * kPairsPerRow + column_pair; ++ float2 value = __half22float2(partials2[index]); ++ float weight = weights[block * rows + row]; ++ accumulator.x = fmaf(weight, value.x, accumulator.x); ++ accumulator.y = fmaf(weight, value.y, accumulator.y); ++ } ++ output2[pair] = __floats2half2_rn(accumulator.x, accumulator.y); ++} ++ ++__global__ void update_prefix_accumulator( ++ __half const* block_output, ++ float const* block_max, ++ float const* block_inv_sum, ++ float* prefix_accumulator, ++ float* prefix_max, ++ float* prefix_sum, ++ int rows, ++ bool initialize) { ++ int row = blockIdx.x; ++ int d = threadIdx.x; ++ if (row >= rows || d >= 256) { ++ return; ++ } ++ __shared__ float scales[3]; ++ if (d == 0) { ++ constexpr float kLog2E = 1.4426950408889634f; ++ float next_max = block_max[row]; ++ float next_mass = 1.0f / block_inv_sum[row]; ++ if (initialize) { ++ scales[0] = 0.0f; ++ scales[1] = next_mass; ++ scales[2] = next_mass; ++ prefix_max[row] = next_max; ++ } else { ++ float old_max = prefix_max[row]; ++ float global_max = fmaxf(old_max, next_max); ++ float old_scale = exp2f((old_max - global_max) * kLog2E); ++ float next_scale = exp2f((next_max - global_max) * kLog2E); ++ scales[0] = old_scale; ++ scales[1] = next_mass * next_scale; ++ scales[2] = prefix_sum[row] * old_scale + scales[1]; ++ prefix_max[row] = global_max; ++ } ++ prefix_sum[row] = scales[2]; ++ } ++ __syncthreads(); ++ int64_t element = int64_t(row) * 256 + d; ++ float block_value = float(block_output[element]); ++ float old_value = initialize ? 0.0f : prefix_accumulator[element]; ++ prefix_accumulator[element] = ++ fmaf(old_value, scales[0], block_value * scales[1]); ++} ++ ++__global__ void prepare_prefix_update( ++ float const* block_max, ++ float const* block_inv_sum, ++ float* prefix_max, ++ float* prefix_sum, ++ float* old_scales, ++ float* block_scales, ++ int rows, ++ bool initialize) { ++ int row = blockIdx.x * blockDim.x + threadIdx.x; ++ if (row >= rows) { ++ return; ++ } ++ constexpr float kLog2E = 1.4426950408889634f; ++ float next_max = block_max[row]; ++ float next_mass = 1.0f / block_inv_sum[row]; ++ if (initialize) { ++ old_scales[row] = 0.0f; ++ block_scales[row] = next_mass; ++ prefix_max[row] = next_max; ++ prefix_sum[row] = next_mass; ++ return; ++ } ++ float old_max = prefix_max[row]; ++ float global_max = fmaxf(old_max, next_max); ++ float old_scale = exp2f((old_max - global_max) * kLog2E); ++ float block_scale = ++ next_mass * exp2f((next_max - global_max) * kLog2E); ++ old_scales[row] = old_scale; ++ block_scales[row] = block_scale; ++ prefix_max[row] = global_max; ++ prefix_sum[row] = prefix_sum[row] * old_scale + block_scale; ++} ++ ++__global__ void apply_prefix_update_half2( ++ __half const* block_output, ++ float* prefix_accumulator, ++ float const* old_scales, ++ float const* block_scales, ++ int rows, ++ bool initialize) { ++ int pair = blockIdx.x * blockDim.x + threadIdx.x; ++ constexpr int kPairsPerRow = 128; ++ if (pair >= rows * kPairsPerRow) { ++ return; ++ } ++ int row = pair / kPairsPerRow; ++ auto const* block2 = reinterpret_cast<__half2 const*>(block_output); ++ auto* accumulator2 = reinterpret_cast(prefix_accumulator); ++ float2 block_value = __half22float2(block2[pair]); ++ float2 old_value = initialize ++ ? make_float2(0.0f, 0.0f) ++ : accumulator2[pair]; ++ float old_scale = old_scales[row]; ++ float block_scale = block_scales[row]; ++ accumulator2[pair] = make_float2( ++ fmaf(old_value.x, old_scale, block_value.x * block_scale), ++ fmaf(old_value.y, old_scale, block_value.y * block_scale)); ++} ++ ++__global__ void finalize_prefix_accumulator( ++ float const* prefix_accumulator, ++ float const* prefix_sum, ++ __half* output, ++ int rows) { ++ int pair = blockIdx.x * blockDim.x + threadIdx.x; ++ constexpr int kPairsPerRow = 128; ++ if (pair >= rows * kPairsPerRow) { ++ return; ++ } ++ int row = pair / kPairsPerRow; ++ int64_t first = int64_t(pair) * 2; ++ float inverse = 1.0f / prefix_sum[row]; ++ output[first] = __float2half_rn(prefix_accumulator[first] * inverse); ++ output[first + 1] = ++ __float2half_rn(prefix_accumulator[first + 1] * inverse); ++} ++ ++__global__ void merge_prefix_accumulator_tail( ++ float const* prefix_accumulator, ++ float const* prefix_max, ++ float const* prefix_sum, ++ __half const* tail_output, ++ float const* tail_max, ++ float const* tail_sum, ++ __half* output, ++ int rows) { ++ int row = blockIdx.x; ++ int d = threadIdx.x; ++ __shared__ float masses[3]; ++ if (d == 0) { ++ constexpr float kLog2E = 1.4426950408889634f; ++ float global_max = fmaxf(prefix_max[row], tail_max[row]); ++ masses[0] = exp2f((prefix_max[row] - global_max) * kLog2E); ++ masses[1] = tail_sum[row] ++ * exp2f((tail_max[row] - global_max) * kLog2E); ++ masses[2] = 1.0f / (prefix_sum[row] * masses[0] + masses[1]); ++ } ++ __syncthreads(); ++ int64_t element = int64_t(row) * 256 + d; ++ float numerator = prefix_accumulator[element] * masses[0] ++ + float(tail_output[element]) * masses[1]; ++ output[element] = __float2half_rn(numerator * masses[2]); ++} ++ ++__global__ void merge_prefix_tail( ++ __half const* prefix_output, ++ float const* prefix_max, ++ float const* prefix_sum, ++ __half const* tail_output, ++ float const* tail_max, ++ float const* tail_sum, ++ __half* output, ++ int rows) { ++ int row = blockIdx.x; ++ int d = threadIdx.x; ++ __shared__ float masses[3]; ++ if (d == 0) { ++ float global_max = fmaxf(prefix_max[row], tail_max[row]); ++ float prefix_scale = exp2f( ++ (prefix_max[row] - global_max) * 1.4426950408889634f); ++ float tail_scale = exp2f( ++ (tail_max[row] - global_max) * 1.4426950408889634f); ++ masses[0] = prefix_sum[row] * prefix_scale; ++ masses[1] = tail_sum[row] * tail_scale; ++ masses[2] = 1.0f / (masses[0] + masses[1]); ++ } ++ __syncthreads(); ++ int64_t element = int64_t(row) * 256 + d; ++ float numerator = float(prefix_output[element]) * masses[0] ++ + float(tail_output[element]) * masses[1]; ++ output[element] = __float2half_rn(numerator * masses[2]); ++} ++ ++using DenseTailRaw = cudaError_t (*)( ++ const void*, ++ const void*, ++ const void*, ++ float*, ++ float*, ++ void*, ++ int, ++ int, ++ int, ++ int, ++ float, ++ cudaStream_t); ++ ++__global__ void fill_pattern(__half* data, size_t elements, uint32_t seed) { ++ size_t index = size_t(blockIdx.x) * blockDim.x + threadIdx.x; ++ if (index >= elements) { ++ return; ++ } ++ uint32_t value = uint32_t(index) ^ seed; ++ value ^= value >> 16; ++ value *= 0x7feb352du; ++ value ^= value >> 15; ++ value *= 0x846ca68bu; ++ value ^= value >> 16; ++ float uniform = float(value & 0xffffu) * (2.0f / 65535.0f) - 1.0f; ++ data[index] = __float2half_rn(uniform); ++} ++ ++struct BlockOperators { ++ int width; ++ QKGemm qk; ++ std::unique_ptr pv; ++}; ++ ++} // namespace ++ ++#if !defined(PREFIX_TORCH_EXTENSION) ++int main(int argc, char** argv) { ++ int block_n = argc > 1 ? std::atoi(argv[1]) : 2048; ++ int iterations = argc > 2 ? std::atoi(argv[2]) : 5; ++#if defined(PREFIX_FULL_ENDPOINT) ++ if (argc <= 3) { ++ std::cerr << "full endpoint requires the exact-tail extension path\n"; ++ return EXIT_FAILURE; ++ } ++#endif ++ constexpr int kRows = 48000; ++ constexpr int kHeadDim = 256; ++ constexpr int kPrefix = 120000; ++ constexpr int kTail = 8000; ++ constexpr int kTotalKV = kPrefix + kTail; ++#if defined(PREFIX_FULL_ENDPOINT) ++ constexpr int kStoredKV = kTotalKV; ++#else ++ constexpr int kStoredKV = kPrefix; ++#endif ++ int blocks = (kPrefix + block_n - 1) / block_n; ++ ++ Element* query = nullptr; ++ Element* key = nullptr; ++ Element* value = nullptr; ++ Element* scores = nullptr; ++ Element* partials = nullptr; ++ Element* output = nullptr; ++ float* qk_norm = nullptr; ++ float* qk_sum = nullptr; ++ float* prefix_accumulator = nullptr; ++ float* prefix_max = nullptr; ++ float* prefix_sum = nullptr; ++ float* old_scales = nullptr; ++ float* block_scales = nullptr; ++#if defined(PREFIX_FULL_ENDPOINT) ++ Element* tail_output = nullptr; ++ Element* final_output = nullptr; ++ Element* exact_output = nullptr; ++ float* tail_max = nullptr; ++ float* tail_sum = nullptr; ++#endif ++ ++ int qk_tiles_n = (block_n + QKThreadblockShape::kN - 1) / ++ QKThreadblockShape::kN; ++ size_t partial_elements = size_t(kRows) * kHeadDim; ++ check(cudaMalloc(&query, size_t(kRows) * kHeadDim * sizeof(Element)), "allocate Q"); ++ check(cudaMalloc(&key, size_t(kStoredKV) * kHeadDim * sizeof(Element)), "allocate K"); ++ check(cudaMalloc(&value, size_t(kStoredKV) * kHeadDim * sizeof(Element)), "allocate V"); ++ check(cudaMalloc(&scores, size_t(kRows) * block_n * sizeof(Element)), "allocate scores"); ++ check(cudaMalloc(&partials, partial_elements * sizeof(Element)), "allocate partials"); ++ check(cudaMalloc(&output, size_t(kRows) * kHeadDim * sizeof(Element)), "allocate output"); ++ check(cudaMalloc(&qk_norm, size_t(qk_tiles_n) * kRows * sizeof(float)), "allocate QK max"); ++ check(cudaMalloc(&qk_sum, size_t(qk_tiles_n) * kRows * sizeof(float)), "allocate QK sum"); ++ check(cudaMalloc( ++ &prefix_accumulator, ++ size_t(kRows) * kHeadDim * sizeof(float)), ++ "allocate prefix accumulator"); ++ check(cudaMalloc(&prefix_max, size_t(kRows) * sizeof(float)), "allocate prefix max"); ++ check(cudaMalloc(&prefix_sum, size_t(kRows) * sizeof(float)), "allocate prefix sum"); ++ check(cudaMalloc(&old_scales, size_t(kRows) * sizeof(float)), ++ "allocate old scales"); ++ check(cudaMalloc(&block_scales, size_t(kRows) * sizeof(float)), ++ "allocate block scales"); ++#if defined(PREFIX_FULL_ENDPOINT) ++ check(cudaMalloc(&tail_output, size_t(kRows) * kHeadDim * sizeof(Element)), ++ "allocate tail output"); ++ check(cudaMalloc(&final_output, size_t(kRows) * kHeadDim * sizeof(Element)), ++ "allocate final output"); ++ check(cudaMalloc(&exact_output, size_t(kRows) * kHeadDim * sizeof(Element)), ++ "allocate exact output"); ++ check(cudaMalloc(&tail_max, size_t(kRows) * sizeof(float)), "allocate tail max"); ++ check(cudaMalloc(&tail_sum, size_t(kRows) * sizeof(float)), "allocate tail sum"); ++#endif ++ size_t query_elements = size_t(kRows) * kHeadDim; ++ size_t kv_elements = size_t(kStoredKV) * kHeadDim; ++ fill_pattern<<<(query_elements + 255) / 256, 256>>>( ++ reinterpret_cast<__half*>(query), query_elements, 0x12345678u); ++ fill_pattern<<<(kv_elements + 255) / 256, 256>>>( ++ reinterpret_cast<__half*>(key), kv_elements, 0x9abcdef0u); ++ fill_pattern<<<(kv_elements + 255) / 256, 256>>>( ++ reinterpret_cast<__half*>(value), kv_elements, 0x31415926u); ++ check(cudaGetLastError(), "initialize deterministic inputs"); ++ ++#if defined(PREFIX_FULL_ENDPOINT) ++ void* tail_handle = dlopen(argv[3], RTLD_NOW | RTLD_LOCAL); ++ if (tail_handle == nullptr) { ++ std::cerr << "load exact-tail extension: " << dlerror() << "\n"; ++ return EXIT_FAILURE; ++ } ++ dlerror(); ++ auto tail_raw = reinterpret_cast( ++ dlsym(tail_handle, "onecat_sm70_d256_dense_state_raw")); ++ char const* symbol_error = dlerror(); ++ if (symbol_error != nullptr || tail_raw == nullptr) { ++ std::cerr << "load exact-tail entry point: " ++ << (symbol_error == nullptr ? "missing symbol" : symbol_error) ++ << "\n"; ++ return EXIT_FAILURE; ++ } ++#endif ++ ++ std::vector operators; ++ operators.reserve(blocks); ++ for (int block = 0; block < blocks; ++block) { ++ int begin = block * block_n; ++ int width = std::min(block_n, kPrefix - begin); ++ BlockOperators operation; ++ operation.width = width; ++ typename QKGemm::Arguments arguments( ++ {kRows, width, kHeadDim}, ++ 1, ++ {query, QKLayoutA(kHeadDim)}, ++ {key + size_t(begin) * kHeadDim, QKLayoutB(kHeadDim)}, ++ {scores, typename QKGemm::LayoutC(width)}, ++ {scores, typename QKGemm::LayoutC(width)}, ++ {Element(0.0625f), Element(0.0f)}, ++ {qk_norm, typename QKGemm::LayoutN(kRows)}, ++ {qk_sum, typename QKGemm::LayoutS(kRows)}, ++ {scores, typename QKGemm::LayoutSoft(width)}); ++ check(operation.qk.initialize(arguments), "initialize QK"); ++ operation.pv = std::make_unique( ++ scores, ++ value + size_t(begin) * kHeadDim, ++ partials, ++ kRows, ++ width); ++ operators.push_back(std::move(operation)); ++ } ++ ++ check(cudaMemcpyToSymbol(g_rows, &kRows, sizeof(kRows)), "set row count"); ++ float const* persistent_max_ptr = qk_norm; ++ check(cudaMemcpyToSymbol( ++ g_row_max, ++ &persistent_max_ptr, ++ sizeof(persistent_max_ptr)), ++ "set persistent PV max pointer"); ++#if defined(PREFIX_QK_FULL_STATS) ++ float const* persistent_inv_sum_ptr = qk_sum; ++ check(cudaMemcpyToSymbol( ++ g_row_inv_sum, ++ &persistent_inv_sum_ptr, ++ sizeof(persistent_inv_sum_ptr)), ++ "set persistent PV inverse-sum pointer"); ++#endif ++ cudaStream_t stream = nullptr; ++ auto launch = [&](cudaEvent_t blocks_done, ++ cudaEvent_t prefix_done, ++ cudaEvent_t tail_done) { ++ for (int block = 0; block < blocks; ++block) { ++ check(operators[block].qk(stream), "launch QK max"); ++#if !defined(PREFIX_QK_FULL_STATS) ++ float* sum_ptr = block_sum + size_t(block) * kRows; ++ check(cudaMemcpyToSymbolAsync( ++ g_row_sum_out, ++ &sum_ptr, ++ sizeof(sum_ptr), ++ 0, ++ cudaMemcpyHostToDevice, ++ stream), ++ "set PV sum pointer"); ++#endif ++ operators[block].pv->launch(stream); ++ prepare_prefix_update<<<(kRows + 255) / 256, 256, 0, stream>>>( ++ qk_norm, ++ qk_sum, ++ prefix_max, ++ prefix_sum, ++ old_scales, ++ block_scales, ++ kRows, ++ block == 0); ++ int pairs = kRows * kHeadDim / 2; ++ apply_prefix_update_half2<<<(pairs + 255) / 256, 256, 0, stream>>>( ++ reinterpret_cast<__half const*>(partials), ++ prefix_accumulator, ++ old_scales, ++ block_scales, ++ kRows, ++ block == 0); ++ } ++ if (blocks_done != nullptr) { ++ check(cudaEventRecord(blocks_done, stream), "record QK/PV blocks done"); ++ } ++#if !defined(PREFIX_FULL_ENDPOINT) ++ int pairs = kRows * kHeadDim / 2; ++ finalize_prefix_accumulator<<<(pairs + 255) / 256, 256, 0, stream>>>( ++ prefix_accumulator, ++ prefix_sum, ++ reinterpret_cast<__half*>(output), ++ kRows); ++#endif ++ if (prefix_done != nullptr) { ++ check(cudaEventRecord(prefix_done, stream), "record prefix merge done"); ++ } ++#if defined(PREFIX_FULL_ENDPOINT) ++ check(tail_raw( ++ query, ++ key + size_t(kPrefix) * kHeadDim, ++ value + size_t(kPrefix) * kHeadDim, ++ tail_max, ++ tail_sum, ++ tail_output, ++ kTail, ++ kTail, ++ 6, ++ 1, ++ 0.0625f, ++ stream), ++ "launch exact causal tail"); ++ if (tail_done != nullptr) { ++ check(cudaEventRecord(tail_done, stream), "record causal tail done"); ++ } ++ merge_prefix_accumulator_tail<<>>( ++ prefix_accumulator, ++ prefix_max, ++ prefix_sum, ++ reinterpret_cast<__half const*>(tail_output), ++ tail_max, ++ tail_sum, ++ reinterpret_cast<__half*>(final_output), ++ kRows); ++#endif ++ }; ++ ++ launch(nullptr, nullptr, nullptr); ++ check(cudaDeviceSynchronize(), "warmup synchronize"); ++ cudaEvent_t start; ++ cudaEvent_t stop; ++ check(cudaEventCreate(&start), "create start"); ++ check(cudaEventCreate(&stop), "create stop"); ++ check(cudaEventRecord(start), "record start"); ++ for (int iteration = 0; iteration < iterations; ++iteration) { ++ launch(nullptr, nullptr, nullptr); ++ } ++ check(cudaEventRecord(stop), "record stop"); ++ check(cudaEventSynchronize(stop), "synchronize stop"); ++ float elapsed_ms = 0.0f; ++ check(cudaEventElapsedTime(&elapsed_ms, start, stop), "elapsed time"); ++ float endpoint_ms = elapsed_ms / float(iterations); ++ constexpr double kFrozenUsefulTflop = 60.0 * 0.101581; ++#if defined(PREFIX_FULL_ENDPOINT) ++ cudaEvent_t phase_start; ++ cudaEvent_t blocks_done; ++ cudaEvent_t prefix_done; ++ cudaEvent_t tail_done; ++ cudaEvent_t phase_stop; ++ check(cudaEventCreate(&phase_start), "create phase start"); ++ check(cudaEventCreate(&blocks_done), "create blocks done"); ++ check(cudaEventCreate(&prefix_done), "create prefix done"); ++ check(cudaEventCreate(&tail_done), "create tail done"); ++ check(cudaEventCreate(&phase_stop), "create phase stop"); ++ check(cudaEventRecord(phase_start, stream), "record phase start"); ++ launch(blocks_done, prefix_done, tail_done); ++ check(cudaEventRecord(phase_stop, stream), "record phase stop"); ++ check(cudaEventSynchronize(phase_stop), "synchronize phase stop"); ++ float blocks_ms = 0.0f; ++ float prefix_merge_ms = 0.0f; ++ float tail_ms = 0.0f; ++ float final_merge_ms = 0.0f; ++ check(cudaEventElapsedTime(&blocks_ms, phase_start, blocks_done), ++ "measure QK/PV blocks"); ++ check(cudaEventElapsedTime(&prefix_merge_ms, blocks_done, prefix_done), ++ "measure prefix merge"); ++ check(cudaEventElapsedTime(&tail_ms, prefix_done, tail_done), ++ "measure causal tail"); ++ check(cudaEventElapsedTime(&final_merge_ms, tail_done, phase_stop), ++ "measure final merge"); ++ check(tail_raw( ++ query, ++ key, ++ value, ++ tail_max, ++ tail_sum, ++ exact_output, ++ kTail, ++ kTotalKV, ++ 6, ++ 1, ++ 0.0625f, ++ stream), ++ "launch monolithic exact reference"); ++ check(cudaDeviceSynchronize(), "synchronize exact reference"); ++ size_t output_elements = size_t(kRows) * kHeadDim; ++ std::vector<__half> candidate_host(output_elements); ++ std::vector<__half> reference_host(output_elements); ++ check(cudaMemcpy( ++ candidate_host.data(), ++ final_output, ++ output_elements * sizeof(__half), ++ cudaMemcpyDeviceToHost), ++ "copy candidate output"); ++ check(cudaMemcpy( ++ reference_host.data(), ++ exact_output, ++ output_elements * sizeof(__half), ++ cudaMemcpyDeviceToHost), ++ "copy exact output"); ++ double squared_error = 0.0; ++ double squared_reference = 0.0; ++ double absolute_error = 0.0; ++ float max_absolute_error = 0.0f; ++ size_t equal_elements = 0; ++ for (size_t index = 0; index < output_elements; ++index) { ++ float candidate_value = __half2float(candidate_host[index]); ++ float reference_value = __half2float(reference_host[index]); ++ float difference = candidate_value - reference_value; ++ float absolute = std::fabs(difference); ++ squared_error += double(difference) * difference; ++ squared_reference += double(reference_value) * reference_value; ++ absolute_error += absolute; ++ max_absolute_error = std::max(max_absolute_error, absolute); ++ equal_elements += candidate_value == reference_value; ++ } ++ double relative_l2 = std::sqrt(squared_error / squared_reference); ++ double mean_absolute_error = absolute_error / double(output_elements); ++ double measured_tflops = kFrozenUsefulTflop / (endpoint_ms * 1.0e-3); ++ std::cout << "{\n" ++ << " \"block_n\": " << block_n << ",\n" ++ << " \"blocks\": " << blocks << ",\n" ++ << " \"full_attention_ms\": " << endpoint_ms << ",\n" ++ << " \"useful_causal_tflops\": " << measured_tflops << ",\n" ++ << " \"prefix_tokens\": " << kPrefix << ",\n" ++ << " \"causal_tail_tokens\": " << kTail << ",\n" ++ << " \"phases_ms\": {\n" ++ << " \"prefix_qk_pv_blocks\": " << blocks_ms << ",\n" ++ << " \"prefix_state_and_partial_merge\": " ++ << prefix_merge_ms << ",\n" ++ << " \"exact_causal_tail\": " << tail_ms << ",\n" ++ << " \"prefix_tail_merge\": " << final_merge_ms << "\n" ++ << " },\n" ++ << " \"numerical\": {\n" ++ << " \"max_abs_diff\": " << max_absolute_error << ",\n" ++ << " \"mean_abs_diff\": " << mean_absolute_error << ",\n" ++ << " \"relative_l2\": " << relative_l2 << ",\n" ++ << " \"equal_elements\": " << equal_elements << ",\n" ++ << " \"elements\": " << output_elements << "\n" ++ << " },\n" ++ << " \"partial_storage_gib\": " ++ << double(partial_elements * sizeof(Element)) / double(1ull << 30) ++ << ",\n" ++ << " \"note\": \"measured serial prefix + exact tail state + final merge endpoint\"\n" ++ << "}\n"; ++ check(cudaEventDestroy(phase_start), "destroy phase start"); ++ check(cudaEventDestroy(blocks_done), "destroy blocks done"); ++ check(cudaEventDestroy(prefix_done), "destroy prefix done"); ++ check(cudaEventDestroy(tail_done), "destroy tail done"); ++ check(cudaEventDestroy(phase_stop), "destroy phase stop"); ++#else ++ double projected_full_ms = endpoint_ms + 4.83; ++ double projected_tflops = kFrozenUsefulTflop / (projected_full_ms * 1.0e-3); ++ std::cout << "{\n" ++ << " \"block_n\": " << block_n << ",\n" ++ << " \"blocks\": " << blocks << ",\n" ++ << " \"prefix_ms_including_merge\": " << endpoint_ms << ",\n" ++ << " \"accepted_tail_ms\": 4.83,\n" ++ << " \"projected_full_ms\": " << projected_full_ms << ",\n" ++ << " \"projected_useful_tflops\": " << projected_tflops << ",\n" ++ << " \"partial_storage_gib\": " ++ << double(partial_elements * sizeof(Element)) / double(1ull << 30) ++ << ",\n" ++ << " \"note\": \"prefix endpoint; prefix/tail state merge excluded\"\n" ++ << "}\n"; ++#endif ++ ++ check(cudaEventDestroy(start), "destroy start"); ++ check(cudaEventDestroy(stop), "destroy stop"); ++ check(cudaFree(query), "free Q"); ++ check(cudaFree(key), "free K"); ++ check(cudaFree(value), "free V"); ++ check(cudaFree(scores), "free scores"); ++ check(cudaFree(partials), "free partials"); ++ check(cudaFree(output), "free output"); ++ check(cudaFree(qk_norm), "free QK max"); ++ check(cudaFree(qk_sum), "free QK sum"); ++ check(cudaFree(prefix_accumulator), "free prefix accumulator"); ++ check(cudaFree(prefix_max), "free prefix max"); ++ check(cudaFree(prefix_sum), "free prefix sum"); ++ check(cudaFree(old_scales), "free old scales"); ++ check(cudaFree(block_scales), "free block scales"); ++#if defined(PREFIX_FULL_ENDPOINT) ++ check(cudaFree(tail_output), "free tail output"); ++ check(cudaFree(final_output), "free final output"); ++ check(cudaFree(exact_output), "free exact output"); ++ check(cudaFree(tail_max), "free tail max"); ++ check(cudaFree(tail_sum), "free tail sum"); ++ dlclose(tail_handle); ++#endif ++ return 0; ++} ++#else ++ ++extern "C" cudaError_t onecat_sm70_d256_dense_state_raw( ++ const void*, ++ const void*, ++ const void*, ++ float*, ++ float*, ++ void*, ++ int, ++ int, ++ int, ++ int, ++ float, ++ cudaStream_t); ++ ++namespace FLASH_NAMESPACE { ++ ++struct Sm70GqaScoreWorkspace { ++ at::Tensor scores; ++ cudaEvent_t completion = nullptr; ++ bool completion_recorded = false; ++ std::mutex launch_mutex; ++ ++ Sm70GqaScoreWorkspace(const at::Tensor& q, int rows, int block_n) ++ : scores(at::empty({rows, block_n}, q.options())) { ++ C10_CUDA_CHECK( ++ cudaEventCreateWithFlags(&completion, cudaEventDisableTiming)); ++ } ++ ++ ~Sm70GqaScoreWorkspace() { ++ if (completion != nullptr) { ++ cudaEventDestroy(completion); ++ } ++ } ++}; ++ ++using Sm70GqaScoreWorkspacePtr = std::shared_ptr; ++ ++std::mutex& sm70_gqa_score_cache_mutex() { ++ static std::mutex cache_mutex; ++ return cache_mutex; ++} ++ ++std::map& sm70_gqa_score_cache() { ++ static std::map cache; ++ return cache; ++} ++ ++Sm70GqaScoreWorkspacePtr get_sm70_gqa_score_workspace( ++ const at::Tensor& q, ++ int rows, ++ int block_n) { ++ std::lock_guard lock(sm70_gqa_score_cache_mutex()); ++ int device = q.get_device(); ++ auto& cache = sm70_gqa_score_cache(); ++ auto& workspace = cache[device]; ++ if (!workspace) { ++ workspace = ++ std::make_shared(q, rows, block_n); ++ } ++ return workspace; ++} ++ ++Sm70GqaScoreWorkspacePtr find_sm70_gqa_score_workspace(int device) { ++ std::lock_guard lock(sm70_gqa_score_cache_mutex()); ++ auto& cache = sm70_gqa_score_cache(); ++ auto found = cache.find(device); ++ return found == cache.end() ? nullptr : found->second; ++} ++ ++at::Tensor sm70_d256_gqa_architecture_fwd( ++ const at::Tensor& q, ++ const at::Tensor& k, ++ const at::Tensor& v, ++ at::Tensor& out, ++ double softmax_scale, ++ bool causal) { ++ constexpr int kQuery = 8000; ++ constexpr int kRows = 48000; ++ constexpr int kHeadDim = 256; ++ constexpr int kHeadsQ = 6; ++ constexpr int kHeadsKV = 1; ++ constexpr int kTail = 8000; ++ constexpr int kMinTotalKV = 40000; ++ constexpr int kMaxTotalKV = 128000; ++ constexpr int kTotalKVStep = 8000; ++ constexpr int kBlockN = 8192; ++ ++ TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && out.is_cuda(), ++ "SM70 GQA architecture requires CUDA tensors"); ++ TORCH_CHECK(q.scalar_type() == at::ScalarType::Half ++ && k.scalar_type() == q.scalar_type() ++ && v.scalar_type() == q.scalar_type() ++ && out.scalar_type() == q.scalar_type(), ++ "SM70 GQA architecture requires FP16 q, k, v, and out"); ++ TORCH_CHECK(q.sizes() == at::IntArrayRef({1, kQuery, kHeadsQ, kHeadDim}) ++ && k.dim() == 4 && k.size(0) == 1 ++ && k.size(2) == kHeadsKV && k.size(3) == kHeadDim ++ && v.sizes() == k.sizes() && out.sizes() == q.sizes(), ++ "SM70 GQA architecture only accepts the validated " ++ "Q8000/Hq6/Hkv1/D256 dense shape family"); ++ const int total_kv = static_cast(k.size(1)); ++ TORCH_CHECK(total_kv >= kMinTotalKV && total_kv <= kMaxTotalKV ++ && total_kv % kTotalKVStep == 0, ++ "SM70 GQA architecture requires KV in [40000, 128000] " ++ "with an 8000-token step, got ", ++ total_kv); ++ const int prefix = total_kv - kTail; ++ const int blocks = (prefix + kBlockN - 1) / kBlockN; ++ TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous() ++ && out.is_contiguous(), ++ "SM70 GQA architecture requires contiguous tensors"); ++ TORCH_CHECK(q.get_device() == k.get_device() ++ && q.get_device() == v.get_device() ++ && q.get_device() == out.get_device(), ++ "SM70 GQA architecture tensors must share one device"); ++ TORCH_CHECK(causal, "SM70 GQA architecture requires causal attention"); ++ TORCH_CHECK(std::abs(softmax_scale - 0.0625) < 1.0e-8, ++ "SM70 GQA architecture requires D256 softmax scale 1/16"); ++ ++ const at::cuda::OptionalCUDAGuard device_guard(q.device()); ++ const int qk_tiles_n = ++ (kBlockN + QKThreadblockShape::kN - 1) ++ / QKThreadblockShape::kN; ++ cudaStream_t stream = at::cuda::getCurrentCUDAStream(); ++ ++ constexpr size_t kScratchAlignment = 256; ++ size_t scratch_bytes = 0; ++ auto reserve_scratch = [&](size_t bytes) { ++ scratch_bytes = ++ (scratch_bytes + kScratchAlignment - 1) & ~(kScratchAlignment - 1); ++ size_t offset = scratch_bytes; ++ scratch_bytes += bytes; ++ return offset; ++ }; ++ size_t partial_offset = reserve_scratch( ++ size_t(kRows) * kHeadDim * sizeof(Element)); ++ size_t prefix_accumulator_offset = reserve_scratch( ++ size_t(kRows) * kHeadDim * sizeof(float)); ++ size_t tail_output_offset = reserve_scratch( ++ size_t(kRows) * kHeadDim * sizeof(Element)); ++ size_t qk_norm_offset = reserve_scratch( ++ size_t(qk_tiles_n) * kRows * sizeof(float)); ++ size_t qk_sum_offset = reserve_scratch( ++ size_t(qk_tiles_n) * kRows * sizeof(float)); ++ size_t prefix_max_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ size_t prefix_sum_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ size_t old_scale_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ size_t block_scale_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ size_t tail_max_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ size_t tail_sum_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ int device = q.get_device(); ++ auto score_workspace = find_sm70_gqa_score_workspace(device); ++ if (!score_workspace) { ++ constexpr size_t kRequiredPostWorkspaceHeadroom = 128 * 1024 * 1024; ++ constexpr size_t kScoreBytes = ++ size_t(kRows) * kBlockN * sizeof(Element); ++ size_t free_bytes = 0; ++ size_t total_bytes = 0; ++ C10_CUDA_CHECK(cudaMemGetInfo(&free_bytes, &total_bytes)); ++ size_t required_bytes = ++ kScoreBytes + scratch_bytes + kRequiredPostWorkspaceHeadroom; ++ TORCH_CHECK_WITH( ++ OutOfMemoryError, ++ free_bytes >= required_bytes, ++ "SM70 GQA architecture requires ", ++ required_bytes / (1024 * 1024), ++ " MiB free before its first workspace allocation, including " ++ "128 MiB of downstream headroom, but only ", ++ free_bytes / (1024 * 1024), ++ " MiB remains out of ", ++ total_bytes / (1024 * 1024), ++ " MiB"); ++ score_workspace = ++ get_sm70_gqa_score_workspace(q, kRows, kBlockN); ++ } ++ std::unique_lock launch_lock(score_workspace->launch_mutex); ++ if (score_workspace->completion_recorded) { ++ C10_CUDA_CHECK( ++ cudaStreamWaitEvent(stream, score_workspace->completion, 0)); ++ } ++ at::Tensor scratch = at::empty( ++ {static_cast(scratch_bytes)}, ++ q.options().dtype(at::ScalarType::Byte)); ++ auto* scratch_base = scratch.data_ptr(); ++ ++ auto* query = reinterpret_cast(q.data_ptr()); ++ auto* key = reinterpret_cast(k.data_ptr()); ++ auto* value = reinterpret_cast(v.data_ptr()); ++ auto* score_ptr = ++ reinterpret_cast(score_workspace->scores.data_ptr()); ++ auto* partial_ptr = ++ reinterpret_cast(scratch_base + partial_offset); ++ float* prefix_accumulator_ptr = ++ reinterpret_cast(scratch_base + prefix_accumulator_offset); ++ auto* tail_output_ptr = ++ reinterpret_cast(scratch_base + tail_output_offset); ++ auto* output_ptr = reinterpret_cast(out.data_ptr()); ++ float* qk_norm_ptr = ++ reinterpret_cast(scratch_base + qk_norm_offset); ++ float* qk_sum_ptr = ++ reinterpret_cast(scratch_base + qk_sum_offset); ++ float* prefix_max_ptr = ++ reinterpret_cast(scratch_base + prefix_max_offset); ++ float* prefix_sum_ptr = ++ reinterpret_cast(scratch_base + prefix_sum_offset); ++ float* old_scale_ptr = ++ reinterpret_cast(scratch_base + old_scale_offset); ++ float* block_scale_ptr = ++ reinterpret_cast(scratch_base + block_scale_offset); ++ float* tail_max_ptr = ++ reinterpret_cast(scratch_base + tail_max_offset); ++ float* tail_sum_ptr = ++ reinterpret_cast(scratch_base + tail_sum_offset); ++ ++ C10_CUDA_CHECK(onecat_sm70_d256_dense_state_raw( ++ query, ++ key + size_t(prefix) * kHeadDim, ++ value + size_t(prefix) * kHeadDim, ++ tail_max_ptr, ++ tail_sum_ptr, ++ tail_output_ptr, ++ kTail, ++ kTail, ++ kHeadsQ, ++ kHeadsKV, ++ static_cast(softmax_scale), ++ stream)); ++ C10_CUDA_CHECK(cudaMemcpyToSymbolAsync( ++ g_rows, ++ &kRows, ++ sizeof(kRows), ++ 0, ++ cudaMemcpyHostToDevice, ++ stream)); ++ float const* persistent_max_ptr = qk_norm_ptr; ++ C10_CUDA_CHECK(cudaMemcpyToSymbolAsync( ++ g_row_max, ++ &persistent_max_ptr, ++ sizeof(persistent_max_ptr), ++ 0, ++ cudaMemcpyHostToDevice, ++ stream)); ++ float const* persistent_inv_sum_ptr = qk_sum_ptr; ++ C10_CUDA_CHECK(cudaMemcpyToSymbolAsync( ++ g_row_inv_sum, ++ &persistent_inv_sum_ptr, ++ sizeof(persistent_inv_sum_ptr), ++ 0, ++ cudaMemcpyHostToDevice, ++ stream)); ++ ++ for (int block = 0; block < blocks; ++block) { ++ int begin = block * kBlockN; ++ int width = std::min(kBlockN, prefix - begin); ++ BlockOperators operation; ++ operation.width = width; ++ typename QKGemm::Arguments arguments( ++ {kRows, width, kHeadDim}, ++ 1, ++ {query, QKLayoutA(kHeadDim)}, ++ {key + size_t(begin) * kHeadDim, QKLayoutB(kHeadDim)}, ++ {score_ptr, typename QKGemm::LayoutC(width)}, ++ {score_ptr, typename QKGemm::LayoutC(width)}, ++ {Element(static_cast(softmax_scale)), Element(0.0f)}, ++ {qk_norm_ptr, typename QKGemm::LayoutN(kRows)}, ++ {qk_sum_ptr, typename QKGemm::LayoutS(kRows)}, ++ {score_ptr, typename QKGemm::LayoutSoft(width)}); ++ TORCH_CHECK(operation.qk.initialize(arguments) == cutlass::Status::kSuccess, ++ "initialize SM70 GQA QK block ", block, " failed"); ++ operation.pv = std::make_unique( ++ score_ptr, ++ value + size_t(begin) * kHeadDim, ++ partial_ptr, ++ kRows, ++ width); ++ TORCH_CHECK(operation.qk(stream) == cutlass::Status::kSuccess, ++ "launch SM70 GQA QK block ", block, " failed"); ++ operation.pv->launch(stream); ++ prepare_prefix_update<<<(kRows + 255) / 256, 256, 0, stream>>>( ++ qk_norm_ptr, ++ qk_sum_ptr, ++ prefix_max_ptr, ++ prefix_sum_ptr, ++ old_scale_ptr, ++ block_scale_ptr, ++ kRows, ++ block == 0); ++ constexpr int kPairs = kRows * kHeadDim / 2; ++ apply_prefix_update_half2<<<(kPairs + 255) / 256, 256, 0, stream>>>( ++ reinterpret_cast<__half const*>(partial_ptr), ++ prefix_accumulator_ptr, ++ old_scale_ptr, ++ block_scale_ptr, ++ kRows, ++ block == 0); ++ } ++ merge_prefix_accumulator_tail<<>>( ++ prefix_accumulator_ptr, ++ prefix_max_ptr, ++ prefix_sum_ptr, ++ reinterpret_cast<__half const*>(tail_output_ptr), ++ tail_max_ptr, ++ tail_sum_ptr, ++ reinterpret_cast<__half*>(output_ptr), ++ kRows); ++ C10_CUDA_KERNEL_LAUNCH_CHECK(); ++ C10_CUDA_CHECK(cudaEventRecord(score_workspace->completion, stream)); ++ score_workspace->completion_recorded = true; ++ return out; ++} ++ ++} // namespace FLASH_NAMESPACE ++ ++#endif +diff --git a/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu b/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu +index 7a0f626..1fd4008 100644 +--- a/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu ++++ b/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu +@@ -398,7 +398,11 @@ __device__ __forceinline__ int64_t paged_kv_thread_offset( + + (tid % kThreadsPerRow) * kElemsPerLoad; + } + +-template ++template < ++ typename Element, ++ bool PagedKV, ++ bool SplitKV3 = false, ++ bool StoreState = false> + __global__ __launch_bounds__(Sm70D256SplitDTraits::kNThreads, 1) + void sm70_d256_splitd_dense_kernel( + const Element *__restrict__ q, +@@ -859,6 +863,29 @@ void sm70_d256_splitd_dense_kernel( + } + } + } else { ++ if constexpr (StoreState) { ++ if ((lane & 0x0e) == 0) { ++#pragma unroll ++ for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { ++ const int row = FLASH_NAMESPACE::sm70_row_slot< ++ Traits::kQkWarpRows>(slot, lane); ++ const int query_row = query_row_base ++ + group_row_base ++ + n_warp * Traits::kQkWarpRows + row; ++ const int64_t state_row = ++ (static_cast(batch) * query_len + query_row) ++ * heads_q ++ + head_q; ++ // The standalone prefix path stores maxima after applying ++ // softmax_scale. Export the dense-tail state in the same ++ // natural-log coordinate so the two states can be merged ++ // directly with exp(prefix_max - global_max). ++ partial_max[state_row] = row_max[slot] ++ * (softmax_scale_log2 * float(M_LN2)); ++ partial_sum[state_row] = row_sum[slot]; ++ } ++ } ++ } + if ((lane & 0x0e) == 0) { + #pragma unroll + for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { +@@ -1042,6 +1069,154 @@ at::Tensor sm70_d256_splitd_dense_fwd( + return out; + } + ++at::Tensor sm70_d256_splitd_dense_state_fwd( ++ const at::Tensor &q, ++ const at::Tensor &k, ++ const at::Tensor &v, ++ at::Tensor &state_max, ++ at::Tensor &state_sum, ++ at::Tensor &out, ++ double softmax_scale, ++ bool causal) { ++ TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() ++ && state_max.is_cuda() && state_sum.is_cuda() ++ && out.is_cuda(), ++ "dense-state inputs must be CUDA tensors"); ++ TORCH_CHECK(q.scalar_type() == at::ScalarType::Half ++ && k.scalar_type() == q.scalar_type() ++ && v.scalar_type() == q.scalar_type() ++ && out.scalar_type() == q.scalar_type(), ++ "dense-state requires FP16 q, k, v, and out"); ++ TORCH_CHECK(state_max.scalar_type() == at::ScalarType::Float ++ && state_sum.scalar_type() == at::ScalarType::Float, ++ "dense-state statistics must use FP32"); ++ TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4 ++ && q.size(0) == k.size(0) && k.sizes() == v.sizes(), ++ "dense-state q, k, and v shapes are invalid"); ++ TORCH_CHECK(q.size(3) == Sm70D256SplitDTraits::kHeadDim ++ && k.size(3) == Sm70D256SplitDTraits::kHeadDim ++ && q.size(2) % k.size(2) == 0, ++ "dense-state requires D256 and valid GQA heads"); ++ TORCH_CHECK(q.size(1) <= k.size(1) ++ && q.size(1) % Sm70D256SplitDTraits::kBlockM == 0 ++ && k.size(1) % Sm70D256SplitDTraits::kBlockN == 0, ++ "dense-state lengths do not match the exact tiles"); ++ TORCH_CHECK(out.sizes() == q.sizes() && out.is_contiguous() ++ && state_max.is_contiguous() && state_sum.is_contiguous(), ++ "dense-state outputs must be contiguous"); ++ const int64_t rows = q.size(0) * q.size(1) * q.size(2); ++ TORCH_CHECK(state_max.numel() == rows && state_sum.numel() == rows, ++ "dense-state statistic shapes are invalid"); ++ TORCH_CHECK(causal, "dense-state requires causal attention"); ++ ++ const at::cuda::OptionalCUDAGuard device_guard(q.device()); ++ const dim3 block(Sm70D256SplitDTraits::kNThreads); ++ const dim3 grid( ++ q.size(1) / Sm70D256SplitDTraits::kBlockM, ++ q.size(0), ++ q.size(2)); ++ auto stream = at::cuda::getCurrentCUDAStream(); ++ auto kernel = sm70_d256_splitd_dense_kernel< ++ cutlass::half_t, false, false, true>; ++ C10_CUDA_CHECK(cudaFuncSetAttribute( ++ kernel, ++ cudaFuncAttributeMaxDynamicSharedMemorySize, ++ Sm70D256SplitDTraits::kSmemBytes)); ++ kernel<<>>( ++ reinterpret_cast(q.data_ptr()), ++ reinterpret_cast(k.data_ptr()), ++ reinterpret_cast(v.data_ptr()), ++ reinterpret_cast(out.data_ptr()), ++ static_cast(q.stride(0)), ++ static_cast(q.stride(1)), ++ static_cast(q.stride(2)), ++ static_cast(k.stride(0)), ++ static_cast(k.stride(1)), ++ static_cast(k.stride(2)), ++ static_cast(v.stride(0)), ++ static_cast(v.stride(1)), ++ static_cast(v.stride(2)), ++ q.size(1), ++ k.size(1), ++ q.size(2), ++ k.size(2), ++ static_cast(softmax_scale * M_LOG2E), ++ nullptr, ++ 0, ++ 0, ++ nullptr, ++ state_max.data_ptr(), ++ state_sum.data_ptr()); ++ C10_CUDA_KERNEL_LAUNCH_CHECK(); ++ return out; ++} ++ ++// Raw fixed-layout entry point for the task-local end-to-end architecture ++// harness. Keeping this in the same cubin guarantees that the measured tail ++// is the exact production candidate kernel rather than a reimplemented proxy. ++extern "C" cudaError_t onecat_sm70_d256_dense_state_raw( ++ const void *q, ++ const void *k, ++ const void *v, ++ float *state_max, ++ float *state_sum, ++ void *out, ++ int query_len, ++ int kv_len, ++ int heads_q, ++ int heads_kv, ++ float softmax_scale, ++ cudaStream_t stream) { ++ if (q == nullptr || k == nullptr || v == nullptr || state_max == nullptr ++ || state_sum == nullptr || out == nullptr || query_len <= 0 ++ || kv_len < query_len || heads_q <= 0 || heads_kv <= 0 ++ || heads_q % heads_kv != 0 ++ || query_len % Sm70D256SplitDTraits::kBlockM != 0 ++ || kv_len % Sm70D256SplitDTraits::kBlockN != 0) { ++ return cudaErrorInvalidValue; ++ } ++ const dim3 block(Sm70D256SplitDTraits::kNThreads); ++ const dim3 grid( ++ query_len / Sm70D256SplitDTraits::kBlockM, ++ 1, ++ heads_q); ++ auto kernel = sm70_d256_splitd_dense_kernel< ++ cutlass::half_t, false, false, true>; ++ cudaError_t result = cudaFuncSetAttribute( ++ kernel, ++ cudaFuncAttributeMaxDynamicSharedMemorySize, ++ Sm70D256SplitDTraits::kSmemBytes); ++ if (result != cudaSuccess) { ++ return result; ++ } ++ kernel<<>>( ++ static_cast(q), ++ static_cast(k), ++ static_cast(v), ++ static_cast(out), ++ query_len * heads_q * Sm70D256SplitDTraits::kHeadDim, ++ heads_q * Sm70D256SplitDTraits::kHeadDim, ++ Sm70D256SplitDTraits::kHeadDim, ++ kv_len * heads_kv * Sm70D256SplitDTraits::kHeadDim, ++ heads_kv * Sm70D256SplitDTraits::kHeadDim, ++ Sm70D256SplitDTraits::kHeadDim, ++ kv_len * heads_kv * Sm70D256SplitDTraits::kHeadDim, ++ heads_kv * Sm70D256SplitDTraits::kHeadDim, ++ Sm70D256SplitDTraits::kHeadDim, ++ query_len, ++ kv_len, ++ heads_q, ++ heads_kv, ++ softmax_scale * float(M_LOG2E), ++ nullptr, ++ 0, ++ 0, ++ nullptr, ++ state_max, ++ state_sum); ++ return cudaPeekAtLastError(); ++} ++ + at::Tensor sm70_d256_splitd_dense_splitkv3_fwd( + const at::Tensor &q, + const at::Tensor &k, +diff --git a/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h b/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h +index 31b2b769..2cb4ba74 100644 +--- a/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h ++++ b/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h +@@ -623,6 +623,7 @@ public: + return cutlass::Status::kErrorInternal; + } + ++#if !defined(PREFIX_QK_SKIP_APPLY) + // + // Launch the SoftmaxApplyKernel + // +@@ -646,6 +647,7 @@ public: + if (result != cudaSuccess) { + return cutlass::Status::kErrorInternal; + } ++#endif + + return cutlass::Status::kSuccess; + } diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 3893a4ab78..f98aef77f5 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -69,6 +69,7 @@ constexpr uint16_t kSm70Tp4PushAllreduceSentinel = 0x7f7f; constexpr int kSm70Tp4PushAllreduceSentinelByte = 0x7f; constexpr size_t kSm70Tp4PushAllreduceBytes = 8 * kSm70GemmaRmsNormHiddenSize * sizeof(half); +constexpr size_t kSm70Tp4PushAllreduce8KiBBytes = 4096 * sizeof(half); constexpr size_t kSm70Tp4PushAllreduceSignalBytes = ((kSm70Tp4PushAllreduceBlocks * sizeof(uint32_t) + 127) / 128) * 128; constexpr size_t kSm70Tp4PushAllreduceBufferBytes = @@ -76,6 +77,13 @@ constexpr size_t kSm70Tp4PushAllreduceBufferBytes = kSm70Tp4PushAllreduceWorldSize * kSm70Tp4PushAllreduceBytes; +inline int sm70_tp4_push_allreduce_blocks(size_t bytes) { + if (bytes == kSm70Tp4PushAllreduceBytes) { + return kSm70Tp4PushAllreduceBlocks; + } + return bytes == kSm70Tp4PushAllreduce8KiBBytes ? 4 : 0; +} + inline int sm70_gemma_rms_norm_threads() { const char* raw = std::getenv("VLLM_SM70_TP2_AR_GEMMA_RMS_THREADS"); if (raw == nullptr) return kSm70GemmaRmsNormThreads; @@ -1519,18 +1527,20 @@ class CustomAllreduce { size /= d; auto bytes = size * sizeof(typename packed_t::P); if constexpr (std::is_same_v) { - // The push protocol amortizes its peer polling across the verifier's - // captured 128-call chain. A lone eager call is materially faster on - // the ordinary registered-buffer pull path. + // The push protocol amortizes peer polling across captured collective + // chains. A lone eager call stays on the ordinary registered-buffer + // pull path. if (sm70_tp4_push_buffers_registered_ && status == cudaStreamCaptureStatusActive && world_size_ == kSm70Tp4PushAllreduceWorldSize && fully_connected_ && - bytes == kSm70Tp4PushAllreduceBytes && custom_allreduce_current_device_is_sm70()) { - sm70_cross_device_reduce_1stage_push - <<>>(sm70_tp4_push_buffers_, input, output, rank_, size); - return; + const int push_blocks = sm70_tp4_push_allreduce_blocks(bytes); + if (push_blocks > 0) { + sm70_cross_device_reduce_1stage_push + <<>>( + sm70_tp4_push_buffers_, input, output, rank_, size); + return; + } } if (sm70_tp8_hierarchical_custom_ar_enabled(world_size_, fully_connected_) && diff --git a/docs/design/sm70_dflash2_acceptance_rd.md b/docs/design/sm70_dflash2_acceptance_rd.md new file mode 100644 index 0000000000..c190844ec0 --- /dev/null +++ b/docs/design/sm70_dflash2_acceptance_rd.md @@ -0,0 +1,77 @@ +# SM70 DFlash2 acceptance research without retraining + +> **Status:** the default-off alignment capture and offline analyzer are +> development tooling. They do not change the production proposal policy and +> do not claim an acceptance or performance improvement by themselves. + +## Scope and frozen contracts + +- Date: 2026-08-25. +- Integration base: `onecat/main@34403018d917054dd7765d5e820ad29c8d342348`. +- Branch: `codex/v100-v100-dflash2-acceptance-rd-20260825-113155`. +- No draft or target retraining and no checkpoint edits. +- Target validation remains the complete MRV2 block-eight workload: seven draft + tokens plus one target bonus row. DDTree index reuse and wider verifier trees + are outside this campaign. +- Target sampling remains exact standard rejection at temperature 1.0, top-p + 0.95, and top-k 20. A candidate proposal distribution is admissible only if + the selector draw and the rejection sampler consume the same exact `q`. + +The score gate remains the retained 128-row mixed32 dataset with SHA256 +`85ce10d84735ec981c001663d5748b1939a81892ef3ccd4549aec66177795607`. +Its private filesystem location is intentionally not committed. +It uses sequential B1, a fixed seed, xhigh reasoning, at most 2,048 generated +tokens, and natural EOS. Corrected MBPP-32 and WikiText-2 PPL retain their +existing independent gates. The historical native-grouped result is 83/96 on +GSM8K/MATH-500/HumanEval, 25/32 on corrected MBPP, request-mean completion +length 4.238430, and pooled acceptance length 3.749104. + +The production performance gate remains the PR #288 practical contract: TP4 +V100, 256K maximum context, `max_num_batched_tokens=4096`, prefix caching, +Mamba align mode, FP8 E5M2 target KV, FP16 draft KV, Flash-V100, probabilistic +DFlash2, tool/reasoning parsers, and FULL CUDA Graph. Its short-context complete +round is 18.465--18.603 ms. Acceptance work may not increase verifier rows or +regress that round by more than measurement noise. + +## First target and promotion gate + +The first research target is at least +0.5 request-mean completion tokens per +verification round on the unchanged mixed32 contract. Promotion additionally +requires: + +1. no aggregate task-score loss against the paired target-only and current + DFlash2 controls; +2. no corrected-MBPP or WikiText PPL regression; +3. exact rejection-distribution tests and normalized proposal rows; +4. unchanged eight-row target validation and at most 0.05 ms incremental + selector latency on SM70; +5. no material acceptance or latency regression at 1K, 32K, 128K, or 256K. + +## Selector headroom audit + +The first implementation is diagnostic-only. With both compact sparse +rejection and `VLLM_SPEC_DUMP_ALIGNMENT=1`, the speculator keeps the checkpoint +top16 IDs, unary logits, full `7 x 16 x 16` selector lattice, and exact realized +proposal rows. The target sampler records its top20 rows and the actual strict +rejection outcome. Disabled diagnostics allocate no shadow lattice and add no +copy or sampler operation. + +`benchmarks/analyze_dflash2_selector_alignment.py` reports, by draft depth: + +- current overlap `sum(min(p, q))`; +- target probability mass covered by the fixed top16 candidate support; +- the gap between support mass and current overlap; +- one-step counterfactual sweeps over proposal temperature, proposal nucleus, + unary/edge calibration, and backward log-sum-exp future messages. + +Even sampler steps tune the candidates and odd steps are a held-out report. +Counterfactual values after a changed path are explicitly only recorded-prefix +overlap proxies, not end-to-end acceptance claims. Survivors must be rerun on +the frozen dataset. + +## Initial verification + +- Focused Ruff lint and format: passed. +- Diagnostic analyzer synthetic checks: passed. +- CPU DFlash2 suite with CUDA hidden: 77 passed, 12 expected CUDA skips. +- No GPU performance or end-to-end acceptance result has been claimed yet. diff --git a/docs/design/sm70_dflash2_ngram_hybrid.md b/docs/design/sm70_dflash2_ngram_hybrid.md new file mode 100644 index 0000000000..bb45e80d93 --- /dev/null +++ b/docs/design/sm70_dflash2_ngram_hybrid.md @@ -0,0 +1,130 @@ +# SM70 DFlash2 + ngram Hybrid + +## Purpose + +Add an optional prompt-ngram assistant to the MRV2 DFlash2 path. An ngram hit +must skip the DFlash2 query/selector while preserving DFlash2 context-KV state, +and a miss must fall back to the existing DFlash2 implementation unchanged. +The target verifier and the probabilistic rejection-sampling contract remain +authoritative, so the optimization cannot change the target distribution. + +This work is based on `onecat/main` at +`d62ef5cb20b48de93a91562e777ac48985f44b76` and is isolated on +`agent/v100-dflash2-ngram-hybrid-20260825-050018`. + +## Upstream audit + +- vLLM prompt lookup: PRs #12193, #22437, #24986, and #29184 provide the CPU + KMP and GPU-vectorized implementations. MRV2 still has no combined ngram + + model-drafter route. +- SGLang: PRs #17260, #21243, and #22737 show that overlap scheduling requires + complete request-token state and explicit accepted-token indexing. A stale + host output list is not a valid lookup source. +- llama.cpp: its comma-separated speculative configuration gives draftless + ngram proposers priority, falls back to DFlash, and still calls `process()` on + every proposer so model-drafter state remains synchronized. This is the + closest reference architecture for the first implementation here. +- Arctic Suffix Decoding and SAM-Decoding are useful follow-ups, but their + confidence policy, external dependency, and longer trees are intentionally + outside this first block-8 implementation. + +## Initial contract + +- The feature is opt-in and only valid for `method=dflash` with a DFlash2 + checkpoint. Standalone `ngram`, Eagle, MTP, DFlash1, and `dflash_ddtree` + routing is unchanged. +- Lookup reads the authoritative MRV2 request-token state and supports normal + synchronous and overlap scheduling without rebuilding history in Python. +- Structured-output requests bypass the ngram assistant until grammar-aware + proposal masking is proved correct. Tool/reasoning parsers without a grammar + are unaffected. +- A hit returns at most the configured DFlash2 draft width (seven tokens in the + official block-8 setup). The normal DFlash2 proposer handles misses. +- Context K/V materialization always runs. Only the DFlash2 query, candidate + projection, selector, and selector walk may be skipped. +- Greedy and probabilistic modes are supported. In probabilistic mode the + ngram proposal is represented as a one-hot draft distribution inside the + existing sparse target-rejection interface, preserving target sampling. +- Prefix-cache hits may rebuild missing DFlash draft K/V as before; this change + must not corrupt either target or draft cache state. + +## Test plan + +1. Unit tests: configuration/routing isolation; KMP hit, miss, truncation, and + overlap cases; per-request mixed-source state; one-hot sparse draft support; + grammar bypass; prefix/state transitions. +2. CPU microbenchmarks: lookup at 1K, 32K, 128K, and 256K contexts. +3. V100 correctness: eager before CUDA graph; batch 1/2/4; hit/miss/mixed; + compare proposed tokens, accepted trajectory, target output, and draft K/V. +4. V100 performance: report lookup, context-KV, query/selector, target verify, + full round, emitted tokens per round, and pure-decode tokens/s for baseline + DFlash2 and DFlash2+ngram at short and long contexts. +5. Quality: compare target-only, DFlash2, and hybrid on coding/general/tool + datasets under the same sampling configuration. Report scores, completion + counts, wall time, hit rate, conditional acceptance length, and failures. + +## Promotion gates + +- No statistically meaningful quality regression against target-only or the + existing DFlash2 route. +- DFlash2-miss acceptance trajectory remains unchanged in deterministic tests. +- Ngram hits actually skip query/selector work in a trace. +- Hybrid pure-decode throughput is non-regressing at every measured context; + otherwise the feature remains opt-in while the losing shape is investigated. + +## Status + +- 2026-08-25: Draft PR #287 implements opt-in MRV2 host lookup over the + authoritative UVA request-token state, full-hit query/selector bypass, + all-hit batch application, and one-hot dense/sparse rejection caches. + Structured-output batches bypass the assistant, and intermediate chunked + prefill materializes draft context K/V before returning without lookup. +- CPU split-history KMP microbenchmark (median/P95): 1K `2.66/2.74 us`, 32K + `50.05/53.22 us`, 128K `190.72/201.03 us`, and 256K `369.57/393.17 us`. + Focused validation reports `93 passed, 9 skipped` for the CPU DFlash2/ngram + set, `10 passed` for MRV2 routing, and `24 passed` for the V100 ngram/AOT + fullgraph set. +- PR #288 first restored the missing production optimization closure to main. + On that restored source, no-assist practical coding measured `18.660 ms` per + complete round, `135.70 tok/s`, and acceptance length `2.532`. Repeated runs + of the historical MBPP item 28 stabilized at `18.484--18.616 ms`, + `224.05--225.66 tok/s`, and acceptance length `4.184`. Older 27--31 ms + branch measurements predate this closure and are invalid baselines. +- With ngram `[5,5]` enabled under the identical TP4 practical contract, the + coding request measured `18.980 ms`, `131.14 tok/s`, and acceptance length + `2.494`; only about 2.2% of eligible rounds were full hits. The matched MBPP + item measured `18.700 ms`, `217.18 tok/s`, and acceptance length `4.071`. + Thus low-hit short requests currently pay roughly `0.1--0.3 ms` per round + and do not pass the default-enable performance gate. +- The 32-case MBPP natural-stop run (16K output cap, fixed per-item seeds) gave + the pure-DFlash control `65,934` output tokens in `389.542 s`, `19,767` + verification rounds, `19.707 ms` wall time per round, and acceptance length + `3.336`. Hybrid produced `53,710` tokens in `327.613 s`, `16,700` rounds, + `19.618 ms` per round, and acceptance length `3.216`. Cumulative full-hit + rate reached about 10.1% and lookup averaged `0.016--0.018 ms`; skipped + queries slightly reduced round cost, but the lower sampled acceptance path + left raw aggregate output throughput at `163.94` versus `169.26 tok/s`. +- EvalPlus on the mapped 31 MBPP cases reports pure DFlash Base `30/31` and + Plus `28/31`; hybrid reports Base `31/31` and Plus `28/31`. This is no score + regression, but it does not override the failed throughput gate. The + probabilistic one-hot proposal preserves the target distribution; it is not + expected to reproduce the same sample-by-sample random trajectory. +- A fresh hybrid 32K chunked-prefill request measured `10.688 s` cold and + `1.416 s` on an identical-prefix hit, versus the restored no-assist evidence + of `10.58--10.65 s` and `1.405 s`. The context-only skip log was present and + no cache corruption occurred. +- Promotion decision: keep `ngram_assist` opt-in. The merge audit removes the + mixed-hit hazard: an n-gram draft is now applied only when every active + request has a full-width hit and the complete DFlash2 query/selector can be + skipped. Mixed-hit batches retain the unchanged DFlash2 proposals because + overriding a row after paying the full query cost cannot improve latency and + can alter sampled acceptance. A future default-on policy still needs matched + evidence that full-query skips repay lookup/probe overhead; blindly reducing + the ngram length is not justified by the current data. +- After the paired run was stopped, fresh graph captures on GPUs 4--7 began + failing inside the draft paged-attention capture with + `cudaErrorStreamCaptureInvalidated`. A clean pre-ngram main worktree failed + at the same point, so this is not an ngram source regression. The NVLink + topology requires resetting all eight GPUs, which was deliberately not done + while the user-facing API remained live on GPUs 0--3. Do not count failed + startups as performance samples. diff --git a/docs/design/sm70_qwen38_fp8_prefill_decay.md b/docs/design/sm70_qwen38_fp8_prefill_decay.md index 62ba587a9e..2ef310aaf6 100644 --- a/docs/design/sm70_qwen38_fp8_prefill_decay.md +++ b/docs/design/sm70_qwen38_fp8_prefill_decay.md @@ -287,8 +287,135 @@ greedy identity fails. Keep Q8000 split-KV3 explicit and default-off until a fixed-text logprob/perplexity or dataset-level quality gate establishes that the classified Type-B reduction-order drift does not reduce model quality. +## 2026-08-25 Q8000/KV40K..128K GQA-packed Attention Architecture + +This experiment freezes the final long-prefill call at causal FP16 +`Q=8000`, `KV=128000`, `Hq=6`, `Hkv=1`, and `D=256` on one +V100-SXM2-32GB. Its metric is 6.094872576 useful causal TFLOPs divided by the +complete Attention elapsed time. It is neither whole-model TOPS nor prompt +tokens/s; the acceptance gate is at most 101.581 ms, or at least 60.0 useful +causal TFLOP/s. + +Tile and barrier tuning of the fused BM64/BN32 kernel was already exhausted: +the long-shape NCU record showed one CTA/SM, 12.5% occupancy, 62.14% of cycles +without an eligible warp, and only 9.06% DRAM throughput. The replacement is a +different scheduling and dataflow architecture: + +- split a fully visible 32K..120K prefix from the exact causal 8K tail; +- pack all six GQA query heads into GEMM M=48,000; +- run wide SM70 Tensor-Core QK GEMMs that emit FP16 scores plus row max/sum; +- transform scores into normalized probabilities during the PV + global-to-shared load, avoiding a materialized probability matrix; +- fold one reusable FP16 PV partial into a FP32 online prefix + numerator/max/sum state; +- export max/sum from the accepted exact 8K tail and merge the two online + softmax states. + +The retained score block is BN8192. Final frozen-extension Torch A/B/A for the +score-cache layout measures `140.00213 -> 100.76979 ms`, or +`43.53414 -> 60.48313` useful causal TFLOP/s: 28.0227% lower latency and a +1.38933x speedup. All 12,288,000 outputs are finite; max/mean absolute error +is `2.2888e-4 / 2.5430e-5` and relative L2 is `6.8629e-3` versus the exact +FP32-accumulator route. BN7168 reaches only 58.42057 TFLOP/s and BN8000 only +59.92825 TFLOP/s, so neither is promoted or rounded into a pass. + +The same bounded workspace was then generalized across the twelve observed +8K-chunk shapes. Each row is a strict A/B bracket against the active split-KV3 +control: + +| KV tokens | split-KV3 ms | architecture ms | latency change | useful TFLOP/s | +|---:|---:|---:|---:|---:| +| 40,000 | 38.41638 | 29.81530 | -22.3891% | 59.34862 | +| 48,000 | 47.25811 | 36.38477 | -23.0084% | 59.44005 | +| 56,000 | 56.53999 | 42.68715 | -24.5010% | 59.87584 | +| 64,000 | 65.17897 | 49.14552 | -24.5991% | 60.00842 | +| 72,000 | 74.00670 | 55.44960 | -25.0749% | 60.27745 | +| 80,000 | 82.74432 | 61.77792 | -25.3388% | 60.46783 | +| 88,000 | 91.98507 | 68.37111 | -25.6715% | 60.38797 | +| 96,000 | 100.88397 | 74.72862 | -25.9262% | 60.51241 | +| 104,000 | 109.80642 | 81.09448 | -26.1478% | 60.61108 | +| 112,000 | 118.42901 | 87.71959 | -25.9307% | 60.51602 | +| 120,000 | 127.35232 | 93.99757 | -26.1909% | 60.65749 | +| 128,000 | 135.93617 | 100.34193 | -26.1845% | 60.74103 | + +All 147,456,000 candidate output elements are finite. Across the family, the +worst max absolute difference is `5.0354e-4` and worst relative L2 is +`6.8983e-3` versus the exact FP32-accumulator route. The three shortest points +remain slightly below 60 because the 43-TFLOP/s exact causal tail is a larger +fixed fraction, but they are still 1.288x..1.325x faster than split-KV3. + +The architecture trace is compute-dense rather than launch-starved: 49 timed +kernels occupy 95.630493 ms of a 95.809850-ms GPU span, or 99.81%. Prefix QK +accounts for 54.22% and prefix PV 39.74%; the exact tail is 4.89%. The next +compute ceiling is therefore QK/PV instruction efficiency, not additional +host launch fusion. + +Real-model integration exposed a separate memory contract. Retaining all +partials or the entire workspace passes the operator gate but leaves too +little memory for the following 80-MiB TP all-reduce. The final layout retains +only the 750-MiB FP16 score matrix per device and packs about 100 MiB of other +state into one transient byte slab. Before first use it requires enough +driver-visible free memory for both allocations plus 128 MiB of downstream +headroom; otherwise it raises a typed OOM before touching the caching allocator +and falls back to the exact route. + +The route is default-on for its exact engine contract. The legacy +`VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL=0` setting is the +rollback, and the route rejects CUDA graph capture. It admits only `Q=8000`, +`KV=40000..128000` in 8000-token steps, `Hq=6/Hkv=1/D=256`, FP16, causal +attention, and scale 1/16. The earlier +endpoint-only `.875` TP4 +Qwen3.8-27B-FP8 control/candidate/control prefill times are +`46.45658 / 45.47797 / 46.11195 s`. Relative to the `46.28426-s` bracketed +control, the candidate lowers latency by 1.7421% and raises prompt throughput +from 2831.89 to 2882.10 token/s (1.7729%). Every candidate rank logs 16 +architecture hits, with no architecture OOM/fallback, and all three runs emit +the same 32 output token IDs and SHA256 +`df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. + +The endpoint-only `.88` run also succeeds at 45.53427-s prefill with the same +hash and 16 hits/rank. Its memory profiler left enough headroom, so the route +passed rather than exercising the fallback branch. + +The widened TP4 Qwen3.8-27B-FP8 gate brackets the candidate with `46.11195-s` +and `46.09222-s` controls. Relative to their `46.10208-s` mean, the +`41.51191-s` candidate lowers prefill latency by 9.9565% and raises prompt +throughput from 2843.08 to 3157.46 token/s (11.0575%). Every rank logs exactly +192 family-route hits with no architecture OOM/fallback. Both controls and the +candidate emit the same 32 token IDs and SHA256 +`df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. + +The 2026-08-26 merge audit rebuilt the four-patch vendored FA2 stack from its +locked clean baseline with CUDA 12.8 and `sm_70`. The production QK/PV kernels +again compiled at 254/119 registers with zero spill. On the rebuilt extension, +the 40K endpoint measured `39.29395 -> 30.27579 ms` (`1.29787x`) with all +outputs finite, max/mean absolute difference `5.4932e-4 / 4.4392e-5`, and +relative L2 `6.6098e-3`. The 128K endpoint measured +`136.07628 -> 100.20284 ms` (`1.35801x`) with all outputs finite, +max/mean absolute difference `2.2888e-4 / 2.5484e-5`, and relative L2 +`6.9166e-3`. Both points pass the FP16 merge envelope of max absolute error +at most `1e-3` and relative L2 at most `1e-2`; greedy or bitwise identity is +not required. The complete Flash-V100 policy suite passes 112 tests on a real +V100, including a direct architecture-OOM-to-dense-fallback test. + ## Artifacts +- D256 GQA architecture artifacts are retained outside Git; their private + filesystem location is intentionally not committed. +- Final operator A/B/A: + `results/torch-architecture-scorecache-final-v1-aba.json` under that task + root. +- Shape-family operator A/B: + `results/torch-architecture-shapefamily-v1-aba.json` under that task root. +- Shape-family TP4 model gate: + `results/tp4-128k-architecture-shapefamily-{candidate,control-after}-0875.json`, + bracketed with the endpoint gate's + `results/tp4-128k-architecture-scorecache-final-control-b-0875.json`. +- Final TP4 `.875` model gate: + `results/tp4-128k-architecture-scorecache-final-{control-a,candidate,control-b}-0875.json`. +- Final high-memory `.88` route-success run: + `results/tp4-128k-architecture-scorecache-final-preflight-fallback-088.json`; + despite the retained filename, the preflight passed and no fallback ran. - K-stage task root: `/data/minimax-h3/task-cache/qwen38-fp8-128k-flashattention-20260824`. - Final clean FA2 binary: diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index a3ea441626..652ef16c83 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -1896,6 +1896,100 @@ Source behavior: - Decode path reads vLLM paged KV cache directly. - Prefill has direct and fallback paths, including paged KV gather fallback. +D256 8K-by-40K..128K GQA architecture checkpoint, 2026-08-25: + +- Frozen single-rank shape is causal FP16 Q/K/V, `Q=8000`, `KV=128000`, + `Hq=6`, `Hkv=1`, `D=256` on one V100-SXM2-32GB. The metric is useful + causal Attention FLOPs divided by complete Attention elapsed time; it is not + model TOPS or prompt throughput. +- The architecture separates a fully visible 32K..120K prefix from the causal + 8K tail. Six GQA heads are packed into GEMM M=48,000. QK writes FP16 logits and + row statistics; PV normalizes logits during its global-to-shared operand + transform. One reusable FP16 block partial is folded into a FP32 online + prefix numerator/state before the next block. The exact causal-tail kernel + exports max/sum state and a final kernel merges prefix and tail. +- The final BN8192 score-cache layout retains only the 750-MiB FP16 score + matrix per device. All other approximately 100 MiB of state is packed into + one transient aligned byte slab. A per-device event and mutex serialize + score reuse across streams. Before its first allocation, the route requires + driver-visible free memory for both workspaces plus 128 MiB of downstream + headroom; failure raises a typed OOM without touching the caching allocator + or launching a kernel, so the Python path can fall back safely. +- The final frozen-extension Torch A/B/A measured the existing exact operator + at `140.00213 ms / 43.53414 TFLOP/s` and the complete architecture at + `100.76979 ms / 60.48313 TFLOP/s`: `28.0227%` lower latency and `1.38933x` + speedup. All 12,288,000 outputs are finite; versus the exact operator, + max absolute difference is `2.2888e-4`, mean absolute difference + `2.5430e-5`, and relative L2 `6.8629e-3`. Peak allocated memory is + `1.11960 GiB`. +- The generalized operator admits all twelve observed `Q=8000` chunk shapes, + `KV=40000..128000` in 8000-token steps. Strict per-length A/B measured + `59.34862..60.74103` useful causal TFLOP/s. Against the active split-KV3 + control, every point is faster: latency is `22.3891%..26.1909%` lower and + speedup is `1.28848x..1.35485x`. All 147,456,000 candidate elements are + finite; the worst max absolute difference is `5.0354e-4` and worst relative + L2 is `6.8983e-3` versus the exact FP32-accumulator route. +- Smaller score blocks do not pass the target: BN7168 is + `104.32751 ms / 58.42057 TFLOP/s`; BN8000 is + `101.70283 ms / 59.92825 TFLOP/s` and must not be rounded to 60. BN8192 is + therefore retained while model KV reservation supplies downstream headroom. +- QK and PV compile at 254 and 119 registers/thread with no spill, stack, or + local storage. The complete operator is loadable through `torch.ops`; the + initial namespace registration mismatch was caught by a static load test and + fixed before timing. Synchronous device-symbol updates were changed to + current-stream asynchronous copies so the integrated measurement includes no + artificial whole-device fence. +- Integration is shape-family bounded and default-on. The legacy + `VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL=0` setting is the + rollback. It rejects CUDA-graph capture and falls back to the exact dense + kernel on workspace OOM. + The final external FA2 patch applies cleanly after the existing D256 + pipeline, split-KV3, and K-ping-pong patches, and a fresh CUDA 12.8 / Torch + cu128 SM70 build passes its operator-schema load test. The full Flash-V100 + policy file passes 112/112 tests on a real V100, including direct typed + architecture-OOM fallback to the exact dense route; Ruff and + `git diff --check` pass. +- Early real-model memory tests at `gpu_memory_utilization=0.88` rejected the + fully persistent and pre-preflight score-cache layouts after their first + successful route call: the following 80-MiB PYNCCL all-reduce had only + 62-74 MiB free. This is a downstream-headroom failure, not an Attention + compute failure. +- The final matched `.875` TP4 Qwen3.8-27B-FP8 + control/candidate/control prefill times are `46.45658 / 45.47797 / + 46.11195 s`. Against the `46.28426-s` bracketed control, the candidate is + `1.7421%` lower latency and `1.01773x` faster (`2831.89 -> 2882.10` + prompt token/s). Every candidate rank logs 16 architecture hits and no + architecture OOM/fallback. All three runs produce the same 32 output token + IDs and SHA256 + `df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. +- A final `.88` run also completes in `45.53427 s`, with the same token hash + and 16 route hits/rank. Its profiler left enough headroom for the route, so + this validates high-memory route success rather than directly exercising + the preflight fallback branch. +- The widened TP4 Qwen3.8-27B-FP8 gate uses a `46.11195-s` control before the + candidate and a `46.09222-s` control after it. Against their `46.10208-s` + mean, the `41.51191-s` candidate lowers prefill latency by `9.9565%` and + raises prompt throughput from `2843.08` to `3157.46` token/s (`11.0575%`). + Every rank logs exactly 192 family-route hits with no architecture OOM or + fallback. Both controls and the candidate emit the same 32 token IDs and + SHA256 + `df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. + Token identity is route-health evidence rather than a promotion + requirement. The stronger numerical gate covers 147,456,000 finite output + elements with worst max absolute error `5.0354e-4` and relative L2 + `6.8983e-3`, inside the accepted FP16 envelope of `1e-3` and `1e-2`. + Artifacts and the experiment ledger are retained outside Git; their private + filesystem location is intentionally not committed. +- The 2026-08-26 merge audit replayed all four patches from the locked clean + vendored baseline and rebuilt `_vllm_fa2_C` with CUDA 12.8 for `sm_70`. + QK/PV again compiled at 254/119 registers with zero spill. Current-source + direct A/B measured `39.29395 -> 30.27579 ms` (`1.29787x`) at KV40K and + `136.07628 -> 100.20284 ms` (`1.35801x`) at KV128K. Both outputs were fully + finite. KV40K had max absolute error `5.4932e-4` and relative L2 + `6.6098e-3`; KV128K had max absolute error `2.2888e-4` and relative L2 + `6.9166e-3`. These pass the FP16 merge envelope without imposing greedy or + bitwise identity and justify keeping the validated shape route default-on. + Latest target state: - No `FLASH_ATTN_V100` enum. @@ -43230,3 +43324,108 @@ Interpretation: async-scheduler unit cases, and a simple CPU-offload store/load round trip. All changed-file pre-commit hooks pass. This is a cache lifetime correctness repair; no throughput claim is attached. + +## 2026-08-25 DFlash2 n-gram hybrid + +- Development is isolated on + `agent/v100-dflash2-ngram-hybrid-20260825-050018` and rebased onto the + restored `onecat/main@d62ef5cb20`; Draft PR #287 contains only the MRV2 + DFlash2 dependency closure. +- The opt-in assistant reads the authoritative UVA token state, uses the same + reverse-KMP policy as standalone vLLM n-gram lookup, skips DFlash2 query and + selector only on a full seven-token hit, and preserves context-KV + materialization. Misses remain on the unchanged DFlash2 route. +- V100 unit coverage is 24/24, including probabilistic one-hot rejection, + mixed rows, overlap state, and the CUDA override kernel. Practical TP4 graph + startup with prefix cache, FP8 E5M2 KV, 256K max context, 4096 batched-token + limit, and tool/reasoning parsers is proven. The earlier 27--31 ms clean-main + measurements predate the merged production closure and are invalid as a + speed baseline; collect `ngram_assist=false/true` numbers from the same + restored build. +- The restored paired result is now complete. No-assist practical coding is + `18.660 ms` per round, `135.70 tok/s`, acceptance `2.532`; ngram `[5,5]` is + `18.980 ms`, `131.14 tok/s`, acceptance `2.494`, with only about 2.2% full + hits. The feature therefore remains opt-in for short/low-hit traffic. +- On the fixed-seed, 16K-cap, natural-stop MBPP32 workload, pure DFlash records + `389.542 s`, `19,767` rounds, `19.707 ms/round`, acceptance `3.336`, and + EvalPlus Base/Plus `30/31` and `28/31`. Hybrid records `327.613 s`, `16,700` + rounds, `19.618 ms/round`, acceptance `3.216`, and Base/Plus `31/31` and + `28/31`. Its cumulative full-hit rate is about 10.1%, but lower sampled + acceptance leaves aggregate output throughput `163.94` versus + `169.26 tok/s`; the quality gate passes and the throughput gate does not. +- Hybrid 32K prefill is `10.688 s` cold and `1.416 s` on an identical-prefix + hit, matching restored no-assist evidence (`10.58--10.65 s`, `1.405 s`) + closely enough to rule out a material prefill regression in this probe. +- Post-run capture failures on GPUs 4--7 reproduce unchanged on the clean + pre-ngram main worktree at the same draft paged-attention capture. They are + runtime/GPU-state evidence, not an ngram regression or speed sample. A full + NVLink-group reset was intentionally deferred because GPUs 0--3 host the + live user API. +- The merge audit removes the mixed-hit quality/performance hazard. N-gram + drafts are now applied only when every active request has a full-width hit + and the complete DFlash2 query/selector is skipped. If any request misses, + the batch keeps the unchanged DFlash2 proposals because overriding a row + after paying the full query cost cannot improve latency and can alter + acceptance. The assistant remains explicit opt-in until matched evidence + shows that full-query skips repay host lookup overhead for a workload. + +## 2026-08-26 PR #283 aggregate-quality acceptance + +- The numerical policy does not require greedy or token-stream identity. + Candidate outputs must remain finite, satisfy dtype-appropriate operator + error bounds, and preserve matched aggregate task quality. Per-example + regressions and improvements remain visible diagnostics rather than an + automatic identity gate. +- The full serialized-FP8 QPN8 candidate improves the matched PP2 x TP4 B1 + no-spec endpoint from 59.248 to 64.359 token/s (`+8.63%`) and reduces mean + TPOT from 16.878 to 15.538 ms. Its paired GSM8K-64 result remains 63/64 with + zero invalid answers: one baseline-correct answer regresses and one + baseline-wrong answer improves, so aggregate quality does not decline. +- Operator checks remain finite with relative L2 at most `6.05e-4`, cosine at + least `0.9999997`, and maximum absolute difference at most `0.00390625` in + the recorded screen. All 365 audited block scales are finite, positive + powers of two and survive the layout scale transform exactly. These are + numerical bounds, not a claim of bitwise equivalence. +- QPN8 is therefore default-on only for its validated engine contract: exact + SM70 block-FP8 operator roles/shapes/layouts, PP2 x TP4, one sequence, no + DBO or explicit ubatching, and no speculative decoding. Missing operators + or workspace allocation fall back to TurboMind. Either + `VLLM_SM70_FP8_QPN8=0` or `VLLM_SM70_FP8_QPN8_PP2_TP4=0` is an explicit + rollback. Admission never reads model name, checkpoint, `model_type`, or + architecture identity. +- Metadata-free PP transfer is also default-on for its exact SM70 B1 schema: + PP2 x TP4, FP16 `[1,4,4096]` contiguous replicated hidden state, CUDA Graph, + no sequence parallelism, no DBO/ubatching/speculation, and one sequence. + Three matched endpoint runs improve QPN8 by about `0.22%` (`0.034 ms/token`) + without changing its output behavior. Non-admitted configurations use the + original metadata plus TP reconstruction path; an admitted-but-invalid + runtime tensor fails fast before communication. The rollback is + `VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER=0`. +- The paired quality tools now pin dataset hashes and the complete evaluation + contract, validate artifact self-consistency, and gate aggregate HumanEval, + LongBench, GSM8K, and needle-retrieval quality. They report directional + sample flips but do not treat greedy identity as task quality. + +## 2026-08-26 PR #299 8-KiB TP4 push all-reduce acceptance + +- The existing two-epoch SM70 TP4 push collective now also admits the exact + 8-KiB FP16 decode payload. Four CTAs cover its 512 packed 16-byte elements; + the existing 80-KiB verifier payload retains its 80-CTA launch. Buffer + sizing already uses the larger payload, so this adds no allocation growth. +- Same-binary TP4 A/B/B/A timing improves from 9.286--9.425 us to + 3.041--3.174 us. A dynamic CUDA Graph gate covers 43 collective nodes, + eight changing input patterns, 64 replays, and all four ranks with zero + element mismatches. Both paths are bitwise equal to an explicit fixed-order + rank-0-through-rank-3 FP32 accumulation oracle. +- The matched PP2 x TP4 endpoint improves from 59.160 to 61.272 token/s + (`+3.57%`) and from 16.903 to 16.321 ms/token (`-0.583 ms/token`). Independent + unchanged controls themselves select different stable greedy streams due to + pre-existing cross-process autotuning, so token hashes are diagnostic only; + the exact operator gate isolates this collective without imposing greedy + identity. +- `VLLM_SM70_TP4_PUSH_ALLREDUCE` is therefore default-on. Runtime admission is + limited to SM70, fully connected TP4, active CUDA Graph capture, FP16, and + the exact 8-KiB or 80-KiB payload. Every other device, topology, dtype, size, + or eager call retains the existing pull collective. Explicit `=0` is the + rollback. No model, checkpoint, `model_type`, or architecture identity is + consulted. diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md index 768e9f78d4..bda540b88c 100644 --- a/docs/features/speculative_decoding/README.md +++ b/docs/features/speculative_decoding/README.md @@ -117,6 +117,33 @@ vllm serve \ }' ``` +#### DFlash2 with n-gram assistance + +Set `ngram_assist` on a DFlash2 configuration to try prompt lookup before the +neural DFlash2 query and selector. When every active request has a full-width +lookup hit, the batch skips that neural proposal work. Mixed-hit and miss +batches keep the unchanged DFlash2 proposals. The target model still verifies +every proposed token, including probabilistic sampling, so this does not +change the target sampling distribution. + +```bash +vllm serve \ + --speculative-config '{ + "method": "dflash", + "model": "your-org/your-dflash2-selector-model", + "num_speculative_tokens": 7, + "draft_sample_method": "probabilistic", + "ngram_assist": true, + "prompt_lookup_min": 5, + "prompt_lookup_max": 5 + }' +``` + +This hybrid route currently requires `method=dflash` and a DFlash2 checkpoint. +It is intentionally not enabled for DFlash1, Eagle, MTP, or +`dflash_ddtree`. Structured-output requests bypass prompt lookup and continue +through DFlash2 until grammar-aware proposal masking is available. + #### Suffix decoding | Key | Type | Default | Meaning | diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0654a59caf..9ffaa6954b 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -539,6 +539,14 @@ See [this page](generative_models.md) for more information on how to use generat These models primarily accept the [`LLM.generate`](./generative_models.md#llmgenerate) API. Chat/Instruct models additionally support the [`LLM.chat`](./generative_models.md#llmchat) API. +!!! note "DeepSeek-OCR image resolution" + `DeepseekOCRForCausalLM` accepts explicit `mm_processor_kwargs` with + `image_mode` set to `tiny`, `small`, `base`, `large`, or `gundam`. + `min_crops` and `max_crops` may also be set with + `1 <= min_crops <= max_crops <= 9`. The default remains the existing + `gundam` processing contract; no mode is selected from a model or + checkpoint identity at runtime. + | Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | | ------------ | ------ | ------ | ----------------- | -------------------- | ------------------------- | | `AriaForConditionalGeneration` | Aria | T + I+ | `rhymes-ai/Aria` | | | diff --git a/tests/benchmarks/test_dsv4_quality_tools.py b/tests/benchmarks/test_dsv4_quality_tools.py new file mode 100644 index 0000000000..09697faec9 --- /dev/null +++ b/tests/benchmarks/test_dsv4_quality_tools.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from copy import deepcopy + +from benchmarks.benchmark_dsv4_gsm8k_api import _last_integer +from benchmarks.compare_dsv4_quality_results import ( + _compare_gsm8k, + _compare_humaneval, + _compare_needle, +) + + +def test_last_integer_normalizes_signed_integral_decimal() -> None: + assert _last_integer("work... #### -1,024") == -1024 + assert _last_integer("answer: 12.0") == 12 + assert _last_integer("answer: 12.5") is None + assert _last_integer("no numeric answer") is None + + +def _gsm8k_artifact(predictions: list[int | None]) -> dict: + expected = [13, 40] + rows = [ + { + "index": index, + "question": f"question-{index}", + "expected": answer, + "predicted": prediction, + "correct": prediction is not None and prediction == answer, + "invalid": prediction is None, + } + for index, (answer, prediction) in enumerate(zip(expected, predictions)) + ] + return { + "contract": { + "questions": 2, + "few_shot": 5, + "temperature": 0.0, + "top_p": 1.0, + "seed": 42, + "max_tokens": 256, + "strictly_sequential": True, + "train_selection": "first_n", + "test_selection": "first_n", + "prompt_format": "gsm8k_question_answer_v1", + "answer_normalization": "last_signed_integral_decimal_v1", + "stop_sequences": ["Question"], + }, + "input_manifest": { + "train": {"sha256": "train-hash"}, + "test": {"sha256": "test-hash"}, + }, + "correct": sum(bool(row["correct"]) for row in rows), + "invalid": sum(bool(row["invalid"]) for row in rows), + "rows": rows, + } + + +def test_gsm8k_gate_allows_balanced_answer_flips_without_greedy_identity() -> None: + reference = _gsm8k_artifact([13, 42]) + candidate = _gsm8k_artifact([12, 40]) + + result = _compare_gsm8k(reference, candidate) + + assert result["passed"] + assert result["correct_delta"] == 0 + assert result["regressions"] == [0] + assert result["improvements"] == [1] + assert result["exact_prediction_matches"] == 0 + + +def test_gsm8k_gate_rejects_aggregate_drop_contract_drift_and_bad_summary() -> None: + reference = _gsm8k_artifact([13, 40]) + candidate = _gsm8k_artifact([12, 40]) + assert not _compare_gsm8k(reference, candidate)["passed"] + + candidate = deepcopy(reference) + candidate["contract"]["seed"] = 7 + assert not _compare_gsm8k(reference, candidate)["passed"] + + candidate = deepcopy(reference) + candidate["correct"] = 0 + assert not _compare_gsm8k(reference, candidate)["passed"] + + +def test_humaneval_gate_uses_aggregate_quality_and_reports_flips() -> None: + reference = { + "human_eval": { + "passed": 1, + "records": [ + {"task_id": "a", "passed": True, "response": "a0"}, + {"task_id": "b", "passed": False, "response": "b0"}, + ], + } + } + candidate = { + "human_eval": { + "passed": 1, + "records": [ + {"task_id": "a", "passed": False, "response": "a1"}, + {"task_id": "b", "passed": True, "response": "b1"}, + ], + } + } + + result = _compare_humaneval(reference, candidate) + + assert result["passed"] + assert result["regressions"] == ["a"] + assert result["improvements"] == ["b"] + + +def test_needle_gate_uses_aggregate_hits_and_reports_flips() -> None: + reference = [ + { + "sample_id": "a", + "target_tokens": 100, + "depth": 0.5, + "code": "x", + "hit": True, + "hit_anywhere": True, + }, + { + "sample_id": "b", + "target_tokens": 100, + "depth": 0.5, + "code": "y", + "hit": False, + "hit_anywhere": False, + }, + ] + candidate = [ + {**reference[0], "hit": False, "hit_anywhere": False}, + {**reference[1], "hit": True, "hit_anywhere": True}, + ] + + result = _compare_needle(reference, candidate) + + assert result["passed"] + assert result["final_hit_regressions"] == ["a"] + assert result["anywhere_hit_regressions"] == ["a"] diff --git a/tests/distributed/test_static_tensor_transfer.py b/tests/distributed/test_static_tensor_transfer.py new file mode 100644 index 0000000000..4a1890f5dc --- /dev/null +++ b/tests/distributed/test_static_tensor_transfer.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import pytest +import torch + +from vllm.distributed.parallel_state import GroupCoordinator + + +class _DummyWork: + def wait(self) -> None: + pass + + +def _group(rank_in_group: int, world_size: int = 2) -> GroupCoordinator: + group = GroupCoordinator.__new__(GroupCoordinator) + group.world_size = world_size + group.rank_in_group = rank_in_group + group.ranks = list(range(world_size)) + group.use_cpu_custom_send_recv = False + group.device_group = None + group.cpu_group = None + return group + + +def test_static_tensor_dict_transfer_uses_caller_owned_buffers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[tuple[str, int, torch.Tensor]] = [] + + def fake_isend(tensor: torch.Tensor, *, dst: int, group: Any) -> _DummyWork: + del group + calls.append(("send", dst, tensor)) + return _DummyWork() + + def fake_irecv(tensor: torch.Tensor, *, src: int, group: Any) -> _DummyWork: + del group + tensor.fill_(7) + calls.append(("recv", src, tensor)) + return _DummyWork() + + monkeypatch.setattr(torch.distributed, "isend", fake_isend) + monkeypatch.setattr(torch.distributed, "irecv", fake_irecv) + + sent = torch.arange(8, dtype=torch.float32) + send_handles = _group(0).isend_tensor_dict_static({"hidden_states": sent}) + + received = torch.empty_like(sent) + recv_handles = _group(1).irecv_tensor_dict_static({"hidden_states": received}) + + assert len(send_handles) == 1 + assert len(recv_handles) == 1 + assert calls == [("send", 1, sent), ("recv", 0, received)] + torch.testing.assert_close(received, torch.full_like(received, 7)) + + +def test_static_tensor_dict_transfer_rejects_non_tensors() -> None: + with pytest.raises(TypeError, match="only accepts tensors"): + _group(0).isend_tensor_dict_static( # type: ignore[arg-type] + {"hidden_states": object()} + ) diff --git a/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py b/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py new file mode 100644 index 0000000000..d65cf59fca --- /dev/null +++ b/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm import envs +from vllm.model_executor.layers.quantization import fp8 + + +def _layer( + suffix: str, + tp_size: int, + k_dim: int, + n_dim: int, + *, + output_partition_sizes: list[int] | None = None, +) -> SimpleNamespace: + return SimpleNamespace( + prefix=f"arbitrary.engine.graph.{suffix}", + tp_size=tp_size, + input_size_per_partition=k_dim, + output_size_per_partition=n_dim, + output_partition_sizes=output_partition_sizes, + weight_block_size=[128, 128], + weight=torch.empty((n_dim, k_dim), device="meta"), + ) + + +def test_pp2_tp4_qpn8_is_default_on_with_explicit_rollback(monkeypatch) -> None: + monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) + monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4", raising=False) + envs.disable_envs_cache() + try: + assert fp8._sm70_fp8_qpn8_pp2_tp4_enabled() + + monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "0") + envs.disable_envs_cache() + assert not fp8._sm70_fp8_qpn8_pp2_tp4_enabled() + + monkeypatch.delenv("VLLM_SM70_FP8_QPN8") + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "0") + envs.disable_envs_cache() + assert not fp8._sm70_fp8_qpn8_pp2_tp4_enabled() + finally: + envs.disable_envs_cache() + + +@pytest.mark.parametrize( + ("layer", "gated_silu", "expected"), + [ + (_layer("fused_wqa_wkv", 1, 4096, 1536), False, (32, 2, False)), + (_layer("wq_b", 4, 1024, 8192), False, (8, 2, False)), + (_layer("wo_b", 4, 2048, 4096), False, (16, 2, False)), + ( + _layer( + "gate_up_proj", + 4, + 4096, + 1024, + output_partition_sizes=[512, 512], + ), + False, + (32, 2, False), + ), + ( + _layer( + "gate_up_proj", + 4, + 4096, + 1024, + output_partition_sizes=[512, 512], + ), + True, + (16, 2, False), + ), + (_layer("down_proj", 4, 512, 4096), False, (16, 2, False)), + ], +) +def test_pp2_tp4_qpn8_exact_operator_contracts( + layer: SimpleNamespace, + gated_silu: bool, + expected: tuple[int, int, bool], +) -> None: + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(layer, gated_silu=gated_silu) == expected + + +def test_pp2_tp4_qpn8_rejects_wrong_tensor_and_concurrency_roles() -> None: + wrong_tp = _layer("wq_b", 8, 1024, 8192) + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(wrong_tp, gated_silu=False) is None + + concurrent_indexer = _layer("wq_b", 1, 1024, 8192) + assert ( + fp8._sm70_fp8_qpn8_pp2_tp4_config(concurrent_indexer, gated_silu=False) is None + ) + + wrong_gate = _layer( + "gate_up_proj", + 4, + 4096, + 1024, + output_partition_sizes=[256, 768], + ) + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(wrong_gate, gated_silu=True) is None + + wrong_layout = _layer("wo_b", 4, 2048, 4096) + wrong_layout.weight_block_size = [64, 128] + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(wrong_layout, gated_silu=False) is None + + +def test_pp2_tp4_qpn8_runtime_and_workspace_contract() -> None: + config = SimpleNamespace( + parallel_config=SimpleNamespace( + pipeline_parallel_size=2, + tensor_parallel_size=4, + enable_dbo=False, + ubatch_size=0, + ), + scheduler_config=SimpleNamespace(max_num_seqs=1), + speculative_config=None, + ) + with patch.object(fp8, "get_current_vllm_config", return_value=config): + assert fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + config.speculative_config = object() + assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + config.speculative_config = None + config.parallel_config.pipeline_parallel_size = 1 + assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + config.parallel_config.pipeline_parallel_size = 2 + config.scheduler_config.max_num_seqs = 2 + assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + config.scheduler_config.max_num_seqs = 1 + config.parallel_config.enable_dbo = True + assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + config.parallel_config.enable_dbo = False + config.parallel_config.ubatch_size = 2 + assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + + assert fp8._SM70_FP8_QPN8_PP2_TP4_WORKSPACE_ELEMENTS * 2 == 16 * 1024 * 1024 + + +def test_pp2_tp4_qpn8_grouped_dispatches_caller_groups() -> None: + layer = _layer("wo_a", 4, 4096, 2048) + layer.is_bmm = True + layer.bmm_batch_size = 2 + assert fp8._sm70_fp8_qpn8_pp2_tp4_bmm_config(layer) == (32, 2, False) + + layer.sm70_fp8_turbomind = True + layer.sm70_fp8_qpn8 = True + layer.sm70_fp8_qpn8_bmm = True + layer.sm70_fp8_bmm_groups = 2 + layer.sm70_fp8_bmm_output_size = 1024 + layer.sm70_fp8_qpn8_split_k = 32 + layer.sm70_fp8_qpn8_nacc = 2 + layer.sm70_fp8_qpn8_prefetch = False + layer.sm70_fp8_prefill_exact_dense_workspace_ptr = 123 + layer.weight = torch.empty((2, 4096, 1024), device="meta") + layer.weight_scale_inv = torch.empty((2, 256, 32), device="meta") + + calls: list[tuple[tuple[int, ...], int, int, bool, bool]] = [] + + def fake_dispatch( + out: torch.Tensor, + workspace_ptr: int, + input_: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + split_k: int, + nacc: int, + prefetch: bool, + gated_silu: bool, + ) -> None: + del codes, scales + calls.append( + (tuple(input_.shape), workspace_ptr, split_k, prefetch, gated_silu) + ) + out.fill_(len(calls)) + assert nacc == 2 + + x = torch.zeros((1, 2, 4096), dtype=torch.float16) + with patch.object(fp8.sm70_ops, "fp8_qpn8_dispatch_sm70_out", fake_dispatch): + out = fp8.Fp8LinearMethod.apply(None, layer, x) + + assert calls == [ + ((1, 4096), 123, 32, False, False), + ((1, 4096), 123, 32, False, False), + ] + assert out.shape == (1, 2, 1024) + torch.testing.assert_close(out[:, 0], torch.ones_like(out[:, 0])) + torch.testing.assert_close(out[:, 1], torch.full_like(out[:, 1], 2)) + + +def test_pp2_tp4_qpn8_default_prepares_matching_layer(monkeypatch) -> None: + monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) + monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4", raising=False) + envs.disable_envs_cache() + layer = _layer("fused_wqa_wkv", 1, 4096, 1536) + layer.orig_dtype = torch.float16 + layer.is_bmm = False + layer.weight_scale_inv = torch.empty((12, 32), device="meta") + method = fp8.Fp8LinearMethod.__new__(fp8.Fp8LinearMethod) + method.use_marlin = False + method.use_sm70_fp8_turbomind = True + method.weight_block_size = [128, 128] + config = SimpleNamespace( + parallel_config=SimpleNamespace( + pipeline_parallel_size=2, + tensor_parallel_size=4, + enable_dbo=False, + ubatch_size=0, + ), + scheduler_config=SimpleNamespace(max_num_seqs=1), + speculative_config=None, + ) + + try: + with ( + patch.object(fp8, "get_current_vllm_config", return_value=config), + patch.object(fp8, "_missing_sm70_fp8_qpn8_ops", return_value=[]), + patch.object( + fp8, + "_get_sm70_fp8_qpn8_pp2_tp4_workspace", + return_value=torch.empty(1), + ), + patch.object( + fp8, + "process_fp8_weight_block_strategy", + side_effect=lambda weight, scales: (weight, scales), + ), + patch.object( + fp8.sm70_ops, + "fp8_qpn8_prepare_sm70", + return_value=( + torch.empty((4096, 1536), dtype=torch.uint8, device="meta"), + torch.empty((256, 48), dtype=torch.float16, device="meta"), + ), + ), + patch.object(fp8, "replace_parameter"), + ): + method.process_weights_after_loading(layer) + finally: + envs.disable_envs_cache() + + assert layer.sm70_fp8_qpn8 + assert layer.sm70_fp8_qpn8_split_k == 32 + assert layer.sm70_fp8_qpn8_nacc == 2 + assert not layer.sm70_fp8_qpn8_prefetch diff --git a/tests/transformers_utils/test_deepseek_ocr_resolution_modes.py b/tests/transformers_utils/test_deepseek_ocr_resolution_modes.py new file mode 100644 index 0000000000..b7f8888ecd --- /dev/null +++ b/tests/transformers_utils/test_deepseek_ocr_resolution_modes.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-OCR resolution modes: per-request ``image_mode`` selection, +crop parametrization, and — critically — the producer/counter consistency +contract (`tokenize_with_images` versus +`count_image_tokens`), whose violation is the "N multimodal tokens vs M +placeholders" crash class.""" + +import pytest +from transformers.processing_utils import ProcessorMixin + +from vllm.transformers_utils.processors.deepseek_ocr import ( + BASE_SIZE, + CROP_MODE, + IMAGE_SIZE, + RESOLUTION_MODES, + DeepseekOCRProcessor, + count_image_tokens_for, +) + +# --------------------------------------------------------------------------- +# Pure-function tests (no tokenizer). +# --------------------------------------------------------------------------- + + +def test_mode_table_matches_official_config(): + # Verbatim from the model repo's config.py (captured in the issue refs). + assert RESOLUTION_MODES == { + "tiny": {"base_size": 512, "image_size": 512, "crop_mode": False}, + "small": {"base_size": 640, "image_size": 640, "crop_mode": False}, + "base": {"base_size": 1024, "image_size": 1024, "crop_mode": False}, + "large": {"base_size": 1280, "image_size": 1280, "crop_mode": False}, + "gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True}, + } + # The module defaults ARE the gundam mode (byte-compat anchor). + assert (BASE_SIZE, IMAGE_SIZE, CROP_MODE) == (1024, 640, True) + + +def _count(mode: str, width: int, height: int, **kw) -> int: + cfg = RESOLUTION_MODES[mode] + return count_image_tokens_for( + image_width=width, + image_height=height, + base_size=cfg["base_size"], + image_size=cfg["image_size"], + cropping=cfg["crop_mode"], + **kw, + ) + + +def test_no_crop_modes_are_size_independent(): + for mode in ("tiny", "small", "base", "large"): + assert _count(mode, 100, 100) == _count(mode, 4000, 3000) + + +def test_gundam_small_image_equals_no_crop_arithmetic(): + # <= image_size on both dims: never cropped, regardless of crop_mode. + assert _count("gundam", 640, 480) == count_image_tokens_for( + image_width=640, + image_height=480, + base_size=1024, + image_size=640, + cropping=False, + ) + + +def test_gundam_large_image_counts_tiles(): + assert _count("gundam", 1700, 2200) > _count("gundam", 100, 100) + + +def test_max_crops_bounds_the_tile_count(): + big = _count("gundam", 4000, 4000, max_crops=9) + small = _count("gundam", 4000, 4000, max_crops=2) + assert big > small + + +# --------------------------------------------------------------------------- +# Processor tests with a minimal tokenizer (CPU-only, no model files). +# --------------------------------------------------------------------------- + + +class _TokenizerStub: + bos_token_id = 1 + eos_token_id = 2 + pad_token_id = 0 + pad_token = None + padding_side = "right" + vocab = {"": 99} + + def add_special_tokens(self, tokens): + self.pad_token = tokens["pad_token"] + + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return [10] if text else [] + + def decode(self, tokens, **kwargs): + del kwargs + return " ".join(map(str, tokens)) + + +@pytest.fixture +def tokenizer(monkeypatch): + monkeypatch.setattr(ProcessorMixin, "__init__", lambda *args, **kwargs: None) + + return _TokenizerStub() + + +def _images(sizes): + from PIL import Image + + return [Image.new("RGB", s, color="white") for s in sizes] + + +def test_default_construction_is_gundam_byte_compat(tokenizer): + default = DeepseekOCRProcessor(tokenizer=tokenizer) + explicit = DeepseekOCRProcessor(tokenizer=tokenizer, image_mode="gundam") + for attr in ("base_size", "image_size", "crop_mode", "min_crops", "max_crops"): + assert getattr(default, attr) == getattr(explicit, attr) + assert default.crop_mode is True and default.image_size == 640 + + +def test_invalid_mode_and_crops_raise(tokenizer): + with pytest.raises(ValueError, match="image_mode"): + DeepseekOCRProcessor(tokenizer=tokenizer, image_mode="giant") + with pytest.raises(ValueError, match="crop bounds"): + DeepseekOCRProcessor(tokenizer=tokenizer, min_crops=5, max_crops=2) + with pytest.raises(ValueError, match="crop bounds"): + DeepseekOCRProcessor(tokenizer=tokenizer, min_crops=2, max_crops=10) + + +def test_image_mode_authoritative_over_trio(tokenizer): + # The processing info always forwards the constant trio; a named mode + # must win over it. + proc = DeepseekOCRProcessor( + tokenizer=tokenizer, + image_size=640, + base_size=1024, + crop_mode=True, + image_mode="base", + ) + assert (proc.base_size, proc.image_size, proc.crop_mode) == (1024, 1024, False) + + +@pytest.mark.parametrize("mode", [None, "base", "large", "tiny", "gundam"]) +@pytest.mark.parametrize("size", [(100, 100), (1200, 800), (1000, 2200)]) +def test_producer_counter_consistency(tokenizer, mode, size): + kwargs = {} if mode is None else {"image_mode": mode} + proc = DeepseekOCRProcessor(tokenizer=tokenizer, **kwargs) + out = proc(prompt="\nFree OCR. ", images=_images([size])) + produced = int(out["num_image_tokens"][0]) + counted = proc.count_image_tokens(image_width=size[0], image_height=size[1]) + assert produced == counted + + +def test_multi_image_producer_counter_consistency(tokenizer): + proc = DeepseekOCRProcessor(tokenizer=tokenizer) + sizes = [(1200, 800), (1000, 2200)] + out = proc(prompt="ab", images=_images(sizes)) + for i, size in enumerate(sizes): + produced = int(out["num_image_tokens"][i]) + counted = proc.count_image_tokens(image_width=size[0], image_height=size[1]) + assert produced == counted diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index 2a54959e57..fb2e195796 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -237,6 +237,19 @@ def test_sm70_fa2_d256_prefill_env_is_default_on(monkeypatch): assert envs.VLLM_FLASH_V100_FA2_D256_PREFILL is False +def test_sm70_d256_gqa_architecture_env_is_default_on(monkeypatch): + import vllm.envs as envs + + name = "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL" + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + assert envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL is True + + monkeypatch.setenv(name, "0") + envs.disable_envs_cache() + assert envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL is False + + def test_sm70_e4m3_batch_xqa_env_contract(monkeypatch): import vllm.envs as envs @@ -604,6 +617,7 @@ def test_prefix_prefill_prioritizes_gathered_exact_dense_over_paged( impl.head_size = 256 impl.scale = 0.0625 impl.sliding_window = None + impl.prefix_anchored_decode_window = None impl.kv_cache_dtype = "auto" impl.use_flash_v100_decode = False impl.use_decode_paged_prefill = False @@ -724,6 +738,36 @@ def test_sm70_splitd_d256_loader_requires_exact_ops(monkeypatch): assert flash_v100._get_sm70_splitd_d256_ops() == (dense, paged, splitkv3) +def test_sm70_d256_gqa_architecture_loader_is_optional(monkeypatch): + import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 + + fake_interface = types.ModuleType("vllm.vllm_flash_attn.flash_attn_interface") + fake_package = types.ModuleType("vllm.vllm_flash_attn") + fake_package.__dict__["flash_attn_interface"] = fake_interface + monkeypatch.setitem(sys.modules, "vllm.vllm_flash_attn", fake_package) + monkeypatch.setitem( + sys.modules, + "vllm.vllm_flash_attn.flash_attn_interface", + fake_interface, + ) + + architecture = object() + fake_ops = SimpleNamespace( + _vllm_fa2_C=SimpleNamespace( + sm70_d256_gqa_architecture_fwd=architecture, + ) + ) + monkeypatch.setattr(flash_v100, "torch", SimpleNamespace(ops=fake_ops)) + monkeypatch.setattr( + flash_v100, + "_sm70_d256_gqa_architecture_op_checked", + False, + ) + monkeypatch.setattr(flash_v100, "_sm70_d256_gqa_architecture_op", None) + + assert flash_v100._get_sm70_d256_gqa_architecture_op() is architecture + + def test_prefill_dense_splitkv3_workspace_reuses_exact_shape(monkeypatch): import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 @@ -824,6 +868,144 @@ def test_prefill_dense_splitkv3_policy_is_exact_shape_bounded(monkeypatch): ) +def test_prefill_d256_gqa_architecture_policy_is_shape_family_bounded(monkeypatch): + import vllm.envs as envs + import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 + + name = "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL" + query = torch.empty((1, 8000, 6, 256), dtype=torch.float16, device="meta") + key = torch.empty((1, 128000, 1, 256), dtype=torch.float16, device="meta") + value = torch.empty_like(key) + + monkeypatch.setenv(name, "0") + envs.disable_envs_cache() + assert not flash_v100._should_use_prefill_d256_gqa_architecture( + query, + key, + value, + max_seqlen_q=8000, + max_seqlen_k=128000, + softmax_scale=0.0625, + architecture_op=object(), + ) + + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + for kv_len in range(40000, 128001, 8000): + family_key = torch.empty( + (1, kv_len, 1, 256), dtype=torch.float16, device="meta" + ) + assert flash_v100._should_use_prefill_d256_gqa_architecture( + query, + family_key, + torch.empty_like(family_key), + max_seqlen_q=8000, + max_seqlen_k=kv_len, + softmax_scale=0.0625, + architecture_op=object(), + ) + assert not flash_v100._should_use_prefill_d256_gqa_architecture( + query[:, :7999], + key, + value, + max_seqlen_q=7999, + max_seqlen_k=128000, + softmax_scale=0.0625, + architecture_op=object(), + ) + assert not flash_v100._should_use_prefill_d256_gqa_architecture( + query, + key[:, :127999], + value[:, :127999], + max_seqlen_q=8000, + max_seqlen_k=127999, + softmax_scale=0.0625, + architecture_op=object(), + ) + assert not flash_v100._should_use_prefill_d256_gqa_architecture( + query, + key[:, :44000], + value[:, :44000], + max_seqlen_q=8000, + max_seqlen_k=44000, + softmax_scale=0.0625, + architecture_op=object(), + ) + assert not flash_v100._should_use_prefill_d256_gqa_architecture( + query, + key, + value, + max_seqlen_q=8000, + max_seqlen_k=128000, + softmax_scale=1.0, + architecture_op=object(), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_prefill_d256_gqa_architecture_oom_uses_dense_fallback(monkeypatch): + import vllm.envs as envs + import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 + + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70/V100 is required") + + name = "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL" + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + monkeypatch.setattr(flash_v100, "_is_cuda_graph_capturing", lambda _: False) + + dense_calls = 0 + + def dense_op(query, key, value, out, softmax_scale, causal): + nonlocal dense_calls + dense_calls += 1 + out.fill_(3) + return out + + def architecture_oom(*args, **kwargs): + raise torch.OutOfMemoryError("expected architecture workspace OOM") + + monkeypatch.setattr( + flash_v100, + "_get_sm70_splitd_d256_ops", + lambda: (dense_op, object(), None), + ) + monkeypatch.setattr( + flash_v100, + "_get_sm70_d256_gqa_architecture_op", + lambda: architecture_oom, + ) + monkeypatch.setattr(flash_v100, "_warned_prefill_d256_gqa_architecture_oom", False) + + query = torch.zeros((1, 8000, 6, 256), dtype=torch.float16, device="cuda") + key = torch.zeros((1, 40000, 1, 256), dtype=torch.float16, device="cuda") + value = torch.zeros_like(key) + out = torch.zeros_like(query) + cu_seqlens_q = torch.tensor([0, 8000], dtype=torch.int32, device="cuda") + cu_seqlens_k = torch.tensor([0, 40000], dtype=torch.int32, device="cuda") + + result = flash_v100._try_sm70_fa2_d256_prefill( + query, + key, + value, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=8000, + max_seqlen_k=40000, + softmax_scale=0.0625, + causal=True, + window_size=(-1, -1), + out=out, + ) + + assert result is not None + assert result.data_ptr() == out.data_ptr() + assert dense_calls == 1 + assert torch.all(out == 3) + assert flash_v100._warned_prefill_d256_gqa_architecture_oom is True + + def test_dense_prefill_restores_uniform_batch_view_for_exact_splitd(monkeypatch): import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index e68dcee145..497cc91ce1 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -48,6 +48,7 @@ from vllm.v1.worker.gpu.spec_decode import init_speculator from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator from vllm.v1.worker.gpu.spec_decode.dflash2.sparse_rejection import ( + _parse_alignment_steps, _supports_sparse_sampling_contract, ) from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( @@ -73,7 +74,6 @@ def test_dflash2_gdn_fastpaths_are_default_off(monkeypatch): "VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS", "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", "VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC", - "VLLM_SM70_TP4_PUSH_ALLREDUCE", ) for name in names: monkeypatch.delenv(name, raising=False) @@ -84,6 +84,18 @@ def test_dflash2_gdn_fastpaths_are_default_off(monkeypatch): envs.disable_envs_cache() +def test_sm70_tp4_push_allreduce_is_default_on_with_rollback(monkeypatch): + monkeypatch.delenv("VLLM_SM70_TP4_PUSH_ALLREDUCE", raising=False) + envs.disable_envs_cache() + try: + assert envs.VLLM_SM70_TP4_PUSH_ALLREDUCE + monkeypatch.setenv("VLLM_SM70_TP4_PUSH_ALLREDUCE", "0") + envs.disable_envs_cache() + assert not envs.VLLM_SM70_TP4_PUSH_ALLREDUCE + finally: + envs.disable_envs_cache() + + def _bare_dflash2_model() -> DFlash2Qwen3Model: model = DFlash2Qwen3Model.__new__(DFlash2Qwen3Model) torch.nn.Module.__init__(model) @@ -260,10 +272,12 @@ def test_selector_default_path_does_not_allocate_sparse_score_cache(monkeypatch) allocated = torch.zeros((2, 7, 31), dtype=torch.float32) _stub_base(monkeypatch, allocated) monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", False) + monkeypatch.setattr(envs, "VLLM_SPEC_DUMP_ALIGNMENT", False) speculator = DFlash2Speculator(None, torch.device("cpu")) assert speculator.draft_logits is allocated assert torch.isneginf(speculator.draft_logits).all() assert speculator.get_sparse_draft_logits() is None + assert speculator.get_selector_alignment_shadow() is None def test_selector_opt_in_allocates_sparse_score_cache(monkeypatch): @@ -279,6 +293,39 @@ def test_selector_opt_in_allocates_sparse_score_cache(monkeypatch): assert candidate_scores.dtype is torch.float32 +def test_selector_alignment_shadow_is_explicit_and_keeps_full_lattice(monkeypatch): + allocated = torch.zeros((2, 7, 31), dtype=torch.float32) + _stub_base(monkeypatch, allocated) + monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", True) + monkeypatch.setattr(envs, "VLLM_SPEC_DUMP_ALIGNMENT", True) + + speculator = DFlash2Speculator(None, torch.device("cpu")) + shadow = speculator.get_selector_alignment_shadow() + + assert shadow is not None + candidate_ids, unary_logits, lattice_scores = shadow + assert candidate_ids.shape == (2, 7, 16) + assert unary_logits.shape == (2, 7, 16) + assert unary_logits.dtype is torch.float32 + assert lattice_scores.shape == (2, 7, 16, 16) + assert lattice_scores.dtype is torch.float32 + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, None), + ("", None), + ("1,3-5", {1, 3, 4, 5}), + ("5-3", set()), + ("bad", set()), + ("-1", set()), + ], +) +def test_selector_alignment_step_filter(raw, expected): + assert _parse_alignment_steps(raw) == expected + + def test_selector_uses_checkpoint_top16_and_fp32_proposal_cache(monkeypatch): _stub_base(monkeypatch, None) speculator = DFlash2Speculator(None, torch.device("cpu")) @@ -972,6 +1019,9 @@ def test_dflash_intermediate_prefill_materializes_context_without_query(monkeypa speculator._context_slot_mappings = torch.zeros(1, 8, dtype=torch.int64) speculator._layer_group_idx = None speculator._context_only_prefill_logged = False + speculator._prepare_ngram_assist = Mock( + side_effect=AssertionError("ngram lookup must not run mid-prefill") + ) speculator.model = SimpleNamespace( precompute_and_store_context_kv=Mock(), ) diff --git a/tests/v1/spec_decode/test_dflash2_ngram_assist.py b/tests/v1/spec_decode/test_dflash2_ngram_assist.py new file mode 100644 index 0000000000..cbf44d08f4 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash2_ngram_assist.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm.config.speculative import SpeculativeConfig +from vllm.v1.spec_decode.ngram_proposer import ( + _find_longest_matched_ngram_and_propose_tokens, +) +from vllm.v1.worker.gpu.spec_decode.dflash2.ngram_assist import ( + DFlash2NgramAssist, + find_split_ngram_proposal, +) +from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( + DFlash2Speculator, + _apply_ngram_draft_kernel, +) +from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( + dflash2_sparse_topk_rejection_sample, + rejection_sample, +) + + +@pytest.mark.parametrize("split", [0, 1, 4, 7, 10]) +@pytest.mark.parametrize(("min_ngram", "max_ngram"), [(1, 1), (2, 3), (3, 5)]) +def test_split_lookup_matches_standalone_ngram( + split: int, min_ngram: int, max_ngram: int +) -> None: + tokens = np.array([1, 2, 3, 4, 1, 2, 3, 5, 1, 2], dtype=np.int32) + expected = _find_longest_matched_ngram_and_propose_tokens( + tokens, + min_ngram, + max_ngram, + max_model_len=64, + k=3, + ) + prefix = tokens[:split] + suffix = tokens[split:].astype(np.int64) + actual = find_split_ngram_proposal( + prefix, + len(prefix), + suffix, + len(suffix), + min_ngram, + max_ngram, + max_model_len=64, + num_draft_tokens=3, + ) + np.testing.assert_array_equal(actual, expected) + + +def test_batch_assist_reports_only_full_width_hits() -> None: + assist = DFlash2NgramAssist(2, 3, num_draft_tokens=2, max_model_len=64) + token_ids = np.zeros((3, 16), dtype=np.int32) + token_ids[2, :6] = [1, 2, 3, 4, 1, 2] + token_ids[0, :4] = [8, 9, 10, 11] + sampled = np.array([[3, -1], [12, -1]], dtype=np.int64) + sampled_lens = np.array([1, 1], dtype=np.int32) + output = np.zeros((2, 2), dtype=np.int64) + output_lens = np.zeros(2, dtype=np.int32) + + hits = assist.propose( + token_ids, + req_state_indices=np.array([2, 0], dtype=np.int32), + prior_lengths=np.array([6, 4], dtype=np.int32), + sampled_token_ids=sampled, + num_sampled_tokens=sampled_lens, + eligible=np.array([True, True]), + output_tokens=output, + output_lengths=output_lens, + ) + + assert hits == 1 + np.testing.assert_array_equal(output[0], [4, 1]) + np.testing.assert_array_equal(output_lens, [2, 0]) + assert assist.num_eligible == 2 + assert assist.num_full_hits == 1 + + +class _CopyEvent: + def __init__(self) -> None: + self.synchronized = False + + def synchronize(self) -> None: + self.synchronized = True + + +def _host_only_speculator() -> DFlash2Speculator: + speculator = object.__new__(DFlash2Speculator) + speculator._ngram_assist = DFlash2NgramAssist( + 2, 3, num_draft_tokens=2, max_model_len=64 + ) + speculator._ngram_num_hits = 0 + speculator._ngram_rounds = 0 + speculator._ngram_skipped_rounds = 0 + speculator.num_speculative_steps = 2 + speculator._ngram_tokens_cpu_tensor = torch.zeros((2, 2), dtype=torch.int64) + speculator._ngram_lengths_cpu_tensor = torch.zeros(2, dtype=torch.int32) + speculator._ngram_tokens_cpu = speculator._ngram_tokens_cpu_tensor.numpy() + speculator._ngram_lengths_cpu = speculator._ngram_lengths_cpu_tensor.numpy() + speculator._ngram_tokens = torch.zeros((2, 2), dtype=torch.int64) + speculator._ngram_lengths = torch.zeros(2, dtype=torch.int32) + return speculator + + +def test_host_token_state_is_requested_only_for_enabled_assist() -> None: + speculator = object.__new__(DFlash2Speculator) + speculator._ngram_assist = None + assert not speculator.requires_host_token_state + + speculator._ngram_assist = object() # type: ignore[assignment] + assert speculator.requires_host_token_state + + +def test_prepare_assist_uses_request_slots_and_skips_only_all_hit() -> None: + speculator = _host_only_speculator() + token_ids = np.zeros((3, 16), dtype=np.int32) + token_ids[2, :6] = [1, 2, 3, 4, 1, 2] + token_ids[0, :6] = [7, 8, 9, 10, 7, 8] + batch = SimpleNamespace( + has_structured_output_reqs=False, + num_reqs=2, + num_draft_tokens_per_req=np.array([2, 2], dtype=np.int32), + seq_lens_cpu_upper_bound=torch.tensor([8, 8], dtype=torch.int32), + is_prefilling_np=np.array([False, False]), + idx_mapping_np=np.array([2, 0], dtype=np.int32), + ) + event = _CopyEvent() + + skip = speculator._prepare_ngram_assist( + batch, + event, + sampled_token_ids_cpu=np.array([[3], [9]], dtype=np.int64), + num_sampled_tokens_cpu=np.array([1, 1], dtype=np.int32), + all_token_ids_cpu=token_ids, + ) + + assert event.synchronized + assert skip + assert speculator._ngram_num_hits == 2 + assert torch.equal( + speculator._ngram_tokens, + torch.tensor([[4, 1], [10, 7]], dtype=torch.int64), + ) + + +def test_prepare_assist_keeps_dflash_draft_for_mixed_hits() -> None: + speculator = _host_only_speculator() + token_ids = np.zeros((3, 16), dtype=np.int32) + token_ids[2, :6] = [1, 2, 3, 4, 1, 2] + token_ids[0, :6] = [7, 8, 9, 10, 11, 12] + batch = SimpleNamespace( + has_structured_output_reqs=False, + num_reqs=2, + num_draft_tokens_per_req=np.array([2, 2], dtype=np.int32), + seq_lens_cpu_upper_bound=torch.tensor([8, 8], dtype=torch.int32), + is_prefilling_np=np.array([False, False]), + idx_mapping_np=np.array([2, 0], dtype=np.int32), + ) + event = _CopyEvent() + + skip = speculator._prepare_ngram_assist( + batch, + event, + sampled_token_ids_cpu=np.array([[3], [13]], dtype=np.int64), + num_sampled_tokens_cpu=np.array([1, 1], dtype=np.int32), + all_token_ids_cpu=token_ids, + ) + + assert event.synchronized + assert not skip + assert speculator._ngram_assist is not None + assert speculator._ngram_assist.num_full_hits == 1 + assert speculator._ngram_num_hits == 0 + assert torch.count_nonzero(speculator._ngram_tokens) == 0 + + +def test_prepare_assist_bypasses_structured_output_without_sync() -> None: + speculator = _host_only_speculator() + batch = SimpleNamespace(has_structured_output_reqs=True) + event = _CopyEvent() + + assert not speculator._prepare_ngram_assist( + batch, + event, + sampled_token_ids_cpu=np.zeros((1, 1), dtype=np.int64), + num_sampled_tokens_cpu=np.ones(1, dtype=np.int32), + all_token_ids_cpu=np.zeros((1, 8), dtype=np.int32), + ) + assert not event.synchronized + + +def test_ngram_assist_rejects_non_dflash_method() -> None: + with pytest.raises(ValueError, match="only supported with method='dflash'"): + SpeculativeConfig( + method="ngram", + num_speculative_tokens=2, + ngram_assist=True, + ) + + +def test_ngram_one_hot_cache_overrides_only_hit_rows() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the Triton cache kernel") + + device = torch.device("cuda") + num_reqs, num_steps, top_k, vocab_size = 2, 3, 4, 32 + ngram_tokens = torch.tensor( + [[11, 12, 13], [21, 22, 23]], dtype=torch.int64, device=device + ) + ngram_lengths = torch.tensor([3, 0], dtype=torch.int32, device=device) + # Batch row 0 belongs to request-state slot 1. + sample_req_state = torch.tensor( + [1, 1, 1, 0, 0, 0], dtype=torch.int32, device=device + ) + draft_tokens = torch.full( + (num_reqs, num_steps), 7, dtype=torch.int64, device=device + ) + cached_ids = ( + torch.arange( + num_reqs * num_steps * top_k, dtype=torch.int64, device=device + ).view(num_reqs, num_steps, top_k) + % vocab_size + ) + cached_scores = torch.randn( + num_reqs, num_steps, top_k, dtype=torch.float32, device=device + ) + draft_logits = torch.full( + (num_reqs, num_steps, vocab_size), + -float("inf"), + dtype=torch.float32, + device=device, + ) + draft_logits.scatter_(2, cached_ids, cached_scores) + miss_ids = cached_ids[0].clone() + miss_scores = cached_scores[0].clone() + miss_logits = draft_logits[0].clone() + + _apply_ngram_draft_kernel[(num_reqs * num_steps,)]( + ngram_tokens, + ngram_lengths, + sample_req_state, + draft_tokens, + draft_tokens.stride(0), + cached_ids, + cached_scores, + cached_ids.stride(0), + cached_ids.stride(1), + draft_logits, + draft_logits.stride(0), + draft_logits.stride(1), + num_steps=num_steps, + top_k=top_k, + BLOCK_K=top_k, + CACHE_DRAFT_LOGITS=True, + CACHE_SCORES=True, + num_warps=1, + ) + + assert torch.equal(draft_tokens[0], ngram_tokens[0]) + assert torch.equal(draft_tokens[1], torch.full((num_steps,), 7, device=device)) + assert torch.equal(cached_ids[0], miss_ids) + assert torch.equal(cached_scores[0], miss_scores) + assert torch.equal(draft_logits[0], miss_logits) + assert torch.equal(cached_ids[1, :, 0], ngram_tokens[0]) + assert torch.equal(cached_scores[1, :, 0], torch.zeros(num_steps, device=device)) + assert torch.isneginf(cached_scores[1, :, 1:]).all() + assert torch.equal( + draft_logits[1].gather(1, ngram_tokens[0, :, None]).squeeze(1), + torch.zeros(num_steps, device=device), + ) + + +@pytest.mark.parametrize("top_p", [1.0, 0.95]) +def test_ngram_one_hot_sparse_rejection_matches_dense(top_p: float) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for rejection sampling") + + from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p + + torch.manual_seed(20260825) + device = torch.device("cuda") + num_reqs, num_steps = 16, 3 + target_top_k, draft_top_k, vocab_size = 20, 16, 97 + rows_per_req = num_steps + 1 + num_logits = num_reqs * rows_per_req + + target_ids = torch.stack( + [ + torch.randperm(vocab_size, device=device)[:target_top_k] + for _ in range(num_logits) + ] + ) + target_logits = ( + torch.randn(num_logits, target_top_k, dtype=torch.float32, device=device) + .sort(dim=-1, descending=True) + .values + ) + target_dense = torch.full( + (num_logits, vocab_size), + -float("inf"), + dtype=torch.float32, + device=device, + ) + target_dense.scatter_(1, target_ids, target_logits) + target_dense = apply_top_k_top_p( + target_dense, + torch.full((num_logits,), target_top_k, dtype=torch.int32, device=device), + torch.full((num_logits,), top_p, dtype=torch.float32, device=device), + ) + + proposals = torch.randint( + 0, vocab_size, (num_reqs, num_steps), dtype=torch.int64, device=device + ) + draft_ids = torch.zeros( + num_reqs, num_steps, draft_top_k, dtype=torch.int64, device=device + ) + draft_ids[:, :, 0] = proposals + draft_logits = torch.full( + (num_reqs, num_steps, draft_top_k), + -float("inf"), + dtype=torch.float32, + device=device, + ) + draft_logits[:, :, 0] = 0.0 + draft_dense = torch.full( + (num_reqs, num_steps, vocab_size), + -float("inf"), + dtype=torch.float32, + device=device, + ) + draft_dense.scatter_(2, proposals[:, :, None], 0.0) + + draft_sampled_2d = torch.zeros( + num_reqs, rows_per_req, dtype=torch.int64, device=device + ) + draft_sampled_2d[:, 1:] = proposals + draft_sampled = draft_sampled_2d.flatten() + cu_num_logits = ( + torch.arange(num_reqs + 1, dtype=torch.int32, device=device) * rows_per_req + ) + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) + expanded_idx_mapping = idx_mapping.repeat_interleave(rows_per_req) + expanded_local_pos = torch.arange( + rows_per_req, dtype=torch.int32, device=device + ).repeat(num_reqs) + positions = torch.arange(num_logits, dtype=torch.int64, device=device) + 4096 + temperature = torch.ones(num_reqs, dtype=torch.float32, device=device) + seeds = torch.arange(100, 100 + num_reqs, dtype=torch.int64, device=device) + + dense_sampled, dense_lengths = rejection_sample( + target_dense, + draft_dense, + draft_sampled, + cu_num_logits, + positions, + idx_mapping, + expanded_idx_mapping, + expanded_local_pos, + temperature, + seeds, + num_steps, + ) + sparse_sampled, sparse_lengths = dflash2_sparse_topk_rejection_sample( + target_ids, + target_logits, + draft_ids, + draft_logits, + draft_sampled, + cu_num_logits, + positions, + idx_mapping, + temperature, + torch.full((num_reqs,), top_p, dtype=torch.float32, device=device), + seeds, + num_steps, + ) + + assert torch.equal(sparse_lengths, dense_lengths) + steps = torch.arange(rows_per_req, device=device).unsqueeze(0) + valid = steps < dense_lengths[:, None] + assert torch.equal(sparse_sampled[valid], dense_sampled[valid]) diff --git a/tests/v1/spec_decode/test_dflash2_selector_analysis.py b/tests/v1/spec_decode/test_dflash2_selector_analysis.py new file mode 100644 index 0000000000..3316e78b6c --- /dev/null +++ b/tests/v1/spec_decode/test_dflash2_selector_analysis.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib.util +import sys +from pathlib import Path + +import torch + +_ANALYZER_PATH = ( + Path(__file__).parents[3] / "benchmarks/analyze_dflash2_selector_alignment.py" +) +_SPEC = importlib.util.spec_from_file_location( + "analyze_dflash2_selector_alignment", _ANALYZER_PATH +) +assert _SPEC is not None and _SPEC.loader is not None +_ANALYZER = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _ANALYZER +_SPEC.loader.exec_module(_ANALYZER) + + +def _record(): + candidate_ids = torch.tensor([[10, 11], [20, 21]]) + lattice = torch.tensor( + [ + [[0.0, 0.0], [0.0, 0.0]], + [[0.0, 4.0], [0.0, -4.0]], + ], + dtype=torch.float64, + ) + return _ANALYZER.AlignmentRecord( + path=Path("synthetic.pt"), + step=1, + temperature=1.0, + top_p=1.0, + target_topk_ids=candidate_ids, + target_topk_logits=torch.zeros((2, 2), dtype=torch.float64), + candidate_ids=candidate_ids, + realized_logits=torch.stack((lattice[0, 0], lattice[1, 0])), + unary_logits=torch.zeros((2, 2), dtype=torch.float64), + lattice_scores=lattice, + draft_sampled=torch.tensor([1, 10, 20]), + num_sampled=2, + ) + + +def test_compact_top_p_keeps_the_crossing_token(): + probs = _ANALYZER._compact_probs(torch.log(torch.tensor([0.6, 0.3, 0.1])), 0.7) + torch.testing.assert_close( + probs, torch.tensor([2 / 3, 1 / 3, 0.0], dtype=torch.float64) + ) + + +def test_beta_zero_reconstructs_realized_selector_rows(): + record = _record() + current = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig(name="current", use_cached_logits=True), + ) + reconstructed = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig(name="beta-zero"), + ) + + for actual, expected in zip(reconstructed, current): + torch.testing.assert_close(actual, expected) + + +def test_future_message_can_prefer_a_better_supported_branch(): + record = _record() + local = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig(name="local"), + )[0] + global_chain = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig(name="global", future_beta=1.0), + )[0] + + assert local[0] == local[1] + assert global_chain[0] > global_chain[1] + + +def test_greedy_mixture_is_normalized_and_moves_exact_mass(): + record = _record() + baseline = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig(name="baseline", use_cached_logits=True), + )[1] + mixed = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig( + name="mixture", + greedy_mix=0.25, + use_cached_logits=True, + ), + )[1] + + expected = baseline * 0.75 + expected[torch.argmax(baseline)] += 0.25 + torch.testing.assert_close(mixed, expected) + torch.testing.assert_close(mixed.sum(), torch.tensor(1.0, dtype=torch.float64)) diff --git a/tests/v1/worker/test_gpu_worker_static_pp.py b/tests/v1/worker/test_gpu_worker_static_pp.py new file mode 100644 index 0000000000..b7e9e8da3a --- /dev/null +++ b/tests/v1/worker/test_gpu_worker_static_pp.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +import vllm.envs as envs +from vllm.config import CUDAGraphMode +from vllm.sequence import IntermediateTensors +from vllm.v1.worker import gpu_worker +from vllm.v1.worker.gpu_model_runner import GPUModelRunner +from vllm.v1.worker.gpu_worker import Worker + + +def _worker() -> SimpleNamespace: + model_config = SimpleNamespace( + dtype=torch.float16, + ) + return SimpleNamespace( + model_runner=SimpleNamespace( + model=SimpleNamespace( + make_empty_intermediate_tensors=lambda batch_size, dtype, device: ( + IntermediateTensors( + { + "hidden_states": torch.empty( + (batch_size, 4, 4096), dtype=dtype, device=device + ) + } + ) + ) + ) + ), + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace( + pipeline_parallel_size=2, + tensor_parallel_size=4, + enable_dbo=False, + ubatch_size=0, + ), + scheduler_config=SimpleNamespace(max_num_seqs=1), + speculative_config=None, + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.FULL, + pass_config=SimpleNamespace(enable_sp=False), + ), + model_config=model_config, + ), + ) + + +def test_static_pp_admission_uses_engine_and_tensor_geometry(monkeypatch) -> None: + monkeypatch.setenv("VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER", "1") + envs.disable_envs_cache() + worker = _worker() + try: + with ( + patch.object(gpu_worker.current_platform, "is_cuda", return_value=True), + patch.object( + gpu_worker.current_platform, + "is_device_capability", + return_value=True, + ), + ): + assert Worker._use_sm70_static_pp_hidden_transfer(worker, 1) + assert not Worker._use_sm70_static_pp_hidden_transfer(worker, 2) + + wrong_schema = _worker() + wrong_schema.model_runner.model.make_empty_intermediate_tensors = ( + lambda batch_size, dtype, device: IntermediateTensors( + { + "hidden_states": torch.empty( + (batch_size, 4096), dtype=dtype, device=device + ) + } + ) + ) + assert not Worker._use_sm70_static_pp_hidden_transfer(wrong_schema, 1) + finally: + envs.disable_envs_cache() + + +def test_static_pp_admission_rejects_unsafe_concurrency(monkeypatch) -> None: + monkeypatch.setenv("VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER", "1") + envs.disable_envs_cache() + try: + for field, value in ( + ("max_num_seqs", 2), + ("pipeline_parallel_size", 1), + ("tensor_parallel_size", 8), + ("enable_dbo", True), + ("ubatch_size", 2), + ): + worker = _worker() + target = ( + worker.vllm_config.scheduler_config + if field == "max_num_seqs" + else worker.vllm_config.parallel_config + ) + setattr(target, field, value) + with ( + patch.object(gpu_worker.current_platform, "is_cuda", return_value=True), + patch.object( + gpu_worker.current_platform, + "is_device_capability", + return_value=True, + ), + ): + assert not Worker._use_sm70_static_pp_hidden_transfer(worker, 1) + finally: + envs.disable_envs_cache() + + +def test_static_pp_schema_is_exact() -> None: + hidden_states = SimpleNamespace( + shape=(1, 4, 4096), + dtype=torch.float16, + is_cuda=True, + is_contiguous=lambda: True, + ) + assert Worker._is_static_pp_hidden_tensor_dict( + {"hidden_states": hidden_states}, + 1, # type: ignore[dict-item] + ) + hidden_states.shape = (1, 4096) + assert not Worker._is_static_pp_hidden_tensor_dict( + {"hidden_states": hidden_states}, + 1, # type: ignore[dict-item] + ) + + +def test_static_pp_receive_buffer_skips_self_copy() -> None: + hidden_states = torch.ones((4, 8), dtype=torch.float16) + runner = SimpleNamespace( + intermediate_tensors=IntermediateTensors({"hidden_states": hidden_states}), + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace(tensor_parallel_size=4) + ), + ) + source = IntermediateTensors({"hidden_states": hidden_states}) + version = hidden_states._version + + with patch( + "vllm.v1.worker.gpu_model_runner.is_residual_scattered_for_sp", + return_value=False, + ): + result = GPUModelRunner.sync_and_gather_intermediate_tensors( + runner, 1, source, True + ) + + assert hidden_states._version == version + assert result["hidden_states"].data_ptr() == hidden_states.data_ptr() diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 06ff86537b..a554e2bf4d 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -19,10 +19,14 @@ def _maybe_load_fp8_qpn8_library() -> None: path lets source experiments add only the QPN8 operators to an otherwise compatible installed build, including in spawned TP workers. """ - if os.getenv("VLLM_SM70_FP8_QPN8", "0") != "1": - return library_path = os.getenv("VLLM_SM70_FP8_QPN8_LIBRARY") - if library_path: + if library_path is None: + return + generic_override = os.getenv("VLLM_SM70_FP8_QPN8") + specific_override = os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") + generic_enabled = generic_override == "1" + default_route_enabled = generic_override != "0" and specific_override != "0" + if generic_enabled or default_route_enabled: torch.ops.load_library(library_path) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index ff5282f315..b89d8b0210 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -158,6 +158,11 @@ class SpeculativeConfig: prompt_lookup_min: int | None = Field(default=None, ge=1) """Minimum size of ngram token window when using Ngram proposer, if provided. Defaults to 1.""" + ngram_assist: bool = False + """Try prompt-ngram lookup before DFlash2 draft generation. Full-width + ngram hits skip the DFlash2 query and selector while preserving its + context-KV state. Only valid with ``method='dflash'`` and a DFlash2 + selector capability.""" # Alternative drafting strategies parallel_drafting: bool = False @@ -665,6 +670,22 @@ def __post_init__(self): if self.method in ("ngram", "[ngram]"): self.method = "ngram" + if self.ngram_assist: + if self.prompt_lookup_min is None and self.prompt_lookup_max is None: + self.prompt_lookup_min = 5 + self.prompt_lookup_max = 5 + elif self.prompt_lookup_min is None: + self.prompt_lookup_min = self.prompt_lookup_max + elif self.prompt_lookup_max is None: + self.prompt_lookup_max = self.prompt_lookup_min + assert self.prompt_lookup_min is not None + assert self.prompt_lookup_max is not None + if self.prompt_lookup_min > self.prompt_lookup_max: + raise ValueError( + f"prompt_lookup_min={self.prompt_lookup_min} must " + f"be <= prompt_lookup_max={self.prompt_lookup_max}" + ) + if self.method in ("ngram", "ngram_gpu"): # Set default values if not provided if self.prompt_lookup_min is None and self.prompt_lookup_max is None: @@ -741,8 +762,9 @@ def __post_init__(self): self.draft_parallel_config = self.target_parallel_config else: - self.prompt_lookup_max = 0 - self.prompt_lookup_min = 0 + if not self.ngram_assist: + self.prompt_lookup_max = 0 + self.prompt_lookup_min = 0 if self.model is not None: self.draft_model_config = ModelConfig( @@ -1163,6 +1185,13 @@ def _verify_args(self) -> Self: "Expected num_speculative_tokens to be greater " f"than zero ({self.num_speculative_tokens})." ) + if self.ngram_assist: + if not self.use_dflash(): + raise ValueError("ngram_assist is only supported with method='dflash'.") + draft_hf_config = getattr(self.draft_model_config, "hf_config", None) + dflash_config = getattr(draft_hf_config, "dflash_config", None) or {} + if int(dflash_config.get("selector_top_k", 0) or 0) <= 0: + raise ValueError("ngram_assist requires DFlash2 selector capability.") if self.use_dflash_ddtree(): if self.ddtree_budget is None: self.ddtree_budget = self.num_speculative_tokens diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 384e3e1529..8d55ceeaff 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -339,7 +339,7 @@ def __init__( ) logger.info( "SM70 TP4 SGLang-style push all-reduce enabled for the " - "FP16 [8, 5120] verifier shape." + "FP16 80-KiB verifier and 8-KiB decode payloads." ) @contextmanager diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 9fd594d9a8..8449f958cb 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1223,6 +1223,68 @@ def isend_tensor_dict( return handles + def isend_tensor_dict_static( + self, + tensor_dict: dict[str, torch.Tensor], + dst: int | None = None, + ) -> list[Handle]: + """Send tensors without metadata when both peers know the schema. + + The receiver must provide the same ordered keys, shapes, dtypes, and + devices to :meth:`irecv_tensor_dict_static`. Callers must fall back to + :meth:`isend_tensor_dict` whenever that contract can change. + """ + if self.world_size <= 1: + return [] + if self.use_cpu_custom_send_recv: + raise RuntimeError("static tensor transfer requires a process group") + + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + assert 0 <= dst < self.world_size, f"Invalid dst rank ({dst})" + + handles: list[Handle] = [] + for tensor in tensor_dict.values(): + if not isinstance(tensor, torch.Tensor): + raise TypeError("static tensor transfer only accepts tensors") + if tensor.numel() == 0: + continue + comm_group = self.cpu_group if tensor.is_cpu else self.device_group + handle = torch.distributed.isend( + tensor, dst=self.ranks[dst], group=comm_group + ) + if tensor.is_cuda: + tensor.record_stream(torch.cuda.current_stream(tensor.device)) + handles.append(handle) + return handles + + def irecv_tensor_dict_static( + self, + tensor_dict: dict[str, torch.Tensor], + src: int | None = None, + ) -> list[Handle]: + """Receive metadata-free tensors into caller-owned buffers.""" + if self.world_size <= 1: + return [] + if self.use_cpu_custom_send_recv: + raise RuntimeError("static tensor transfer requires a process group") + + if src is None: + src = (self.rank_in_group - 1) % self.world_size + assert 0 <= src < self.world_size, f"Invalid src rank ({src})" + + handles: list[Handle] = [] + for tensor in tensor_dict.values(): + if not isinstance(tensor, torch.Tensor): + raise TypeError("static tensor transfer only accepts tensors") + if tensor.numel() == 0: + continue + comm_group = self.cpu_group if tensor.is_cpu else self.device_group + handles.append( + torch.distributed.irecv(tensor, src=self.ranks[src], group=comm_group) + ) + return handles + def recv_tensor_dict( self, src: int | None = None, diff --git a/vllm/envs.py b/vllm/envs.py index 5eaf5161b4..a9b9332d74 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -163,6 +163,7 @@ VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS_ONLY: bool = False VLLM_SM70_FP8_PREFILL_EXACT_DENSE: bool = True VLLM_SM70_FP8_QPN8: bool = False + VLLM_SM70_FP8_QPN8_PP2_TP4: bool = True VLLM_SM70_FP8_QPN8_LIBRARY: str | None = None VLLM_SM70_SAMPLER_LIBRARY: str | None = None VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM: bool = False @@ -182,6 +183,7 @@ VLLM_SM70_NVFP4_DENSE_TUNE_MAX_M: int = 16 VLLM_SM70_DSV4_FP16_GEMV: bool = False VLLM_SM70_DSV4_MHC_FP32_STAGE: bool = True + VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER: bool = True VLLM_SM70_AWQ_MOE_TUNE_MAX_TOKENS: int = 128 VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS: int = 128 VLLM_SM70_ENABLE_DENSE_F16_FASTPATH: bool = False @@ -204,7 +206,7 @@ VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS: bool = False VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION: bool = False VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC: bool = False - VLLM_SM70_TP4_PUSH_ALLREDUCE: bool = False + VLLM_SM70_TP4_PUSH_ALLREDUCE: bool = True VLLM_SM70_TOP1_CUSTOM_AR: bool = False VLLM_SM70_GREEDY_TOKEN_FASTPATH: bool = True VLLM_SM70_GREEDY_TOKEN_FASTPATH_TRACE: bool = False @@ -346,6 +348,7 @@ VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3: bool = True VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV: int = 32768 VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL: bool = False + VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL: bool = True VLLM_FLASH_V100_PREFILL_SPLIT_KV: bool = False VLLM_FLASH_V100_PREFILL_SPLIT_KV_TOKENS: int = 32768 VLLM_FLASH_V100_PREFILL_SPLIT_KV_MIN_Q: int = 1 @@ -1679,6 +1682,13 @@ def _resolve_rust_frontend_path() -> str | None: # a mixed NVFP4 checkpoint may select its separately validated default in # the compressed-tensors scheme. Explicit 0 disables both routes. "VLLM_SM70_FP8_QPN8": lambda: bool(int(os.getenv("VLLM_SM70_FP8_QPN8", "0"))), + # Default-on QPN8 route for the validated serialized PP2 x TP4 contract. + # Admission additionally requires exact operator shapes/layouts, B1, + # no speculative decoding, no DBO, and no explicit ubatching. Set this to + # 0 (or the generic QPN8 flag to 0) to retain TurboMind everywhere. + "VLLM_SM70_FP8_QPN8_PP2_TP4": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "1")) + ), # Optional source-built QPN8-only extension. Production builds leave this # unset because the same operators are linked into vllm._C. "VLLM_SM70_FP8_QPN8_LIBRARY": lambda: os.getenv("VLLM_SM70_FP8_QPN8_LIBRARY", None), @@ -1799,6 +1809,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_DSV4_MHC_FP32_STAGE": lambda: bool( int(os.getenv("VLLM_SM70_DSV4_MHC_FP32_STAGE", "1")) ), + # Skip PP metadata and TP reconstruction only for the exact, replicated + # SM70 B1 hidden-state schema validated by the worker on both stages. + "VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER": lambda: bool( + int(os.getenv("VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER", "1")) + ), "VLLM_SM70_AWQ_MOE_TUNE_MAX_TOKENS": lambda: int( os.getenv("VLLM_SM70_AWQ_MOE_TUNE_MAX_TOKENS", "128") ), @@ -1937,11 +1952,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC": lambda: bool( int(os.getenv("VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC", "0")) ), - # Opt-in SGLang-style push collective for the exact FP16 [8, 5120] - # verifier shape on fully-connected SM70 TP4. The communicator allocates - # dedicated two-epoch push storage only when this gate is enabled. + # Default-on SGLang-style push collective for the validated FP16 80-KiB + # verifier and 8-KiB decode payloads on fully-connected SM70 TP4 CUDA + # Graphs. Other devices, topologies, sizes, and eager calls retain the + # ordinary pull path; explicit 0 is the rollback. "VLLM_SM70_TP4_PUSH_ALLREDUCE": lambda: bool( - int(os.getenv("VLLM_SM70_TP4_PUSH_ALLREDUCE", "0")) + int(os.getenv("VLLM_SM70_TP4_PUSH_ALLREDUCE", "1")) ), # Safe greedy-only shortcut: avoid full vocab all-gather/sampler work when # the request batch is pure greedy and has no penalties, logprobs, grammar, @@ -2455,6 +2471,14 @@ def _resolve_rust_frontend_path() -> str | None: ) ) ), + "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL": lambda: bool( + int( + os.getenv( + "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL", + "1", + ) + ) + ), "VLLM_FLASH_V100_PREFILL_SPLIT_KV": lambda: bool( int(os.getenv("VLLM_FLASH_V100_PREFILL_SPLIT_KV", "0")) ), diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 797ec93def..64b386b6cb 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -131,6 +131,29 @@ "in_proj_qkvz": (5120, 4096), "qkv_proj": (5120, 3584), } +_SM70_FP8_QPN8_PP2_TP4_CONFIGS = { + # Measured B1 winners: (K, N, fused gated-SiLU) maps to + # (split-K, accumulator chains, prefetch codes). + (4096, 1536, False): (32, 2, False), + (1024, 8192, False): (8, 2, False), + (2048, 4096, False): (16, 2, False), + (4096, 1024, False): (32, 2, False), + (4096, 1024, True): (16, 2, False), + (512, 4096, False): (16, 2, False), +} +_SM70_FP8_QPN8_PP2_TP4_SHAPES = { + # Operator role: accepted (layer TP size, K, N) tuples. The replicated + # indexer wq_b is deliberately excluded by TP size: its long-prefill work + # may overlap the main wq_b and cannot share one dense fallback workspace. + "fused_wqa_wkv": {(1, 4096, 1536)}, + "wq_b": {(4, 1024, 8192)}, + "wo_b": {(4, 2048, 4096)}, + "gate_up_proj": {(4, 4096, 1024)}, + "down_proj": {(4, 512, 4096)}, +} +_SM70_FP8_QPN8_PP2_TP4_WORKSPACE_ELEMENTS = max( + k * n for k, n, _ in _SM70_FP8_QPN8_PP2_TP4_CONFIGS +) _SM70_FP8_QPN8_REQUIRED_OPS = ( "fp8_qpn8_prepare_sm70", "fp8_qpn8_dequantize_sm70_out", @@ -142,6 +165,7 @@ _SM70_FP8_QPN8_MAX_NUM_SEQS = 8 # Layers retain only data_ptr(), so this cache owns each allocation's lifetime. _sm70_fp8_prefill_dense_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} +_sm70_fp8_qpn8_pp2_tp4_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} def _is_sm70_fp8_exact_8k_prefill_layer(layer: torch.nn.Module) -> bool: @@ -206,6 +230,78 @@ def _is_sm70_fp8_qpn8_runtime_contract() -> bool: return os.getenv("VLLM_SM70_FP8_QPN8") is not None +def _sm70_fp8_qpn8_pp2_tp4_enabled() -> bool: + """Resolve the validated default while retaining explicit rollback.""" + generic_override = os.getenv("VLLM_SM70_FP8_QPN8") + if generic_override is not None and not envs.VLLM_SM70_FP8_QPN8: + return False + specific_override = os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") + if specific_override is not None: + return envs.VLLM_SM70_FP8_QPN8_PP2_TP4 + if generic_override is not None: + return envs.VLLM_SM70_FP8_QPN8 + return envs.VLLM_SM70_FP8_QPN8_PP2_TP4 + + +def _is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() -> bool: + """Require the measured serialized PP2 x TP4 single-request route.""" + vllm_config = get_current_vllm_config() + parallel_config = vllm_config.parallel_config + scheduler_config = vllm_config.scheduler_config + return bool( + parallel_config.pipeline_parallel_size == 2 + and parallel_config.tensor_parallel_size == 4 + and scheduler_config.max_num_seqs == 1 + and not getattr(parallel_config, "enable_dbo", False) + and int(getattr(parallel_config, "ubatch_size", 0)) <= 1 + and getattr(vllm_config, "speculative_config", None) is None + ) + + +def _sm70_fp8_qpn8_pp2_tp4_config( + layer: torch.nn.Module, *, gated_silu: bool +) -> tuple[int, int, bool] | None: + """Select by operator/tensor contract, never model or checkpoint identity.""" + if getattr(layer, "weight_block_size", None) != [128, 128]: + return None + suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] + accepted = _SM70_FP8_QPN8_PP2_TP4_SHAPES.get(suffix) + if accepted is None: + return None + k_dim = int(getattr(layer, "input_size_per_partition", 0)) + n_dim = int(getattr(layer, "output_size_per_partition", 0)) + layer_contract = (int(getattr(layer, "tp_size", 1)), k_dim, n_dim) + if layer_contract not in accepted: + return None + if tuple(reversed(layer.weight.shape)) != (k_dim, n_dim): + return None + if gated_silu: + output_partitions = getattr(layer, "output_partition_sizes", None) + if ( + suffix != "gate_up_proj" + or not isinstance(output_partitions, list) + or output_partitions != [n_dim // 2, n_dim // 2] + ): + return None + return _SM70_FP8_QPN8_PP2_TP4_CONFIGS.get((k_dim, n_dim, gated_silu)) + + +def _sm70_fp8_qpn8_pp2_tp4_bmm_config( + layer: torch.nn.Module, +) -> tuple[int, int, bool] | None: + if ( + getattr(layer, "prefix", "").rsplit(".", 1)[-1] != "wo_a" + or getattr(layer, "weight_block_size", None) != [128, 128] + or int(getattr(layer, "tp_size", 1)) != 4 + or int(getattr(layer, "bmm_batch_size", 0)) != 2 + or int(getattr(layer, "input_size_per_partition", 0)) != 4096 + or int(getattr(layer, "output_size_per_partition", 0)) != 2048 + or tuple(layer.weight.shape) != (2048, 4096) + ): + return None + return _SM70_FP8_QPN8_PP2_TP4_CONFIGS[(4096, 1024, False)] + + def _missing_sm70_fp8_qpn8_ops() -> list[str]: return [ name for name in _SM70_FP8_QPN8_REQUIRED_OPS if not hasattr(torch.ops._C, name) @@ -238,6 +334,33 @@ def _get_sm70_fp8_prefill_exact_dense_workspace( return workspace +def _get_sm70_fp8_qpn8_pp2_tp4_workspace( + weight: torch.Tensor, +) -> torch.Tensor | None: + """Allocate one bounded FP16 prefill fallback per device.""" + device_index = weight.device.index + if device_index is None: + device_index = torch.accelerator.current_device_index() + cache_key = (device_index, torch.float16) + workspace = _sm70_fp8_qpn8_pp2_tp4_workspaces.get(cache_key) + if workspace is not None: + return workspace + try: + workspace = torch.empty( + (_SM70_FP8_QPN8_PP2_TP4_WORKSPACE_ELEMENTS,), + dtype=torch.float16, + device=weight.device, + ) + except torch.OutOfMemoryError: + logger.warning_once( + "Insufficient memory for the bounded SM70 PP2 x TP4 QPN8 " + "prefill workspace; retaining TurboMind FP8." + ) + return None + _sm70_fp8_qpn8_pp2_tp4_workspaces[cache_key] = workspace + return workspace + + def _sm70_fp8_prefill_visible_dense_mm( input: torch.Tensor, weight: torch.Tensor, @@ -665,6 +788,15 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: if weight_scale_inv.dtype != torch.float32: weight_scale_inv = weight_scale_inv.to(torch.float32) if getattr(layer, "is_bmm", False): + qpn8_bmm_config = ( + _sm70_fp8_qpn8_pp2_tp4_bmm_config(layer) + if _sm70_fp8_qpn8_pp2_tp4_enabled() + else None + ) + qpn8_bmm_runtime = bool( + qpn8_bmm_config is not None + and _is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + ) group_count = int(getattr(layer, "bmm_batch_size", 0)) if group_count <= 0 or weight.shape[0] % group_count != 0: raise RuntimeError( @@ -688,6 +820,73 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: f"{weight_scale_inv.shape[0]}." ) + if qpn8_bmm_config is not None and not qpn8_bmm_runtime: + logger.info_once( + "Grouped SM70 QPN8 retains TurboMind outside the " + "serialized PP2 x TP4 single-request contract." + ) + if qpn8_bmm_runtime: + missing_ops = _missing_sm70_fp8_qpn8_ops() + explicitly_enabled = ( + os.getenv("VLLM_SM70_FP8_QPN8") == "1" + or os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") == "1" + ) + if missing_ops and explicitly_enabled: + raise RuntimeError( + "The explicitly enabled SM70 PP2 x TP4 QPN8 route " + f"requires source-built operators; missing: {missing_ops}." + ) + if missing_ops: + logger.warning_once( + "The default SM70 PP2 x TP4 QPN8 route is unavailable " + "in the loaded vllm._C; retaining TurboMind FP8." + ) + workspace = ( + None + if missing_ops + else _get_sm70_fp8_qpn8_pp2_tp4_workspace(weight) + ) + if workspace is not None: + qpn8_weights = [] + qpn8_scales = [] + for group_idx in range(group_count): + row_start = group_idx * rows_per_group + scale_start = group_idx * scale_rows_per_group + qpn8_weight, qpn8_scale = sm70_ops.fp8_qpn8_prepare_sm70( + weight[ + row_start : row_start + rows_per_group + ].contiguous(), + weight_scale_inv[ + scale_start : scale_start + scale_rows_per_group + ].contiguous(), + ) + qpn8_weights.append(qpn8_weight) + qpn8_scales.append(qpn8_scale) + replace_parameter(layer, "weight", torch.stack(qpn8_weights)) + replace_parameter( + layer, "weight_scale_inv", torch.stack(qpn8_scales) + ) + assert qpn8_bmm_config is not None + split_k, nacc, prefetch = qpn8_bmm_config + layer.input_scale = None + layer.sm70_fp8_turbomind = True + layer.sm70_fp8_qpn8 = True + layer.sm70_fp8_qpn8_bmm = True + layer.sm70_fp8_bmm = True + layer.sm70_fp8_bmm_groups = group_count + layer.sm70_fp8_bmm_output_size = rows_per_group + layer.sm70_fp8_qpn8_split_k = split_k + layer.sm70_fp8_qpn8_nacc = nacc + layer.sm70_fp8_qpn8_prefetch = prefetch + layer.sm70_fp8_prefill_exact_dense_workspace_ptr = ( + workspace.data_ptr() + ) + logger.info_once( + "Default SM70 grouped QPN8 enabled for the validated " + "serialized PP2 x TP4 tensor contract." + ) + return + prepared_weights = [] prepared_scales = [] metas = [] @@ -733,43 +932,73 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: return is_gated_silu_layer = self._is_sm70_gated_silu_layer(layer) use_gated_silu = is_gated_silu_layer and envs.VLLM_SM70_FP8_DENSE_GATED_SILU - qpn8_candidate_layer = envs.VLLM_SM70_FP8_QPN8 and _is_sm70_fp8_qpn8_layer( - layer + generic_qpn8_candidate = ( + envs.VLLM_SM70_FP8_QPN8 and _is_sm70_fp8_qpn8_layer(layer) ) - qpn8_concurrency = ( - _is_sm70_fp8_qpn8_runtime_contract() if qpn8_candidate_layer else False + pp2_tp4_qpn8_config = ( + _sm70_fp8_qpn8_pp2_tp4_config(layer, gated_silu=False) + if _sm70_fp8_qpn8_pp2_tp4_enabled() + else None ) - if qpn8_candidate_layer and not qpn8_concurrency: + pp2_tp4_qpn8_candidate = pp2_tp4_qpn8_config is not None + qpn8_candidate_layer = generic_qpn8_candidate or pp2_tp4_qpn8_candidate + qpn8_runtime = bool( + ( + pp2_tp4_qpn8_candidate + and _is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + ) + or (generic_qpn8_candidate and _is_sm70_fp8_qpn8_runtime_contract()) + ) + if qpn8_candidate_layer and not qpn8_runtime: logger.info_once( "The SM70 FP8 QPN8 route retains TurboMind unless its " "bounded-concurrency runtime contract is explicit." ) - if qpn8_candidate_layer and qpn8_concurrency: + if qpn8_candidate_layer and qpn8_runtime: + pp2_tp4_gated_config = None + if pp2_tp4_qpn8_candidate and use_gated_silu: + pp2_tp4_gated_config = _sm70_fp8_qpn8_pp2_tp4_config( + layer, gated_silu=True + ) + if pp2_tp4_gated_config is None: + raise RuntimeError( + "The SM70 PP2 x TP4 QPN8 gate/up layer violated " + "its fused-SiLU tensor contract." + ) missing_ops = _missing_sm70_fp8_qpn8_ops() if missing_ops: - if os.getenv("VLLM_SM70_FP8_QPN8") is not None: + explicitly_enabled = ( + os.getenv("VLLM_SM70_FP8_QPN8") == "1" + or os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") == "1" + ) + if explicitly_enabled: raise RuntimeError( - "VLLM_SM70_FP8_QPN8=1 requires the source-built SM70 " - f"QPN8 extension; missing ops: {missing_ops}." + "The explicitly enabled SM70 QPN8 route requires " + f"source-built operators; missing: {missing_ops}." ) logger.warning_once( - "The automatic SM70 FP8 QPN8 route is unavailable in " + "The default SM70 FP8 QPN8 route is unavailable in " "the loaded vllm._C; retaining the TurboMind layout." ) - workspace = ( - None - if missing_ops - else _get_sm70_fp8_prefill_exact_dense_workspace(weight) - ) + if missing_ops: + workspace = None + elif pp2_tp4_qpn8_candidate: + workspace = _get_sm70_fp8_qpn8_pp2_tp4_workspace(weight) + else: + workspace = _get_sm70_fp8_prefill_exact_dense_workspace(weight) if not missing_ops and workspace is not None: qpn8_codes, qpn8_scales = sm70_ops.fp8_qpn8_prepare_sm70( weight, weight_scale_inv ) k_dim, n_dim = (int(dim) for dim in qpn8_codes.shape) - split_k, nacc, prefetch = _SM70_FP8_QPN8_CONFIGS[ - (k_dim, n_dim, False) - ] + if pp2_tp4_qpn8_candidate: + assert pp2_tp4_qpn8_config is not None + split_k, nacc, prefetch = pp2_tp4_qpn8_config + else: + split_k, nacc, prefetch = _SM70_FP8_QPN8_CONFIGS[ + (k_dim, n_dim, False) + ] replace_parameter(layer, "weight", qpn8_codes) replace_parameter(layer, "weight_scale_inv", qpn8_scales) layer.input_scale = None @@ -782,18 +1011,30 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: workspace.data_ptr() ) if use_gated_silu: - gated_split_k, gated_nacc, gated_prefetch = ( - _SM70_FP8_QPN8_CONFIGS[(k_dim, n_dim, True)] - ) + if pp2_tp4_qpn8_candidate: + assert pp2_tp4_gated_config is not None + gated_split_k, gated_nacc, gated_prefetch = ( + pp2_tp4_gated_config + ) + else: + gated_split_k, gated_nacc, gated_prefetch = ( + _SM70_FP8_QPN8_CONFIGS[(k_dim, n_dim, True)] + ) layer.sm70_fp8_gated_silu = True layer.sm70_fp8_gated_silu_primary = True layer.sm70_fp8_qpn8_gated_split_k = gated_split_k layer.sm70_fp8_qpn8_gated_nacc = gated_nacc layer.sm70_fp8_qpn8_gated_prefetch = gated_prefetch - logger.info_once( - "Memory-neutral SM70 FP8 QPN8 path enabled for accepted " - "TP4 block-FP8 operator shapes." - ) + if pp2_tp4_qpn8_candidate: + logger.info_once( + "Default SM70 QPN8 enabled for the validated " + "serialized PP2 x TP4 operator contract." + ) + else: + logger.info_once( + "Memory-neutral SM70 FP8 QPN8 path enabled for " + "accepted TP4 block-FP8 operator shapes." + ) return if not missing_ops: logger.warning_once( @@ -983,6 +1224,44 @@ def apply( "SM70 FP8 QPN8 currently requires float16 activations, " f"got {x.dtype}." ) + if getattr(layer, "sm70_fp8_qpn8_bmm", False): + group_count = int(layer.sm70_fp8_bmm_groups) + output_size = int(layer.sm70_fp8_bmm_output_size) + if x.ndim < 2 or x.shape[-2] != group_count: + raise RuntimeError( + "SM70 grouped QPN8 input must end in [groups, K], got " + f"{tuple(x.shape)} for groups={group_count}." + ) + x_grouped = x.reshape(-1, group_count, x.shape[-1]) + x_by_group = x_grouped.transpose(0, 1).contiguous() + out_by_group = torch.empty( + (group_count, x_grouped.shape[0], output_size), + device=x.device, + dtype=x.dtype, + ) + if x_grouped.shape[0] == 0: + return out_by_group.transpose(0, 1).reshape( + *x.shape[:-2], group_count, output_size + ) + for group_idx in range(group_count): + sm70_ops.fp8_qpn8_dispatch_sm70_out( + out_by_group[group_idx], + int(layer.sm70_fp8_prefill_exact_dense_workspace_ptr), + x_by_group[group_idx], + layer.weight[group_idx], + layer.weight_scale_inv[group_idx], + int(layer.sm70_fp8_qpn8_split_k), + int(layer.sm70_fp8_qpn8_nacc), + bool(layer.sm70_fp8_qpn8_prefetch), + False, + ) + out = out_by_group.transpose(0, 1).reshape( + *x.shape[:-2], group_count, output_size + ) + if bias is not None: + out.add_(bias.view(group_count, output_size)) + return out + out_shape = (*x.shape[:-1], layer.output_size_per_partition) x_2d = x.reshape(-1, x.shape[-1]) if x_2d.stride(-1) != 1: diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 2575d3dcd4..78ad964a3c 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Deepseek-OCR model compatible with HuggingFace weights.""" -import math from collections.abc import Iterable, Mapping, Sequence from typing import Annotated, Literal @@ -53,7 +52,7 @@ BASE_SIZE, CROP_MODE, DeepseekOCRProcessor, - count_tiles, + count_image_tokens_for, ) from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.v1.sample.logits_processor import ( @@ -190,12 +189,15 @@ def get_hf_config(self): return self.ctx.get_hf_config(DeepseekVLV2Config) def get_hf_processor(self, **kwargs: object): - v1_processor_config = dict( - image_size=IMAGE_SIZE, - base_size=BASE_SIZE, - crop_mode=CROP_MODE, - strategy="v1", - ) + v1_processor_config: dict[str, object] = dict(strategy="v1") + if "image_mode" not in kwargs: + # Only inject the constant trio when no named mode is requested; + # the processor treats ``image_mode`` as authoritative. + v1_processor_config.update( + image_size=IMAGE_SIZE, + base_size=BASE_SIZE, + crop_mode=CROP_MODE, + ) return self.ctx.get_hf_processor( DeepseekOCRProcessor, @@ -208,35 +210,13 @@ def get_supported_mm_limits(self) -> Mapping[str, int | None]: def get_num_image_tokens( self, *, image_width: int, image_height: int, cropping: bool = True ) -> int: - image_size = IMAGE_SIZE - base_size = BASE_SIZE - patch_size = 16 - downsample_ratio = 4 - - if CROP_MODE: - if image_width <= 640 and image_height <= 640: - crop_ratio = [1, 1] - else: - # find the closest aspect ratio to the target - crop_ratio = count_tiles( - image_width, image_height, image_size=IMAGE_SIZE - ) - - num_width_tiles, num_height_tiles = crop_ratio - else: - num_width_tiles = num_height_tiles = 1 - - h = w = math.ceil((base_size // patch_size) / downsample_ratio) - - h2 = w2 = math.ceil((image_size // patch_size) / downsample_ratio) - - global_views_tokens = h * (w + 1) - if num_width_tiles > 1 or num_height_tiles > 1: - local_views_tokens = (num_height_tiles * h2) * (num_width_tiles * w2 + 1) - else: - local_views_tokens = 0 - - return global_views_tokens + local_views_tokens + 1 + return count_image_tokens_for( + image_width=image_width, + image_height=image_height, + base_size=BASE_SIZE, + image_size=IMAGE_SIZE, + cropping=cropping and CROP_MODE, + ) def get_image_size_with_most_features(self) -> ImageSize: if IMAGE_SIZE == 1024 and BASE_SIZE == 1280: @@ -334,10 +314,12 @@ def get_replacement_deepseek_vl2(item_idx: int): else: size = images.get_image_size(item_idx) - num_image_tokens = self.info.get_num_image_tokens( + # Count with the SAME processor instance the request's + # mm kwargs construct (mode/crops) — the producer/counter + # pair must never diverge. + num_image_tokens = hf_processor.count_image_tokens( image_width=size.width, image_height=size.height, - cropping=CROP_MODE, ) return [image_token_id] * num_image_tokens diff --git a/vllm/transformers_utils/processors/deepseek_ocr.py b/vllm/transformers_utils/processors/deepseek_ocr.py index 68a2b1aaaa..62972ad242 100644 --- a/vllm/transformers_utils/processors/deepseek_ocr.py +++ b/vllm/transformers_utils/processors/deepseek_ocr.py @@ -3,7 +3,7 @@ # adapted from https://github.com/deepseek-ai/DeepSeek-OCR/blob/main/DeepSeek-OCR-master/DeepSeek-OCR-vllm/process/image_process.py # and https://github.com/deepseek-ai/DeepSeek-OCR-2/blob/main/DeepSeek-OCR2-master/DeepSeek-OCR2-vllm/process/image_process.py import math -from typing import Literal +from typing import Literal, TypedDict import torch import torchvision.transforms as T @@ -11,18 +11,30 @@ from transformers import BatchFeature, LlamaTokenizerFast from transformers.processing_utils import ProcessorMixin -# TODO(Isotr0py): change modes for variants -# see: https://github.com/deepseek-ai/DeepSeek-OCR/blob/8cf003d38821fa1b19c73da3bd1b0dc262ea8136/DeepSeek-OCR-master/DeepSeek-OCR-vllm/config.py#L1-L6 -# Tiny: base_size = 512, image_size = 512, crop_mode = False -# Small: base_size = 640, image_size = 640, crop_mode = False -# Base: base_size = 1024, image_size = 1024, crop_mode = False -# Large: base_size = 1280, image_size = 1280, crop_mode = False -# Gundam: base_size = 1024, image_size = 640, crop_mode = True + +# Official resolution modes, verbatim from the model's config.py: +# https://github.com/deepseek-ai/DeepSeek-OCR/blob/8cf003d38821fa1b19c73da3bd1b0dc262ea8136/DeepSeek-OCR-master/DeepSeek-OCR-vllm/config.py#L1-L6 +# Selectable per request via the ``image_mode`` mm_processor_kwarg; the +# default stays the Gundam configuration below (byte-compatible with the +# previous hardcoded constants). +class ResolutionMode(TypedDict): + base_size: int + image_size: int + crop_mode: bool + + +RESOLUTION_MODES: dict[str, ResolutionMode] = { + "tiny": {"base_size": 512, "image_size": 512, "crop_mode": False}, + "small": {"base_size": 640, "image_size": 640, "crop_mode": False}, + "base": {"base_size": 1024, "image_size": 1024, "crop_mode": False}, + "large": {"base_size": 1280, "image_size": 1280, "crop_mode": False}, + "gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True}, +} + BASE_SIZE = 1024 IMAGE_SIZE = 640 CROP_MODE = True -# TODO(Isotr0py): Expose as mm_kwargs MIN_CROPS = 2 MAX_CROPS = 6 # max:9; If your GPU memory is small, it is recommended to set it to 6. @@ -117,6 +129,53 @@ def dynamic_preprocess( return processed_images, target_aspect_ratio +def count_image_tokens_for( + *, + image_width: int, + image_height: int, + base_size: int, + image_size: int, + cropping: bool, + min_crops: int = MIN_CROPS, + max_crops: int = MAX_CROPS, + strategy: Literal["v1", "v2"] = "v1", + patch_size: int = 16, + downsample_ratio: int = 4, +) -> int: + """Placeholder-token count mirroring ``tokenize_with_images``.""" + if image_width <= image_size and image_height <= image_size: + crop_ratio = (1, 1) + elif cropping: + crop_ratio = count_tiles( + image_width, + image_height, + min_num=min_crops, + max_num=max_crops, + image_size=image_size, + ) + else: + crop_ratio = (1, 1) + num_width_tiles, num_height_tiles = crop_ratio + + num_queries = math.ceil((image_size // patch_size) / downsample_ratio) + num_queries_base = math.ceil((base_size // patch_size) / downsample_ratio) + + num_tokens_base = ( + (num_queries_base * (num_queries_base + 1)) + if strategy == "v1" + else num_queries_base * num_queries_base + ) + total = num_tokens_base + 1 + if num_width_tiles > 1 or num_height_tiles > 1: + num_tokens_per_row = ( + num_queries * num_width_tiles + 1 + if strategy == "v1" + else num_queries * num_width_tiles + ) + total += num_tokens_per_row * (num_queries * num_height_tiles) + return total + + class ImageTransform: def __init__( self, @@ -160,22 +219,63 @@ def __init__( ignore_id: int = -100, image_size: int = IMAGE_SIZE, base_size: int = BASE_SIZE, + crop_mode: bool = CROP_MODE, + image_mode: str | None = None, + min_crops: int = MIN_CROPS, + max_crops: int = MAX_CROPS, strategy: Literal["v1", "v2"] = "v1", **kwargs, ): + if image_mode is not None: + # The named mode is authoritative over any individually passed + # base_size/image_size/crop_mode (the vLLM processing info always + # forwards the defaults, so a raise-on-conflict would never let a + # mode through). + if not isinstance(image_mode, str): + raise ValueError( + f"image_mode must be a string, got {type(image_mode).__name__}." + ) + try: + mode = RESOLUTION_MODES[image_mode] + except KeyError: + raise ValueError( + f"Unknown image_mode {image_mode!r}; valid modes: " + f"{sorted(RESOLUTION_MODES)}" + ) from None + base_size = mode["base_size"] + image_size = mode["image_size"] + crop_mode = mode["crop_mode"] + if not ( + type(min_crops) is int + and type(max_crops) is int + and 1 <= min_crops <= max_crops <= 9 + ): + raise ValueError( + f"Invalid crop bounds: min_crops={min_crops!r}, " + f"max_crops={max_crops!r} (need 1 <= min <= max <= 9)." + ) + if type(crop_mode) is not bool: + raise ValueError(f"crop_mode must be a bool, got {crop_mode!r}.") + if patch_size <= 0 or downsample_ratio <= 0: + raise ValueError("patch_size and downsample_ratio must be positive.") self.image_size = image_size self.base_size = base_size + self.crop_mode = crop_mode + self.image_mode = image_mode + self.min_crops = min_crops + self.max_crops = max_crops # image token calculation strategy for # Deepseek-OCR and Deepseek-OCR-2 self.strategy = strategy - assert strategy in ["v1", "v2"], "Only 'v1' and 'v2' strategies are supported." + if strategy not in ("v1", "v2"): + raise ValueError("Only 'v1' and 'v2' strategies are supported.") - self.patch_size = 16 + self.patch_size = patch_size self.image_mean = image_mean self.image_std = image_std self.normalize = normalize - self.downsample_ratio = 4 + self.downsample_ratio = downsample_ratio self.image_transform = ImageTransform( mean=image_mean, std=image_std, normalize=normalize @@ -203,6 +303,28 @@ def __init__( **kwargs, ) + def count_image_tokens(self, *, image_width: int, image_height: int) -> int: + """Number of placeholder tokens for one image. + + This mirrors ``tokenize_with_images`` exactly (same instance + parameters and same crop arithmetic) so the prompt-update counting + can never diverge from what processing + produces — a mismatch here is the "N multimodal tokens vs M + placeholders" crash class. + """ + return count_image_tokens_for( + image_width=image_width, + image_height=image_height, + base_size=self.base_size, + image_size=self.image_size, + cropping=self.crop_mode, + min_crops=self.min_crops, + max_crops=self.max_crops, + strategy=self.strategy, + patch_size=self.patch_size, + downsample_ratio=self.downsample_ratio, + ) + @property def bos_id(self): return self.tokenizer.bos_token_id @@ -230,7 +352,7 @@ def process_one( self, prompt: str, images: list[Image.Image], - crop_mode: bool = CROP_MODE, + crop_mode: bool | None = None, ): """ @@ -253,6 +375,8 @@ def process_one( ) sft_format = prompt + if crop_mode is None: + crop_mode = self.crop_mode ( input_ids, @@ -288,7 +412,7 @@ def __call__( *, prompt: str, images: list[Image.Image], - crop_mode: bool = CROP_MODE, + crop_mode: bool | None = None, **kwargs, ): prepare = self.process_one( @@ -332,7 +456,10 @@ def tokenize_with_images( crop_ratio = [1, 1] elif cropping: images_crop_raw, crop_ratio = dynamic_preprocess( - image, image_size=self.image_size + image, + min_num=self.min_crops, + max_num=self.max_crops, + image_size=self.image_size, ) else: crop_ratio = [1, 1] diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index e3954627ae..e6a6f15e65 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -502,6 +502,8 @@ def _sm70_profile_trace(message: str, *args: object) -> None: _flash_attn_prefill_paged_splitkv = None _sm70_splitd_d256_ops = None _sm70_splitd_d256_ops_checked = False +_sm70_d256_gqa_architecture_op = None +_sm70_d256_gqa_architecture_op_checked = False _sm70_fa2_cu_seqlens_cache: dict[ tuple[int, int, int, int], tuple[torch.Tensor, torch.Tensor] ] = {} @@ -515,6 +517,7 @@ def _sm70_profile_trace(message: str, *args: object) -> None: _warned_decode_strict_fallback = False _warned_prefill_gather_oom = False _warned_prefill_dense_splitkv3_oom = False +_warned_prefill_d256_gqa_architecture_oom = False _logged_prefill_flash = False _logged_prefill_prefix_flash = False _logged_prefill_prefix_contig_dense = False @@ -527,6 +530,7 @@ def _sm70_profile_trace(message: str, *args: object) -> None: _logged_prefill_smallq_grouped_verify_gate = False _logged_prefill_fa2_d256 = False _logged_prefill_dense_splitkv3 = False +_logged_prefill_d256_gqa_architecture = False _logged_prefill_triton_safe = False _logged_decode_flash = False _logged_decode_dense_reference = False @@ -1295,6 +1299,35 @@ def _get_sm70_splitd_d256_ops(): return _sm70_splitd_d256_ops +def _get_sm70_d256_gqa_architecture_op(): + """Load the optional SM70 GQA long-prefill architecture operator.""" + global _sm70_d256_gqa_architecture_op + global _sm70_d256_gqa_architecture_op_checked + if _sm70_d256_gqa_architecture_op_checked: + return _sm70_d256_gqa_architecture_op + + _sm70_d256_gqa_architecture_op_checked = True + try: + # Importing the interface loads the vendored FA2 torch library. + from vllm.vllm_flash_attn import flash_attn_interface # noqa: F401 + + _sm70_d256_gqa_architecture_op = getattr( + torch.ops._vllm_fa2_C, + "sm70_d256_gqa_architecture_fwd", + None, + ) + except (AttributeError, ImportError, RuntimeError) as exc: + _sm70_d256_gqa_architecture_op = None + if envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL: + logger.warning_once( + "SM70 D256 GQA architecture operator is unavailable " + "(%s: %s); using the exact dense prefill kernel.", + type(exc).__name__, + exc, + ) + return _sm70_d256_gqa_architecture_op + + def _uniform_cu_seqlens( tensor: torch.Tensor, *, @@ -1408,6 +1441,42 @@ def _should_use_prefill_dense_splitkv3( ) +def _should_use_prefill_d256_gqa_architecture( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + max_seqlen_q: int, + max_seqlen_k: int, + softmax_scale: float, + architecture_op: Callable[..., torch.Tensor] | None, +) -> bool: + """Gate the measured Q8000/KV40K..128K/Hq6/Hkv1/D256 family.""" + return ( + envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL + and architecture_op is not None + and query.shape == (1, 8000, 6, 256) + and key.ndim == 4 + and key.shape[0] == 1 + and key.shape[2:] == (1, 256) + and value.shape == key.shape + and max_seqlen_q == 8000 + and max_seqlen_k == key.shape[1] + and 40000 <= max_seqlen_k <= 128000 + and max_seqlen_k % 8000 == 0 + and query.dtype == torch.float16 + and key.dtype == query.dtype + and value.dtype == query.dtype + and query.device == key.device + and query.device == value.device + and query.is_contiguous() + and key.is_contiguous() + and value.is_contiguous() + and abs(softmax_scale - 0.0625) <= 1.0e-8 + and not _is_cuda_graph_capturing(query) + ) + + def _try_sm70_fa2_d256_prefill( query: torch.Tensor, key: torch.Tensor, @@ -1424,6 +1493,9 @@ def _try_sm70_fa2_d256_prefill( seqused_k: torch.Tensor | None = None, block_table: torch.Tensor | None = None, ) -> torch.Tensor | None: + global _logged_prefill_d256_gqa_architecture + global _warned_prefill_d256_gqa_architecture_oom + int32_max = torch.iinfo(torch.int32).max if not envs.VLLM_FLASH_V100_FA2_D256_PREFILL: return None @@ -1524,7 +1596,47 @@ def _try_sm70_fa2_d256_prefill( ) if splitd_eligible: splitd_out = out if out is not None else torch.empty_like(query) - if _should_use_prefill_dense_splitkv3( + architecture_op = ( + _get_sm70_d256_gqa_architecture_op() + if envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL + else None + ) + if _should_use_prefill_d256_gqa_architecture( + query, + key, + value, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + architecture_op=architecture_op, + ): + assert architecture_op is not None + try: + splitd_result = architecture_op( + query, + key, + value, + splitd_out, + softmax_scale, + True, + ) + except torch.OutOfMemoryError: + if not _warned_prefill_d256_gqa_architecture_oom: + logger.warning( + "Insufficient memory for the default-on " + "SM70 D256 GQA long-prefill architecture; " + "falling back to the exact dense kernel." + ) + _warned_prefill_d256_gqa_architecture_oom = True + if splitd_result is not None: + if not _logged_prefill_d256_gqa_architecture: + logger.info( + "FLASH_ATTN_V100 SM70 D256 GQA " + "8K-by-40K..128K architecture route active." + ) + _logged_prefill_d256_gqa_architecture = True + _record_route("prefill_dense_d256_gqa_arch_long") + if splitd_result is None and _should_use_prefill_dense_splitkv3( query, key, max_seqlen_q=max_seqlen_q, diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 5963790a77..2dd6bf92b8 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -137,6 +137,12 @@ def __init__( self.write_starts = new_buffer(self.num_rows, dtype=torch.int32) self.write_cu_lens = new_buffer(self.num_rows, dtype=torch.int32) + def get_cpu_view(self) -> torch.Tensor: + """Return the host view when this tensor is backed by mapped memory.""" + if not hasattr(self, "_uva_buf"): + raise RuntimeError("This StagedWriteTensor is not UVA-backed") + return self._uva_buf.cpu + def stage_write( self, index: int, start: int, x: Iterable[int] | Iterable[float] ) -> None: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 561e6aa411..47d1da107e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1420,6 +1420,9 @@ def sample_tokens( if self.speculator is not None: assert self.sampler is not None + requires_host_token_state = bool( + getattr(self.speculator, "requires_host_token_state", False) + ) # Let the target override the hidden state fed to the drafter # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The # target returns a persistent buffer sized at max_num_batched_tokens; @@ -1441,6 +1444,24 @@ def sample_tokens( self.sampler.sampling_states.temperature.gpu, self.sampler.sampling_states.seeds.gpu, mm_inputs=mm_inputs, + output_copy_event=( + async_output.copy_event if requires_host_token_state else None + ), + sampled_token_ids_cpu=( + async_output.sampled_token_ids + if requires_host_token_state + else None + ), + num_sampled_tokens_cpu=( + async_output.num_sampled_tokens_np + if requires_host_token_state + else None + ), + all_token_ids_cpu=( + self.req_states.all_token_ids.get_cpu_view().numpy() + if requires_host_token_state + else None + ), ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 8c6dc8506e..9d08b77ee5 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -325,6 +325,20 @@ def _build_draft_attn_metadata( dcp_local_seq_lens=dcp_local_seq_lens, ) + def _prepare_ngram_assist( + self, + input_batch: InputBatch, + output_copy_event: torch.cuda.Event | None, + sampled_token_ids_cpu: np.ndarray | None, + num_sampled_tokens_cpu: np.ndarray | None, + all_token_ids_cpu: np.ndarray | None, + ) -> bool: + """Prepare an optional draftless proposal; return whether it is complete.""" + return False + + def _apply_ngram_assist(self, num_reqs: int) -> None: + """Override model proposals for ngram-hit rows, if configured.""" + @torch.inference_mode() def propose( self, @@ -352,6 +366,10 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, + output_copy_event: torch.cuda.Event | None = None, + sampled_token_ids_cpu: np.ndarray | None = None, + num_sampled_tokens_cpu: np.ndarray | None = None, + all_token_ids_cpu: np.ndarray | None = None, ) -> torch.Tensor: num_reqs = input_batch.num_reqs num_target_tokens = input_batch.num_tokens @@ -467,6 +485,16 @@ def propose( self._context_only_prefill_logged = True return self.draft_tokens[:num_reqs] + if self._prepare_ngram_assist( + input_batch, + output_copy_event, + sampled_token_ids_cpu, + num_sampled_tokens_cpu, + all_token_ids_cpu, + ): + self._apply_ngram_assist(num_reqs) + return self.draft_tokens[:num_reqs] + with record_function_or_nullcontext("dflash: query and selector"): # Every DFlash step has exactly num_query_per_req tokens, so we can # use FULL CUDA graphs. @@ -511,6 +539,8 @@ def propose( cudagraph_runtime_mode=batch_desc.cg_mode, ) + self._apply_ngram_assist(num_reqs) + return self.draft_tokens[:num_reqs] diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/ngram_assist.py b/vllm/v1/worker/gpu/spec_decode/dflash2/ngram_assist.py new file mode 100644 index 0000000000..fe89c55512 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/ngram_assist.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host prompt-ngram lookup for the MRV2 DFlash2 assistant.""" + +from __future__ import annotations + +import time + +import numpy as np +from numba import njit + + +@njit(inline="always") +def _split_token( + prefix: np.ndarray, + prefix_len: int, + suffix: np.ndarray, + index: int, +) -> int: + if index < prefix_len: + return int(prefix[index]) + return int(suffix[index - prefix_len]) + + +@njit(cache=True) +def find_split_ngram_proposal( + prefix: np.ndarray, + prefix_len: int, + suffix: np.ndarray, + suffix_len: int, + min_ngram: int, + max_ngram: int, + max_model_len: int, + num_draft_tokens: int, +) -> np.ndarray: + """Run the prompt-lookup KMP over ``prefix + suffix`` without copying it.""" + total_tokens = prefix_len + suffix_len + if total_tokens < min_ngram: + return np.empty(0, dtype=np.int32) + + num_draft_tokens = min(num_draft_tokens, max_model_len - total_tokens) + if num_draft_tokens <= 0: + return np.empty(0, dtype=np.int32) + + # This is the same reverse-KMP policy as the standalone vLLM ngram + # proposer: prefer the longest suffix match and the earliest occurrence. + lps = np.zeros(max_ngram, dtype=np.int32) + longest_ngram = 0 + position = 0 + prev_lps = 0 + i = 1 + while i < total_tokens: + prefix_token = _split_token( + prefix, + prefix_len, + suffix, + total_tokens - 1 - prev_lps, + ) + current_token = _split_token( + prefix, + prefix_len, + suffix, + total_tokens - 1 - i, + ) + if prefix_token == current_token: + prev_lps += 1 + if prev_lps >= longest_ngram: + longest_ngram = prev_lps + position = i + if i < max_ngram: + lps[i] = prev_lps + if prev_lps == max_ngram: + prev_lps = lps[max_ngram - 1] + i += 1 + elif prev_lps != 0: + prev_lps = lps[prev_lps - 1] + else: + i += 1 + + if longest_ngram < min_ngram: + return np.empty(0, dtype=np.int32) + + start = total_tokens - 1 - position + longest_ngram + proposal_len = min(num_draft_tokens, total_tokens - start) + proposal = np.empty(proposal_len, dtype=np.int32) + for j in range(proposal_len): + proposal[j] = _split_token(prefix, prefix_len, suffix, start + j) + return proposal + + +class DFlash2NgramAssist: + """Batch adapter and counters around the split prompt lookup.""" + + def __init__( + self, + min_ngram: int, + max_ngram: int, + num_draft_tokens: int, + max_model_len: int, + ) -> None: + self.min_ngram = min_ngram + self.max_ngram = max_ngram + self.num_draft_tokens = num_draft_tokens + self.max_model_len = max_model_len + self.num_eligible = 0 + self.num_full_hits = 0 + self.lookup_seconds = 0.0 + + # Compile the small split-input kernel at startup, not on the first + # user-visible decode token. + find_split_ngram_proposal( + np.zeros(1, dtype=np.int32), + 1, + np.zeros(1, dtype=np.int64), + 0, + min_ngram, + max_ngram, + max_model_len, + num_draft_tokens, + ) + + def propose( + self, + token_ids_cpu: np.ndarray, + req_state_indices: np.ndarray, + prior_lengths: np.ndarray, + sampled_token_ids: np.ndarray, + num_sampled_tokens: np.ndarray, + eligible: np.ndarray, + output_tokens: np.ndarray, + output_lengths: np.ndarray, + ) -> int: + """Fill fixed-width outputs and return the number of full-width hits.""" + num_reqs = len(req_state_indices) + output_tokens[:num_reqs].fill(0) + output_lengths[:num_reqs].fill(0) + full_hits = 0 + started = time.perf_counter() + for batch_idx in range(num_reqs): + if not eligible[batch_idx]: + continue + self.num_eligible += 1 + req_state_idx = int(req_state_indices[batch_idx]) + proposal = find_split_ngram_proposal( + token_ids_cpu[req_state_idx], + int(prior_lengths[batch_idx]), + sampled_token_ids[batch_idx], + int(num_sampled_tokens[batch_idx]), + self.min_ngram, + self.max_ngram, + self.max_model_len, + self.num_draft_tokens, + ) + proposal_len = len(proposal) + output_lengths[batch_idx] = proposal_len + if proposal_len: + output_tokens[batch_idx, :proposal_len] = proposal + if proposal_len == self.num_draft_tokens: + full_hits += 1 + self.num_full_hits += full_hits + self.lookup_seconds += time.perf_counter() - started + return full_hits diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py index 6c04c86961..01b42e3968 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py @@ -4,6 +4,7 @@ from __future__ import annotations +import os from typing import TYPE_CHECKING, Any import numpy as np @@ -26,6 +27,125 @@ logger = init_logger(__name__) _TARGET_TOP_K = 20 +_SELECTOR_ALIGNMENT_DUMP_COUNT = 0 +_SELECTOR_ALIGNMENT_STEP = 0 + + +def _parse_alignment_steps(raw_steps: str | None) -> set[int] | None: + if not raw_steps: + return None + steps: set[int] = set() + try: + for item in raw_steps.split(","): + item = item.strip() + if not item: + continue + if "-" in item: + start_text, end_text = item.split("-", 1) + start = int(start_text) + end = int(end_text) + if start < 0 or end < start: + return set() + steps.update(range(start, end + 1)) + else: + step = int(item) + if step < 0: + return set() + steps.add(step) + except ValueError: + return set() + return steps + + +def _safe_dump_tag(raw_tag: str) -> str: + return "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in raw_tag) + + +def _diagnostic_rank() -> int: + return int(os.getenv("RANK", os.getenv("LOCAL_RANK", "0"))) + + +def _maybe_dump_selector_alignment( + *, + speculator: DFlash2Speculator, + rejection_sampler: RejectionSampler, + input_batch: InputBatch, + target_topk_ids: torch.Tensor, + target_topk_logits: torch.Tensor, + draft_topk_ids: torch.Tensor, + draft_topk_logits: torch.Tensor, + draft_sampled: torch.Tensor, + pos: torch.Tensor, + sampled: torch.Tensor, + num_sampled: torch.Tensor, +) -> None: + """Dump one exact B1 selector/target alignment record when requested.""" + if not envs.VLLM_SPEC_DUMP_ALIGNMENT: + return + if _diagnostic_rank() != 0: + return + + global _SELECTOR_ALIGNMENT_DUMP_COUNT, _SELECTOR_ALIGNMENT_STEP + _SELECTOR_ALIGNMENT_STEP += 1 + if _SELECTOR_ALIGNMENT_DUMP_COUNT >= envs.VLLM_SPEC_DUMP_ALIGNMENT_LIMIT: + return + selected_steps = _parse_alignment_steps(envs.VLLM_SPEC_DUMP_ALIGNMENT_STEPS) + if selected_steps is not None and _SELECTOR_ALIGNMENT_STEP not in selected_steps: + return + + shadow = speculator.get_selector_alignment_shadow() + if shadow is None: + return + shadow_ids, unary_logits, lattice_scores = shadow + req_state = int(input_batch.idx_mapping_np[0]) + packed_row = 0 + sampling_states = rejection_sampler.sampler.sampling_states + + with torch.no_grad(): + draft_sampled_cpu = draft_sampled.detach().cpu() + if bool(torch.all(draft_sampled_cpu == 0).item()): + return + payload = { + "format": "dflash2_selector_alignment_v1", + "rank": _diagnostic_rank(), + "step": _SELECTOR_ALIGNMENT_STEP, + "request_state": req_state, + "selector_top_k": speculator.selector_top_k, + "num_speculative_steps": speculator.num_speculative_steps, + "target_topk_ids": target_topk_ids.detach().cpu(), + "target_topk_logits": target_topk_logits.detach().float().cpu(), + "draft_candidate_ids": draft_topk_ids[req_state].detach().cpu(), + "draft_realized_logits": ( + draft_topk_logits[req_state].detach().float().cpu() + ), + "selector_candidate_ids": shadow_ids[packed_row].detach().cpu(), + "selector_unary_logits": (unary_logits[packed_row].detach().float().cpu()), + "selector_lattice_scores": ( + lattice_scores[packed_row].detach().float().cpu() + ), + "draft_sampled": draft_sampled_cpu, + "positions": pos.detach().cpu(), + "cu_num_logits": input_batch.cu_num_logits.detach().cpu(), + "idx_mapping": input_batch.idx_mapping.detach().cpu(), + "temperature": float(sampling_states.temperature.np[req_state]), + "top_p": float(sampling_states.top_p.np[req_state]), + "top_k": int(sampling_states.top_k.np[req_state]), + "sampled_token_ids": sampled.detach().cpu(), + "num_sampled": num_sampled.detach().cpu(), + } + _SELECTOR_ALIGNMENT_DUMP_COUNT += 1 + dump_dir = os.getenv("VLLM_SPEC_DUMP_ALIGNMENT_DIR", "/tmp") + os.makedirs(dump_dir, exist_ok=True) + tag = _safe_dump_tag(os.getenv("VLLM_SPEC_DUMP_ALIGNMENT_TAG", "")) + tag_part = f"{tag}_" if tag else "" + dump_path = os.path.join( + dump_dir, + f"spec_alignment_dflash2_selector_{tag_part}pid{os.getpid()}_" + f"step{_SELECTOR_ALIGNMENT_STEP:06d}_" + f"{_SELECTOR_ALIGNMENT_DUMP_COUNT}.pt", + ) + torch.save(payload, dump_path) + logger.warning("Dumped DFlash2 selector alignment diagnostics to %s", dump_path) def _supports_sparse_sampling_contract( @@ -118,6 +238,20 @@ def try_dflash2_sparse_target_rejection( rejection_sampler.num_speculative_steps, use_fp64=rejection_sampler.sampler.use_fp64_gumbel, ) + if envs.VLLM_SPEC_DUMP_ALIGNMENT: + _maybe_dump_selector_alignment( + speculator=speculator, + rejection_sampler=rejection_sampler, + input_batch=input_batch, + target_topk_ids=target_topk_ids, + target_topk_logits=target_topk_logits, + draft_topk_ids=draft_topk_ids, + draft_topk_logits=draft_topk_logits, + draft_sampled=draft_sampled, + pos=pos, + sampled=sampled, + num_sampled=num_sampled, + ) logger.info_once("Using SM70 DFlash2 compact target top-k rejection sampling.") return SamplerOutput( sampled_token_ids=sampled, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index cac1d3e9ba..64f987238a 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -3,14 +3,19 @@ from typing import Any +import numpy as np import torch from vllm import envs from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger from vllm.triton_utils import tl, triton from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator +from vllm.v1.worker.gpu.spec_decode.dflash2.ngram_assist import DFlash2NgramAssist + +logger = init_logger(__name__) def _requires_sm70_tail(device: torch.device, num_steps: int) -> bool: @@ -198,6 +203,64 @@ def _cache_draft_logits_kernel( tl.store(cached_score_ptr + cache_base + offsets, scores, mask=mask) +@triton.jit +def _apply_ngram_draft_kernel( + ngram_tokens_ptr, + ngram_lengths_ptr, + sample_req_state_ptr, + draft_tokens_ptr, + draft_tokens_stride, + cached_candidate_ptr, + cached_score_ptr, + cache_stride_0, + cache_stride_1, + draft_logits_ptr, + draft_logits_stride_0, + draft_logits_stride_1, + num_steps: tl.constexpr, + top_k: tl.constexpr, + BLOCK_K: tl.constexpr, + CACHE_DRAFT_LOGITS: tl.constexpr, + CACHE_SCORES: tl.constexpr, +): + flat = tl.program_id(0) + batch_idx = flat // num_steps + step = flat % num_steps + req_state = tl.load(sample_req_state_ptr + flat) + valid = (req_state >= 0) & (tl.load(ngram_lengths_ptr + batch_idx) == num_steps) + token = tl.load(ngram_tokens_ptr + flat, mask=valid, other=0).to(tl.int64) + tl.store( + draft_tokens_ptr + batch_idx * draft_tokens_stride + step, + token, + mask=valid, + ) + + if CACHE_DRAFT_LOGITS: + offsets = tl.arange(0, BLOCK_K) + topk_mask = valid & (offsets < top_k) + cache_base = ( + cached_candidate_ptr + req_state * cache_stride_0 + step * cache_stride_1 + ) + old_ids = tl.load(cache_base + offsets, mask=topk_mask, other=0) + logits_base = ( + draft_logits_ptr + + req_state * draft_logits_stride_0 + + step * draft_logits_stride_1 + ) + tl.store(logits_base + old_ids, -float("inf"), mask=topk_mask) + + is_proposal = offsets == 0 + new_ids = tl.where(is_proposal, token, 0) + new_scores = tl.where(is_proposal, 0.0, -float("inf")) + tl.store(cache_base + offsets, new_ids, mask=topk_mask) + if CACHE_SCORES: + score_base = ( + cached_score_ptr + req_state * cache_stride_0 + step * cache_stride_1 + ) + tl.store(score_base + offsets, new_scores, mask=topk_mask) + tl.store(logits_base + token, 0.0, mask=valid) + + class DFlash2Speculator(DFlashSpeculator): _speculator_name = "DFlash2" @@ -228,12 +291,89 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self._selector_path_state = torch.empty( self.max_num_reqs, dtype=torch.int32, device=device ) + self._alignment_candidate_ids: torch.Tensor | None = None + self._alignment_unary_logits: torch.Tensor | None = None + self._alignment_lattice_scores: torch.Tensor | None = None + if ( + envs.VLLM_SPEC_DUMP_ALIGNMENT + and envs.VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION + ): + packed_shape = ( + self.max_num_reqs, + self.num_speculative_steps, + self.selector_top_k, + ) + self._alignment_candidate_ids = torch.empty( + packed_shape, dtype=torch.int64, device=device + ) + self._alignment_unary_logits = torch.empty( + packed_shape, dtype=torch.float32, device=device + ) + self._alignment_lattice_scores = torch.empty( + (*packed_shape, self.selector_top_k), + dtype=torch.float32, + device=device, + ) self._use_sm70_tail = _requires_sm70_tail(device, self.num_speculative_steps) if self.draft_logits is not None: # The cache kernel writes only K columns; all other vocabulary # columns must remain impossible. self.draft_logits.fill_(-float("inf")) + self._ngram_assist: DFlash2NgramAssist | None = None + self._ngram_num_hits = 0 + self._ngram_rounds = 0 + self._ngram_skipped_rounds = 0 + speculative_config = getattr(self, "speculative_config", None) + if speculative_config is not None and getattr( + speculative_config, "ngram_assist", False + ): + min_ngram = speculative_config.prompt_lookup_min + max_ngram = speculative_config.prompt_lookup_max + assert min_ngram is not None and max_ngram is not None + self._ngram_assist = DFlash2NgramAssist( + min_ngram=min_ngram, + max_ngram=max_ngram, + num_draft_tokens=self.num_speculative_steps, + max_model_len=self.max_model_len, + ) + self._ngram_tokens_cpu_tensor = torch.zeros( + self.max_num_reqs, + self.num_speculative_steps, + dtype=torch.int64, + device="cpu", + pin_memory=True, + ) + self._ngram_lengths_cpu_tensor = torch.zeros( + self.max_num_reqs, + dtype=torch.int32, + device="cpu", + pin_memory=True, + ) + self._ngram_tokens_cpu = self._ngram_tokens_cpu_tensor.numpy() + self._ngram_lengths_cpu = self._ngram_lengths_cpu_tensor.numpy() + self._ngram_tokens = torch.zeros( + self.max_num_reqs, + self.num_speculative_steps, + dtype=torch.int64, + device=device, + ) + self._ngram_lengths = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + logger.info( + "Enabled DFlash2 ngram assist with prompt lookup [%d, %d] " + "and draft width %d.", + min_ngram, + max_ngram, + self.num_speculative_steps, + ) + + @property + def requires_host_token_state(self) -> bool: + """Whether the runner must expose async samples and request history.""" + return self._ngram_assist is not None + def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]: # The selector walk and rejection sampler must consume identical scores. # BF16 rounding measurably changes candidate order, so keep this FP32. @@ -322,6 +462,120 @@ def get_sparse_draft_logits( return None return self._cached_candidate_ids, self._cached_candidate_scores + def get_selector_alignment_shadow( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + """Return packed selector tensors for an explicitly enabled diagnostic.""" + if self._alignment_candidate_ids is None: + return None + assert self._alignment_unary_logits is not None + assert self._alignment_lattice_scores is not None + return ( + self._alignment_candidate_ids, + self._alignment_unary_logits, + self._alignment_lattice_scores, + ) + + def _prepare_ngram_assist( + self, + input_batch, + output_copy_event: torch.cuda.Event | None, + sampled_token_ids_cpu: np.ndarray | None, + num_sampled_tokens_cpu: np.ndarray | None, + all_token_ids_cpu: np.ndarray | None, + ) -> bool: + assist = self._ngram_assist + self._ngram_num_hits = 0 + if ( + assist is None + or input_batch.has_structured_output_reqs + or output_copy_event is None + or sampled_token_ids_cpu is None + or num_sampled_tokens_cpu is None + or all_token_ids_cpu is None + ): + return False + + # The copy stream only depends on target sampling. The main stream can + # materialize DFlash context K/V while the host waits here, so lookup + # does not serialize the context projection. + output_copy_event.synchronize() + num_reqs = input_batch.num_reqs + num_draft_tokens = input_batch.num_draft_tokens_per_req + if num_draft_tokens is None: + num_draft_tokens = np.zeros(num_reqs, dtype=np.int32) + prior_lengths = ( + input_batch.seq_lens_cpu_upper_bound[:num_reqs].numpy() - num_draft_tokens + ) + eligible = (~input_batch.is_prefilling_np[:num_reqs]) & ( + num_sampled_tokens_cpu[:num_reqs] > 0 + ) + full_hits = assist.propose( + all_token_ids_cpu, + input_batch.idx_mapping_np, + prior_lengths, + sampled_token_ids_cpu, + num_sampled_tokens_cpu, + eligible, + self._ngram_tokens_cpu, + self._ngram_lengths_cpu, + ) + self._ngram_rounds += 1 + skip_query = num_reqs > 0 and full_hits == num_reqs + self._ngram_num_hits = full_hits if skip_query else 0 + if skip_query: + self._ngram_tokens[:num_reqs].copy_( + self._ngram_tokens_cpu_tensor[:num_reqs], non_blocking=True + ) + self._ngram_lengths[:num_reqs].copy_( + self._ngram_lengths_cpu_tensor[:num_reqs], non_blocking=True + ) + + self._ngram_skipped_rounds += int(skip_query) + if ( + envs.VLLM_DFLASH_PROFILE + and self._ngram_rounds % envs.VLLM_DFLASH_PROFILE_LOG_INTERVAL == 0 + ): + eligible_count = max(assist.num_eligible, 1) + logger.info( + "DFLASH2_NGRAM_PROFILE rounds=%d eligible=%d full_hits=%d " + "hit_rate=%.4f skipped_query_rounds=%d lookup_avg_ms=%.4f", + self._ngram_rounds, + assist.num_eligible, + assist.num_full_hits, + assist.num_full_hits / eligible_count, + self._ngram_skipped_rounds, + assist.lookup_seconds * 1000.0 / self._ngram_rounds, + ) + return skip_query + + def _apply_ngram_assist(self, num_reqs: int) -> None: + if self._ngram_assist is None or self._ngram_num_hits == 0: + return + draft_logits = self.draft_logits + cached_scores = self._cached_candidate_scores + block_k = triton.next_power_of_2(self.selector_top_k) + _apply_ngram_draft_kernel[(num_reqs * self.num_speculative_steps,)]( + self._ngram_tokens, + self._ngram_lengths, + self.sample_idx_mapping, + self.draft_tokens, + self.draft_tokens.stride(0), + self._cached_candidate_ids, + self._selector_scores if cached_scores is None else cached_scores, + self._cached_candidate_ids.stride(0), + self._cached_candidate_ids.stride(1), + self._selector_scores if draft_logits is None else draft_logits, + 0 if draft_logits is None else draft_logits.stride(0), + 0 if draft_logits is None else draft_logits.stride(1), + num_steps=self.num_speculative_steps, + top_k=self.selector_top_k, + BLOCK_K=block_k, + CACHE_DRAFT_LOGITS=draft_logits is not None, + CACHE_SCORES=cached_scores is not None, + num_warps=1, + ) + def _generate_draft( self, num_reqs: int, @@ -356,6 +610,12 @@ def _generate_draft( hidden_states, anchor_token_ids, ) + if self._alignment_candidate_ids is not None: + assert self._alignment_unary_logits is not None + assert self._alignment_lattice_scores is not None + self._alignment_candidate_ids[:num_reqs].copy_(candidate_ids) + self._alignment_unary_logits[:num_reqs].copy_(unary_logits) + self._alignment_lattice_scores[:num_reqs].copy_(scores) self._sample_path(candidate_ids, scores, num_reqs) if self.draft_logits is not None: self._cache_draft_logits(candidate_ids, num_sample) diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 6ae3fe793b..8376a51558 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Any +import numpy as np import torch import torch.nn as nn @@ -481,6 +482,10 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, + output_copy_event: torch.cuda.Event | None = None, + sampled_token_ids_cpu: np.ndarray | None = None, + num_sampled_tokens_cpu: np.ndarray | None = None, + all_token_ids_cpu: np.ndarray | None = None, ) -> torch.Tensor: num_tokens = input_batch.num_tokens_after_padding num_reqs = input_batch.num_reqs diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 006f0179f0..0775cfdca6 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -64,6 +64,10 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, + output_copy_event: torch.cuda.Event | None = None, + sampled_token_ids_cpu: np.ndarray | None = None, + num_sampled_tokens_cpu: np.ndarray | None = None, + all_token_ids_cpu: np.ndarray | None = None, ) -> torch.Tensor: pass diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 308931f172..c150cf791e 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6465,9 +6465,14 @@ def sync_and_gather_intermediate_tensors( local_len = num_tokens // tp v = get_tp_group().all_gather(v[:local_len], dim=0) - self.intermediate_tensors[k][:num_tokens].copy_( - v[:num_tokens], non_blocking=True - ) + destination = self.intermediate_tensors[k][:num_tokens] + source = v[:num_tokens] + if ( + destination.data_ptr() != source.data_ptr() + or destination.shape != source.shape + or destination.stride() != source.stride() + ): + destination.copy_(source, non_blocking=True) return IntermediateTensors( {k: v[:num_tokens] for k, v in self.intermediate_tensors.items()} diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 6ca3c687f7..c131a019b1 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -220,9 +220,9 @@ def _scoped_allocator_max_split(self, max_split_size_mb: int): yield return - set_allocator_settings = getattr(torch._C, - "_accelerator_setAllocatorSettings", - None) + set_allocator_settings = getattr( + torch._C, "_accelerator_setAllocatorSettings", None + ) if set_allocator_settings is None: yield return @@ -788,6 +788,92 @@ def sample_tokens( ) -> ModelRunnerOutput | AsyncModelRunnerOutput: return self.model_runner.sample_tokens(grammar_output) + def _use_sm70_static_pp_hidden_transfer(self, num_tokens: int) -> bool: + """Whether this step has the exact metadata-free PP tensor contract.""" + if not envs.VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER or num_tokens != 1: + return False + cached = getattr(self, "_sm70_static_pp_hidden_contract", None) + if cached is not None: + return bool(cached) + + config = self.vllm_config + parallel_config = config.parallel_config + compilation_config = config.compilation_config + model_config = config.model_config + enabled = not ( + parallel_config.pipeline_parallel_size != 2 + or parallel_config.tensor_parallel_size != 4 + or parallel_config.enable_dbo + or parallel_config.ubatch_size > 1 + or config.scheduler_config.max_num_seqs != 1 + or getattr(config, "speculative_config", None) is not None + or compilation_config.cudagraph_mode == CUDAGraphMode.NONE + or compilation_config.pass_config.enable_sp + or model_config.dtype != torch.float16 + ) + + schema_matches = False + if enabled: + try: + schema = self.model_runner.model.make_empty_intermediate_tensors( + batch_size=1, + dtype=model_config.dtype, + device=torch.device("meta"), + ).tensors + hidden_states = schema.get("hidden_states") + schema_matches = bool( + tuple(schema) == ("hidden_states",) + and isinstance(hidden_states, torch.Tensor) + and hidden_states.shape == (1, 4, 4096) + and hidden_states.dtype == torch.float16 + and hidden_states.is_contiguous() + ) + except (AttributeError, RuntimeError, TypeError): + # This is an optional fast path; unsupported model-provided + # schema construction retains the metadata protocol. + schema_matches = False + enabled = bool( + enabled + and schema_matches + and current_platform.is_cuda() + and current_platform.is_device_capability((7, 0)) + ) + self._sm70_static_pp_hidden_contract = enabled + return enabled + + @staticmethod + def _is_static_pp_hidden_tensor_dict( + tensor_dict: dict[str, torch.Tensor], num_tokens: int + ) -> bool: + if tuple(tensor_dict) != ("hidden_states",): + return False + hidden_states = tensor_dict["hidden_states"] + return bool( + hidden_states.shape == (num_tokens, 4, 4096) + and hidden_states.dtype == torch.float16 + and hidden_states.is_cuda + and hidden_states.is_contiguous() + ) + + def _static_pp_hidden_recv_buffers( + self, num_tokens: int + ) -> dict[str, torch.Tensor] | None: + if not self._use_sm70_static_pp_hidden_transfer(num_tokens): + return None + buffers = self.model_runner.intermediate_tensors + if buffers is None or tuple(buffers.tensors) != ("hidden_states",): + raise RuntimeError( + "static PP hidden transfer requires one persistent " + "hidden_states input buffer" + ) + hidden_states = buffers.tensors["hidden_states"][:num_tokens] + tensor_dict = {"hidden_states": hidden_states} + if not self._is_static_pp_hidden_tensor_dict(tensor_dict, num_tokens): + raise RuntimeError( + "static PP hidden transfer input buffer violates its fixed schema" + ) + return tensor_dict + @torch.inference_mode() def execute_model( self, scheduler_output: "SchedulerOutput" @@ -835,12 +921,20 @@ def execute_model( } if forward_pass and not get_pp_group().is_first_rank: - tensor_dict, comm_handles, comm_postprocess = ( - get_pp_group().irecv_tensor_dict( - all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, + tensor_dict = self._static_pp_hidden_recv_buffers(num_scheduled_tokens) + if tensor_dict is not None: + comm_handles = get_pp_group().irecv_tensor_dict_static(tensor_dict) + comm_postprocess: list[Callable[[], None]] = [] + logger.info_once( + "Default SM70 metadata-free PP hidden transfer enabled." + ) + else: + tensor_dict, comm_handles, comm_postprocess = ( + get_pp_group().irecv_tensor_dict( + all_gather_group=get_tp_group(), + all_gather_tensors=all_gather_tensors, + ) ) - ) assert tensor_dict is not None intermediate_tensors = AsyncIntermediateTensors( tensor_dict, @@ -870,12 +964,23 @@ def execute_model( and not get_pp_group().is_last_rank ) - # launch non-blocking send of intermediate tensors - self._pp_send_work = get_pp_group().isend_tensor_dict( - output.tensors, - all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, - ) + # The exact SM70 B1 path sends the replicated tensor directly into + # the next stage's persistent graph-input buffer without CPU metadata. + if self._use_sm70_static_pp_hidden_transfer(num_scheduled_tokens): + if not self._is_static_pp_hidden_tensor_dict( + output.tensors, num_scheduled_tokens + ): + raise RuntimeError( + "static PP hidden transfer output violates its fixed schema" + ) + self._pp_send_work = get_pp_group().isend_tensor_dict_static(output.tensors) + logger.info_once("Default SM70 metadata-free PP hidden transfer enabled.") + else: + self._pp_send_work = get_pp_group().isend_tensor_dict( + output.tensors, + all_gather_group=get_tp_group(), + all_gather_tensors=all_gather_tensors, + ) return None