[Perf] Add opt-in DFlash2 schedules and numerical audits

benchmarks/compare_sm70_dflash2_natural_audit.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Locate the first observed target/proposal difference with real acceptance.
+
+This is a diagnostic comparison, not a quality or acceptance noninferiority
+gate. Inputs after the first differing proposal need not be the same. A first
+observed state difference still requires an operator-level causality check.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+
+from benchmarks.compare_sm70_dflash2_state_audit import (
+    sampling_difference,
+    tensor_difference,
+)
+
+
+def _load(directory: Path) -> dict:
+    records = {}
+    for path in directory.glob("*-rank*-step*.pt"):
+        row = torch.load(path, map_location="cpu", weights_only=True, mmap=True)
+        case, rank, step = (row[k] for k in ("case", "rank", "step"))
+        key = case, step, rank
+        if key in records or row.get("control") != "natural_sampling":
+            raise ValueError(f"{path}: duplicate or non-natural observation")
+        if row["phase"] != ("prefill" if step == 0 else "verify"):
+            raise ValueError(f"{path}: incorrect phase")
+        if not row.get("expected_layers") or "capture_epoch" not in row:
+            raise ValueError(f"{path}: missing snapshot provenance")
+        if rank == 0 and not torch.isfinite(row["native_logits"]).all():
+            raise ValueError(f"{path}: nonfinite native logits")
+        for field in ("aux_hidden_states", "sampling", "states"):
+            if not row.get(field):
+                raise ValueError(f"{path}: missing {field}")
+        if row["num_sampled"].numel() != 1 or row["num_rejected"].numel() != 1:
+            raise ValueError(f"{path}: expected B1 counts")
+        n = int(row["num_sampled"].item())
+        rejected = int(row["num_rejected"].item())
+        if not 1 <= n <= 8 or not 0 <= rejected <= 7:
+            raise ValueError(f"{path}: invalid acceptance counts")
+        if n > row["sampled_token_ids"].shape[1]:
+            raise ValueError(f"{path}: missing accepted output IDs")
+        proposal = directory / f"proposal-{case}-tp{rank}-forward{step}.pt"
+        if not proposal.exists():
+            raise ValueError(f"{path}: missing proposal observation")
+        draft = torch.load(proposal, map_location="cpu", weights_only=True, mmap=True)
+        if (draft["case"], draft["step"], draft["rank"]) != key:
+            raise ValueError(f"{proposal}: proposal identity differs")
+        if "draft_tokens" not in draft or "projected_context" not in draft:
+            raise ValueError(f"{proposal}: incomplete proposal")
+        records[key] = row, draft
+    if not records:
+        raise ValueError(f"{directory}: empty captures")
+    for case in {k[0] for k in records}:
+        steps = {k[1] for k in records if k[0] == case}
+        if len(steps) < 2 or steps != set(range(max(steps) + 1)):
+            raise ValueError(f"{case}: missing prefill or verifier steps")
+        for step in steps:
+            if {k[2] for k in records if k[:2] == (case, step)} != {0, 1, 2, 3}:
+                raise ValueError(f"{case}/{step}: every step requires four TP ranks")
+    return records
+
+
+def _target_tensors(row: dict) -> dict[str, torch.Tensor]:
+    result = {
+        k: row[k]
+        for k in ("positions", "input_ids", "hidden", "num_sampled", "num_rejected")
+    }
+    for k in ("native_logits", "draft_logits"):
+        if row.get(k) is not None:
+            result[k] = row[k]
+    for k in ("sampling", "states"):
+        result.update({f"{k}/{label}": t for label, t in row[k].items()})
+    result.update({f"aux/{i}": t for i, t in enumerate(row["aux_hidden_states"])})
+    result.update(
+        {
+            f"layer/{v['layer_idx']}/{v['label']}": v["tensor"]
+            for v in row["tensors"].values()
+        }
+    )
+    # Unwritten output padding is not an emitted or accepted token.
+    result["accepted_output"] = row["sampled_token_ids"][:, : row["num_sampled"].item()]
+    return result
+
+
+def _proposal_tensors(row: dict) -> dict[str, torch.Tensor]:
+    values = {
+        k: v
+        for k, v in row.items()
+        if isinstance(v, torch.Tensor) and k != "idx_mapping"
+    }
+    if "idx_mapping" in row and row.get("sampling_layout") != "request_gathered_v1":
+        # Early captures retained complete arrays indexed by request slot.
+        indices = row["idx_mapping"].to(torch.int64)
+        for name in ("temperature", "seeds"):
+            if name in values:
+                values[name] = values[name].index_select(0, indices)
+    return values
+
+
+def compare_natural(left_dir: Path, right_dir: Path) -> dict:
+    left, right = _load(left_dir), _load(right_dir)
+    cases = {k[0] for k in left}
+    if cases != {k[0] for k in right}:
+        raise ValueError("Case coverage differs")
+    result = {"left": str(left_dir), "right": str(right_dir), "cases": []}
+    for case in sorted(cases):
+        lengths = [len({k[1] for k in arm if k[0] == case}) for arm in (left, right)]
+        first = None
+        mappings = []
+        for step in range(min(lengths)):
+            for phase_index, phase in enumerate(("target", "proposal")):
+                differences = []
+                for rank in range(4):
+                    rows = [arm[case, step, rank][phase_index] for arm in (left, right)]
+                    values = [
+                        _target_tensors(row)
+                        if phase == "target"
+                        else _proposal_tensors(row)
+                        for row in rows
+                    ]
+                    if phase == "proposal" and "idx_mapping" in rows[0]:
+                        slots = [row["idx_mapping"].tolist() for row in rows]
+                        if slots[0] != slots[1]:
+                            mappings.append(
+                                {"step": step, "rank": rank, "slots": slots}
+                            )
+                    if values[0].keys() != values[1].keys():
+                        raise ValueError(
+                            f"{case}/{step}/{rank}: tensor coverage differs"
+                        )
+                    for name in sorted(values[0]):
+                        a, b = (v[name] for v in values)
+                        if a.shape != b.shape or a.dtype != b.dtype:
+                            diff = {
+                                "contract_changed": [
+                                    str(a.shape),
+                                    str(b.shape),
+                                    str(a.dtype),
+                                    str(b.dtype),
+                                ]
+                            }
+                        else:
+                            diff = tensor_difference(a, b)
+                            if diff["bitwise_equal"]:
+                                continue
+                            if name == "native_logits":
+                                diff.update(sampling_difference(a, b))
+                        differences.append({"rank": rank, "name": name, **diff})
+                if differences:
+                    first = {"step": step, "phase": phase, "differences": differences}
+                    break
+            if first is not None:
+                break
+        result["cases"].append(
+            {
+                "case": case,
+                "steps_per_arm": lengths,
+                "first_observed_difference": first,
+                "different_request_slot_mappings": mappings,
+                "all_logical_tensors_equal": first is None and lengths[0] == lengths[1],
+            }
+        )
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("left", type=Path)
+    parser.add_argument("right", type=Path)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    torch.set_num_threads(4)
+    result = compare_natural(args.left, args.right)
+    args.output.write_text(json.dumps(result, indent=2) + "\n")
+    for case in result["cases"]:
+        first = case["first_observed_difference"]
+        print(
+            case["case"], "equal" if first is None else (first["step"], first["phase"])
+        )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/compare_sm70_dflash2_state_audit.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Compare complete four-rank StateAuditExtension captures, failing on gaps."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+
+
+def tensor_difference(left: torch.Tensor, right: torch.Tensor) -> dict:
+    if left.shape != right.shape or left.dtype != right.dtype:
+        raise ValueError(
+            f"Tensor contract differs: {left.shape}/{left.dtype}, "
+            f"{right.shape}/{right.dtype}"
+        )
+
+    def raw_bytes(tensor):
+        flat = tensor.contiguous().reshape(-1)
+        # A size-one tensor can be "contiguous" with stride 8. Reinterpret
+        # only its logical storage, not padding or neighboring metadata rows.
+        return flat.as_strided((flat.numel(),), (1,)).view(torch.uint8)
+
+    byte_equal = torch.equal(raw_bytes(left), raw_bytes(right))
+    if byte_equal:
+        return {"bitwise_equal": True}
+    difference = (left.double() - right.double()).abs()
+    return {
+        "bitwise_equal": False,
+        "different_elements": int((left != right).sum()),
+        "max_abs": float(difference.max()) if difference.numel() else 0.0,
+        "left_nonfinite": int((~torch.isfinite(left)).sum()),
+        "right_nonfinite": int((~torch.isfinite(right)).sum()),
+    }
+
+
+def sampling_difference(left: torch.Tensor, right: torch.Tensor) -> dict:
+    from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch
+
+    if left.shape != right.shape:
+        raise ValueError("Logit shapes differ")
+    count = left.shape[0]
+
+    def probabilities(logits):
+        return apply_top_k_top_p_pytorch(
+            logits.float().clone(),
+            torch.full((count,), 20, dtype=torch.int32),
+            torch.full((count,), 0.95),
+        ).softmax(-1)
+
+    p, q = probabilities(left), probabilities(right)
+    full_tv = (left.float().softmax(-1) - right.float().softmax(-1)).abs().sum(-1) / 2
+    return {
+        "full_softmax_tv": full_tv.tolist(),
+        "sampling_tv": ((p - q).abs().sum(-1) / 2).tolist(),
+        "support_changed": ((p > 0) != (q > 0)).any(-1).tolist(),
+        "top1_changed": (left.argmax(-1) != right.argmax(-1)).tolist(),
+    }
+
+
+def compare(
+    left_dir: Path, right_dir: Path, *, right_verifier_route: str | None = None
+) -> dict:
+    left_files = {p.name: p for p in left_dir.glob("*-rank*-step*.pt")}
+    right_files = {p.name: p for p in right_dir.glob("*-rank*-step*.pt")}
+    if not left_files or left_files.keys() != right_files.keys():
+        raise ValueError("Missing or mismatched capture files")
+    result = {
+        "left": str(left_dir),
+        "right": str(right_dir),
+        "comparisons": [],
+        "logits": [],
+    }
+    coverage: dict[tuple[str, int], set[int]] = {}
+    seen_states: dict[tuple[str, int, str], set[str]] = {}
+    for name in sorted(left_files):
+        left = torch.load(left_files[name], map_location="cpu", weights_only=True)
+        right = torch.load(right_files[name], map_location="cpu", weights_only=True)
+        for key in ("case", "rank", "step", "phase", "num_draft_tokens"):
+            if left[key] != right[key]:
+                raise ValueError(f"{name}: {key} differs")
+        identity = {key: left[key] for key in ("case", "rank", "step", "phase")}
+        if right_verifier_route is not None and right["phase"] == "verify":
+            expected = {
+                f"route/verify/layer{layer}/{right_verifier_route}"
+                for layer in right["expected_layers"]
+            }
+            if set(right.get("verifier_routes", ())) != expected:
+                raise ValueError(f"{name}: missing {right_verifier_route} route hit")
+        if left.get("expected_layers") != right.get("expected_layers"):
+            raise ValueError(f"{name}: requested layers differ")
+        if not left.get("expected_layers") or "capture_epoch" not in left:
+            raise ValueError(f"{name}: missing current-forward snapshot provenance")
+        for side in (left, right):
+            for layer in side["expected_layers"]:
+                prefix = f"{side['phase']}/layer{layer}"
+                for required in (
+                    "/conv/input:",
+                    "/conv/output:",
+                    "/recurrent/q:",
+                    "/recurrent/input_state",
+                    "/recurrent/output:",
+                ):
+                    if not any(k.startswith(prefix + required) for k in side["states"]):
+                        raise ValueError(f"{name}: missing {prefix + required}")
+        coverage.setdefault((left["case"], left["step"]), set()).add(left["rank"])
+        for key in ("positions", "input_ids"):
+            if not torch.equal(left[key], right[key]):
+                raise ValueError(f"{name}: forced {key} differs")
+        groups = {
+            "boundary": ({"hidden": left["hidden"]}, {"hidden": right["hidden"]}),
+            "state": (left["states"], right["states"]),
+            "layer": tuple(
+                {
+                    f"layer{v['layer_idx']}/{v['label']}": v["tensor"]
+                    for v in d["tensors"].values()
+                }
+                for d in (left, right)
+            ),
+            "sampling": (left["sampling"], right["sampling"]),
+        }
+        seen_states.setdefault(
+            (left["case"], left["rank"], left["phase"]), set()
+        ).update(left["states"])
+        for group, (lvalues, rvalues) in groups.items():
+            if lvalues.keys() != rvalues.keys():
+                raise ValueError(f"{name}: {group} snapshot coverage differs")
+            for label in sorted(lvalues):
+                difference = tensor_difference(lvalues[label], rvalues[label])
+                if not difference["bitwise_equal"]:
+                    result["comparisons"].append(
+                        {**identity, "group": group, "label": label, **difference}
+                    )
+        if left["rank"] == 0:
+            if not all(torch.isfinite(d["native_logits"]).all() for d in (left, right)):
+                raise ValueError(f"{name}: nonfinite native logits")
+            result["logits"].append(
+                {
+                    **identity,
+                    "positions": left["positions"].tolist(),
+                    **tensor_difference(left["native_logits"], right["native_logits"]),
+                    **sampling_difference(
+                        left["native_logits"], right["native_logits"]
+                    ),
+                }
+            )
+    if any(ranks != {0, 1, 2, 3} for ranks in coverage.values()):
+        raise ValueError("Every step requires four TP ranks")
+    for key, labels in seen_states.items():
+        for required in (
+            "/conv/input",
+            "/conv/output",
+            "/recurrent/input_state",
+            "/recurrent/output",
+        ):
+            if not any(required in label for label in labels):
+                raise ValueError(f"{key}: missing {required} evidence")
+    for case in {key[0] for key in coverage}:
+        steps = {key[1] for key in coverage if key[0] == case}
+        if steps != set(range(max(steps) + 1)) or len(steps) < 2:
+            raise ValueError(f"{case}: missing prefill or verifier steps")
+        for rank in range(4):
+            if any(
+                (case, rank, phase) not in seen_states
+                for phase in ("prefill", "verify")
+            ):
+                raise ValueError(f"{case}/rank{rank}: missing prefill/verifier states")
+    result["summary"] = {
+        "files_per_arm": len(left_files),
+        "differing_intermediates": len(result["comparisons"]),
+        "max_sampling_tv": max(max(row["sampling_tv"]) for row in result["logits"]),
+        "support_changed_rows": sum(
+            sum(row["support_changed"]) for row in result["logits"]
+        ),
+        "top1_changed_rows": sum(sum(row["top1_changed"]) for row in result["logits"]),
+        "all_logits_bitwise_equal": all(
+            row["bitwise_equal"] for row in result["logits"]
+        ),
+    }
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("left", type=Path)
+    parser.add_argument("right", type=Path)
+    parser.add_argument("--output", required=True, type=Path)
+    parser.add_argument("--right-verifier-route", choices=("split", "packed"))
+    args = parser.parse_args()
+    torch.set_num_threads(4)
+    result = compare(
+        args.left, args.right, right_verifier_route=args.right_verifier_route
+    )
+    args.output.write_text(json.dumps(result, indent=2) + "\n")
+    print(json.dumps(result["summary"], indent=2))
+
+
+if __name__ == "__main__":
+    main()