diff --git a/benchmarks/glm_prefill_checkpoints.py b/benchmarks/glm_prefill_checkpoints.py new file mode 100644 index 000000000000..681894079c76 --- /dev/null +++ b/benchmarks/glm_prefill_checkpoints.py @@ -0,0 +1,603 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Run serial GLM prefill/cache smoke checks and fixed cold TTFT samples. + +Requires /health, /metrics, /tokenize and OpenAI-compatible chat endpoints. +This client controls no model processes or hardware. Readiness and exclusive +access are operator attestations; feature activation is not verified here. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import statistics +import sys +import tempfile +import time +import urllib.error +import urllib.parse +import urllib.request +import uuid +from datetime import datetime, timezone +from pathlib import Path + +import regex as re + +SIZES = (8192, 16384, 32768) +REPEATS = 3 +SEMANTIC_MAX_TOKENS = 384 +MEASURED_PROTOCOL_SOURCE_SHA256 = ( + "6e4bf5ff62379bb2db4dc439d27a85eda0837959e983ef5f92ab3605898e69ac" +) + + +def normalize_base_url(value): + """Normalize the API prefix without accepting credentials in journaled URLs.""" + base = value.rstrip("/") + if base.endswith("/v1"): + base = base[:-3] + parsed = urllib.parse.urlsplit(base) + if ( + parsed.scheme not in ("http", "https") + or not parsed.netloc + or parsed.username is not None + or parsed.password is not None + or parsed.query + or parsed.fragment + ): + raise ValueError( + "Use an HTTP(S) base URL without credentials, query or fragment" + ) + return base + + +def load_conditions(path): + """Read optional operator metadata without promoting it to verification.""" + if path is None: + return None, None + raw = path.read_bytes() + conditions = json.loads(raw) + if ( + not isinstance(conditions, dict) + or conditions.get("schema") != "glm-prefill-reproduction-conditions/v1" + or set(conditions) - {"schema", "sources", "settings"} + or not isinstance(conditions.get("sources", {}), dict) + or not isinstance(conditions.get("settings", {}), dict) + ): + raise ValueError( + "Conditions require schema glm-prefill-reproduction-conditions/v1 " + "and optional sources/settings objects" + ) + if any( + not isinstance(value, str) or not value + for value in conditions.get("sources", {}).values() + ): + raise ValueError("Condition source identities must be nonempty strings") + return conditions, hashlib.sha256(raw).hexdigest() + + +def utc_now(): + return datetime.now(timezone.utc).isoformat() + + +def cached_tokens(usage): + details = usage.get("prompt_tokens_details") + count = details.get("cached_tokens") if isinstance(details, dict) else None + if type(count) is not int or count < 0: + raise ValueError("Response must provide explicit nonnegative cached_tokens") + return count + + +def validate_usage(usage, expected_tokens, cache): + if ( + type(usage.get("prompt_tokens")) is not int + or usage["prompt_tokens"] != expected_tokens + ): + raise ValueError( + f"Prompt usage differs: expected {expected_tokens}, " + f"got {usage.get('prompt_tokens')}" + ) + count = cached_tokens(usage) + if cache == "cold" and count != 0: + raise ValueError(f"Cold request reused {count} tokens") + if cache == "reuse" and count <= 0: + raise ValueError( + "Expected cache reuse was absent; recomputation is not a reuse pass" + ) + return count + + +def request_gauges(metrics): + result = {} + for metric in ("num_requests_running", "num_requests_waiting"): + pattern = re.compile( + r"^vllm:" + metric + r"(?:\{[^\n]*\})?\s+(\S+)(?:\s+\S+)?$" + ) + values = [ + float(match.group(1)) + for line in metrics.splitlines() + if (match := pattern.match(line)) + ] + if not values or any(not math.isfinite(value) or value < 0 for value in values): + raise ValueError(f"Missing or invalid vllm:{metric} metrics") + result[metric] = sum(values) + return result + + +def first_token_delta(chunk): + return any( + any( + choice.get("delta", {}).get(key) + for key in ("content", "reasoning", "reasoning_content") + ) + for choice in chunk.get("choices", []) + ) + + +def summarize(cases): + result = [] + for size in SIZES: + selected = [ + row + for row in cases + if row["phase"] == "measured" + and row.get("accepted") + and row["tokens"] == size + ] + if not selected: + continue + times = [row["ttft_seconds"] for row in selected] + if any(not math.isfinite(value) or value <= 0 for value in times): + raise ValueError("Invalid TTFT cannot enter a summary") + median = statistics.median(times) + result.append( + { + "tokens": size, + "samples": len(times), + "median_ttft_seconds": median, + "min_ttft_seconds": min(times), + "max_ttft_seconds": max(times), + "tokens_per_second": size / median, + } + ) + return result + + +class Journal: + def __init__(self, path, record): + self.path, self.record = path, record + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("x", encoding="utf-8", newline="\n") as stream: + json.dump(record, stream, indent=2) + + def save(self): + self.record["updated_at_utc"] = utc_now() + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + newline="\n", + dir=self.path.parent, + prefix=self.path.name + ".", + suffix=".tmp", + delete=False, + ) as stream: + json.dump(self.record, stream, indent=2) + stream.write("\n") + temporary = Path(stream.name) + temporary.replace(self.path) + + +class PrefillChecks: + def __init__(self, args, journal): + self.args, self.journal = args, journal + self.base = normalize_base_url(args.base_url) + self.headers = {"Content-Type": "application/json"} + if args.api_key_env: + key = os.environ.get(args.api_key_env) + if not key: + raise ValueError("Selected API-key environment variable is empty") + self.headers["Authorization"] = "Bearer " + key + self.model = args.model + + def control(self, path): + request = urllib.request.Request(self.base + path, headers=self.headers) + with urllib.request.urlopen( + request, timeout=self.args.request_timeout + ) as response: + return response.read().decode("utf-8") + + def idle(self, *, drain=False, stable=1): + deadline = time.monotonic() + self.args.idle_timeout + consecutive = 0 + observations = [] + while True: + self.control("/health") + gauges = request_gauges(self.control("/metrics")) + observations.append({"at_utc": utc_now(), **gauges}) + if all(value == 0 for value in gauges.values()): + consecutive += 1 + if consecutive >= stable: + break + else: + consecutive = 0 + if not drain: + raise RuntimeError("Service is busy before an exclusive check") + if time.monotonic() >= deadline: + raise TimeoutError("Service did not reach idle before the deadline") + time.sleep(0.5) + self.journal.record["idle_checks"].append(observations) + self.journal.save() + + def post(self, path, body, *, stream=False): + row = { + "path": path, + "request": body, + "started_at_utc": utc_now(), + "complete": False, + } + index = len(self.journal.record["requests"]) + self.journal.record["requests"].append(row) + self.journal.save() + data = json.dumps(body).encode("utf-8") + request = urllib.request.Request( + self.base + path, data=data, headers=self.headers + ) + started = time.perf_counter() + try: + with urllib.request.urlopen( + request, timeout=self.args.request_timeout + ) as response: + row["http_status"] = response.status + if not stream: + payload = json.loads(response.read()) + row["response"] = payload + row["elapsed_seconds"] = time.perf_counter() - started + else: + chunks, first, usage, done = [], None, None, False + row["sse_chunks"] = chunks + for line in response: + if not line.startswith(b"data:"): + continue + raw = line[5:].strip() + if raw == b"[DONE]": + done = True + break + chunk = json.loads(raw) + chunks.append(chunk) + if chunk.get("error"): + raise RuntimeError(f"Streaming API error: {chunk['error']}") + if first is None and first_token_delta(chunk): + first = time.perf_counter() - started + if chunk.get("usage") is not None: + usage = chunk["usage"] + row.update( + ttft_seconds=first, + usage=usage, + saw_done=done, + elapsed_seconds=time.perf_counter() - started, + ) + if first is None or usage is None or not done: + raise RuntimeError( + "Stream lacks a token delta, final usage, or DONE marker" + ) + payload = row + row["complete"] = True + return payload, index + except urllib.error.HTTPError as exc: + row.update( + http_status=exc.code, + error_response=exc.read().decode("utf-8", errors="replace"), + ) + raise + except Exception as exc: + row["error"] = f"{type(exc).__name__}: {exc}" + raise + finally: + row["finished_at_utc"] = utc_now() + self.journal.save() + + def count(self, messages): + response, index = self.post( + "/tokenize", + {"model": self.model, "messages": messages, "add_generation_prompt": True}, + ) + count = response.get("count") + if type(count) is not int or count <= 0: + raise ValueError("Tokenizer must return a positive integer count") + return count, index + + def calibrate(self, tokens, fact): + nonce = uuid.uuid4().hex + prefix = f"Test {nonce}. The project code is {fact}. Remember it.\n" + suffix = "\nWhat is the project code? Reply with just the code." + words = tokens + indices = [] + for _ in range(12): + if words < 0: + raise ValueError("Exact-token calibration exhausted prompt padding") + text = ( + prefix + + " ".join( + (["alpha", "beta", "gamma", "delta"] * ((words + 3) // 4))[:words] + ) + + suffix + ) + messages = [{"role": "user", "content": text}] + count, index = self.count(messages) + indices.append(index) + if count == tokens: + return messages, indices + words += tokens - count + raise ValueError("Exact prompt calibration did not converge in twelve attempts") + + def prefill(self, tokens, phase, sample): + self.idle() + messages, calibration = self.calibrate(tokens, "STONE-7482") + self.idle() + case = { + "phase": phase, + "sample": sample, + "tokens": tokens, + "calibration_requests": calibration, + "accepted": False, + "request_index": len(self.journal.record["requests"]), + } + self.journal.record["cases"].append(case) + response, index = self.post( + "/v1/chat/completions", + { + "model": self.model, + "messages": messages, + "max_tokens": 1, + "temperature": 0, + "stream": True, + "stream_options": {"include_usage": True}, + }, + stream=True, + ) + case["request_index"] = index + case["cached_tokens"] = validate_usage(response["usage"], tokens, "cold") + case["ttft_seconds"] = response["ttft_seconds"] + self.idle(drain=True) + case["accepted"] = True + self.journal.save() + print( + json.dumps( + { + key: case[key] + for key in ("phase", "sample", "tokens", "ttft_seconds", "accepted") + } + ), + flush=True, + ) + + def semantic(self, tokens, *, reuse): + self.idle() + fact = "RIVER-" + str(1000 + int(uuid.uuid4().hex[:4], 16) % 9000) + messages, calibration = self.calibrate(tokens, fact) + for kind in ("cold", "repeated", "extended") if reuse else ("cold",): + self.idle() + query = ( + messages + if kind != "extended" + else [ + { + "role": "user", + "content": messages[0]["content"] + + "\nFinal instruction: give exactly the project code.", + } + ] + ) + expected_tokens, count_index = self.count(query) + if kind != "extended" and expected_tokens != tokens: + raise ValueError( + "Exact semantic prompt count changed after calibration" + ) + if kind == "extended" and expected_tokens <= tokens: + raise ValueError("Extended prompt did not increase the token count") + self.idle() + case = { + "phase": "semantic", + "kind": kind, + "tokens": expected_tokens, + "base_tokens": tokens, + "expected_answer": fact, + "calibration_requests": calibration, + "count_request": count_index, + "accepted": False, + "request_index": len(self.journal.record["requests"]), + } + self.journal.record["cases"].append(case) + response, index = self.post( + "/v1/chat/completions", + { + "model": self.model, + "messages": query, + "max_tokens": self.args.semantic_max_tokens, + "temperature": 0, + "top_p": 1, + }, + ) + case["request_index"] = index + choice = response["choices"][0] + answer = (choice["message"].get("content") or "").strip() + case.update(answer=answer, finish_reason=choice.get("finish_reason")) + case["cached_tokens"] = validate_usage( + response["usage"], + expected_tokens, + "cold" if kind == "cold" else "reuse", + ) + if answer != fact or choice.get("finish_reason") != "stop": + raise RuntimeError( + "Exact-answer check failed; inspect the preserved response" + ) + self.idle(drain=True) + case["accepted"] = True + self.journal.save() + print( + json.dumps( + { + key: case[key] + for key in ( + "phase", + "kind", + "tokens", + "answer", + "cached_tokens", + "accepted", + ) + } + ), + flush=True, + ) + + def run(self): + self.journal.record["models_response"] = json.loads(self.control("/v1/models")) + models = self.journal.record["models_response"]["data"] + if self.model is None: + if len(models) != 1: + raise ValueError("Specify --model when the API serves multiple models") + self.model = models[0]["id"] + if self.model not in {row["id"] for row in models}: + raise ValueError("Requested model is absent from /v1/models") + self.journal.record["model"] = self.model + self.idle(drain=True, stable=3) + for size in SIZES: + self.prefill(size, "warmup", 0) + self.semantic(8192, reuse=True) + self.semantic(16384, reuse=False) + self.semantic(32768, reuse=False) + self.journal.save() + for sample in range(REPEATS): + for size in SIZES: + self.prefill(size, "measured", sample) + self.journal.record["summary"] = summarize(self.journal.record["cases"]) + + +def parser(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--base-url", required=True) + p.add_argument( + "--model", help="Served model ID; inferred only if /v1/models lists one model" + ) + p.add_argument( + "--conditions", + type=Path, + help="Optional operator-supplied sources/settings JSON", + ) + p.add_argument("--output", type=Path, required=True) + p.add_argument( + "--label", + required=True, + help="Operator label, not evidence that a feature executed", + ) + p.add_argument( + "--ready-confirmed", + action="store_true", + help="Attest that all-rank model warmup completed", + ) + p.add_argument( + "--exclusive-window", + action="store_true", + help="Attest that this client owns the inference window", + ) + p.add_argument("--request-timeout", type=float, default=180) + p.add_argument("--idle-timeout", type=float, default=60) + p.add_argument( + "--api-key-env", + help="API-key environment variable; its value is never journaled", + ) + p.set_defaults(semantic_max_tokens=SEMANTIC_MAX_TOKENS) + return p + + +def main(): + args = parser().parse_args() + if not args.ready_confirmed or not args.exclusive_window: + raise ValueError( + "Execution requires completed model warmup and an exclusive inference " + "window, attested by both flags" + ) + if any( + not math.isfinite(value) or value <= 0 + for value in (args.request_timeout, args.idle_timeout) + ): + raise ValueError("Timeouts must be finite and positive") + # Reject credentials before creating a journal that includes the command. + args.base_url = normalize_base_url(args.base_url) + conditions, conditions_sha = load_conditions(args.conditions) + clock = time.get_clock_info("perf_counter") + record = { + "schema": "glm-prefill-checkpoints-reproduction/v1", + "qualification": "bounded semantic/cache smoke checks and TTFT observations", + "execution_status": "running", + "label": args.label, + "run_id": uuid.uuid4().hex, + "started_at_utc": utc_now(), + "command": sys.argv, + "base_url": args.base_url, + "operator_attestations": { + "all_rank_warmup_complete": True, + "exclusive_window": True, + }, + "conditions": conditions, + "conditions_sha256": conditions_sha, + "conditions_provenance": "operator_supplied" + if conditions is not None + else "not_supplied", + "feature_activation": { + "status": "not_verified", + "source": "separate runtime evidence required", + }, + "model_quality_qualified": False, + "script_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "measured_protocol_source_sha256": MEASURED_PROTOCOL_SOURCE_SHA256, + "clock": { + name: getattr(clock, name) + for name in ("implementation", "resolution", "monotonic", "adjustable") + }, + "settings": { + "sizes": list(SIZES), + "repeats": REPEATS, + "inference_concurrency": 1, + "prefill_output_tokens": 1, + "semantic_max_tokens": SEMANTIC_MAX_TOKENS, + }, + "requests": [], + "cases": [], + "idle_checks": [], + "limitations": [ + "Synthetic exact-answer/cache checks do not establish full model quality " + "or numerical equivalence.", + "TTFT includes API/client/network time; it is not GPU kernel time.", + "Idle snapshots and a serial client cannot exclude outside requests; " + "the operator controls exclusivity.", + "Source/settings metadata and completed warmup are operator attestations, " + "not remotely verified.", + "Feature activation is not checked; collect request-associated per-rank " + "dispatch logs separately before attributing timings to an optimization.", + "The measured protocol used an external activation gate between semantic " + "checks and timing; this standalone client omits that collection step.", + "Positive reuse proves a cache hit, not every checkpoint destination " + "or persistent-cache tier.", + ], + } + journal = Journal(args.output, record) + try: + PrefillChecks(args, journal).run() + record["execution_status"] = "passed" + except Exception as exc: + record["execution_status"] = "failed" + record["error"] = f"{type(exc).__name__}: {exc}" + raise + finally: + record["finished_at_utc"] = utc_now() + journal.save() + + +if __name__ == "__main__": + main() diff --git a/docs/benchmarking/glm-kda-checkpoints-20260907/README.md b/docs/benchmarking/glm-kda-checkpoints-20260907/README.md new file mode 100644 index 000000000000..5fbdff1c43e1 --- /dev/null +++ b/docs/benchmarking/glm-kda-checkpoints-20260907/README.md @@ -0,0 +1,25 @@ +# Continuation-prefill coalescing: serving evidence + +Status: **research-only**. Four GB10 GPUs running GLM-5.3-Flash with +TP4/DCP4 were tested under all four combinations of continuation coalescing +and token-sharded mHC. [Results](report.md) summarize 36 cold prefill samples, +20 exact-answer/cache checks, and 16 decode cells. [Measurements](evidence.json) +retain individual timings, activation witnesses, settings, and artifact hashes. + +[Source composition](manifest.json) identifies the standalone feature revision +and the combined runtime used for measurements. Combining the pinned feature +revisions in that manifest reproduces the measured runtime and test subtrees. +The separately supplied benchmark client and its tests are reproduction tools, +not part of the measured server image. + +These observations cover bounded serving checks, not full model-quality +qualification. The table includes disabled and enabled feature combinations +on the same image. Native 512-token split-page geometry, diagnostic logging, +and fused SparkRing transport are constant. Reboots between configurations and +short, sequential measurement windows limit small-difference conclusions. + +The JSON preserves private-source artifact hashes for audit while excluding +request text, private addresses and deployment credentials. The component +implementation does not require the private deployment tooling. + +[Reproduction instructions](reproduction.md) pin the model and benchmark inputs. diff --git a/docs/benchmarking/glm-kda-checkpoints-20260907/decode-benchmark.patch b/docs/benchmarking/glm-kda-checkpoints-20260907/decode-benchmark.patch new file mode 100644 index 000000000000..0d524f3533a8 --- /dev/null +++ b/docs/benchmarking/glm-kda-checkpoints-20260907/decode-benchmark.patch @@ -0,0 +1,294 @@ +diff --git a/llm_decode_bench.py b/llm_decode_bench.py +index d4e5508..d53f270 100644 +--- a/llm_decode_bench.py ++++ b/llm_decode_bench.py +@@ -33,19 +33,30 @@ import signal + import string + import subprocess + import sys +-import termios ++try: ++ import termios ++ import tty ++except ImportError: # Windows has no POSIX terminal modules. ++ termios = None ++ tty = None + import threading + import time +-import tty + import zipfile + import zlib +-from dataclasses import dataclass, field, asdict, fields as dataclass_fields ++from dataclasses import dataclass, field, asdict, fields as dataclass_fields, replace + from datetime import datetime + from pathlib import Path + from statistics import mean, median, pstdev + from typing import Optional + from urllib.parse import urlparse + ++if os.name == "nt": ++ for _stream in (sys.stdout, sys.stderr): ++ try: ++ _stream.reconfigure(encoding="utf-8") ++ except (AttributeError, OSError): ++ pass ++ + import httpx + from rich import box + from rich.columns import Columns +@@ -4649,6 +4660,8 @@ class CompletionStatsRun: + @dataclass + class GpuStats: + index: int = 0 ++ label: str = "" ++ host: str = "" + temp_c: float = 0.0 + gpu_util_pct: float = 0.0 + mem_util_pct: float = 0.0 +@@ -4662,6 +4675,33 @@ class GpuStats: + pcie_width: float = 0.0 + pcie_rx_mb_s: float = 0.0 + pcie_tx_mb_s: float = 0.0 ++ stale: bool = False ++ ++ ++@dataclass(frozen=True) ++class RemoteHardwareTarget: ++ label: str ++ ssh_target: str ++ ++ ++@dataclass ++class RemoteHardwareSample: ++ target: RemoteHardwareTarget ++ gpu: Optional[GpuStats] = None ++ gpus: list[GpuStats] = field(default_factory=list) ++ cpu_total_ticks: float = 0.0 ++ cpu_idle_ticks: float = 0.0 ++ cpu_freq_mhz: float = 0.0 ++ cpu_temp_c: float = 0.0 ++ net_rx_bytes: float = 0.0 ++ net_tx_bytes: float = 0.0 ++ error: str = "" ++ ++ def __post_init__(self) -> None: ++ if self.gpu is not None and not self.gpus: ++ self.gpus = [self.gpu] ++ elif self.gpu is None and len(self.gpus) == 1: ++ self.gpu = self.gpus[0] + + + @dataclass +@@ -4828,6 +4868,108 @@ def _to_float(value: str) -> float: + return 0.0 + + ++def resolve_prompt_run_id(requested: str) -> str: ++ if requested: ++ if not re.fullmatch(r"[a-z]{12}", requested): ++ raise ValueError("prompt run id must contain exactly 12 lowercase ASCII letters") ++ return requested ++ return "".join(random.choices(string.ascii_lowercase, k=12)) ++ ++ ++REMOTE_HARDWARE_SCRIPT = """ ++LC_ALL=C ++nvidia-smi --query-gpu=index,temperature.gpu,utilization.gpu,utilization.memory,memory.used,memory.total,power.draw,power.limit,clocks.sm,clocks.mem --format=csv,noheader,nounits 2>/dev/null | while IFS= read -r line; do printf 'GPU|%s\\n' "$line"; done ++printf 'PROCMB|'; nvidia-smi --query-compute-apps=used_memory --format=csv,noheader,nounits 2>/dev/null | awk 'BEGIN{s=0} $1 ~ /^[0-9.]+$/ {s+=$1} END{printf "%.3f\\n",s}' ++awk '/MemTotal:/ {t=$2} /MemAvailable:/ {a=$2} END {printf "MEM|%.3f|%.3f\\n",t/1024,(t-a)/1024}' /proc/meminfo ++awk '/^cpu / {idle=$5+$6; total=0; for(i=2;i<=NF;i++) total+=$i; printf "CPU|%.0f|%.0f\\n",total,idle; exit}' /proc/stat ++awk '/cpu MHz/ {s+=$4;n++} END {printf "FREQ|%.1f\\n",n?s/n:0}' /proc/cpuinfo ++printf 'TEMP|0\\nNET|0|0\\n' ++""".strip() ++ ++ ++def parse_remote_hardware_targets(value: str) -> list[RemoteHardwareTarget]: ++ targets: list[RemoteHardwareTarget] = [] ++ labels: set[str] = set() ++ for index, raw in enumerate((value or "").split(",")): ++ raw = raw.strip() ++ if not raw: ++ continue ++ if "=" in raw: ++ label, ssh_target = (part.strip() for part in raw.split("=", 1)) ++ else: ++ label, ssh_target = f"S{index}", raw ++ if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_.-]{0,15}", label): ++ raise ValueError(f"invalid remote hardware label {label!r}") ++ if not ssh_target or ssh_target.startswith("-") or not re.fullmatch(r"[A-Za-z0-9_.:@%+\-\[\]]+", ssh_target): ++ raise ValueError(f"invalid SSH target {ssh_target!r}") ++ if label in labels: ++ raise ValueError(f"duplicate remote hardware label {label!r}") ++ labels.add(label) ++ targets.append(RemoteHardwareTarget(label, ssh_target)) ++ return targets ++ ++ ++def _parse_remote_hardware_output(target: RemoteHardwareTarget, index: int, output: str) -> RemoteHardwareSample: ++ values: dict[str, list[list[str]]] = {} ++ for raw_line in output.splitlines(): ++ if "|" not in raw_line: ++ continue ++ tag, *parts = raw_line.strip().split("|") ++ values.setdefault(tag.upper(), []).append([part.strip() for part in parts]) ++ gpu_rows = values.get("GPU", []) ++ if not gpu_rows: ++ return RemoteHardwareSample(target=target, error="remote nvidia-smi returned no GPU sample") ++ proc_used_mb = _to_float((values.get("PROCMB") or [["0"]])[0][0]) ++ mem_parts = (values.get("MEM") or [["0", "0"]])[0] ++ physical_total_mb = _to_float(mem_parts[0]) ++ is_uma = len(gpu_rows) == 1 and physical_total_mb > 0 ++ gpus: list[GpuStats] = [] ++ for row in gpu_rows: ++ parts = row[0].split(",") ++ if len(parts) < 9: ++ continue ++ vals = [_to_float(part) for part in parts] ++ vals.extend([0.0] * (10 - len(vals))) ++ gpu_index = int(vals[0]) ++ gpus.append(GpuStats(index=index * 100 + gpu_index, label=target.label if is_uma else f"{target.label}-G{gpu_index}", host=target.ssh_target, temp_c=vals[1], gpu_util_pct=vals[2], mem_util_pct=vals[3], mem_used_mb=proc_used_mb if is_uma else vals[4], mem_total_mb=physical_total_mb if is_uma else vals[5], power_w=vals[6], power_limit_w=vals[7], sm_clock_mhz=vals[8], mem_clock_mhz=vals[9])) ++ cpu = (values.get("CPU") or [["0", "0"]])[0] ++ net = (values.get("NET") or [["0", "0"]])[0] ++ return RemoteHardwareSample(target=target, gpus=gpus, cpu_total_ticks=_to_float(cpu[0]), cpu_idle_ticks=_to_float(cpu[1]) if len(cpu) > 1 else 0.0, cpu_freq_mhz=_to_float((values.get("FREQ") or [["0"]])[0][0]), cpu_temp_c=_to_float((values.get("TEMP") or [["0"]])[0][0]), net_rx_bytes=_to_float(net[0]), net_tx_bytes=_to_float(net[1]) if len(net) > 1 else 0.0) ++ ++ ++def _sample_remote_hardware_target(target: RemoteHardwareTarget, index: int, timeout: float = 10.0) -> RemoteHardwareSample: ++ if shutil.which("ssh") is None: ++ return RemoteHardwareSample(target=target, error="ssh executable not found") ++ try: ++ proc = subprocess.run(["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=3", "-o", "ConnectionAttempts=1", "-o", "LogLevel=ERROR", target.ssh_target, REMOTE_HARDWARE_SCRIPT], check=False, capture_output=True, text=True, timeout=timeout) ++ except subprocess.TimeoutExpired: ++ return RemoteHardwareSample(target=target, error=f"SSH sample timed out after {timeout:g}s") ++ except Exception as exc: ++ return RemoteHardwareSample(target=target, error=f"{type(exc).__name__}: {exc}") ++ if proc.returncode != 0: ++ return RemoteHardwareSample(target=target, error=(proc.stderr or proc.stdout or f"ssh exit {proc.returncode}").strip()[:160]) ++ return _parse_remote_hardware_output(target, index, proc.stdout) ++ ++ ++def _sample_remote_hardware(targets, previous, last_good=None): ++ last_good = last_good or {} ++ samples = [_sample_remote_hardware_target(target, index) for index, target in enumerate(targets)] ++ gpus: list[GpuStats] = [] ++ errors: list[str] = [] ++ for sample in samples: ++ if sample.error or not sample.gpus: ++ errors.append(f"{sample.target.label}: {sample.error or 'no sample'}") ++ for label, cached in last_good.items(): ++ if label == sample.target.label or label.startswith(f"{sample.target.label}-G"): ++ gpus.append(replace(cached, stale=True)) ++ continue ++ for gpu in sample.gpus: ++ gpu.stale = False ++ last_good[gpu.label] = replace(gpu) ++ gpus.append(gpu) ++ return gpus, 0.0, 0.0, [], errors ++ ++ + def _sample_gpu_query() -> list[GpuStats]: + if shutil.which("nvidia-smi") is None: + return [] +@@ -7211,34 +7353,49 @@ def collect_startup_diagnostics(args, base_url: str) -> dict: + return diagnostics + + +-def start_hardware_monitor(state: TUIState, interval: float) -> None: ++def start_hardware_monitor(state: TUIState, interval: float, ssh_hosts: str = "") -> None: + if interval <= 0: + state.hw_monitor_enabled = False + state.hw_last_error = "hardware monitor disabled" + return +- if shutil.which("nvidia-smi") is None: ++ try: ++ remote_targets = parse_remote_hardware_targets(ssh_hosts) ++ except ValueError as exc: ++ state.hw_monitor_enabled = False ++ state.hw_last_error = str(exc) ++ return ++ if remote_targets and shutil.which("ssh") is None: ++ state.hw_monitor_enabled = False ++ state.hw_last_error = "ssh executable not found" ++ return ++ if not remote_targets and shutil.which("nvidia-smi") is None: + state.hw_monitor_enabled = False + state.hw_last_error = "nvidia-smi not found" + return + state.hw_monitor_enabled = True + + def loop() -> None: ++ previous = {} ++ last_good = {} + while True: + try: +- gpus = _sample_gpu_query() +- pcie = _sample_gpu_pcie() +- for gpu in gpus: +- rx, tx = pcie.get(gpu.index, (0.0, 0.0)) +- gpu.pcie_rx_mb_s = rx +- gpu.pcie_tx_mb_s = tx +- cpu_util, cpu_freq = _sample_cpu_stats() +- cpu_temps = _sample_cpu_temperatures() ++ if remote_targets: ++ gpus, cpu_util, cpu_freq, cpu_temps, errors = _sample_remote_hardware(remote_targets, previous, last_good) ++ else: ++ gpus = _sample_gpu_query() ++ pcie = _sample_gpu_pcie() ++ for gpu in gpus: ++ rx, tx = pcie.get(gpu.index, (0.0, 0.0)) ++ gpu.pcie_rx_mb_s = rx ++ gpu.pcie_tx_mb_s = tx ++ cpu_util, cpu_freq = _sample_cpu_stats() ++ cpu_temps = _sample_cpu_temperatures() + state.cpu_util_pct = cpu_util + state.cpu_freq_mhz = cpu_freq + state.cpu_temps = cpu_temps + state.gpu_stats = gpus + state.hw_available = bool(gpus) +- state.hw_last_error = "" if gpus else "no GPU samples" ++ state.hw_last_error = "; ".join(errors) if remote_targets and errors else ("" if gpus else "no GPU samples") + state.hw_last_update = time.monotonic() + if gpus: + state.hw_history.append( +@@ -13029,11 +13186,11 @@ async def run_completion_stats_benchmark(args) -> dict: + if not hw_monitor_requested: + hw_state.hw_monitor_enabled = False + hw_state.hw_last_error = "hardware monitor disabled" +- elif shutil.which("nvidia-smi") is None: ++ elif not args.hw_ssh_hosts and shutil.which("nvidia-smi") is None: + hw_state.hw_monitor_enabled = False + hw_state.hw_last_error = "nvidia-smi not found" + else: +- start_hardware_monitor(hw_state, args.hw_monitor_interval) ++ start_hardware_monitor(hw_state, args.hw_monitor_interval, args.hw_ssh_hosts) + + payload = { + "model": args.model, +@@ -14004,7 +14161,7 @@ async def run_benchmark(args): + state.hw_monitor_enabled = False + state.hw_last_error = "hardware monitor disabled" + add_event(state, "startup hardware monitor disabled") +- elif shutil.which("nvidia-smi") is None: ++ elif not args.hw_ssh_hosts and shutil.which("nvidia-smi") is None: + state.hw_monitor_enabled = False + state.hw_last_error = "nvidia-smi not found" + console.print( +@@ -14018,7 +14175,7 @@ async def run_benchmark(args): + "(benchmark may be running off the GPU server/container)", + ) + else: +- start_hardware_monitor(state, args.hw_monitor_interval) ++ start_hardware_monitor(state, args.hw_monitor_interval, args.hw_ssh_hosts) + add_event(state, f"hardware monitor interval={args.hw_monitor_interval:g}s") + + # Mark skipped decode cells +@@ -16412,6 +16569,10 @@ def parse_args(): + help="Maximum GPUs to show in the live hardware panel. All sampled GPUs " + "still contribute to aggregate PCIe rx/tx. (default: 8)" + ) ++ parser.add_argument( ++ "--hw-ssh-hosts", default="", ++ help="Comma-separated remote hardware targets, optionally labeled as S0=user@host." ++ ) + parser.add_argument( + "--p2pmark", action="store_true", + help="Run the bundled CUDA/NCCL P2P diagnostic before the LLM benchmark " diff --git a/docs/benchmarking/glm-kda-checkpoints-20260907/evidence.json b/docs/benchmarking/glm-kda-checkpoints-20260907/evidence.json new file mode 100644 index 000000000000..aa390048cf75 --- /dev/null +++ b/docs/benchmarking/glm-kda-checkpoints-20260907/evidence.json @@ -0,0 +1,1534 @@ +{ + "schema": "glm-prefill-four-arm-public-evidence/v1", + "exporter_sha256": "290fede850bc4f296bd6b689e4dfcc3ef71a641bee23184a1b64ef61b16bf285", + "status": "complete", + "qualification": "bounded semantic/cache smoke checks and performance observations; not full numerical/model-quality qualification", + "hardware": "four NVIDIA GB10 Sparks", + "model_family": "GLM-5.3-Flash-NVFP4-Spark", + "runtime_sources": { + "vllm": "abb715f132bdccb592a34b2596a3d3a8d757ffbc", + "b12x": "70fe41974ef4b18f61caaa2579c81cdc05d1265f" + }, + "image_digest": "sha256:52b207e716a285c16e5e1b14ec2a41f6208b9450c617e7d1cde5a507ea879d7f", + "settings": { + "tp": 4, + "dcp": 4, + "pp": 1, + "max_batched_tokens": 8192, + "max_model_len": 1048576, + "block_size": 512, + "mamba_block_size": 512, + "mtp_tokens": 3, + "kv_bytes_per_rank": 25769803776, + "prefix_cache_retention_interval": 0, + "recurrent_checkpoint_policy": "aligned", + "speculative_method": "mtp" + }, + "native_split_settings": { + "VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE": "512", + "VLLM_GLM53_SPLIT_MAMBA_BLOCK_SIZE": "512" + }, + "diagnostics_enabled": true, + "operational_comparison": { + "recovery_events": [ + { + "event": "one rank rebooted between experiment arms", + "after_arm": "neither", + "before_arm": "continuation", + "reason": "recover memory allocation capacity after a failed startup preflight", + "recorded_settings_matched_before_after": true, + "recorded_settings": { + "driver": "580.173.02", + "maximum_sm_clock_mhz": 3003, + "application_graphics_clock_mhz": 2418, + "persistence": "enabled", + "cpu0_governor": "performance" + }, + "receipt_sha256": "64254bb4ee8a54428b5695ef4f4096754d8921a685dc9bddf3db5294e9a2fc45" + }, + { + "event": "one rank rebooted between experiment arms", + "after_arm": "continuation", + "before_arm": "mhc", + "reason": "recover memory allocation capacity after a failed startup preflight", + "recorded_settings_matched_before_after": true, + "recorded_settings": { + "driver": "580.173.02", + "maximum_sm_clock_mhz": 3003, + "application_graphics_clock_mhz": 2418, + "persistence": "enabled", + "cpu0_governor": "performance" + }, + "receipt_sha256": "093455e00436e5b0f87b6de3d1972dffc4e34e18d9631b416a28d48bca1a620e" + } + ], + "continuous_operating_condition_equivalence_verified": false + }, + "measurement_order": [ + "both", + "neither", + "continuation", + "mhc" + ], + "arms": [ + { + "arm": "neither", + "flags": { + "VLLM_B12X_KDA_PREFILL_COALESCING": "0", + "VLLM_GLM53_MHC_PREFILL_SHARD": "0" + }, + "started_at_utc": "2026-09-07T13:01:15.745976+00:00", + "prefill_samples": [ + { + "prompt_tokens": 8192, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 4.196101099980297, + "tokens_per_second": 1952.2885185103064 + }, + { + "prompt_tokens": 8192, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 4.203437799995299, + "tokens_per_second": 1948.8809849902289 + }, + { + "prompt_tokens": 8192, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 4.201881200016942, + "tokens_per_second": 1949.6029540213965 + }, + { + "prompt_tokens": 16384, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 6.943599899997935, + "tokens_per_second": 2359.582959266543 + }, + { + "prompt_tokens": 16384, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 6.9457052999932785, + "tokens_per_second": 2358.8677164313112 + }, + { + "prompt_tokens": 16384, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 6.966370399983134, + "tokens_per_second": 2351.8703513151795 + }, + { + "prompt_tokens": 32768, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 12.541768900002353, + "tokens_per_second": 2612.709599520196 + }, + { + "prompt_tokens": 32768, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 12.532975899986923, + "tokens_per_second": 2614.5426482495823 + }, + { + "prompt_tokens": 32768, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 12.522141299996292, + "tokens_per_second": 2616.8048431149473 + } + ], + "prefill_summary": [ + { + "prompt_tokens": 8192, + "samples": 3, + "median_ttft_seconds": 4.201881200016942, + "min_ttft_seconds": 4.196101099980297, + "max_ttft_seconds": 4.203437799995299, + "tokens_per_second": 1949.6029540213965, + "throughput_change_vs_neither_pct": 0.0, + "ttft_reduction_vs_neither_pct": 0.0 + }, + { + "prompt_tokens": 16384, + "samples": 3, + "median_ttft_seconds": 6.9457052999932785, + "min_ttft_seconds": 6.943599899997935, + "max_ttft_seconds": 6.966370399983134, + "tokens_per_second": 2358.8677164313112, + "throughput_change_vs_neither_pct": 0.0, + "ttft_reduction_vs_neither_pct": 0.0 + }, + { + "prompt_tokens": 32768, + "samples": 3, + "median_ttft_seconds": 12.532975899986923, + "min_ttft_seconds": 12.522141299996292, + "max_ttft_seconds": 12.541768900002353, + "tokens_per_second": 2614.5426482495823, + "throughput_change_vs_neither_pct": 0.0, + "ttft_reduction_vs_neither_pct": 0.0 + } + ], + "semantic_smoke_checks": [ + { + "kind": "cold", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "repeated", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 7168, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "extended", + "prompt_tokens": 8201, + "base_tokens": 8192, + "cached_tokens": 4096, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 16384, + "base_tokens": 16384, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 32768, + "base_tokens": 32768, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + } + ], + "activation": { + "checkpoint_enabled": false, + "checkpoint_scheduled": false, + "checkpoint_dispatch_shape": null, + "checkpoint_capacity": 1, + "configured_checkpoint_grid": null, + "mhc_enabled": false, + "mhc_collectives_per_dispatch": null, + "mhc_request_dispatches_by_rank": [ + 0, + 0, + 0, + 0 + ], + "mhc_common_dispatch_count": 0, + "mhc_rows": null, + "per_kernel_gpu_completion_verified": false, + "completed_requests_corroborate_dispatch": true + }, + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tokens_per_second": 49.16553901666799, + "mtp_accept_length": 2.752808988764045, + "mtp_normalized_steps_per_second": 17.86013458156511, + "reported_server_steps_per_second": 17.892046308818017, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.953, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 981, + "server_output_tokens": 981, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3599999999860302, + "latency": 20.578999999997905, + "input_tokens": 8192, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.3900000000139698, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 151, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.0, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.0, + "reported_server_steps_per_second_change_vs_neither_pct": 0.0 + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tokens_per_second": 122.18713977842667, + "mtp_accept_length": 2.7958715596330275, + "mtp_normalized_steps_per_second": 43.702701348149326, + "reported_server_steps_per_second": 43.702701348149326, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.953, + "measurement_wall_seconds": 20.015, + "client_output_tokens": 2438, + "server_output_tokens": 2438, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3589999999967404, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 711, + "completed": false + }, + { + "ttft": 3.2030000000086147, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 707, + "completed": false + }, + { + "ttft": 3.2030000000086147, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 725, + "completed": false + }, + { + "ttft": 3.2030000000086147, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 758, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.0, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.0, + "reported_server_steps_per_second_change_vs_neither_pct": 0.0 + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tokens_per_second": 46.18463847889143, + "mtp_accept_length": 2.591036414565826, + "mtp_normalized_steps_per_second": 17.82477398590729, + "reported_server_steps_per_second": 17.86339754817361, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.985, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 923, + "server_output_tokens": 923, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 21.34299999999348, + "input_tokens": 32768, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.4070000000065193, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 163, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.0, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.0, + "reported_server_steps_per_second_change_vs_neither_pct": 0.0 + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tokens_per_second": 123.43883232178055, + "mtp_accept_length": 2.7839366515837103, + "mtp_normalized_steps_per_second": 44.339669960363274, + "reported_server_steps_per_second": 44.339669960363274, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.937, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 2461, + "server_output_tokens": 2461, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4530000000086147, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 764, + "completed": false + }, + { + "ttft": 3.312000000005355, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 786, + "completed": false + }, + { + "ttft": 3.312000000005355, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 709, + "completed": false + }, + { + "ttft": 3.312000000005355, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 734, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.0, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.0, + "reported_server_steps_per_second_change_vs_neither_pct": 0.0 + } + ], + "artifact_sha256": { + "conditions": "106930faaea7c14960599399e280bd5a2930124ebb9ed81a181ece2373bf5746", + "prefill": "c01c9662b422ccd7926589c69136b41a26871ea4f8401a32cb7da23ad6f8d65d", + "activation": "af7926a719fcc0b52535bc850d639daf681bf2ad843eb9ccff582377e7100372", + "complete": "7fe568efe23bc6c96dec974c0f01678b5d6f11daa2e4b8b29b5621e87ee5eb3d", + "decode": "1ff4deea8f1d2d9d4bbf18fd1b28a40a48f340b6f1fcf3ccf55d4fe3933872ea", + "decode_receipt": "bd636073241c3003bd1fcc955c7bba8b79b758c4b131686c3b4edf1c67e3f0b7" + }, + "protocol_sources": { + "prefill_harness_sha256": "6e4bf5ff62379bb2db4dc439d27a85eda0837959e983ef5f92ab3605898e69ac", + "activation_validator_sha256": "3359d4d9b5cf28b95c68a4989ac11cd8119dbeaf39058b9c1c27bcf95f218a09", + "decode_harness_sha256": "46ace1dad13c245807bc1b4ccf4ab6b90e95d12a99dd729ee24caa51034194c6", + "decode_wrapper_sha256": "5c2f7ecaef4c8e1ea60b7c44d3e39f710c660387fb2c0caa0871682f4ec712ce" + }, + "source_manifest_sha256": "082301aa98d9a4e0a84e52d32ff1c4efab29bd888c0fa8606182ba55a6a72e85", + "package_versions": { + "torch": "2.13.0+cu130", + "triton": "3.7.1", + "transformers": "5.16.1", + "nvidia-cutlass-dsl": "4.6.2", + "flashinfer-python": "0.6.17" + } + }, + { + "arm": "continuation", + "flags": { + "VLLM_B12X_KDA_PREFILL_COALESCING": "1", + "VLLM_GLM53_MHC_PREFILL_SHARD": "0" + }, + "started_at_utc": "2026-09-07T13:24:22.409490+00:00", + "prefill_samples": [ + { + "prompt_tokens": 8192, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 2.7674208999960683, + "tokens_per_second": 2960.1568738646292 + }, + { + "prompt_tokens": 8192, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 2.7794476999843027, + "tokens_per_second": 2947.3481368425337 + }, + { + "prompt_tokens": 8192, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 2.7874553999863565, + "tokens_per_second": 2938.8811028295186 + }, + { + "prompt_tokens": 16384, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 5.595004900009371, + "tokens_per_second": 2928.3263004778705 + }, + { + "prompt_tokens": 16384, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 5.563162100006593, + "tokens_per_second": 2945.087650776989 + }, + { + "prompt_tokens": 16384, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 5.603458099998534, + "tokens_per_second": 2923.908719867877 + }, + { + "prompt_tokens": 32768, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 11.21369129998493, + "tokens_per_second": 2922.1421495742475 + }, + { + "prompt_tokens": 32768, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 11.208513199992012, + "tokens_per_second": 2923.4921184750315 + }, + { + "prompt_tokens": 32768, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 11.190499399992405, + "tokens_per_second": 2928.1981821135028 + } + ], + "prefill_summary": [ + { + "prompt_tokens": 8192, + "samples": 3, + "median_ttft_seconds": 2.7794476999843027, + "min_ttft_seconds": 2.7674208999960683, + "max_ttft_seconds": 2.7874553999863565, + "tokens_per_second": 2947.3481368425337, + "throughput_change_vs_neither_pct": 51.176839918256874, + "ttft_reduction_vs_neither_pct": 33.85230167923129 + }, + { + "prompt_tokens": 16384, + "samples": 3, + "median_ttft_seconds": 5.595004900009371, + "min_ttft_seconds": 5.563162100006593, + "max_ttft_seconds": 5.603458099998534, + "tokens_per_second": 2928.3263004778705, + "throughput_change_vs_neither_pct": 24.141183504265463, + "ttft_reduction_vs_neither_pct": 19.446554981035757 + }, + { + "prompt_tokens": 32768, + "samples": 3, + "median_ttft_seconds": 11.208513199992012, + "min_ttft_seconds": 11.190499399992405, + "max_ttft_seconds": 11.21369129998493, + "tokens_per_second": 2923.4921184750315, + "throughput_change_vs_neither_pct": 11.81657795608302, + "ttft_reduction_vs_neither_pct": 10.567822922218273 + } + ], + "semantic_smoke_checks": [ + { + "kind": "cold", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "repeated", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 7168, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "extended", + "prompt_tokens": 8201, + "base_tokens": 8192, + "cached_tokens": 4096, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 16384, + "base_tokens": 16384, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 32768, + "base_tokens": 32768, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + } + ], + "activation": { + "checkpoint_enabled": true, + "checkpoint_scheduled": true, + "checkpoint_dispatch_shape": { + "tokens": 8192, + "interior_checkpoints": 4 + }, + "checkpoint_capacity": 4, + "configured_checkpoint_grid": [ + 512, + 512, + 2048 + ], + "mhc_enabled": false, + "mhc_collectives_per_dispatch": null, + "mhc_request_dispatches_by_rank": [ + 0, + 0, + 0, + 0 + ], + "mhc_common_dispatch_count": 0, + "mhc_rows": null, + "per_kernel_gpu_completion_verified": false, + "completed_requests_corroborate_dispatch": true + }, + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tokens_per_second": 47.3734288147036, + "mtp_accept_length": 2.6452513966480447, + "mtp_normalized_steps_per_second": 17.908856933119207, + "reported_server_steps_per_second": 17.92778807152631, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.969, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 946, + "server_output_tokens": 946, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3440000000118744, + "latency": 20.875, + "input_tokens": 8192, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.360000000015134, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 144, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -3.6450535025291497, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.27279946481695294, + "reported_server_steps_per_second_change_vs_neither_pct": 0.19976341493526117 + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tokens_per_second": 120.17856247175385, + "mtp_accept_length": 2.698198198198198, + "mtp_normalized_steps_per_second": 44.540301951134154, + "reported_server_steps_per_second": 44.54030195113415, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.937, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 2396, + "server_output_tokens": 2396, + "num_errors": 0, + "request_samples": [ + { + "ttft": 3.172000000020489, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 747, + "completed": false + }, + { + "ttft": 1.3590000000258442, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 732, + "completed": false + }, + { + "ttft": 3.172000000020489, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 685, + "completed": false + }, + { + "ttft": 3.172000000020489, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 729, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -1.6438532813806472, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 1.9165877100187467, + "reported_server_steps_per_second_change_vs_neither_pct": 1.9165877100187245 + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tokens_per_second": 45.92117782560242, + "mtp_accept_length": 2.5903954802259888, + "mtp_normalized_steps_per_second": 17.72747759025437, + "reported_server_steps_per_second": 17.72747759025437, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.969, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 917, + "server_output_tokens": 917, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 22.312999999994645, + "input_tokens": 32768, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.4219999999913853, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 107, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -0.570450829466651, + "mtp_normalized_steps_per_second_change_vs_neither_pct": -0.5458492530106884, + "reported_server_steps_per_second_change_vs_neither_pct": -0.7608852546257983 + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tokens_per_second": 125.2001601281229, + "mtp_accept_length": 2.7434210526315788, + "mtp_normalized_steps_per_second": 45.63650920737334, + "reported_server_steps_per_second": 45.63650920737334, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.984, + "measurement_wall_seconds": 20.015, + "client_output_tokens": 2502, + "server_output_tokens": 2502, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 773, + "completed": false + }, + { + "ttft": 3.25, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 762, + "completed": false + }, + { + "ttft": 3.25, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 781, + "completed": false + }, + { + "ttft": 3.25, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 728, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 1.4268830749718298, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 2.924783265570885, + "reported_server_steps_per_second_change_vs_neither_pct": 2.924783265570885 + } + ], + "artifact_sha256": { + "conditions": "f412918e008077386ab15b64593acb2ece6d1e236388b1eeac1a8ad7f9f64031", + "prefill": "802eae6d9bca97cc529933d186a50c3ae5d2d82e711888f8889986a2789d86ce", + "activation": "311fd5b30308386959bd719f5684b1e667af8ab454dafa9db02562dded25a310", + "complete": "0a01ca5a333609d6c3aa99a124b695ba203bba47c42d813e1a2cf5e2e14a1ac0", + "decode": "fccb47fc2141ec0cf0e1e1ba412706b2db96719d48301b38a9fb0c3a716a12bf", + "decode_receipt": "218a37c92ed7f70a40acf571a0fd39688602f50a6a5702c0201cce89fdd9b872" + }, + "protocol_sources": { + "prefill_harness_sha256": "6e4bf5ff62379bb2db4dc439d27a85eda0837959e983ef5f92ab3605898e69ac", + "activation_validator_sha256": "3359d4d9b5cf28b95c68a4989ac11cd8119dbeaf39058b9c1c27bcf95f218a09", + "decode_harness_sha256": "46ace1dad13c245807bc1b4ccf4ab6b90e95d12a99dd729ee24caa51034194c6", + "decode_wrapper_sha256": "5c2f7ecaef4c8e1ea60b7c44d3e39f710c660387fb2c0caa0871682f4ec712ce" + }, + "source_manifest_sha256": "082301aa98d9a4e0a84e52d32ff1c4efab29bd888c0fa8606182ba55a6a72e85", + "package_versions": { + "torch": "2.13.0+cu130", + "triton": "3.7.1", + "transformers": "5.16.1", + "nvidia-cutlass-dsl": "4.6.2", + "flashinfer-python": "0.6.17" + } + }, + { + "arm": "mhc", + "flags": { + "VLLM_B12X_KDA_PREFILL_COALESCING": "0", + "VLLM_GLM53_MHC_PREFILL_SHARD": "1" + }, + "started_at_utc": "2026-09-07T13:51:05.058507+00:00", + "prefill_samples": [ + { + "prompt_tokens": 8192, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 4.192637299973285, + "tokens_per_second": 1953.9014262102278 + }, + { + "prompt_tokens": 8192, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 4.173298399982741, + "tokens_per_second": 1962.9557282637347 + }, + { + "prompt_tokens": 8192, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 4.182778299989877, + "tokens_per_second": 1958.5068613413782 + }, + { + "prompt_tokens": 16384, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 6.833160300011514, + "tokens_per_second": 2397.7192515112506 + }, + { + "prompt_tokens": 16384, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 6.852885000000242, + "tokens_per_second": 2390.8178818117362 + }, + { + "prompt_tokens": 16384, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 6.870854299981147, + "tokens_per_second": 2384.5651915577596 + }, + { + "prompt_tokens": 32768, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 12.249805500003276, + "tokens_per_second": 2674.981247660727 + }, + { + "prompt_tokens": 32768, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 12.213526299980003, + "tokens_per_second": 2682.9270429502126 + }, + { + "prompt_tokens": 32768, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 12.218640999984927, + "tokens_per_second": 2681.8039747661314 + } + ], + "prefill_summary": [ + { + "prompt_tokens": 8192, + "samples": 3, + "median_ttft_seconds": 4.182778299989877, + "min_ttft_seconds": 4.173298399982741, + "max_ttft_seconds": 4.192637299973285, + "tokens_per_second": 1958.5068613413782, + "throughput_change_vs_neither_pct": 0.4567036227358878, + "ttft_reduction_vs_neither_pct": 0.4546273232805209 + }, + { + "prompt_tokens": 16384, + "samples": 3, + "median_ttft_seconds": 6.852885000000242, + "min_ttft_seconds": 6.833160300011514, + "max_ttft_seconds": 6.870854299981147, + "tokens_per_second": 2390.8178818117362, + "throughput_change_vs_neither_pct": 1.3544704163725552, + "ttft_reduction_vs_neither_pct": 1.3363696843447403 + }, + { + "prompt_tokens": 32768, + "samples": 3, + "median_ttft_seconds": 12.218640999984927, + "min_ttft_seconds": 12.213526299980003, + "max_ttft_seconds": 12.249805500003276, + "tokens_per_second": 2681.8039747661314, + "throughput_change_vs_neither_pct": 2.5725847907503097, + "ttft_reduction_vs_neither_pct": 2.5080627499038166 + } + ], + "semantic_smoke_checks": [ + { + "kind": "cold", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "repeated", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 7168, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "extended", + "prompt_tokens": 8201, + "base_tokens": 8192, + "cached_tokens": 4096, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 16384, + "base_tokens": 16384, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 32768, + "base_tokens": 32768, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + } + ], + "activation": { + "checkpoint_enabled": false, + "checkpoint_scheduled": false, + "checkpoint_dispatch_shape": null, + "checkpoint_capacity": 1, + "configured_checkpoint_grid": null, + "mhc_enabled": true, + "mhc_collectives_per_dispatch": { + "reduce_scatter": 90, + "all_gather": 90, + "auxiliary_gathers": 0 + }, + "mhc_request_dispatches_by_rank": [ + 20, + 20, + 20, + 20 + ], + "mhc_common_dispatch_count": 20, + "mhc_rows": { + "first_pre": { + "8192": 1 + }, + "attention_post_pre": { + "2048": 44 + }, + "ffn_post_pre": { + "2048": 45 + }, + "final_post": { + "2048": 1 + } + }, + "per_kernel_gpu_completion_verified": false, + "completed_requests_corroborate_dispatch": true + }, + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tokens_per_second": 48.463652260841044, + "mtp_accept_length": 2.7142857142857144, + "mtp_normalized_steps_per_second": 17.85502978030986, + "reported_server_steps_per_second": 17.886585061217623, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 20.015, + "measurement_wall_seconds": 20.015, + "client_output_tokens": 970, + "server_output_tokens": 970, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3599999999860302, + "latency": 20.60999999998603, + "input_tokens": 8192, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.3589999999967404, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 161, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -1.4275990253844184, + "mtp_normalized_steps_per_second_change_vs_neither_pct": -0.028582098482721197, + "reported_server_steps_per_second_change_vs_neither_pct": -0.030523325874143303 + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tokens_per_second": 121.17257303488826, + "mtp_accept_length": 2.6933333333333334, + "mtp_normalized_steps_per_second": 44.98981672087435, + "reported_server_steps_per_second": 44.98981672087435, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.922, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 2414, + "server_output_tokens": 2424, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3599999999860302, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 777, + "completed": false + }, + { + "ttft": 3.187999999994645, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 765, + "completed": false + }, + { + "ttft": 3.187999999994645, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 728, + "completed": false + }, + { + "ttft": 3.187999999994645, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 730, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -0.83033840171578, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 2.9451620449533866, + "reported_server_steps_per_second_change_vs_neither_pct": 2.9451620449533866 + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tokens_per_second": 47.31118127597817, + "mtp_accept_length": 2.684659090909091, + "mtp_normalized_steps_per_second": 17.622789216025733, + "reported_server_steps_per_second": 17.641457424941013, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.953, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 944, + "server_output_tokens": 944, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 21.280999999988126, + "input_tokens": 32768, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.422000000020489, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 135, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 2.439215362921221, + "mtp_normalized_steps_per_second_change_vs_neither_pct": -1.1331687573780669, + "reported_server_steps_per_second_change_vs_neither_pct": -1.2424295133894425 + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tokens_per_second": 123.89911929545656, + "mtp_accept_length": 2.7030567685589517, + "mtp_normalized_steps_per_second": 45.836669335475854, + "reported_server_steps_per_second": 45.83666933547585, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.984, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 2476, + "server_output_tokens": 2476, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 761, + "completed": false + }, + { + "ttft": 3.2339999999967404, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 734, + "completed": false + }, + { + "ttft": 3.2339999999967404, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 763, + "completed": false + }, + { + "ttft": 3.2339999999967404, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 775, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.37288668810162573, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 3.3762077535777646, + "reported_server_steps_per_second_change_vs_neither_pct": 3.3762077535777646 + } + ], + "artifact_sha256": { + "conditions": "c16558f03cd14b7978ae4a1ef3c30a743777a268c7957c4e31f0b19728db98d8", + "prefill": "dece9ad530835a6cb49a128995379156b688864bc343d3c012c71933506fb92e", + "activation": "a7776dc5d9ea45771af66b50d6bd7dba330f2b5633ad868fdca75c1eee2b1236", + "complete": "097d63814b64631fe9e7ddb226b85671a7b5cef4503613e4dde1251931b1b895", + "decode": "804885007b0ccef0ef113543b02800459f3354516be6b19d57dbf507b3c9148f", + "decode_receipt": "23938dd51cc74b63ffa659c7deb0ebc968b856810e068f3b6f9f4e212db50cf9" + }, + "protocol_sources": { + "prefill_harness_sha256": "6e4bf5ff62379bb2db4dc439d27a85eda0837959e983ef5f92ab3605898e69ac", + "activation_validator_sha256": "3359d4d9b5cf28b95c68a4989ac11cd8119dbeaf39058b9c1c27bcf95f218a09", + "decode_harness_sha256": "46ace1dad13c245807bc1b4ccf4ab6b90e95d12a99dd729ee24caa51034194c6", + "decode_wrapper_sha256": "5c2f7ecaef4c8e1ea60b7c44d3e39f710c660387fb2c0caa0871682f4ec712ce" + }, + "source_manifest_sha256": "082301aa98d9a4e0a84e52d32ff1c4efab29bd888c0fa8606182ba55a6a72e85", + "package_versions": { + "torch": "2.13.0+cu130", + "triton": "3.7.1", + "transformers": "5.16.1", + "nvidia-cutlass-dsl": "4.6.2", + "flashinfer-python": "0.6.17" + } + }, + { + "arm": "both", + "flags": { + "VLLM_B12X_KDA_PREFILL_COALESCING": "1", + "VLLM_GLM53_MHC_PREFILL_SHARD": "1" + }, + "started_at_utc": "2026-09-07T12:41:23.139314+00:00", + "prefill_samples": [ + { + "prompt_tokens": 8192, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 2.6545916999748442, + "tokens_per_second": 3085.973635824157 + }, + { + "prompt_tokens": 8192, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 2.6740572999988217, + "tokens_per_second": 3063.509521655953 + }, + { + "prompt_tokens": 8192, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 2.6768195999902673, + "tokens_per_second": 3060.348183355272 + }, + { + "prompt_tokens": 16384, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 5.342545599996811, + "tokens_per_second": 3066.7028841101105 + }, + { + "prompt_tokens": 16384, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 5.352795799990417, + "tokens_per_second": 3060.83037952416 + }, + { + "prompt_tokens": 16384, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 5.355028399993898, + "tokens_per_second": 3059.554268660586 + }, + { + "prompt_tokens": 32768, + "sample": 0, + "cached_tokens": 0, + "ttft_seconds": 10.720633199991425, + "tokens_per_second": 3056.5358770064263 + }, + { + "prompt_tokens": 32768, + "sample": 1, + "cached_tokens": 0, + "ttft_seconds": 10.74581809999654, + "tokens_per_second": 3049.3722948847003 + }, + { + "prompt_tokens": 32768, + "sample": 2, + "cached_tokens": 0, + "ttft_seconds": 10.745053999999072, + "tokens_per_second": 3049.5891411995535 + } + ], + "prefill_summary": [ + { + "prompt_tokens": 8192, + "samples": 3, + "median_ttft_seconds": 2.6740572999988217, + "min_ttft_seconds": 2.6545916999748442, + "max_ttft_seconds": 2.6768195999902673, + "tokens_per_second": 3063.509521655953, + "throughput_change_vs_neither_pct": 57.13504718162894, + "ttft_reduction_vs_neither_pct": 36.360473494870824 + }, + { + "prompt_tokens": 16384, + "samples": 3, + "median_ttft_seconds": 5.352795799990417, + "min_ttft_seconds": 5.342545599996811, + "max_ttft_seconds": 5.355028399993898, + "tokens_per_second": 3060.83037952416, + "throughput_change_vs_neither_pct": 29.75845818750851, + "ttft_reduction_vs_neither_pct": 22.933732878134094 + }, + { + "prompt_tokens": 32768, + "samples": 3, + "median_ttft_seconds": 10.745053999999072, + "min_ttft_seconds": 10.720633199991425, + "max_ttft_seconds": 10.74581809999654, + "tokens_per_second": 3049.5891411995535, + "throughput_change_vs_neither_pct": 16.639487339831007, + "ttft_reduction_vs_neither_pct": 14.265741147636902 + } + ], + "semantic_smoke_checks": [ + { + "kind": "cold", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "repeated", + "prompt_tokens": 8192, + "base_tokens": 8192, + "cached_tokens": 7168, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "extended", + "prompt_tokens": 8201, + "base_tokens": 8192, + "cached_tokens": 4096, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 16384, + "base_tokens": 16384, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + }, + { + "kind": "cold", + "prompt_tokens": 32768, + "base_tokens": 32768, + "cached_tokens": 0, + "exact_answer_passed": true, + "finish_reason": "stop" + } + ], + "activation": { + "checkpoint_enabled": true, + "checkpoint_scheduled": true, + "checkpoint_dispatch_shape": { + "tokens": 8192, + "interior_checkpoints": 4 + }, + "checkpoint_capacity": 4, + "configured_checkpoint_grid": [ + 512, + 512, + 2048 + ], + "mhc_enabled": true, + "mhc_collectives_per_dispatch": { + "reduce_scatter": 90, + "all_gather": 90, + "auxiliary_gathers": 0 + }, + "mhc_request_dispatches_by_rank": [ + 35, + 35, + 35, + 35 + ], + "mhc_common_dispatch_count": 35, + "mhc_rows": { + "first_pre": { + "8192": 1 + }, + "attention_post_pre": { + "2048": 44 + }, + "ffn_post_pre": { + "2048": 45 + }, + "final_post": { + "2048": 1 + } + }, + "per_kernel_gpu_completion_verified": false, + "completed_requests_corroborate_dispatch": true + }, + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tokens_per_second": 46.28702962370651, + "mtp_accept_length": 2.5865921787709496, + "mtp_normalized_steps_per_second": 17.894985534867097, + "reported_server_steps_per_second": 17.894985534867097, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.984, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 925, + "server_output_tokens": 926, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3590000000258442, + "latency": 21.90600000001723, + "input_tokens": 8192, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.4059999999881256, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 101, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -5.854729655227853, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 0.19513264663726382, + "reported_server_steps_per_second_change_vs_neither_pct": 0.01642755668271878 + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tokens_per_second": 119.36552554952085, + "mtp_accept_length": 2.7767441860465114, + "mtp_normalized_steps_per_second": 42.98758457813565, + "reported_server_steps_per_second": 42.98758457813565, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.922, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 2378, + "server_output_tokens": 2388, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.3599999999860302, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 722, + "completed": false + }, + { + "ttft": 3.1719999999913853, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 751, + "completed": false + }, + { + "ttft": 3.1719999999913853, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 646, + "completed": false + }, + { + "ttft": 3.1719999999913853, + "latency": 0.0, + "input_tokens": 8192, + "output_tokens": 697, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": -2.309256304732654, + "mtp_normalized_steps_per_second_change_vs_neither_pct": -1.636321664230389, + "reported_server_steps_per_second_change_vs_neither_pct": -1.636321664230389 + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tokens_per_second": 47.10834920317223, + "mtp_accept_length": 2.6818181818181817, + "mtp_normalized_steps_per_second": 17.565825126606597, + "reported_server_steps_per_second": 17.565825126606597, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 19.954, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 940, + "server_output_tokens": 944, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4059999999881256, + "latency": 21.655999999988126, + "input_tokens": 32768, + "output_tokens": 1024, + "completed": true + }, + { + "ttft": 1.422000000020489, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 127, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 2.0000388759197074, + "mtp_normalized_steps_per_second_change_vs_neither_pct": -1.4527469436943385, + "reported_server_steps_per_second_change_vs_neither_pct": -1.665822085437696 + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tokens_per_second": 124.25, + "mtp_accept_length": 2.761111111111111, + "mtp_normalized_steps_per_second": 45.0, + "reported_server_steps_per_second": 45.0, + "aggregate_source": "openai_continuous_usage", + "measurement_seconds": 20.0, + "measurement_wall_seconds": 20.0, + "client_output_tokens": 2485, + "server_output_tokens": 2485, + "num_errors": 0, + "request_samples": [ + { + "ttft": 1.4219999999913853, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 709, + "completed": false + }, + { + "ttft": 3.265999999974156, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 732, + "completed": false + }, + { + "ttft": 3.265999999974156, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 733, + "completed": false + }, + { + "ttft": 3.265999999974156, + "latency": 0.0, + "input_tokens": 32768, + "output_tokens": 761, + "completed": false + } + ], + "aggregate_tokens_per_second_change_vs_neither_pct": 0.6571414059595826, + "mtp_normalized_steps_per_second_change_vs_neither_pct": 1.4892533936924135, + "reported_server_steps_per_second_change_vs_neither_pct": 1.4892533936924135 + } + ], + "artifact_sha256": { + "conditions": "984fe189c370b3e4d49904d8dac1fffbee22694c627bc5e6e79b8112bae58c37", + "prefill": "4a79c47b00dd8d20768a728b75b2f4214275829be171d75ab996a8e388313b3e", + "activation": "f65297db47d7da94ff2ce93bcc0667dc71b317f55535f12dc9e2ba2a85644f67", + "complete": "ce055a3f82345da35462862d5157fa42d39d01cd019a6445f8f25048d4ff972c", + "decode": "62ccbd2cfbd7e3b8f15115fbff8c1a38f508d8045edcb474ff53d25d606127d5", + "decode_receipt": "31c66c9129073ba86c4f16769834bb34db0415626f353f55ad01c6b9392fa4b4" + }, + "protocol_sources": { + "prefill_harness_sha256": "6e4bf5ff62379bb2db4dc439d27a85eda0837959e983ef5f92ab3605898e69ac", + "activation_validator_sha256": "3359d4d9b5cf28b95c68a4989ac11cd8119dbeaf39058b9c1c27bcf95f218a09", + "decode_harness_sha256": "46ace1dad13c245807bc1b4ccf4ab6b90e95d12a99dd729ee24caa51034194c6", + "decode_wrapper_sha256": "5c2f7ecaef4c8e1ea60b7c44d3e39f710c660387fb2c0caa0871682f4ec712ce" + }, + "source_manifest_sha256": "082301aa98d9a4e0a84e52d32ff1c4efab29bd888c0fa8606182ba55a6a72e85", + "package_versions": { + "torch": "2.13.0+cu130", + "triton": "3.7.1", + "transformers": "5.16.1", + "nvidia-cutlass-dsl": "4.6.2", + "flashinfer-python": "0.6.17" + } + } + ], + "limitations": [ + "Three cold samples per prefill size; one 20-second decode observation per cell.", + "Sequential arm order and stochastic decode acceptance limit small-difference conclusions.", + "Two recovery reboots occurred between arms: one between Neither and Coalescing, and another between Coalescing and Sharded mHC. Each receipt records matching driver, GPU clock settings, persistence and CPU0 governor before and after its reboot. Reboots can change allocator/cache state and thermal conditions; this was not an interleaved A/B experiment.", + "Prefill throughput is prompt tokens divided by first-token latency, including request overhead.", + "MTP-normalized steps/s is aggregate tok/s divided by measured acceptance length; server steps/s is retained separately.", + "Dispatch logs plus completed requests corroborate asynchronous execution without per-kernel GPU completion events.", + "Exact-answer and cache-reuse tests are smoke coverage, not full model-quality or numerical-equivalence evaluation.", + "Diagnostic logging remains enabled in all arms and adds host work.", + "Deployment readiness/source checks and exclusive access are recorded attestations; public output excludes private deployment identifiers." + ] +} diff --git a/docs/benchmarking/glm-kda-checkpoints-20260907/manifest.json b/docs/benchmarking/glm-kda-checkpoints-20260907/manifest.json new file mode 100644 index 000000000000..64b1f7675fbe --- /dev/null +++ b/docs/benchmarking/glm-kda-checkpoints-20260907/manifest.json @@ -0,0 +1,32 @@ +{ + "schema": "glm-prefill-source-composition/v1", + "feature": "Continuation-prefill coalescing", + "feature_source_head": "1bf41ecb19ff82bd4bca1195cd1ab8f63e7e0717", + "vllm_base": "2a979314dc97b03173a0a76fc15664ec924db32b", + "composition": { + "image_source": "abb715f132bdccb592a34b2596a3d3a8d757ffbc", + "merged_tree": "b92b08bdd69e24d17ea1795b70d1736749a444f0", + "different_paths_from_image_source": [ + "docs/design/glm_mhc_prefill_sharding.md", + "docs/features/quantization/b12x.md" + ], + "vllm_subtree_in_both": "5f6b014449c94344bde776d086b0b895fad683ec", + "tests_subtree_in_both": "ee93b94ed702d40a8ff855987ff424f5b0210585", + "all_non_documentation_paths_identical": true, + "combined_checkout_unchanged": true, + "image_unchanged": true + }, + "files": { + "decode-benchmark.patch": "edaab1cef67ce511a93a67168444366db1f6b383e908f0e2bdda0f3ca8de5cb3", + "evidence.json": "20f4bab99a7c8bcd1ac564d8f62aca3c4628da014458be48829ba7ea65d7091b", + "README.md": "8564c8795985e7c01ce7f9588a464420667abf5c526fb4a3f6c0e8b0c2996067", + "report.md": "063a667f62b9b5da1de115ba2f55eeb1ac4f92ea08f52c3e2a8f09ba636ec3e8", + "reproduction.md": "5267816bedbe9a61656d246068beff0d9086a60c4b8db93bdb2e5e1244ffeb41", + "source-reproduction.md": "b4457f1149da788685bb5d8d35edd4fcd12118b4bb8e2e85c75916f9e4a4cc68" + }, + "standalone_prefill_harness": { + "path": "benchmarks/glm_prefill_checkpoints.py", + "sha256": "92d66779d42d9b50e1f432bb6e7247be97930a972e920ba8920246131a5f8f9e", + "activation_verified_by_client": false + } +} diff --git a/docs/benchmarking/glm-kda-checkpoints-20260907/report.md b/docs/benchmarking/glm-kda-checkpoints-20260907/report.md new file mode 100644 index 000000000000..da4dcc3afe1f --- /dev/null +++ b/docs/benchmarking/glm-kda-checkpoints-20260907/report.md @@ -0,0 +1,51 @@ +# GLM prefill: four configurations + +bounded semantic/cache smoke checks and performance observations; not full numerical/model-quality qualification. + +Four GB10 Sparks, TP4/DCP4, MTP3, fixed 8192-token budget. Runtime vLLM `abb715f132bdccb592a34b2596a3d3a8d757ffbc`, B12X `70fe41974ef4b18f61caaa2579c81cdc05d1265f`. + +Both native split-page settings are 512 in every arm. Diagnostics are enabled throughout. + +## Cold prefill + +Three samples per size; prompt tokens / median first-token latency. Positive percentages indicate higher throughput than Neither. + +| Prompt | Neither tok/s | Coalescing tok/s | Sharded mHC tok/s | Both tok/s | +| --- | ---: | ---: | ---: | ---: | +| 8192 | 1949.6 (+0.0%) | 2947.3 (+51.2%) | 1958.5 (+0.5%) | 3063.5 (+57.1%) | +| 16384 | 2358.9 (+0.0%) | 2928.3 (+24.1%) | 2390.8 (+1.4%) | 3060.8 (+29.8%) | +| 32768 | 2614.5 (+0.0%) | 2923.5 (+11.8%) | 2681.8 (+2.6%) | 3049.6 (+16.6%) | + +## Decode + +One 20-second observation per cell. Each entry is aggregate tok/s / acceptance-normalized steps/s. + +| Context / concurrency | Neither | Coalescing | Sharded mHC | Both | +| --- | ---: | ---: | ---: | ---: | +| 8192 / C1 | 49.17 / 17.86 | 47.37 / 17.91 | 48.46 / 17.86 | 46.29 / 17.89 | +| 8192 / C4 | 122.19 / 43.70 | 120.18 / 44.54 | 121.17 / 44.99 | 119.37 / 42.99 | +| 32768 / C1 | 46.18 / 17.82 | 45.92 / 17.73 | 47.31 / 17.62 | 47.11 / 17.57 | +| 32768 / C4 | 123.44 / 44.34 | 125.20 / 45.64 | 123.90 / 45.84 | 124.25 / 45.00 | + +## Execution and correctness evidence + +| Configuration | Exact-answer/cache checks | Four-checkpoint dispatch | mHC common dispatches across four ranks | +| --- | ---: | --- | ---: | +| Neither | 5/5 passed | Disabled | 0 | +| Coalescing | 5/5 passed | Observed | 0 | +| Sharded mHC | 5/5 passed | Disabled | 20 | +| Both | 5/5 passed | Observed | 35 | + +The JSON contains every prefill sample, sanitized decode request measurements, source/artifact hashes and deltas. No prompts, answers, private addresses or container identifiers are included. + +## Limits + +- Three cold samples per prefill size; one 20-second decode observation per cell. +- Sequential arm order and stochastic decode acceptance limit small-difference conclusions. +- Two recovery reboots occurred between arms: one between Neither and Coalescing, and another between Coalescing and Sharded mHC. Each receipt records matching driver, GPU clock settings, persistence and CPU0 governor before and after its reboot. Reboots can change allocator/cache state and thermal conditions; this was not an interleaved A/B experiment. +- Prefill throughput is prompt tokens divided by first-token latency, including request overhead. +- MTP-normalized steps/s is aggregate tok/s divided by measured acceptance length; server steps/s is retained separately. +- Dispatch logs plus completed requests corroborate asynchronous execution without per-kernel GPU completion events. +- Exact-answer and cache-reuse tests are smoke coverage, not full model-quality or numerical-equivalence evaluation. +- Diagnostic logging remains enabled in all arms and adds host work. +- Deployment readiness/source checks and exclusive access are recorded attestations; public output excludes private deployment identifiers. diff --git a/docs/benchmarking/glm-kda-checkpoints-20260907/reproduction.md b/docs/benchmarking/glm-kda-checkpoints-20260907/reproduction.md new file mode 100644 index 000000000000..bb16e474caaa --- /dev/null +++ b/docs/benchmarking/glm-kda-checkpoints-20260907/reproduction.md @@ -0,0 +1,95 @@ +# Reproducing the serving measurements + +Status: **research-only**. The results cover bounded correctness/cache probes +and short performance observations. Use an otherwise working four-GB10 TP4/DCP4 +GLM deployment. [Source reconstruction](source-reproduction.md) identifies public +commits that reproduce the measured runtime; it does not supply the private image. + +## Model and tokenizer + +Weights: [local-inference-lab/GLM-5.3-Flash-NVFP4-Spark](https://huggingface.co/local-inference-lab/GLM-5.3-Flash-NVFP4-Spark/tree/df116c4fb16b1d37ae43d2cfd624de26ffbc832e), +revision `df116c4fb16b1d37ae43d2cfd624de26ffbc832e`. All 16 candidate +container specifications mounted that snapshot. No separate tokenizer override +was configured; the tokenizer is supplied by the model snapshot. + +Use BF16 activations, FP8 KV cache, TP4/DCP4/PP1, native MTP with three +speculative tokens, 24 GiB KV budget per rank, 16 maximum sequences, +1,048,576 maximum model tokens and 8,192 maximum batched tokens. Enable native +prefix caching with aligned recurrent checkpoints and retention interval zero. +Keep both split-page settings at 512 on every worker: +`VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE=512` and +`VLLM_GLM53_SPLIT_MAMBA_BLOCK_SIZE=512`. Keep +`VLLM_GLM53_MHC_PREFILL_DIAGNOSTICS=1` in every comparison configuration. + +| Configuration | VLLM_B12X_KDA_PREFILL_COALESCING | VLLM_GLM53_MHC_PREFILL_SHARD | +| --- | ---: | ---: | +| Neither | 0 | 0 | +| Coalescing only | 1 | 0 | +| mHC only | 0 | 1 | +| Both | 1 | 1 | + +The measured deployment held fused SparkRing transport constant. A switch-based +cluster can test the same feature logic with its own working TP communicator, +but changing transport does not reproduce the absolute recorded timings. +Complete automatic model warmup and reserve exclusive inference access before +running either benchmark. Serving preparation is outside these client scripts. + +## Decode + +The measured client was a local merge of the public benchmark. Reconstruct its +exact bytes from the public parent and the adjacent patch: + +```bash +git clone https://github.com/local-inference-lab/llm-inference-bench.git +git -C llm-inference-bench checkout bd88816e9e7bcc97e1bcfd954c3053528f31af69 +# Run from this evidence directory, or substitute its absolute path. +git -C llm-inference-bench apply "$PWD/decode-benchmark.patch" +sha256sum llm-inference-bench/llm_decode_bench.py +``` + +Expected SHA-256: +`46ace1dad13c245807bc1b4ccf4ab6b90e95d12a99dd729ee24caa51034194c6`. +The patch contains Windows portability and display/hardware-monitor integration +used by the measured client. Hardware monitoring was disabled for these tests. + +```bash +uv venv .benchmark-venv +uv pip install --python .benchmark-venv/bin/python httpx==0.28.1 rich==15.0.0 +.benchmark-venv/bin/python llm-inference-bench/llm_decode_bench.py \ + --host http://SERVER --port PORT --model SERVED_MODEL_ID \ + --dcp-size 4 --temperature 1.0 --token-targeting exact \ + --kv-budget 524288 --display-mode plain --no-hw-monitor --skip-prefill \ + --concurrency 1,4 --contexts 8k,32k --max-tokens 1024 --duration 20 \ + --decode-warmup-seconds 5 --cell-warmup-timeout-seconds 180 \ + --output decode-results.json +``` + +Replace SERVER, PORT and SERVED_MODEL_ID with the deployment endpoint. Repeat +for each flag combination. The benchmark generates run-specific padding and +uses the server tokenizer for exact context sizes. Stochastic outputs need not +match token-for-token. Report raw output tok/s, mean MTP acceptance length, +and their quotient; retain reported server step rates separately. + +## Prefill and cache checks + +Run the standalone client from the vLLM checkout after model warmup: + +```bash +.venv/bin/python benchmarks/glm_prefill_checkpoints.py \ + --base-url http://SERVER:PORT --model SERVED_MODEL_ID \ + --label coalescing-only --output prefill-results.json \ + --ready-confirmed --exclusive-window +``` + +The client preserves the measured protocol's prompt generator, exact tokenizer +calibration, zero-cache checks, streaming first-token timing and phase order: +three excluded shape warmups, five exact-answer/cache checks, then three cold +samples each at 8K/16K/32K. The synthetic record queries are generated in the +script; no private prompt corpus is needed. Repeated and extended prompts test +cache reuse. Sampling and token limits are encoded in the client. + +The standalone client records feature activation as **not verified**. The +recorded four-configuration experiment separately gated timing on request-bound +all-rank checkpoint/mHC diagnostics. Collect those diagnostics as described in +the feature documentation before attributing a reproduced rate to either +optimization. The client neither accesses SSH nor starts/stops the model. diff --git a/docs/benchmarking/glm-kda-checkpoints-20260907/source-reproduction.md b/docs/benchmarking/glm-kda-checkpoints-20260907/source-reproduction.md new file mode 100644 index 000000000000..4f5dbd38edf9 --- /dev/null +++ b/docs/benchmarking/glm-kda-checkpoints-20260907/source-reproduction.md @@ -0,0 +1,151 @@ +# Reconstructing the measured prefill source from public forks + +The experiment records vLLM `abb715f132bdccb592a34b2596a3d3a8d757ffbc` as its +historical build-source identifier. That commit is not retrievable from the +public fork. Reconstruct the equivalent runtime from the two published feature +commits below; no private Git repository, deployment snapshot, or image is +needed to obtain this source. + +| Component | Public fork branch | Frozen public commit | +| --- | --- | --- | +| vLLM continuation | `FujitsuPolycom/vllm:feat/gb10-continuation-prefill` | [`8b5a65e23c43695565802fd5a9408322ab300294`](https://github.com/FujitsuPolycom/vllm/commit/8b5a65e23c43695565802fd5a9408322ab300294) | +| vLLM mHC | `FujitsuPolycom/vllm:feat/gb10-mhc-token-sharding` | [`7a43439c18f6c1bde921762e9ed0d40c6334909c`](https://github.com/FujitsuPolycom/vllm/commit/7a43439c18f6c1bde921762e9ed0d40c6334909c) | +| B12X publication | `FujitsuPolycom/b12x:feat/gb10-kda-checkpoint-export` | [`b8044279b346095761d45f3a5db59afbfc2f4929`](https://github.com/FujitsuPolycom/b12x/commit/b8044279b346095761d45f3a5db59afbfc2f4929) | +| B12X tested implementation | Ancestor of the published B12X commit | [`70fe41974ef4b18f61caaa2579c81cdc05d1265f`](https://github.com/FujitsuPolycom/b12x/commit/70fe41974ef4b18f61caaa2579c81cdc05d1265f) | + +These public commit endpoints and branch mappings were verified on +2026-09-07. Pin the commits rather than moving branch tips. The vLLM commits +share base `2a979314dc97b03173a0a76fc15664ec924db32b`; B12X's review base is +`06b4de7c723e6f166d65abf5909c5b7d0f8acc68`. + +## Obtain and verify the source + +Run the following Bash commands from a directory where `vllm-source` and +`b12x-source` do not exist. Git must support `merge-tree --write-tree`. +The merge is left uncommitted; it does not push or alter either public branch. + +```bash +set -euo pipefail + +CP_SOURCE=8b5a65e23c43695565802fd5a9408322ab300294 +MHC_SOURCE=7a43439c18f6c1bde921762e9ed0d40c6334909c +B12X_PUBLIC=b8044279b346095761d45f3a5db59afbfc2f4929 +B12X_TESTED=70fe41974ef4b18f61caaa2579c81cdc05d1265f + +git clone --filter=blob:none --no-checkout \ + https://github.com/FujitsuPolycom/vllm.git vllm-source +git -C vllm-source fetch --no-tags origin "$CP_SOURCE" "$MHC_SOURCE" + +SOURCE_TREE=$(git -C vllm-source merge-tree --write-tree \ + "$CP_SOURCE" "$MHC_SOURCE") +test "$SOURCE_TREE" = 3769f7f9754421b2f330031255e33a446af34b1a +test "$(git -C vllm-source rev-parse "$SOURCE_TREE:vllm")" = \ + 5f6b014449c94344bde776d086b0b895fad683ec +test "$(git -C vllm-source rev-parse "$SOURCE_TREE:tests")" = \ + ee93b94ed702d40a8ff855987ff424f5b0210585 + +git -C vllm-source checkout --detach "$CP_SOURCE" +git -C vllm-source merge --no-commit --no-ff "$MHC_SOURCE" +test "$(git -C vllm-source write-tree)" = "$SOURCE_TREE" +git -C vllm-source diff --exit-code + +git clone --filter=blob:none --no-checkout \ + https://github.com/FujitsuPolycom/b12x.git b12x-source +git -C b12x-source fetch --no-tags origin "$B12X_PUBLIC" +git -C b12x-source merge-base --is-ancestor "$B12X_TESTED" "$B12X_PUBLIC" +git -C b12x-source checkout --detach "$B12X_TESTED" +test "$(git -C b12x-source rev-parse HEAD:b12x)" = \ + be623cf8abea2e5dc6771232246edf7ac0a39712 +test "$(git -C b12x-source rev-parse HEAD:tests)" = \ + 4db2bd8301af46711aa408b11f7fd0c95629a7ee +git -C b12x-source diff --exit-code +``` + +The reconstructed vLLM `vllm/` and `tests/` tree hashes exactly match the +historical `abb715f1` build source. A local comparison also found no differences +outside `docs/`. The complete reconstructed tree differs because the public +commits include additional documentation and evidence. + +B12X `b8044279` updates the module docstring in +`b12x/sequence/kda_prefill/__init__.py` as well as publishing evidence. Its +`b12x/` tree is therefore not byte-identical to the tested source, although the +checkpoint implementation is unchanged. Checking out its public `70fe4197` +ancestor, as above, recovers the exact tested package tree. Its tests are +identical at both commits. + +`HEAD` in `vllm-source` still names the CP commit while the merged files are +staged. The authoritative source identity for this uncommitted reconstruction +is `SOURCE_TREE` and the verified subtree hashes, not `git rev-parse HEAD`. +Keep both source directories after an editable installation. + +## Integrate into a Linux CUDA build + +Source reconstruction and hash equality were verified. A new Linux/aarch64 +clean build was **not** executed as part of this reconstruction check, and no +public base-image recipe has been qualified to reproduce image digest +`52b207e716a285c16e5e1b14ec2a41f6208b9450c617e7d1cde5a507ea879d7f`. +The following is the pinned repository's documented existing-PyTorch build +procedure, applied to the reconstructed source; a successful rebuild receives +its own binary and image identities. + +Use a Linux/aarch64 CUDA development environment appropriate for GB10, with +Python3.12, a working CUDA13.0 PyTorch installation, CUDA development tools, +and the C++ compiler required by the pinned vLLM source. The measured runtime +used PyTorch `2.13.0+cu130`, Triton `3.7.1`, CUTLASS DSL `4.6.2`, and FlashInfer +`0.6.17`. Provisioning that platform-compatible PyTorch/toolchain installation +is a prerequisite, not a verified download step in this note. Follow the +[pinned CUDA build instructions](https://github.com/FujitsuPolycom/vllm/blob/8b5a65e23c43695565802fd5a9408322ab300294/docs/getting_started/installation/gpu.cuda.inc.md) +for the native build prerequisites. + +From the parent directory containing both source checkouts, where the chosen +Python3.12 interpreter already has the required PyTorch available: + +```bash +set -euo pipefail +cd vllm-source +uv venv --python 3.12 --system-site-packages .venv +uv run --no-project -- .venv/bin/python -c \ + 'import torch; assert str(torch.__version__) == "2.13.0+cu130"; assert torch.version.cuda == "13.0"' + +# This repository helper adjusts dependency declarations for the installed torch. +# It does not change the verified vllm/ or tests/ source trees. +uv run --no-project -- .venv/bin/python use_existing_torch.py +uv pip install --python .venv/bin/python -r requirements/build/cuda.txt +uv pip install --python .venv/bin/python --no-build-isolation --editable . +uv pip install --python .venv/bin/python --editable ../b12x-source +uv pip check --python .venv/bin/python + +uv run --no-project -- .venv/bin/python - <<'PY' +from pathlib import Path +from importlib.metadata import version +import torch +import vllm +import vllm._C +import b12x +from b12x.sequence.kda_prefill import Caps + +root = Path.cwd().resolve() +assert Path(vllm.__file__).resolve().is_relative_to(root / "vllm") +assert Path(b12x.__file__).resolve().is_relative_to(root.parent / "b12x-source" / "b12x") +assert "max_checkpoints" in Caps.__dataclass_fields__ +assert str(torch.__version__) == "2.13.0+cu130" +assert torch.version.cuda == "13.0" +for package in ("triton", "nvidia-cutlass-dsl", "flashinfer-python"): + print(package, version(package)) +print("vllm:", vllm.__file__) +print("b12x:", b12x.__file__) +PY +``` + +Record the resolved dependency versions and native artifact hashes. If they +differ from the measured toolchain, describe the run as a new build rather +than an identical binary reproduction. Do not replace the source build with +an arbitrary upstream-main precompiled wheel and assume its native ABI matches +the maintained GLM source. B12X's CuTe kernels compile on first use; installing +the package or importing `Caps` does not qualify GPU execution. + +Run the committed CPU and GB10 component checks before evaluating a rebuilt +serving stack. Source equality and import success do not establish native ABI, +GPU numerical, graph-replay, model-quality, or performance equivalence. The +model checkpoint/tokenizer pin and benchmark workload remain separate inputs; +this document reconstructs and integrates the source only. diff --git a/docs/benchmarking/glm-tp4-dcp-prefill/README.md b/docs/benchmarking/glm-tp4-dcp-prefill/README.md new file mode 100644 index 000000000000..6d7fb234f0a2 --- /dev/null +++ b/docs/benchmarking/glm-tp4-dcp-prefill/README.md @@ -0,0 +1,102 @@ +# TP4 GLM prefill at DCP1 and DCP2 + +Status: **qualified** for the bounded combined-runtime checks described below; **research-only** for broader model accuracy and DCP2 cache-invariant probabilities. + +## Recorded serving composition + +These measurements enable or disable **both continuation coalescing and token-sharded mHC together**. They establish no isolated speedup for either PR. The standalone coalescing implementation at vLLM `a6c8407645cf5e751711883f24e8822e73dba9a0` and mHC implementation at vLLM `138ffcdbc97141ef4600d33ae12077a7b5bf926a` have CPU coverage; neither complete standalone revision has GPU or full-model validation. + +- Model: `local-inference-lab/GLM-5.3-Flash-NVFP4-Spark`, revision `df116c4fb16b1d37ae43d2cfd624de26ffbc832e`. +- Image: `sha256:e585e8b3ebeeea2852011d319f4a4d781bf2ae33388210b7cd16ff16d3765d9b`. +- vLLM: `8f8ea47be212bbdd91b2172d5958ea2aae2b0e50`; B12X: `0b6d61c37c87ae49d2f9d20d38b9da023146e243`. +- Four GB10 Sparks, TP4, MTP3, BF16 activations, FP8 KV, B12X attention/KDA/MoE/linear, 8192-token scheduler budget, explicit 512-token target and recurrent cache grids, aligned checkpoint policy, prefix retention interval 0, 24 GiB KV per rank, max 16 sequences. +- `FULL_AND_PIECEWISE` graphs with capture sizes 4,8,...,64; async scheduling; prefill schedule interval 2. DCP1 uses cache interleave 1; DCP2 uses 4. These choices are fixed within each on/off pair. +- TP4 mesh and dual-domain NCCL are retained. NCCL library SHA-256: `768a450b5eb84bf3d1191795350e43c96de75aeba4783ec314d47672fe6e1fc6`. CKV gathering uses the DCP2 group and bypasses at DCP1. SparkCache, compact index cache and DCP4-only top-k owner exchange/fused endpoints are disabled. +- Both flags are 0 in controls and 1 in enabled arms: `VLLM_B12X_KDA_PREFILL_COALESCING`, `VLLM_GLM53_MHC_PREFILL_SHARD`. Both native geometry overrides are 512: `VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE`, `VLLM_GLM53_SPLIT_MAMBA_BLOCK_SIZE`. + +## Prefill + +Median of three cold exact-token requests per size after excluded warmups and semantic checks; each allows one generated token. Tokens/s is prompt tokens divided by TTFT. Arms ran sequentially, not interleaved. + +| DCP | Coalescing + mHC | 8K tok/s | 16K tok/s | 32K tok/s | +| --- | --- | ---: | ---: | ---: | +| 1 | Off | 2,263.0 | 2,615.7 | 2,828.4 | +| 1 | On | 3,616.3 | 3,615.1 | 3,586.1 | +| 2 | Off | 1,938.7 | 2,345.5 | 2,619.0 | +| 2 | On | 3,493.9 | 3,461.8 | 3,445.9 | + +## MTP-normalized decode + +Aggregate emitted tokens/s divided by observed speculative acceptance length. Each cell is one approximately 10-second observation at temperature 1 with a 2048-token output cap. These short observations do not demonstrate repeatable decode gains. At C4, the value aggregates request decode steps rather than counting GPU launches. + +| DCP | Coalescing + mHC | Context | C1 steps/s | C4 aggregate steps/s | +| --- | --- | ---: | ---: | ---: | +| 1 | Off | 8,192 | 21.17 | 47.93 | +| 1 | Off | 32,768 | 21.23 | 46.21 | +| 1 | On | 8,192 | 21.40 | 48.00 | +| 1 | On | 32,768 | 21.17 | 49.93 | +| 2 | Off | 8,192 | 19.73 | 45.41 | +| 2 | Off | 32,768 | 19.86 | 45.01 | +| 2 | On | 8,192 | 20.05 | 46.62 | +| 2 | On | 32,768 | 19.80 | 44.47 | + +## Correctness and numerical limits + +The measured composition passed 18 GB10 recurrent-export, continuation/reload and convolution-history component tests, with zero skips. Each serving arm has six exact-answer cases, nine cold prefill timings and four decode cells; four-rank logs prove enabled paths executed and disabled paths did not. Those component tests do not validate distributed mHC intermediate tensors or broad model accuracy. + +Fixed-input comparisons require identical calibrated prompt token IDs, generation options and output token sequences. The preselected generated-token logprob tolerance is `abs(on-off) <= 0.1 + 0.02*abs(off)`. All six matched on/off comparisons pass. DCP1 also passes cached-versus-cold comparisons. + +**DCP2 cached-versus-cold behavior remains research-only.** With both patches enabled, the maximum generated-token logprob difference is 0.127015 and exceeds that comparison's bound; the disabled comparison is 0.018938. DCP2's enabled/disabled extended-cold maximum difference is 0.100147, narrowly within its approximately 0.100867 bound. Exact 13-token answers match, but this does not prove the numerical difference harmless for other prompts. Separate on/off passes do not establish cache invariance, and the enabled cached-versus-cold fixture has status `needs-review`. + +The DCP2 enabled 8K/C1 decode cell records 575 server and 571 client output tokens, within the benchmark's accepted boundary tolerance. Small normalized differences should not be described as speedups. Mixed-request scheduling, preemption, every checkpoint boundary, long-context capacity, four-rank cache restoration and broad accuracy are outside this performance matrix. The DCP2 disabled numerical fixture uses a separate correctness-only deployment requiring a node reboot. The prefill/decode tables use the four performance deployments identified in the measurement artifact. + +AI assistance was used for implementation, testing and evidence preparation. + +## Checkpoint and collective contracts + +With TP4 and 512-token physical cache blocks, the scheduler grids are 512, +1024 and 2048 tokens for DCP1, DCP2 and DCP4 respectively. Retention requires +two, three and four active checkpoint destinations respectively. B12X keeps +planned capacity four; three active destinations do not require capacity three. +The checkpoint kernel receives local head geometry, packed ranges and destinations; +the vLLM adapter determines distributed cache boundaries. + +Token-sharded mHC always uses the four-rank TP communicator. Each rank owns +2048 rows of an eligible 8192-token prefill, independently of DCP group size. +The CKV attention gather uses the two-rank DCP subgroup at DCP2 and bypasses +at DCP1. These group identities do not change mHC token ownership. + +## Standalone CPU validation + +Coalescing revision `a6c8407645cf5e751711883f24e8822e73dba9a0` passes 62 +checkpoint/allocator and worker-binding cases. Its scheduler configuration +interface matches the PR base `2a979314dc97b03173a0a76fc15664ec924db32b`. +It requires the four-checkpoint B12X API and contains no mHC implementation. + +mHC revision `138ffcdbc97141ef4600d33ae12077a7b5bf926a` passes 62 cases, +including all four TP ranks with DCP1/2/4. Tests use B12X +`06b4de7c723e6f166d65abf5909c5b7d0f8acc68`, whose capacity type lacks +`max_checkpoints`, demonstrating independence from the checkpoint extension. + +Both suites run on Windows with Python 3.12 and Torch 2.13.0+cpu. The launcher +maps `uvloop` to `winloop`; GPU operations in the suites use CPU doubles. +Commands executed from the corresponding standalone checkout: + +```powershell +# Continuation coalescing +.venv/Scripts/python.exe -c "import sys,winloop; sys.modules['uvloop']=winloop; from b12x.sequence.kda_prefill._impl import Caps; assert 'max_checkpoints' in Caps.__dataclass_fields__; import pytest; raise SystemExit(pytest.main(['tests/v1/core/test_recurrent_prefill_checkpoint.py','tests/v1/worker/test_kda_prefill_checkpoint_binding.py','--confcutdir=tests/v1','-q','--tb=short']))" + +# Token-sharded mHC +.venv/Scripts/python.exe -c "import sys,winloop; sys.modules['uvloop']=winloop; from b12x.sequence.kda_prefill._impl import Caps; assert 'max_checkpoints' not in Caps.__dataclass_fields__; import pytest; raise SystemExit(pytest.main(['tests/models/test_glm5next_mhc_prefill.py','--confcutdir=tests/models','-q','--tb=short']))" +``` + +## Measurement artifact + +[Serving evidence](evidence.json), schema `dcp-prefill-review-evidence/v1`, +contains numeric samples, semantic serving arguments, runtime identities and +source receipt hashes. Configuration labels `dcp1-off`, `dcp1-on`, `dcp2-off` +and `dcp2-on` identify the DCP size and whether both prefill features are disabled +or enabled. Hashes identify source receipts; the artifact does not include the +complete private deployment receipts or request-associated worker logs. +This artifact supports inspection of the reported measurements; it is not a +complete image-build or cluster-deployment recipe. diff --git a/docs/benchmarking/glm-tp4-dcp-prefill/evidence.json b/docs/benchmarking/glm-tp4-dcp-prefill/evidence.json new file mode 100644 index 000000000000..1b858c0c7025 --- /dev/null +++ b/docs/benchmarking/glm-tp4-dcp-prefill/evidence.json @@ -0,0 +1,705 @@ +{ + "schema": "dcp-prefill-review-evidence/v1", + "scope": "Combined coalescing and mHC on/off in measured serving composition; not standalone PR-head GPU qualification", + "image": "sha256:e585e8b3ebeeea2852011d319f4a4d781bf2ae33388210b7cd16ff16d3765d9b", + "vllm": "8f8ea47be212bbdd91b2172d5958ea2aae2b0e50", + "b12x": "0b6d61c37c87ae49d2f9d20d38b9da023146e243", + "model": { + "repository": "local-inference-lab/GLM-5.3-Flash-NVFP4-Spark", + "revision": "df116c4fb16b1d37ae43d2cfd624de26ffbc832e" + }, + "configurations": { + "1": { + "semantic_serving_arguments": [ + "--tensor-parallel-size", + "4", + "--pipeline-parallel-size", + "1", + "--decode-context-parallel-size", + "1", + "--cp-kv-cache-interleave-size", + "1", + "--distributed-executor-backend", + "mp", + "--nnodes", + "4", + "--disable-custom-all-reduce", + "--mamba-cache-mode", + "align", + "--limit-mm-per-prompt", + "{\"image\":4,\"video\":1}", + "--enable-chunked-prefill", + "--dtype", + "bfloat16", + "--kv-cache-dtype", + "fp8", + "--quantization", + "modelopt_mixed", + "--attention-backend", + "B12X", + "--block-size", + "512", + "--mamba-block-size", + "512", + "--recurrent-checkpoint-policy", + "aligned", + "--prefix-cache-retention-interval", + "0", + "--moe-backend", + "b12x", + "--linear-backend", + "b12x", + "--no-enable-flashinfer-autotune", + "--load-format", + "fastsafetensors", + "--enable-auto-tool-choice", + "--tool-call-parser", + "glm47", + "--reasoning-parser", + "glm45", + "--kda-prefill-backend", + "b12x", + "--gpu-memory-utilization", + "0.80", + "--kv-cache-memory-bytes", + "25769803776", + "--max-model-len", + "1048576", + "--max-num-seqs", + "16", + "--max-num-batched-tokens", + "8192", + "--prefill-schedule-interval", + "2", + "--speculative-config", + "{\"method\":\"mtp\",\"num_speculative_tokens\":3,\"draft_tensor_parallel_size\":4,\"kv_cache_dtype\":\"auto\",\"draft_sample_method\":\"probabilistic\",\"rejection_sample_method\":\"standard\",\"draft_load_config\":{\"load_format\":\"safetensors\"},\"attention_backend\":\"B12X\"}", + "--compilation-config", + "{\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"cudagraph_capture_sizes\":[4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64],\"custom_ops\":[\"all\"],\"pass_config\":{\"fuse_allreduce_rms\":false}}", + "--max-cudagraph-capture-size", + "64", + "--async-scheduling", + "--enable-prefix-caching", + "--cudagraph-metrics", + "--enable-prompt-tokens-details" + ], + "feature_environment": { + "VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE": "512", + "VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS": "524288", + "VLLM_GLM53_SPLIT_MAMBA_BLOCK_SIZE": "512", + "NCCL_MIN_NCHANNELS": "4", + "VLLM_B12X_MLA_CKV_GATHER": "1", + "NCCL_MAX_NCHANNELS": "4", + "NCCL_IB_HCA": "=rocep1s0f0:1,rocep1s0f1:1,roceP2p1s0f0:1,roceP2p1s0f1:1", + "VLLM_B12X_KDA_PREFILL_COALESCING": "1", + "VLLM_GLM53_MHC_PREFILL_DIAGNOSTICS": "1", + "VLLM_GLM53_MHC_PREFILL_SHARD": "1", + "VLLM_DCP_OWNER_FUSED_ENDPOINTS": "0", + "VLLM_DCP_TOPK_OWNER_MERGE": "0", + "VLLM_DCP_COMPACT_INDEX_CACHE_OWNER": "0" + }, + "nccl_sha256": "768a450b5eb84bf3d1191795350e43c96de75aeba4783ec314d47672fe6e1fc6" + }, + "2": { + "semantic_serving_arguments": [ + "--tensor-parallel-size", + "4", + "--pipeline-parallel-size", + "1", + "--decode-context-parallel-size", + "2", + "--cp-kv-cache-interleave-size", + "4", + "--distributed-executor-backend", + "mp", + "--nnodes", + "4", + "--disable-custom-all-reduce", + "--mamba-cache-mode", + "align", + "--limit-mm-per-prompt", + "{\"image\":4,\"video\":1}", + "--enable-chunked-prefill", + "--dtype", + "bfloat16", + "--kv-cache-dtype", + "fp8", + "--quantization", + "modelopt_mixed", + "--attention-backend", + "B12X", + "--block-size", + "512", + "--mamba-block-size", + "512", + "--recurrent-checkpoint-policy", + "aligned", + "--prefix-cache-retention-interval", + "0", + "--moe-backend", + "b12x", + "--linear-backend", + "b12x", + "--no-enable-flashinfer-autotune", + "--load-format", + "fastsafetensors", + "--enable-auto-tool-choice", + "--tool-call-parser", + "glm47", + "--reasoning-parser", + "glm45", + "--kda-prefill-backend", + "b12x", + "--gpu-memory-utilization", + "0.80", + "--kv-cache-memory-bytes", + "25769803776", + "--max-model-len", + "1048576", + "--max-num-seqs", + "16", + "--max-num-batched-tokens", + "8192", + "--prefill-schedule-interval", + "2", + "--speculative-config", + "{\"method\":\"mtp\",\"num_speculative_tokens\":3,\"draft_tensor_parallel_size\":4,\"kv_cache_dtype\":\"auto\",\"draft_sample_method\":\"probabilistic\",\"rejection_sample_method\":\"standard\",\"draft_load_config\":{\"load_format\":\"safetensors\"},\"attention_backend\":\"B12X\"}", + "--compilation-config", + "{\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"cudagraph_capture_sizes\":[4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64],\"custom_ops\":[\"all\"],\"pass_config\":{\"fuse_allreduce_rms\":false}}", + "--max-cudagraph-capture-size", + "64", + "--async-scheduling", + "--enable-prefix-caching", + "--cudagraph-metrics", + "--enable-prompt-tokens-details" + ], + "feature_environment": { + "VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE": "512", + "VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS": "524288", + "VLLM_GLM53_SPLIT_MAMBA_BLOCK_SIZE": "512", + "NCCL_MIN_NCHANNELS": "4", + "VLLM_B12X_MLA_CKV_GATHER": "1", + "NCCL_MAX_NCHANNELS": "4", + "NCCL_IB_HCA": "=rocep1s0f0:1,rocep1s0f1:1,roceP2p1s0f0:1,roceP2p1s0f1:1", + "VLLM_B12X_KDA_PREFILL_COALESCING": "1", + "VLLM_GLM53_MHC_PREFILL_DIAGNOSTICS": "1", + "VLLM_GLM53_MHC_PREFILL_SHARD": "1", + "VLLM_DCP_OWNER_FUSED_ENDPOINTS": "0", + "VLLM_DCP_TOPK_OWNER_MERGE": "0", + "VLLM_DCP_COMPACT_INDEX_CACHE_OWNER": "0" + }, + "nccl_sha256": "768a450b5eb84bf3d1191795350e43c96de75aeba4783ec314d47672fe6e1fc6" + } + }, + "rows": [ + { + "dcp": 2, + "enabled": false, + "prefill": [ + { + "tokens": 8192, + "samples": 3, + "median_ttft_seconds": 4.225529100003769, + "min_ttft_seconds": 4.215828300002613, + "max_ttft_seconds": 4.231428599989158, + "tokens_per_second": 1938.6921273344665 + }, + { + "tokens": 16384, + "samples": 3, + "median_ttft_seconds": 6.985251599995536, + "min_ttft_seconds": 6.965220199999749, + "max_ttft_seconds": 6.9966285999980755, + "tokens_per_second": 2345.5132238916735 + }, + { + "tokens": 32768, + "samples": 3, + "median_ttft_seconds": 12.511659099996905, + "min_ttft_seconds": 12.46474899999157, + "max_ttft_seconds": 12.524254100004327, + "tokens_per_second": 2618.9971879914874 + } + ], + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tps": 56.39022435899277, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.8578680203045685, + "server_steps_per_s": 19.731570512826956, + "server_output_tokens": 563, + "client_output_tokens": 563, + "server_spec_drafts": 197, + "server_spec_draft_tokens": 591, + "server_spec_accepted_tokens": 366, + "server_accept_len_effective": 2.8578680203045685, + "server_engine_steps": 197.0, + "aggregate_div_spec_accept_length": 19.731570512826956, + "effective_normalized_steps_per_s": 19.731570512826956, + "measurement_seconds": 9.984, + "measurement_wall_seconds": 10.015, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tps": 53.96729862567551, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.717171717171717, + "server_steps_per_s": 19.861570869672402, + "server_output_tokens": 538, + "client_output_tokens": 538, + "server_spec_drafts": 198, + "server_spec_draft_tokens": 594, + "server_spec_accepted_tokens": 340, + "server_accept_len_effective": 2.717171717171717, + "server_engine_steps": 198.0, + "aggregate_div_spec_accept_length": 19.861570869672402, + "effective_normalized_steps_per_s": 19.861570869672402, + "measurement_seconds": 9.969, + "measurement_wall_seconds": 10.016, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tps": 129.0063297497125, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.84070796460177, + "server_steps_per_s": 45.4134431829206, + "server_output_tokens": 1284, + "client_output_tokens": 1284, + "server_spec_drafts": 452, + "server_spec_draft_tokens": 1356, + "server_spec_accepted_tokens": 832, + "server_accept_len_effective": 2.84070796460177, + "server_engine_steps": 452.0, + "aggregate_div_spec_accept_length": 45.4134431829206, + "effective_normalized_steps_per_s": 45.4134431829206, + "measurement_seconds": 9.953, + "measurement_wall_seconds": 10.0, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tps": 124.18366321701298, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.758928571428571, + "server_steps_per_s": 45.01155430519564, + "server_output_tokens": 1236, + "client_output_tokens": 1236, + "server_spec_drafts": 448, + "server_spec_draft_tokens": 1344, + "server_spec_accepted_tokens": 788, + "server_accept_len_effective": 2.7589285714285716, + "server_engine_steps": 448.0, + "aggregate_div_spec_accept_length": 45.011554305195645, + "effective_normalized_steps_per_s": 45.01155430519564, + "measurement_seconds": 9.953, + "measurement_wall_seconds": 10.015, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + } + ] + }, + { + "dcp": 2, + "enabled": true, + "prefill": [ + { + "tokens": 8192, + "samples": 3, + "median_ttft_seconds": 2.3446302999946056, + "min_ttft_seconds": 2.3418036000075517, + "max_ttft_seconds": 2.360843100002967, + "tokens_per_second": 3493.9410277257134 + }, + { + "tokens": 16384, + "samples": 3, + "median_ttft_seconds": 4.732827400002861, + "min_ttft_seconds": 4.727047099993797, + "max_ttft_seconds": 4.737775899993721, + "tokens_per_second": 3461.778470938977 + }, + { + "tokens": 32768, + "samples": 3, + "median_ttft_seconds": 9.50934899999993, + "min_ttft_seconds": 9.49472020000394, + "max_ttft_seconds": 9.536788299999898, + "tokens_per_second": 3445.872057067234 + } + ], + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tps": 57.36963729532207, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.8606965174129355, + "server_steps_per_s": 20.05442973279954, + "server_output_tokens": 575, + "client_output_tokens": 571, + "server_spec_drafts": 201, + "server_spec_draft_tokens": 603, + "server_spec_accepted_tokens": 374, + "server_accept_len_effective": 2.8606965174129355, + "server_engine_steps": 201.0, + "aggregate_div_spec_accept_length": 20.05442973279954, + "effective_normalized_steps_per_s": 20.05442973279954, + "measurement_seconds": 9.953, + "measurement_wall_seconds": 10.016, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tps": 51.5, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.601010101010101, + "server_steps_per_s": 19.799999999999997, + "server_output_tokens": 515, + "client_output_tokens": 515, + "server_spec_drafts": 198, + "server_spec_draft_tokens": 594, + "server_spec_accepted_tokens": 317, + "server_accept_len_effective": 2.601010101010101, + "server_engine_steps": 198.0, + "aggregate_div_spec_accept_length": 19.799999999999997, + "effective_normalized_steps_per_s": 19.799999999999997, + "measurement_seconds": 10.0, + "measurement_wall_seconds": 10.0, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tps": 133.32663518545078, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.8599137931034484, + "server_steps_per_s": 46.61910981616365, + "server_output_tokens": 1327, + "client_output_tokens": 1327, + "server_spec_drafts": 464, + "server_spec_draft_tokens": 1392, + "server_spec_accepted_tokens": 863, + "server_accept_len_effective": 2.8599137931034484, + "server_engine_steps": 464.0, + "aggregate_div_spec_accept_length": 46.61910981616365, + "effective_normalized_steps_per_s": 46.61910981616365, + "measurement_seconds": 9.953, + "measurement_wall_seconds": 10.016, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tps": 126.50240384619515, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.8445945945945947, + "server_steps_per_s": 44.47115384616836, + "server_output_tokens": 1263, + "client_output_tokens": 1263, + "server_spec_drafts": 444, + "server_spec_draft_tokens": 1332, + "server_spec_accepted_tokens": 819, + "server_accept_len_effective": 2.8445945945945947, + "server_engine_steps": 444.0, + "aggregate_div_spec_accept_length": 44.47115384616836, + "effective_normalized_steps_per_s": 44.47115384616836, + "measurement_seconds": 9.984, + "measurement_wall_seconds": 10.016, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + } + ] + }, + { + "dcp": 1, + "enabled": false, + "prefill": [ + { + "tokens": 8192, + "samples": 3, + "median_ttft_seconds": 3.619972600004985, + "min_ttft_seconds": 3.612405099993339, + "max_ttft_seconds": 3.6210150999977486, + "tokens_per_second": 2263.0005541999735 + }, + { + "tokens": 16384, + "samples": 3, + "median_ttft_seconds": 6.263822799999616, + "min_ttft_seconds": 6.244891500013182, + "max_ttft_seconds": 6.276198699997622, + "tokens_per_second": 2615.6550916480273 + }, + { + "tokens": 32768, + "samples": 3, + "median_ttft_seconds": 11.58518439999898, + "min_ttft_seconds": 11.57795970000734, + "max_ttft_seconds": 11.618868300007307, + "tokens_per_second": 2828.440089395805 + } + ], + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tps": 58.90575079878845, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.783018867924528, + "server_steps_per_s": 21.166134185327376, + "server_output_tokens": 590, + "client_output_tokens": 590, + "server_spec_drafts": 212, + "server_spec_draft_tokens": 636, + "server_spec_accepted_tokens": 378, + "server_accept_len_effective": 2.7830188679245285, + "server_engine_steps": 212.0, + "aggregate_div_spec_accept_length": 21.16613418532738, + "effective_normalized_steps_per_s": 21.166134185327376, + "measurement_seconds": 10.016, + "measurement_wall_seconds": 10.016, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tps": 59.395032051214876, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.797169811320755, + "server_steps_per_s": 21.23397435895034, + "server_output_tokens": 593, + "client_output_tokens": 593, + "server_spec_drafts": 212, + "server_spec_draft_tokens": 636, + "server_spec_accepted_tokens": 381, + "server_accept_len_effective": 2.797169811320755, + "server_engine_steps": 212.0, + "aggregate_div_spec_accept_length": 21.23397435895034, + "effective_normalized_steps_per_s": 21.23397435895034, + "measurement_seconds": 9.984, + "measurement_wall_seconds": 10.0, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tps": 134.59810284573922, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.8083333333333336, + "server_steps_per_s": 47.928107838245424, + "server_output_tokens": 1348, + "client_output_tokens": 1348, + "server_spec_drafts": 480, + "server_spec_draft_tokens": 1440, + "server_spec_accepted_tokens": 868, + "server_accept_len_effective": 2.808333333333333, + "server_engine_steps": 480.0, + "aggregate_div_spec_accept_length": 47.92810783824542, + "effective_normalized_steps_per_s": 47.928107838245424, + "measurement_seconds": 10.015, + "measurement_wall_seconds": 10.015, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tps": 126.88366485837511, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.7456521739130437, + "server_steps_per_s": 46.21257785815721, + "server_output_tokens": 1263, + "client_output_tokens": 1263, + "server_spec_drafts": 460, + "server_spec_draft_tokens": 1380, + "server_spec_accepted_tokens": 803, + "server_accept_len_effective": 2.7456521739130433, + "server_engine_steps": 460.0, + "aggregate_div_spec_accept_length": 46.2125778581572, + "effective_normalized_steps_per_s": 46.21257785815721, + "measurement_seconds": 9.954, + "measurement_wall_seconds": 10.016, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + } + ] + }, + { + "dcp": 1, + "enabled": true, + "prefill": [ + { + "tokens": 8192, + "samples": 3, + "median_ttft_seconds": 2.265330099995481, + "min_ttft_seconds": 2.2511383999953978, + "max_ttft_seconds": 2.2724249999882886, + "tokens_per_second": 3616.2500114293907 + }, + { + "tokens": 16384, + "samples": 3, + "median_ttft_seconds": 4.532102700002724, + "min_ttft_seconds": 4.531001600000309, + "max_ttft_seconds": 4.537505600004806, + "tokens_per_second": 3615.0990135307725 + }, + { + "tokens": 32768, + "samples": 3, + "median_ttft_seconds": 9.137594200001331, + "min_ttft_seconds": 9.114788200007752, + "max_ttft_seconds": 9.137752499998896, + "tokens_per_second": 3586.0642618595634 + } + ], + "decode": [ + { + "concurrency": 1, + "context_tokens": 8192, + "aggregate_tps": 57.1, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.668224299065421, + "server_steps_per_s": 21.400000000000002, + "server_output_tokens": 571, + "client_output_tokens": 571, + "server_spec_drafts": 214, + "server_spec_draft_tokens": 642, + "server_spec_accepted_tokens": 357, + "server_accept_len_effective": 2.6682242990654204, + "server_engine_steps": 214.0, + "aggregate_div_spec_accept_length": 21.4, + "effective_normalized_steps_per_s": 21.400000000000002, + "measurement_seconds": 10.0, + "measurement_wall_seconds": 10.0, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 1, + "context_tokens": 32768, + "aggregate_tps": 58.68191393320866, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.7725118483412325, + "server_steps_per_s": 21.165613401550477, + "server_output_tokens": 585, + "client_output_tokens": 585, + "server_spec_drafts": 211, + "server_spec_draft_tokens": 633, + "server_spec_accepted_tokens": 374, + "server_accept_len_effective": 2.772511848341232, + "server_engine_steps": 211.0, + "aggregate_div_spec_accept_length": 21.165613401550473, + "effective_normalized_steps_per_s": 21.165613401550477, + "measurement_seconds": 9.969, + "measurement_wall_seconds": 10.0, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 4, + "context_tokens": 8192, + "aggregate_tps": 133.2, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.775, + "server_steps_per_s": 48.0, + "server_output_tokens": 1332, + "client_output_tokens": 1332, + "server_spec_drafts": 480, + "server_spec_draft_tokens": 1440, + "server_spec_accepted_tokens": 852, + "server_accept_len_effective": 2.775, + "server_engine_steps": 480.0, + "aggregate_div_spec_accept_length": 48.0, + "effective_normalized_steps_per_s": 48.0, + "measurement_seconds": 10.0, + "measurement_wall_seconds": 10.0, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + }, + { + "concurrency": 4, + "context_tokens": 32768, + "aggregate_tps": 142.98552171743216, + "mtp_normalization_available": true, + "server_spec_accept_length": 2.864, + "server_steps_per_s": 49.925112331505645, + "server_output_tokens": 1432, + "client_output_tokens": 1432, + "server_spec_drafts": 500, + "server_spec_draft_tokens": 1500, + "server_spec_accepted_tokens": 932, + "server_accept_len_effective": 2.864, + "server_engine_steps": 500.0, + "aggregate_div_spec_accept_length": 49.925112331505645, + "effective_normalized_steps_per_s": 49.925112331505645, + "measurement_seconds": 10.015, + "measurement_wall_seconds": 10.015, + "aggregate_source": "openai_continuous_usage", + "counter_normalization_verified": true + } + ] + } + ], + "receipt_audit": [ + { + "arm": "dcp1-off", + "passed": true, + "cold_samples": 9, + "semantic_cases": 6, + "decode_cells": 4, + "prefill_sha256": "30ea18031931c939c613a59ddb80de340961e11056fca04c76741472010bba43", + "decode_sha256": "42a6f64eae58c4acffcbccdd87ff0d16e549278e9d9d20cee6cd93a3ff6321cf" + }, + { + "arm": "dcp1-on", + "passed": true, + "cold_samples": 9, + "semantic_cases": 6, + "decode_cells": 4, + "prefill_sha256": "49a4dec267de7ddf953899206f7515786081fe5e18b8cd32b5b254e46dc3462e", + "decode_sha256": "c458014b8deda50576a4e8f4add8c5965f6ce6d70c92cd395cf16bb7c242894a" + }, + { + "arm": "dcp2-off", + "passed": true, + "cold_samples": 9, + "semantic_cases": 6, + "decode_cells": 4, + "prefill_sha256": "ec11d53952a425a04e45148d1d8383ed784d1002a52e4ff8f21d89864d82ccca", + "decode_sha256": "cfd8bf506e212b4338f8e9e3e7ee53251d1cbd77edc75891bf5794f25cbbfb30" + }, + { + "arm": "dcp2-on", + "passed": true, + "cold_samples": 9, + "semantic_cases": 6, + "decode_cells": 4, + "prefill_sha256": "4c7cc53c4022c469ec462375b8f2b50158991f2d03ee60ef2e5b5c403b510756", + "decode_sha256": "2fed41df8506584930cf369b3895feed0dd51f4fcc0ea4835a4a3d21c217c532" + } + ], + "input_sha256": { + "comparison.json": "443322cc57f8ed6c752167677eb9199e46638a5c7632d507838a66192489a54b", + "RESULTS.md": "cb167d452dfa4f27ca50134bcb341352915677ba6f5c4279878b380e9773efbe", + "QUALIFICATION-AUDIT.md": "ca2ef84c0de7385458baeb2783439f73e1684a386fa546a788eac10d2159dca8", + "continuation-fixture-comparison.md": "26d2beea0340588e42b97566fa8ed89461ec7ce9265b08138dfb74bedce054ba", + "continuation-fixture-comparison.json": "4f22a04f814630b08a66d3bad4a4ce18d4cb08286f6e0dd06f4a37aca5de824c" + }, + "numerical_limit": "DCP2 enabled cached-vs-cold max logprob difference 0.127015 fails the preselected tolerance; matched patch on/off cases pass, including extended-cold 0.100147 near allowed 0.100867. Exact-answer tokens match. Cache-invariant numerical behavior remains research-only." +} diff --git a/docs/features/quantization/b12x.md b/docs/features/quantization/b12x.md index f636dce05422..807abfb2df74 100644 --- a/docs/features/quantization/b12x.md +++ b/docs/features/quantization/b12x.md @@ -113,3 +113,113 @@ This crossover is topology-dependent; set the variable to `off` when the DMA ring has not been qualified against the fallback collective on the deployment hardware. `VLLM_PCIE_DMA_FP8` selects an optional compressed wire format; leaving it unset preserves lossless transport. + +## Bounded KDA prefill checkpoints + +A recurrent checkpoint saves the KDA state after a specific token offset. +`VLLM_B12X_KDA_PREFILL_COALESCING=1` lets an aligned prefill chunk produce +up to four required checkpoints without splitting the model forward at those +positions. The scheduler derives the positions from the cache manager's +retention policy and retains an 8192-token ceiling. With 512-token physical +recurrent blocks, 2048-token scheduler alignment and MTP, an 8192-token prompt +retains states at 4096, 6144, 7168 and 7680. The four-slot plan preserves both +fine prefix hits and scheduler-aligned fallback hits in one forward pass. + +Status: research-only. CPU contracts, GB10 checkpoint GPU tests, and bounded +TP4/DCP4 serving checks passed. The [four-configuration serving evidence](../../benchmarking/glm-kda-checkpoints-20260907/README.md) records +36 cold prefill samples, 16 decode cells and 20 exact-answer/cache checks. +These results do not establish full numerical or model-quality equivalence. + +The opt-in requires GLM5Next with BF16 activations on NVIDIA GB10 / SM121 / +48 SMs, model runner V2, TP4/DCP4/PP1/DP1, B12X KDA prefill, aligned prefix +caching and retention interval 0. Static MTP3 or no speculation is supported. +LoRA, expert/prefill-context parallelism, fairness scheduling and request-boundary +checkpoint mode are rejected. Coalescing applies when one cold prompt is the +only queued/running request. Cache-hit, resumed, preempted and mixed service +retains ordinary chunk splitting. Tensor-parallel collectives and topology are +unchanged; a switch-connected four-Spark cluster does not require SparkRing. + +Install a B12X source checkout whose `sequence.kda_prefill.Caps` accepts +`max_checkpoints=4` (four-column checkpoint metadata with transactional +validation). The implementation is maintained separately from token-sharded +mHC; neither optimization requires the other. With coalescing disabled, vLLM +uses one-checkpoint vectors and does not pass the added capability keyword, so +a B12X API without that keyword remains usable. + +Select the feature in the worker environment on all four ranks and retain the +cluster's serving launcher and model/communication settings. These arguments +define the checkpoint-specific configuration. The native split-page settings +preserve independent attention and recurrent page sizes. Without them, the +platform may enlarge both token blocks to fit one recurrent state in an +attention page, even when the CLI requests 512-token blocks. + +```bash +VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE=512 \ +VLLM_GLM53_SPLIT_MAMBA_BLOCK_SIZE=512 \ +VLLM_USE_V2_MODEL_RUNNER=1 VLLM_B12X_KDA_PREFILL_COALESCING=1 \ +vllm serve \ + --dtype bfloat16 \ + --tensor-parallel-size 4 --decode-context-parallel-size 4 \ + --pipeline-parallel-size 1 --max-num-batched-tokens 8192 \ + --enable-prefix-caching --mamba-cache-mode align \ + --block-size 512 --mamba-block-size 512 \ + --recurrent-checkpoint-policy aligned --prefix-cache-retention-interval 0 \ + --additional-config '{"kda_prefill_backend":"b12x"}' +``` + +Preserve unrelated entries in `--additional-config` when adding the KDA backend. +The command describes one head-node invocation; use the same environment and +source revisions in the cluster's worker launch configuration. A B12X build +without the checkpoint capability or a mismatched device/configuration raises +an initialization error. + +Keep both split-page settings at 512 in every performance-comparison arm, +including coalescing-disabled and mHC-only runs. They select a native vLLM +cache layout and do not require a SparkCache connector. The `auto` value +selects a different geometry from the scheduler budget and must not replace +512 for this configuration. Check the startup record for physical/lookup/ +scheduler grids `(512, 512, 2048)`, then verify a scheduled 8192-token span with +four checkpoints. Enabled flags alone do not establish execution of that path. + +The allocator keeps checkpoint and final-state destinations disjoint from the +retained source state. Only private, unhashed speculative reservations may move, +and worker block-table updates remain append-only. Model runner V2 maps the +scheduler plan to packed request rows; the GLM metadata adapter forwards it to +every GDN layer. Required NULL destinations fail before state consumers run. +B12X exports recurrent states from the FP32 accumulator, and the convolution +export stores its causal history rather than speculative buffer capacity. +With 16 local recurrent heads of 128 by 128 FP32 elements, each exported +state writes 1 MiB per GDN layer. Four checkpoints write 4 MiB: 3 MiB more +than one checkpoint, or 2 MiB more than two checkpoints. State storage and metadata capacities are +reserved before execution. + +The scheduler logs `KDA_PREFILL_COALESCING configured` with physical, lookup +and scheduler alignment, then `KDA_PREFILL_COALESCING scheduled` for the first +materialized plan of each span/count. These records distinguish configuration +from scheduled checkpoint export; GPU correctness still requires validation. + +Before model evaluation, run the CPU contracts and the GPU checkpoint-store +regression in a built Linux vLLM environment: + +```bash +.venv/bin/python -m pytest -q \ + tests/v1/core/test_recurrent_prefill_checkpoint.py \ + tests/v1/core/test_mamba_align_chunk_split.py \ + tests/v1/core/test_single_type_kv_cache_manager.py \ + tests/v1/attention/test_gdn_metadata_builder.py \ + tests/v1/worker/test_kda_prefill_checkpoint_binding.py +.venv/bin/python -m pytest -q tests/kernels/mamba/test_causal_conv1d.py \ + -k test_kda_checkpoint_history_excludes_speculative_cells +``` + +Evaluate the exact composed vLLM/B12X sources with the flag off and on. Use cold +8K/10K/16K/32K/64K/128K token-ID prompts, then repeat and extend each prompt to +exercise checkpoint reuse. Compare greedy output token IDs/logprobs and exact +record-retrieval answers before interpreting timings. Include mixed decode/ +prefill, concurrent cold requests, cache hits, preemption and capacity-failure +cases. Warm each prompt shape before performance measurements, preserve raw +samples and source/device identities, and report component and whole-model +results separately. The linked serving evidence covers 8K/16K/32K requests +and bounded exact-answer/cache checks. Full logit/state comparisons, +64K/128K serving qualification, preemption/capacity-failure model tests and +GSM8K/MRCR evaluations remain unqualified. diff --git a/tests/benchmarks/test_glm_prefill_checkpoints.py b/tests/benchmarks/test_glm_prefill_checkpoints.py new file mode 100644 index 000000000000..e292735d85f7 --- /dev/null +++ b/tests/benchmarks/test_glm_prefill_checkpoints.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline response/gate tests; never contact an inference endpoint.""" + +import importlib.util +import io +import json +from pathlib import Path +from types import SimpleNamespace as NS + +import pytest + +PATH = Path(__file__).resolve().parents[2] / "benchmarks/glm_prefill_checkpoints.py" +SPEC = importlib.util.spec_from_file_location("glm_prefill_checks", PATH) +assert SPEC is not None and SPEC.loader is not None +CHECKS = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(CHECKS) + + +@pytest.mark.parametrize( + "details", + [None, {}, {"cached_tokens": False}, {"cached_tokens": -1}, {"cached_tokens": "0"}], +) +def test_unknown_or_malformed_cache_usage_is_rejected(details): + with pytest.raises(ValueError, match="cached_tokens"): + CHECKS.cached_tokens({"prompt_tokens_details": details}) + + +def test_exact_usage_and_expected_cache_hit_are_both_required(): + details = {"cached_tokens": 0} + usage = {"prompt_tokens": 8192, "prompt_tokens_details": details} + assert CHECKS.validate_usage(usage, 8192, "cold") == 0 + with pytest.raises(ValueError, match="Prompt usage"): + CHECKS.validate_usage(usage, 16384, "cold") + with pytest.raises(ValueError, match="reuse was absent"): + CHECKS.validate_usage(usage, 8192, "reuse") + details["cached_tokens"] = 4096 + assert CHECKS.validate_usage(usage, 8192, "reuse") == 4096 + with pytest.raises(ValueError, match="Cold request"): + CHECKS.validate_usage(usage, 8192, "cold") + + +def test_metrics_require_present_finite_gauges(): + assert CHECKS.request_gauges( + 'vllm:num_requests_running{model="m"} 0\nvllm:num_requests_waiting 0\n' + ) == {"num_requests_running": 0, "num_requests_waiting": 0} + for text in ( + "", + "vllm:num_requests_running 0\n", + "vllm:num_requests_running NaN\nvllm:num_requests_waiting 0", + ): + with pytest.raises(ValueError): + CHECKS.request_gauges(text) + + +def test_summary_excludes_shape_warmup_and_failed_samples(): + cases = [ + {"phase": phase, "accepted": accepted, "tokens": 8192, "ttft_seconds": elapsed} + for phase, accepted, elapsed in [ + ("warmup", True, 100), + ("measured", True, 2), + ("measured", True, 4), + ("measured", False, 1), + ] + ] + summary = CHECKS.summarize(cases) + assert summary[0]["samples"] == 2 + assert summary[0]["median_ttft_seconds"] == 3 + assert summary[0]["tokens_per_second"] == 8192 / 3 + + +@pytest.mark.parametrize("done", [True, False]) +def test_stream_records_full_chunks_and_requires_final_completion(monkeypatch, done): + journal = NS(record={"requests": []}, save=lambda: None) + args = NS( + base_url="http://unit.invalid/v1", + api_key_env=None, + model="m", + request_timeout=1, + ) + client = CHECKS.PrefillChecks(args, journal) + chunks = [ + {"choices": [{"delta": {"role": "assistant", "content": ""}}]}, + {"choices": [{"delta": {"reasoning": "x"}}]}, + { + "choices": [], + "usage": { + "prompt_tokens": 8192, + "prompt_tokens_details": {"cached_tokens": 0}, + }, + }, + ] + raw = b"".join(b"data: " + json.dumps(chunk).encode() + b"\n" for chunk in chunks) + if done: + raw += b"data: [DONE]\n" + + class Response(io.BytesIO): + status = 200 + + monkeypatch.setattr( + CHECKS.urllib.request, "urlopen", lambda *a, **kw: Response(raw) + ) + values = iter([100.0, 100.25, 100.5]) + monkeypatch.setattr(CHECKS.time, "perf_counter", lambda: next(values)) + if done: + response, index = client.post( + "/v1/chat/completions", {"model": "m"}, stream=True + ) + assert index == 0 and response["ttft_seconds"] == 0.25 + assert response["complete"] and response["sse_chunks"] == chunks + else: + with pytest.raises(RuntimeError, match="DONE"): + client.post("/v1/chat/completions", {"model": "m"}, stream=True) + assert journal.record["requests"][0]["sse_chunks"] == chunks + assert not journal.record["requests"][0]["complete"] + + +def test_journal_refuses_overwrite_and_persists_updates_atomically(tmp_path): + path = tmp_path / "result.json" + journal = CHECKS.Journal(path, {"execution_status": "running"}) + with pytest.raises(FileExistsError): + CHECKS.Journal(path, {}) + journal.record["execution_status"] = "failed" + journal.save() + assert json.loads(path.read_text())["execution_status"] == "failed" + assert not list(tmp_path.glob("*.tmp")) + + +def test_fixed_protocol_orders_excluded_warmups_cache_checks_then_timing(monkeypatch): + args = NS( + base_url="http://unit.invalid", + api_key_env=None, + model="m", + request_timeout=1, + idle_timeout=1, + semantic_max_tokens=384, + ) + journal = NS(record={"requests": [], "cases": []}, save=lambda: None) + client = CHECKS.PrefillChecks(args, journal) + phases = [] + monkeypatch.setattr( + client, "control", lambda path: json.dumps({"data": [{"id": "m"}]}) + ) + monkeypatch.setattr(client, "idle", lambda **kw: None) + monkeypatch.setattr( + client, + "prefill", + lambda size, phase, sample: phases.append((phase, size, sample)), + ) + monkeypatch.setattr( + client, "semantic", lambda size, reuse: phases.append(("semantic", size, reuse)) + ) + client.run() + assert phases == ( + [("warmup", size, 0) for size in CHECKS.SIZES] + + [ + ("semantic", 8192, True), + ("semantic", 16384, False), + ("semantic", 32768, False), + ] + + [("measured", size, sample) for sample in range(3) for size in CHECKS.SIZES] + ) + + +def test_prompt_padding_calibrates_exact_token_count(monkeypatch): + journal = NS(record={"requests": []}, save=lambda: None) + client = CHECKS.PrefillChecks( + NS(base_url="http://unit.invalid", api_key_env=None, model="m"), journal + ) + monkeypatch.setattr(CHECKS.uuid, "uuid4", lambda: NS(hex="1" * 32)) + counts = [] + + def count(messages): + words = messages[0]["content"].split() + observed = ( + sum(word in ("alpha", "beta", "gamma", "delta") for word in words) + 37 + ) + counts.append(observed) + return observed, len(counts) - 1 + + monkeypatch.setattr(client, "count", count) + messages, requests = client.calibrate(8192, "STONE-7482") + assert counts == [8229, 8192] + assert requests == [0, 1] + assert messages[0]["content"].startswith("Test " + "1" * 32) + assert "The project code is STONE-7482. Remember it." in messages[0]["content"] + + +def test_generic_conditions_cannot_claim_activation(tmp_path): + path = tmp_path / "conditions.json" + path.write_text( + json.dumps( + { + "schema": "glm-prefill-reproduction-conditions/v1", + "sources": {"vllm": "revision"}, + "settings": {"tp": 4}, + } + ), + encoding="utf-8", + ) + conditions, sha = CHECKS.load_conditions(path) + assert conditions["sources"] == {"vllm": "revision"} and len(sha) == 64 + conditions["activation_verified"] = True + path.write_text(json.dumps(conditions), encoding="utf-8") + with pytest.raises(ValueError, match="Conditions require"): + CHECKS.load_conditions(path) + assert CHECKS.load_conditions(None) == (None, None) + + +def test_successful_run_does_not_assert_feature_activation(tmp_path, monkeypatch): + output = tmp_path / "result.json" + monkeypatch.setattr( + CHECKS.sys, + "argv", + [ + str(PATH), + "--base-url", + "http://unit.invalid", + "--output", + str(output), + "--label", + "both", + "--ready-confirmed", + "--exclusive-window", + ], + ) + monkeypatch.setattr(CHECKS.PrefillChecks, "run", lambda self: None) + CHECKS.main() + result = json.loads(output.read_text(encoding="utf-8")) + assert result["execution_status"] == "passed" + assert result["feature_activation"]["status"] == "not_verified" + assert result["conditions_provenance"] == "not_supplied" + assert result["settings"]["repeats"] == 3 + assert result["model_quality_qualified"] is False + + +@pytest.mark.parametrize( + "url", ["http://user:secret@unit.invalid", "http://:secret@unit.invalid"] +) +def test_url_credentials_are_rejected_before_journal_creation( + tmp_path, monkeypatch, url +): + output = tmp_path / "result.json" + monkeypatch.setattr( + CHECKS.sys, + "argv", + [ + str(PATH), + "--base-url", + url, + "--output", + str(output), + "--label", + "case", + "--ready-confirmed", + "--exclusive-window", + ], + ) + with pytest.raises(ValueError, match="without credentials"): + CHECKS.main() + assert not output.exists() + + +def test_api_key_is_sent_but_never_journaled(monkeypatch): + monkeypatch.setenv("PREFILL_UNIT_KEY", "synthetic-secret") + journal = NS(record={"requests": []}, save=lambda: None) + args = NS( + base_url="http://unit.invalid", + api_key_env="PREFILL_UNIT_KEY", + model="m", + request_timeout=1, + ) + client = CHECKS.PrefillChecks(args, journal) + + class Response(io.BytesIO): + status = 200 + + def urlopen(request, **kwargs): + assert request.get_header("Authorization") == "Bearer synthetic-secret" + return Response(b'{"count": 8192}') + + monkeypatch.setattr(CHECKS.urllib.request, "urlopen", urlopen) + assert client.count([{"role": "user", "content": "synthetic input"}])[0] == 8192 + assert "synthetic-secret" not in json.dumps(journal.record) diff --git a/tests/kernels/mamba/test_causal_conv1d.py b/tests/kernels/mamba/test_causal_conv1d.py index dad3dec3c9d5..25a31199b282 100644 --- a/tests/kernels/mamba/test_causal_conv1d.py +++ b/tests/kernels/mamba/test_causal_conv1d.py @@ -516,3 +516,66 @@ def run(): assert torch.equal(out_h, out_g) assert torch.equal(state_h, state_g) assert torch.isfinite(out_h.float()).all() + + +@pytest.mark.parametrize("checkpoint_count", [1, 2, 4]) +@pytest.mark.parametrize("metadata_error", [0, 1]) +def test_kda_checkpoint_history_excludes_speculative_cells( + checkpoint_count, metadata_error +): + """Convolution exports match the raw three-token history at each offset.""" + from types import SimpleNamespace + + from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import ( + KimiGatedDeltaNetAttention, + ) + + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for checkpoint-store kernel execution") + device = torch.device("cuda") + width = 12 + raw = torch.arange(64 * width, device=device, dtype=torch.float32).reshape( + 64, width + ) + state = torch.full((8, width, 6), -7.0, device=device) + offsets = torch.tensor( + [[16, 32, 48, 64]] if checkpoint_count == 4 else [[16, 48]], + dtype=torch.int32, + device=device, + ) + destinations = torch.tensor( + [[2, 3, 4, 5]] if checkpoint_count == 4 else [[2, 5]], + dtype=torch.int32, + device=device, + ) + if checkpoint_count == 1: + offsets = offsets[:, 0].contiguous() + destinations = destinations[:, 0].contiguous() + layer = KimiGatedDeltaNetAttention.__new__(KimiGatedDeltaNetAttention) + torch.nn.Module.__init__(layer) + layer.conv1d = torch.nn.Conv1d( + width, width, 4, groups=width, bias=False, device=device + ) + checkpoint = SimpleNamespace(checkpoint_offsets=offsets, state_indices=destinations) + KimiGatedDeltaNetAttention._store_kda_conv_checkpoint( + layer, + mixed_qkv=raw, + conv_state=state, + recurrent_state=torch.empty(8, 1, 128, 128, device=device), + query_start_loc=torch.tensor([0, 64], dtype=torch.int32, device=device), + checkpoint=checkpoint, + error_code=torch.tensor([metadata_error], dtype=torch.int32, device=device), + ) + torch.accelerator.synchronize() + if metadata_error: + assert torch.all(state == -7.0) + return + for slot, offset in zip( + destinations.flatten().tolist(), offsets.flatten().tolist() + ): + torch.testing.assert_close( + state[slot, :, :3], raw[offset - 3 : offset].T, rtol=0, atol=0 + ) + assert torch.all(state[slot, :, 3:] == -7.0) + untouched = sorted(set(range(8)) - set(destinations.flatten().tolist())) + assert torch.all(state[untouched] == -7.0) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index f1fdc5be4f25..7c06d367e66c 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -3,6 +3,7 @@ """Tests for mixed speculative and non-speculative GDN metadata.""" from dataclasses import dataclass, replace +from types import SimpleNamespace import pytest import torch @@ -10,9 +11,7 @@ from tests.v1.attention.utils import ( BatchSpec, create_common_attn_metadata, - create_vllm_config, ) -from vllm.config import SpeculativeConfig from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backends.gdn_attn import ( GDNAttentionMetadata, @@ -25,6 +24,13 @@ DEVICE = torch.device("cpu") +@pytest.fixture(autouse=True) +def _metadata_uses_unpinned_host_storage(monkeypatch): + """CPU metadata tests do not require accelerator-pinned allocations.""" + monkeypatch.setattr("vllm.utils.torch_utils.PIN_MEMORY", False) + monkeypatch.setattr("vllm.v1.attention.backends.utils.PIN_MEMORY", False) + + @dataclass class GDNBuildTestCase: """Specification for a GDN metadata builder classification test.""" @@ -126,18 +132,32 @@ def _create_gdn_builder( num_prefill_checkpoint_blocks: int = 0, ) -> GDNAttentionMetadataBuilder: """Create a GDNAttentionMetadataBuilder with minimal config.""" - vllm_config = create_vllm_config( - model_name="Qwen/Qwen3.5-0.8B", - block_size=BLOCK_SIZE, - max_num_batched_tokens=4096, - ) - if full_cuda_graph: - vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE - if num_speculative_tokens > 0: - vllm_config.speculative_config = SpeculativeConfig( - method="ngram", - num_speculative_tokens=num_speculative_tokens, - ) + # Metadata construction needs capacity and layout, not model weights or a + # platform-specific serving configuration. + vllm_config = SimpleNamespace( + compilation_config=SimpleNamespace( + cudagraph_mode=( + CUDAGraphMode.FULL_AND_PIECEWISE + if full_cuda_graph + else CUDAGraphMode.NONE + ), + max_cudagraph_capture_size=None, + ), + speculative_config=( + SimpleNamespace( + num_speculative_tokens=num_speculative_tokens, parallel_drafting=False + ) + if num_speculative_tokens + else None + ), + scheduler_config=SimpleNamespace(max_num_seqs=256), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + cache_config=SimpleNamespace(mamba_cache_mode="align"), + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace(linear_key_head_dim=128) + ), + additional_config={"gdn_prefill_backend": "triton"}, + ) mamba_spec = MambaSpec( block_size=BLOCK_SIZE, shapes=((16, 64),), @@ -715,3 +735,140 @@ def test_gdn_mixed_spec_update_selects_group_specific_state_indices() -> None: metadata_b.prefill_state_indices, builder_b.mamba_aligned_state_indices[0:1, 0], ) + + +def test_gdn_two_checkpoint_plan_uses_explicit_columns_and_refreshes_destinations(): + builder = _create_gdn_builder(num_prefill_checkpoint_blocks=2) + builder.vllm_config.cache_config.mamba_cache_mode = "align" + common = create_common_attn_metadata( + BatchSpec(seq_lens=[128], query_lens=[64]), + BLOCK_SIZE, + DEVICE, + arange_block_indices=True, + ).replace( + is_prefilling=torch.tensor([True]), + recurrent_prefill_checkpoint_plans_cpu=[(64, 128, (80, 112))], + ) + metadata = builder.build(common_prefix_len=0, common_attn_metadata=common) + checkpoint = metadata.prefill_checkpoint + assert checkpoint is not None + torch.testing.assert_close( + checkpoint.checkpoint_offsets, torch.tensor([[16, 48]], dtype=torch.int32) + ) + torch.testing.assert_close( + checkpoint.block_table_columns, torch.tensor([[4, 6]], dtype=torch.int64) + ) + assert checkpoint.required_mask is not None and checkpoint.required_mask.all() + updated = builder.update_block_table( + metadata, common.block_table_tensor + 100, torch.zeros(64, dtype=torch.int64) + ) + assert updated.prefill_checkpoint is not None + torch.testing.assert_close( + updated.prefill_checkpoint.state_indices, checkpoint.state_indices + 100 + ) + + +def test_gdn_planned_checkpoint_rejects_missing_required_state_storage(): + builder = _create_gdn_builder(num_prefill_checkpoint_blocks=2) + builder.vllm_config.cache_config.mamba_cache_mode = "align" + common = create_common_attn_metadata( + BatchSpec(seq_lens=[128], query_lens=[64]), + BLOCK_SIZE, + DEVICE, + arange_block_indices=True, + ).replace( + is_prefilling=torch.tensor([True]), + recurrent_prefill_checkpoint_plans_cpu=[(64, 128, (80, 112))], + ) + common.block_table_tensor[0, 6] = 0 + with pytest.raises(RuntimeError, match="NULL state block"): + builder.build(common_prefix_len=0, common_attn_metadata=common) + + +def test_unpadded_common_metadata_preserves_live_checkpoint_plan_rows(): + plan = (0, 64, (16, 48)) + common = create_common_attn_metadata( + BatchSpec(seq_lens=[64, 64], query_lens=[64, 64]), BLOCK_SIZE, DEVICE + ).replace(recurrent_prefill_checkpoint_plans_cpu=[plan, None]) + assert common.unpadded(64, 1).recurrent_prefill_checkpoint_plans_cpu == [plan] + + +@pytest.mark.parametrize("missing_column", [None, 6]) +def test_gdn_four_checkpoint_plan_refreshes_all_fine_and_coarse_destinations( + missing_column, +): + builder = _create_gdn_builder(num_prefill_checkpoint_blocks=4) + builder.vllm_config.cache_config.mamba_cache_mode = "align" + common = create_common_attn_metadata( + BatchSpec(seq_lens=[160], query_lens=[96]), + BLOCK_SIZE, + DEVICE, + arange_block_indices=True, + ).replace( + is_prefilling=torch.tensor([True]), + recurrent_prefill_checkpoint_plans_cpu=[(64, 160, (80, 96, 112, 144))], + ) + if missing_column is not None: + common.block_table_tensor[0, missing_column] = 0 + with pytest.raises(RuntimeError, match="NULL state block"): + builder.build(common_prefix_len=0, common_attn_metadata=common) + return + metadata = builder.build(common_prefix_len=0, common_attn_metadata=common) + checkpoint = metadata.prefill_checkpoint + assert checkpoint is not None + assert checkpoint.checkpoint_offsets.tolist() == [[16, 32, 48, 80]] + assert checkpoint.block_table_columns.tolist() == [[4, 5, 6, 8]] + assert checkpoint.required_mask is not None and checkpoint.required_mask.all() + updated = builder.update_block_table( + metadata, common.block_table_tensor + 100, torch.zeros(64, dtype=torch.int64) + ) + assert updated.prefill_checkpoint is not None + torch.testing.assert_close( + updated.prefill_checkpoint.state_indices, checkpoint.state_indices + 100 + ) + + +@pytest.mark.parametrize("prompt", [8192, 16384]) +def test_dcp4_allocator_plan_reaches_all_gdn_checkpoint_destinations(prompt): + from tests.v1.core.test_recurrent_prefill_checkpoint import cache_fixture + + cache, manager, scheduler, request = cache_fixture(prompt, speculative=3) + start = prompt - 8192 + if start: + assert cache.allocate_slots(request, 8192, num_lookahead_tokens=3) is not None + scheduler._record_coalescing_origin(request, 0, 0, 0, False) + request.num_computed_tokens = start + plan = scheduler._recurrent_checkpoint_plan(request, start, prompt) + assert plan is not None and len(plan[2]) == 4 + assert ( + cache.allocate_slots( + request, 8192, num_lookahead_tokens=3, recurrent_checkpoint_plan=plan + ) + is not None + ) + blocks = manager.req_to_blocks[request.request_id] + builder = _create_gdn_builder( + num_speculative_tokens=3, num_prefill_checkpoint_blocks=4 + ) + builder.kv_cache_spec = manager.kv_cache_spec + builder.vllm_config.parallel_config.decode_context_parallel_size = 4 + common = create_common_attn_metadata( + BatchSpec(seq_lens=[prompt], query_lens=[8192]), 512, DEVICE + ).replace( + is_prefilling=torch.tensor([True]), + block_table_tensor=torch.tensor( + [[b.block_id for b in blocks]], dtype=torch.int32 + ), + recurrent_prefill_checkpoint_plans_cpu=[plan], + ) + metadata = builder.build(common_prefix_len=0, common_attn_metadata=common) + checkpoint = metadata.prefill_checkpoint + assert checkpoint is not None + assert checkpoint.checkpoint_offsets.tolist() == [[4096, 6144, 7168, 7680]] + columns = [position // 512 - 1 for position in plan[2]] + expected_ids = [blocks[column].block_id for column in columns] + assert checkpoint.state_indices.tolist() == [expected_ids] + assert checkpoint.required_mask is not None and checkpoint.required_mask.all() + final_and_reserve = [block.block_id for block in blocks[prompt // 512 - 1 :]] + assert len(final_and_reserve) == 4 + assert len(set(expected_ids + final_and_reserve)) == 8 diff --git a/tests/v1/core/test_recurrent_prefill_checkpoint.py b/tests/v1/core/test_recurrent_prefill_checkpoint.py new file mode 100644 index 000000000000..359692a6f3ae --- /dev/null +++ b/tests/v1/core/test_recurrent_prefill_checkpoint.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Checkpoint planning, append-only ownership and worker metadata contracts.""" + +from types import SimpleNamespace as NS +from typing import TYPE_CHECKING + +import numpy as np +import pytest +import torch + +from vllm.v1.core.recurrent_prefill_checkpoint import ( + checkpoint_metadata, + checkpoint_plan_rows, + continuation_layout, + prefill_checkpoint_plan, + validate_coalescing_config, +) + +if TYPE_CHECKING: + from vllm.v1.core.kv_cache_utils import KVCacheBlock + +pytestmark = pytest.mark.cpu_test + + +def test_fresh_and_continuation_plans_export_exact_retained_boundaries(): + assert prefill_checkpoint_plan( + start=0, + end=8192, + prompt=8192, + num_tokens=8192, + block_size=512, + publications=(6144, 7168), + ) == (0, 8192, (6144, 7168)) + assert prefill_checkpoint_plan( + start=8192, + end=16384, + prompt=16384, + num_tokens=16384, + block_size=512, + publications=(14336, 15360), + ) == (8192, 16384, (14336, 15360)) + assert checkpoint_metadata((8192, 16384, (14336, 15360)), 8192, 16384, 512, 2) == ( + [6144, 7168], + [27, 29], + ) + + +@pytest.mark.parametrize("end", [8193, 16385, 24576]) +def test_unrepresentable_or_over_budget_final_spans_use_ordinary_splitting(end): + assert ( + prefill_checkpoint_plan( + start=8192, + end=end, + prompt=end, + num_tokens=end, + block_size=512, + publications=(14336,), + ) + is None + ) + + +def test_worker_checkpoint_rows_follow_request_order_and_reject_span_drift(): + batch = NS( + req_ids=["b", "a"], + num_reqs=2, + num_computed_tokens_np=np.array([8192, 0]), + num_scheduled_tokens=np.array([8192, 8192]), + query_start_loc_np=np.array([0, 8192, 16384]), + ) + plans = {"a": (0, 8192, (6144, 7168)), "b": (8192, 16384, (14336, 15360))} + assert checkpoint_plan_rows(batch, plans, 3) == [plans["b"], plans["a"], None] + with pytest.raises(ValueError, match="capture"): + checkpoint_plan_rows(batch, plans, 3, for_capture=True) + batch.num_computed_tokens_np[0] += 1 + with pytest.raises(ValueError, match="query span"): + checkpoint_plan_rows(batch, plans, 3) + + +def supported_config(): + return NS( + model_config=NS( + hf_text_config=NS(model_type="glm5_next_text"), + dtype=torch.bfloat16, + enable_sleep_mode=False, + enable_return_routed_experts=False, + ), + cache_config=NS( + enable_prefix_caching=True, + mamba_cache_mode="align", + prefix_cache_retention_interval=0, + ), + parallel_config=NS( + tensor_parallel_size=4, + decode_context_parallel_size=4, + pipeline_parallel_size=1, + data_parallel_size=1, + prefill_context_parallel_size=1, + enable_expert_parallel=False, + enable_eplb=False, + ), + scheduler_config=NS( + max_num_batched_tokens=8192, + max_num_scheduled_tokens=None, + long_prefill_token_threshold=0, + fairness_engine=None, + enable_chunked_prefill=True, + ), + use_v2_model_runner=True, + lora_config=None, + use_request_boundary_checkpoints=False, + additional_config={"kda_prefill_backend": "b12x"}, + speculative_config=None, + ) + + +def test_coalescing_default_off_does_not_require_model_capabilities(monkeypatch): + monkeypatch.setenv("VLLM_B12X_KDA_PREFILL_COALESCING", "0") + assert not validate_coalescing_config(object()) + + +@pytest.mark.parametrize("dcp", [1, 2, 4]) +def test_coalescing_accepts_context_groups_within_tp4(monkeypatch, dcp): + monkeypatch.setenv("VLLM_B12X_KDA_PREFILL_COALESCING", "1") + config = supported_config() + config.parallel_config.decode_context_parallel_size = dcp + assert validate_coalescing_config(config) + + +@pytest.mark.parametrize( + "section,field,value", + [ + ("parallel_config", "tensor_parallel_size", 2), + ("parallel_config", "decode_context_parallel_size", 3), + ("scheduler_config", "max_num_batched_tokens", 4096), + ("cache_config", "prefix_cache_retention_interval", 2048), + ("model_config", "dtype", torch.float16), + ], +) +def test_explicit_coalescing_rejects_unsupported_configuration( + monkeypatch, section, field, value +): + monkeypatch.setenv("VLLM_B12X_KDA_PREFILL_COALESCING", "1") + config = supported_config() + assert validate_coalescing_config(config) + setattr(getattr(config, section), field, value) + with pytest.raises(ValueError, match="requires GLM5Next"): + validate_coalescing_config(config) + + +def cache_fixture(prompt, speculative=3, *, checkpoints=4, dcp=4): + from tests.v1.core.utils import create_requests + from vllm.v1.core.kv_cache_manager import KVCacheManager + from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes + from vllm.v1.core.sched.scheduler import Scheduler + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, + ) + + config = KVCacheConfig( + num_blocks=1000, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=512, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["recurrent"], + MambaSpec( + block_size=512, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + num_speculative_blocks=speculative, + num_prefill_checkpoint_blocks=checkpoints, + ), + ), + ], + ) + scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes( + config, + NS( + cache_config=NS( + block_size=512, enable_prefix_caching=True, prefix_match_unit=None + ), + parallel_config=NS(decode_context_parallel_size=dcp), + kv_transfer_config=None, + ), + ) + assert (scheduler_block_size, hash_block_size) == (512 * dcp, 512) + cache = KVCacheManager( + config, + max_model_len=131072, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + enable_caching=True, + use_eagle=True, + dcp_world_size=dcp, + ) + (request,) = create_requests(1, num_tokens=prompt, block_size=512) + scheduler = Scheduler.__new__(Scheduler) + scheduler._kda_coalescing_enabled = True + scheduler._kda_coalescing_exclusive = True + scheduler._kda_coalescing_origins = {} + scheduler.cache_config = NS(block_size=512, prefix_cache_retention_interval=0) + scheduler.kv_cache_manager = cache + scheduler.max_num_scheduled_tokens = 8192 + scheduler.scheduler_config = NS(long_prefill_token_threshold=0) + scheduler.mamba_has_prefill_checkpoint_blocks = True + scheduler.mamba_partial_cache_hit = cache.coordinator.enable_partial_hash_hits + scheduler.hash_block_size = 512 + scheduler.drop_last_prefix_cache_block = True + scheduler.use_eagle = True + return cache, cache.coordinator.single_type_managers[1], scheduler, request + + +@pytest.mark.parametrize("prompt", [8192, 10240, 16384, 32768]) +@pytest.mark.parametrize("speculative", [0, 3]) +@pytest.mark.parametrize("dcp", [1, 2, 4]) +def test_scheduler_and_allocator_keep_8k_chunks_and_worker_column_ownership( + prompt, speculative, dcp +): + from vllm.v1.request import RequestStatus + + cache, manager, scheduler, request = cache_fixture(prompt, speculative, dcp=dcp) + worker: list[KVCacheBlock] = [] + chunks = [] + while request.num_computed_tokens < prompt: + start = request.num_computed_tokens + size = scheduler._mamba_block_aligned_split(request, min(8192, prompt - start)) + plan = scheduler._recurrent_checkpoint_plan(request, start, start + size) + prefix = tuple(worker) + result = cache.allocate_slots( + request, + size, + num_lookahead_tokens=speculative, + recurrent_checkpoint_plan=plan, + ) + assert result is not None + worker.extend(result.blocks[1]) + assert worker[: len(prefix)] == list(prefix) + if plan: + active = ([start // 512 - 1] if start else []) + [ + p // 512 - 1 for p in plan[2] + ] + active += list(range((start + size) // 512 - 1, len(worker))) + assert all( + worker[c] is manager.req_to_blocks[request.request_id][c] + for c in active + ) + assert len(active) == len({worker[c].block_id for c in active}) + assert manager._planned_recurrent_checkpoints == {} + if start == 0: + scheduler._record_coalescing_origin(request, 0, 0, 0, False) + request.status = RequestStatus.RUNNING + request.num_computed_tokens += size + chunks.append(size) + assert chunks == [8192] * (prompt // 8192) + ( + [prompt % 8192] if prompt % 8192 else [] + ) + + +@pytest.mark.parametrize("dcp", [1, 2, 4]) +def test_continuation_admission_failure_releases_only_its_temporary_plan( + monkeypatch, dcp +): + cache, manager, scheduler, request = cache_fixture(16384, dcp=dcp) + assert cache.allocate_slots(request, 8192) is not None + scheduler._record_coalescing_origin(request, 0, 0, 0, False) + request.num_computed_tokens = 8192 + plan = scheduler._recurrent_checkpoint_plan(request, 8192, 16384) + assert plan is not None + retained = tuple(manager.req_to_blocks[request.request_id]) + monkeypatch.setattr(cache.block_pool, "get_num_free_blocks", lambda: 0) + assert cache.allocate_slots(request, 8192, recurrent_checkpoint_plan=plan) is None + assert manager._planned_recurrent_checkpoints == {} + assert tuple(manager.req_to_blocks[request.request_id]) == retained + + +@pytest.mark.parametrize("dcp", [1, 2, 4]) +def test_cache_hit_preemption_and_shared_speculative_slots_disable_coalescing(dcp): + cache, manager, scheduler, request = cache_fixture(16384, dcp=dcp) + assert cache.allocate_slots(request, 8192) is not None + request.num_computed_tokens = 8192 + assert scheduler._recurrent_checkpoint_plan(request, 8192, 16384) is None + scheduler._kda_coalescing_origins[request.request_id] = request + assert scheduler._recurrent_checkpoint_plan(request, 8192, 16384) is not None + request.num_preemptions = 1 + assert scheduler._recurrent_checkpoint_plan(request, 8192, 16384) is None + request.num_preemptions = 0 + manager.req_to_blocks[request.request_id][-1].ref_cnt += 1 + assert scheduler._recurrent_checkpoint_plan(request, 8192, 16384) is None + with pytest.raises(ValueError, match="private"): + continuation_layout( + manager.req_to_blocks[request.request_id], + (8192, 16384, (14336, 15360)), + 512, + 3, + ) + + +def test_full_prompt_admission_reserves_checkpoint_peak_for_internal_8k_chunk(): + cache, manager, scheduler, request = cache_fixture(10240) + plan = scheduler._recurrent_checkpoint_plan(request, 0, 8192) + assert plan == (0, 8192, (6144,)) + manager._planned_recurrent_checkpoints[request.request_id] = plan + try: + peak = manager.get_num_blocks_to_allocate( + request.request_id, 10240, [], 0, 0, 10240, apply_admission_cap=True + ) + assert peak == 9 + assert request.request_id not in manager._num_checkpoint_blocks + finally: + manager._planned_recurrent_checkpoints.clear() + result = cache.allocate_slots( + request, 8192, recurrent_checkpoint_plan=plan, full_sequence_must_fit=True + ) + assert result is not None + assert manager._planned_recurrent_checkpoints == {} + assert not manager.req_to_blocks[request.request_id][11].is_null + + +@pytest.mark.parametrize("prompt", [8192, 16384, 32768]) +@pytest.mark.parametrize("dcp", [1, 2, 4]) +def test_dcp_geometry_retains_required_states_without_extra_passes(prompt, dcp): + cache, manager, scheduler, request = cache_fixture(prompt, dcp=dcp) + assert manager.hit_alignment_tokens == 512 + assert manager.scheduler_block_size == 512 * dcp + offsets = {1: (1024, 512), 2: (2048, 1024, 512), 4: (4096, 2048, 1024, 512)} + expected = tuple(prompt - offset for offset in offsets[dcp]) + assert tuple(sorted(manager._expand_reachable_boundaries([prompt - 1]))) == expected + start = prompt - 8192 + if start: + while request.num_computed_tokens < start: + assert cache.allocate_slots(request, 8192) is not None + if request.num_computed_tokens == 0: + scheduler._record_coalescing_origin(request, 0, 0, 0, False) + request.num_computed_tokens += 8192 + plan = scheduler._recurrent_checkpoint_plan(request, start, prompt) + assert plan == (start, prompt, expected) + assert scheduler._mamba_block_aligned_split(request, 8192) == 8192 + assert ( + cache.allocate_slots( + request, 8192, num_lookahead_tokens=3, recurrent_checkpoint_plan=plan + ) + is not None + ) + blocks = manager.req_to_blocks[request.request_id] + columns = [p // 512 - 1 for p in expected] + [prompt // 512 - 1] + assert all(not blocks[column].is_null for column in columns) + assert len({blocks[column].block_id for column in columns}) == len(columns) + + +def test_two_checkpoint_capacity_keeps_safe_dcp4_fallback(): + _, manager, scheduler, request = cache_fixture(8192, checkpoints=2) + assert manager.hit_alignment_tokens == 512 + assert scheduler._recurrent_checkpoint_plan(request, 0, 8192) is None + assert scheduler._mamba_block_aligned_split(request, 8192) == 4096 diff --git a/tests/v1/worker/test_kda_prefill_checkpoint_binding.py b/tests/v1/worker/test_kda_prefill_checkpoint_binding.py new file mode 100644 index 000000000000..5718ea2b29e9 --- /dev/null +++ b/tests/v1/worker/test_kda_prefill_checkpoint_binding.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""KDA checkpoint binding and convolution-history wrapper contracts.""" + +from types import SimpleNamespace as NS + +import pytest +import torch + +pytestmark = pytest.mark.cpu_test + + +def test_disabled_coalescing_uses_one_checkpoint_caps_without_added_keyword( + monkeypatch, +): + from vllm.model_executor.layers.mamba.gdn import kimi_gdn_linear_attn as module + + class OneCheckpointCaps: + def __init__( + self, + *, + device, + max_tokens, + max_seqs, + max_state_slots, + heads, + head_dim, + model_dtype, + state_dtype, + qk_l2norm, + checkpoint_export, + null_state_index, + metadata_validation, + ): + self.metadata_validation = metadata_validation + self.max_tokens = max_tokens + + api = NS(Caps=OneCheckpointCaps, plan=lambda caps: caps) + monkeypatch.setattr(module, "current_platform", NS(current_device=lambda: "cpu")) + layer = NS( + _b12x_prefill_api=api, + _b12x_prefill_checkpoint_capacity=1, + _b12x_prefill_max_tokens=8192, + _b12x_prefill_max_seqs=16, + local_num_heads=16, + head_dim=128, + model_config=NS(dtype=torch.bfloat16), + get_state_dtype=lambda: (torch.float32, torch.float32), + ) + plan = module.KimiGatedDeltaNetAttention._make_b12x_kda_prefill_plan(layer, 8) + assert isinstance(plan, OneCheckpointCaps) + assert plan.metadata_validation == "trusted" + + +def test_one_checkpoint_binding_retains_vector_metadata_without_mhc_state(): + from vllm.model_executor.layers.mamba.gdn import kimi_gdn_linear_attn as module + + observed = {} + + def bind(plan, **kwargs): + observed.update(kwargs) + return NS(error_code=torch.zeros(1, dtype=torch.int32)) + + layer = NS( + _b12x_prefill_api=NS(bind=bind, run=lambda *args, **kwargs: None), + _b12x_prefill_plan=object(), + _b12x_prefill_checkpoint_capacity=1, + _b12x_prefill_max_tokens=32, + _b12x_prefill_max_seqs=2, + _b12x_prefill_initial_indices=torch.zeros(2, dtype=torch.int32), + _b12x_prefill_null_indices=torch.zeros(2, dtype=torch.int32), + _b12x_prefill_zero_offsets=torch.zeros(2, dtype=torch.int32), + _b12x_prefill_num_seqs=torch.zeros(1, dtype=torch.int32), + _b12x_prefill_num_tokens=torch.zeros(1, dtype=torch.int32), + A_log=torch.zeros(1), + dt_bias=torch.zeros(1, 128), + head_dim=128, + gate_lower_bound=-5.0, + ) + rows = torch.empty(16, 1, 128, dtype=torch.bfloat16) + result = module.KimiGatedDeltaNetAttention._run_b12x_kda_prefill( + layer, + scratch=torch.empty(1024, dtype=torch.uint8), + q=rows, + k=rows, + v=rows, + raw_g=rows, + raw_beta=torch.zeros(16, 1, dtype=torch.bfloat16), + cu_seqlens=torch.tensor([0, 16], dtype=torch.int32), + state_indices=torch.tensor([1], dtype=torch.int32), + has_initial_state=torch.tensor([True]), + checkpoint=None, + recurrent_state=torch.empty(4, 1, 128, 128), + output=torch.empty_like(rows), + ) + assert result is None + assert observed["checkpoint_state_indices"].shape == (1,) + assert observed["checkpoint_offsets"].shape == (1,) + assert observed["initial_state_indices"].tolist() == [1] + + +@pytest.mark.parametrize("capacity", [2, 4]) +def test_checkpoint_convolution_stores_causal_history_not_speculative_capacity( + monkeypatch, + capacity, +): + from vllm.model_executor.layers.mamba.gdn import kimi_gdn_linear_attn as module + + observed = {} + + class CaptureKernel: + def __getitem__(self, grid): + observed["grid"] = grid + + def launch(*args): + observed["args"] = args + + return launch + + monkeypatch.setattr(module, "_store_cache_checkpoints_kernel", CaptureKernel()) + layer = NS(conv1d=NS(weight=torch.empty(6, 1, 4))) + checkpoint = NS( + checkpoint_offsets=torch.arange(1, capacity + 1, dtype=torch.int32).reshape( + 1, capacity + ) + * 16, + state_indices=torch.arange(2, capacity + 2, dtype=torch.int32).reshape( + 1, capacity + ), + ) + error = torch.zeros(1, dtype=torch.int32) + module.KimiGatedDeltaNetAttention._store_kda_conv_checkpoint( + layer, + mixed_qkv=torch.empty(capacity * 16, 6), + conv_state=torch.empty(8, 6, 6), + recurrent_state=torch.empty(8, 1, 128, 128), + query_start_loc=torch.tensor([0, capacity * 16], dtype=torch.int32), + checkpoint=checkpoint, + error_code=error, + ) + assert observed["grid"][0] == capacity + assert observed["args"][15] == 3 + assert observed["args"][21] == capacity + assert observed["args"][22] is error and observed["args"][23] is True + + +@pytest.mark.parametrize("capacity", [1, 2, 4]) +def test_prefill_warmup_binds_the_planned_checkpoint_capacity(monkeypatch, capacity): + from vllm.model_executor.layers.mamba.gdn import kimi_gdn_linear_attn as module + + shape = (2, capacity) if capacity > 1 else (2,) + observed = {} + + def bind(plan, **kwargs): + observed.update(kwargs) + return object() + + caps = NS( + device=torch.device("cpu"), + heads=1, + head_dim=128, + chunk_tokens=16, + model_dtype=torch.bfloat16, + max_state_slots=4, + ) + layer = NS( + _b12x_prefill_plan=NS(caps=caps), + _b12x_prefill_api=NS(bind=bind, prewarm=lambda binding: None), + _b12x_prefill_null_indices=torch.zeros(shape, dtype=torch.int32), + _b12x_prefill_zero_offsets=torch.zeros(shape, dtype=torch.int32), + _b12x_prefill_num_seqs=torch.zeros(1, dtype=torch.int32), + _b12x_prefill_num_tokens=torch.zeros(1, dtype=torch.int32), + _b12x_prefill_max_tokens=8192, + _b12x_prefill_max_seqs=2, + _b12x_prefill_checkpoint_capacity=capacity, + local_num_heads=1, + head_dim=128, + A_log=torch.zeros(1), + dt_bias=torch.zeros(1, 128), + kv_cache=(torch.empty(4, 384, 6), torch.empty(4, 1, 128, 128)), + ) + monkeypatch.setattr( + module, + "get_b12x_scratch_buffers", + lambda plan: [torch.empty(4096, dtype=torch.uint8)], + ) + unit = module._B12xKdaPrefillWarmup().get_b12x_warmup_unit( + layer, (16,), torch.bfloat16 + ) + unit.compile() + expected = (1, capacity) if capacity > 1 else (1,) + assert observed["checkpoint_state_indices"].shape == expected + assert observed["checkpoint_offsets"].shape == expected diff --git a/vllm/envs.py b/vllm/envs.py index 90dafdc75173..3fb116d37f5c 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -189,6 +189,7 @@ VLLM_HUMMING_INPUT_QUANT_CONFIG: dict[str, Any] | None = None VLLM_HUMMING_USE_F16_ACCUM: bool = False VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None + VLLM_B12X_KDA_PREFILL_COALESCING: bool = False VLLM_B12X_MOE_FP4_FORCE_A16: bool = False VLLM_B12X_DENSE_ACTIVATION_MODE: Literal["auto", "a16", "quantized"] = "auto" VLLM_B12X_NVFP4_ACTIVATION_MODE: Literal["auto", "a16", "quantized"] | None = None @@ -1639,6 +1640,10 @@ def _resolve_rust_cli_path() -> str | None: int(os.getenv("VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER", "1")) ), # Force b12x FP4 MoE to use BF16 activations. + # Export recurrent checkpoints inside bounded GLM KDA prefills. + "VLLM_B12X_KDA_PREFILL_COALESCING": lambda: bool( + int(os.getenv("VLLM_B12X_KDA_PREFILL_COALESCING", "0")) + ), "VLLM_B12X_MOE_FP4_FORCE_A16": lambda: bool( int(os.getenv("VLLM_B12X_MOE_FP4_FORCE_A16", "0")) ), diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index c717f7685e6d..4ae6b8ebf281 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -157,18 +157,26 @@ def _store_cache_checkpoints_kernel( NULL_STATE_IDX: tl.constexpr, BLOCK_SIZE: tl.constexpr, STORE_RECURRENT: tl.constexpr, + CHECKPOINTS: tl.constexpr = 1, + error_code_ptr=None, + CHECK_ERROR: tl.constexpr = False, ): """Store FlashKDA recurrent and convolution state at an internal boundary.""" - seq_idx = tl.program_id(0) + checkpoint_idx = tl.program_id(0) + seq_idx = checkpoint_idx // CHECKPOINTS cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) seq_idx_i64 = seq_idx.to(tl.int64) cols_i64 = cols.to(tl.int64) - state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx_i64) + state_idx = tl.load(checkpoint_state_indices_ptr + checkpoint_idx.to(tl.int64)) state_idx_i64 = state_idx.to(tl.int64) checkpoint_offset = tl.load( - checkpoint_offsets_ptr + seq_idx_i64 * checkpoint_offset_stride + checkpoint_offsets_ptr + + seq_idx_i64 * checkpoint_offset_stride + + (checkpoint_idx % CHECKPOINTS).to(tl.int64) ) valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0) + if CHECK_ERROR: + valid_checkpoint = valid_checkpoint & (tl.load(error_code_ptr) == 0) valid_conv = ( (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN) ) @@ -389,8 +397,8 @@ def compile() -> None: cu_seqlens=torch.tensor([0, tokens], dtype=torch.int32, device=device), initial_state_indices=indices, final_state_indices=indices, - checkpoint_state_indices=indices, - checkpoint_offsets=torch.zeros(1, dtype=torch.int32, device=device), + checkpoint_state_indices=layer._b12x_prefill_null_indices[:1], + checkpoint_offsets=layer._b12x_prefill_zero_offsets[:1], num_seqs=layer._b12x_prefill_num_seqs, num_tokens=layer._b12x_prefill_num_tokens, output=torch.zeros_like(rows), @@ -407,6 +415,7 @@ def compile() -> None: layer.head_dim, layer._b12x_prefill_max_tokens, layer._b12x_prefill_max_seqs, + layer._b12x_prefill_checkpoint_capacity, None if caps is None else caps.max_state_slots, ), compile=compile, @@ -443,7 +452,8 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec: assert isinstance(spec, MambaSpec) return replace( spec, - num_prefill_checkpoint_blocks=int( + num_prefill_checkpoint_blocks=self._b12x_prefill_checkpoint_capacity + * int( self.kda_prefill_backend in ("flashkda", "b12x") and not vllm_config.use_request_boundary_checkpoints ), @@ -724,6 +734,15 @@ def unbind_kv_cache(self) -> None: def _initialize_b12x_kda_prefill(self, vllm_config: VllmConfig) -> None: """Hold the b12x prefill op and its per-request metadata buffers.""" + from vllm.v1.core.recurrent_prefill_checkpoint import ( + COALESCED_CHECKPOINT_CAPACITY, + validate_coalescing_config, + ) + + coalescing = validate_coalescing_config(vllm_config) + self._b12x_prefill_checkpoint_capacity = ( + COALESCED_CHECKPOINT_CAPACITY if coalescing else 1 + ) if self.kda_prefill_backend != "b12x": return api = get_b12x_kda_prefill() @@ -732,11 +751,28 @@ def _initialize_b12x_kda_prefill(self, vllm_config: VllmConfig) -> None: "The b12x KDA prefill backend requires the b12x package." ) device = torch.device(current_platform.current_device()) + if coalescing: + if "max_checkpoints" not in getattr(api.Caps, "__dataclass_fields__", {}): + raise RuntimeError( + "KDA coalescing requires B12X four-checkpoint support" + ) + properties = torch.cuda.get_device_properties(device) + if ( + properties.name.strip().lower() != "nvidia gb10" + or (properties.major, properties.minor) != (12, 1) + or properties.multi_processor_count != 48 + ): + raise ValueError("KDA coalescing requires NVIDIA GB10 / SM121 / 48 SMs") scheduler_config = vllm_config.scheduler_config self._b12x_prefill_api = api self._b12x_prefill_max_tokens = int(scheduler_config.max_num_batched_tokens) self._b12x_prefill_max_seqs = int(scheduler_config.max_num_seqs) max_seqs = self._b12x_prefill_max_seqs + checkpoint_shape = ( + (max_seqs, self._b12x_prefill_checkpoint_capacity) + if coalescing + else (max_seqs,) + ) self.register_buffer( "_b12x_prefill_num_seqs", torch.zeros(1, dtype=torch.int32, device=device), @@ -754,12 +790,14 @@ def _initialize_b12x_kda_prefill(self, vllm_config: VllmConfig) -> None: ) self.register_buffer( "_b12x_prefill_null_indices", - torch.full((max_seqs,), NULL_BLOCK_ID, dtype=torch.int32, device=device), + torch.full( + checkpoint_shape, NULL_BLOCK_ID, dtype=torch.int32, device=device + ), persistent=False, ) self.register_buffer( "_b12x_prefill_zero_offsets", - torch.zeros(max_seqs, dtype=torch.int32, device=device), + torch.zeros(checkpoint_shape, dtype=torch.int32, device=device), persistent=False, ) self.b12x_warmup_provider = _B12xKdaPrefillWarmup() @@ -781,7 +819,16 @@ def _make_b12x_kda_prefill_plan(self, max_state_slots: int): qk_l2norm=True, checkpoint_export=True, null_state_index=NULL_BLOCK_ID, - metadata_validation="trusted", + metadata_validation=( + "transactional" + if self._b12x_prefill_checkpoint_capacity > 1 + else "trusted" + ), + **( + {"max_checkpoints": self._b12x_prefill_checkpoint_capacity} + if self._b12x_prefill_checkpoint_capacity > 1 + else {} + ), ) ) @@ -821,7 +868,7 @@ def _run_b12x_kda_prefill( checkpoint: Any | None, recurrent_state: torch.Tensor, output: torch.Tensor, - ) -> None: + ) -> torch.Tensor | None: """Run packed KDA prefill straight against the recurrent-state pool. The op reads each request's initial state and writes its final state, @@ -903,6 +950,13 @@ def _run_b12x_kda_prefill( max_live_tokens=num_tokens, max_live_seqs=num_requests, ) + if self._b12x_prefill_checkpoint_capacity > 1: + # Consumers share this stream and must not publish invalid states. + torch._assert_async( + binding.error_code == 0, "invalid recurrent checkpoint metadata" + ) + return binding.error_code + return None def _store_kda_conv_checkpoint( self, @@ -912,11 +966,30 @@ def _store_kda_conv_checkpoint( recurrent_state: torch.Tensor, query_start_loc: torch.Tensor, checkpoint: Any, + error_code: torch.Tensor | None = None, ) -> None: """Store the convolution history at each request's checkpoint offset.""" - state_len = conv_state.shape[-1] + # Speculative storage includes future-token cells. A reusable prefill + # checkpoint stores only the causal kernel's history in the first cells. + state_len = self.conv1d.weight.shape[-1] - 1 + if not 1 <= state_len <= conv_state.shape[-1]: + raise ValueError("Convolution checkpoint history exceeds state storage") width = mixed_qkv.shape[-1] store_block_size = 256 + offsets = checkpoint.checkpoint_offsets + indices = checkpoint.state_indices + checkpoint_count = offsets.shape[1] if offsets.ndim == 2 else 1 + if ( + checkpoint_count not in (1, 2, 4) + or not offsets.is_contiguous() + or not indices.is_contiguous() + ): + raise ValueError( + "Checkpoint convolution metadata must be contiguous " + "with capacity one, two or four" + ) + if tuple(offsets.shape) != tuple(indices.shape): + raise ValueError("Checkpoint offset and destination shapes differ") _store_cache_checkpoints_kernel[ ( checkpoint.checkpoint_offsets.numel(), @@ -944,6 +1017,9 @@ def _store_kda_conv_checkpoint( NULL_BLOCK_ID, store_block_size, False, + checkpoint_count, + error_code, + error_code is not None, ) def rearrange_mixed_qkv( @@ -1421,7 +1497,7 @@ def _prefill_conv( num_prefill_tokens = int(q_ns.shape[1]) b12x_scratch, b12x_out = self._get_b12x_prefill_workspace() b12x_out = b12x_out[:num_prefill_tokens] - self._run_b12x_kda_prefill( + checkpoint_error = self._run_b12x_kda_prefill( scratch=b12x_scratch, q=q_ns[0], k=k_ns[0], @@ -1445,6 +1521,7 @@ def _prefill_conv( recurrent_state=recurrent_state, query_start_loc=prefill_query_start_loc, checkpoint=prefill_checkpoint, + error_code=checkpoint_error, ) elif self.kda_prefill_backend == "flashkda": assert initial_state is not None diff --git a/vllm/models/glm5next/model_state.py b/vllm/models/glm5next/model_state.py index d71219d0c9e2..851397fbfca5 100644 --- a/vllm/models/glm5next/model_state.py +++ b/vllm/models/glm5next/model_state.py @@ -12,6 +12,7 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.triton_utils import tl, triton +from vllm.v1.core.recurrent_prefill_checkpoint import checkpoint_plan_rows from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.utils import CpuGpuBuffer @@ -267,6 +268,7 @@ def prepare_attn( attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, for_capture: bool = False, + recurrent_prefill_checkpoint_plans: dict | None = None, ) -> dict[str, Any]: # This is the MambaHybridModelState construction with only the metadata # object specialized. Keeping it package-local avoids a GLM hook in the @@ -331,7 +333,14 @@ def prepare_attn( block_tables, ) + checkpoint_rows = checkpoint_plan_rows( + input_batch, + recurrent_prefill_checkpoint_plans, + num_reqs, + for_capture=for_capture, + ) model_metadata = Glm5NextAttnMetadata( + recurrent_prefill_checkpoint_plans_cpu=checkpoint_rows, is_prefilling=is_prefilling, num_accepted_tokens=num_accepted_tokens, num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, @@ -358,6 +367,21 @@ def prepare_attn( for_cudagraph_capture=for_capture, rswa_prefix_lens=input_batch.prompt_lens, ) + if checkpoint_rows is not None: + gdn_metadata = [ + item + for item in attn_metadata.values() + if hasattr(item, "prefill_checkpoint") + ] + if not gdn_metadata or any( + item.prefill_checkpoint is None + or item.prefill_checkpoint.required_mask is None + for item in gdn_metadata + ): + raise RuntimeError( + "GLM checkpoint plans require GDN export metadata " + "on every recurrent layer" + ) if self.recoverssm is not None: self.recoverssm.record_step( attn_metadata, diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 32f76190050f..c7829d1ae228 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -423,6 +423,9 @@ class CommonAttentionMetadata: (num_computed_tokens < num_prompt_tokens). Used by some backends to distinguish actual decodes from short extends.""" + recurrent_prefill_checkpoint_plans_cpu: ( + list[tuple[int, int, tuple[int, ...]] | None] | None + ) = None seq_lens_cpu_upper_bound: torch.Tensor | None = None """(batch_size,) CPU upper bound on seq_lens. Precise for prefill rows and for all rows outside async spec decode; optimistic for async-spec @@ -562,6 +565,9 @@ def unpadded( is_prefilling=maybe_slice_reqs(self.is_prefilling), rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens), replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu), + recurrent_prefill_checkpoint_plans_cpu=maybe_slice_reqs( + self.recurrent_prefill_checkpoint_plans_cpu + ), ) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index ac494843b93c..99afaa20921d 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -23,6 +23,10 @@ mamba_get_block_table_tensor, split_decodes_and_prefills, ) +from vllm.v1.core.recurrent_prefill_checkpoint import ( + COALESCED_CHECKPOINT_CAPACITY, + checkpoint_metadata, +) from vllm.v1.kv_cache_interface import MambaSpec @@ -42,7 +46,7 @@ def is_ssm(cls) -> bool: @dataclass class GDNPrefillCheckpointMetadata: - """One recurrent-state checkpoint inside each packed prefill sequence. + """Bounded recurrent-state checkpoints inside each packed prefill sequence. ``checkpoint_offsets`` are relative to the corresponding packed query. ``request_rows`` and ``block_table_columns`` identify the cache slots that @@ -54,6 +58,7 @@ class GDNPrefillCheckpointMetadata: state_indices: torch.Tensor request_rows: torch.Tensor block_table_columns: torch.Tensor + required_mask: torch.Tensor | None = None @dataclass @@ -574,9 +579,7 @@ def build( # type: ignore[override] and self.kv_cache_spec.num_prefill_checkpoint_blocks > 0 and self.vllm_config.cache_config.mamba_cache_mode == "align" ): - # FlashKDA can materialize one state at a cache-block boundary - # without splitting the target-model forward. Only prefill rows - # participate, in the same order as prefill_query_start_loc. + # Checkpoint rows use the packed prefill query order. assert m.seq_lens_cpu_upper_bound is not None all_query_lens = query_start_loc_cpu.diff().tolist() if spec_sequence_masks_cpu is None: @@ -591,23 +594,36 @@ def build( # type: ignore[override] seq_lens = m.seq_lens_cpu_upper_bound.tolist() block_size = self.kv_cache_spec.block_size - checkpoint_offsets: list[int] = [] - checkpoint_columns: list[int] = [] + capacity = min( + self.kv_cache_spec.num_prefill_checkpoint_blocks, + COALESCED_CHECKPOINT_CAPACITY, + ) + plans = getattr(m, "recurrent_prefill_checkpoint_plans_cpu", None) + if plans is not None and len(plans) < len(all_query_lens): + raise ValueError("checkpoint plan rows do not cover the packed batch") + checkpoint_offsets: list[list[int]] = [] + checkpoint_columns: list[list[int]] = [] + checkpoint_required: list[list[bool]] = [] for row in request_rows: query_len = all_query_lens[row] seq_len = seq_lens[row] - offset = seq_len // block_size * block_size - (seq_len - query_len) - valid = ( - seq_len % block_size != 0 - and 0 < offset < query_len - # FlashKDA checkpoint outputs are produced on its - # 16-token recurrence boundary. - and offset % 16 == 0 + offsets, columns = checkpoint_metadata( + None if plans is None else plans[row], + seq_len - query_len, + seq_len, + block_size, + capacity, ) - checkpoint_offsets.append(offset if valid else 0) - checkpoint_columns.append(seq_len // block_size - 1 if valid else -1) + checkpoint_offsets.append(offsets) + checkpoint_columns.append(columns) + required = [ + plans is not None and plans[row] is not None and column >= 0 + for column in columns + ] + checkpoint_required.append(required) - if any(checkpoint_offsets): + any_checkpoint = any(any(row) for row in checkpoint_offsets) + if any_checkpoint: checkpoint_offsets_tensor = async_tensor_h2d( checkpoint_offsets, dtype=torch.int32, @@ -623,19 +639,42 @@ def build( # type: ignore[override] dtype=torch.int64, device=query_start_loc.device, ) + if capacity == 1: + checkpoint_offsets_tensor = checkpoint_offsets_tensor[:, 0] + checkpoint_columns_tensor = checkpoint_columns_tensor[:, 0] checkpoint_state_indices = m.block_table_tensor[ - request_rows_tensor, checkpoint_columns_tensor + request_rows_tensor[:, None] + if capacity > 1 + else request_rows_tensor, + checkpoint_columns_tensor, ] checkpoint_state_indices = torch.where( checkpoint_columns_tensor >= 0, checkpoint_state_indices, NULL_BLOCK_ID, ) + has_required = any(any(row) for row in checkpoint_required) + required_mask = None + if has_required: + required_mask = async_tensor_h2d( + checkpoint_required, + dtype=torch.bool, + device=query_start_loc.device, + ) + if capacity == 1: + required_mask = required_mask[:, 0] + torch._assert_async( + torch.all( + ~required_mask | (checkpoint_state_indices != NULL_BLOCK_ID) + ), + "planned recurrent checkpoint refers to a NULL state block", + ) prefill_checkpoint = GDNPrefillCheckpointMetadata( checkpoint_offsets=checkpoint_offsets_tensor, state_indices=checkpoint_state_indices, request_rows=request_rows_tensor, block_table_columns=checkpoint_columns_tensor, + required_mask=required_mask, ) # Function code counted on either presency non-spec decode or spec decode, @@ -802,7 +841,9 @@ def update_block_table( prefill_checkpoint = metadata.prefill_checkpoint if prefill_checkpoint is not None: checkpoint_state_indices = blk_table[ - prefill_checkpoint.request_rows, + prefill_checkpoint.request_rows[:, None] + if prefill_checkpoint.block_table_columns.ndim == 2 + else prefill_checkpoint.request_rows, prefill_checkpoint.block_table_columns, ] checkpoint_state_indices = torch.where( @@ -810,6 +851,14 @@ def update_block_table( checkpoint_state_indices, NULL_BLOCK_ID, ) + if prefill_checkpoint.required_mask is not None: + torch._assert_async( + torch.all( + ~prefill_checkpoint.required_mask + | (checkpoint_state_indices != NULL_BLOCK_ID) + ), + "planned recurrent checkpoint refers to a NULL state block", + ) prefill_checkpoint = replace( prefill_checkpoint, state_indices=checkpoint_state_indices, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index e68f5458d012..728f6db20e47 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -4,7 +4,7 @@ import itertools from collections.abc import Sequence from dataclasses import dataclass -from typing import Literal, overload +from typing import Any, Literal, overload from vllm.distributed.kv_events import BlockStored, KVCacheEvent from vllm.logger import init_logger @@ -23,6 +23,12 @@ ) from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy +from vllm.v1.core.recurrent_prefill_checkpoint import ( + CheckpointPlan, + continuation_layout, + validate_plan, +) +from vllm.v1.core.single_type_kv_cache_manager import MambaManager from vllm.v1.kv_cache_interface import ( AttentionSpec, CrossAttentionSpec, @@ -384,6 +390,68 @@ def get_computed_blocks_for_connector( # Per-group lookups do not detect an uncached shared prefix (boundary 0). return blocks, num_local, 0, min(per_group_hits) < num_local + def _allocate_checkpoint_slots( + self, request: Request, plan: CheckpointPlan, **kwargs: Any + ) -> KVCacheBlocks | None: + """Keep checkpoint allocation plans scoped to one admission attempt.""" + start, end, targets = plan + cached = kwargs["new_computed_blocks"] + if ( + request.num_computed_tokens != start + or start + kwargs["num_new_tokens"] != end + or end > request.num_prompt_tokens + or request.num_tokens != request.num_prompt_tokens + or kwargs["num_new_computed_tokens"] + or kwargs["num_external_computed_tokens"] + or kwargs["delay_cache_blocks"] + or kwargs["num_encoder_tokens"] + or (cached is not None and any(cached.blocks)) + or request.num_preemptions + or request.spec_token_ids + or request.has_encoder_inputs + or request.resumable + or request.use_boundary_checkpoints + ): + raise ValueError("checkpoint allocation requires a cold pure prompt chunk") + managers = [ + manager + for manager in self.coordinator.single_type_managers + if isinstance(manager, MambaManager) + ] + if not managers: + raise ValueError("checkpoint allocation requires recurrent cache groups") + for manager in managers: + assert isinstance(manager.kv_cache_spec, MambaSpec) + validate_plan(plan, start, end, manager.block_size) + if ( + manager.mamba_cache_mode != "align" + or manager.kv_cache_spec.num_prefill_checkpoint_blocks < len(targets) + or request.request_id in manager._partial_hit_reqs + ): + raise ValueError( + "recurrent cache group cannot produce the checkpoint plan" + ) + if start: + if request.request_id not in manager._allocated_block_reqs: + raise ValueError("continuation has no retained running state") + continuation_layout( + manager.req_to_blocks[request.request_id], + plan, + manager.block_size, + manager.num_speculative_blocks, + ) + elif manager.req_to_blocks.get(request.request_id): + raise ValueError( + "cold checkpoint request already owns recurrent blocks" + ) + try: + for manager in managers: + manager._planned_recurrent_checkpoints[request.request_id] = plan + return self.allocate_slots(request, **kwargs) + finally: + for manager in managers: + manager._planned_recurrent_checkpoints.pop(request.request_id, None) + def allocate_slots( self, request: Request, @@ -397,6 +465,7 @@ def allocate_slots( full_sequence_must_fit: bool = False, reserved_blocks: int = 0, has_scheduled_reqs: bool = True, + recurrent_checkpoint_plan: CheckpointPlan | None = None, ) -> KVCacheBlocks | None: """Add slots for a request with new tokens to append. @@ -480,6 +549,21 @@ def allocate_slots( Returns: A list of new allocated blocks. """ + if recurrent_checkpoint_plan is not None: + return self._allocate_checkpoint_slots( + request, + recurrent_checkpoint_plan, + num_new_tokens=num_new_tokens, + num_new_computed_tokens=num_new_computed_tokens, + new_computed_blocks=new_computed_blocks, + num_lookahead_tokens=num_lookahead_tokens, + num_external_computed_tokens=num_external_computed_tokens, + delay_cache_blocks=delay_cache_blocks, + num_encoder_tokens=num_encoder_tokens, + full_sequence_must_fit=full_sequence_must_fit, + reserved_blocks=reserved_blocks, + has_scheduled_reqs=has_scheduled_reqs, + ) # When loading KV data asynchronously, we may have zero new tokens to # compute while still allocating slots for externally computed tokens. if num_new_tokens == 0 and num_external_computed_tokens == 0: diff --git a/vllm/v1/core/recurrent_prefill_checkpoint.py b/vllm/v1/core/recurrent_prefill_checkpoint.py new file mode 100644 index 000000000000..f64180cab63e --- /dev/null +++ b/vllm/v1/core/recurrent_prefill_checkpoint.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Checkpoint plans and physical ownership for bounded KDA prefill.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.core.kv_cache_utils import KVCacheBlock + from vllm.v1.worker.gpu.input_batch import InputBatch + +CheckpointPlan = tuple[int, int, tuple[int, ...]] + +# DCP4 retention preserves two fine-grid and two scheduler-grid states. +COALESCED_CHECKPOINT_CAPACITY = 4 + + +def validate_plan( + plan: CheckpointPlan, start: int, end: int, block_size: int +) -> tuple[int, ...]: + planned_start, planned_end, targets = plan + if (planned_start, planned_end) != (start, end): + raise ValueError( + "recurrent checkpoint plan does not match the actual query span" + ) + if ( + not isinstance(targets, tuple) + or not 1 <= len(targets) <= COALESCED_CHECKPOINT_CAPACITY + or targets != tuple(sorted(set(targets))) + ): + raise ValueError( + "recurrent checkpoint targets must be one to four sorted unique positions" + ) + if any( + type(p) is not int or not start < p < end or p % block_size or (p - start) % 16 + for p in targets + ): + raise ValueError( + "recurrent checkpoint target is not a representable interior state" + ) + return targets + + +def prefill_checkpoint_plan( + *, + start: int, + end: int, + prompt: int, + num_tokens: int, + block_size: int, + publications: tuple[int, ...], + shared_prefix_boundary: int = 0, +) -> CheckpointPlan | None: + """Export exact retained states inside a bounded pure-prompt chunk.""" + if ( + start < 0 + or start % block_size + or not start < end <= prompt + or num_tokens != prompt + or end - start > 8192 + or end % block_size + ): + return None + required = {position for position in publications if start < position < end} + if start < shared_prefix_boundary < end: + required.add(shared_prefix_boundary // block_size * block_size) + required.discard(0) + plan = (start, end, tuple(sorted(required))) + try: + validate_plan(plan, start, end, block_size) + except ValueError: + return None + return plan + + +def checkpoint_metadata( + plan: CheckpointPlan | None, start: int, end: int, block_size: int, capacity: int +) -> tuple[list[int], list[int]]: + if capacity not in (1, 2, COALESCED_CHECKPOINT_CAPACITY): + raise ValueError("unsupported recurrent checkpoint capacity") + if plan is not None: + targets = validate_plan(plan, start, end, block_size) + if capacity < len(targets): + raise ValueError("metadata checkpoint capacity smaller than scheduled plan") + else: + boundary = end // block_size * block_size + targets = ( + (boundary,) + if end % block_size + and start < boundary < end + and (boundary - start) % 16 == 0 + else () + ) + return ( + [p - start for p in targets] + [0] * (capacity - len(targets)), + [p // block_size - 1 for p in targets] + [-1] * (capacity - len(targets)), + ) + + +def continuation_layout( + blocks: Sequence[KVCacheBlock], + plan: CheckpointPlan, + block_size: int, + speculative_blocks: int, +) -> tuple[int, tuple[int, ...]]: + """Return retained column indices, -1 for NULL and -2 for allocation. + + Checkpoint columns retain their physical pages. Unused private speculative + pages move to appended columns so worker block tables remain append-only. + """ + start, end, targets = plan + validate_plan(plan, start, end, block_size) + if start <= 0 or start % block_size or end % block_size: + raise ValueError( + "continuation requires aligned nonzero source and final states" + ) + source = start // block_size - 1 + final = end // block_size - 1 + if speculative_blocks < 0 or len(blocks) != source + 1 + speculative_blocks: + raise ValueError("continuation table does not end at its expected reserve") + if blocks[source].is_null or blocks[source].ref_cnt < 1: + raise ValueError("continuation source is not retained") + reserve_columns = range(source + 1, len(blocks)) + for column in reserve_columns: + block = blocks[column] + if block.is_null or block.ref_cnt != 1 or block.block_hash is not None: + raise ValueError("continuation reserve must be private and unhashed") + owned = [block.block_id for block in blocks if not block.is_null] + if len(owned) != len(set(owned)): + raise ValueError("recurrent table aliases physical ownership") + outputs = [p // block_size - 1 for p in targets] + [final] + reserves = list(range(final + 1, final + 1 + speculative_blocks)) + desired = set(outputs + reserves) + if len(desired) != len(outputs) + len(reserves): + raise ValueError("continuation outputs alias") + layout = list(range(len(blocks))) + [-1] * ( + final + 1 + speculative_blocks - len(blocks) + ) + movable = [] + for column in reserve_columns: + if column not in desired: + movable.append(column) + layout[column] = -1 + for column in outputs: + if column >= len(blocks): + layout[column] = -2 + for column in reserves: + if column >= len(blocks): + layout[column] = movable.pop(0) if movable else -2 + if movable or layout.count(-2) != len(targets) + 1: + raise ValueError("continuation reserve accounting differs") + for column in desired: + if column < len(blocks) and layout[column] != column: + raise ValueError("continuation would replace an existing worker column") + return source, tuple(layout) + + +def validate_coalescing_config(config: VllmConfig) -> bool: + """Validate the opt-in serving contract without querying a GPU.""" + from vllm import envs + + if not envs.VLLM_B12X_KDA_PREFILL_COALESCING: + return False + model = config.model_config + cache = config.cache_config + parallel = config.parallel_config + scheduler = config.scheduler_config + additional = config.additional_config + spec = config.speculative_config + supported = ( + model is not None + and model.hf_text_config.model_type in ("glm5_next", "glm5_next_text") + and str(model.dtype) == "torch.bfloat16" + and config.use_v2_model_runner + and config.lora_config is None + and not model.enable_sleep_mode + and not model.enable_return_routed_experts + and parallel.tensor_parallel_size == 4 + and parallel.decode_context_parallel_size in (1, 2, 4) + and parallel.pipeline_parallel_size == 1 + and parallel.data_parallel_size == 1 + and parallel.prefill_context_parallel_size == 1 + and not parallel.enable_expert_parallel + and not parallel.enable_eplb + and scheduler.max_num_batched_tokens == 8192 + and scheduler.max_num_scheduled_tokens in (None, 8192) + and scheduler.long_prefill_token_threshold in (0, 8192) + and scheduler.fairness_engine is None + and scheduler.enable_chunked_prefill + and cache.enable_prefix_caching + and cache.mamba_cache_mode == "align" + and cache.prefix_cache_retention_interval == 0 + and not config.use_request_boundary_checkpoints + and isinstance(additional, dict) + and additional.get("kda_prefill_backend") == "b12x" + and ( + spec is None + or ( + spec.method == "mtp" + and spec.num_speculative_tokens == 3 + and not spec.uses_dynamic_speculative_decoding() + ) + ) + ) + if not supported: + raise ValueError( + "VLLM_B12X_KDA_PREFILL_COALESCING requires GLM5Next BF16, V2, " + "TP4 with DCP1/2/4, PP1/DP1, B12X KDA, an 8192-token scheduler budget, " + "align-mode prefix caching with retention interval 0, and static " + "MTP3 or no speculation; LoRA, EP, PCP and fairness engines are unsupported" + ) + return True + + +def checkpoint_plan_rows( + input_batch: InputBatch, + plans: dict[str, CheckpointPlan] | None, + num_reqs: int, + *, + for_capture: bool = False, +) -> list[CheckpointPlan | None] | None: + """Map scheduler plans to the worker's packed request order.""" + if not plans: + return None + if for_capture: + raise ValueError("recurrent checkpoint plans cannot enter graph capture") + if ( + len(input_batch.req_ids) != input_batch.num_reqs + or num_reqs < input_batch.num_reqs + ): + raise ValueError("checkpoint request counts differ from the worker batch") + positions = {request_id: row for row, request_id in enumerate(input_batch.req_ids)} + if len(positions) != input_batch.num_reqs: + raise ValueError("checkpoint batch contains duplicate request IDs") + rows: list[CheckpointPlan | None] = [None] * num_reqs + for request_id, plan in plans.items(): + row = positions.get(request_id) + if row is None: + raise ValueError("checkpoint request is absent from the worker batch") + start, end, targets = plan + validate_plan(plan, start, end, 16) + if ( + type(start) is not int + or type(end) is not int + or start < 0 + or end <= start + or int(input_batch.num_computed_tokens_np[row]) != start + or int(input_batch.num_scheduled_tokens[row]) != end - start + or int( + input_batch.query_start_loc_np[row + 1] + - input_batch.query_start_loc_np[row] + ) + != end - start + ): + raise ValueError("checkpoint plan differs from the worker query span") + rows[row] = (start, end, targets) + return rows diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index fec086b92be3..0f39751f6841 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -262,6 +262,9 @@ class SchedulerOutput: scheduled_encoder_input_stats: ScheduledEncoderInputStats | None = None # This batch samples saved final hidden states without a target forward. boundary_logits_only: bool = False + recurrent_prefill_checkpoint_plans: ( + dict[str, tuple[int, int, tuple[int, ...]]] | None + ) = None # Request IDs that are preempted in this step. # Only used for v2 model runner. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index b5a18109a647..0e3f694f274f 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -39,6 +39,13 @@ from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import KVCacheBlock +from vllm.v1.core.recurrent_prefill_checkpoint import ( + COALESCED_CHECKPOINT_CAPACITY, + CheckpointPlan, + continuation_layout, + prefill_checkpoint_plan, + validate_coalescing_config, +) from vllm.v1.core.sched.compute_fairness import ( ComputeServiceClass, PrefillComputeShareController, @@ -62,6 +69,7 @@ create_request_queue, ) from vllm.v1.core.sched.utils import check_stop, remove_all +from vllm.v1.core.single_type_kv_cache_manager import MambaManager from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.metrics.perf import ModelMetrics, PerfStats @@ -113,6 +121,9 @@ def _build_kv_connector_block_state( class Scheduler(SchedulerInterface): + _kda_coalescing_enabled = False + _kda_coalescing_exclusive = False + def __init__( self, vllm_config: VllmConfig, @@ -510,6 +521,133 @@ def __init__( # In-flight requests still prefilling (prefill chunks + in-progress # async KV loads). Their remaining-block reservation gates async loads. self._inflight_prefills: set[Request] = set() + self._kda_coalescing_enabled = validate_coalescing_config(vllm_config) + self._kda_coalescing_origins: dict[str, Request] = {} + self._kda_coalescing_exclusive = False + if self._kda_coalescing_enabled: + recurrent_groups = [ + group.kv_cache_spec + for group in kv_cache_config.kv_cache_groups + if isinstance(group.kv_cache_spec, MambaSpec) + ] + if not recurrent_groups or any( + spec.num_prefill_checkpoint_blocks != COALESCED_CHECKPOINT_CAPACITY + or spec.block_size != self.cache_config.block_size + for spec in recurrent_groups + ): + raise ValueError( + "KDA coalescing requires four-checkpoint recurrent groups " + "on one block grid" + ) + managers = self.kv_cache_manager.coordinator.single_type_managers + logger.info( + "KDA_PREFILL_COALESCING configured capacity=%d grids=%s " + "(physical, lookup, scheduler)", + COALESCED_CHECKPOINT_CAPACITY, + sorted( + { + ( + manager.block_size, + manager.hit_alignment_tokens, + manager.scheduler_block_size, + ) + for manager in managers + if isinstance(manager, MambaManager) + } + ), + ) + + def _recurrent_checkpoint_plan( + self, request: Request, start: int, end: int + ) -> CheckpointPlan | None: + """Plan exact retained states for one cold, unmixed prompt chunk.""" + if ( + not self._kda_coalescing_enabled + or not self._kda_coalescing_exclusive + or request.use_boundary_checkpoints + or request.num_preemptions + or request.has_encoder_inputs + or request.resumable + or request.num_tokens != request.num_prompt_tokens + or request.spec_token_ids + or request.num_computed_tokens != start + ): + return None + if ( + start + and self._kda_coalescing_origins.get(request.request_id) is not request + ): + return None + managers = [ + manager + for manager in self.kv_cache_manager.coordinator.single_type_managers + if isinstance(manager, MambaManager) + ] + boundaries = [request.num_prompt_tokens - 1] + if request.shared_prefix_boundary: + boundaries.append(request.shared_prefix_boundary) + publications = tuple( + sorted( + { + position + for manager in managers + for position in manager._expand_reachable_boundaries(boundaries) + } + ) + ) + plan = prefill_checkpoint_plan( + start=start, + end=end, + prompt=request.num_prompt_tokens, + num_tokens=request.num_tokens, + block_size=self.cache_config.block_size, + publications=publications, + shared_prefix_boundary=request.shared_prefix_boundary, + ) + if plan is None: + return None + for manager in managers: + assert isinstance(manager.kv_cache_spec, MambaSpec) + if ( + request.request_id in manager._partial_hit_reqs + or manager.kv_cache_spec.num_prefill_checkpoint_blocks < len(plan[2]) + ): + return None + if start: + if request.request_id not in manager._allocated_block_reqs: + return None + try: + continuation_layout( + manager.req_to_blocks[request.request_id], + plan, + manager.block_size, + manager.num_speculative_blocks, + ) + except ValueError: + return None + elif manager.req_to_blocks.get(request.request_id): + return None + return plan + + def _record_coalescing_origin( + self, request: Request, computed: int, local: int, external: int, loading: bool + ) -> None: + if not self._kda_coalescing_enabled: + return + self._kda_coalescing_origins.pop(request.request_id, None) + if ( + self._kda_coalescing_enabled + and not loading + and computed == 0 + and local == 0 + and external == 0 + and not request.num_preemptions + and request.status == RequestStatus.WAITING + and not request.resumable + and not request.has_encoder_inputs + and request.num_tokens == request.num_prompt_tokens + ): + self._kda_coalescing_origins[request.request_id] = request def _mamba_block_aligned_split( self, @@ -531,6 +669,14 @@ def _mamba_block_aligned_split( + num_new_local_computed_tokens + num_external_computed_tokens ) + if ( + getattr(self, "_kda_coalescing_enabled", False) + and not num_new_local_computed_tokens + and not num_external_computed_tokens + and self._recurrent_checkpoint_plan(request, start, start + num_new_tokens) + is not None + ): + return num_new_tokens if request.use_boundary_checkpoints: # Running-state migration still happens in the worker, but these # requests publish only semantic endpoints. Stop exactly at the @@ -744,6 +890,10 @@ def _has_waiting_boundary_logits(self) -> bool: def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.current_step += 1 + self._kda_coalescing_exclusive = ( + self._kda_coalescing_enabled + and len(self.running) + len(self.waiting) + len(self.skipped_waiting) == 1 + ) # NOTE(woosuk) on the scheduling algorithm: # There's no "decoding phase" nor "prefill phase" in the scheduler. # Each request just has the num_computed_tokens and @@ -762,6 +912,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: req_to_new_blocks: dict[str, KVCacheBlocks] = {} num_scheduled_tokens: dict[str, int] = {} + recurrent_checkpoint_plans: dict[str, CheckpointPlan] = {} token_budget = self.max_num_scheduled_tokens spec = self.vllm_config.speculative_config separate_draft_input_tokens = 0 @@ -1064,10 +1215,16 @@ def schedule_running_requests( # Schedule newly needed KV blocks for the request. with record_function_or_nullcontext("schedule: allocate_slots"): while True: + checkpoint_plan = self._recurrent_checkpoint_plan( + request, + request.num_computed_tokens, + request.num_computed_tokens + num_new_tokens, + ) new_blocks = self.kv_cache_manager.allocate_slots( request, num_new_tokens, num_lookahead_tokens=self.num_lookahead_tokens, + recurrent_checkpoint_plan=checkpoint_plan, ) if new_blocks is not None: @@ -1142,6 +1299,8 @@ def schedule_running_requests( request_id = request.request_id req_to_new_blocks[request_id] = new_blocks num_scheduled_tokens[request_id] = num_new_tokens + if checkpoint_plan is not None: + recurrent_checkpoint_plans[request_id] = checkpoint_plan token_budget -= num_new_tokens input_budget -= num_new_tokens + draft_slots draft_input_budget -= separate_draft_input_tokens @@ -1556,6 +1715,17 @@ def schedule_running_requests( # avoid deadlock and predictable preemptions. reserved_blocks = self._inflight_prefill_reserved_blocks() + checkpoint_plan = ( + None + if load_kv_async + or num_new_local_computed_tokens + or num_external_computed_tokens + else self._recurrent_checkpoint_plan( + request, + num_computed_tokens, + num_computed_tokens + num_new_tokens, + ) + ) new_blocks = self.kv_cache_manager.allocate_slots( request, num_new_tokens, @@ -1568,6 +1738,7 @@ def schedule_running_requests( full_sequence_must_fit=self.scheduler_reserve_full_isl, reserved_blocks=reserved_blocks, has_scheduled_reqs=bool(self.running), + recurrent_checkpoint_plan=checkpoint_plan, ) if new_blocks is None: @@ -1579,6 +1750,15 @@ def schedule_running_requests( self.encoder_cache_manager.free(request) break + if checkpoint_plan is not None: + recurrent_checkpoint_plans[request_id] = checkpoint_plan + self._record_coalescing_origin( + request, + num_computed_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + load_kv_async, + ) # KVTransfer: the connector uses this info to determine # if a load is needed. Note that # This information is used to determine if a load is @@ -1903,6 +2083,11 @@ def schedule_running_requests( ) scheduler_output = SchedulerOutput( + recurrent_prefill_checkpoint_plans={ + request_id: plan + for request_id, plan in recurrent_checkpoint_plans.items() + if request_id in num_scheduled_tokens + }, scheduled_new_reqs=new_reqs_data, boundary_logits_only=bool( new_reqs_data @@ -1980,6 +2165,13 @@ def schedule_running_requests( scheduler_output.compute_service_class, contended=scheduler_output.compute_contention, ) + if scheduler_output.recurrent_prefill_checkpoint_plans: + for plan in scheduler_output.recurrent_prefill_checkpoint_plans.values(): + logger.info_once( + "KDA_PREFILL_COALESCING scheduled span=%d checkpoints=%d", + plan[1] - plan[0], + len(plan[2]), + ) return scheduler_output def record_compute_time( @@ -3308,6 +3500,8 @@ def _free_request_blocks(self, request: Request): """Free the request's KV blocks, deferring the return to the block pool when an in-flight GPU step may still write them. """ + if self._kda_coalescing_enabled: + self._kda_coalescing_origins.pop(request.request_id, None) if not self.defer_block_free or ( # Last scheduled step already processed: no in-flight write remains # (always the case for a normal finish), so free now. diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 6d5a5e9808ef..9cd78ce976c8 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -16,6 +16,11 @@ KVCacheBlock, resolve_block_hashes, ) +from vllm.v1.core.recurrent_prefill_checkpoint import ( + CheckpointPlan, + continuation_layout, + validate_plan, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, ChunkedLocalAttentionSpec, @@ -1633,6 +1638,7 @@ def __init__( self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks self.cached_blocks_this_step: set[BlockHashWithGroupId] = set() + self._planned_recurrent_checkpoints: dict[str, CheckpointPlan] = {} if self.mamba_cache_mode == "align": # Mapping from request ID to the index of the block # allocated in the previous step @@ -1854,6 +1860,43 @@ def get_num_blocks_to_allocate( apply_admission_cap: bool = False, ) -> int: assert isinstance(self.kv_cache_spec, MambaSpec) + plan = self._planned_recurrent_checkpoints.get(request_id) + if plan is not None: + start, end, targets = plan + validate_plan(plan, start, end, self.block_size) + if apply_admission_cap: + # Full-prompt admission bounds the recurrent resident peak; + # actual allocation below is restricted to this chunk's span. + resident = ( + 2 + + self.num_speculative_blocks + + self.kv_cache_spec.num_prefill_checkpoint_blocks + ) + owned = sum( + not block.is_null for block in self.req_to_blocks[request_id] + ) + return max(resident - owned, 0) + if ( + new_computed_blocks + or total_computed_tokens != start + or num_tokens_main_model != end + ): + raise ValueError( + "checkpoint admission differs from the planned query span" + ) + if start: + _, layout = continuation_layout( + self.req_to_blocks[request_id], + plan, + self.block_size, + self.num_speculative_blocks, + ) + return layout.count(-2) + if self.req_to_blocks.get(request_id): + raise ValueError( + "cold checkpoint allocation requires an empty block table" + ) + return len(targets) + 1 + self.num_speculative_blocks if ( len(new_computed_blocks) > 0 and new_computed_blocks[-1].block_hash in self.cached_blocks_this_step @@ -1926,6 +1969,47 @@ def allocate_new_blocks( self, request_id: str, num_tokens: int, num_tokens_main_model: int ) -> list[KVCacheBlock]: assert isinstance(self.kv_cache_spec, MambaSpec) + plan = self._planned_recurrent_checkpoints.get(request_id) + if plan is not None: + start, end, targets = plan + if num_tokens_main_model != end: + raise ValueError("checkpoint allocation differs from the planned end") + blocks = self.req_to_blocks[request_id] + if start: + source, layout = continuation_layout( + blocks, plan, self.block_size, self.num_speculative_blocks + ) + retained = tuple(blocks) + allocated = iter(self.block_pool.get_new_blocks(layout.count(-2))) + replacement = [ + retained[index] + if index >= 0 + else self._null_block + if index == -1 + else next(allocated) + for index in layout + ] + self.last_state_block_idx[request_id] = source + blocks[:] = replacement + self._allocated_block_reqs.add(request_id) + return replacement[len(retained) :] + if blocks: + raise ValueError( + "cold checkpoint allocation requires an empty block table" + ) + final_column = cdiv(end, self.block_size) - 1 + columns = [target // self.block_size - 1 for target in targets] + columns.extend( + range(final_column, final_column + 1 + self.num_speculative_blocks) + ) + physical = self.block_pool.get_new_blocks(len(columns)) + blocks.extend( + [self._null_block] * (final_column + 1 + self.num_speculative_blocks) + ) + for column, block in zip(columns, physical): + blocks[column] = block + self._allocated_block_reqs.add(request_id) + return blocks[:] if self.mamba_cache_mode != "align": # Allocate extra `num_speculative_blocks` blocks for # speculative decoding (MTP/EAGLE) with linear attention. @@ -2058,6 +2142,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: self._allocated_block_reqs.discard(request_id) self.last_state_block_idx.pop(request_id, None) self._num_checkpoint_blocks.pop(request_id, None) + self._planned_recurrent_checkpoints.pop(request_id, None) self._producer_partial_tail_reqs.pop(request_id, None) # An offer is only guaranteed to hold committed bytes until the end # of the pass that made it. This request's blocks are going back to diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 08a828b1484f..716c2be258e4 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1747,6 +1747,20 @@ def execute_model( [g for g in groups if not isinstance(g.kv_cache_spec, MambaSpec)] for groups in attn_groups ] + checkpoint_kwargs = {} + checkpoint_plans = ( + None + if dummy_run + else scheduler_output.recurrent_prefill_checkpoint_plans + ) + if checkpoint_plans: + if batch_desc.cg_mode != CUDAGraphMode.NONE: + raise ValueError( + "coalesced recurrent checkpoints require eager prefill" + ) + checkpoint_kwargs["recurrent_prefill_checkpoint_plans"] = ( + checkpoint_plans + ) attn_metadata = self.model_state.prepare_attn( input_batch, batch_desc.cg_mode, @@ -1758,6 +1772,7 @@ def execute_model( # from the zeroed dummy block tables instead of retaining state # indices from the previous real batch. for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, + **checkpoint_kwargs, ) input_ids = input_batch.input_ids diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 622442c59bec..637aedf6247f 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -13,6 +13,7 @@ from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadataBuilder +from vllm.v1.core.recurrent_prefill_checkpoint import checkpoint_plan_rows from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.utils import CpuGpuBuffer @@ -37,13 +38,21 @@ class MambaHybridAttnMetadata(ModelSpecificAttnMetadata): is_prefilling: torch.Tensor num_accepted_tokens: torch.Tensor | None = None num_decode_draft_tokens_cpu: torch.Tensor | None = None + recurrent_prefill_checkpoint_plans_cpu: list | None = None def get_extra_common_attn_kwargs( self, kv_cache_group_id: int, num_reqs: int, ) -> dict[str, Any]: - return {"is_prefilling": self.is_prefilling[:num_reqs]} + return { + "is_prefilling": self.is_prefilling[:num_reqs], + "recurrent_prefill_checkpoint_plans_cpu": ( + None + if self.recurrent_prefill_checkpoint_plans_cpu is None + else self.recurrent_prefill_checkpoint_plans_cpu[:num_reqs] + ), + } def get_extra_attn_kwargs( self, @@ -267,6 +276,7 @@ def prepare_attn( attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, for_capture: bool = False, + recurrent_prefill_checkpoint_plans: dict | None = None, ) -> dict[str, Any]: if cudagraph_mode == CUDAGraphMode.FULL: num_reqs = input_batch.num_reqs_after_padding @@ -323,7 +333,14 @@ def prepare_attn( block_tables, ) + checkpoint_rows = checkpoint_plan_rows( + input_batch, + recurrent_prefill_checkpoint_plans, + num_reqs, + for_capture=for_capture, + ) mamba_attn_metadata = MambaHybridAttnMetadata( + recurrent_prefill_checkpoint_plans_cpu=checkpoint_rows, is_prefilling=is_prefilling, num_accepted_tokens=num_accepted_tokens, num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,