diff --git a/examples/vllm_serve/calibrate_mask_reuse.py b/examples/vllm_serve/calibrate_mask_reuse.py new file mode 100644 index 00000000000..c4535082140 --- /dev/null +++ b/examples/vllm_serve/calibrate_mask_reuse.py @@ -0,0 +1,486 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Build a fail-closed schema-v3 mask-reuse candidate from compact captures. + +The input is compact capture JSONL emitted by ``collect_mask_reuse.py``. Its +streaming selector never expands the full consumer-head by donor-head matrix +into repeated row objects. This command cannot promote a serving policy until +the grouped inner/outer protocol is implemented and its preregistered gates pass. + +Example:: + + python examples/vllm_serve/calibrate_mask_reuse.py \ + --checkpoint /path/to/checkpoint \ + --compact-captures compact-captures.jsonl \ + --capture-manifest compact-captures.jsonl.manifest.json \ + --vanilla-config sparse_attention_config.json \ + --topology topology.json \ + --calibration-plan calibration-plan.json \ + --family-registry family-registry.json \ + --grouped-fit grouped-fit.json \ + --outer-report outer-report.json \ + --max-anchor-dropped-mass 0.02 \ + --reuse-dropped-mass-report-threshold 0.02 \ + --target-bmm1-skip-ratio 0.10 + +The final output is candidate-only and must be rejected by serving. +""" + +from __future__ import annotations + +import argparse +import json +import os +import tempfile +from collections.abc import Mapping +from hashlib import sha256 +from pathlib import Path + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + StableFileSnapshot, + read_stable_file_snapshot, + stable_file_sha256, + verify_checkpoint_manifest, +) +from modelopt.torch.sparsity.attention_sparsity.calibration.mask_reuse_compact import ( + calibrate_compact_mask_reuse_policy, + load_compact_mask_reuse_captures, +) + +_EVIDENCE_ARTIFACTS = { + "calibration_plan_sha256": "calibration_plan", + "family_registry_sha256": "family_registry", + "grouped_fit_sha256": "grouped_fit", + "outer_report_sha256": "outer_report", +} + +_CAPTURE_MANIFEST_FIELDS = frozenset( + { + "capture_manifest_schema_version", + "capture_protocol", + "model", + "checkpoint_manifest_sha256", + "checkpoint_manifest_path", + "checkpoint_file_count", + "checkpoint_total_size_bytes", + "plan", + "fa4_source", + "fa4_source_commit", + "fa4_source_git_tree", + "fa4_source_git_archive_sha256", + "fa4_source_manifest_path", + "fa4_source_manifest_sha256", + "fa4_source_directory_count", + "fa4_source_file_count", + "fa4_source_total_size_bytes", + "engine_kwargs", + "dense_shadow_validation_requested", + "target_sparsity_hex", + "vanilla_threshold_scale_factor", + "vanilla_fit_sha256", + "vanilla_config_file_sha256", + "prompt_plan_file_sha256", + "compact_capture_file_sha256", + "capture_count", + "candidate_cell_count", + "captures", + } +) + + +def _reject_duplicate_json_keys(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + +def _parse_json_object(payload: bytes, *, path: Path, label: str) -> dict[str, object]: + try: + raw = json.loads( + payload, + object_pairs_hook=_reject_duplicate_json_keys, + ) + except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as error: + raise ValueError(f"could not load {label} from {path}: {error}") from error + if not isinstance(raw, dict): + raise ValueError(f"{label} must contain a JSON object") + return raw + + +def _load_json_snapshot(path: Path, *, label: str) -> tuple[dict[str, object], StableFileSnapshot]: + snapshot = read_stable_file_snapshot(path, label=label) + return _parse_json_object(snapshot.payload, path=path, label=label), snapshot + + +def _load_json_object(path: Path, *, label: str) -> dict[str, object]: + """Load strict JSON from one stable no-follow byte snapshot.""" + return _load_json_snapshot(path, label=label)[0] + + +def _stable_file_sha256(path: Path, *, label: str) -> str: + return stable_file_sha256(path, label=label) + + +def _evidence_artifacts( + args: argparse.Namespace, *, vanilla_fit_sha256: str +) -> tuple[dict[str, str], dict[str, Path]]: + paths = { + field: Path(getattr(args, attribute)) for field, attribute in _EVIDENCE_ARTIFACTS.items() + } + paths["vanilla_fit_sha256"] = args.vanilla_config + paths["reuse_bundle_sha256"] = args.compact_captures + evidence = { + field: ( + vanilla_fit_sha256 + if field == "vanilla_fit_sha256" + else _stable_file_sha256(path, label=field) + ) + for field, path in paths.items() + } + return evidence, paths + + +def _canonical_json_bytes(value: object) -> bytes: + return ( + json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode("utf-8") + + +def _fsync_directory(path: Path) -> None: + directory_flag = getattr(os, "O_DIRECTORY", None) + if directory_flag is None: + # Windows lacks portable directory fsync; no-clobber linking remains + # atomic, while crash durability of the directory entry is best effort. + return + descriptor = os.open(path, os.O_RDONLY | directory_flag) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def _temporary_payload(path: Path, payload: bytes) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + mode="wb", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as handle: + temporary = Path(handle.name) + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + return temporary + + +def _unlink_if_identity(path: Path, identity: tuple[int, int]) -> None: + if identity[1] == 0: + return + try: + observed = path.stat(follow_symlinks=False) + except FileNotFoundError: + return + if observed.st_ino != 0 and (observed.st_dev, observed.st_ino) == identity: + path.unlink() + _fsync_directory(path.parent) + + +def _publish_no_clobber(temporary: Path, destination: Path) -> tuple[int, int]: + observed = temporary.stat(follow_symlinks=False) + if observed.st_ino == 0: + raise RuntimeError("candidate temporary file has no stable identity") + identity = observed.st_dev, observed.st_ino + os.link(temporary, destination, follow_symlinks=False) + try: + published = destination.stat(follow_symlinks=False) + if published.st_ino == 0 or (published.st_dev, published.st_ino) != identity: + raise RuntimeError("candidate destination changed during publication") + temporary.unlink() + _fsync_directory(destination.parent) + except BaseException: + _unlink_if_identity(destination, identity) + raise + return identity + + +def _publish_candidate_outputs( + policy_path: Path, policy_payload: bytes, report_path: Path, report_payload: bytes +) -> None: + if policy_path.exists() or report_path.exists(): + raise FileExistsError("candidate outputs already exist; refusing to overwrite them") + policy_temporary: Path | None = None + report_temporary: Path | None = None + report_identity: tuple[int, int] | None = None + try: + policy_temporary = _temporary_payload(policy_path, policy_payload) + report_temporary = _temporary_payload(report_path, report_payload) + report_identity = _publish_no_clobber(report_temporary, report_path) + report_temporary = None + _publish_no_clobber(policy_temporary, policy_path) + policy_temporary = None + except BaseException: + if report_identity is not None: + _unlink_if_identity(report_path, report_identity) + if policy_temporary is not None: + policy_temporary.unlink(missing_ok=True) + if report_temporary is not None: + report_temporary.unlink(missing_ok=True) + raise + + +def _validate_capture_manifest( + raw: Mapping[str, object], + *, + checkpoint_sha256: str, + model: str, + compact_capture_sha256: str, + vanilla_config_sha256: str, +) -> None: + missing = _CAPTURE_MANIFEST_FIELDS - raw.keys() + extra = raw.keys() - _CAPTURE_MANIFEST_FIELDS + if missing or extra: + raise ValueError( + "capture manifest fields do not match schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + expected = { + "capture_manifest_schema_version": 4, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v4", + "model": model, + "checkpoint_manifest_sha256": checkpoint_sha256, + "compact_capture_file_sha256": compact_capture_sha256, + "vanilla_config_file_sha256": vanilla_config_sha256, + } + for field, value in expected.items(): + if raw[field] != value: + raise ValueError(f"capture manifest {field} does not match its verified input") + for field, length in { + "fa4_source_commit": 40, + "fa4_source_git_tree": 40, + "fa4_source_git_archive_sha256": 64, + "fa4_source_manifest_sha256": 64, + }.items(): + value = raw[field] + if ( + not isinstance(value, str) + or len(value) != length + or any(character not in "0123456789abcdef" for character in value) + ): + raise ValueError(f"capture manifest {field} is not canonical hexadecimal evidence") + for field in ("fa4_source", "fa4_source_manifest_path"): + if not isinstance(raw[field], str) or not raw[field]: + raise ValueError(f"capture manifest {field} must be a non-empty path") + for field in ( + "fa4_source_directory_count", + "fa4_source_file_count", + "fa4_source_total_size_bytes", + ): + value = raw[field] + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"capture manifest {field} must be an integer >= 0") + if raw["fa4_source_file_count"] == 0: + raise ValueError("capture manifest must bind at least one FA4 source file") + if not isinstance(raw["engine_kwargs"], Mapping): + raise ValueError("capture manifest engine_kwargs must be an object") + if not isinstance(raw["dense_shadow_validation_requested"], bool): + raise ValueError("capture manifest dense_shadow_validation_requested must be boolean") + if isinstance(raw["capture_count"], bool) or not isinstance(raw["capture_count"], int): + raise ValueError("capture manifest capture_count must be an integer") + if raw["capture_count"] <= 0: + raise ValueError("capture manifest must contain at least one capture") + captures = raw["captures"] + if not isinstance(captures, list) or len(captures) != raw["capture_count"]: + raise ValueError("capture manifest captures do not match capture_count") + candidate_cell_count = raw["candidate_cell_count"] + if ( + isinstance(candidate_cell_count, bool) + or not isinstance(candidate_cell_count, int) + or candidate_cell_count <= 0 + ): + raise ValueError("capture manifest candidate_cell_count must be positive") + observed_cells = 0 + for index, capture in enumerate(captures): + if not isinstance(capture, Mapping) or not isinstance( + capture.get("candidate_cell_count"), int + ): + raise ValueError( + f"capture manifest captures[{index}].candidate_cell_count must be an integer" + ) + observed_cells += int(capture["candidate_cell_count"]) + if observed_cells != candidate_cell_count: + raise ValueError("capture manifest candidate-cell total is inconsistent") + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Build a fail-closed schema-v3 candidate from schema-v4 mask-reuse capture evidence" + ), + allow_abbrev=False, + ) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument( + "--compact-captures", + type=Path, + required=True, + help="Compact capture JSONL emitted by collect_mask_reuse.py (recommended)", + ) + parser.add_argument("--capture-manifest", type=Path, required=True) + parser.add_argument( + "--vanilla-config", + type=Path, + required=True, + help="ModelOpt sparse_attention_config JSON or checkpoint config.json", + ) + parser.add_argument( + "--topology", + type=Path, + required=True, + help="JSON object containing anchors and nearest layer mappings", + ) + parser.add_argument("--calibration-plan", type=Path, required=True) + parser.add_argument("--family-registry", type=Path, required=True) + parser.add_argument("--grouped-fit", type=Path, required=True) + parser.add_argument("--outer-report", type=Path, required=True) + parser.add_argument( + "--max-anchor-dropped-mass", + type=float, + required=True, + help="Maximum allowed anchor dropped mass", + ) + parser.add_argument( + "--reuse-dropped-mass-report-threshold", + type=float, + required=True, + help="Diagnostic reuse threshold; does not affect selection", + ) + parser.add_argument( + "--target-bmm1-skip-ratio", + type=float, + required=True, + help="Minimum model-wide BMM1 tile skip ratio required in every context bucket", + ) + parser.add_argument( + "--output-policy", + type=Path, + default=Path("mask_reuse_candidate.json"), + help="Fail-closed schema-v3 candidate output path", + ) + parser.add_argument( + "--output-report", + type=Path, + default=Path("mask_reuse_calibration_report.json"), + help="Standalone calibration-report output path", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + if args.output_policy.resolve() == args.output_report.resolve(): + parser.error("--output-policy and --output-report must be different paths") + + try: + checkpoint = verify_checkpoint_manifest(args.checkpoint) + vanilla_config, vanilla_snapshot = _load_json_snapshot( + args.vanilla_config, label="vanilla config" + ) + topology, topology_snapshot = _load_json_snapshot(args.topology, label="topology") + capture_manifest, capture_manifest_snapshot = _load_json_snapshot( + args.capture_manifest, label="capture manifest" + ) + evidence, evidence_paths = _evidence_artifacts( + args, vanilla_fit_sha256=vanilla_snapshot.sha256 + ) + capture_manifest_sha256 = capture_manifest_snapshot.sha256 + topology_sha256 = topology_snapshot.sha256 + _validate_capture_manifest( + capture_manifest, + checkpoint_sha256=checkpoint.sha256, + model=checkpoint.model, + compact_capture_sha256=evidence["reuse_bundle_sha256"], + vanilla_config_sha256=evidence["vanilla_fit_sha256"], + ) + artifact = calibrate_compact_mask_reuse_policy( + load_compact_mask_reuse_captures(args.compact_captures), + vanilla_calibration=vanilla_config, + topology=topology, + checkpoint_manifest=checkpoint, + evidence=evidence, + max_anchor_dropped_mass=args.max_anchor_dropped_mass, + reuse_dropped_mass_report_threshold=(args.reuse_dropped_mass_report_threshold), + target_bmm1_skip_ratio=args.target_bmm1_skip_ratio, + source_provenance={ + "capture_manifest_sha256": capture_manifest_sha256, + "topology_file_sha256": topology_sha256, + }, + ) + provenance = artifact.get("provenance") + if not isinstance(provenance, Mapping): + raise ValueError("calibrator returned no provenance object") + if provenance.get("input_capture_count") != capture_manifest["capture_count"]: + raise ValueError("calibrator capture count does not match capture manifest") + if provenance.get("candidate_cell_count") != capture_manifest["candidate_cell_count"]: + raise ValueError("calibrator candidate-cell count does not match capture manifest") + for field, path in evidence_paths.items(): + if _stable_file_sha256(path, label=field) != evidence[field]: + raise ValueError(f"{field} artifact changed during calibration") + if ( + _stable_file_sha256(args.capture_manifest, label="capture manifest") + != capture_manifest_sha256 + ): + raise ValueError("capture manifest changed during calibration") + if _stable_file_sha256(args.topology, label="topology") != topology_sha256: + raise ValueError("topology changed during calibration") + if verify_checkpoint_manifest(args.checkpoint) != checkpoint: + raise ValueError("checkpoint changed during calibration") + except (OSError, ValueError) as error: + parser.error(str(error)) + + if ( + artifact.get("promotion_status") != "candidate_only" + or artifact.get("deployment_geometry_validated") is not False + ): + parser.error("calibrator did not return a fail-closed candidate-only artifact") + + report = artifact.get("calibration_report") + if not isinstance(report, Mapping): + parser.error("calibrator returned no calibration_report object") + + try: + policy_payload = _canonical_json_bytes(artifact) + report_payload = _canonical_json_bytes(report) + _publish_candidate_outputs( + args.output_policy, policy_payload, args.output_report, report_payload + ) + except (OSError, FileExistsError) as error: + parser.error(f"could not write calibration outputs: {error}") + + policy_digest = sha256(policy_payload).hexdigest() + print(f"[ModelOpt] Wrote fail-closed mask-reuse candidate to {args.output_policy.resolve()}") + print(f"[ModelOpt] Wrote calibration report to {args.output_report.resolve()}") + print(f"MASK_REUSE_FA4_CANDIDATE_SHA256={policy_digest}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/vllm_serve/create_checkpoint_manifest.py b/examples/vllm_serve/create_checkpoint_manifest.py new file mode 100644 index 00000000000..8fe276f0d0f --- /dev/null +++ b/examples/vllm_serve/create_checkpoint_manifest.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Create a deterministic, no-clobber checkpoint manifest for calibration.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + create_checkpoint_manifest, +) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(allow_abbrev=False) + parser.add_argument("checkpoint", type=Path) + parser.add_argument("--model-id", required=True) + args = parser.parse_args(argv) + try: + manifest = create_checkpoint_manifest(args.checkpoint, model=args.model_id) + except (OSError, ValueError) as error: + parser.error(str(error)) + print(f"[ModelOpt] Wrote {manifest.manifest_path}") + print(f"CHECKPOINT_MANIFEST_SHA256={manifest.sha256}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py b/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py index 87088f805bd..c8b55c03a3a 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py @@ -17,10 +17,54 @@ from .calibrate import calibrate_sparse_attention from .calibrator import DynamicThresholdCalibrator +from .checkpoint_manifest import ( + CHECKPOINT_MANIFEST_NAME, + CheckpointManifestError, + StableFileSnapshot, + VerifiedCheckpointManifest, + create_checkpoint_manifest, + read_stable_file_snapshot, + stable_file_sha256, + verify_checkpoint_manifest, +) +from .mask_reuse import ( + AnchorLayerStats, + MaskReuseCalibrationError, + MaskReuseObservation, + calibrate_mask_reuse_policy, + canonical_prefill_threshold_scale_factor, + load_mask_reuse_observations, + parse_mask_reuse_observations, +) +from .mask_reuse_compact import ( + CompactMaskReuseCapture, + CompactMaskReuseCaptureSource, + calibrate_compact_mask_reuse_policy, + load_compact_mask_reuse_captures, +) from .ruler_dataset import RulerDatasetBuilder __all__ = [ + "CHECKPOINT_MANIFEST_NAME", + "AnchorLayerStats", + "CheckpointManifestError", + "CompactMaskReuseCapture", + "CompactMaskReuseCaptureSource", "DynamicThresholdCalibrator", + "MaskReuseCalibrationError", + "MaskReuseObservation", "RulerDatasetBuilder", + "StableFileSnapshot", + "VerifiedCheckpointManifest", + "calibrate_compact_mask_reuse_policy", + "calibrate_mask_reuse_policy", "calibrate_sparse_attention", + "canonical_prefill_threshold_scale_factor", + "create_checkpoint_manifest", + "load_compact_mask_reuse_captures", + "load_mask_reuse_observations", + "parse_mask_reuse_observations", + "read_stable_file_snapshot", + "stable_file_sha256", + "verify_checkpoint_manifest", ] diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/checkpoint_manifest.py b/modelopt/torch/sparsity/attention_sparsity/calibration/checkpoint_manifest.py new file mode 100644 index 00000000000..dd79737fb25 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/checkpoint_manifest.py @@ -0,0 +1,439 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Strict, content-addressed checkpoint identity for mask-reuse calibration.""" + +from __future__ import annotations + +import json +import os +import stat +import tempfile +import unicodedata +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path, PurePosixPath + +__all__ = [ + "CHECKPOINT_MANIFEST_NAME", + "CheckpointManifestError", + "StableFileSnapshot", + "VerifiedCheckpointManifest", + "create_checkpoint_manifest", + "read_stable_file_snapshot", + "stable_file_sha256", + "verify_checkpoint_manifest", +] + +CHECKPOINT_MANIFEST_NAME = "checkpoint_manifest.json" +_MANIFEST_FIELDS = frozenset({"checkpoint_manifest_schema_version", "model", "files"}) +_FILE_FIELDS = frozenset({"path", "size_bytes", "sha256"}) +_WEIGHT_SUFFIXES = frozenset({".bin", ".pt", ".safetensors"}) +_FILE_ATTRIBUTE_REPARSE_POINT = getattr(stat, "FILE_ATTRIBUTE_REPARSE_POINT", 0x400) + + +class CheckpointManifestError(ValueError): + """Raised when a checkpoint cannot be bound to its exact file contents.""" + + +def _strict_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise CheckpointManifestError(f"checkpoint manifest repeats JSON key {key!r}") + result[key] = value + return result + + +def _exact_fields(raw: dict[str, object], expected: frozenset[str], label: str) -> None: + missing = expected - raw.keys() + extra = raw.keys() - expected + if missing or extra: + raise CheckpointManifestError( + f"{label} fields do not match the schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + + +def _text(value: object, label: str) -> str: + if ( + not isinstance(value, str) + or not value + or value != value.strip() + or unicodedata.normalize("NFC", value) != value + or any(ord(character) < 32 for character in value) + ): + raise CheckpointManifestError(f"{label} must be non-empty canonical NFC text") + return value + + +def _sha256(value: object, label: str) -> str: + if ( + not isinstance(value, str) + or len(value) != 64 + or any(character not in "0123456789abcdef" for character in value) + ): + raise CheckpointManifestError(f"{label} must be a lowercase SHA256") + return value + + +def _canonical_json_bytes(value: object) -> bytes: + return ( + json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + + +def _is_link_like(value: os.stat_result) -> bool: + """Return whether a no-follow stat identifies a symlink or Windows reparse point.""" + return stat.S_ISLNK(value.st_mode) or bool( + getattr(value, "st_file_attributes", 0) & _FILE_ATTRIBUTE_REPARSE_POINT + ) + + +def _same_file(left: os.stat_result, right: os.stat_result) -> bool: + """Compare file identities, failing closed when an inode is unavailable.""" + return left.st_ino != 0 and right.st_ino != 0 and os.path.samestat(left, right) + + +def _open_stable_regular(path: Path, label: str) -> tuple[int, os.stat_result]: + try: + named_before = path.stat(follow_symlinks=False) + except OSError as error: + raise CheckpointManifestError(f"could not inspect {label}") from error + if _is_link_like(named_before): + raise CheckpointManifestError(f"could not open {label} without following symlinks") + if not stat.S_ISREG(named_before.st_mode): + raise CheckpointManifestError(f"{label} must be one stable regular file, not a symlink") + try: + # Windows has no O_NOFOLLOW. The no-follow pre/post stats and handle + # identity checks keep that fallback fail-closed before any bytes are read. + descriptor = os.open( + path, + os.O_RDONLY | getattr(os, "O_BINARY", 0) | getattr(os, "O_NOFOLLOW", 0), + ) + except OSError as error: + raise CheckpointManifestError( + f"could not open {label} without following symlinks" + ) from error + try: + opened = os.fstat(descriptor) + named = path.stat(follow_symlinks=False) + except OSError: + os.close(descriptor) + raise + if _is_link_like(named): + os.close(descriptor) + raise CheckpointManifestError(f"could not open {label} without following symlinks") + if ( + not stat.S_ISREG(opened.st_mode) + or not stat.S_ISREG(named.st_mode) + or not _same_file(named_before, opened) + or not _same_file(opened, named) + ): + os.close(descriptor) + raise CheckpointManifestError(f"{label} must be one stable regular file, not a symlink") + return descriptor, opened + + +def _hash_stable_regular( + path: Path, label: str, *, capture_payload: bool = False +) -> tuple[int, str, bytes | None]: + descriptor, before = _open_stable_regular(path, label) + digest = sha256() + payload = bytearray() if capture_payload else None + observed_size = 0 + try: + for chunk in iter(lambda: os.read(descriptor, 1024 * 1024), b""): + observed_size += len(chunk) + digest.update(chunk) + if payload is not None: + payload.extend(chunk) + after = os.fstat(descriptor) + named_after = path.stat(follow_symlinks=False) + except OSError as error: + raise CheckpointManifestError(f"could not hash stable {label}") from error + finally: + os.close(descriptor) + if ( + _is_link_like(named_after) + or not stat.S_ISREG(named_after.st_mode) + or not _same_file(before, after) + or not _same_file(after, named_after) + or (before.st_size, before.st_mtime_ns) != (after.st_size, after.st_mtime_ns) + or (after.st_size, after.st_mtime_ns) != (named_after.st_size, named_after.st_mtime_ns) + ): + raise CheckpointManifestError(f"{label} changed while it was being hashed") + return observed_size, digest.hexdigest(), None if payload is None else bytes(payload) + + +@dataclass(frozen=True, slots=True) +class StableFileSnapshot: + """Exact bytes and SHA256 read from one stable no-follow descriptor.""" + + path: Path + payload: bytes + sha256: str + + +def read_stable_file_snapshot(path: str | Path, *, label: str) -> StableFileSnapshot: + """Read and hash identical bytes from one stable regular file.""" + source = Path(path) + _, digest, payload = _hash_stable_regular(source, label, capture_payload=True) + assert payload is not None + return StableFileSnapshot(source, payload, digest) + + +def stable_file_sha256(path: str | Path, *, label: str) -> str: + """Hash one stable regular file without retaining its contents.""" + _, digest, _ = _hash_stable_regular(Path(path), label) + return digest + + +def _checkpoint_files(root: Path, manifest_path: Path) -> set[str]: + files: set[str] = set() + pending = [root] + while pending: + directory = pending.pop() + try: + with os.scandir(directory) as iterator: + entries = sorted(iterator, key=lambda entry: entry.name) + except OSError as error: + raise CheckpointManifestError( + f"could not traverse checkpoint directory {directory}" + ) from error + for entry in entries: + path = Path(entry.path) + relative = path.relative_to(root).as_posix() + try: + observed = entry.stat(follow_symlinks=False) + except OSError as error: + raise CheckpointManifestError( + f"could not inspect checkpoint path {relative!r}" + ) from error + if _is_link_like(observed): + raise CheckpointManifestError( + f"checkpoint contains forbidden symlink or reparse point {relative!r}" + ) + if stat.S_ISDIR(observed.st_mode): + pending.append(path) + elif stat.S_ISREG(observed.st_mode): + if path != manifest_path: + files.add(relative) + else: + raise CheckpointManifestError(f"checkpoint contains non-regular path {relative!r}") + return files + + +def _fsync_directory(path: Path) -> None: + directory_flag = getattr(os, "O_DIRECTORY", None) + if directory_flag is None: + # Python on Windows cannot portably open and fsync a directory. The + # complete temporary file is still fsynced before its no-clobber link. + return + descriptor = os.open(path, os.O_RDONLY | directory_flag) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def create_checkpoint_manifest(checkpoint: str | Path, *, model: str) -> VerifiedCheckpointManifest: + """Create the deterministic checkpoint manifest without replacing any file.""" + root = Path(checkpoint).expanduser().resolve() + if not root.is_dir(): + raise CheckpointManifestError("checkpoint must be a local directory") + manifest_path = root / CHECKPOINT_MANIFEST_NAME + if os.path.lexists(manifest_path): + raise CheckpointManifestError( + f"{CHECKPOINT_MANIFEST_NAME} already exists; refusing to overwrite it" + ) + files = _checkpoint_files(root, manifest_path) + entries = [] + for relative in sorted(files): + size, digest, _ = _hash_stable_regular(root / relative, f"checkpoint file {relative!r}") + entries.append({"path": relative, "size_bytes": size, "sha256": digest}) + if _checkpoint_files(root, manifest_path) != files: + raise CheckpointManifestError("checkpoint file set changed while building manifest") + payload = _canonical_json_bytes( + { + "checkpoint_manifest_schema_version": 1, + "model": _text(model, "model"), + "files": entries, + } + ) + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=root, + prefix=f".{CHECKPOINT_MANIFEST_NAME}.", + suffix=".tmp", + delete=False, + ) as handle: + temporary = Path(handle.name) + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + observed = temporary.stat(follow_symlinks=False) + if observed.st_ino == 0: + raise CheckpointManifestError( + "checkpoint manifest temporary file has no stable identity" + ) + identity = observed.st_dev, observed.st_ino + os.link(temporary, manifest_path, follow_symlinks=False) + try: + published = manifest_path.stat(follow_symlinks=False) + if ( + _is_link_like(published) + or not stat.S_ISREG(published.st_mode) + or not _same_file(observed, published) + ): + raise CheckpointManifestError( + "checkpoint manifest destination changed during publication" + ) + temporary.unlink() + temporary = None + _fsync_directory(root) + except BaseException: + try: + published = manifest_path.stat(follow_symlinks=False) + if published.st_ino != 0 and (published.st_dev, published.st_ino) == identity: + manifest_path.unlink() + _fsync_directory(root) + finally: + raise + except FileExistsError as error: + raise CheckpointManifestError( + f"{CHECKPOINT_MANIFEST_NAME} appeared during publication; refusing to overwrite it" + ) from error + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + return verify_checkpoint_manifest(root, expected_model=model) + + +def _relative_path(value: object, label: str) -> str: + text = _text(value, label) + path = PurePosixPath(text) + if ( + path.is_absolute() + or text != path.as_posix() + or "\\" in text + or any(part in {"", ".", ".."} for part in path.parts) + or text == CHECKPOINT_MANIFEST_NAME + ): + raise CheckpointManifestError(f"{label} must be a canonical relative POSIX path") + return text + + +@dataclass(frozen=True, slots=True) +class VerifiedCheckpointManifest: + """Identity of a checkpoint whose complete file set was SHA256-verified.""" + + checkpoint_root: Path + manifest_path: Path + model: str + sha256: str + file_count: int + total_size_bytes: int + + +def verify_checkpoint_manifest( + checkpoint: str | Path, *, expected_model: str | None = None +) -> VerifiedCheckpointManifest: + """Verify the fixed manifest under ``checkpoint`` and every declared file. + + The manifest must enumerate every regular file below the loaded checkpoint + directory except itself. This prevents a manifest that binds only a subset + of weights or remote-code/tokenizer inputs from naming the checkpoint. + """ + root = Path(checkpoint).expanduser().resolve() + if not root.is_dir(): + raise CheckpointManifestError("checkpoint must be a local directory") + manifest_path = root / CHECKPOINT_MANIFEST_NAME + _, manifest_digest, payload = _hash_stable_regular( + manifest_path, "checkpoint manifest", capture_payload=True + ) + assert payload is not None + try: + raw = json.loads(payload, object_pairs_hook=_strict_object) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise CheckpointManifestError("checkpoint manifest is not strict UTF-8 JSON") from error + if not isinstance(raw, dict): + raise CheckpointManifestError("checkpoint manifest must be a JSON object") + _exact_fields(raw, _MANIFEST_FIELDS, "checkpoint manifest") + if raw["checkpoint_manifest_schema_version"] != 1: + raise CheckpointManifestError("checkpoint_manifest_schema_version must be 1") + if payload != _canonical_json_bytes(raw): + raise CheckpointManifestError("checkpoint manifest bytes are not canonical JSON") + model = _text(raw["model"], "checkpoint manifest.model") + if expected_model is not None and model != expected_model: + raise CheckpointManifestError( + f"checkpoint manifest model {model!r} does not match requested model {expected_model!r}" + ) + raw_files = raw["files"] + if not isinstance(raw_files, list) or not raw_files: + raise CheckpointManifestError("checkpoint manifest.files must be a non-empty list") + + declared: dict[str, tuple[int, str]] = {} + for index, item in enumerate(raw_files): + label = f"checkpoint manifest.files[{index}]" + if not isinstance(item, dict): + raise CheckpointManifestError(f"{label} must be an object") + _exact_fields(item, _FILE_FIELDS, label) + relative = _relative_path(item["path"], f"{label}.path") + size = item["size_bytes"] + if isinstance(size, bool) or not isinstance(size, int) or size < 0: + raise CheckpointManifestError(f"{label}.size_bytes must be an integer >= 0") + digest = _sha256(item["sha256"], f"{label}.sha256") + if relative in declared: + raise CheckpointManifestError(f"checkpoint manifest repeats file {relative!r}") + declared[relative] = (size, digest) + if list(declared) != sorted(declared): + raise CheckpointManifestError("checkpoint manifest files must be sorted by path") + if "config.json" not in declared or not any( + Path(relative).suffix in _WEIGHT_SUFFIXES for relative in declared + ): + raise CheckpointManifestError( + "checkpoint manifest must bind config.json and at least one model weight file" + ) + + actual = _checkpoint_files(root, manifest_path) + if actual != set(declared): + raise CheckpointManifestError( + "checkpoint manifest does not exactly cover checkpoint files; " + f"missing={sorted(actual - set(declared))}, extra={sorted(set(declared) - actual)}" + ) + total_size = 0 + for relative, (expected_size, expected_digest) in declared.items(): + path = root / relative + observed_size, observed_digest, _ = _hash_stable_regular( + path, f"checkpoint file {relative!r}" + ) + if observed_size != expected_size or observed_digest != expected_digest: + raise CheckpointManifestError( + f"checkpoint file {relative!r} does not match its size/SHA256 manifest entry" + ) + total_size += observed_size + if _checkpoint_files(root, manifest_path) != actual: + raise CheckpointManifestError("checkpoint file set changed during verification") + return VerifiedCheckpointManifest( + checkpoint_root=root, + manifest_path=manifest_path, + model=model, + sha256=manifest_digest, + file_count=len(declared), + total_size_bytes=total_size, + ) diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py new file mode 100644 index 00000000000..e7b9f6bbce3 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py @@ -0,0 +1,1702 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline calibration and schema-v3 export for cross-layer mask reuse. + +The selector consumes prompt-level observations measured at target sparsities +derived from an existing ModelOpt skip-softmax fit. Calibration observations +alone select one target sparsity per context bucket and one donor head (or an +exact fallback) per consumer head. Held-out observations only evaluate the +frozen policy. + +This module intentionally has no serving-backend dependency. It exports the +JSON-safe schema consumed by the mask-reuse attention backend. +""" + +from __future__ import annotations + +import json +import math +from collections import defaultdict +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import cast + +import pulp + +import modelopt + +from .checkpoint_manifest import VerifiedCheckpointManifest + +__all__ = [ + "AnchorLayerStats", + "MaskReuseCalibrationError", + "MaskReuseObservation", + "calibrate_mask_reuse_policy", + "canonical_prefill_threshold_scale_factor", + "load_mask_reuse_observations", + "parse_mask_reuse_observations", +] + + +_FORMULA = "a * exp(b * target_sparsity)" +_SPLITS = frozenset({"calibration", "heldout"}) +_OBSERVATION_FIELDS = frozenset( + { + "model", + "min_kv_tokens", + "max_kv_tokens", + "target_sparsity", + "sample_length", + "threshold_lambda", + "threshold_log2", + "q_tokens", + "kv_tokens", + "q_start_tokens", + "split", + "prompt_id", + "source_capture_sha256", + "anchor_layer", + "consumer_layer", + "consumer_head", + "donor_head", + "retained_tiles", + "eligible_tiles", + "anchor_dropped_mass", + "anchor_stats_by_layer", + "dropped_mass", + } +) +_CALIBRATION_PROTOCOL = "modelopt_mask_reuse_target_sparsity_v1" +_SOLVER_LEXICOGRAPHIC_ATOL = 1e-7 +_EVIDENCE_FIELDS = frozenset( + { + "calibration_plan_sha256", + "family_registry_sha256", + "vanilla_fit_sha256", + "reuse_bundle_sha256", + "grouped_fit_sha256", + "outer_report_sha256", + } +) +_DEPLOYMENT_GEOMETRY_CONTRACT: dict[str, object] = { + "schema_version": 1, + "batch_size": 1, + "max_query_chunk_tokens": 8192, + "query_block_tokens": 128, + "key_block_tokens": 128, + "qstage2_query_pair_tokens": 256, + "kv_page_tokens": 16, + "head_dim": 128, + "causal": True, + "bottom_right_aligned": True, + "query_chunk_start_alignment_tokens": 128, + "attention_dtype": "bfloat16", + "kv_cache_dtype": "bfloat16", + "common_prefix": False, + "cascade_attention": False, + "context_parallel_size": 1, + "pipeline_parallel_size": 1, +} + +Bucket = tuple[int, int | None] +ConsumerHead = tuple[int, int] +ObservationKey = tuple[str, str, float, int, int, int] +AnchorKey = tuple[str, str, float, int, int] + + +class MaskReuseCalibrationError(ValueError): + """Raised when observations cannot produce a trustworthy reuse policy.""" + + +@dataclass(frozen=True, slots=True) +class AnchorLayerStats: + """Per-head BLASST mask statistics for one topology anchor layer.""" + + retained_tiles: tuple[int, ...] + dropped_mass: tuple[float, ...] + + +def _integer(value: object, name: str, *, minimum: int) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise MaskReuseCalibrationError(f"{name} must be an integer >= {minimum}") + return value + + +def _number( + value: object, + name: str, + *, + minimum: float, + maximum: float | None = None, + minimum_inclusive: bool = True, +) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise MaskReuseCalibrationError(f"{name} must be a finite number") + result = float(value) + below = result < minimum if minimum_inclusive else result <= minimum + if not math.isfinite(result) or below or (maximum is not None and result > maximum): + raise MaskReuseCalibrationError(f"{name} is outside its valid range") + return result + + +def _sha256(value: object, name: str) -> str: + if not isinstance(value, str): + raise MaskReuseCalibrationError(f"{name} must be a lowercase SHA256") + normalized = value.strip().lower() + if len(normalized) != 64 or any( + character not in "0123456789abcdef" for character in normalized + ): + raise MaskReuseCalibrationError(f"{name} must be a lowercase SHA256") + return normalized + + +def _parse_anchor_stats_by_layer( + value: object, + *, + require_canonical_string_keys: bool, +) -> dict[int, AnchorLayerStats]: + if not isinstance(value, Mapping) or not value: + raise MaskReuseCalibrationError("anchor_stats_by_layer must be a non-empty object") + parsed: dict[int, AnchorLayerStats] = {} + for raw_layer, raw_stats in value.items(): + if isinstance(raw_layer, bool): + raise MaskReuseCalibrationError("anchor_stats_by_layer has a non-integer layer key") + try: + layer = int(raw_layer) + except (TypeError, ValueError) as error: + raise MaskReuseCalibrationError( + "anchor_stats_by_layer has a non-integer layer key" + ) from error + if layer < 0 or ( + require_canonical_string_keys + and (not isinstance(raw_layer, str) or raw_layer != str(layer)) + ): + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer layer key {raw_layer!r} is not canonical" + ) + if layer in parsed: + raise MaskReuseCalibrationError(f"anchor_stats_by_layer repeats layer {layer}") + if isinstance(raw_stats, AnchorLayerStats): + raw_retained = raw_stats.retained_tiles + raw_dropped = raw_stats.dropped_mass + elif isinstance(raw_stats, Mapping): + missing = {"retained_tiles", "dropped_mass"} - raw_stats.keys() + extra = raw_stats.keys() - {"retained_tiles", "dropped_mass"} + if missing or extra: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}] requires exactly retained_tiles " + f"and dropped_mass; missing={sorted(missing)}, extra={sorted(extra)}" + ) + raw_retained = raw_stats["retained_tiles"] + raw_dropped = raw_stats["dropped_mass"] + else: + raise MaskReuseCalibrationError(f"anchor_stats_by_layer[{layer}] must be an object") + if not isinstance(raw_retained, list | tuple) or not raw_retained: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}].retained_tiles must be a non-empty list" + ) + if not isinstance(raw_dropped, list | tuple) or not raw_dropped: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}].dropped_mass must be a non-empty list" + ) + retained = tuple( + _integer( + item, + f"anchor_stats_by_layer[{layer}].retained_tiles[{head}]", + minimum=0, + ) + for head, item in enumerate(raw_retained) + ) + dropped = tuple( + _number( + item, + f"anchor_stats_by_layer[{layer}].dropped_mass[{head}]", + minimum=0.0, + maximum=1.0, + ) + for head, item in enumerate(raw_dropped) + ) + if len(retained) != len(dropped): + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}] head arrays differ in width" + ) + parsed[layer] = AnchorLayerStats(retained, dropped) + return dict(sorted(parsed.items())) + + +@dataclass(frozen=True, slots=True) +class MaskReuseObservation: + """One prompt, target-sparsity, consumer-head, and donor-head observation.""" + + model: str + min_kv_tokens: int + max_kv_tokens: int | None + target_sparsity: float + sample_length: int + threshold_lambda: float + threshold_log2: float + q_tokens: int + kv_tokens: int + q_start_tokens: int + split: str + prompt_id: str + source_capture_sha256: str + anchor_layer: int + consumer_layer: int + consumer_head: int + donor_head: int + retained_tiles: int + eligible_tiles: int + anchor_dropped_mass: float + anchor_stats_by_layer: Mapping[int, AnchorLayerStats] + dropped_mass: float + + def __post_init__(self) -> None: + if not isinstance(self.model, str) or not self.model.strip(): + raise MaskReuseCalibrationError("model must be a non-empty string") + if not isinstance(self.prompt_id, str) or not self.prompt_id.strip(): + raise MaskReuseCalibrationError("prompt_id must be a non-empty string") + if self.split not in _SPLITS: + raise MaskReuseCalibrationError(f"split must be one of {sorted(_SPLITS)}") + minimum = _integer(self.min_kv_tokens, "min_kv_tokens", minimum=1) + maximum = self.max_kv_tokens + if maximum is not None: + maximum = _integer(maximum, "max_kv_tokens", minimum=minimum) + sample_length = _integer(self.sample_length, "sample_length", minimum=1) + if sample_length < minimum or (maximum is not None and sample_length > maximum): + raise MaskReuseCalibrationError("sample_length lies outside its context bucket") + q_tokens = _integer(self.q_tokens, "q_tokens", minimum=129) + if q_tokens > int(cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["max_query_chunk_tokens"])): + raise MaskReuseCalibrationError("q_tokens exceeds the deployment geometry limit") + kv_tokens = _integer(self.kv_tokens, "kv_tokens", minimum=1) + q_start_tokens = _integer(self.q_start_tokens, "q_start_tokens", minimum=0) + if sample_length != kv_tokens: + raise MaskReuseCalibrationError("sample_length must equal kv_tokens") + if q_start_tokens + q_tokens != kv_tokens: + raise MaskReuseCalibrationError("q_start_tokens + q_tokens must equal kv_tokens") + alignment = int( + cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["query_chunk_start_alignment_tokens"]) + ) + if q_start_tokens % alignment: + raise MaskReuseCalibrationError("q_start_tokens must be 128-token aligned") + anchor_layer = _integer(self.anchor_layer, "anchor_layer", minimum=0) + consumer_layer = _integer(self.consumer_layer, "consumer_layer", minimum=0) + if anchor_layer >= consumer_layer: + raise MaskReuseCalibrationError("anchor_layer must precede consumer_layer") + _integer(self.consumer_head, "consumer_head", minimum=0) + _integer(self.donor_head, "donor_head", minimum=0) + retained = _integer(self.retained_tiles, "retained_tiles", minimum=0) + eligible = _integer(self.eligible_tiles, "eligible_tiles", minimum=1) + if retained > eligible: + raise MaskReuseCalibrationError("retained_tiles must not exceed eligible_tiles") + + object.__setattr__(self, "model", self.model.strip()) + object.__setattr__(self, "prompt_id", self.prompt_id.strip()) + object.__setattr__( + self, + "source_capture_sha256", + _sha256(self.source_capture_sha256, "source_capture_sha256"), + ) + object.__setattr__( + self, + "target_sparsity", + _number( + self.target_sparsity, + "target_sparsity", + minimum=0.0, + maximum=1.0, + minimum_inclusive=False, + ), + ) + if self.target_sparsity >= 1.0: + raise MaskReuseCalibrationError("target_sparsity must be in (0, 1)") + object.__setattr__( + self, + "threshold_lambda", + _number( + self.threshold_lambda, + "threshold_lambda", + minimum=0.0, + maximum=1.0, + minimum_inclusive=False, + ), + ) + if self.threshold_lambda >= 1.0: + raise MaskReuseCalibrationError("threshold_lambda must be in (0, 1)") + object.__setattr__( + self, + "threshold_log2", + _number(self.threshold_log2, "threshold_log2", minimum=-math.inf, maximum=0.0), + ) + object.__setattr__( + self, + "anchor_dropped_mass", + _number(self.anchor_dropped_mass, "anchor_dropped_mass", minimum=0.0, maximum=1.0), + ) + object.__setattr__( + self, + "anchor_stats_by_layer", + _parse_anchor_stats_by_layer( + self.anchor_stats_by_layer, + require_canonical_string_keys=False, + ), + ) + object.__setattr__( + self, + "dropped_mass", + _number(self.dropped_mass, "dropped_mass", minimum=0.0, maximum=1.0), + ) + + @classmethod + def from_mapping(cls, raw: Mapping[str, object]) -> MaskReuseObservation: + """Build a validated observation from normalized JSON.""" + missing = _OBSERVATION_FIELDS - raw.keys() + extra = raw.keys() - _OBSERVATION_FIELDS + if missing or extra: + raise MaskReuseCalibrationError( + f"observation fields do not match the schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + values = dict(raw) + values["anchor_stats_by_layer"] = _parse_anchor_stats_by_layer( + raw["anchor_stats_by_layer"], + require_canonical_string_keys=True, + ) + return cls(**values) # type: ignore[arg-type] + + def to_mapping(self) -> dict[str, object]: + """Return the normalized JSON representation.""" + result = { + field: getattr(self, field) + for field in _OBSERVATION_FIELDS + if field != "anchor_stats_by_layer" + } + result["anchor_stats_by_layer"] = { + str(layer): { + "retained_tiles": list(stats.retained_tiles), + "dropped_mass": list(stats.dropped_mass), + } + for layer, stats in self.anchor_stats_by_layer.items() + } + return result + + +def _reject_duplicate_json_keys(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise MaskReuseCalibrationError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + +def parse_mask_reuse_observations(lines: Iterable[str]) -> list[MaskReuseObservation]: + """Parse strict normalized observation JSONL.""" + observations: list[MaskReuseObservation] = [] + for line_number, line in enumerate(lines, start=1): + if not line.strip(): + continue + try: + raw = json.loads(line, object_pairs_hook=_reject_duplicate_json_keys) + except json.JSONDecodeError as error: + raise MaskReuseCalibrationError( + f"line {line_number}: invalid JSON: {error.msg}" + ) from error + except MaskReuseCalibrationError as error: + raise MaskReuseCalibrationError(f"line {line_number}: {error}") from error + if not isinstance(raw, dict): + raise MaskReuseCalibrationError(f"line {line_number}: observation must be an object") + try: + observations.append(MaskReuseObservation.from_mapping(raw)) + except MaskReuseCalibrationError as error: + raise MaskReuseCalibrationError(f"line {line_number}: {error}") from error + if not observations: + raise MaskReuseCalibrationError("input contains no mask-reuse observations") + return observations + + +def load_mask_reuse_observations(path: str | Path) -> list[MaskReuseObservation]: + """Load normalized mask-reuse observations from JSONL.""" + with Path(path).open(encoding="utf-8") as handle: + return parse_mask_reuse_observations(handle) + + +def _find_skip_softmax_group(raw: Mapping[str, object]) -> Mapping[str, object]: + current: object = raw + if "sparse_attention_config" in raw: + current = raw["sparse_attention_config"] + if not isinstance(current, Mapping): + raise MaskReuseCalibrationError("sparse_attention_config must be an object") + if "config_groups" in current: + groups = current["config_groups"] + if not isinstance(groups, Mapping): + raise MaskReuseCalibrationError("config_groups must be an object") + matches = [ + group + for group in groups.values() + if isinstance(group, Mapping) + and ( + group.get("algorithm") == "skip_softmax" + or group.get("sparse_algo") == "softmax_skip" + ) + ] + if len(matches) != 1: + raise MaskReuseCalibrationError( + "vanilla config must contain exactly one skip_softmax config group" + ) + selected = matches[0] + if selected.get("sparse_algo") == "softmax_skip" and "threshold_scale_factor" in current: + # Older ModelOpt serving calibration stored the fit beside a + # ``sparse_algo: softmax_skip`` group instead of inside it. + current = current["threshold_scale_factor"] + else: + current = selected + if isinstance(current, Mapping) and "threshold_scale_factor" in current: + current = current["threshold_scale_factor"] + if not isinstance(current, Mapping): + raise MaskReuseCalibrationError("threshold_scale_factor must be an object") + return current + + +def canonical_prefill_threshold_scale_factor( + vanilla_calibration: Mapping[str, object], +) -> dict[str, object]: + """Canonicalize ModelOpt fit parameters or exported skip-softmax metadata.""" + raw = _find_skip_softmax_group(vanilla_calibration) + if "calibration_params" in raw: + raw = _find_skip_softmax_group(raw["calibration_params"]) # type: ignore[arg-type] + formula = raw.get("formula", _FORMULA) + if formula != _FORMULA: + raise MaskReuseCalibrationError("vanilla calibration uses an unsupported formula") + params = raw.get("prefill") + if not isinstance(params, Mapping): + raise MaskReuseCalibrationError("vanilla calibration requires prefill fit parameters") + unknown = params.keys() - { + "a", + "b", + "min_observed_sparsity", + "max_observed_sparsity", + } + if unknown or not {"a", "b"} <= params.keys(): + raise MaskReuseCalibrationError( + f"prefill fit requires a and b and contains unknown fields {sorted(unknown)}" + ) + prefill: dict[str, float] = { + "a": _number(params["a"], "prefill.a", minimum=0.0, minimum_inclusive=False), + "b": _number(params["b"], "prefill.b", minimum=0.0, maximum=20.0), + } + bounds = {"min_observed_sparsity", "max_observed_sparsity"} & params.keys() + if bounds and len(bounds) != 2: + raise MaskReuseCalibrationError("observed sparsity bounds must appear together") + if bounds: + lower = _number( + params["min_observed_sparsity"], + "prefill.min_observed_sparsity", + minimum=0.0, + maximum=1.0, + ) + upper = _number( + params["max_observed_sparsity"], + "prefill.max_observed_sparsity", + minimum=0.0, + maximum=1.0, + ) + if lower > upper: + raise MaskReuseCalibrationError("observed sparsity range is reversed") + prefill.update(min_observed_sparsity=lower, max_observed_sparsity=upper) + return {"formula": _FORMULA, "prefill": prefill} + + +def _normalize_topology(raw: Mapping[str, object]) -> tuple[tuple[int, ...], dict[int, int]]: + if set(raw) != {"anchors", "nearest"}: + raise MaskReuseCalibrationError("topology must contain exactly anchors and nearest") + raw_anchors = raw["anchors"] + if not isinstance(raw_anchors, list) or not raw_anchors: + raise MaskReuseCalibrationError("topology anchors must be a non-empty list") + anchors = tuple(sorted(_integer(value, "topology anchor", minimum=0) for value in raw_anchors)) + if len(anchors) != len(set(anchors)): + raise MaskReuseCalibrationError("topology anchors must be unique") + raw_nearest = raw["nearest"] + if not isinstance(raw_nearest, Mapping): + raise MaskReuseCalibrationError("topology nearest must be an object") + nearest: dict[int, int] = {} + for raw_layer, raw_anchor in raw_nearest.items(): + try: + layer = int(raw_layer) + except (TypeError, ValueError) as error: + raise MaskReuseCalibrationError("topology nearest has a non-integer key") from error + if str(layer) != str(raw_layer) or layer in nearest: + raise MaskReuseCalibrationError("topology nearest keys must be canonical and unique") + nearest[layer] = _integer(raw_anchor, f"topology nearest[{layer}]", minimum=0) + anchor_set = set(anchors) + if not anchor_set <= nearest.keys(): + raise MaskReuseCalibrationError("topology nearest must include every anchor") + for layer, anchor in nearest.items(): + if anchor not in anchor_set or (layer != anchor and anchor >= layer): + raise MaskReuseCalibrationError(f"topology layer {layer} has invalid anchor {anchor}") + if layer in anchor_set and anchor != layer: + raise MaskReuseCalibrationError(f"topology anchor {layer} must map to itself") + if not any(layer != anchor for layer, anchor in nearest.items()): + raise MaskReuseCalibrationError("topology must contain at least one reuse layer") + return anchors, dict(sorted(nearest.items())) + + +def _bucket_key(bucket: Bucket) -> tuple[int, float]: + return bucket[0], math.inf if bucket[1] is None else float(bucket[1]) + + +@dataclass(frozen=True, slots=True) +class _Choice: + donor_head: int + fallback: bool + retained_tiles: int + + +@dataclass(frozen=True, slots=True) +class _Selection: + target_sparsity: float | None + choices: Mapping[ConsumerHead, _Choice] + frontier: tuple[Mapping[str, object], ...] + exact_reason: str | None = None + bmm1_eligible_tiles: int = 0 + bmm1_skipped_tiles: int = 0 + target_bmm1_skip_ratio_met: bool = False + worst_prompt_reuse_dropped_mass: float = 0.0 + mean_prompt_reuse_dropped_mass: float = 0.0 + worst_individual_reuse_dropped_mass: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class _DonorOption: + choice: _Choice + bmm1_skipped_tiles: int + risk_by_prompt: tuple[float, ...] + + +@dataclass(frozen=True, slots=True) +class _BucketIndex: + observations: Mapping[ObservationKey, MaskReuseObservation] + prompts: Mapping[str, tuple[str, ...]] + target_menus: Mapping[tuple[str, str], frozenset[float]] + donor_menus: Mapping[tuple[str, str, float, ConsumerHead], frozenset[int]] + eligible: Mapping[tuple[str, str, ConsumerHead], int] + anchor_masks: Mapping[AnchorKey, tuple[int, int, float]] + anchors: tuple[int, ...] + + +@dataclass(slots=True) +class _ReuseEvaluation: + eligible_tiles: int = 0 + retained_tiles: int = 0 + sparse_observations: int = 0 + violations: int = 0 + dropped_mass_sum: float = 0.0 + worst_dropped_mass: float = 0.0 + + def add(self, other: _ReuseEvaluation) -> None: + self.eligible_tiles += other.eligible_tiles + self.retained_tiles += other.retained_tiles + self.sparse_observations += other.sparse_observations + self.violations += other.violations + self.dropped_mass_sum += other.dropped_mass_sum + self.worst_dropped_mass = max(self.worst_dropped_mass, other.worst_dropped_mass) + + def to_mapping(self) -> dict[str, object]: + return { + "eligible_tiles": self.eligible_tiles, + "retained_tiles": self.retained_tiles, + "bmm1_tile_savings_fraction": ( + 1.0 - self.retained_tiles / self.eligible_tiles if self.eligible_tiles else 0.0 + ), + "sparse_head_prompt_observations": self.sparse_observations, + "report_threshold_exceedance_count": self.violations, + "report_threshold_exceedance_rate": ( + self.violations / self.sparse_observations if self.sparse_observations else 0.0 + ), + "mean_dropped_mass": ( + self.dropped_mass_sum / self.sparse_observations + if self.sparse_observations + else 0.0 + ), + "worst_dropped_mass": self.worst_dropped_mass, + } + + +@dataclass(slots=True) +class _AnchorEvaluation: + eligible_tiles: int = 0 + retained_tiles: int = 0 + prompt_count: int = 0 + violations: int = 0 + prompt_mean_sum: float = 0.0 + worst_prompt_mean: float = 0.0 + + def add(self, other: _AnchorEvaluation) -> None: + self.eligible_tiles += other.eligible_tiles + self.retained_tiles += other.retained_tiles + self.prompt_count += other.prompt_count + self.violations += other.violations + self.prompt_mean_sum += other.prompt_mean_sum + self.worst_prompt_mean = max(self.worst_prompt_mean, other.worst_prompt_mean) + + def to_mapping(self, *, exact: bool = False) -> dict[str, object]: + return { + "policy_exact": exact, + "constraint_statistic": "worst_prompt_mean_anchor_dropped_mass", + "eligible_tiles": self.eligible_tiles, + "retained_tiles": self.retained_tiles, + "bmm2_tile_savings_fraction": ( + 1.0 - self.retained_tiles / self.eligible_tiles if self.eligible_tiles else 0.0 + ), + "evaluated_prompt_count": self.prompt_count, + "constraint_violation_count": self.violations, + "constraint_violation_rate": ( + self.violations / self.prompt_count if self.prompt_count else 0.0 + ), + "mean_prompt_mean_anchor_dropped_mass": ( + self.prompt_mean_sum / self.prompt_count if self.prompt_count else 0.0 + ), + "worst_prompt_mean_anchor_dropped_mass": self.worst_prompt_mean, + } + + +def _normalize_observations( + values: Sequence[MaskReuseObservation | Mapping[str, object]], +) -> list[MaskReuseObservation]: + if not values: + raise MaskReuseCalibrationError("at least one observation is required") + return [ + value + if isinstance(value, MaskReuseObservation) + else MaskReuseObservation.from_mapping(value) + for value in values + ] + + +def _validate_thresholds( + observations: Sequence[MaskReuseObservation], threshold_scale_factor: Mapping[str, object] +) -> None: + params = threshold_scale_factor["prefill"] + assert isinstance(params, Mapping) + a = float(params["a"]) + b = float(params["b"]) + lower = params.get("min_observed_sparsity") + upper = params.get("max_observed_sparsity") + for observation in observations: + target = observation.target_sparsity + if (lower is not None and target < float(lower)) or ( + upper is not None and target > float(upper) + ): + raise MaskReuseCalibrationError( + f"target_sparsity={target} is outside the observed vanilla calibration range" + ) + expected_log2 = ( + math.log2(a) + b * target * math.log2(math.e) - math.log2(observation.sample_length) + ) + expected_lambda = 2.0**expected_log2 + if not 0.0 < expected_lambda < 1.0: + raise MaskReuseCalibrationError( + "vanilla calibration derives a threshold outside (0, 1) for " + f"prompt={observation.prompt_id!r}, target_sparsity={target}" + ) + if observation.threshold_log2.hex() != expected_log2.hex(): + raise MaskReuseCalibrationError( + "threshold_log2 does not match the log-domain vanilla fit for " + f"prompt={observation.prompt_id!r}, target_sparsity={target}: " + f"observed={observation.threshold_log2.hex()}, " + f"expected={expected_log2.hex()}" + ) + if observation.threshold_lambda.hex() != expected_lambda.hex(): + raise MaskReuseCalibrationError( + "threshold_lambda does not match " + "exp2(log2(a) + b * target_sparsity * log2(e) - log2(sample_length)) for " + f"prompt={observation.prompt_id!r}, target_sparsity={target}: " + f"observed={observation.threshold_lambda.hex()}, " + f"expected={expected_lambda.hex()}" + ) + + +def _validate_dataset( + observations: Sequence[MaskReuseObservation], + *, + nearest: Mapping[int, int], +) -> tuple[ + str, int, tuple[Bucket, ...], tuple[ConsumerHead, ...], dict[Bucket, list[MaskReuseObservation]] +]: + by_split = { + split: [observation for observation in observations if observation.split == split] + for split in _SPLITS + } + if any(not rows for rows in by_split.values()): + raise MaskReuseCalibrationError("observations require calibration and heldout splits") + models = {observation.model for observation in observations} + if len(models) != 1: + raise MaskReuseCalibrationError("observations must contain exactly one model") + calibration_prompts = {row.prompt_id for row in by_split["calibration"]} + heldout_prompts = {row.prompt_id for row in by_split["heldout"]} + if calibration_prompts & heldout_prompts: + raise MaskReuseCalibrationError("prompt IDs overlap calibration and heldout splits") + calibration_sources = {row.source_capture_sha256 for row in by_split["calibration"]} + heldout_sources = {row.source_capture_sha256 for row in by_split["heldout"]} + if calibration_sources & heldout_sources: + raise MaskReuseCalibrationError("source captures overlap calibration and heldout splits") + + consumer_to_anchor = {layer: anchor for layer, anchor in nearest.items() if layer != anchor} + max_head = max(max(row.consumer_head, row.donor_head) for row in observations) + global_num_heads = max_head + 1 + targets = tuple( + (layer, head) for layer in sorted(consumer_to_anchor) for head in range(global_num_heads) + ) + expected_targets = set(targets) + seen: set[tuple[object, ...]] = set() + capture_sources: dict[tuple[str, str, Bucket], str] = {} + by_bucket: dict[Bucket, list[MaskReuseObservation]] = defaultdict(list) + for row in observations: + expected_anchor = consumer_to_anchor.get(row.consumer_layer) + if expected_anchor != row.anchor_layer: + raise MaskReuseCalibrationError( + f"consumer layer {row.consumer_layer} does not match the explicit topology" + ) + bucket = (row.min_kv_tokens, row.max_kv_tokens) + by_bucket[bucket].append(row) + identity = ( + row.model, + bucket, + row.split, + row.prompt_id, + row.target_sparsity, + row.consumer_layer, + row.consumer_head, + row.donor_head, + ) + if identity in seen: + raise MaskReuseCalibrationError("observations contain a duplicate candidate row") + seen.add(identity) + capture = (row.split, row.prompt_id, bucket) + previous_source = capture_sources.setdefault(capture, row.source_capture_sha256) + if previous_source != row.source_capture_sha256: + raise MaskReuseCalibrationError("one prompt/context capture has multiple fingerprints") + + split_buckets = { + split: {(row.min_kv_tokens, row.max_kv_tokens) for row in rows} + for split, rows in by_split.items() + } + if split_buckets["calibration"] != split_buckets["heldout"]: + raise MaskReuseCalibrationError("heldout context buckets must match calibration buckets") + buckets = tuple(sorted(split_buckets["calibration"], key=_bucket_key)) + previous_max: int | None = 0 + for minimum, maximum in buckets: + if previous_max is None or minimum <= previous_max: + raise MaskReuseCalibrationError("context buckets must be ordered and non-overlapping") + previous_max = maximum + for bucket in buckets: + for split in _SPLITS: + actual = { + (row.consumer_layer, row.consumer_head) + for row in by_bucket[bucket] + if row.split == split + } + if actual != expected_targets: + raise MaskReuseCalibrationError( + f"{split} bucket {bucket} does not cover every consumer head" + ) + return next(iter(models)), global_num_heads, buckets, targets, dict(by_bucket) + + +def _index_bucket( + rows: Sequence[MaskReuseObservation], + *, + anchors: tuple[int, ...], + global_num_heads: int, + targets: Sequence[ConsumerHead], +) -> _BucketIndex: + observations: dict[ObservationKey, MaskReuseObservation] = {} + prompts: dict[str, set[str]] = defaultdict(set) + target_menus: dict[tuple[str, str], set[float]] = defaultdict(set) + donor_menus: dict[tuple[str, str, float, ConsumerHead], set[int]] = defaultdict(set) + eligible: dict[tuple[str, str, ConsumerHead], int] = {} + anchor_masks: dict[AnchorKey, tuple[int, int, float]] = {} + prompt_targets: dict[tuple[str, str, float], tuple[int, float, float]] = {} + anchor_payloads: dict[tuple[str, str, float], Mapping[int, AnchorLayerStats]] = {} + payload_eligible: dict[tuple[str, str, float], int] = {} + for row in rows: + target = (row.consumer_layer, row.consumer_head) + key: ObservationKey = ( + row.split, + row.prompt_id, + row.target_sparsity, + row.consumer_layer, + row.consumer_head, + row.donor_head, + ) + observations[key] = row + prompts[row.split].add(row.prompt_id) + target_menus[(row.split, row.prompt_id)].add(row.target_sparsity) + donor_menus[(row.split, row.prompt_id, row.target_sparsity, target)].add(row.donor_head) + eligible_key = (row.split, row.prompt_id, target) + previous_eligible = eligible.setdefault(eligible_key, row.eligible_tiles) + if previous_eligible != row.eligible_tiles: + raise MaskReuseCalibrationError("eligible_tiles differs across candidates") + prompt_target = (row.split, row.prompt_id, row.target_sparsity) + sample_and_threshold = ( + row.sample_length, + row.threshold_lambda, + row.threshold_log2, + ) + if prompt_targets.setdefault(prompt_target, sample_and_threshold) != sample_and_threshold: + raise MaskReuseCalibrationError( + "sample_length, threshold_lambda, or threshold_log2 differs within a capture" + ) + previous_payload = anchor_payloads.setdefault(prompt_target, row.anchor_stats_by_layer) + if previous_payload != row.anchor_stats_by_layer: + raise MaskReuseCalibrationError( + "anchor_stats_by_layer differs across repeated candidate rows" + ) + previous_payload_eligible = payload_eligible.setdefault(prompt_target, row.eligible_tiles) + if previous_payload_eligible != row.eligible_tiles: + raise MaskReuseCalibrationError( + "eligible_tiles differs within one capture and target_sparsity" + ) + + expected_anchors = set(anchors) + for prompt_target, payload in anchor_payloads.items(): + actual_anchors = set(payload) + if actual_anchors != expected_anchors: + raise MaskReuseCalibrationError( + "anchor_stats_by_layer does not exactly cover topology anchors; " + f"missing={sorted(expected_anchors - actual_anchors)}, " + f"extra={sorted(actual_anchors - expected_anchors)}" + ) + eligible_tiles = payload_eligible[prompt_target] + for anchor, stats in payload.items(): + if len(stats.retained_tiles) != global_num_heads: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{anchor}] has head width " + f"{len(stats.retained_tiles)}, expected {global_num_heads}" + ) + for head, (retained, dropped) in enumerate( + zip(stats.retained_tiles, stats.dropped_mass, strict=True) + ): + if retained > eligible_tiles: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{anchor}].retained_tiles[{head}] " + "exceeds eligible_tiles" + ) + split, prompt, target_sparsity = prompt_target + anchor_masks[(split, prompt, target_sparsity, anchor, head)] = ( + retained, + eligible_tiles, + dropped, + ) + + for row in rows: + stats = row.anchor_stats_by_layer[row.anchor_layer] + payload_candidate = ( + stats.retained_tiles[row.donor_head], + stats.dropped_mass[row.donor_head], + ) + if payload_candidate != (row.retained_tiles, row.anchor_dropped_mass): + raise MaskReuseCalibrationError( + "candidate retained_tiles/anchor_dropped_mass does not match anchor_stats_by_layer" + ) + + first_prompt = min(prompts["calibration"]) + expected_targets = target_menus[("calibration", first_prompt)] + full_donors = set(range(global_num_heads)) + for split in _SPLITS: + for prompt in prompts[split]: + if target_menus[(split, prompt)] != expected_targets: + raise MaskReuseCalibrationError("target_sparsity menu differs across captures") + for target_sparsity in expected_targets: + for target in targets: + if donor_menus[(split, prompt, target_sparsity, target)] != full_donors: + raise MaskReuseCalibrationError("candidate donor menu is incomplete") + for anchor in anchors: + for head in range(global_num_heads): + if (split, prompt, target_sparsity, anchor, head) not in anchor_masks: + raise MaskReuseCalibrationError( + "anchor/head mask observations are incomplete" + ) + return _BucketIndex( + observations=observations, + prompts={split: tuple(sorted(values)) for split, values in prompts.items()}, + target_menus={key: frozenset(values) for key, values in target_menus.items()}, + donor_menus={key: frozenset(values) for key, values in donor_menus.items()}, + eligible=eligible, + anchor_masks=anchor_masks, + anchors=anchors, + ) + + +def _evaluate_anchor( + index: _BucketIndex, + *, + split: str, + target_sparsity: float | None, + global_num_heads: int, + maximum: float, +) -> _AnchorEvaluation: + result = _AnchorEvaluation() + for prompt in index.prompts[split]: + selected_target = ( + min(index.target_menus[(split, prompt)]) if target_sparsity is None else target_sparsity + ) + values: list[float] = [] + retained_tiles = 0 + eligible_tiles = 0 + for anchor in index.anchors: + for head in range(global_num_heads): + retained, eligible, dropped = index.anchor_masks[ + (split, prompt, selected_target, anchor, head) + ] + values.append(0.0 if target_sparsity is None else dropped) + retained_tiles += eligible if target_sparsity is None else retained + eligible_tiles += eligible + prompt_mean = sum(values) / len(values) + result.eligible_tiles += eligible_tiles + result.retained_tiles += retained_tiles + result.prompt_count += 1 + result.prompt_mean_sum += prompt_mean + result.worst_prompt_mean = max(result.worst_prompt_mean, prompt_mean) + result.violations += int(prompt_mean > maximum) + return result + + +def _select_bucket( + index: _BucketIndex, + *, + targets: Sequence[ConsumerHead], + global_num_heads: int, + max_anchor_dropped_mass: float, + target_bmm1_skip_ratio: float, +) -> _Selection: + prompts = index.prompts["calibration"] + target_menu = tuple(sorted(index.target_menus[("calibration", prompts[0])])) + frontier: list[Mapping[str, object]] = [] + candidates: list[tuple[tuple[object, ...], _Selection]] = [] + maximum_candidates: list[tuple[tuple[object, ...], _Selection]] = [] + + def donor_options(target_sparsity: float, target: ConsumerHead) -> tuple[_DonorOption, ...]: + eligible = sum(index.eligible[("calibration", prompt, target)] for prompt in prompts) + options = [ + _DonorOption( + _Choice(0, True, eligible), + 0, + tuple(0.0 for _ in prompts), + ) + ] + for donor in range(global_num_heads): + rows = [ + index.observations[ + ( + "calibration", + prompt, + target_sparsity, + target[0], + target[1], + donor, + ) + ] + for prompt in prompts + ] + retained = sum(row.retained_tiles for row in rows) + options.append( + _DonorOption( + _Choice(donor, False, retained), + eligible - retained, + tuple(row.dropped_mass for row in rows), + ) + ) + pareto = [] + for candidate_index, candidate in enumerate(options): + dominated = False + for other_index, other in enumerate(options): + if candidate_index == other_index: + continue + no_worse = other.bmm1_skipped_tiles >= candidate.bmm1_skipped_tiles and all( + other_risk <= candidate_risk + for other_risk, candidate_risk in zip( + other.risk_by_prompt, + candidate.risk_by_prompt, + strict=True, + ) + ) + strictly_better = other.bmm1_skipped_tiles > candidate.bmm1_skipped_tiles or any( + other_risk < candidate_risk + for other_risk, candidate_risk in zip( + other.risk_by_prompt, + candidate.risk_by_prompt, + strict=True, + ) + ) + canonical_tie = not strictly_better and ( + (other.choice.fallback and not candidate.choice.fallback) + or ( + other.choice.fallback == candidate.choice.fallback + and other.choice.donor_head < candidate.choice.donor_head + ) + ) + if no_worse and (strictly_better or canonical_tie): + dominated = True + break + if not dominated: + pareto.append(candidate) + return tuple(pareto) + + def solve_target( + target_sparsity: float, + *, + minimum_bmm1_skipped_tiles: int | None, + maximize_bmm1_skipped_tiles: bool, + target_met: bool, + ) -> _Selection | None: + problem = pulp.LpProblem("legacy_mask_reuse", pulp.LpMinimize) + variables: dict[tuple[ConsumerHead, int], pulp.LpVariable] = {} + options: dict[tuple[ConsumerHead, int], _DonorOption] = {} + prompt_risk_terms: dict[str, list[object]] = defaultdict(list) + for target_index, target in enumerate(targets): + menu = donor_options(target_sparsity, target) + choice_variables = [] + for option_index, option in enumerate(menu): + key = (target, option_index) + variable = pulp.LpVariable( + f"choice_{target_index}_{option_index}", + lowBound=0, + upBound=1, + cat="Binary", + ) + variables[key] = variable + options[key] = option + choice_variables.append(variable) + for prompt, risk in zip(prompts, option.risk_by_prompt, strict=True): + prompt_risk_terms[prompt].append(risk * variable) + problem += pulp.lpSum(choice_variables) == 1, f"choose_{target_index}" + + bmm1_skipped = pulp.lpSum( + options[key].bmm1_skipped_tiles * variable for key, variable in variables.items() + ) + retained_reuse = pulp.lpSum( + options[key].choice.retained_tiles * variable for key, variable in variables.items() + ) + reuse_count = pulp.lpSum( + int(not options[key].choice.fallback) * variable for key, variable in variables.items() + ) + all_layer_head_count = len(index.anchors) * global_num_heads + len(targets) + worst_prompt_risk = pulp.LpVariable("worst_prompt_reuse_dropped_mass", lowBound=0.0) + for prompt_index, prompt in enumerate(prompts): + problem += ( + pulp.lpSum(prompt_risk_terms[prompt]) <= all_layer_head_count * worst_prompt_risk, + f"reuse_risk_{prompt_index}", + ) + if minimum_bmm1_skipped_tiles is not None: + problem += bmm1_skipped >= minimum_bmm1_skipped_tiles, "minimum_bmm1_skips" + solver = pulp.PULP_CBC_CMD(msg=False, threads=1, options=["randomSeed 0"]) + warm_solver = pulp.PULP_CBC_CMD( + msg=False, + threads=1, + options=["randomSeed 0"], + warmStart=True, + ) + has_incumbent = False + + def minimize(expression: object) -> bool: + nonlocal has_incumbent + problem.setObjective(expression) + status = problem.solve(warm_solver if has_incumbent else solver) + has_incumbent = status == pulp.LpStatusOptimal + return has_incumbent + + def minimize_and_fix(expression: object, name: str, *, integral: bool) -> bool: + nonlocal problem + if not minimize(expression): + return False + raw_value = pulp.value(expression) + value = 0.0 if raw_value is None else float(raw_value) + if integral: + problem += expression == round(value), name + else: + problem += expression <= value + _SOLVER_LEXICOGRAPHIC_ATOL, name + return True + + if maximize_bmm1_skipped_tiles and not minimize_and_fix( + -bmm1_skipped, "fix_maximum_bmm1_skips", integral=True + ): + return None + if not minimize_and_fix(worst_prompt_risk, "fix_worst_prompt_risk", integral=False): + return None + if not minimize_and_fix(retained_reuse, "fix_retained_reuse", integral=True): + raise MaskReuseCalibrationError( + "legacy selector lost feasibility after fixing worst-prompt reuse risk" + ) + donor_signature = pulp.lpSum( + (options[key].choice.donor_head + 1) * variable + for key, variable in variables.items() + if not options[key].choice.fallback + ) + donor_base = global_num_heads * len(targets) + 1 + if not minimize(reuse_count * donor_base + donor_signature): + raise MaskReuseCalibrationError( + "legacy selector lost deterministic tie-break feasibility" + ) + + choices: dict[ConsumerHead, _Choice] = {} + prompt_totals = dict.fromkeys(prompts, 0.0) + worst_individual = 0.0 + skipped_tiles = 0 + for key, variable in variables.items(): + if variable.value() <= 0.5: + continue + target, _ = key + option = options[key] + choices[target] = option.choice + skipped_tiles += option.bmm1_skipped_tiles + if not option.choice.fallback: + for prompt, risk in zip(prompts, option.risk_by_prompt, strict=True): + prompt_totals[prompt] += risk + worst_individual = max(worst_individual, risk) + eligible_per_head = sum( + index.eligible[("calibration", prompt, targets[0])] for prompt in prompts + ) + eligible_tiles = eligible_per_head * all_layer_head_count + prompt_risks = [prompt_totals[prompt] / all_layer_head_count for prompt in prompts] + return _Selection( + target_sparsity, + choices, + (), + bmm1_eligible_tiles=eligible_tiles, + bmm1_skipped_tiles=skipped_tiles, + target_bmm1_skip_ratio_met=target_met, + worst_prompt_reuse_dropped_mass=max(prompt_risks, default=0.0), + mean_prompt_reuse_dropped_mass=( + sum(prompt_risks) / len(prompt_risks) if prompt_risks else 0.0 + ), + worst_individual_reuse_dropped_mass=worst_individual, + ) + + for target_sparsity in target_menu: + anchor = _evaluate_anchor( + index, + split="calibration", + target_sparsity=target_sparsity, + global_num_heads=global_num_heads, + maximum=max_anchor_dropped_mass, + ) + eligible_per_head = sum( + index.eligible[("calibration", prompt, targets[0])] for prompt in prompts + ) + eligible_tiles = eligible_per_head * (len(index.anchors) * global_num_heads + len(targets)) + required_tiles = math.ceil(target_bmm1_skip_ratio * eligible_tiles) + selected = None + if anchor.violations == 0: + selected = solve_target( + target_sparsity, + minimum_bmm1_skipped_tiles=required_tiles, + maximize_bmm1_skipped_tiles=False, + target_met=True, + ) + total_retained = ( + None + if selected is None + else sum(choice.retained_tiles for choice in selected.choices.values()) + ) + fallback_count = ( + None + if selected is None + else sum(choice.fallback for choice in selected.choices.values()) + ) + combined_tile_cost = ( + None if total_retained is None else 2 * total_retained + anchor.retained_tiles + ) + frontier.append( + { + "target_sparsity": target_sparsity, + "anchor_safe": anchor.violations == 0, + "target_bmm1_skip_ratio": target_bmm1_skip_ratio, + "target_bmm1_skip_ratio_feasible": selected is not None, + "retained_reuse_tiles": total_retained, + "retained_anchor_tiles": anchor.retained_tiles, + "combined_tile_cost": combined_tile_cost, + "fallback_head_count": fallback_count, + "anchor_calibration": anchor.to_mapping(), + } + ) + if selected is not None: + rank = ( + selected.worst_prompt_reuse_dropped_mass, + combined_tile_cost, + target_sparsity, + ) + candidates.append((rank, selected)) + elif anchor.violations == 0: + maximum = solve_target( + target_sparsity, + minimum_bmm1_skipped_tiles=None, + maximize_bmm1_skipped_tiles=True, + target_met=False, + ) + if maximum is not None: + maximum_candidates.append( + ( + ( + -maximum.bmm1_skipped_tiles, + maximum.worst_prompt_reuse_dropped_mass, + target_sparsity, + ), + maximum, + ) + ) + if not candidates and not maximum_candidates: + dense_choices = { + target: _Choice( + 0, True, sum(index.eligible[("calibration", prompt, target)] for prompt in prompts) + ) + for target in targets + } + return _Selection( + None, + dense_choices, + tuple(frontier), + "no_target_sparsity_satisfied_anchor_calibration_constraint", + bmm1_eligible_tiles=( + sum(index.eligible[("calibration", prompt, targets[0])] for prompt in prompts) + * (len(index.anchors) * global_num_heads + len(targets)) + ), + ) + if candidates: + _, selected = min(candidates, key=lambda item: item[0]) + else: + _, selected = min(maximum_candidates, key=lambda item: item[0]) + return _Selection( + selected.target_sparsity, + selected.choices, + tuple(frontier), + bmm1_eligible_tiles=selected.bmm1_eligible_tiles, + bmm1_skipped_tiles=selected.bmm1_skipped_tiles, + target_bmm1_skip_ratio_met=selected.target_bmm1_skip_ratio_met, + worst_prompt_reuse_dropped_mass=selected.worst_prompt_reuse_dropped_mass, + mean_prompt_reuse_dropped_mass=selected.mean_prompt_reuse_dropped_mass, + worst_individual_reuse_dropped_mass=selected.worst_individual_reuse_dropped_mass, + ) + + +def _evaluate_reuse( + selection: _Selection, + index: _BucketIndex, + *, + split: str, + maximum: float, +) -> _ReuseEvaluation: + result = _ReuseEvaluation() + for prompt in index.prompts[split]: + for target, choice in sorted(selection.choices.items()): + eligible = index.eligible[(split, prompt, target)] + result.eligible_tiles += eligible + if choice.fallback: + result.retained_tiles += eligible + continue + assert selection.target_sparsity is not None + row = index.observations[ + ( + split, + prompt, + selection.target_sparsity, + target[0], + target[1], + choice.donor_head, + ) + ] + result.retained_tiles += row.retained_tiles + result.sparse_observations += 1 + result.dropped_mass_sum += row.dropped_mass + result.worst_dropped_mass = max(result.worst_dropped_mass, row.dropped_mass) + result.violations += int(row.dropped_mass > maximum) + return result + + +def _canonical_digest(observations: Sequence[MaskReuseObservation]) -> str: + rows = [ + json.dumps(row.to_mapping(), sort_keys=True, separators=(",", ":")).encode() + for row in observations + ] + digest = sha256() + for row in sorted(rows): + digest.update(sha256(row).digest()) + return digest.hexdigest() + + +def _deployment_geometry( + observations: Sequence[MaskReuseObservation], +) -> dict[str, object]: + by_capture: dict[tuple[object, ...], tuple[int, int, int]] = {} + for row in observations: + identity = ( + row.split, + row.prompt_id, + row.source_capture_sha256, + row.min_kv_tokens, + row.max_kv_tokens, + ) + geometry = (row.q_tokens, row.kv_tokens, row.q_start_tokens) + if by_capture.setdefault(identity, geometry) != geometry: + raise MaskReuseCalibrationError( + "q_tokens, kv_tokens, or q_start_tokens differs within one capture" + ) + geometry_rows = [ + { + "split": identity[0], + "prompt_id": identity[1], + "source_capture_sha256": identity[2], + "min_kv_tokens": identity[3], + "max_kv_tokens": identity[4], + "q_tokens": geometry[0], + "kv_tokens": geometry[1], + "q_start_tokens": geometry[2], + } + for identity, geometry in sorted(by_capture.items(), key=lambda item: str(item[0])) + ] + return { + "contract": dict(_DEPLOYMENT_GEOMETRY_CONTRACT), + "observations": geometry_rows, + } + + +def _canonical_evidence(raw: Mapping[str, object]) -> dict[str, str]: + missing = _EVIDENCE_FIELDS - raw.keys() + extra = raw.keys() - _EVIDENCE_FIELDS + if missing or extra: + raise MaskReuseCalibrationError( + f"evidence fields do not match schema v3; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + return {field: _sha256(raw[field], f"evidence.{field}") for field in sorted(raw)} + + +def calibrate_mask_reuse_policy( + observations: Sequence[MaskReuseObservation | Mapping[str, object]], + *, + vanilla_calibration: Mapping[str, object], + topology: Mapping[str, object], + checkpoint_manifest: VerifiedCheckpointManifest, + evidence: Mapping[str, object], + max_anchor_dropped_mass: float, + reuse_dropped_mass_report_threshold: float, + target_bmm1_skip_ratio: float, + source_provenance: Mapping[str, object] | None = None, +) -> dict[str, object]: + """Select a minimum-risk legacy candidate under a per-bucket BMM1 target.""" + if not isinstance(checkpoint_manifest, VerifiedCheckpointManifest): + raise MaskReuseCalibrationError( + "checkpoint_manifest must be returned by verify_checkpoint_manifest" + ) + rows = _normalize_observations(observations) + threshold_scale_factor = canonical_prefill_threshold_scale_factor(vanilla_calibration) + prefill_params = threshold_scale_factor["prefill"] + assert isinstance(prefill_params, Mapping) + if not {"min_observed_sparsity", "max_observed_sparsity"} <= prefill_params.keys(): + raise MaskReuseCalibrationError( + "schema-v3 export requires observed prefill sparsity bounds" + ) + _validate_thresholds(rows, threshold_scale_factor) + anchors, nearest = _normalize_topology(topology) + checkpoint_identity = checkpoint_manifest.sha256 + geometry = _deployment_geometry(rows) + canonical_evidence = _canonical_evidence(evidence) + anchor_bound = _number( + max_anchor_dropped_mass, "max_anchor_dropped_mass", minimum=0.0, maximum=1.0 + ) + reuse_report_threshold = _number( + reuse_dropped_mass_report_threshold, + "reuse_dropped_mass_report_threshold", + minimum=0.0, + maximum=1.0, + ) + bmm1_target = _number( + target_bmm1_skip_ratio, + "target_bmm1_skip_ratio", + minimum=0.0, + maximum=1.0, + ) + model, global_num_heads, buckets, targets, by_bucket = _validate_dataset(rows, nearest=nearest) + if model != checkpoint_manifest.model: + raise MaskReuseCalibrationError( + "observation model does not match the verified checkpoint manifest" + ) + + context_policies: list[dict[str, object]] = [] + bucket_reports: list[dict[str, object]] = [] + target_menus: list[dict[str, object]] = [] + overall_reuse_calibration = _ReuseEvaluation() + overall_reuse_heldout = _ReuseEvaluation() + overall_anchor_calibration = _AnchorEvaluation() + overall_anchor_heldout = _AnchorEvaluation() + total_fallback = 0 + total_bmm1_eligible = 0 + total_bmm1_skipped = 0 + all_bmm1_targets_met = True + for bounds in buckets: + index = _index_bucket( + by_bucket[bounds], + anchors=anchors, + global_num_heads=global_num_heads, + targets=targets, + ) + selection = _select_bucket( + index, + targets=targets, + global_num_heads=global_num_heads, + max_anchor_dropped_mass=anchor_bound, + target_bmm1_skip_ratio=bmm1_target, + ) + if selection.target_sparsity is not None and bounds[1] is None: + raise MaskReuseCalibrationError( + "a deployment-qualified sparse context bucket requires a finite maximum" + ) + reuse_calibration = _evaluate_reuse( + selection, index, split="calibration", maximum=reuse_report_threshold + ) + reuse_heldout = _evaluate_reuse( + selection, index, split="heldout", maximum=reuse_report_threshold + ) + anchor_calibration = _evaluate_anchor( + index, + split="calibration", + target_sparsity=selection.target_sparsity, + global_num_heads=global_num_heads, + maximum=anchor_bound, + ) + anchor_heldout = _evaluate_anchor( + index, + split="heldout", + target_sparsity=selection.target_sparsity, + global_num_heads=global_num_heads, + maximum=anchor_bound, + ) + overall_reuse_calibration.add(reuse_calibration) + overall_reuse_heldout.add(reuse_heldout) + overall_anchor_calibration.add(anchor_calibration) + overall_anchor_heldout.add(anchor_heldout) + + policy: dict[str, object] + if selection.target_sparsity is None: + policy = { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "exact": True, + } + headmaps: dict[str, list[int]] = {} + fallback_heads: dict[str, list[int]] = {} + else: + headmaps = { + str(layer): [ + selection.choices[(layer, head)].donor_head for head in range(global_num_heads) + ] + for layer in sorted({target[0] for target in targets}) + } + fallback_heads = { + str(layer): [ + head + for head in range(global_num_heads) + if selection.choices[(layer, head)].fallback + ] + for layer in sorted({target[0] for target in targets}) + } + policy = { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "target_sparsity": selection.target_sparsity, + "headmaps": headmaps, + "fallback_heads": fallback_heads, + } + context_policies.append(policy) + fallback_count = sum(len(heads) for heads in fallback_heads.values()) + total_fallback += fallback_count + total_bmm1_eligible += selection.bmm1_eligible_tiles + total_bmm1_skipped += selection.bmm1_skipped_tiles + all_bmm1_targets_met &= selection.target_bmm1_skip_ratio_met + menu = [row["target_sparsity"] for row in selection.frontier] + target_menus.append( + { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "target_sparsities": menu, + } + ) + bucket_reports.append( + { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "selected_target_sparsity": selection.target_sparsity, + "selection_status": ( + "target_bmm1_skip_ratio_met" + if selection.target_sparsity is not None + and selection.target_bmm1_skip_ratio_met + else ( + "target_bmm1_skip_ratio_unmet_maximum_feasible" + if selection.target_sparsity is not None + else selection.exact_reason + ) + ), + "target_sparsity_frontier": list(selection.frontier), + "fallback_head_count": fallback_count, + "bmm1_skip_objective": { + "target": bmm1_target, + "target_met": selection.target_bmm1_skip_ratio_met, + "eligible_tiles": selection.bmm1_eligible_tiles, + "skipped_tiles": selection.bmm1_skipped_tiles, + "achieved": ( + selection.bmm1_skipped_tiles / selection.bmm1_eligible_tiles + if selection.bmm1_eligible_tiles + else 0.0 + ), + }, + "reuse_selection_objective": { + "hard_maximum": None, + "worst_development_prompt_model_wide_dropped_mass": ( + selection.worst_prompt_reuse_dropped_mass + ), + "mean_development_prompt_model_wide_dropped_mass": ( + selection.mean_prompt_reuse_dropped_mass + ), + "worst_individual_dropped_mass": ( + selection.worst_individual_reuse_dropped_mass + ), + }, + "reuse_calibration": reuse_calibration.to_mapping(), + "reuse_heldout": reuse_heldout.to_mapping(), + "anchor_calibration": anchor_calibration.to_mapping( + exact=selection.target_sparsity is None + ), + "anchor_heldout": anchor_heldout.to_mapping( + exact=selection.target_sparsity is None + ), + } + ) + + provenance: dict[str, object] = { + "calibrator": "modelopt.mask_reuse", + "observation_schema_version": 1, + "input_observation_count": len(rows), + "canonical_input_sha256": _canonical_digest(rows), + "calibration_prompt_ids": sorted( + {row.prompt_id for row in rows if row.split == "calibration"} + ), + "heldout_prompt_ids": sorted({row.prompt_id for row in rows if row.split == "heldout"}), + "selection_split": "calibration", + "evaluation_split": "heldout", + "threshold_semantics": "a * exp(b * target_sparsity) / sample_length", + "threshold_implementation": { + "threshold_log2": ("log2(a) + b * target_sparsity * log2(e) - log2(sample_length)"), + "threshold_lambda": "exp2(threshold_log2)", + "validation": "exact IEEE-754 binary64 hex equality", + }, + "checkpoint_manifest_sha256": checkpoint_identity, + "target_sparsity_menus": target_menus, + "constraints": { + "anchor": { + "metric": "worst_prompt_mean_anchor_dropped_mass", + "comparison": "<=", + "maximum": anchor_bound, + }, + "reuse": { + "selection_metric": ( + "per_prompt_mean_across_all_attention_layers_and_heads_reuse_dropped_mass" + ), + "selection_hard_maximum": None, + "report_metric": "per_prompt_candidate_reuse_dropped_mass", + "report_threshold": reuse_report_threshold, + "report_threshold_affects_selection": False, + }, + "bmm1_skip_ratio": { + "metric": "per_context_bucket_model_wide_eligible_bmm1_tile_skip_ratio", + "comparison": ">=", + "target": bmm1_target, + }, + }, + "tie_breaks": [ + "reject target sparsities exceeding the calibration anchor bound", + "meet the requested BMM1 skip ratio in every context bucket", + "minimum worst-prompt model-wide reuse dropped mass", + "minimum equal-BMM combined tile cost 2*A_R + A_A", + "fewest reused consumer heads", + "smallest target sparsity", + "lexicographically smallest donor-head map", + ], + "selection_cost": { + "formula": "2 * retained_reuse_tiles + retained_anchor_tiles", + "bmm1_weight": 1.0, + "bmm2_weight": 1.0, + }, + } + if source_provenance is not None: + provenance["source"] = dict(source_provenance) + + report = { + "model": model, + "checkpoint_manifest_sha256": checkpoint_identity, + "constraints": provenance["constraints"], + "selection_unit": "context_bucket", + "by_bucket": bucket_reports, + "overall": { + "consumer_head_bucket_count": len(targets) * len(buckets), + "fallback_head_bucket_count": total_fallback, + "fallback_fraction": total_fallback / (len(targets) * len(buckets)), + "bmm1_skip_objective": { + "target_per_context_bucket": bmm1_target, + "all_context_buckets_met": all_bmm1_targets_met, + "eligible_tiles": total_bmm1_eligible, + "skipped_tiles": total_bmm1_skipped, + "achieved_model_wide": ( + total_bmm1_skipped / total_bmm1_eligible if total_bmm1_eligible else 0.0 + ), + }, + "reuse_calibration": overall_reuse_calibration.to_mapping(), + "reuse_heldout": overall_reuse_heldout.to_mapping(), + "anchor_calibration": overall_anchor_calibration.to_mapping(), + "anchor_heldout": overall_anchor_heldout.to_mapping(), + }, + "promotion": { + "status": "candidate_only", + "eligible": False, + "reasons": [ + "legacy observations are not capture-schema-v2 checkpoint-bound records", + "grouped inner-fold and preregistered outer gates were not evaluated", + "deployment rectangular-geometry promotion gate was not evaluated", + ], + }, + } + return { + "version": 3, + "promotion_status": "candidate_only", + "phase": "prefill", + "decode": {"mode": "dense"}, + "calibration_protocol": _CALIBRATION_PROTOCOL, + "producer": {"name": "modelopt", "version": modelopt.__version__}, + "evidence": canonical_evidence, + "threshold_scale_factor": threshold_scale_factor, + "model": model, + "checkpoint_manifest_sha256": checkpoint_identity, + "global_num_heads": global_num_heads, + "target_bmm1_skip_ratio": bmm1_target, + "anchors": list(anchors), + "nearest": {str(layer): anchor for layer, anchor in nearest.items()}, + "deployment_geometry_validated": False, + "deployment_geometry": geometry, + "context_policies": context_policies, + "provenance": provenance, + "calibration_report": report, + } diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py new file mode 100644 index 00000000000..5c7c38a51fd --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py @@ -0,0 +1,1545 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Streaming selector for compact mask-reuse capture JSONL. + +One compact record stores the anchor vectors and consumer ``[H, H]`` risk +matrices for a single prompt and target sparsity. The selector makes three +semantic passes (validation, calibration selection, frozen-policy evaluation) +and hashes the exact file bytes before and after selection/evaluation to detect +concurrent mutation. Selection retains development risk matrices long enough to +optimize the worst prompt under the requested BMM1 target; held-out captures +remain evaluation-only. +""" + +from __future__ import annotations + +import json +import math +import unicodedata +from collections import defaultdict +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import cast + +import numpy as np +import pulp + +import modelopt + +from .checkpoint_manifest import VerifiedCheckpointManifest +from .mask_reuse import ( + _CALIBRATION_PROTOCOL, + _DEPLOYMENT_GEOMETRY_CONTRACT, + _SOLVER_LEXICOGRAPHIC_ATOL, + AnchorLayerStats, + Bucket, + ConsumerHead, + MaskReuseCalibrationError, + _AnchorEvaluation, + _bucket_key, + _canonical_evidence, + _Choice, + _normalize_topology, + _number, + _ReuseEvaluation, + _Selection, + _sha256, + canonical_prefill_threshold_scale_factor, +) + +__all__ = [ + "CompactMaskReuseCapture", + "CompactMaskReuseCaptureSource", + "calibrate_compact_mask_reuse_policy", + "load_compact_mask_reuse_captures", +] + +_CAPTURE_FIELDS = frozenset( + { + "compact_capture_schema_version", + "invocation", + "geometry", + "global_num_heads", + "eligible_tiles", + "anchor_stats_by_layer", + "consumer_layers", + } +) +_INVOCATION_FIELDS = frozenset( + { + "capture_schema_version", + "model", + "checkpoint_manifest_sha256", + "split", + "partition", + "inner_fold", + "prompt_id", + "source", + "source_group_sha256", + "source_capture_sha256", + "min_kv_tokens", + "max_kv_tokens", + "target_sparsity_hex", + "sample_length", + "threshold_log2_hex", + "threshold_lambda_hex", + "expected_geometry", + } +) +_GEOMETRY_FIELDS = frozenset({"q_tokens", "kv_tokens", "q_start_tokens"}) +_ANCHOR_FIELDS = frozenset({"retained_tiles", "dropped_mass"}) +_CONSUMER_FIELDS = frozenset({"anchor_layer", "dropped_mass"}) +_SPLITS = ("calibration", "heldout") +_PARTITIONS = ("development", "outer_test") +_MONOTONIC_ATOL = 1e-7 + + +def _exact_fields(raw: Mapping[str, object], expected: frozenset[str], label: str) -> None: + missing = expected - raw.keys() + extra = raw.keys() - expected + if missing or extra: + raise MaskReuseCalibrationError( + f"{label} fields do not match the schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + + +def _integer(value: object, label: str, *, minimum: int = 0) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise MaskReuseCalibrationError(f"{label} must be an integer >= {minimum}") + return value + + +def _text(value: object, label: str) -> str: + if not isinstance(value, str) or not value: + raise MaskReuseCalibrationError(f"{label} must be a non-empty string") + if unicodedata.normalize("NFC", value) != value or any( + ord(character) < 32 for character in value + ): + raise MaskReuseCalibrationError(f"{label} must be NFC text without control characters") + return value + + +def _float_hex(value: object, label: str) -> float: + if not isinstance(value, str): + raise MaskReuseCalibrationError(f"{label} must be a canonical float.hex string") + try: + parsed = float.fromhex(value) + except ValueError as error: + raise MaskReuseCalibrationError(f"{label} must be a canonical float.hex string") from error + if not math.isfinite(parsed) or parsed.hex() != value: + raise MaskReuseCalibrationError(f"{label} must be a canonical finite float.hex string") + return parsed + + +def _geometry(raw: object, label: str) -> dict[str, int]: + if not isinstance(raw, Mapping): + raise MaskReuseCalibrationError(f"{label} must be an object") + _exact_fields(raw, _GEOMETRY_FIELDS, label) + q_tokens = _integer(raw["q_tokens"], f"{label}.q_tokens", minimum=129) + if q_tokens > int(cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["max_query_chunk_tokens"])): + raise MaskReuseCalibrationError(f"{label}.q_tokens exceeds the deployment maximum") + kv_tokens = _integer(raw["kv_tokens"], f"{label}.kv_tokens", minimum=q_tokens) + q_start = _integer(raw["q_start_tokens"], f"{label}.q_start_tokens") + alignment = int( + cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["query_chunk_start_alignment_tokens"]) + ) + if q_start % alignment or q_start + q_tokens != kv_tokens: + raise MaskReuseCalibrationError(f"{label} is not a bottom-right aligned final chunk") + return {"q_tokens": q_tokens, "kv_tokens": kv_tokens, "q_start_tokens": q_start} + + +def _eligible_tiles(geometry: Mapping[str, int]) -> int: + q_blocks = (geometry["q_tokens"] + 127) // 128 + first_eligible = geometry["q_start_tokens"] // 128 + 1 + return q_blocks * (2 * first_eligible + q_blocks - 1) // 2 + + +@dataclass(frozen=True, slots=True) +class CompactConsumerStats: + """One consumer layer's global consumer-by-donor dropped-mass matrix.""" + + anchor_layer: int + dropped_mass: Sequence[Sequence[float]] + + +@dataclass(frozen=True, slots=True) +class CompactMaskReuseCapture: + """One prompt/target capture decoded without expanding candidate rows.""" + + model: str + checkpoint_manifest_sha256: str + split: str + partition: str + inner_fold: int | None + prompt_id: str + source: str + source_group_sha256: str + source_capture_sha256: str + min_kv_tokens: int + max_kv_tokens: int | None + target_sparsity: float + sample_length: int + threshold_log2: float + threshold_lambda: float + geometry: Mapping[str, int] + global_num_heads: int + eligible_tiles: int + anchor_stats_by_layer: Mapping[int, AnchorLayerStats] + consumer_layers: Mapping[int, CompactConsumerStats] + + @property + def bucket(self) -> Bucket: + """Return this capture's context bounds.""" + return self.min_kv_tokens, self.max_kv_tokens + + @classmethod + def from_mapping(cls, raw: Mapping[str, object]) -> CompactMaskReuseCapture: + """Parse one strict compact-capture object.""" + _exact_fields(raw, _CAPTURE_FIELDS, "compact capture") + if raw["compact_capture_schema_version"] != 1: + raise MaskReuseCalibrationError("compact_capture_schema_version must be 1") + invocation = raw["invocation"] + if not isinstance(invocation, Mapping): + raise MaskReuseCalibrationError("compact capture invocation must be an object") + _exact_fields(invocation, _INVOCATION_FIELDS, "compact capture invocation") + if invocation["capture_schema_version"] != 2: + raise MaskReuseCalibrationError("capture_schema_version must be 2") + split = invocation["split"] + if split not in _SPLITS: + raise MaskReuseCalibrationError("capture split must be calibration or heldout") + partition = invocation["partition"] + if partition not in _PARTITIONS: + raise MaskReuseCalibrationError("capture partition must be development or outer_test") + expected_split = "calibration" if partition == "development" else "heldout" + if split != expected_split: + raise MaskReuseCalibrationError("capture split and partition disagree") + raw_fold = invocation["inner_fold"] + if partition == "development": + inner_fold = _integer(raw_fold, "inner_fold") + elif raw_fold is not None: + raise MaskReuseCalibrationError("outer_test capture must have null inner_fold") + else: + inner_fold = None + minimum = _integer(invocation["min_kv_tokens"], "min_kv_tokens", minimum=1) + maximum = invocation["max_kv_tokens"] + if maximum is not None: + maximum = _integer(maximum, "max_kv_tokens", minimum=minimum) + sample_length = _integer(invocation["sample_length"], "sample_length", minimum=1) + if sample_length < minimum or (maximum is not None and sample_length > maximum): + raise MaskReuseCalibrationError("sample_length lies outside its context bucket") + target = _float_hex(invocation["target_sparsity_hex"], "target_sparsity_hex") + threshold_log2 = _float_hex(invocation["threshold_log2_hex"], "threshold_log2_hex") + threshold_lambda = _float_hex(invocation["threshold_lambda_hex"], "threshold_lambda_hex") + if not 0.0 < target < 1.0: + raise MaskReuseCalibrationError("target_sparsity must be in (0, 1)") + if threshold_log2 >= 0.0 or not 0.0 < threshold_lambda < 1.0: + raise MaskReuseCalibrationError("threshold must be in (0, 1)") + if (2.0**threshold_log2).hex() != threshold_lambda.hex(): + raise MaskReuseCalibrationError("threshold lambda and log2 fields disagree") + expected_geometry = _geometry(invocation["expected_geometry"], "expected_geometry") + observed_geometry = _geometry(raw["geometry"], "geometry") + if ( + expected_geometry != observed_geometry + or observed_geometry["kv_tokens"] != sample_length + ): + raise MaskReuseCalibrationError("capture geometry does not match its invocation") + global_num_heads = _integer(raw["global_num_heads"], "global_num_heads", minimum=1) + eligible_tiles = _integer(raw["eligible_tiles"], "eligible_tiles", minimum=1) + if eligible_tiles != _eligible_tiles(observed_geometry): + raise MaskReuseCalibrationError( + "eligible_tiles does not match 128x128 bottom-right causal geometry" + ) + + raw_anchors = raw["anchor_stats_by_layer"] + if not isinstance(raw_anchors, Mapping) or not raw_anchors: + raise MaskReuseCalibrationError("anchor_stats_by_layer must be non-empty") + anchors: dict[int, AnchorLayerStats] = {} + for raw_layer, raw_stats in raw_anchors.items(): + if ( + not isinstance(raw_layer, str) + or not raw_layer.isdigit() + or raw_layer != str(int(raw_layer)) + ): + raise MaskReuseCalibrationError("anchor layer keys must be canonical integers") + layer = int(raw_layer) + if not isinstance(raw_stats, Mapping): + raise MaskReuseCalibrationError(f"anchor_stats_by_layer[{layer}] must be an object") + _exact_fields(raw_stats, _ANCHOR_FIELDS, f"anchor_stats_by_layer[{layer}]") + retained_raw = raw_stats["retained_tiles"] + dropped_raw = raw_stats["dropped_mass"] + if not isinstance(retained_raw, list) or not isinstance(dropped_raw, list): + raise MaskReuseCalibrationError(f"anchor_stats_by_layer[{layer}] must use arrays") + if len(retained_raw) != global_num_heads or len(dropped_raw) != global_num_heads: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}] does not cover all global heads" + ) + retained = tuple( + _integer(value, f"anchor {layer} retained[{head}]") + for head, value in enumerate(retained_raw) + ) + if any(value > eligible_tiles for value in retained): + raise MaskReuseCalibrationError(f"anchor {layer} retained tiles exceed eligible") + dropped = tuple( + _number(value, f"anchor {layer} dropped[{head}]", minimum=0.0, maximum=1.0) + for head, value in enumerate(dropped_raw) + ) + anchors[layer] = AnchorLayerStats(retained, dropped) + + raw_consumers = raw["consumer_layers"] + if not isinstance(raw_consumers, Mapping) or not raw_consumers: + raise MaskReuseCalibrationError("consumer_layers must be non-empty") + consumers: dict[int, CompactConsumerStats] = {} + for raw_layer, raw_stats in raw_consumers.items(): + if ( + not isinstance(raw_layer, str) + or not raw_layer.isdigit() + or raw_layer != str(int(raw_layer)) + ): + raise MaskReuseCalibrationError("consumer layer keys must be canonical integers") + layer = int(raw_layer) + if not isinstance(raw_stats, Mapping): + raise MaskReuseCalibrationError(f"consumer_layers[{layer}] must be an object") + _exact_fields(raw_stats, _CONSUMER_FIELDS, f"consumer_layers[{layer}]") + anchor = _integer(raw_stats["anchor_layer"], f"consumer_layers[{layer}].anchor") + matrix = raw_stats["dropped_mass"] + if not isinstance(matrix, list) or len(matrix) != global_num_heads: + raise MaskReuseCalibrationError( + f"consumer_layers[{layer}] must have {global_num_heads} consumer rows" + ) + for consumer_head, row in enumerate(matrix): + if not isinstance(row, list) or len(row) != global_num_heads: + raise MaskReuseCalibrationError( + f"consumer_layers[{layer}][{consumer_head}] must cover all donors" + ) + for donor_head, value in enumerate(row): + _number( + value, + f"consumer_layers[{layer}][{consumer_head}][{donor_head}]", + minimum=0.0, + maximum=1.0, + ) + consumers[layer] = CompactConsumerStats(anchor, matrix) + return cls( + model=_text(invocation["model"], "model"), + checkpoint_manifest_sha256=_sha256( + invocation["checkpoint_manifest_sha256"], "checkpoint_manifest_sha256" + ), + split=split, + partition=partition, + inner_fold=inner_fold, + prompt_id=_text(invocation["prompt_id"], "prompt_id"), + source=_text(invocation["source"], "source"), + source_group_sha256=_sha256(invocation["source_group_sha256"], "source_group_sha256"), + source_capture_sha256=_sha256( + invocation["source_capture_sha256"], "source_capture_sha256" + ), + min_kv_tokens=minimum, + max_kv_tokens=maximum, + target_sparsity=target, + sample_length=sample_length, + threshold_log2=threshold_log2, + threshold_lambda=threshold_lambda, + geometry=observed_geometry, + global_num_heads=global_num_heads, + eligible_tiles=eligible_tiles, + anchor_stats_by_layer=dict(sorted(anchors.items())), + consumer_layers=dict(sorted(consumers.items())), + ) + + +def _reject_duplicate_json_keys(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise MaskReuseCalibrationError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + +@dataclass(frozen=True, slots=True) +class CompactMaskReuseCaptureSource: + """Re-iterable strict JSONL source used by the three-pass selector.""" + + path: Path + + def __iter__(self) -> Iterator[CompactMaskReuseCapture]: + with self.path.open(encoding="utf-8") as handle: + seen = False + for line_number, line in enumerate(handle, start=1): + if not line.strip(): + continue + seen = True + try: + raw = json.loads(line, object_pairs_hook=_reject_duplicate_json_keys) + except json.JSONDecodeError as error: + raise MaskReuseCalibrationError( + f"line {line_number}: invalid JSON: {error.msg}" + ) from error + except MaskReuseCalibrationError as error: + raise MaskReuseCalibrationError(f"line {line_number}: {error}") from error + if not isinstance(raw, dict): + raise MaskReuseCalibrationError( + f"line {line_number}: compact capture must be an object" + ) + try: + yield CompactMaskReuseCapture.from_mapping(raw) + except MaskReuseCalibrationError as error: + raise MaskReuseCalibrationError(f"line {line_number}: {error}") from error + if not seen: + raise MaskReuseCalibrationError("compact capture input is empty") + + def sha256(self) -> str: + """Hash exact file bytes without loading the capture bundle.""" + digest = sha256() + with self.path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def load_compact_mask_reuse_captures(path: str | Path) -> CompactMaskReuseCaptureSource: + """Return a lazy, re-iterable compact-capture source.""" + source = CompactMaskReuseCaptureSource(Path(path)) + if not source.path.is_file(): + raise MaskReuseCalibrationError(f"compact capture file does not exist: {source.path}") + return source + + +@dataclass(frozen=True, slots=True) +class _Dataset: + model: str + checkpoint_manifest_sha256: str + global_num_heads: int + buckets: tuple[Bucket, ...] + anchors: tuple[int, ...] + nearest: Mapping[int, int] + consumer_layers: tuple[int, ...] + targets: tuple[ConsumerHead, ...] + prompts: Mapping[tuple[Bucket, str], tuple[str, ...]] + menus: Mapping[Bucket, tuple[float, ...]] + deployment_geometry: Mapping[str, object] + prompt_sources: tuple[Mapping[str, object], ...] + calibration_prompt_ids: tuple[str, ...] + heldout_prompt_ids: tuple[str, ...] + capture_count: int + + +def _validate_threshold(capture: CompactMaskReuseCapture, fit: Mapping[str, object]) -> None: + params = fit["prefill"] + assert isinstance(params, Mapping) + lower = params.get("min_observed_sparsity") + upper = params.get("max_observed_sparsity") + if (lower is not None and capture.target_sparsity < float(lower)) or ( + upper is not None and capture.target_sparsity > float(upper) + ): + raise MaskReuseCalibrationError( + f"target_sparsity={capture.target_sparsity} is outside the vanilla fit range" + ) + expected_log2 = ( + math.log2(float(params["a"])) + + float(params["b"]) * capture.target_sparsity * math.log2(math.e) + - math.log2(capture.sample_length) + ) + expected_lambda = 2.0**expected_log2 + if capture.threshold_log2.hex() != expected_log2.hex(): + raise MaskReuseCalibrationError("compact capture threshold_log2 differs from vanilla fit") + if capture.threshold_lambda.hex() != expected_lambda.hex(): + raise MaskReuseCalibrationError("compact capture threshold_lambda differs from vanilla fit") + + +def _capture_order_key(capture: CompactMaskReuseCapture) -> tuple[object, ...]: + return ( + capture.min_kv_tokens, + math.inf if capture.max_kv_tokens is None else capture.max_kv_tokens, + _SPLITS.index(capture.split), + capture.prompt_id, + capture.target_sparsity, + ) + + +def _validate_monotonic_pair( + previous: CompactMaskReuseCapture, current: CompactMaskReuseCapture +) -> None: + identity = (current.bucket, current.split, current.prompt_id) + if identity != (previous.bucket, previous.split, previous.prompt_id): + return + if current.target_sparsity <= previous.target_sparsity: + raise MaskReuseCalibrationError("target sparsities must be strictly increasing per prompt") + for layer, current_stats in current.anchor_stats_by_layer.items(): + previous_stats = previous.anchor_stats_by_layer[layer] + if any( + current_value > previous_value + for current_value, previous_value in zip( + current_stats.retained_tiles, + previous_stats.retained_tiles, + strict=True, + ) + ): + raise MaskReuseCalibrationError( + f"anchor {layer} retained counts increase with target sparsity" + ) + if any( + current_value + _MONOTONIC_ATOL < previous_value + for current_value, previous_value in zip( + current_stats.dropped_mass, + previous_stats.dropped_mass, + strict=True, + ) + ): + raise MaskReuseCalibrationError( + f"anchor {layer} dropped mass decreases with target sparsity" + ) + for layer, current_consumer_stats in current.consumer_layers.items(): + previous_consumer_stats = previous.consumer_layers[layer] + for consumer_head, (current_row, previous_row) in enumerate( + zip( + current_consumer_stats.dropped_mass, + previous_consumer_stats.dropped_mass, + strict=True, + ) + ): + if any( + current_value + _MONOTONIC_ATOL < previous_value + for current_value, previous_value in zip(current_row, previous_row, strict=True) + ): + raise MaskReuseCalibrationError( + f"consumer {layer} head {consumer_head} dropped mass decreases " + "with target sparsity" + ) + + +def _validate_dataset( + source: CompactMaskReuseCaptureSource, + *, + fit: Mapping[str, object], + anchors: tuple[int, ...], + nearest: Mapping[int, int], +) -> _Dataset: + models: set[str] = set() + checkpoint_identities: set[str] = set() + head_counts: set[int] = set() + split_counts: dict[str, int] = defaultdict(int) + split_buckets: dict[str, set[Bucket]] = defaultdict(set) + prompts: dict[tuple[Bucket, str], set[str]] = defaultdict(set) + menus: dict[tuple[Bucket, str, str], set[float]] = defaultdict(set) + capture_sources: dict[tuple[Bucket, str, str], tuple[object, ...]] = {} + fingerprint_owner: dict[tuple[str, str], tuple[Bucket, str]] = {} + seen: set[tuple[Bucket, str, str, float]] = set() + geometry_by_prompt: dict[tuple[Bucket, str, str], tuple[object, ...]] = {} + prompt_source_rows: dict[tuple[Bucket, str, str], Mapping[str, object]] = {} + split_prompt_ids: dict[str, set[str]] = defaultdict(set) + prompt_buckets: dict[tuple[str, str], Bucket] = {} + split_fingerprints: dict[str, set[str]] = defaultdict(set) + group_assignments: dict[str, tuple[str, int | None]] = {} + fingerprint_groups: dict[str, str] = {} + expected_anchor_set = set(anchors) + consumer_to_anchor = {layer: anchor for layer, anchor in nearest.items() if layer != anchor} + expected_consumer_set = set(consumer_to_anchor) + capture_count = 0 + previous_order: tuple[object, ...] | None = None + previous_capture: CompactMaskReuseCapture | None = None + for capture in source: + capture_count += 1 + order = _capture_order_key(capture) + if previous_order is not None and order <= previous_order: + raise MaskReuseCalibrationError("compact capture records are not in canonical order") + if previous_capture is not None: + _validate_monotonic_pair(previous_capture, capture) + previous_order = order + previous_capture = capture + _validate_threshold(capture, fit) + if set(capture.anchor_stats_by_layer) != expected_anchor_set: + raise MaskReuseCalibrationError("compact capture does not cover every topology anchor") + if set(capture.consumer_layers) != expected_consumer_set: + raise MaskReuseCalibrationError("compact capture does not cover every reuse layer") + for layer, stats in capture.consumer_layers.items(): + if stats.anchor_layer != consumer_to_anchor[layer]: + raise MaskReuseCalibrationError( + f"consumer layer {layer} does not match the explicit topology" + ) + models.add(capture.model) + checkpoint_identities.add(capture.checkpoint_manifest_sha256) + head_counts.add(capture.global_num_heads) + split_counts[capture.split] += 1 + split_buckets[capture.split].add(capture.bucket) + prompts[(capture.bucket, capture.split)].add(capture.prompt_id) + menus[(capture.bucket, capture.split, capture.prompt_id)].add(capture.target_sparsity) + split_prompt_ids[capture.split].add(capture.prompt_id) + prompt_identity = (capture.split, capture.prompt_id) + if prompt_buckets.setdefault(prompt_identity, capture.bucket) != capture.bucket: + raise MaskReuseCalibrationError("a prompt ID is assigned to multiple context buckets") + split_fingerprints[capture.split].add(capture.source_capture_sha256) + assignment = (capture.partition, capture.inner_fold) + previous_assignment = group_assignments.setdefault(capture.source_group_sha256, assignment) + if previous_assignment != assignment: + raise MaskReuseCalibrationError( + "one source group is assigned to multiple partitions or inner folds" + ) + previous_group = fingerprint_groups.setdefault( + capture.source_capture_sha256, capture.source_group_sha256 + ) + if previous_group != capture.source_group_sha256: + raise MaskReuseCalibrationError( + "one rendered source capture is assigned to multiple source groups" + ) + key = (capture.bucket, capture.split, capture.prompt_id, capture.target_sparsity) + if key in seen: + raise MaskReuseCalibrationError("compact captures contain a duplicate prompt target") + seen.add(key) + prompt_key = (capture.bucket, capture.split, capture.prompt_id) + source_identity = ( + capture.source, + capture.source_group_sha256, + capture.partition, + capture.inner_fold, + capture.source_capture_sha256, + capture.sample_length, + tuple(sorted(capture.geometry.items())), + capture.eligible_tiles, + ) + if capture_sources.setdefault(prompt_key, source_identity) != source_identity: + raise MaskReuseCalibrationError("prompt metadata differs across target sparsities") + owner_key = (capture.split, capture.source_capture_sha256) + owner = (capture.bucket, capture.prompt_id) + if fingerprint_owner.setdefault(owner_key, owner) != owner: + raise MaskReuseCalibrationError("one source fingerprint names multiple prompt captures") + geometry_identity = ( + capture.source_capture_sha256, + capture.geometry["q_tokens"], + capture.geometry["kv_tokens"], + capture.geometry["q_start_tokens"], + ) + if geometry_by_prompt.setdefault(prompt_key, geometry_identity) != geometry_identity: + raise MaskReuseCalibrationError("prompt geometry differs across target sparsities") + prompt_source_rows[prompt_key] = { + "min_kv_tokens": capture.min_kv_tokens, + "max_kv_tokens": capture.max_kv_tokens, + "split": capture.split, + "prompt_id": capture.prompt_id, + "source": capture.source, + "source_group_sha256": capture.source_group_sha256, + "partition": capture.partition, + "inner_fold": capture.inner_fold, + "source_capture_sha256": capture.source_capture_sha256, + } + if capture_count == 0 or any(not split_counts[split] for split in _SPLITS): + raise MaskReuseCalibrationError("compact captures require calibration and heldout splits") + if len(models) != 1 or len(head_counts) != 1 or len(checkpoint_identities) != 1: + raise MaskReuseCalibrationError( + "compact captures must use one model, checkpoint, and head count" + ) + if split_prompt_ids["calibration"] & split_prompt_ids["heldout"]: + raise MaskReuseCalibrationError("prompt IDs overlap calibration and heldout splits") + if split_fingerprints["calibration"] & split_fingerprints["heldout"]: + raise MaskReuseCalibrationError("source captures overlap calibration and heldout splits") + if split_buckets["calibration"] != split_buckets["heldout"]: + raise MaskReuseCalibrationError("heldout context buckets must match calibration buckets") + buckets = tuple(sorted(split_buckets["calibration"], key=_bucket_key)) + previous_max: int | None = 0 + for minimum, maximum in buckets: + if previous_max is None or minimum <= previous_max: + raise MaskReuseCalibrationError("context buckets must be ordered and non-overlapping") + previous_max = maximum + canonical_menus: dict[Bucket, tuple[float, ...]] = {} + for bucket in buckets: + calibration_prompts = sorted(prompts[(bucket, "calibration")]) + if not calibration_prompts: + raise MaskReuseCalibrationError(f"bucket {bucket} has no calibration prompts") + menu = menus[(bucket, "calibration", calibration_prompts[0])] + if not menu: + raise MaskReuseCalibrationError(f"bucket {bucket} has no target menu") + for split in _SPLITS: + for prompt in prompts[(bucket, split)]: + if menus[(bucket, split, prompt)] != menu: + raise MaskReuseCalibrationError("target-sparsity menu differs across prompts") + canonical_menus[bucket] = tuple(sorted(menu)) + + geometry_rows = [] + for (bucket, split, prompt), identity in sorted( + geometry_by_prompt.items(), key=lambda item: str(item[0]) + ): + geometry_rows.append( + { + "split": split, + "prompt_id": prompt, + "source_capture_sha256": identity[0], + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "q_tokens": identity[1], + "kv_tokens": identity[2], + "q_start_tokens": identity[3], + } + ) + global_num_heads = next(iter(head_counts)) + consumer_layers = tuple(sorted(expected_consumer_set)) + targets = tuple((layer, head) for layer in consumer_layers for head in range(global_num_heads)) + return _Dataset( + model=next(iter(models)), + checkpoint_manifest_sha256=next(iter(checkpoint_identities)), + global_num_heads=global_num_heads, + buckets=buckets, + anchors=anchors, + nearest=nearest, + consumer_layers=consumer_layers, + targets=targets, + prompts={key: tuple(sorted(value)) for key, value in prompts.items()}, + menus=canonical_menus, + deployment_geometry={ + "contract": dict(_DEPLOYMENT_GEOMETRY_CONTRACT), + "observations": geometry_rows, + }, + prompt_sources=tuple( + prompt_source_rows[key] for key in sorted(prompt_source_rows, key=str) + ), + calibration_prompt_ids=tuple(sorted(split_prompt_ids["calibration"])), + heldout_prompt_ids=tuple(sorted(split_prompt_ids["heldout"])), + capture_count=capture_count, + ) + + +@dataclass(slots=True) +class _SelectionAccumulator: + prompt_mass: Mapping[float, dict[str, np.ndarray]] + retained_by_anchor: Mapping[float, np.ndarray] + eligible_sum: dict[float, int] + anchor_evaluation: Mapping[float, _AnchorEvaluation] + + +@dataclass(frozen=True, slots=True) +class _DonorOption: + choice: _Choice + bmm1_skipped_tiles: int + risk_by_prompt: tuple[float, ...] + + +def _selection_pass( + source: CompactMaskReuseCaptureSource, + dataset: _Dataset, + *, + max_anchor_dropped_mass: float, + target_bmm1_skip_ratio: float, +) -> dict[Bucket, _Selection]: + consumer_index = {layer: index for index, layer in enumerate(dataset.consumer_layers)} + anchor_index = {layer: index for index, layer in enumerate(dataset.anchors)} + accumulators: dict[Bucket, _SelectionAccumulator] = {} + for bucket in dataset.buckets: + menus = dataset.menus[bucket] + accumulators[bucket] = _SelectionAccumulator( + prompt_mass={target: {} for target in menus}, + retained_by_anchor={ + target: np.zeros((len(dataset.anchors), dataset.global_num_heads), dtype=np.int64) + for target in menus + }, + eligible_sum=dict.fromkeys(menus, 0), + anchor_evaluation={target: _AnchorEvaluation() for target in menus}, + ) + for capture in source: + if capture.split != "calibration": + continue + accumulator = accumulators[capture.bucket] + target = capture.target_sparsity + accumulator.eligible_sum[target] += capture.eligible_tiles + for layer, stats in capture.anchor_stats_by_layer.items(): + accumulator.retained_by_anchor[target][anchor_index[layer]] += np.asarray( + stats.retained_tiles, dtype=np.int64 + ) + evaluation = accumulator.anchor_evaluation[target] + dropped = [ + value + for stats in capture.anchor_stats_by_layer.values() + for value in stats.dropped_mass + ] + prompt_mean = sum(dropped) / len(dropped) + evaluation.eligible_tiles += ( + capture.eligible_tiles * len(dataset.anchors) * dataset.global_num_heads + ) + evaluation.retained_tiles += sum( + sum(stats.retained_tiles) for stats in capture.anchor_stats_by_layer.values() + ) + evaluation.prompt_count += 1 + evaluation.prompt_mean_sum += prompt_mean + evaluation.worst_prompt_mean = max(evaluation.worst_prompt_mean, prompt_mean) + evaluation.violations += int(prompt_mean > max_anchor_dropped_mass) + for layer, stats in capture.consumer_layers.items(): + if capture.prompt_id not in accumulator.prompt_mass[target]: + accumulator.prompt_mass[target][capture.prompt_id] = np.empty( + ( + len(dataset.consumer_layers), + dataset.global_num_heads, + dataset.global_num_heads, + ), + dtype=np.float64, + ) + accumulator.prompt_mass[target][capture.prompt_id][consumer_index[layer]] = np.asarray( + stats.dropped_mass, dtype=np.float64 + ) + + def donor_options( + accumulator: _SelectionAccumulator, + *, + target: float, + layer: int, + head: int, + retained: np.ndarray, + prompts: tuple[str, ...], + ) -> tuple[_DonorOption, ...]: + options = [ + _DonorOption( + _Choice(0, True, accumulator.eligible_sum[target]), + 0, + tuple(0.0 for _ in prompts), + ) + ] + layer_row = consumer_index[layer] + for donor in range(dataset.global_num_heads): + retained_tiles = int(retained[donor]) + options.append( + _DonorOption( + _Choice(donor, False, retained_tiles), + accumulator.eligible_sum[target] - retained_tiles, + tuple( + float(accumulator.prompt_mass[target][prompt][layer_row, head, donor]) + for prompt in prompts + ), + ) + ) + pareto = [] + for candidate_index, candidate in enumerate(options): + dominated = False + for other_index, other in enumerate(options): + if candidate_index == other_index: + continue + no_worse = other.bmm1_skipped_tiles >= candidate.bmm1_skipped_tiles and all( + other_risk <= candidate_risk + for other_risk, candidate_risk in zip( + other.risk_by_prompt, + candidate.risk_by_prompt, + strict=True, + ) + ) + strictly_better = other.bmm1_skipped_tiles > candidate.bmm1_skipped_tiles or any( + other_risk < candidate_risk + for other_risk, candidate_risk in zip( + other.risk_by_prompt, + candidate.risk_by_prompt, + strict=True, + ) + ) + canonical_tie = not strictly_better and ( + (other.choice.fallback and not candidate.choice.fallback) + or ( + other.choice.fallback == candidate.choice.fallback + and other.choice.donor_head < candidate.choice.donor_head + ) + ) + if no_worse and (strictly_better or canonical_tie): + dominated = True + break + if not dominated: + pareto.append(candidate) + return tuple(pareto) + + def solve_target( + bucket: Bucket, + target: float, + *, + minimum_bmm1_skipped_tiles: int | None, + maximize_bmm1_skipped_tiles: bool, + target_met: bool, + ) -> _Selection | None: + accumulator = accumulators[bucket] + prompts = dataset.prompts[(bucket, "calibration")] + problem = pulp.LpProblem("compact_mask_reuse", pulp.LpMinimize) + variables: dict[tuple[int, int, int], pulp.LpVariable] = {} + options: dict[tuple[int, int, int], _DonorOption] = {} + prompt_risk_terms: dict[str, list[object]] = defaultdict(list) + for layer in dataset.consumer_layers: + anchor = dataset.nearest[layer] + retained = accumulator.retained_by_anchor[target][anchor_index[anchor]] + for head in range(dataset.global_num_heads): + menu = donor_options( + accumulator, + target=target, + layer=layer, + head=head, + retained=retained, + prompts=prompts, + ) + choice_variables = [] + for option_index, option in enumerate(menu): + key = (layer, head, option_index) + variable = pulp.LpVariable( + f"choice_{layer}_{head}_{option_index}", + lowBound=0, + upBound=1, + cat="Binary", + ) + variables[key] = variable + options[key] = option + choice_variables.append(variable) + for prompt, risk in zip(prompts, option.risk_by_prompt, strict=True): + prompt_risk_terms[prompt].append(risk * variable) + problem += pulp.lpSum(choice_variables) == 1, f"choose_{layer}_{head}" + + bmm1_skipped = pulp.lpSum( + options[key].bmm1_skipped_tiles * variable for key, variable in variables.items() + ) + retained_reuse = pulp.lpSum( + options[key].choice.retained_tiles * variable for key, variable in variables.items() + ) + reuse_count = pulp.lpSum( + int(not options[key].choice.fallback) * variable for key, variable in variables.items() + ) + normalizer = len(dataset.nearest) * dataset.global_num_heads + worst_prompt_risk = pulp.LpVariable("worst_prompt_reuse_dropped_mass", lowBound=0.0) + for prompt_index, prompt in enumerate(prompts): + problem += ( + pulp.lpSum(prompt_risk_terms[prompt]) <= normalizer * worst_prompt_risk, + f"reuse_risk_{prompt_index}", + ) + if minimum_bmm1_skipped_tiles is not None: + problem += bmm1_skipped >= minimum_bmm1_skipped_tiles, "minimum_bmm1_skips" + solver = pulp.PULP_CBC_CMD(msg=False, threads=1, options=["randomSeed 0"]) + warm_solver = pulp.PULP_CBC_CMD( + msg=False, + threads=1, + options=["randomSeed 0"], + warmStart=True, + ) + has_incumbent = False + + def minimize(expression: object) -> bool: + nonlocal has_incumbent + problem.setObjective(expression) + status = problem.solve(warm_solver if has_incumbent else solver) + has_incumbent = status == pulp.LpStatusOptimal + return has_incumbent + + def minimize_and_fix(expression: object, name: str, *, integral: bool) -> bool: + nonlocal problem + if not minimize(expression): + return False + raw_value = pulp.value(expression) + value = 0.0 if raw_value is None else float(raw_value) + if integral: + problem += expression == round(value), name + else: + problem += expression <= value + _SOLVER_LEXICOGRAPHIC_ATOL, name + return True + + if maximize_bmm1_skipped_tiles and not minimize_and_fix( + -bmm1_skipped, "fix_maximum_bmm1_skips", integral=True + ): + return None + if not minimize_and_fix(worst_prompt_risk, "fix_worst_prompt_risk", integral=False): + return None + if not minimize_and_fix(retained_reuse, "fix_retained_reuse", integral=True): + raise MaskReuseCalibrationError( + "compact selector lost feasibility after fixing worst-prompt reuse risk" + ) + donor_signature = pulp.lpSum( + (options[key].choice.donor_head + 1) * variable + for key, variable in variables.items() + if not options[key].choice.fallback + ) + donor_base = dataset.global_num_heads * len(dataset.targets) + 1 + if not minimize(reuse_count * donor_base + donor_signature): + raise MaskReuseCalibrationError( + "compact selector lost deterministic tie-break feasibility" + ) + + choices: dict[ConsumerHead, _Choice] = {} + prompt_totals = dict.fromkeys(prompts, 0.0) + worst_individual = 0.0 + skipped_tiles = 0 + retained_tiles = 0 + for key, variable in variables.items(): + if variable.value() <= 0.5: + continue + layer, head, _ = key + option = options[key] + choices[(layer, head)] = option.choice + skipped_tiles += option.bmm1_skipped_tiles + retained_tiles += option.choice.retained_tiles + if not option.choice.fallback: + for prompt, risk in zip(prompts, option.risk_by_prompt, strict=True): + prompt_totals[prompt] += risk + worst_individual = max(worst_individual, risk) + prompt_risks = [prompt_totals[prompt] / normalizer for prompt in prompts] + eligible_tiles = ( + accumulator.eligible_sum[target] * dataset.global_num_heads * len(dataset.nearest) + ) + return _Selection( + target, + choices, + (), + bmm1_eligible_tiles=eligible_tiles, + bmm1_skipped_tiles=skipped_tiles, + target_bmm1_skip_ratio_met=target_met, + worst_prompt_reuse_dropped_mass=max(prompt_risks, default=0.0), + mean_prompt_reuse_dropped_mass=( + sum(prompt_risks) / len(prompt_risks) if prompt_risks else 0.0 + ), + worst_individual_reuse_dropped_mass=worst_individual, + ) + + selections: dict[Bucket, _Selection] = {} + for bucket in dataset.buckets: + accumulator = accumulators[bucket] + frontier: list[dict[str, object]] = [] + candidates: list[tuple[tuple[object, ...], _Selection]] = [] + maximum_candidates: list[tuple[tuple[object, ...], _Selection]] = [] + for target in dataset.menus[bucket]: + anchor_evaluation = accumulator.anchor_evaluation[target] + eligible_tiles = ( + accumulator.eligible_sum[target] * dataset.global_num_heads * len(dataset.nearest) + ) + required_tiles = math.ceil(target_bmm1_skip_ratio * eligible_tiles) + selected = None + if anchor_evaluation.violations == 0: + selected = solve_target( + bucket, + target, + minimum_bmm1_skipped_tiles=required_tiles, + maximize_bmm1_skipped_tiles=False, + target_met=True, + ) + fallback_count = ( + None + if selected is None + else sum(choice.fallback for choice in selected.choices.values()) + ) + retained_reuse = ( + None + if selected is None + else sum(choice.retained_tiles for choice in selected.choices.values()) + ) + combined_tile_cost = ( + None + if retained_reuse is None + else 2 * retained_reuse + anchor_evaluation.retained_tiles + ) + frontier.append( + { + "target_sparsity": target, + "anchor_safe": anchor_evaluation.violations == 0, + "target_bmm1_skip_ratio": target_bmm1_skip_ratio, + "target_bmm1_skip_ratio_feasible": selected is not None, + "retained_reuse_tiles": retained_reuse, + "retained_anchor_tiles": anchor_evaluation.retained_tiles, + "combined_tile_cost": combined_tile_cost, + "fallback_head_count": fallback_count, + "anchor_calibration": anchor_evaluation.to_mapping(), + } + ) + if selected is not None: + rank = ( + selected.worst_prompt_reuse_dropped_mass, + combined_tile_cost, + target, + ) + candidates.append((rank, selected)) + frontier[-1].update( + { + "bmm1_skipped_tiles": selected.bmm1_skipped_tiles, + "achieved_bmm1_skip_ratio": ( + selected.bmm1_skipped_tiles / selected.bmm1_eligible_tiles + ), + "worst_prompt_model_wide_reuse_dropped_mass": ( + selected.worst_prompt_reuse_dropped_mass + ), + } + ) + elif anchor_evaluation.violations == 0: + maximum = solve_target( + bucket, + target, + minimum_bmm1_skipped_tiles=None, + maximize_bmm1_skipped_tiles=True, + target_met=False, + ) + if maximum is not None: + maximum_candidates.append( + ( + ( + -maximum.bmm1_skipped_tiles, + maximum.worst_prompt_reuse_dropped_mass, + target, + ), + maximum, + ) + ) + frontier[-1].update( + { + "maximum_feasible_bmm1_skipped_tiles": (maximum.bmm1_skipped_tiles), + "maximum_feasible_bmm1_skip_ratio": ( + maximum.bmm1_skipped_tiles / maximum.bmm1_eligible_tiles + ), + } + ) + if candidates: + _, selected = min(candidates, key=lambda item: item[0]) + selections[bucket] = _Selection( + selected.target_sparsity, + selected.choices, + tuple(frontier), + bmm1_eligible_tiles=selected.bmm1_eligible_tiles, + bmm1_skipped_tiles=selected.bmm1_skipped_tiles, + target_bmm1_skip_ratio_met=True, + worst_prompt_reuse_dropped_mass=selected.worst_prompt_reuse_dropped_mass, + mean_prompt_reuse_dropped_mass=selected.mean_prompt_reuse_dropped_mass, + worst_individual_reuse_dropped_mass=(selected.worst_individual_reuse_dropped_mass), + ) + elif maximum_candidates: + _, selected = min(maximum_candidates, key=lambda item: item[0]) + selections[bucket] = _Selection( + selected.target_sparsity, + selected.choices, + tuple(frontier), + bmm1_eligible_tiles=selected.bmm1_eligible_tiles, + bmm1_skipped_tiles=selected.bmm1_skipped_tiles, + target_bmm1_skip_ratio_met=False, + worst_prompt_reuse_dropped_mass=selected.worst_prompt_reuse_dropped_mass, + mean_prompt_reuse_dropped_mass=selected.mean_prompt_reuse_dropped_mass, + worst_individual_reuse_dropped_mass=(selected.worst_individual_reuse_dropped_mass), + ) + else: + target = dataset.menus[bucket][0] + dense_choices = { + item: _Choice(0, True, accumulators[bucket].eligible_sum[target]) + for item in dataset.targets + } + selections[bucket] = _Selection( + None, + dense_choices, + tuple(frontier), + "no_target_sparsity_satisfied_anchor_calibration_constraint", + bmm1_eligible_tiles=( + accumulator.eligible_sum[target] + * dataset.global_num_heads + * len(dataset.nearest) + ), + ) + return selections + + +def _evaluation_pass( + source: CompactMaskReuseCaptureSource, + dataset: _Dataset, + selections: Mapping[Bucket, _Selection], + *, + max_anchor_dropped_mass: float, + reuse_dropped_mass_report_threshold: float, +) -> tuple[ + dict[tuple[Bucket, str], _ReuseEvaluation], + dict[tuple[Bucket, str], _AnchorEvaluation], +]: + reuse = {(bucket, split): _ReuseEvaluation() for bucket in dataset.buckets for split in _SPLITS} + anchor = { + (bucket, split): _AnchorEvaluation() for bucket in dataset.buckets for split in _SPLITS + } + for capture in source: + selection = selections[capture.bucket] + evaluation_target = ( + dataset.menus[capture.bucket][0] + if selection.target_sparsity is None + else selection.target_sparsity + ) + if capture.target_sparsity != evaluation_target: + continue + reuse_result = reuse[(capture.bucket, capture.split)] + for layer, stats in capture.consumer_layers.items(): + anchor_stats = capture.anchor_stats_by_layer[stats.anchor_layer] + for head in range(dataset.global_num_heads): + choice = selection.choices[(layer, head)] + reuse_result.eligible_tiles += capture.eligible_tiles + if choice.fallback: + reuse_result.retained_tiles += capture.eligible_tiles + continue + donor = choice.donor_head + dropped = float(stats.dropped_mass[head][donor]) + reuse_result.retained_tiles += anchor_stats.retained_tiles[donor] + reuse_result.sparse_observations += 1 + reuse_result.dropped_mass_sum += dropped + reuse_result.worst_dropped_mass = max(reuse_result.worst_dropped_mass, dropped) + reuse_result.violations += int(dropped > reuse_dropped_mass_report_threshold) + + anchor_result = anchor[(capture.bucket, capture.split)] + dropped_values: list[float] = [] + retained_tiles = 0 + eligible_tiles = 0 + for anchor_stats in capture.anchor_stats_by_layer.values(): + eligible_tiles += capture.eligible_tiles * dataset.global_num_heads + if selection.target_sparsity is None: + retained_tiles += capture.eligible_tiles * dataset.global_num_heads + dropped_values.extend([0.0] * dataset.global_num_heads) + else: + retained_tiles += sum(anchor_stats.retained_tiles) + dropped_values.extend(anchor_stats.dropped_mass) + prompt_mean = sum(dropped_values) / len(dropped_values) + anchor_result.eligible_tiles += eligible_tiles + anchor_result.retained_tiles += retained_tiles + anchor_result.prompt_count += 1 + anchor_result.prompt_mean_sum += prompt_mean + anchor_result.worst_prompt_mean = max(anchor_result.worst_prompt_mean, prompt_mean) + anchor_result.violations += int(prompt_mean > max_anchor_dropped_mass) + for bucket in dataset.buckets: + for split in _SPLITS: + expected = len(dataset.prompts[(bucket, split)]) + if anchor[(bucket, split)].prompt_count != expected: + raise MaskReuseCalibrationError( + f"evaluation pass did not cover every {split} prompt in bucket {bucket}" + ) + return reuse, anchor + + +def calibrate_compact_mask_reuse_policy( + captures: CompactMaskReuseCaptureSource | str | Path, + *, + vanilla_calibration: Mapping[str, object], + topology: Mapping[str, object], + checkpoint_manifest: VerifiedCheckpointManifest, + evidence: Mapping[str, object], + max_anchor_dropped_mass: float, + reuse_dropped_mass_report_threshold: float, + target_bmm1_skip_ratio: float, + source_provenance: Mapping[str, object] | None = None, +) -> dict[str, object]: + """Select a minimum-risk schema-v3 candidate under a per-bucket BMM1 target.""" + if not isinstance(checkpoint_manifest, VerifiedCheckpointManifest): + raise MaskReuseCalibrationError( + "checkpoint_manifest must be returned by verify_checkpoint_manifest" + ) + source = ( + captures + if isinstance(captures, CompactMaskReuseCaptureSource) + else load_compact_mask_reuse_captures(captures) + ) + threshold_scale_factor = canonical_prefill_threshold_scale_factor(vanilla_calibration) + params = threshold_scale_factor["prefill"] + assert isinstance(params, Mapping) + if not {"min_observed_sparsity", "max_observed_sparsity"} <= params.keys(): + raise MaskReuseCalibrationError("compact schema-v3 export requires vanilla fit bounds") + anchors, nearest = _normalize_topology(topology) + dataset = _validate_dataset( + source, + fit=threshold_scale_factor, + anchors=anchors, + nearest=nearest, + ) + checkpoint_identity = checkpoint_manifest.sha256 + if dataset.checkpoint_manifest_sha256 != checkpoint_identity: + raise MaskReuseCalibrationError( + "compact captures do not match the verified checkpoint manifest" + ) + if dataset.model != checkpoint_manifest.model: + raise MaskReuseCalibrationError( + "compact capture model does not match the verified checkpoint manifest" + ) + canonical_evidence = _canonical_evidence(evidence) + input_sha256 = source.sha256() + if canonical_evidence["reuse_bundle_sha256"] != input_sha256: + raise MaskReuseCalibrationError( + "evidence.reuse_bundle_sha256 does not match the compact capture file" + ) + anchor_bound = _number( + max_anchor_dropped_mass, "max_anchor_dropped_mass", minimum=0.0, maximum=1.0 + ) + reuse_report_threshold = _number( + reuse_dropped_mass_report_threshold, + "reuse_dropped_mass_report_threshold", + minimum=0.0, + maximum=1.0, + ) + bmm1_target = _number( + target_bmm1_skip_ratio, + "target_bmm1_skip_ratio", + minimum=0.0, + maximum=1.0, + ) + selections = _selection_pass( + source, + dataset, + max_anchor_dropped_mass=anchor_bound, + target_bmm1_skip_ratio=bmm1_target, + ) + reuse_evaluations, anchor_evaluations = _evaluation_pass( + source, + dataset, + selections, + max_anchor_dropped_mass=anchor_bound, + reuse_dropped_mass_report_threshold=reuse_report_threshold, + ) + if source.sha256() != input_sha256: + raise MaskReuseCalibrationError( + "compact capture file changed during calibration; discard this result" + ) + + context_policies = [] + bucket_reports = [] + target_menus = [] + overall_reuse_calibration = _ReuseEvaluation() + overall_reuse_heldout = _ReuseEvaluation() + overall_anchor_calibration = _AnchorEvaluation() + overall_anchor_heldout = _AnchorEvaluation() + total_fallback = 0 + total_bmm1_eligible = 0 + total_bmm1_skipped = 0 + for bucket in dataset.buckets: + selection = selections[bucket] + if selection.target_sparsity is not None and bucket[1] is None: + raise MaskReuseCalibrationError( + "a deployment-qualified sparse context bucket requires a finite maximum" + ) + reuse_calibration = reuse_evaluations[(bucket, "calibration")] + reuse_heldout = reuse_evaluations[(bucket, "heldout")] + anchor_calibration = anchor_evaluations[(bucket, "calibration")] + anchor_heldout = anchor_evaluations[(bucket, "heldout")] + overall_reuse_calibration.add(reuse_calibration) + overall_reuse_heldout.add(reuse_heldout) + overall_anchor_calibration.add(anchor_calibration) + overall_anchor_heldout.add(anchor_heldout) + policy: dict[str, object] + if selection.target_sparsity is None: + policy = {"min_kv_tokens": bucket[0], "max_kv_tokens": bucket[1], "exact": True} + headmaps: dict[str, list[int]] = {} + fallback_heads: dict[str, list[int]] = {} + else: + headmaps = { + str(layer): [ + selection.choices[(layer, head)].donor_head + for head in range(dataset.global_num_heads) + ] + for layer in dataset.consumer_layers + } + fallback_heads = { + str(layer): [ + head + for head in range(dataset.global_num_heads) + if selection.choices[(layer, head)].fallback + ] + for layer in dataset.consumer_layers + } + policy = { + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "target_sparsity": selection.target_sparsity, + "headmaps": headmaps, + "fallback_heads": fallback_heads, + } + context_policies.append(policy) + fallback_count = sum(len(heads) for heads in fallback_heads.values()) + total_fallback += fallback_count + total_bmm1_eligible += selection.bmm1_eligible_tiles + total_bmm1_skipped += selection.bmm1_skipped_tiles + target_menus.append( + { + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "target_sparsities": [row["target_sparsity"] for row in selection.frontier], + } + ) + bucket_reports.append( + { + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "selected_target_sparsity": selection.target_sparsity, + "selection_status": ( + "target_bmm1_skip_ratio_met" + if selection.target_sparsity is not None + and selection.target_bmm1_skip_ratio_met + else ( + "target_bmm1_skip_ratio_unmet_maximum_feasible" + if selection.target_sparsity is not None + else selection.exact_reason + ) + ), + "target_sparsity_frontier": list(selection.frontier), + "fallback_head_count": fallback_count, + "bmm1_skip_objective": { + "target": bmm1_target, + "target_met": selection.target_bmm1_skip_ratio_met, + "eligible_tiles": selection.bmm1_eligible_tiles, + "skipped_tiles": selection.bmm1_skipped_tiles, + "achieved": ( + selection.bmm1_skipped_tiles / selection.bmm1_eligible_tiles + if selection.bmm1_eligible_tiles + else 0.0 + ), + }, + "reuse_selection_objective": { + "hard_maximum": None, + "worst_development_prompt_model_wide_dropped_mass": ( + selection.worst_prompt_reuse_dropped_mass + ), + "mean_development_prompt_model_wide_dropped_mass": ( + selection.mean_prompt_reuse_dropped_mass + ), + "worst_individual_dropped_mass": ( + selection.worst_individual_reuse_dropped_mass + ), + }, + "reuse_calibration": reuse_calibration.to_mapping(), + "reuse_heldout": reuse_heldout.to_mapping(), + "anchor_calibration": anchor_calibration.to_mapping( + exact=selection.target_sparsity is None + ), + "anchor_heldout": anchor_heldout.to_mapping( + exact=selection.target_sparsity is None + ), + } + ) + + candidate_cell_count = ( + dataset.capture_count + * len(dataset.consumer_layers) + * dataset.global_num_heads + * dataset.global_num_heads + ) + constraints = { + "anchor": { + "metric": "worst_prompt_mean_anchor_dropped_mass", + "comparison": "<=", + "maximum": anchor_bound, + }, + "reuse": { + "selection_metric": ( + "per_prompt_mean_across_all_attention_layers_and_heads_reuse_dropped_mass" + ), + "selection_hard_maximum": None, + "report_metric": "per_prompt_candidate_reuse_dropped_mass", + "report_threshold": reuse_report_threshold, + "report_threshold_affects_selection": False, + }, + "bmm1_skip_ratio": { + "metric": "per_context_bucket_model_wide_eligible_bmm1_tile_skip_ratio", + "comparison": ">=", + "target": bmm1_target, + }, + } + provenance: dict[str, object] = { + "calibrator": "modelopt.mask_reuse.compact_streaming", + "compact_capture_schema_version": 1, + "input_capture_count": dataset.capture_count, + "candidate_cell_count": candidate_cell_count, + "canonical_input_sha256": input_sha256, + "calibration_prompt_ids": list(dataset.calibration_prompt_ids), + "heldout_prompt_ids": list(dataset.heldout_prompt_ids), + "prompt_sources": list(dataset.prompt_sources), + "development_source_group_sha256": sorted( + { + str(row["source_group_sha256"]) + for row in dataset.prompt_sources + if row["partition"] == "development" + } + ), + "outer_test_source_group_sha256": sorted( + { + str(row["source_group_sha256"]) + for row in dataset.prompt_sources + if row["partition"] == "outer_test" + } + ), + "selection_split": "calibration", + "evaluation_split": "heldout", + "streaming_passes": ["validation", "calibration_selection", "frozen_evaluation"], + "threshold_semantics": "a * exp(b * target_sparsity) / sample_length", + "threshold_implementation": { + "threshold_log2": "log2(a) + b * target_sparsity * log2(e) - log2(sample_length)", + "threshold_lambda": "exp2(threshold_log2)", + "validation": "exact IEEE-754 binary64 hex equality", + }, + "checkpoint_manifest_sha256": checkpoint_identity, + "target_sparsity_menus": target_menus, + "constraints": constraints, + "tie_breaks": [ + "reject target sparsities exceeding the calibration anchor bound", + "meet the requested BMM1 skip ratio in every context bucket", + "minimum worst-prompt model-wide reuse dropped mass", + "minimum equal-BMM combined tile cost 2*A_R + A_A", + "fewest reused consumer heads", + "smallest target sparsity", + "lexicographically smallest donor-head map", + ], + "selection_cost": { + "formula": "2 * retained_reuse_tiles + retained_anchor_tiles", + "bmm1_weight": 1.0, + "bmm2_weight": 1.0, + }, + } + if source_provenance is not None: + provenance["source"] = dict(source_provenance) + denominator = len(dataset.targets) * len(dataset.buckets) + report = { + "model": dataset.model, + "checkpoint_manifest_sha256": checkpoint_identity, + "constraints": constraints, + "selection_unit": "context_bucket", + "by_bucket": bucket_reports, + "overall": { + "consumer_head_bucket_count": denominator, + "fallback_head_bucket_count": total_fallback, + "fallback_fraction": total_fallback / denominator, + "bmm1_skip_objective": { + "target_per_context_bucket": bmm1_target, + "all_context_buckets_met": all( + selection.target_bmm1_skip_ratio_met for selection in selections.values() + ), + "eligible_tiles": total_bmm1_eligible, + "skipped_tiles": total_bmm1_skipped, + "achieved_model_wide": ( + total_bmm1_skipped / total_bmm1_eligible if total_bmm1_eligible else 0.0 + ), + }, + "reuse_calibration": overall_reuse_calibration.to_mapping(), + "reuse_heldout": overall_reuse_heldout.to_mapping(), + "anchor_calibration": overall_anchor_calibration.to_mapping(), + "anchor_heldout": overall_anchor_heldout.to_mapping(), + }, + "promotion": { + "status": "candidate_only", + "eligible": False, + "reasons": [ + "grouped inner-fold modal and safety stability not evaluated", + "preregistered outer gate with at least 99 independent groups per family cell not evaluated", + "deployment rectangular-geometry promotion gate not evaluated", + ], + }, + } + return { + "version": 3, + "promotion_status": "candidate_only", + "phase": "prefill", + "decode": {"mode": "dense"}, + "calibration_protocol": _CALIBRATION_PROTOCOL, + "producer": {"name": "modelopt", "version": modelopt.__version__}, + "evidence": canonical_evidence, + "threshold_scale_factor": threshold_scale_factor, + "model": dataset.model, + "checkpoint_manifest_sha256": checkpoint_identity, + "global_num_heads": dataset.global_num_heads, + "target_bmm1_skip_ratio": bmm1_target, + "anchors": list(dataset.anchors), + "nearest": {str(layer): anchor for layer, anchor in dataset.nearest.items()}, + "deployment_geometry_validated": False, + "deployment_geometry": dataset.deployment_geometry, + "context_policies": context_policies, + "provenance": provenance, + "calibration_report": report, + } diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_calibrate_mask_reuse_cli.py b/tests/unit/torch/sparsity/attention_sparsity/test_calibrate_mask_reuse_cli.py new file mode 100644 index 00000000000..2c5244b256c --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_calibrate_mask_reuse_cli.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Focused tests for fail-closed mask-reuse candidate publication.""" + +import importlib.util +import json +import os +from hashlib import sha256 +from pathlib import Path + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + create_checkpoint_manifest, +) + +_SCRIPT_PATH = Path(__file__).parents[5] / "examples/vllm_serve/calibrate_mask_reuse.py" +_SPEC = importlib.util.spec_from_file_location("calibrate_mask_reuse_cli", _SCRIPT_PATH) +assert _SPEC is not None and _SPEC.loader is not None +calibrate_mask_reuse_cli = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(calibrate_mask_reuse_cli) + + +def _canonical(value): + return ( + json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + + +def _base_args(tmp_path: Path): + checkpoint_root = tmp_path / "checkpoint" + checkpoint_root.mkdir() + (checkpoint_root / "config.json").write_text("{}\n", encoding="utf-8") + (checkpoint_root / "model.safetensors").write_bytes(b"weights") + checkpoint = create_checkpoint_manifest(checkpoint_root, model="test-model") + + captures = tmp_path / "captures.jsonl" + captures.write_text("unused\n", encoding="utf-8") + vanilla = tmp_path / "vanilla.json" + vanilla.write_text('{"vanilla":true}\n', encoding="utf-8") + topology = tmp_path / "topology.json" + topology.write_text('{"anchors":[0],"nearest":{"0":0}}\n', encoding="utf-8") + artifacts = {} + for name in ("calibration_plan", "family_registry", "grouped_fit", "outer_report"): + path = tmp_path / f"{name}.json" + path.write_text(f'{{"artifact":"{name}"}}\n', encoding="utf-8") + artifacts[name] = path + capture_manifest = tmp_path / "capture-manifest.json" + capture_manifest.write_bytes( + _canonical( + { + "capture_manifest_schema_version": 4, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v4", + "model": "test-model", + "checkpoint_manifest_sha256": checkpoint.sha256, + "checkpoint_manifest_path": str(checkpoint.manifest_path), + "checkpoint_file_count": checkpoint.file_count, + "checkpoint_total_size_bytes": checkpoint.total_size_bytes, + "plan": "test_stride2", + "fa4_source": "/source", + "fa4_source_commit": "a" * 40, + "fa4_source_git_tree": "b" * 40, + "fa4_source_git_archive_sha256": "c" * 64, + "fa4_source_manifest_path": "/source-manifest.json", + "fa4_source_manifest_sha256": "d" * 64, + "fa4_source_directory_count": 4, + "fa4_source_file_count": 2, + "fa4_source_total_size_bytes": 42, + "engine_kwargs": {"tensor_parallel_size": 2}, + "dense_shadow_validation_requested": True, + "target_sparsity_hex": [(0.7).hex()], + "vanilla_threshold_scale_factor": {"formula": "unused"}, + "vanilla_fit_sha256": sha256(b"normalized-fit").hexdigest(), + "vanilla_config_file_sha256": sha256(vanilla.read_bytes()).hexdigest(), + "prompt_plan_file_sha256": sha256(b"prompts").hexdigest(), + "compact_capture_file_sha256": sha256(captures.read_bytes()).hexdigest(), + "capture_count": 1, + "candidate_cell_count": 4, + "captures": [{"candidate_cell_count": 4}], + } + ) + ) + policy = tmp_path / "candidate.json" + report = tmp_path / "report.json" + args = [ + "--checkpoint", + str(checkpoint_root), + "--compact-captures", + str(captures), + "--capture-manifest", + str(capture_manifest), + "--vanilla-config", + str(vanilla), + "--topology", + str(topology), + "--calibration-plan", + str(artifacts["calibration_plan"]), + "--family-registry", + str(artifacts["family_registry"]), + "--grouped-fit", + str(artifacts["grouped_fit"]), + "--outer-report", + str(artifacts["outer_report"]), + "--max-anchor-dropped-mass", + "0.02", + "--reuse-dropped-mass-report-threshold", + "0.03", + "--target-bmm1-skip-ratio", + "0.10", + "--output-policy", + str(policy), + "--output-report", + str(report), + ] + return args, policy, report, captures, vanilla, artifacts + + +def test_main_verifies_artifacts_and_atomically_writes_candidate(tmp_path, monkeypatch, capsys): + args, policy_path, report_path, captures, vanilla, artifacts = _base_args(tmp_path) + source = object() + monkeypatch.setattr( + calibrate_mask_reuse_cli, "load_compact_mask_reuse_captures", lambda path: source + ) + captured = {} + artifact = { + "version": 3, + "promotion_status": "candidate_only", + "deployment_geometry_validated": False, + "provenance": {"input_capture_count": 1, "candidate_cell_count": 4}, + "calibration_report": {"promotion": {"eligible": False}}, + } + + def fake_calibrate(compact_source, **kwargs): + captured["source"] = compact_source + captured.update(kwargs) + return artifact + + monkeypatch.setattr( + calibrate_mask_reuse_cli, "calibrate_compact_mask_reuse_policy", fake_calibrate + ) + + assert calibrate_mask_reuse_cli.main(args) == 0 + + expected_policy = _canonical(artifact) + assert policy_path.read_bytes() == expected_policy + assert report_path.read_bytes() == _canonical(artifact["calibration_report"]) + assert captured["source"] is source + assert captured["evidence"] == { + "calibration_plan_sha256": sha256(artifacts["calibration_plan"].read_bytes()).hexdigest(), + "family_registry_sha256": sha256(artifacts["family_registry"].read_bytes()).hexdigest(), + "grouped_fit_sha256": sha256(artifacts["grouped_fit"].read_bytes()).hexdigest(), + "outer_report_sha256": sha256(artifacts["outer_report"].read_bytes()).hexdigest(), + "vanilla_fit_sha256": sha256(vanilla.read_bytes()).hexdigest(), + "reuse_bundle_sha256": sha256(captures.read_bytes()).hexdigest(), + } + assert captured["reuse_dropped_mass_report_threshold"] == 0.03 + assert captured["target_bmm1_skip_ratio"] == 0.10 + assert "MASK_REUSE_FA4_CANDIDATE_SHA256=" in capsys.readouterr().out + + +def test_vanilla_mutation_after_semantic_snapshot_cannot_publish(tmp_path, monkeypatch, capsys): + args, policy_path, report_path, _, vanilla, _ = _base_args(tmp_path) + original_payload = vanilla.read_bytes() + real_evidence_artifacts = calibrate_mask_reuse_cli._evidence_artifacts + + def mutate_after_snapshot(namespace, *, vanilla_fit_sha256): + vanilla.write_text('{"vanilla":false}\n', encoding="utf-8") + return real_evidence_artifacts(namespace, vanilla_fit_sha256=vanilla_fit_sha256) + + monkeypatch.setattr(calibrate_mask_reuse_cli, "_evidence_artifacts", mutate_after_snapshot) + monkeypatch.setattr( + calibrate_mask_reuse_cli, + "load_compact_mask_reuse_captures", + lambda path: object(), + ) + + def fake_calibrate(compact_source, **kwargs): + assert kwargs["vanilla_calibration"] == {"vanilla": True} + assert kwargs["evidence"]["vanilla_fit_sha256"] == sha256(original_payload).hexdigest() + return { + "promotion_status": "candidate_only", + "deployment_geometry_validated": False, + "provenance": {"input_capture_count": 1, "candidate_cell_count": 4}, + "calibration_report": {"promotion": {"eligible": False}}, + } + + monkeypatch.setattr( + calibrate_mask_reuse_cli, + "calibrate_compact_mask_reuse_policy", + fake_calibrate, + ) + + with pytest.raises(SystemExit, match="2"): + calibrate_mask_reuse_cli.main(args) + + assert "vanilla_fit_sha256 artifact changed during calibration" in capsys.readouterr().err + assert not policy_path.exists() + assert not report_path.exists() + + +def test_main_rejects_capture_manifest_hash_mismatch(tmp_path, capsys): + args, _, _, captures, _, _ = _base_args(tmp_path) + captures.write_text("changed\n", encoding="utf-8") + + with pytest.raises(SystemExit, match="2"): + calibrate_mask_reuse_cli.main(args) + + assert "compact_capture_file_sha256" in capsys.readouterr().err + + +def test_load_json_object_rejects_duplicate_keys(tmp_path): + path = tmp_path / "duplicate.json" + path.write_text('{"anchors":[0],"anchors":[1]}', encoding="utf-8") + + with pytest.raises(ValueError, match="duplicate JSON key 'anchors'"): + calibrate_mask_reuse_cli._load_json_object(path, label="topology") + + +def test_candidate_publication_rolls_back_report_when_policy_race_wins(tmp_path, monkeypatch): + policy = tmp_path / "candidate.json" + report = tmp_path / "report.json" + real_link = os.link + call_count = 0 + + def racing_link(source, target, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 2: + Path(target).write_bytes(b"racer") + return real_link(source, target, **kwargs) + + monkeypatch.setattr(calibrate_mask_reuse_cli.os, "link", racing_link) + + with pytest.raises(FileExistsError): + calibrate_mask_reuse_cli._publish_candidate_outputs(policy, b"ours", report, b"report") + + assert policy.read_bytes() == b"racer" + assert not report.exists() + + +def test_candidate_publication_refuses_existing_output(tmp_path): + policy = tmp_path / "candidate.json" + report = tmp_path / "report.json" + policy.write_bytes(b"existing") + + with pytest.raises(FileExistsError, match="refusing to overwrite"): + calibrate_mask_reuse_cli._publish_candidate_outputs(policy, b"ours", report, b"report") + + assert policy.read_bytes() == b"existing" + assert not report.exists() diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py b/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py new file mode 100644 index 00000000000..8e90f4feb97 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for content-addressed checkpoint identity.""" + +from hashlib import sha256 +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration import checkpoint_manifest +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + CHECKPOINT_MANIFEST_NAME, + CheckpointManifestError, + create_checkpoint_manifest, + read_stable_file_snapshot, + verify_checkpoint_manifest, +) + + +def _toy_checkpoint(path: Path) -> Path: + path.mkdir() + (path / "config.json").write_text('{"model_type":"toy"}\n', encoding="utf-8") + (path / "model.safetensors").write_bytes(b"toy-weights") + (path / "tokenizer.json").write_text("{}\n", encoding="utf-8") + return path + + +def test_manifest_generation_is_deterministic_verified_and_no_clobber(tmp_path): + first_root = _toy_checkpoint(tmp_path / "first") + second_root = _toy_checkpoint(tmp_path / "second") + + first = create_checkpoint_manifest(first_root, model="toy-model") + second = create_checkpoint_manifest(second_root, model="toy-model") + + assert first.sha256 == second.sha256 + assert first.manifest_path.read_bytes() == second.manifest_path.read_bytes() + assert verify_checkpoint_manifest(first_root, expected_model="toy-model") == first + with pytest.raises(CheckpointManifestError, match="refusing to overwrite"): + create_checkpoint_manifest(first_root, model="toy-model") + + +def test_verifier_rejects_content_mutation_and_symlinks(tmp_path): + root = _toy_checkpoint(tmp_path / "checkpoint") + create_checkpoint_manifest(root, model="toy-model") + (root / "model.safetensors").write_bytes(b"mutated") + with pytest.raises(CheckpointManifestError, match="does not match"): + verify_checkpoint_manifest(root) + + other = _toy_checkpoint(tmp_path / "symlinked") + (other / "alias.bin").symlink_to(other / "model.safetensors") + with pytest.raises(CheckpointManifestError, match="forbidden symlink"): + create_checkpoint_manifest(other, model="toy-model") + + +def test_verifier_rejects_symlink_manifest(tmp_path): + root = _toy_checkpoint(tmp_path / "checkpoint") + target = tmp_path / "outside.json" + target.write_text("{}\n", encoding="utf-8") + (root / CHECKPOINT_MANIFEST_NAME).symlink_to(target) + + with pytest.raises(CheckpointManifestError, match="without following symlinks"): + verify_checkpoint_manifest(root) + + +def test_portable_open_fallback_still_rejects_symlink(tmp_path, monkeypatch): + target = tmp_path / "target.json" + target.write_text("{}\n", encoding="utf-8") + alias = tmp_path / "alias.json" + alias.symlink_to(target) + monkeypatch.delattr(checkpoint_manifest.os, "O_NOFOLLOW", raising=False) + + with pytest.raises(CheckpointManifestError, match="without following symlinks"): + read_stable_file_snapshot(alias, label="fallback input") + + +def test_portable_open_fallback_preserves_exact_binary_bytes(tmp_path, monkeypatch): + path = tmp_path / "binary.dat" + payload = b"line-1\r\nline-2\x00\n" + path.write_bytes(payload) + monkeypatch.delattr(checkpoint_manifest.os, "O_NOFOLLOW", raising=False) + + snapshot = read_stable_file_snapshot(path, label="binary input") + + assert snapshot.payload == payload + assert snapshot.sha256 == sha256(payload).hexdigest() + + +def test_portable_open_uses_binary_flag_when_available(tmp_path, monkeypatch): + path = tmp_path / "binary.dat" + path.write_bytes(b"payload") + binary_flag = 1 << 29 + observed_flags = [] + real_open = checkpoint_manifest.os.open + monkeypatch.setattr(checkpoint_manifest.os, "O_BINARY", binary_flag, raising=False) + + def recording_open(source, flags, *args, **kwargs): + observed_flags.append(flags) + return real_open(source, flags & ~binary_flag, *args, **kwargs) + + monkeypatch.setattr(checkpoint_manifest.os, "open", recording_open) + + snapshot = read_stable_file_snapshot(path, label="binary input") + + assert snapshot.payload == b"payload" + assert observed_flags and observed_flags[0] & binary_flag + + +def test_portable_open_rejects_path_swap_before_read(tmp_path, monkeypatch): + path = tmp_path / "input.dat" + replacement = tmp_path / "replacement.dat" + path.write_bytes(b"expected") + replacement.write_bytes(b"attacker") + real_open = checkpoint_manifest.os.open + swapped = False + monkeypatch.delattr(checkpoint_manifest.os, "O_NOFOLLOW", raising=False) + + def swapping_open(source, flags, *args, **kwargs): + nonlocal swapped + if Path(source) == path and not swapped: + swapped = True + path.unlink() + replacement.rename(path) + return real_open(source, flags, *args, **kwargs) + + monkeypatch.setattr(checkpoint_manifest.os, "open", swapping_open) + + with pytest.raises(CheckpointManifestError, match="stable regular file"): + read_stable_file_snapshot(path, label="swapped input") + + +def test_windows_reparse_attribute_is_link_like(): + observed = SimpleNamespace( + st_mode=0, + st_file_attributes=checkpoint_manifest._FILE_ATTRIBUTE_REPARSE_POINT, + ) + + assert checkpoint_manifest._is_link_like(observed) + + +def test_manifest_creation_without_directory_fsync_support(tmp_path, monkeypatch): + root = _toy_checkpoint(tmp_path / "checkpoint") + monkeypatch.delattr(checkpoint_manifest.os, "O_DIRECTORY", raising=False) + + created = create_checkpoint_manifest(root, model="toy-model") + + assert verify_checkpoint_manifest(root, expected_model="toy-model") == created + + +def test_manifest_publication_preserves_destination_created_by_racer(tmp_path, monkeypatch): + root = _toy_checkpoint(tmp_path / "checkpoint") + manifest = root / CHECKPOINT_MANIFEST_NAME + real_link = checkpoint_manifest.os.link + + def racing_link(source, target, **kwargs): + Path(target).write_text("racer\n", encoding="utf-8") + return real_link(source, target, **kwargs) + + monkeypatch.setattr(checkpoint_manifest.os, "link", racing_link) + + with pytest.raises(CheckpointManifestError, match="appeared during publication"): + create_checkpoint_manifest(root, model="toy-model") + + assert manifest.read_text(encoding="utf-8") == "racer\n" diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py new file mode 100644 index 00000000000..8f2b95623a5 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py @@ -0,0 +1,322 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Deterministic tests for ModelOpt-owned mask-reuse calibration.""" + +import json +import math +from dataclasses import replace +from hashlib import sha256 +from pathlib import Path + +import pytest + +import modelopt +from modelopt.torch.sparsity.attention_sparsity.calibration import ( + AnchorLayerStats, + MaskReuseCalibrationError, + MaskReuseObservation, + calibrate_mask_reuse_policy, + canonical_prefill_threshold_scale_factor, + parse_mask_reuse_observations, +) +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + VerifiedCheckpointManifest, +) + +A = 14.47 +B = 10.91 +VANILLA_FIT = { + "prefill": { + "a": A, + "b": B, + "min_observed_sparsity": 0.4, + "max_observed_sparsity": 0.8, + } +} +TOPOLOGY = {"anchors": [0, 2], "nearest": {"0": 0, "1": 0, "2": 2}} +CHECKPOINT = sha256(b"checkpoint").hexdigest() +VERIFIED_CHECKPOINT = VerifiedCheckpointManifest( + checkpoint_root=Path("/verified-checkpoint"), + manifest_path=Path("/verified-checkpoint/checkpoint_manifest.json"), + model="toy", + sha256=CHECKPOINT, + file_count=2, + total_size_bytes=1, +) +EVIDENCE = { + field: sha256(field.encode()).hexdigest() + for field in ( + "calibration_plan_sha256", + "family_registry_sha256", + "vanilla_fit_sha256", + "reuse_bundle_sha256", + "grouped_fit_sha256", + "outer_report_sha256", + ) +} + + +def _source(prompt: str) -> str: + return sha256(prompt.encode()).hexdigest() + + +def _thresholds(target_sparsity: float, sample_length: int) -> tuple[float, float]: + threshold_log2 = ( + math.log2(A) + B * target_sparsity * math.log2(math.e) - math.log2(sample_length) + ) + return 2.0**threshold_log2, threshold_log2 + + +def _observations() -> list[MaskReuseObservation]: + rows = [] + prompts = { + "calibration": (("cal-0", 65_536), ("cal-1", 98_304)), + "heldout": (("held-0", 65_536), ("held-1", 98_304)), + } + for split, samples in prompts.items(): + for prompt, sample_length in samples: + for target_sparsity in (0.5, 0.7): + threshold, threshold_log2 = _thresholds(target_sparsity, sample_length) + anchor_retained = { + 0.5: (80, 90), + 0.7: (40, 20), + }[target_sparsity] + anchor_dropped = (0.03, 0.03) if target_sparsity == 0.5 else (0.07, 0.08) + anchor_stats = { + 0: AnchorLayerStats(anchor_retained, anchor_dropped), + 2: AnchorLayerStats( + (60, 60) if target_sparsity == 0.5 else (30, 30), + (0.02, 0.02), + ), + } + for consumer_head in range(2): + for donor_head in range(2): + retained = anchor_retained[donor_head] + if target_sparsity == 0.7 and consumer_head == 1: + dropped_mass = 0.07 + 0.01 * donor_head + else: + dropped_mass = 0.02 + 0.01 * donor_head + rows.append( + MaskReuseObservation( + model="toy", + min_kv_tokens=65_536, + max_kv_tokens=131_072, + target_sparsity=target_sparsity, + sample_length=sample_length, + threshold_lambda=threshold, + threshold_log2=threshold_log2, + q_tokens=8192, + kv_tokens=sample_length, + q_start_tokens=sample_length - 8192, + split=split, + prompt_id=prompt, + source_capture_sha256=_source(prompt), + anchor_layer=0, + consumer_layer=1, + consumer_head=consumer_head, + donor_head=donor_head, + retained_tiles=retained, + eligible_tiles=100, + anchor_dropped_mass=( + 0.03 if target_sparsity == 0.5 else 0.07 + 0.01 * donor_head + ), + anchor_stats_by_layer=anchor_stats, + dropped_mass=dropped_mass, + ) + ) + return rows + + +def _calibrate(rows): + return calibrate_mask_reuse_policy( + rows, + vanilla_calibration=VANILLA_FIT, + topology=TOPOLOGY, + checkpoint_manifest=VERIFIED_CHECKPOINT, + evidence=EVIDENCE, + max_anchor_dropped_mass=0.1, + reuse_dropped_mass_report_threshold=0.1, + target_bmm1_skip_ratio=0.1, + ) + + +def test_selects_target_sparsity_and_exports_backend_v3(): + artifact = _calibrate(_observations()) + + assert artifact["version"] == 3 + assert artifact["phase"] == "prefill" + assert artifact["decode"] == {"mode": "dense"} + assert artifact["calibration_protocol"] == "modelopt_mask_reuse_target_sparsity_v1" + assert artifact["producer"] == {"name": "modelopt", "version": modelopt.__version__} + assert artifact["evidence"] == EVIDENCE + assert artifact["threshold_scale_factor"] == { + "formula": "a * exp(b * target_sparsity)", + "prefill": VANILLA_FIT["prefill"], + } + assert artifact["context_policies"] == [ + { + "min_kv_tokens": 65_536, + "max_kv_tokens": 131_072, + "target_sparsity": 0.7, + "headmaps": {"1": [0, 0]}, + "fallback_heads": {"1": [1]}, + } + ] + assert artifact["promotion_status"] == "candidate_only" + assert artifact["deployment_geometry_validated"] is False + assert artifact["deployment_geometry"]["contract"]["kv_page_tokens"] == 16 + assert len(artifact["deployment_geometry"]["observations"]) == 4 + assert ( + artifact["calibration_report"]["overall"]["reuse_heldout"][ + "report_threshold_exceedance_rate" + ] + == 0.0 + ) + json.dumps(artifact) + + +def test_heldout_values_evaluate_but_cannot_change_selection(): + baseline = _calibrate(_observations()) + hostile_rows = [ + replace(row, dropped_mass=0.9) + if row.split == "heldout" + and row.target_sparsity == 0.7 + and row.consumer_head == 0 + and row.donor_head == 0 + else row + for row in _observations() + ] + + hostile = _calibrate(hostile_rows) + + assert hostile["context_policies"] == baseline["context_policies"] + assert ( + hostile["calibration_report"]["overall"]["reuse_heldout"][ + "report_threshold_exceedance_rate" + ] + == 1.0 + ) + assert hostile["calibration_report"]["overall"]["reuse_heldout"]["worst_dropped_mass"] == 0.9 + + +def test_rejects_relabelled_fixed_lambda_observation(): + rows = _observations() + rows[0] = replace( + rows[0], + threshold_lambda=math.nextafter(rows[0].threshold_lambda, math.inf), + ) + + with pytest.raises(MaskReuseCalibrationError, match="threshold_lambda does not match"): + _calibrate(rows) + + +def test_rejects_inexact_log2_launch_argument(): + rows = _observations() + rows[0] = replace( + rows[0], + threshold_log2=math.nextafter(rows[0].threshold_log2, math.inf), + ) + + with pytest.raises(MaskReuseCalibrationError, match="threshold_log2 does not match"): + _calibrate(rows) + + +def test_anchor_gate_includes_anchor_without_consumer_layer(): + rows = [ + replace( + row, + anchor_stats_by_layer={ + **row.anchor_stats_by_layer, + 2: AnchorLayerStats((30, 30), (0.2, 0.2)), + }, + ) + if row.target_sparsity == 0.7 + else row + for row in _observations() + ] + + artifact = _calibrate(rows) + + assert artifact["context_policies"][0]["target_sparsity"] == 0.5 + + +def test_rejects_inconsistent_repeated_anchor_payload(): + rows = _observations() + rows[0] = replace( + rows[0], + anchor_stats_by_layer={ + **rows[0].anchor_stats_by_layer, + 2: AnchorLayerStats((59, 60), (0.02, 0.02)), + }, + ) + + with pytest.raises(MaskReuseCalibrationError, match="anchor_stats_by_layer differs"): + _calibrate(rows) + + +def test_jsonl_rejects_duplicate_keys(): + with pytest.raises(MaskReuseCalibrationError, match="duplicate JSON key 'model'"): + parse_mask_reuse_observations(['{"model":"first","model":"second"}']) + + +def test_prefill_fit_rejects_b_above_backend_limit(): + invalid = {"prefill": {**VANILLA_FIT["prefill"], "b": 20.000_001}} + + with pytest.raises(MaskReuseCalibrationError, match=r"prefill\.b"): + canonical_prefill_threshold_scale_factor(invalid) + + +def test_canonicalizes_existing_sparse_attention_config(): + exported = { + "sparse_attention_config": { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "threshold_scale_factor": { + "formula": "a * exp(b * target_sparsity)", + "prefill": VANILLA_FIT["prefill"], + }, + } + } + } + } + + assert canonical_prefill_threshold_scale_factor(exported) == { + "formula": "a * exp(b * target_sparsity)", + "prefill": VANILLA_FIT["prefill"], + } + + +def test_canonicalizes_legacy_modelopt_serving_calibration(): + exported = { + "config_groups": { + "group_0": { + "sparse_algo": "softmax_skip", + "targets": ["Attention"], + } + }, + "threshold_scale_factor": { + "formula": "a * exp(b * target_sparsity)", + "prefill": {"a": 1.6771257955393728, "b": 8.894668875002724}, + "decode": {"a": 0.006180090552526715, "b": 10.23399476354776}, + }, + "target_sparse_ratio": {"prefill": 0.5, "decode": 0.5}, + } + + assert canonical_prefill_threshold_scale_factor(exported) == { + "formula": "a * exp(b * target_sparsity)", + "prefill": {"a": 1.6771257955393728, "b": 8.894668875002724}, + } diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py new file mode 100644 index 00000000000..67ab4214107 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py @@ -0,0 +1,408 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for streaming compact-capture policy selection.""" + +import inspect +import json +import math +from hashlib import sha256 + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration import mask_reuse_compact +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + create_checkpoint_manifest, +) +from modelopt.torch.sparsity.attention_sparsity.calibration.mask_reuse import ( + MaskReuseCalibrationError, + calibrate_mask_reuse_policy, +) +from modelopt.torch.sparsity.attention_sparsity.calibration.mask_reuse_compact import ( + calibrate_compact_mask_reuse_policy, + load_compact_mask_reuse_captures, +) + +_FIT = { + "threshold_scale_factor": { + "formula": "a * exp(b * target_sparsity)", + "prefill": { + "a": 1.0, + "b": 1.0, + "min_observed_sparsity": 0.4, + "max_observed_sparsity": 0.8, + }, + } +} +_TOPOLOGY = {"anchors": [0], "nearest": {"0": 0, "1": 0}} + + +def _checkpoint(tmp_path): + root = tmp_path / "checkpoint" + root.mkdir() + (root / "config.json").write_text("{}\n", encoding="utf-8") + (root / "model.safetensors").write_bytes(b"weights") + return create_checkpoint_manifest(root, model="test-model") + + +def _capture(split, prompt_id, target, checkpoint_sha256): + threshold_log2 = math.log2(1.0) + target * math.log2(math.e) - math.log2(256) + if target == 0.5: + retained = [2, 3] + anchor_dropped = [0.005, 0.005] + matrix = [[0.001, 0.001], [0.001, 0.001]] + else: + retained = [1, 2] + anchor_dropped = [0.02, 0.02] + matrix = ( + [[0.01, 0.03], [0.04, 0.01]] + if split == "calibration" + else [[0.015, 0.05], [0.05, 0.015]] + ) + invocation = { + "capture_schema_version": 2, + "model": "test-model", + "checkpoint_manifest_sha256": checkpoint_sha256, + "split": split, + "partition": "development" if split == "calibration" else "outer_test", + "inner_fold": 0 if split == "calibration" else None, + "prompt_id": prompt_id, + "source": f"dataset/{split}", + "source_group_sha256": sha256(f"group/{prompt_id}".encode()).hexdigest(), + "source_capture_sha256": sha256(prompt_id.encode()).hexdigest(), + "min_kv_tokens": 129, + "max_kv_tokens": 512, + "target_sparsity_hex": target.hex(), + "sample_length": 256, + "threshold_log2_hex": threshold_log2.hex(), + "threshold_lambda_hex": (2.0**threshold_log2).hex(), + "expected_geometry": {"q_tokens": 256, "kv_tokens": 256, "q_start_tokens": 0}, + } + return { + "compact_capture_schema_version": 1, + "invocation": invocation, + "geometry": invocation["expected_geometry"], + "global_num_heads": 2, + "eligible_tiles": 3, + "anchor_stats_by_layer": { + "0": {"retained_tiles": retained, "dropped_mass": anchor_dropped} + }, + "consumer_layers": {"1": {"anchor_layer": 0, "dropped_mass": matrix}}, + } + + +def _write_captures(path, checkpoint_sha256): + captures = [ + _capture(split, prompt, target, checkpoint_sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + payload = b"".join( + ( + json.dumps(capture, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + for capture in captures + ) + path.write_bytes(payload) + return captures, sha256(payload).hexdigest() + + +def _write_payload(path, captures): + payload = b"".join( + ( + json.dumps(capture, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + for capture in captures + ) + path.write_bytes(payload) + return sha256(payload).hexdigest() + + +def _evidence(reuse_bundle_sha256): + fields = ( + "calibration_plan_sha256", + "family_registry_sha256", + "vanilla_fit_sha256", + "reuse_bundle_sha256", + "grouped_fit_sha256", + "outer_report_sha256", + ) + return { + field: reuse_bundle_sha256 + if field == "reuse_bundle_sha256" + else sha256(field.encode()).hexdigest() + for field in fields + } + + +def _expanded_rows(captures): + rows = [] + for capture in captures: + invocation = capture["invocation"] + anchors = capture["anchor_stats_by_layer"] + anchor = anchors["0"] + matrix = capture["consumer_layers"]["1"]["dropped_mass"] + for consumer_head in range(2): + rows.extend( + { + "model": invocation["model"], + "min_kv_tokens": invocation["min_kv_tokens"], + "max_kv_tokens": invocation["max_kv_tokens"], + "target_sparsity": float.fromhex(invocation["target_sparsity_hex"]), + "sample_length": invocation["sample_length"], + "threshold_lambda": float.fromhex(invocation["threshold_lambda_hex"]), + "threshold_log2": float.fromhex(invocation["threshold_log2_hex"]), + "q_tokens": 256, + "kv_tokens": 256, + "q_start_tokens": 0, + "split": invocation["split"], + "prompt_id": invocation["prompt_id"], + "source_capture_sha256": invocation["source_capture_sha256"], + "anchor_layer": 0, + "consumer_layer": 1, + "consumer_head": consumer_head, + "donor_head": donor_head, + "retained_tiles": anchor["retained_tiles"][donor_head], + "eligible_tiles": 3, + "anchor_dropped_mass": anchor["dropped_mass"][donor_head], + "anchor_stats_by_layer": anchors, + "dropped_mass": matrix[consumer_head][donor_head], + } + for donor_head in range(2) + ) + return rows + + +def test_streaming_compact_selector_matches_legacy_row_policy(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures, digest = _write_captures(path, checkpoint.sha256) + kwargs = { + "vanilla_calibration": _FIT, + "topology": _TOPOLOGY, + "checkpoint_manifest": checkpoint, + "evidence": _evidence(digest), + "max_anchor_dropped_mass": 0.03, + "reuse_dropped_mass_report_threshold": 0.025, + "target_bmm1_skip_ratio": 0.25, + } + + compact = calibrate_compact_mask_reuse_policy(load_compact_mask_reuse_captures(path), **kwargs) + legacy = calibrate_mask_reuse_policy(_expanded_rows(captures), **kwargs) + + assert compact["context_policies"] == legacy["context_policies"] + assert compact["context_policies"][0]["target_sparsity"] == 0.7 + assert compact["context_policies"][0]["headmaps"] == {"1": [0, 1]} + compact_report = dict(compact["calibration_report"]) + legacy_report = dict(legacy["calibration_report"]) + compact_report.pop("promotion") + legacy_report.pop("promotion") + assert compact_report["constraints"] == legacy_report["constraints"] + assert compact_report["overall"] == legacy_report["overall"] + assert compact["provenance"]["input_capture_count"] == 4 + assert compact["provenance"]["candidate_cell_count"] == 16 + assert compact["provenance"]["streaming_passes"] == [ + "validation", + "calibration_selection", + "frozen_evaluation", + ] + assert compact["promotion_status"] == "candidate_only" + assert compact["deployment_geometry_validated"] is False + + +def test_compact_public_interface_has_no_hard_reuse_risk_gate(): + parameters = inspect.signature(calibrate_compact_mask_reuse_policy).parameters + + assert parameters["target_bmm1_skip_ratio"].default is inspect.Parameter.empty + assert "max_reuse_selection_dropped_mass" not in parameters + + +def test_compact_selector_binds_reuse_bundle_sha(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + _, digest = _write_captures(path, checkpoint.sha256) + evidence = _evidence(digest) + evidence["reuse_bundle_sha256"] = sha256(b"wrong").hexdigest() + + with pytest.raises(MaskReuseCalibrationError, match="does not match"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=evidence, + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + ) + + +def test_compact_selector_binds_verified_checkpoint_and_disjoint_groups(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures = [ + _capture(split, prompt, target, checkpoint.sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + captures[0]["invocation"]["checkpoint_manifest_sha256"] = "0" * 64 + digest = _write_payload(path, captures) + with pytest.raises(MaskReuseCalibrationError, match="one model, checkpoint"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + ) + + shared_group = captures[1]["invocation"]["source_group_sha256"] + for capture in captures: + capture["invocation"]["checkpoint_manifest_sha256"] = checkpoint.sha256 + capture["invocation"]["source_group_sha256"] = shared_group + digest = _write_payload(path, captures) + with pytest.raises(MaskReuseCalibrationError, match="multiple partitions"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + ) + + +def test_compact_selector_rejects_file_changed_during_evaluation(tmp_path, monkeypatch): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + _, digest = _write_captures(path, checkpoint.sha256) + real_evaluation_pass = mask_reuse_compact._evaluation_pass + + def mutate_after_evaluation(*args, **kwargs): + result = real_evaluation_pass(*args, **kwargs) + with path.open("ab") as handle: + handle.write(b"\n") + return result + + monkeypatch.setattr(mask_reuse_compact, "_evaluation_pass", mutate_after_evaluation) + + with pytest.raises(MaskReuseCalibrationError, match="changed during calibration"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + ) + + +def test_selector_meets_bmm1_target_before_minimizing_reuse_risk(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures = [ + _capture(split, prompt, target, checkpoint.sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + for capture in captures: + target = float.fromhex(capture["invocation"]["target_sparsity_hex"]) + if target == 0.5: + capture["anchor_stats_by_layer"]["0"] = { + "retained_tiles": [1, 3], + "dropped_mass": [0.001, 0.001], + } + capture["anchor_stats_by_layer"]["2"] = { + "retained_tiles": [3, 3], + "dropped_mass": [0.001, 0.001], + } + capture["consumer_layers"]["1"]["dropped_mass"] = [ + [0.001, 0.001], + [0.001, 0.001], + ] + else: + capture["anchor_stats_by_layer"]["0"] = { + "retained_tiles": [0, 2], + "dropped_mass": [0.002, 0.002], + } + capture["anchor_stats_by_layer"]["2"] = { + "retained_tiles": [0, 0], + "dropped_mass": [0.002, 0.002], + } + capture["consumer_layers"]["1"]["dropped_mass"] = [ + [0.03, 0.01], + [0.03, 0.01], + ] + digest = _write_payload(path, captures) + + candidate = calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology={"anchors": [0, 2], "nearest": {"0": 0, "1": 0, "2": 2}}, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + ) + + assert candidate["context_policies"][0]["target_sparsity"] == 0.7 + frontier = candidate["calibration_report"]["by_bucket"][0]["target_sparsity_frontier"] + assert [row["target_bmm1_skip_ratio_feasible"] for row in frontier] == [False, True] + assert frontier[1]["combined_tile_cost"] == 2 + + +def test_validation_pass_does_not_retain_capture_objects(): + implementation = inspect.getsource(mask_reuse_compact._validate_dataset) + + assert ".append(capture)" not in implementation + assert "list[CompactMaskReuseCapture]" not in implementation + + +@pytest.mark.parametrize("corruption", ["eligible", "retained_monotonic", "mass_monotonic"]) +def test_compact_validation_rejects_impossible_geometry_or_sparsity_trend(tmp_path, corruption): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures = [ + _capture(split, prompt, target, checkpoint.sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + if corruption == "eligible": + captures[0]["eligible_tiles"] = 4 + elif corruption == "retained_monotonic": + captures[1]["anchor_stats_by_layer"]["0"]["retained_tiles"][0] = 3 + else: + captures[1]["consumer_layers"]["1"]["dropped_mass"][0][0] = 0.0 + digest = _write_payload(path, captures) + + with pytest.raises(MaskReuseCalibrationError): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + )