diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6687ebd31ea..a2019000088 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,10 @@ Changelog **New Features** +*Sparsity* + +- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. Configurations that cannot compile the 128x128 tile are rejected rather than re-tiled, and skip-softmax cannot be combined with attention quantization. **Behavior change for sparse-only serving installs:** ``install_vllm_sparse_attention_from_checkpoint`` now runs the same engine-level compatibility validation as quantized installs (decode context parallelism, DBO, speculative decoding, and FULL mixed-batch CUDA graphs are rejected; prefix caching remains supported), and checkpoints with a calibrated ``decode`` ``threshold_scale_factor`` are rejected under a FULL decode CUDA graph mode — including vLLM's default ``FULL_AND_PIECEWISE`` — because the captured graph would replay one request's stale threshold. Previously such deployments started and silently served miscalibrated sparsity; pass ``--enforce-eager`` (or select a non-FULL decode graph mode). See ``examples/vllm_serve/calibrate_sparse_attn.py``. + *Quantization* - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported. diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index fc4e8a0ebcc..c5596003cd3 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -178,6 +178,26 @@ Workflow: If the checkpoint has no `sparse_attention_config`, the sparse-only installer passes through and vLLM runs unchanged. Whole-model fakequant flows remain handled by `vllm_serve_fakequant.py`; the compact attention-only path is below. +### Calibrate skip-softmax thresholds through vLLM + +Instead of the HF path in step 1, thresholds can be calibrated directly through vLLM — over the paged KV cache, for both prefill and decode, with tensor parallelism: + +```bash +# One-time: fetch the RULER essay haystack +bash ../llm_sparsity/attention_sparsity/download_ruler_data.sh + +python calibrate_sparse_attn.py \ + --calib_data_dir ../llm_sparsity/attention_sparsity/data \ + --target_sparse_ratio 0.5 \ + --decode_tokens 32 --tensor_parallel_size 8 --update_checkpoint_config +``` + +Calibration prompts default to the **RULER dataset** via the same `RulerDatasetBuilder` the HF calibration path uses (`--calib_samples` / `--calib_max_seqlen` mirror the HF defaults of 24 / 32768), so vLLM- and PyTorch-calibrated thresholds are fit on identical data. `--prompts_file` (one prompt per line) substitutes custom calibration data. + +`install_vllm_skip_softmax_calibration` (called by `sparse_attn_worker.SkipSoftmaxCalibWorker` at model load) swaps calibration adapters onto every attention layer after validating all of them — eager execution is required, model and KV-cache dtypes must be fp16/bf16, and no attention Q/K/P/V fakequant may be active. During `llm.generate`, the paged Triton calibration kernel computes full dense attention — no sparsification is applied to generation, though the dense kernel's numerics differ slightly from the native backend's — while counting, per candidate threshold, how many KV tiles the skip criterion would drop. The driver then collects **raw tile counts from every TP rank** (each rank only measures its head shard), merges them, fits `scale_factor = a * exp(b * sparsity)` once per phase, and writes the same canonical `sparse_attention_config` block the HF export produces — preserving any exported N:M sparse-softmax groups — so the serving workflow above picks it up unchanged. + +Calibration and serving measure skipping at the same fixed 128x128 tile geometry (active skip-softmax launches bypass the autotuner), so the sparsity realized at serve time matches the calibrated `(a, b)` model. + The reusable serving policies live in `modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py`. `install_vllm_sparse_attention_from_checkpoint` installs checkpoint-driven sparse-only attention, while `install_vllm_nvfp4_attention` installs fixed NVFP4 Q/K/P/V with optional checkpoint sparsity. Both validate every selected layer before publishing any replacement implementation and return a `VllmAttentionInstallReport` with the installed layer names and backend counts. `sparse_attn_worker.py` only invokes these APIs after vLLM loads the model. It retains `SparseAttnWorker` as the launcher's default and provides `QuantSparseAttnWorker` for the compact NVFP4 policy. Other vLLM integrations can invoke the same library APIs directly: @@ -192,8 +212,9 @@ report = install_vllm_nvfp4_attention(model_runner, sparse_cfg="checkpoint") Limitations: -- vLLM V1 chunked prefill and prefix-cache suffix attention are supported by offsetting query positions into the longer KV span. -- `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`. +- vLLM V1 chunked prefill and prefix-cache suffix attention are supported by offsetting query positions into the longer KV span. This applies to sparse-only serving; quantized attention installs and skip-softmax calibration reject `enable_prefix_caching` (quantize-on-write and per-request measurement both require uncached prefills). +- `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`. Checkpoints with a calibrated `decode` `threshold_scale_factor` are rejected at install under a FULL decode CUDA graph mode (including vLLM's default `FULL_AND_PIECEWISE`): the captured graph would replay one request's stale threshold. +- Sparse-only installs validate engine-level compatibility like quantized installs do: decode context parallelism, DBO, speculative decoding, and FULL mixed-batch CUDA graphs are rejected (prefix caching remains supported, per the bullet above). ### Compact NVFP4 attention worker @@ -209,7 +230,7 @@ python vllm_serve_sparse_attn.py -tp 8 \ The installer supports both FlashInfer and FlashAttention, and the worker prints the installed adapter counts. Pass `--attention-backend FLASHINFER` or `--attention-backend FLASH_ATTN` only when an explicit override is needed. -This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant format to Q/K/P/V. Q is dynamic; missing K/V scales default to global scale 1.0, and P defaults to amax 1.0. Existing scalar attention amax values are preserved, but this path does not calibrate or restore them itself. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored. +This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant format to Q/K/P/V. Q is dynamic; missing K/V scales default to global scale 1.0, and P defaults to amax 1.0. Existing scalar attention amax values are preserved, but this path does not calibrate or restore them itself. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored for N:M sparse softmax; calibrated skip-softmax groups are rejected in combination with attention quantization, because quantized Q/K/P change the score distribution the skip thresholds were calibrated on. Decode uses a fixed 32-split, 128-key-tile schedule. P QDQ consumes split-local, unnormalized online-softmax probabilities, so changing that schedule can change 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/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py new file mode 100644 index 00000000000..42c4d36a972 --- /dev/null +++ b/examples/vllm_serve/calibrate_sparse_attn.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. + +"""Calibrate skip-softmax thresholds *through vLLM* and write the serving config. + +Runs calibration prompts through a vLLM ``LLM`` whose attention layers carry +the ModelOpt calibration adapters (installed by +``sparse_attn_worker.SkipSoftmaxCalibWorker`` via +``install_vllm_skip_softmax_calibration``). The paged Triton calibration +kernel measures, per candidate threshold, how many KV tiles would be skipped — +over the paged KV cache, for both prefill and decode — then this driver +aggregates the raw counts from every tensor-parallel rank and fits the +exponential model ``scale_factor = a * exp(b * sparsity)`` once per phase. + +The fitted ``(a, b)`` are written as a canonical ``sparse_attention_config`` +block (the same schema ModelOpt's HF export produces), so the serving path +(``vllm_serve_sparse_attn.py`` / ``install_vllm_sparse_attention_from_checkpoint``) +loads it without changes. Any exported N:M sparse-softmax groups already in +the checkpoint config are preserved. + +Usage: + python calibrate_sparse_attn.py \ + --target_sparse_ratio 0.5 \ + --decode_tokens 32 \ + --update_checkpoint_config + +Calibration prompts default to the RULER dataset — the same +``RulerDatasetBuilder`` the PyTorch (HF) calibration path uses — so both paths +calibrate on identical data. NIAH tasks need the essay haystack downloaded by +``examples/llm_sparsity/attention_sparsity/download_ruler_data.sh`` (point +``--calib_data_dir`` at its ``data`` directory). ``--prompts_file`` (one prompt +per line) overrides the RULER set with custom calibration data. +""" + +import argparse +import json +import os +import sys +from pathlib import Path + +from modelopt.torch.sparsity.attention_sparsity.calibration.ruler_dataset import RulerDatasetBuilder +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_calibration import ( + DEFAULT_THRESHOLD_TRIALS, + build_sparse_attention_config, + fit_from_counts, + merge_phase_counts, +) + + +def _load_prompts(llm, args) -> list[str]: + """Load override prompts from a file, or build the default RULER set.""" + if args.prompts_file is not None: + lines = [ + ln.strip() for ln in Path(args.prompts_file).read_text().splitlines() if ln.strip() + ] + if not lines: + raise ValueError(f"No prompts found in {args.prompts_file}") + print(f"[ModelOpt] Loaded {len(lines)} calibration prompts from {args.prompts_file}") + return lines + + # Same dataset as the HF calibration path (calibration/calibrate.py), so the + # vLLM- and PyTorch-calibrated thresholds are fit on identical data. + builder = RulerDatasetBuilder( + samples=args.calib_samples, + max_seqlen=args.calib_max_seqlen, + tokenizer_name_or_path=llm.get_tokenizer(), + max_length_filter=int(args.calib_max_seqlen * 1.5), + data_dir=args.calib_data_dir, + ) + samples = builder.build_calibration_dataset() + if not samples: + raise ValueError( + "RULER produced no calibration samples (all candidates exceeded " + f"max_length_filter={int(args.calib_max_seqlen * 1.5)} tokens). " + "Adjust --calib_max_seqlen / --calib_samples, or pass --prompts_file." + ) + prompts = [sample["input"] for sample in samples] + lengths = sorted(sample["length"] for sample in samples) + print( + f"[ModelOpt] Built {len(prompts)} RULER calibration prompts " + f"(token lengths {lengths[0]}..{lengths[-1]})" + ) + return prompts + + +def _existing_sparse_config(ckpt: str) -> dict | None: + """Read the checkpoint's sparse_attention_config so non-skip groups survive.""" + config_json = Path(ckpt) / "config.json" + if not config_json.is_file(): + return None + existing = json.loads(config_json.read_text()).get("sparse_attention_config") + return existing if isinstance(existing, dict) else None + + +def _write_config(ckpt: str, sparse_config: dict, update_checkpoint: bool) -> None: + """Dump the sparse_attention_config and optionally merge into config.json.""" + out_path = Path("sparse_attention_config.json") + out_path.write_text(json.dumps(sparse_config, indent=2)) + print(f"[ModelOpt] Wrote calibrated config to {out_path.resolve()}") + + if not update_checkpoint: + print( + "[ModelOpt] Re-run with --update_checkpoint_config to merge this into " + f"{ckpt}/config.json (required for vllm_serve_sparse_attn.py to pick it up)." + ) + return + + config_json = Path(ckpt) / "config.json" + config = json.loads(config_json.read_text()) + config["sparse_attention_config"] = sparse_config + # Atomic replace: a crash mid-write must not truncate the checkpoint's + # config.json (write_text would rewrite it in place). + tmp_path = config_json.with_name(config_json.name + ".tmp") + tmp_path.write_text(json.dumps(config, indent=2)) + os.replace(tmp_path, config_json) + print(f"[ModelOpt] Merged sparse_attention_config into {config_json}") + + +def main(): + parser = argparse.ArgumentParser(description="Calibrate skip-softmax thresholds via vLLM") + parser.add_argument("model", type=str, help="Path to the HF checkpoint to calibrate") + parser.add_argument( + "--prompts_file", + type=str, + default=None, + help="Optional custom calibration prompts (one per line), overriding the " + "default RULER dataset", + ) + parser.add_argument( + "--calib_samples", + type=int, + default=24, + help="Total RULER samples, distributed across length bins (HF-path default: 24)", + ) + parser.add_argument( + "--calib_max_seqlen", + type=int, + default=32768, + help="Maximum RULER sequence length; length bins descend in powers of 2. " + "Must fit within --max_model_len together with --decode_tokens.", + ) + parser.add_argument( + "--calib_data_dir", + type=str, + default=None, + help="RULER data directory containing the 'essays' haystack (populated by " + "examples/llm_sparsity/attention_sparsity/download_ruler_data.sh)", + ) + parser.add_argument( + "--target_sparse_ratio", + type=float, + default=0.5, + help="Target sparsity baked into the exported config (applied to both phases)", + ) + parser.add_argument( + "--decode_tokens", + type=int, + default=32, + help="Decode attention steps per prompt (drives decode-phase calibration). " + "Generation runs decode_tokens + 1 output tokens: the first output token " + "comes from the prefill forward and performs no decode attention.", + ) + parser.add_argument( + "--max_model_len", type=int, default=None, help="vLLM max_model_len override" + ) + parser.add_argument( + "--tensor_parallel_size", type=int, default=1, help="vLLM tensor-parallel size" + ) + parser.add_argument( + "--gpu_memory_utilization", + type=float, + default=None, + help="vLLM GPU memory utilization fraction", + ) + parser.add_argument( + "--trust_remote_code", + action="store_true", + help="Trust remote code for custom model classes (e.g. NemotronH)", + ) + parser.add_argument("--dtype", type=str, default=None, help="Model dtype, e.g. bfloat16") + parser.add_argument( + "--attention_backend", + type=str, + default=None, + help="Force the vLLM attention backend, e.g. FLASH_ATTN or FLASHINFER. " + "Default: let vLLM choose (the installer supports whichever of FlashAttention " + "/ FlashInfer is selected).", + ) + parser.add_argument( + "--engine_kwargs", + type=str, + default=None, + help="JSON dict of extra vLLM engine kwargs, e.g. " + '\'{"enable_expert_parallel": true, "mamba_cache_mode": "align"}\' ' + "for hybrid MoE/Mamba models", + ) + parser.add_argument( + "--fit_logspace", + action="store_true", + help="Fit the exponential model in log space (wide scale_factor ranges)", + ) + parser.add_argument( + "--update_checkpoint_config", + action="store_true", + help="Merge the calibrated config into /config.json in place", + ) + args = parser.parse_args() + + if args.update_checkpoint_config and not (Path(args.model) / "config.json").is_file(): + # Fail before the (expensive, multi-GPU) calibration run, not after: + # merging requires a local checkpoint directory, not a HF hub ID. + parser.error( + f"--update_checkpoint_config requires a local checkpoint directory " + f"containing config.json; {args.model!r} has none" + ) + + # Workers run in separate processes and must import the calibration worker. + repo_root = str(Path(__file__).resolve().parent) + if repo_root not in sys.path: + sys.path.insert(0, repo_root) + current = os.environ.get("PYTHONPATH") + os.environ["PYTHONPATH"] = os.pathsep.join([current, repo_root]) if current else repo_root + + # Deferred heavy import: keep argparse/--help (and arg errors) fast, and + # only import vLLM after the PYTHONPATH setup above. + from vllm import LLM, SamplingParams + + llm_kwargs = { + "model": args.model, + "worker_cls": "sparse_attn_worker.SkipSoftmaxCalibWorker", + # The calibration installer requires eager execution: the per-request + # calibration loop cannot be CUDA-graph captured. + "enforce_eager": True, + # Shared-prefix reuse would make prefill measurements cover only the + # non-cached suffix of each prompt; the installer rejects it. + "enable_prefix_caching": False, + } + if args.max_model_len is not None: + llm_kwargs["max_model_len"] = args.max_model_len + if args.tensor_parallel_size and args.tensor_parallel_size > 1: + llm_kwargs["tensor_parallel_size"] = args.tensor_parallel_size + if args.gpu_memory_utilization is not None: + llm_kwargs["gpu_memory_utilization"] = args.gpu_memory_utilization + if args.trust_remote_code: + llm_kwargs["trust_remote_code"] = True + if args.dtype is not None: + llm_kwargs["dtype"] = args.dtype + if args.attention_backend is not None: + llm_kwargs["attention_backend"] = args.attention_backend + if args.engine_kwargs: + extra = json.loads(args.engine_kwargs) + if not isinstance(extra, dict): + raise ValueError("--engine_kwargs must be a JSON object") + llm_kwargs.update(extra) + llm = LLM(**llm_kwargs) + + # Built after engine init so the RULER builder reuses the engine's tokenizer. + prompts = _load_prompts(llm, args) + + trials = list(DEFAULT_THRESHOLD_TRIALS) + n_layers = llm.collective_rpc("sparse_calib_enable", args=(trials,))[0] + status = llm.collective_rpc("sparse_calib_status")[0] + print(f"[ModelOpt] Calibration enabled on {n_layers} attention layers") + print(f"[ModelOpt] Active sparse impls: {status['impl_types']}") + + # generate() drives prefill (prefill-phase stats) then decode steps + # (decode-phase stats). No sparsification is applied during calibration — + # the kernel computes full dense attention while recording tile-skip + # counts. ignore_eos forces the full decode length so early EOS cannot + # thin the decode-phase statistics. max_tokens is decode_tokens + 1: the + # first output token comes from the prefill forward, so decode_tokens + # decode-attention steps need one extra output token. + sampling = SamplingParams(temperature=0.0, max_tokens=args.decode_tokens + 1, ignore_eos=True) + llm.generate(prompts, sampling) + + # Aggregate RAW counts from every TP rank (each rank only measures its + # attention-head shard), then fit once per phase on the global counts. + rank_counts = llm.collective_rpc("sparse_calib_counts") + merged = merge_phase_counts(rank_counts) + calibration_params = fit_from_counts(merged, trials, fit_logspace=args.fit_logspace) + + requested_phases = ["prefill"] + (["decode"] if args.decode_tokens > 0 else []) + missing = [phase for phase in requested_phases if phase not in calibration_params] + if missing: + print( + f"[ModelOpt] Calibration FAILED: no valid fit for phase(s) {', '.join(missing)}. " + "No config was written — a partially calibrated export would silently serve " + "the missing phase dense. Try more/longer prompts (and more decode tokens) " + "so observed sparsity spans the (10%, 90%) fitting window." + ) + sys.exit(1) + # Export only requested phases: a stray record (e.g. a scheduling corner + # case classified into an unrequested phase) must not bake an + # uncalibrated-by-intent phase into the config. + calibration_params = { + phase: params for phase, params in calibration_params.items() if phase in requested_phases + } + + sparse_config = build_sparse_attention_config( + calibration_params, + {"prefill": args.target_sparse_ratio, "decode": args.target_sparse_ratio}, + existing_config=_existing_sparse_config(args.model), + ) + print("[ModelOpt] Calibrated threshold_scale_factor:") + print(json.dumps(sparse_config["config_groups"]["group_0"]["threshold_scale_factor"], indent=2)) + _write_config(args.model, sparse_config, args.update_checkpoint_config) + + +if __name__ == "__main__": + 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/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 7b5fb28bf68..27e1eebdf97 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -17,12 +17,22 @@ from vllm.v1.worker.gpu_worker import Worker as BaseWorker +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_calibration import ( + DEFAULT_THRESHOLD_TRIALS, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( + collect_calibration_counts, + disable_calibration, + enable_calibration, + iter_sparse_impls, +) from modelopt.torch.sparsity.attention_sparsity.plugins.vllm_runtime import ( install_vllm_nvfp4_attention, + install_vllm_skip_softmax_calibration, install_vllm_sparse_attention_from_checkpoint, ) -__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker"] # noqa: RUF022 +__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker", "SkipSoftmaxCalibWorker"] # noqa: RUF022 _QUANT_FORMAT_KEYS = ("q_format", "k_format", "p_format", "v_format") @@ -69,6 +79,57 @@ def load_model(self, *args, **kwargs) -> None: _print_install_report("Sparse attention", report) +class SkipSoftmaxCalibWorker(BaseWorker): + """Calibrate skip-softmax thresholds through the engine. + + Unlike :class:`SparseAttnWorker` (which serves an already-calibrated + ``sparse_attention_config``), this worker *produces* that config. The + library installer swaps calibration-capable adapters onto every attention + layer at load; measurement starts only when the driver calls + ``sparse_calib_enable`` (so warmup launches are never recorded) and raw + per-threshold tile counts are harvested with ``sparse_calib_counts`` for + the driver to aggregate across TP ranks and fit. + """ + + def load_model(self, *args, **kwargs) -> None: + """Load the model, then install calibration adapters on every layer.""" + super().load_model(*args, **kwargs) + report = install_vllm_skip_softmax_calibration(self.model_runner) + print( + f"[ModelOpt] Skip-softmax calibration installed on {report.installed_count} " + f"attention layers: {dict(report.backend_counts)}" + ) + + # -- RPC methods (invoked via LLM.collective_rpc) ---------------------- + + def sparse_calib_enable(self, threshold_trials: list[float] | None = None) -> int: + """Enter calibration mode on all installed impls; returns layer count.""" + impls = list(iter_sparse_impls(_unwrapped_model(self))) + enable_calibration(impls, list(threshold_trials or DEFAULT_THRESHOLD_TRIALS)) + return len(impls) + + def sparse_calib_status(self) -> dict: + """Report active impls and record counts, so the backend is verifiable.""" + impls = list(iter_sparse_impls(_unwrapped_model(self))) + impl_types: dict[str, int] = {} + total_records = 0 + for impl in impls: + impl_types[type(impl).__name__] = impl_types.get(type(impl).__name__, 0) + 1 + total_records += len(getattr(impl, "_calib_records", [])) + return { + "num_sparse_layers": len(impls), + "impl_types": impl_types, + "calibrating": any(getattr(impl, "_calibrate", False) for impl in impls), + "total_records": total_records, + } + + def sparse_calib_counts(self) -> dict[str, list[dict]]: + """Stop measuring and return this rank's layer-merged raw tile counts.""" + model = _unwrapped_model(self) + disable_calibration(list(iter_sparse_impls(model))) + return collect_calibration_counts(model) + + class QuantSparseAttnWorker(BaseWorker): """Install quantized attention plus optional checkpoint sparsity. diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 5acc9fb787c..4a223d4c2ab 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -87,7 +87,6 @@ def _load_qdq_helpers() -> None: ] _MEASURE_BLOCK_M = 128 -_P_QDQ_MEASURE_BLOCK_M = 16 # 128 so the kernel sparsity-measurement block matches the PyTorch # calibration/reference granularity. This is deliberately independent of the # autotuned compute tile. @@ -944,6 +943,16 @@ def forward( else: apply_skip = False skip_threshold_log2 = 0.0 + if apply_skip and (p_qdq_mode or v_qdq_mode): + # Quantized operands change what the calibrated skip thresholds mean, + # and P-QDQ additionally uses a different measurement tile geometry. + # The vLLM installers reject this composition at plan time; the raw + # kernel API rejects it here so no path can serve it. + raise ValueError( + "skip-softmax cannot be combined with attention quantization " + "(P/V QDQ): the calibrated tile-skip contract does not hold " + "under quantized operands" + ) o = torch.empty_like(q) lse = torch.empty(q.shape[0], num_q_heads, device=q.device, dtype=torch.float32) @@ -1030,17 +1039,40 @@ def grid(META): # kernel dereferences the right pointers instead of triggering an # illegal memory access. with torch.cuda.device(q.device): - if do_measure: - # Runtime counters mutate global tensors, so do not run them through - # autotune candidate trials. Use one stable config for measurement. - _attn_fwd.fn[grid]( - *fwd_args, - **fwd_kwargs, - BLOCK_M=_P_QDQ_MEASURE_BLOCK_M if p_qdq_mode else _MEASURE_BLOCK_M, - BLOCK_N=_MEASURE_BLOCK_N, - num_warps=_MEASURE_NUM_WARPS, - num_stages=_MEASURE_NUM_STAGES, - ) + if do_measure or apply_skip: + # Fixed-tile launches, bypassing autotune: + # - Measurement: runtime counters mutate global tensors, so they + # must not run through autotune candidate trials. + # - Active skip-softmax: the tile-skip decision depends on the + # (BLOCK_M, BLOCK_N) geometry, and thresholds are calibrated at + # the 128x128 measurement granularity (attention_calibrate and + # flash_skip_softmax both use 128x128 blocks). Autotuned tiles + # (e.g. BLOCK_N=32) would realize a different sparsity than + # calibrated, so skip launches always use the calibration tile. + # + # The tile is a contract, not a preference: configurations that + # cannot compile it (e.g. fp32 inputs on ~100KB-shared-memory + # GPUs) are rejected rather than re-tiled, because a different + # tile realizes a different sparsity than was calibrated. + try: + # P/V QDQ is rejected above when skip is active, so the tile + # here is unconditionally the 128x128 calibration geometry. + _attn_fwd.fn[grid]( + *fwd_args, + **fwd_kwargs, + BLOCK_M=_MEASURE_BLOCK_M, + BLOCK_N=_MEASURE_BLOCK_N, + num_warps=_MEASURE_NUM_WARPS, + num_stages=_MEASURE_NUM_STAGES, + ) + except triton.runtime.errors.OutOfResources as err: + raise RuntimeError( + "skip-softmax requires the fixed 128x128 calibration tile, " + f"which exceeds this GPU's shared memory for {q.dtype} " + f"inputs ({err}). Use fp16/bf16 inputs or a device with " + "more shared memory; re-tiling would change the " + "calibrated sparsity contract." + ) from err else: _attn_fwd[grid]( *fwd_args, diff --git a/modelopt/torch/kernels/sparsity/attention/calibrate.py b/modelopt/torch/kernels/sparsity/attention/calibrate.py index d26e781d48c..3655c4cbfcd 100644 --- a/modelopt/torch/kernels/sparsity/attention/calibrate.py +++ b/modelopt/torch/kernels/sparsity/attention/calibrate.py @@ -22,13 +22,19 @@ ``modelopt.torch.sparsity.attention_sparsity`` to fit a skip threshold. """ +import functools import math import torch import triton import triton.language as tl -from modelopt.torch.kernels.common.attention.triton_fa import LOG2E, _apply_mask +from modelopt.torch.kernels.common.attention.triton_fa import ( + LOG2E, + _apply_mask, + _load_paged_k_tile, + _load_paged_v_tile, +) # --------------------------------------------------------------------------- @@ -64,6 +70,19 @@ def _attn_fwd_calibrate( HEAD_DIM: tl.constexpr, NUM_THRESHOLDS: tl.constexpr, PADDED_THRESHOLDS: tl.constexpr, # next_power_of_2(NUM_THRESHOLDS) for tl.arange + Q_IS_FP32: tl.constexpr = False, # match the serving kernel's IEEE fp32 QK dot + IS_PAGED: tl.constexpr = False, # Whether K/V are read from a paged KV cache + K_cache=None, # [num_blocks, page_size, num_kv_heads, head_dim] paged K + V_cache=None, # [num_blocks, page_size, num_kv_heads, head_dim] paged V + Block_table=None, # [batch, max_blocks_per_seq] page table + stride_kc_block=0, + stride_kc_pos=0, + stride_kc_head=0, + stride_vc_block=0, + stride_vc_pos=0, + stride_vc_head=0, + PAGE_SIZE: tl.constexpr = 16, + max_blocks_per_seq=0, ): """Forward kernel with multi-threshold sparsity measurement. @@ -126,14 +145,41 @@ def _attn_fwd_calibrate( for kv_start in range(0, kv_bound, BLOCK_N): kv_start = tl.multiple_of(kv_start, BLOCK_N) - k_offs = (kv_offset + kv_start + kv_pos[None, :]) * stride_kbs + dim_pos[:, None] - k = tl.load( - k_base + k_offs, - mask=((kv_start + kv_pos[None, :]) < seq_len_kv) & d_mask[:, None], - other=0.0, - ) - - scores = tl.dot(q, k) * qk_scale + # Load K^T [BLOCK_D, BLOCK_N] from paged cache or contiguous K. + if IS_PAGED: + k = _load_paged_k_tile( + K_cache, + Block_table, + batch_idx, + kv_head_idx, + kv_start, + kv_pos, + dim_pos, + seq_len_kv, + stride_kc_block, + stride_kc_pos, + stride_kc_head, + PAGE_SIZE, + BLOCK_N, + BLOCK_D, + HEAD_DIM, + max_blocks_per_seq, + ) + else: + k_offs = (kv_offset + kv_start + kv_pos[None, :]) * stride_kbs + dim_pos[:, None] + k = tl.load( + k_base + k_offs, + mask=((kv_start + kv_pos[None, :]) < seq_len_kv) & d_mask[:, None], + other=0.0, + ) + + # Match the serving kernel's QK precision: fp32 Q uses the IEEE dot + # (default tl.dot is TF32 for fp32 inputs), so near-threshold scores + # round to the same skip decisions in calibration and serving. + if Q_IS_FP32: + scores = tl.dot(q, k.to(tl.float32), input_precision="ieee") * qk_scale + else: + scores = tl.dot(q, k) * qk_scale scores = _apply_mask(scores, q_pos, kv_pos, seq_len_q, seq_len_kv, kv_start, IS_CAUSAL) tile_row_max = tl.max(scores, 1) @@ -164,12 +210,32 @@ def _attn_fwd_calibrate( row_sum = row_sum * correction + l_new acc = acc * correction[:, None] - v_offs = (kv_offset + kv_start + kv_pos[:, None]) * stride_vbs + dim_pos[None, :] - v = tl.load( - v_base + v_offs, - mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], - other=0.0, - ) + if IS_PAGED: + v = _load_paged_v_tile( + V_cache, + Block_table, + batch_idx, + kv_head_idx, + kv_start, + kv_pos, + dim_pos, + seq_len_kv, + stride_vc_block, + stride_vc_pos, + stride_vc_head, + PAGE_SIZE, + BLOCK_N, + BLOCK_D, + HEAD_DIM, + max_blocks_per_seq, + ) + else: + v_offs = (kv_offset + kv_start + kv_pos[:, None]) * stride_vbs + dim_pos[None, :] + v = tl.load( + v_base + v_offs, + mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], + other=0.0, + ) acc = tl.dot(p.to(v.dtype), v, acc) row_max = m_new @@ -198,6 +264,22 @@ def _attn_fwd_calibrate( tl.store(Out + o_ptrs, acc, mask=(q_pos[:, None] < seq_len_q) & d_mask[None, :]) +@functools.lru_cache(maxsize=64) +def _log2_threshold_tensor( + threshold_trials: tuple[float, ...], device: torch.device +) -> torch.Tensor: + """Build the log2-space threshold tensor, cached per (trials, device). + + Scores already include sm_scale and LOG2E; convert lambda to log2 space + only. Trials are constant for a whole calibration run, and the vLLM path + calls :func:`attention_calibrate` once per request per layer per step, so + rebuilding (and re-uploading) the tensor per call would be pure waste. + """ + return torch.tensor( + [math.log2(t) for t in threshold_trials], dtype=torch.float32, device=device + ) + + def attention_calibrate( q: torch.Tensor, k: torch.Tensor, @@ -212,6 +294,10 @@ def attention_calibrate( max_input_len_k: int | None = None, *, threshold_trials: list[float] | None = None, + k_cache: torch.Tensor | None = None, + v_cache: torch.Tensor | None = None, + block_table: torch.Tensor | None = None, + page_size: int = 16, ) -> tuple[torch.Tensor, torch.Tensor]: """Flash attention with multi-threshold skip-softmax sparsity measurement. @@ -219,12 +305,21 @@ def attention_calibrate( measuring how many KV tiles would be skipped at each threshold in ``threshold_trials``. No autograd — forward only. - All arguments except ``threshold_trials`` match + All positional arguments match :func:`modelopt.torch.kernels.common.attention.attention`. Args: threshold_trials: List of threshold values to measure sparsity for. Each value is converted to log2-scaled space for the kernel. + k_cache: Paged K cache ``[num_blocks, page_size, num_kv_heads, head_dim]``. + When provided, K/V are read from the paged cache via ``block_table`` + (vLLM NHD layout) instead of from the contiguous ``k``/``v`` tensors. + ``k``/``v`` are then dummies whose only meaningful dimension is + ``shape[1] == num_kv_heads`` (used to compute the GQA ratio). + v_cache: Paged V cache ``[num_blocks, page_size, num_kv_heads, head_dim]``. + block_table: Page table ``[batch, max_blocks_per_seq]`` mapping each + sequence's block indices to global page IDs. + page_size: Number of tokens per page in the KV cache. Returns: Tuple of ``(output, sparsity_counters)``: @@ -237,6 +332,10 @@ def attention_calibrate( if threshold_trials is None or len(threshold_trials) == 0: raise ValueError("threshold_trials must be a non-empty list") + is_paged = k_cache is not None + if is_paged and block_table is None: + raise ValueError("block_table is required when k_cache/v_cache are provided.") + # Calibration has only been validated with uniform-length batches (current # diffusion + RULER paths). Varlen inputs would exercise code paths in the # kernel that have not been tested — fail loudly rather than silently @@ -281,14 +380,22 @@ def attention_calibrate( b_seq_len_k = b_seq_len b_start_loc_k = b_start_loc + if b_start_loc_k is None: + if not is_paged: + # A zeros dummy here would silently read every sequence's K/V from + # offset 0 — fail loudly instead (contiguous K/V needs real offsets). + raise ValueError( + "b_start_loc_k is required when b_seq_len_k is provided for " + "contiguous (non-paged) K/V" + ) + # Paged mode: KV positions come from block_table, so the contiguous KV + # offsets are unused. Alias b_start_loc (same shape/dtype/device) so + # Triton can compile the tl.load without allocating a dummy per call. + b_start_loc_k = b_start_loc + num_thresholds = len(threshold_trials) - # Scores already include sm_scale and LOG2E; convert lambda to log2 space only. - threshold_tensor = torch.tensor( - [math.log2(t) for t in threshold_trials], - dtype=torch.float32, - device=q.device, - ) + threshold_tensor = _log2_threshold_tensor(tuple(threshold_trials), q.device) o = torch.empty_like(q) @@ -304,6 +411,18 @@ def attention_calibrate( num_programs * num_thresholds, dtype=torch.int32, device=q.device ) + # Paged KV cache strides (zeros when not paged; computed here so the type + # narrowing of k_cache/v_cache/block_table is explicit for the kernel call). + if is_paged: + assert k_cache is not None and v_cache is not None and block_table is not None + kc_strides = (k_cache.stride(0), k_cache.stride(1), k_cache.stride(2)) + vc_strides = (v_cache.stride(0), v_cache.stride(1), v_cache.stride(2)) + max_blocks_per_seq = block_table.shape[1] + else: + kc_strides = (0, 0, 0) + vc_strides = (0, 0, 0) + max_blocks_per_seq = 0 + # Triton launches on torch.cuda.current_device(), which is not necessarily # the device the tensors live on (e.g. under accelerate device_map="auto" # sharding). Activate the tensor's device so the kernel dereferences the @@ -338,6 +457,19 @@ def attention_calibrate( HEAD_DIM=HEAD_DIM, NUM_THRESHOLDS=num_thresholds, PADDED_THRESHOLDS=triton.next_power_of_2(num_thresholds), + Q_IS_FP32=q.dtype == torch.float32, + IS_PAGED=is_paged, + K_cache=k_cache, + V_cache=v_cache, + Block_table=block_table, + stride_kc_block=kc_strides[0], + stride_kc_pos=kc_strides[1], + stride_kc_head=kc_strides[2], + stride_vc_block=vc_strides[0], + stride_vc_pos=vc_strides[1], + stride_vc_head=vc_strides[2], + PAGE_SIZE=page_size, + max_blocks_per_seq=max_blocks_per_seq, num_warps=4, num_stages=1, ) diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..9e5b6c186ec 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -180,7 +180,10 @@ def _dynamic_block_quantize_impl( and not DISABLE_TRITON_KERNEL and amax is not None ): - return triton_kernel.fp4_fake_quant_block(inputs, amax) + # Forward the configured block size: the kernel defaults to 16, and + # silently ignoring a non-16 NVFP4 block config here would diverge + # from the cuda_ext fallback below (which honors block_size). + return triton_kernel.fp4_fake_quant_block(inputs, amax, block_size=block_size) cuda_ext_mx = get_cuda_ext_mx(raise_if_failed=True) return cuda_ext_mx.fused_amax_convert( inputs, 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/calibrator.py b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py index aded26fefdc..ef46758ead0 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py @@ -28,6 +28,33 @@ from ..stats_manager import SparseAttentionStatsManager from ..utils import get_sparse_attention_modules +# Canonical skip-softmax threshold sweep — should span sparsities from ~10% to +# ~95%. Shared by the HF calibration path (this class's default) and the vLLM +# calibration path (``plugins/sparse_attn_calibration.py``), so both fit on the +# same trial grid. +DEFAULT_THRESHOLD_TRIALS = [ + 1e-6, + 5e-6, + 1e-5, + 5e-5, + 1e-4, + 5e-4, + 1e-3, + 5e-3, + 1e-2, + 2e-2, + 5e-2, + 1e-1, + 2e-1, + 3e-1, + 5e-1, + 7e-1, + 8e-1, + 9e-1, + 9.5e-1, + 9.9e-1, +] + class DynamicThresholdCalibrator: """Dynamic threshold calibrator using Exponential model. @@ -67,28 +94,7 @@ def __init__( where scale_factors span many orders of magnitude. """ # Default threshold trials if not provided - self.threshold_trials = threshold_trials or [ - 1e-6, - 5e-6, - 1e-5, - 5e-5, - 1e-4, - 5e-4, - 1e-3, - 5e-3, - 1e-2, - 2e-2, - 5e-2, - 1e-1, - 2e-1, - 3e-1, - 5e-1, - 7e-1, - 8e-1, - 9e-1, - 9.5e-1, - 9.9e-1, - ] + self.threshold_trials = threshold_trials or list(DEFAULT_THRESHOLD_TRIALS) self.fit_logspace = fit_logspace def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dict[str, Any]: @@ -130,8 +136,6 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic # with one entry per threshold, eliminating the need for repeated forward passes. print(f"\nStage 1: Collecting {phase} sparsity data for all thresholds in one pass...") - all_data_points = [] # List of {"threshold", "length", "scale_factor", "sparsity"} - self._set_thresholds(attention_modules, self.threshold_trials) self._enable_calibration_mode(attention_modules) with torch.no_grad(): @@ -139,9 +143,38 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic per_sample_stats = self._extract_calibration_stats(attention_modules, phase=phase) self._disable_calibration_mode(attention_modules) + return self.calibrate_from_stats(per_sample_stats, phase) + + def calibrate_from_stats(self, per_sample_stats: list[dict], phase: str) -> dict[str, Any]: + """Fit the exponential model from already-collected per-sample stats. + + This is the backend-agnostic Stage 2/3 of :meth:`calibrate`. The HF and + diffusion paths reach it through :meth:`calibrate` (which runs a + ``forward_loop`` to collect the stats first); the vLLM path collects the + stats itself — one record per scheduled request — and calls this directly + so both paths share the same exponential fit. + + Args: + per_sample_stats: List of ``{"sparsity": [s_0, ..., s_n], "sample_length": L}`` + records, one per calibration sample. ``sparsity`` holds the + skipped-tile fraction at each threshold in ``threshold_trials`` + (same order, same length). + phase: Phase being calibrated ('prefill' or 'decode'). + + Returns: + Dict with calibration results including a, b, r_squared, and num_data_points. + """ + all_data_points = [] # List of {"threshold", "length", "scale_factor", "sparsity"} + for sample_stat in per_sample_stats: length = sample_stat["sample_length"] sparsity_list = sample_stat["sparsity"] + if len(sparsity_list) != len(self.threshold_trials): + # A silent zip would misattribute sparsities to thresholds. + raise ValueError( + f"per-sample sparsity has {len(sparsity_list)} entries but " + f"{len(self.threshold_trials)} threshold trials are configured" + ) for threshold, sparsity in zip(self.threshold_trials, sparsity_list): scale_factor = threshold * length all_data_points.append( @@ -153,6 +186,12 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic } ) + # Per-sample measured sparsity (one row per calibration sample: its + # skipped-tile fraction at every threshold). Printed before the fit- + # validity guard so the raw per-sample data is visible even when the fit + # bails (e.g. degenerate near-zero sparsity). + self._print_per_sample_sparsity(per_sample_stats, phase) + if len(all_data_points) < 10: warnings.warn( f"Not enough data points for {phase} calibration. " @@ -286,11 +325,32 @@ def exponential(sparsity, a, b): "fit_logspace": self.fit_logspace, "min_observed_sparsity": min_observed_sparsity, "max_observed_sparsity": max_observed_sparsity, + # Raw per-sample measured sparsity, so callers can audit the spread + # across samples (not just the fitted average). + "per_sample_sparsity": [ + { + "sample_length": s.get("sample_length", 0), + "sparsity": list(s.get("sparsity", [])), + } + for s in per_sample_stats + ], } if self.fit_logspace: result["log_a"] = float(log_a) return result + def _print_per_sample_sparsity(self, per_sample_stats: list[dict], phase: str) -> None: + """Print each sample's measured skipped-tile fraction at every threshold.""" + if not per_sample_stats: + return + print(f"\nPer-sample {phase} sparsity (skipped-tile fraction per threshold):") + header = " ".join(f"{t:>7.0e}" for t in self.threshold_trials) + print(f" {'sample':>6} {'length':>8} {header}") + for idx, stat in enumerate(per_sample_stats): + sparsity = stat.get("sparsity", []) + row = " ".join(f"{s:>7.2%}" for s in sparsity) + print(f" {idx:>6} {stat.get('sample_length', 0):>8} {row}") + def _enable_calibration_mode(self, modules: list[nn.Module]): """Enable calibration mode on sparse attention modules.""" for idx, module in enumerate(modules): 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/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index 8c41b895a8a..b2acd137930 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py @@ -39,6 +39,29 @@ ) +def export_threshold_scale_factor(calibration_params: dict[str, Any]) -> dict[str, Any]: + """Build the canonical per-phase ``threshold_scale_factor`` export block. + + Single source of the exported skip-softmax schema fragment, shared by the + HF exporter (:func:`export_sparse_attention_config`) and the vLLM + calibration path (``plugins.sparse_attn_calibration``), so the serving + loader sees one format regardless of which path produced the checkpoint. + """ + block: dict[str, Any] = {"formula": "a * exp(b * target_sparsity)"} + for phase in ("prefill", "decode"): + if phase in calibration_params: + block[phase] = { + "a": float(calibration_params[phase]["a"]), + "b": float(calibration_params[phase]["b"]), + } + return block + + +def export_config_producer() -> dict[str, str]: + """Build the canonical ``producer`` block of ``sparse_attention_config``.""" + return {"name": "modelopt", "version": mo_version} + + def _set_attn_implementation(model: nn.Module, config: SparseAttentionConfig) -> None: """Set the correct attn_implementation based on the sparse attention method/backend. @@ -469,16 +492,7 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: skip_group["initial_disabled_steps"] = initial_disabled_steps # threshold_scale_factor (a * exp(b * target_sparsity)) and target_sparsity are # skip-softmax-specific, so they live in this group. - threshold_scale_factor: dict[str, Any] = { - "formula": "a * exp(b * target_sparsity)", - } - for phase in ["prefill", "decode"]: - if phase in calibration_params: - threshold_scale_factor[phase] = { - "a": calibration_params[phase]["a"], - "b": calibration_params[phase]["b"], - } - skip_group["threshold_scale_factor"] = threshold_scale_factor + skip_group["threshold_scale_factor"] = export_threshold_scale_factor(calibration_params) if target_sparse_ratio is not None: skip_group["target_sparsity"] = target_sparse_ratio config_groups[f"group_{group_idx}"] = skip_group @@ -495,10 +509,7 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: return { "config_groups": config_groups, - "producer": { - "name": "modelopt", - "version": mo_version, - }, + "producer": export_config_producer(), } diff --git a/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py b/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py index c1d6465ba66..baefbd7058d 100644 --- a/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py +++ b/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py @@ -193,6 +193,16 @@ def calc_correction_factor_and_p( dense_blocks_list = [] block_mask_0 = None block_diff = block_max - block_max_cummax + # Exclude padded query rows from the keep decision. _reshape_to_blocks + # pads the last block row with ``dtype.min``; a fully-padded row then + # has ``block_diff == 0`` (min - min), which passes ``> log_threshold`` + # and forces every block in the last partial block row to be kept + # (never skipped) — under-counting sparsity by up to one block row. + # Mask those rows to -inf so they vote "skip", matching the Triton + # kernel, which drops padding rows from its tile-skip reduction. + pad_q = padded_seq_q - seq_q + if pad_q > 0: + block_diff[:, :, -1, self.br - pad_q :, :] = float("-inf") for i, log_threshold in enumerate(log_thresholds): block_mask = (block_diff > log_threshold).any(dim=-2) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py new file mode 100644 index 00000000000..e0bd2250d87 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -0,0 +1,286 @@ +# 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. + +"""vLLM-free helpers for skip-softmax calibration through a serving engine. + +The serving adapters (``plugins/vllm.py``) record **raw per-threshold tile +counts** per scheduled request. These helpers merge those counts — across the +layers of one rank and across tensor-parallel ranks — then fit the exponential +threshold model and build the canonical ``sparse_attention_config`` block. + +Counts are additive, so aggregation is a plain sum: every layer of a rank and +every TP rank observes the same launches in the same order (TP ranks each see +their head shard), which makes records align by index within a phase. Fitting +happens once, per phase, on the globally merged counts — never per rank +(head-sharded counts are incomplete) and never by averaging independently +fitted coefficients (the fit is nonlinear). + +Everything here operates on plain Python data and is unit-testable without +vLLM installed. +""" + +from typing import Any + +# One canonical sweep for both calibration paths: re-exported from the HF-path +# calibrator so the vLLM path fits on the identical trial grid. +from ..calibration.calibrator import DEFAULT_THRESHOLD_TRIALS, DynamicThresholdCalibrator + +# One canonical schema for both calibration paths: the exported skip-softmax +# blocks come from the HF exporter's helpers so the formats cannot drift. +from ..conversion import export_config_producer, export_threshold_scale_factor + +__all__ = [ + "DEFAULT_THRESHOLD_TRIALS", + "build_sparse_attention_config", + "fit_from_counts", + "merge_count_records", + "merge_phase_counts", + "split_records_by_phase", + "stats_from_counts", +] + +_PHASES = ("prefill", "decode") + + +def split_records_by_phase(records: list[dict]) -> dict[str, list[dict]]: + """Group one impl's ordered calibration records by phase, preserving order.""" + per_phase: dict[str, list[dict]] = {phase: [] for phase in _PHASES} + for record in records: + per_phase.setdefault(record["phase"], []).append(record) + return per_phase + + +def merge_count_records(sources: list[list[dict]]) -> list[dict]: + """Element-wise sum aligned raw-count records from multiple sources. + + ``sources`` is a list over sources — the layers of one rank, or the + already layer-merged records of each TP rank — where each source is an + ordered list of ``{"sample_length", "total_tiles", "skipped_tiles"}`` + records for one phase. All sources observe the same launches in the same + order, so records align by index; tile counts are additive across both + layers and head-sharded TP ranks. + + Alignment is a contract: every source must report the same number of + samples, the same per-sample lengths, and the same threshold-vector width. + Any mismatch indicates a collection bug and raises rather than silently + dropping records. + """ + if not sources: + return [] + sample_counts = {len(source) for source in sources} + if len(sample_counts) != 1: + raise ValueError( + "Misaligned calibration records: sources disagree on sample count " + f"({sorted(sample_counts)})" + ) + num_samples = sample_counts.pop() + merged = [] + for i in range(num_samples): + base = sources[0][i] + width = len(base["total_tiles"]) + total = [0] * width + skipped = [0] * width + for record in (source[i] for source in sources): + if record["sample_length"] != base["sample_length"]: + raise ValueError( + "Misaligned calibration records: sample lengths differ across " + f"sources at index {i} ({record['sample_length']} vs " + f"{base['sample_length']})" + ) + if len(record["total_tiles"]) != width or len(record["skipped_tiles"]) != width: + raise ValueError( + "Misaligned calibration records: threshold-vector widths differ " + f"across sources at index {i} " + f"({len(record['total_tiles'])}/{len(record['skipped_tiles'])} vs {width})" + ) + total = [a + b for a, b in zip(total, record["total_tiles"])] + skipped = [a + b for a, b in zip(skipped, record["skipped_tiles"])] + merged.append( + { + "sample_length": base["sample_length"], + "total_tiles": total, + "skipped_tiles": skipped, + } + ) + return merged + + +def merge_phase_counts( + rank_counts: list[dict[str, list[dict]]], *, source_desc: str = "rank" +) -> dict[str, list[dict]]: + """Merge per-phase raw-count records collected from every TP rank. + + ``rank_counts`` is the list of per-rank results (one + ``{"prefill": [...], "decode": [...]}`` dict per rank, as returned by + ``collect_calibration_counts``). Use ALL ranks: with tensor parallelism + each rank only measures its attention-head shard, so any single rank's + counts are incomplete. A phase recorded by some ranks but not others + indicates a collection bug and raises. The same merge also sums one + rank's per-layer splits (``collect_calibration_counts`` delegates here + with ``source_desc="attention layer"``). + """ + phases = {phase for rank in rank_counts for phase in rank} + merged: dict[str, list[dict]] = {} + for phase in phases: + sources = [rank.get(phase, []) for rank in rank_counts] + empty = sum(1 for source in sources if not source) + if empty and empty != len(sources): + raise ValueError( + f"Misaligned calibration records: {empty}/{len(sources)} {source_desc}(s) " + f"recorded no {phase!r} samples while others did" + ) + # All-empty sources merge to [] (merge_count_records sums zero samples). + merged[phase] = merge_count_records(sources) + return merged + + +def stats_from_counts(count_records: list[dict]) -> list[dict]: + """Convert merged raw-count records into per-sample sparsity-ratio stats. + + Returns ``{"sample_length", "sparsity"}`` records in the shape + :meth:`DynamicThresholdCalibrator.calibrate_from_stats` consumes. + """ + stats = [] + for record in count_records: + sparsity = [ + (skipped / total if total else 0.0) + for skipped, total in zip(record["skipped_tiles"], record["total_tiles"]) + ] + stats.append({"sample_length": record["sample_length"], "sparsity": sparsity}) + return stats + + +def fit_from_counts( + per_phase_counts: dict[str, list[dict]], + threshold_trials: list[float], + *, + fit_logspace: bool = False, +) -> dict[str, dict[str, float]]: + """Fit the exponential skip-softmax model from globally merged counts. + + Reuses :class:`DynamicThresholdCalibrator` so vLLM-calibrated ``(a, b)`` + are identical in form to the HF path and export unchanged via + ``threshold_scale_factor``. One fit per phase, on counts already merged + across all TP ranks and layers. + + Returns: + ``{phase: {"a", "b", "min_observed_sparsity", "max_observed_sparsity"}}`` + for each phase that produced a valid fit. + """ + calibration_params: dict[str, dict[str, float]] = {} + for phase, records in per_phase_counts.items(): + if not records: + continue + for record in records: + if len(record["total_tiles"]) != len(threshold_trials): + raise ValueError( + f"{phase} record has {len(record['total_tiles'])} counters but " + f"{len(threshold_trials)} threshold trials are configured" + ) + calibrator = DynamicThresholdCalibrator( + threshold_trials=list(threshold_trials), fit_logspace=fit_logspace + ) + result = calibrator.calibrate_from_stats(stats_from_counts(records), phase=phase) + if "a" in result and "b" in result: + params = {"a": result["a"], "b": result["b"]} + for key in ("min_observed_sparsity", "max_observed_sparsity"): + if key in result: + params[key] = result[key] + calibration_params[phase] = params + return calibration_params + + +def _normalize_target_sparsity(target_sparsity: dict[str, float] | float) -> dict[str, float]: + if isinstance(target_sparsity, int | float): + values = {phase: float(target_sparsity) for phase in _PHASES} + else: + values = {phase: float(target_sparsity.get(phase, 0.5)) for phase in _PHASES} + for phase, value in values.items(): + # Same range the HF calibration config enforces. + if not 0.0 <= value <= 1.0: + raise ValueError( + f"target_sparsity for phase {phase!r} must be between 0.0 and 1.0, got {value}" + ) + return values + + +def build_sparse_attention_config( + calibration_params: dict[str, dict[str, float]], + target_sparsity: dict[str, float] | float = 0.5, + *, + existing_config: dict | None = None, +) -> dict[str, Any]: + """Build the canonical ``sparse_attention_config`` block for a checkpoint. + + Emits the same schema as + ``modelopt.torch.sparsity.attention_sparsity.conversion.export_sparse_attention_config`` + — a ``config_groups`` entry with ``algorithm: skip_softmax`` holding the + group-local ``threshold_scale_factor`` and ``target_sparsity`` — so + ``load_from_checkpoint_metadata`` (the serving loader) round-trips it + without changes. + + Non-skip groups from ``existing_config`` (e.g. exported N:M + ``sparse_softmax`` metadata) are preserved after the skip group; an + existing ``skip_softmax`` group is replaced by the new calibration, + carrying over its layer policy (``ignore`` — layers deliberately kept + dense — and ``initial_disabled_steps``): recalibration replaces the + fitted thresholds, not which layers the export sparsifies. + """ + # target_sparsity covers only fitted phases: claiming a target for a phase + # without calibrated (a, b) would advertise sparsity the serving path + # silently serves dense (it needs the per-phase scale factors). + target_sparsity_by_phase = { + phase: value + for phase, value in _normalize_target_sparsity(target_sparsity).items() + if phase in calibration_params + } + + skip_group: dict[str, Any] = { + "algorithm": "skip_softmax", + "targets": ["Attention"], + "threshold_scale_factor": export_threshold_scale_factor(calibration_params), + "target_sparsity": target_sparsity_by_phase, + } + + config_groups: dict[str, Any] = {"group_0": skip_group} + existing_groups = (existing_config or {}).get("config_groups") + if isinstance(existing_groups, dict): + preserved = [] + for group in existing_groups.values(): + if not isinstance(group, dict): + continue + if group.get("algorithm") == "skip_softmax": + # Keep the replaced group's layer policy: dropping ``ignore`` + # would sparsify layers the original export deliberately kept + # dense (e.g. first/last blocks). + for key in ("ignore", "initial_disabled_steps"): + if key in group and key not in skip_group: + skip_group[key] = group[key] + else: + preserved.append(group) + for idx, group in enumerate(preserved, start=1): + config_groups[f"group_{idx}"] = group + + result: dict[str, Any] = { + "config_groups": config_groups, + "producer": export_config_producer(), + } + # Legacy checkpoints carry N:M parameters as a top-level ``sparse_softmax`` + # dict (read by the serving loader ahead of group params) — preserve it so + # recalibration does not silently reset N:M settings to defaults. + legacy_sparse_softmax = (existing_config or {}).get("sparse_softmax") + if isinstance(legacy_sparse_softmax, dict): + result["sparse_softmax"] = legacy_sparse_softmax + return result diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 243db16a2bc..a6d3abffa4e 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -28,6 +28,7 @@ """ import functools +import importlib import inspect import math import warnings @@ -45,6 +46,18 @@ ) from modelopt.torch.kernels.common.attention.triton_fa import attention as triton_attention from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite +from modelopt.torch.kernels.sparsity.attention.calibrate import attention_calibrate + +from .sparse_attn_calibration import merge_phase_counts, split_records_by_phase + +__all__ = [ + "ModelOptSparseAttentionBackend", + "ModelOptSparseAttentionImpl", + "collect_calibration_counts", + "disable_calibration", + "enable_calibration", + "iter_sparse_impls", +] @functools.cache @@ -252,6 +265,161 @@ def _resolve_forward( ) +def _calibration_active(impl) -> bool: + """Return whether skip-softmax calibration mode is enabled on an impl.""" + return bool(getattr(impl, "_calibrate", False)) and bool( + getattr(impl, "_calib_threshold_trials", None) + ) + + +def _flashinfer_kv_cache_layout() -> str | None: + """Best-effort query of vLLM's configured FlashInfer KV-cache layout. + + Returns ``"NHD"`` / ``"HND"`` when the running vLLM exposes the layout + (newer releases select HND for some Blackwell FlashInfer paths), or + ``None`` when it cannot be determined — callers then fall back to the + shape-based check in :func:`_forward_calibrate`. + """ + for module_name in ( + "vllm.v1.attention.backends.utils", + "vllm.attention.backends.utils", + ): + try: + module = importlib.import_module(module_name) + except ImportError: + continue + getter = getattr(module, "get_kv_cache_layout", None) + if getter is None: + continue + try: + value = getter() + except Exception: + return None + # Preserve a genuine None (layout unset) so the shape fallback runs; + # str(None) would become the truthy string "None" and hard-reject. + return None if value is None else str(value) + return None + + +def _forward_calibrate( + impl, + *, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + num_actual_tokens: int, + output: torch.Tensor, +) -> torch.Tensor: + """Measure per-request tile-skip stats via the paged Triton calibration kernel. + + Each scheduled request is calibrated independently (batch=1) so its KV + length is the per-sample length the exponential fit needs, and so the + kernel keeps the uniform-length contract it was validated against. The + kernel computes full attention, so ``output`` is written densely — no + sparsification is applied to generation (the dense Triton kernel's + numerics may differ slightly from the native backend's). + + Phase and causality are decided per request: ``q_len > 1`` is (chunked) + prefill (causal — the kernel offsets the query into the KV span). A + ``q_len == 1`` row is a decode step (full-cache, non-causal) only when + its KV span exceeds the request's prompt (at least one generated token); + a 1-token row still inside the prompt is the final chunk of a chunked + prefill and is recorded as prefill. Prompt lengths come from the runner's + input batch (the installer attaches the runner as ``_calib_model_runner``; + the input batch is resolved per forward because vLLM can rebuild it after + install, same request order as the metadata rows); without it, + ``q_len == 1`` falls back to decode. A mixed prefill/decode batch + therefore contributes correctly to both phase fits. + + Records raw per-threshold tile counts (not ratios) on + ``impl._calib_records`` so tensor-parallel workers can be aggregated by + summing counts before the fit. + """ + if key_cache.dtype not in (torch.float16, torch.bfloat16): + raise NotImplementedError( + f"skip-softmax calibration requires an fp16/bf16 KV cache, got {key_cache.dtype}" + ) + if key_cache.shape[2] != impl.num_kv_heads: + # NHD is the only supported paged layout: [blocks, page, kv_heads, dim]. + # An HND FlashInfer cache would put kv_heads on axis 1. + raise NotImplementedError( + f"KV cache layout is not NHD (shape {tuple(key_cache.shape)}, " + f"expected axis 2 == num_kv_heads == {impl.num_kv_heads}); " + "HND caches are unsupported for calibration" + ) + page_size = key_cache.shape[1] + trials = impl._calib_threshold_trials + batch = seq_lens.shape[0] + # Hoist per-request tensors out of the loop: kernel args are sliced views + # of these, so the loop performs no allocations or casts. + b_seq_len_i32 = (cu_seqlens_q[1 : batch + 1] - cu_seqlens_q[:batch]).to(torch.int32) + seq_lens_i32 = seq_lens[:batch].to(torch.int32) + b_start_loc_zero = torch.zeros(1, device=query.device, dtype=torch.int32) + # vLLM v1 metadata tensors are GPU-resident: take one host copy per launch + # instead of stalling the stream with per-request ``.item()`` syncs. + cu_seqlens_q_cpu = cu_seqlens_q[: batch + 1].cpu() + seq_lens_cpu = seq_lens[:batch].cpu() + # Per-request prompt lengths (same request order as the metadata rows) + # distinguish decode steps from 1-token final chunks of a chunked + # prefill. Resolved from the runner per forward: vLLM can replace + # input_batch after install (KV-cache init for hybrid models). + input_batch = getattr(getattr(impl, "_calib_model_runner", None), "input_batch", None) + + q = query[:num_actual_tokens].contiguous() + # Dummy K/V: in paged mode KV is read from the cache via block_table. + # Only shape[1] (num_kv_heads) is consulted, to compute the GQA ratio. + k_dummy = torch.empty(0, impl.num_kv_heads, impl.head_size, device=q.device, dtype=q.dtype) + + for i in range(batch): + q_start = int(cu_seqlens_q_cpu[i]) + q_len = int(cu_seqlens_q_cpu[i + 1]) - q_start + if q_len <= 0: + continue + seq_k = int(seq_lens_cpu[i]) + if q_len > 1: + phase = "prefill" + elif input_batch is not None and seq_k <= int(input_batch.num_prompt_tokens[i]): + # 1-token final chunk of a chunked prefill: still inside the prompt. + phase = "prefill" + else: + phase = "decode" + + oi, counters = attention_calibrate( + q[q_start : q_start + q_len], + k_dummy, + k_dummy, + b_start_loc=b_start_loc_zero, + b_seq_len=b_seq_len_i32[i : i + 1], + max_input_len=q_len, + is_causal=q_len > 1, + softmax_scale=impl.scale, + b_seq_len_k=seq_lens_i32[i : i + 1], + max_input_len_k=seq_k, + threshold_trials=trials, + k_cache=key_cache, + v_cache=value_cache, + block_table=block_table[i : i + 1], + page_size=page_size, + ) + output[q_start : q_start + q_len] = oi + + # One host transfer for both counter columns (counters is GPU-resident). + counters_cpu = counters.cpu() + impl._calib_records.append( + { + "phase": phase, + "sample_length": seq_k, + "total_tiles": counters_cpu[:, 0].tolist(), + "skipped_tiles": counters_cpu[:, 1].tolist(), + } + ) + + return output + + # Resolution guards raw configured transforms; dispatch rechecks effective # sparse work after calibration and decode-only pruning. def _forward_modelopt( @@ -388,6 +556,8 @@ def _dispatch_modelopt( num_prefills: int, num_decode_tokens: int, num_prefill_tokens: int, + max_seq_len_decode: int | None = None, + max_seq_len_prefill: int | None = None, **common_kw, ) -> torch.Tensor: """Run the ModelOpt path, splitting mixed decode+prefill batches by phase. @@ -397,8 +567,14 @@ def _dispatch_modelopt( ``q_len==1`` decode rows with ``q_len>1`` (chunked-)prefill rows, ``max_query_len > 1`` and the whole batch would otherwise take the prefill skip-softmax path. Split so each phase runs its own schedule -- decode rows - always take the fixed decode path. Both the FlashAttention and FlashInfer - adapters share this dispatch. + always take the fixed decode path. + + Both adapters route through this dispatch, but the split is live only on + FlashInfer: its metadata carries the ``num_decodes``/``num_prefills`` + counts (and vLLM reorders those batches decode-first). vLLM's + FlashAttention metadata has no phase counts, so FA mixed batches fall + through to the whole-batch path and are classified by ``max_query_len`` + alone (decode rows then follow the prefill contract for that launch). """ if not (num_decodes and num_prefills): return _forward_modelopt( @@ -424,6 +600,18 @@ def _dispatch_modelopt( if not common_kw.get("quant_active", False): common_kw["dense_fallback"]() + # Each phase derives its skip threshold from its own KV maximum: reusing + # the batch-global max_seq_len (e.g. a co-scheduled 32k prefill next to 2k + # decodes) would shrink the decode threshold far below — much denser than + # — the calibrated target. Fall back to the batch-global value only when + # the builder did not provide per-phase maxima. + decode_kw = dict(common_kw) + if max_seq_len_decode is not None: + decode_kw["max_seq_len"] = max_seq_len_decode + prefill_kw = dict(common_kw) + if max_seq_len_prefill is not None: + prefill_kw["max_seq_len"] = max_seq_len_prefill + _forward_modelopt( impl, query=query[:num_decode_tokens], @@ -433,7 +621,7 @@ def _dispatch_modelopt( num_actual_tokens=num_decode_tokens, max_query_len=num_decode_tokens // num_decodes, output=output[:num_decode_tokens], - **common_kw, + **decode_kw, ) prefill_start = num_decode_tokens prefill_cu_seqlens_q = cu_seqlens_q[num_decodes:] - cu_seqlens_q[num_decodes] @@ -446,7 +634,7 @@ def _dispatch_modelopt( num_actual_tokens=num_prefill_tokens, max_query_len=max_query_len, output=output[prefill_start : prefill_start + num_prefill_tokens], - **common_kw, + **prefill_kw, ) return output @@ -517,6 +705,26 @@ def native_forward(): ) return native_result + if _calibration_active(self): + if getattr(attn_metadata, "use_cascade", False): + # Cascade splits shared prefixes across requests, so per-request + # KV lengths are unavailable; skip measurement for this launch. + return native_forward() + # vLLM >= 0.15 writes the current K/V to the paged cache before + # impl.forward, so the calibrate kernel reads a complete cache. + key_cache, value_cache = kv_cache.unbind(0) + return _forward_calibrate( + self, + query=query, + key_cache=key_cache, + value_cache=value_cache, + block_table=attn_metadata.block_table, + seq_lens=attn_metadata.seq_lens, + cu_seqlens_q=attn_metadata.query_start_loc, + num_actual_tokens=attn_metadata.num_actual_tokens, + output=output, + ) + resolved = _resolve_forward( self, layer, @@ -628,6 +836,23 @@ def build(*args, **kwargs): common = build_sig.bind(*args, **kwargs).arguments["common_attn_metadata"] for target, source in _FLASHINFER_METADATA_FIELDS.items(): setattr(metadata, target, getattr(common, source)) + # Per-phase KV maxima for the mixed-batch split (batch is reordered + # decode-first): computed once per build — not per layer forward — so + # the split's threshold derivation neither reuses the batch-global max + # nor syncs the stream inside every layer. + num_decodes = getattr(metadata, "num_decodes", 0) + num_prefills = getattr(metadata, "num_prefills", 0) + max_seq_len_decode = max_seq_len_prefill = None + if num_decodes and num_prefills: + # Prefer the host-resident copy the runner may already carry; + # fall back to one device->host copy per mixed-batch build. + seq_lens_cpu = getattr(common, "_seq_lens_cpu", None) + if seq_lens_cpu is None: + seq_lens_cpu = common.seq_lens.cpu() + max_seq_len_decode = int(seq_lens_cpu[:num_decodes].max()) + max_seq_len_prefill = int(seq_lens_cpu[num_decodes : num_decodes + num_prefills].max()) + metadata._modelopt_max_seq_len_decode = max_seq_len_decode + metadata._modelopt_max_seq_len_prefill = max_seq_len_prefill return metadata setattr(build, "_modelopt_sparse_metadata_patch", True) @@ -705,6 +930,44 @@ def prepare_modelopt(): _maybe_update_flashinfer_cache(layer, key, value, kv_cache, attn_metadata, impl) cache_prepared = True + if _calibration_active(impl): + if getattr(attn_metadata, "use_cascade", False): + # Cascade splits shared prefixes across requests, so per-request + # KV lengths are unavailable; skip measurement for this launch. + return dense_fallback() + missing = [name for name in _FLASHINFER_METADATA_FIELDS if not hasattr(attn_metadata, name)] + if missing: + raise NotImplementedError( + "FlashInfer metadata is missing the ModelOpt calibration " + f"fields: {', '.join(missing)}" + ) + if kv_cache.ndim != 5 or kv_cache.shape[1] != 2: + raise ValueError( + "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" + ) + layout = _flashinfer_kv_cache_layout() + if layout is not None and layout.upper() != "NHD": + # Authoritative layout metadata beats the shape heuristic (which is + # ambiguous when page_size equals the per-rank KV-head count). + raise NotImplementedError( + f"FlashInfer KV-cache layout {layout!r} is unsupported for " + "skip-softmax calibration; only NHD is supported" + ) + # Order matters: releases that update the KV cache inside forward must + # write the current K/V before the calibrate kernel reads the cache. + prepare_modelopt() + return _forward_calibrate( + impl, + query=query, + key_cache=kv_cache[:, 0], + value_cache=kv_cache[:, 1], + block_table=attn_metadata._modelopt_block_table, + seq_lens=attn_metadata._modelopt_seq_lens, + cu_seqlens_q=attn_metadata._modelopt_query_start_loc, + num_actual_tokens=attn_metadata._modelopt_num_actual_tokens, + output=output, + ) + resolved = _resolve_forward( impl, layer, @@ -752,6 +1015,8 @@ def prepare_modelopt(): num_prefills=getattr(attn_metadata, "num_prefills", 0), num_decode_tokens=getattr(attn_metadata, "num_decode_tokens", 0), num_prefill_tokens=getattr(attn_metadata, "num_prefill_tokens", 0), + max_seq_len_decode=getattr(attn_metadata, "_modelopt_max_seq_len_decode", None), + max_seq_len_prefill=getattr(attn_metadata, "_modelopt_max_seq_len_prefill", None), **common_kw, ) @@ -833,3 +1098,58 @@ def _clone_sparse_impl(old_impl, new_cls=None): new_impl = object.__new__(new_cls) new_impl.__dict__.update(old_state) return new_impl + + +def iter_sparse_impls(model): + """Yield every ModelOpt sparse attention impl reachable from a vLLM model. + + Walks ``model.named_modules()`` and returns the swapped ``impl`` of each + attention layer (FlashAttention or FlashInfer adapter). Used by the + calibration installer and RPC methods to toggle calibration mode and + harvest stats without knowing vLLM's module layout. + """ + for _, module in model.named_modules(): + impl = getattr(module, "impl", None) + if impl is None: + continue + if isinstance(impl, ModelOptSparseAttentionImpl) or ( + _FLASHINFER_IMPL_CLS is not None and isinstance(impl, _FLASHINFER_IMPL_CLS) + ): + yield impl + + +def enable_calibration(impls, threshold_trials: list[float]) -> None: + """Put a set of sparse impls into calibration mode and clear prior records.""" + if not threshold_trials: + raise ValueError("threshold_trials must be a non-empty list for calibration.") + for impl in impls: + impl._calibrate = True + impl._calib_threshold_trials = list(threshold_trials) + impl._calib_records = [] + + +def disable_calibration(impls) -> None: + """Turn off calibration mode (collected records are left intact).""" + for impl in impls: + impl._calibrate = False + + +def collect_calibration_counts(model) -> dict[str, list[dict]]: + """Harvest one rank's raw per-phase tile counts from every calibrating impl. + + Sums counts across the rank's layers per aligned sample (every layer sees + the same launches in the same order), keeping raw + ``{"sample_length", "total_tiles", "skipped_tiles"}`` records per phase. + The driver merges these across TP ranks with + :func:`~.sparse_attn_calibration.merge_phase_counts` and fits once per + phase with :func:`~.sparse_attn_calibration.fit_from_counts` — sparsity + ratios are only formed after the global merge. + """ + splits = [ + split_records_by_phase(getattr(impl, "_calib_records", [])) + for impl in iter_sparse_impls(model) + ] + # Same merge as the cross-rank aggregation: every layer sees every launch, + # so a layer with no records for a phase others measured indicates a + # collection bug (merge_phase_counts raises). + return merge_phase_counts(splits, source_desc="attention layer") diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index b4141740b64..6e2ad911288 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -32,6 +32,7 @@ __all__ = [ "VllmAttentionInstallReport", "install_vllm_nvfp4_attention", + "install_vllm_skip_softmax_calibration", "install_vllm_sparse_attention_from_checkpoint", ] @@ -155,7 +156,7 @@ def _cudagraph_mode(model_runner): return mode if mode is not None else CUDAGraphMode.NONE -def _global_errors(model_runner) -> list[str]: +def _global_errors(model_runner, *, sparse_only: bool = False) -> list[str]: config = getattr(model_runner, "vllm_config", None) if config is None: return ["model_runner.vllm_config is required"] @@ -173,10 +174,16 @@ def _global_errors(model_runner) -> list[str]: errors.append("decode_context_parallel_size must be 1") if getattr(parallel, "enable_dbo", False) or getattr(parallel, "use_ubatching", False): errors.append("DBO/ubatching is unsupported") - if getattr(cache_config, "enable_prefix_caching", False): - errors.append("prefix caching is unsupported") - if getattr(config, "kv_transfer_config", None) is not None: - errors.append("KV transfer is unsupported") + if not sparse_only: + # Prefix caching and KV transfer break only flows that quantize the + # cache on write or measure per-request prefills (quantized installs, + # skip-softmax calibration). Sparse-only serving reads the cache + # unmodified and supports prefix-cache suffix attention by offsetting + # query positions (see the vllm_serve README limitations). + if getattr(cache_config, "enable_prefix_caching", False): + errors.append("prefix caching is unsupported") + if getattr(config, "kv_transfer_config", None) is not None: + errors.append("KV transfer is unsupported") if getattr(config, "speculative_config", None) is not None: errors.append("speculative decoding is unsupported") if _cudagraph_mode(model_runner).mixed_mode() == CUDAGraphMode.FULL: @@ -248,6 +255,13 @@ def _device_capability_error(device) -> str | None: return None +def _skip_softmax_active(sparse_kw: dict[str, Any]) -> bool: + """Return whether a layer's sparse config contains skip-softmax work.""" + return bool(sparse_kw) and ( + "skip_softmax_threshold" in sparse_kw or "threshold_scale_factor" in sparse_kw + ) + + def _sparse_graph_error(sparse_kw: dict[str, Any], mode) -> str | None: from vllm.config.compilation import CUDAGraphMode @@ -347,12 +361,35 @@ def _plan_vllm_attention( ) _require_supported_vllm() - errors = _global_errors(model_runner) if quantize else [] - mode = _cudagraph_mode(model_runner) if quantize else None + # Engine-level checks apply to sparse-only installs too: the ModelOpt + # kernel path silently ignores decode context parallelism, DBO, and + # speculative decoding, and FULL mixed-batch graphs would capture stale + # per-launch thresholds (same rationale as the decode graph guard below). + # Sparse-only installs skip only the cache-mutation checks. + errors = _global_errors(model_runner, sparse_only=not quantize) + mode = _cudagraph_mode(model_runner) quant_plugin: Any = _load_quant_plugin() if quantize else None plans = [] for name, module, sparse_kw in candidates: reasons = _layer_errors(module) + if _skip_softmax_active(sparse_kw): + # Quantized Q/K/P change the attention-score distribution the skip + # thresholds were calibrated on, so the calibrated sparsity contract + # no longer holds. This guards both installation directions: quantized + # installs adding skip, and sparse-only installs onto layers that + # already carry active attention quantizers. N:M sparse softmax has + # no calibrated threshold and composes with quantization. + active = ( + "attention quantization is being installed" + if quantize + else _active_attention_quantization(module) + ) + if active: + reasons.append( + f"skip-softmax cannot be combined with attention quantization ({active}); " + "serve skip-softmax unquantized or drop the skip_softmax group " + "(N:M sparse softmax composes with quantization)" + ) device = dtype = None if quantize: device, dtype = quant_plugin._get_device_dtype(module) @@ -365,9 +402,12 @@ def _plan_vllm_attention( reasons.append(f"resolved dtype {dtype} must be fp16 or bf16") if capability_error := _device_capability_error(device): reasons.append(capability_error) - if quantize: - if graph_error := _sparse_graph_error(sparse_kw, mode): - reasons.append(graph_error) + # Calibrated decode skip-softmax replays through the decode kernel path, + # which a FULL decode CUDA graph would capture with a stale threshold. + # This holds for sparse-only installs exactly as for quantized ones, so + # the guard is not gated on ``quantize``. + if graph_error := _sparse_graph_error(sparse_kw, mode): + reasons.append(graph_error) new_impl, requires_flashinfer_patch, backend_error = _select_new_impl(module) if backend_error: reasons.append(backend_error) @@ -492,6 +532,97 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor return _build_report(plan) +def _active_attention_quantization(module) -> str | None: + """Describe any active attention Q/K/P/V quantization on a layer, or None.""" + for attr in ("q_bmm_quantizer", "k_bmm_quantizer", "p_bmm_quantizer", "v_bmm_quantizer"): + if getattr(getattr(module, attr, None), "is_enabled", False): + return f"{attr} is enabled" + if getattr(module, "_query_quant_in_kernel", False) or getattr( + module, "_value_quant_in_kernel", False + ): + return "in-kernel attention quantization flags are set" + return None + + +def _attention_quant_error(module) -> str | None: + """Reject calibration on layers with any active attention Q/K/P/V fakequant.""" + if active := _active_attention_quantization(module): + return f"{active}; skip-softmax calibration requires unquantized attention" + return None + + +def install_vllm_skip_softmax_calibration(model_runner) -> VllmAttentionInstallReport: + """Install skip-softmax calibration adapters into a loaded vLLM model. + + Swaps the backend-matched ModelOpt adapter onto every attention layer and + disables cascade attention, following validation-before-mutation: every + known compatibility error — across all layers — is collected and raised + before any module is changed. Calibration itself starts separately via + :func:`~.vllm.enable_calibration` (typically over a worker RPC), so engine + warmup/profiling launches after install are served natively and never + pollute the measurement; until then the adapters delegate every forward to + the backend's native implementation. + + Requirements validated here: eager execution (``enforce_eager=True`` — + the per-request calibration loop cannot be CUDA-graph captured), fp16/bf16 + model and KV-cache dtypes, no active attention Q/K/P/V fakequant, and a + FlashAttention or FlashInfer backend per layer. + """ + from vllm.config.compilation import CUDAGraphMode + + model = _unwrapped_model(model_runner) + candidates = [ + (name, module) + for name, module in model.named_modules() + if isinstance(module, _VLLM_ATTENTION) + ] + + _require_supported_vllm() + errors = _global_errors(model_runner) + if _cudagraph_mode(model_runner) != CUDAGraphMode.NONE: + errors.append( + "skip-softmax calibration requires eager execution (enforce_eager=True); " + "the per-request calibration loop cannot be CUDA-graph captured" + ) + if not candidates: + errors.append("no attention layers were found") + + plans = [] + for name, module in candidates: + reasons = _layer_errors(module) + if quant_error := _attention_quant_error(module): + reasons.append(quant_error) + new_impl, requires_flashinfer_patch, backend_error = _select_new_impl(module) + if backend_error: + reasons.append(backend_error) + if reasons: + errors.extend(f"{name or ''}: {reason}" for reason in reasons) + continue + plans.append( + _AttentionPlan(name, module, new_impl, {}, None, None, requires_flashinfer_patch) + ) + if any(plan.requires_flashinfer_patch for plan in plans): + layout = attention_plugin._flashinfer_kv_cache_layout() + if layout is not None and layout.upper() != "NHD": + errors.append( + f"FlashInfer KV-cache layout {layout!r} is unsupported for calibration (NHD only)" + ) + _raise_unsupported(errors, "skip-softmax calibration") + + plan = _InstallPlan(model_runner, tuple(plans), False, "SKIP_SOFTMAX_CALIBRATION") + # Per-request prompt lengths (same request order as the attention-metadata + # rows, kept aligned by vLLM's in-place batch reorder) let the adapter + # classify q_len == 1 rows: a decode step's KV span exceeds the prompt, + # while a 1-token final chunk of a chunked prefill is still inside it. + # The runner — not its input_batch — is attached: vLLM can rebuild + # input_batch after load_model (may_reinitialize_input_batch during KV + # cache init, e.g. hybrid Mamba/attention models), so the adapter must + # resolve the live object per forward. + for attention_plan in plans: + attention_plan.new_impl._calib_model_runner = model_runner + return _apply_vllm_attention_plans(plan) + + def install_vllm_sparse_attention_from_checkpoint( model_runner, ) -> VllmAttentionInstallReport: diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index 60523d44457..99e8a3a83a4 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -20,7 +20,7 @@ import pytest import torch -from conftest import make_qkv, make_varlen_meta, sdpa_reference +from conftest import make_qkv, make_varlen_meta from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor, e2m1_values @@ -454,9 +454,12 @@ def test_invalid_amax_raises(self): with pytest.raises(ValueError, match="p_qdq_amax"): attention(q, k, v, locs, lens, 8, p_qdq="fp8", p_qdq_amax=0.0) - @requires_native_e4m3 - def test_composes_with_skip_softmax(self): - """p_qdq composes with the skip-softmax feature in one launch.""" + def test_rejects_skip_softmax(self): + """p_qdq cannot combine with active skip-softmax in one launch. + + Quantized P changes the score distribution the skip thresholds were + calibrated on, so the kernel rejects the composition (pre-launch). + """ seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 scale = 1.0 / (head_dim**0.5) @@ -464,19 +467,18 @@ def test_composes_with_skip_softmax(self): q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.float16) locs, lens = make_varlen_meta([seq_len]) - o = attention( - q, - k, - v, - locs, - lens, - seq_len, - softmax_scale=scale, - p_qdq="fp8", - skip_softmax_threshold=1e-3, - ) - ref = sdpa_reference(q, k, v, locs, lens) - torch.testing.assert_close(o, ref, rtol=5e-2, atol=5e-2) + with pytest.raises(ValueError, match="cannot be combined with attention quantization"): + attention( + q, + k, + v, + locs, + lens, + seq_len, + softmax_scale=scale, + p_qdq="fp8", + skip_softmax_threshold=1e-3, + ) def test_invalid_mode_raises(self): q, k, v = make_qkv(8, 2, 2, 32, dtype=torch.float16) diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py new file mode 100644 index 00000000000..c7390382098 --- /dev/null +++ b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py @@ -0,0 +1,232 @@ +# 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. + +"""Paged-cache calibration kernel tests and the calibration/serving tile contract.""" + +import pytest +import torch +from conftest import make_qkv, make_varlen_meta + +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE + +if TRITON_KERNEL_AVAILABLE: + from modelopt.torch.kernels.common.attention import attention + from modelopt.torch.kernels.sparsity.attention.calibrate import attention_calibrate + +pytestmark = pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") + +TRIALS = [1e-3, 1e-2, 1e-1, 3e-1] + + +def _pack_paged(k, v, page_size, *, shuffle=True, num_spare_blocks=7): + """Pack one sequence's contiguous K/V into a (optionally shuffled) paged cache.""" + seq = k.shape[0] + num_blocks = (seq + page_size - 1) // page_size + order = torch.randperm(num_blocks) if shuffle else torch.arange(num_blocks) + k_cache = k.new_zeros(num_blocks + num_spare_blocks, page_size, k.shape[1], k.shape[2]) + v_cache = torch.zeros_like(k_cache) + block_table = torch.zeros(1, num_blocks, device=k.device, dtype=torch.int32) + for i in range(num_blocks): + page = int(order[i]) + num_spare_blocks # keep low pages unused + ts, te = i * page_size, min((i + 1) * page_size, seq) + k_cache[page, : te - ts] = k[ts:te] + v_cache[page, : te - ts] = v[ts:te] + block_table[0, i] = page + return k_cache, v_cache, block_table + + +class TestPagedCalibrate: + @pytest.mark.parametrize("seq_len", [256, 300, 512]) # 300: non-128-aligned padding + def test_paged_matches_contiguous_prefill(self, seq_len): + """Paged and contiguous calibration agree exactly on counters and output.""" + torch.manual_seed(0) + num_heads, num_kv_heads, head_dim, page_size = 8, 2, 64, 16 + q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + locs, lens = make_varlen_meta([seq_len]) + + out_ref, counters_ref = attention_calibrate( + q, k, v, locs, lens, seq_len, is_causal=True, threshold_trials=TRIALS + ) + + k_cache, v_cache, block_table = _pack_paged(k, v, page_size) + k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) + out_paged, counters_paged = attention_calibrate( + q, + k_dummy, + k_dummy, + locs, + lens, + seq_len, + is_causal=True, + threshold_trials=TRIALS, + b_seq_len_k=lens, + max_input_len_k=seq_len, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + page_size=page_size, + ) + + assert torch.equal(counters_ref, counters_paged) + torch.testing.assert_close(out_paged, out_ref, rtol=1e-3, atol=1e-3) + + def test_paged_decode_measures_full_cache(self): + """A one-row decode query measures every KV tile of the paged cache.""" + torch.manual_seed(1) + num_heads, num_kv_heads, head_dim, page_size = 8, 2, 64, 16 + ctx = 384 + q, k, v = make_qkv(ctx, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + k_cache, v_cache, block_table = _pack_paged(k, v, page_size) + k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) + locs = torch.zeros(1, device="cuda", dtype=torch.int32) + + _, counters = attention_calibrate( + q[:1], + k_dummy, + k_dummy, + locs, + torch.ones(1, device="cuda", dtype=torch.int32), + 1, + is_causal=False, + threshold_trials=TRIALS, + b_seq_len_k=torch.tensor([ctx], device="cuda", dtype=torch.int32), + max_input_len_k=ctx, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + page_size=page_size, + ) + + num_kv_tiles = -(-ctx // 128) + assert counters[:, 0].tolist() == [num_heads * num_kv_tiles] * len(TRIALS) + + def test_high_block_id_pointer_arithmetic(self): + """Block IDs whose int32 byte offsets would wrap still read correctly.""" + num_kv_heads, head_dim, page_size = 2, 64, 16 + block_elems = page_size * num_kv_heads * head_dim + # Smallest block ID whose element offset exceeds int32. V aliases the K + # cache storage (same values on both operands), halving the allocation. + high_block = (2**31) // block_elems + 1 + bytes_needed = (high_block + 1) * block_elems * 2 # one shared K/V cache, bf16 + free, _ = torch.cuda.mem_get_info() + if free < bytes_needed + (2 << 30): + pytest.skip(f"needs ~{bytes_needed / 2**30:.1f} GiB free GPU memory") + + torch.manual_seed(2) + num_heads = 4 + q, k, _ = make_qkv(page_size, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + k_cache = torch.zeros( + high_block + 1, page_size, num_kv_heads, head_dim, device="cuda", dtype=torch.bfloat16 + ) + v_cache = k_cache # alias: V reads the same storage (and the same values) + k_cache[high_block] = k + block_table = torch.tensor([[high_block]], device="cuda", dtype=torch.int32) + locs, lens = make_varlen_meta([page_size]) + + out_ref, counters_ref = attention_calibrate( + q, k, k, locs, lens, page_size, is_causal=True, threshold_trials=TRIALS + ) + k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) + out_paged, counters_paged = attention_calibrate( + q, + k_dummy, + k_dummy, + locs, + lens, + page_size, + is_causal=True, + threshold_trials=TRIALS, + b_seq_len_k=lens, + max_input_len_k=page_size, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + page_size=page_size, + ) + del k_cache, v_cache + + assert torch.equal(counters_ref, counters_paged) + torch.testing.assert_close(out_paged, out_ref, rtol=1e-3, atol=1e-3) + + +class TestCalibrationServingTileContract: + """Active skip launches and calibration must count identically (same tiles).""" + + def _contrasty_qkv(self, seq_len, num_heads, num_kv_heads, head_dim): + """K with a dominant head-of-sequence so later tiles are skippable.""" + torch.manual_seed(3) + q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + k = k * 0.05 + k[:32] = k[:32] * 600.0 # first tile dominates the running max by >> log2(threshold) + return q, k, v + + @pytest.mark.parametrize("threshold", [1e-3, 1e-2]) + def test_serve_skip_counts_equal_calibrate_counts(self, threshold): + seq_len, num_heads, num_kv_heads, head_dim = 512, 8, 2, 64 + q, k, v = self._contrasty_qkv(seq_len, num_heads, num_kv_heads, head_dim) + locs, lens = make_varlen_meta([seq_len]) + scale = 1.0 / (head_dim**0.5) + + _, counters = attention_calibrate( + q, + k, + v, + locs, + lens, + seq_len, + is_causal=True, + softmax_scale=scale, + threshold_trials=[threshold], + ) + + out = attention( + q, + k, + v, + locs, + lens, + seq_len, + is_causal=True, + softmax_scale=scale, + skip_softmax_threshold=threshold, + measure_sparsity=True, + ) + + calib_total, calib_skipped = int(counters[0, 0]), int(counters[0, 1]) + assert calib_skipped > 0, "test data must produce skippable tiles" + # Same 128x128 tile geometry and same prefix-max criterion => the serve + # kernel must skip exactly the tiles calibration predicted. + assert out._sparsity_total == calib_total + assert out._sparsity_skipped == calib_skipped + + @pytest.mark.parametrize("qdq_kw", [{"p_qdq": "nvfp4"}, {"v_qdq": "nvfp4", "v_qdq_amax": 1.0}]) + def test_skip_rejects_pv_qdq(self, qdq_kw): + """Active skip rejects P/V QDQ: quantized operands break the calibrated contract.""" + seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 + q, k, v = self._contrasty_qkv(seq_len, num_heads, num_kv_heads, head_dim) + locs, lens = make_varlen_meta([seq_len]) + + with pytest.raises(ValueError, match="cannot be combined with attention quantization"): + attention( + q, + k, + v, + locs, + lens, + seq_len, + is_causal=True, + skip_softmax_threshold=1e-2, + **qdq_kw, + ) diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py index 2806e39b4a5..598f97a3fda 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py @@ -356,11 +356,16 @@ class TestBackwardWithSparsity: """Backward pass with skip-softmax (covers _attn_bwd_dq / _attn_bwd_dkdv).""" def test_backward_with_skip_softmax(self): - """Backward pass runs without error when skip-softmax is active.""" + """Backward pass runs without error when skip-softmax is active. + + fp16 rather than fp32: active skip launches require the fixed 128x128 + calibration tile, which fp32 inputs cannot compile on ~100KB-shared- + memory GPUs (such configurations are rejected by design). + """ seq_len, num_heads, head_dim = 128, 4, 64 scale = 1.0 / (head_dim**0.5) torch.manual_seed(7) - q, k, v = make_qkv(seq_len, num_heads, num_heads, head_dim, dtype=torch.float32) + q, k, v = make_qkv(seq_len, num_heads, num_heads, head_dim, dtype=torch.float16) q.requires_grad_(True) k.requires_grad_(True) v.requires_grad_(True) diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py index fc26c5db17c..033952d34eb 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py @@ -217,17 +217,26 @@ def test_triton_matches_pytorch_reference(self): locs = torch.arange(batch, device="cuda", dtype=torch.int32) * seq_len lens = torch.full((batch,), seq_len, device="cuda", dtype=torch.int32) - triton_out = attention( - q_flat, - k_flat, - v_flat, - locs, - lens, - seq_len, - is_causal=True, - softmax_scale=scale, - skip_softmax_threshold=threshold, - ) + try: + triton_out = attention( + q_flat, + k_flat, + v_flat, + locs, + lens, + seq_len, + is_causal=True, + softmax_scale=scale, + skip_softmax_threshold=threshold, + ) + except RuntimeError as err: + if "shared memory" in str(err): + # Active skip requires the fixed 128x128 calibration tile; fp32 + # inputs cannot compile it on ~100KB-shared-memory GPUs and the + # configuration is rejected by design. fp32 is kept here for a + # tight reference comparison on GPUs that support it. + pytest.skip("fp32 skip tile exceeds this GPU's shared memory") + raise triton_out_4d = triton_out.view(batch, seq_len, num_heads, head_dim).permute(0, 2, 1, 3) # Both outputs should be close — same algorithm, different implementations. @@ -271,5 +280,8 @@ def test_skip_softmax_via_sparsify(self, tiny_llama_dir): assert not torch.isnan(logits_skip).any(), "NaN in skip-softmax logits" assert not torch.isinf(logits_skip).any(), "Inf in skip-softmax logits" - # On short sequences (64 tokens), no tiles are skipped — output should match dense - torch.testing.assert_close(logits_skip, logits_dense, rtol=1e-3, atol=1e-3) + # On short sequences (64 tokens), no tiles are skipped — output should match + # dense up to bf16 accumulation-order noise: active skip launches run on the + # fixed 128x128 calibration tile, so their summation order differs from the + # HF dense reference (and from the autotuned dense Triton tile). + torch.testing.assert_close(logits_skip, logits_dense, rtol=1e-2, atol=8e-3) diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index cf802abaf40..2e2a0105846 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -354,3 +354,36 @@ def test_static_vs_dynamic_fp4_kernels(self, set_torch_dtype, block_size, num_bl f"Mean abs diff: {(output_static - output_dynamic).abs().mean()}\n" f"Max relative diff: {((output_static - output_dynamic).abs() / (output_dynamic.abs() + 1e-8)).max()}" ) + + +def test_dynamic_block_quantize_forwards_block_size(monkeypatch): + """NVFP4 dynamic dispatch must pass the configured block size to the Triton kernel. + + The Triton kernel defaults to block_size=16; dropping the argument would + silently quantize non-16 NVFP4 configs with 16-element blocks while the + cuda_ext fallback honors the configured size (device-dependent numerics). + """ + recorded = {} + + def fake_fp4_fake_quant_block(inputs, amax, block_size=16): + recorded["block_size"] = block_size + return inputs.clone() + + monkeypatch.setattr(tensor_quant, "DISABLE_TRITON_KERNEL", False) + monkeypatch.setattr(triton_kernel, "IS_AVAILABLE", True, raising=False) + monkeypatch.setattr( + triton_kernel, "fp4_fake_quant_block", fake_fp4_fake_quant_block, raising=False + ) + + inputs = torch.randn(4, 64, device="cuda") + amax = inputs.abs().amax() + tensor_quant._dynamic_block_quantize_impl( + inputs, + 32, # non-default NVFP4 block size + amax, + 4, # num_bits total (E2M1) + 2, # exponent_bits -> (2, 1) + 8, # scale_num_bits -> (4, 3) + 4, # scale_exponent_bits + ) + assert recorded["block_size"] == 32 diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 578922db077..f87a8cbb26b 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -80,7 +80,11 @@ def guarded_import(name, *args, **kwargs): monkeypatch.setattr(builtins, "__import__", guarded_import) worker_module = _load_worker_module("sparse_attn_worker_import_test") - assert worker_module.__all__ == ["SparseAttnWorker", "QuantSparseAttnWorker"] + assert worker_module.__all__ == [ + "SparseAttnWorker", + "QuantSparseAttnWorker", + "SkipSoftmaxCalibWorker", + ] @pytest.mark.parametrize( diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py new file mode 100644 index 00000000000..d3563d226cf --- /dev/null +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -0,0 +1,542 @@ +# 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 skip-softmax calibration through the vLLM adapters and installer.""" + +import sys +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.attention.backends.flashinfer import FlashInferImpl + +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.quantization.plugins import vllm as quant_plugin +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as attention_plugin +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm_runtime +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( + ModelOptSparseAttentionImpl, + collect_calibration_counts, + disable_calibration, + enable_calibration, + iter_sparse_impls, +) + +TRIALS = [1e-3, 1e-1, 5e-1] + + +# --------------------------------------------------------------------------- +# Installer fakes (mirroring test_vllm_runtime.py) +# --------------------------------------------------------------------------- +def _bare_attention(impl_cls=FlashAttentionImpl): + module = object.__new__(vllm_runtime._VLLM_ATTENTION) + nn.Module.__init__(module) + module.attn_type = "decoder" + module.head_size = 64 + module.device = torch.device("cpu") + module.dtype = torch.float16 + module.impl = object.__new__(impl_cls) + module.impl.sinks = None + return module + + +def _model_runner(model, *, sparse_metadata=None, cudagraph_mode=CUDAGraphMode.NONE): + hf_config = SimpleNamespace(sparse_attention_config=sparse_metadata) + model_config = SimpleNamespace(hf_config=hf_config, dtype=torch.float16) + return SimpleNamespace( + model=model, + model_config=model_config, + cascade_attn_enabled=True, + vllm_config=SimpleNamespace( + model_config=model_config, + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + enable_dbo=False, + use_ubatching=False, + ), + cache_config=SimpleNamespace(enable_prefix_caching=False, cache_dtype="auto"), + compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), + kv_transfer_config=None, + speculative_config=None, + ), + ) + + +class TestCalibrationInstaller: + def test_installs_adapters_without_enabling_measurement(self): + first = _bare_attention() + second = _bare_attention() + runner = _model_runner(nn.ModuleDict({"a_attn": first, "b_attn": second})) + + report = vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + assert report.installed_count == 2 + assert report.sparse_algorithm == "SKIP_SOFTMAX_CALIBRATION" + assert report.cascade_disabled and runner.cascade_attn_enabled is False + for module in (first, second): + assert isinstance(module.impl, ModelOptSparseAttentionImpl) + # Measurement starts only via enable_calibration, so warmup + # launches after install are never recorded. + assert not attention_plugin._calibration_active(module.impl) + + def test_rejects_non_eager_execution(self): + runner = _model_runner( + nn.ModuleDict({"attn": _bare_attention()}), + cudagraph_mode=CUDAGraphMode.PIECEWISE, + ) + with pytest.raises(NotImplementedError, match="enforce_eager"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + def test_rejects_active_attention_quantizers_atomically(self): + quantized = _bare_attention() + quantized.q_bmm_quantizer = SimpleNamespace(is_enabled=True) + clean = _bare_attention() + clean_impl = clean.impl + runner = _model_runner(nn.ModuleDict({"q_attn": quantized, "c_attn": clean})) + + with pytest.raises(NotImplementedError, match="requires unquantized attention"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + # Validation-before-mutation: the clean layer must not be touched either. + assert clean.impl is clean_impl + assert runner.cascade_attn_enabled is True + + def test_rejects_fp8_kv_cache(self): + attention = _bare_attention() + attention.kv_cache_dtype = "fp8" + runner = _model_runner(nn.ModuleDict({"attn": attention})) + with pytest.raises(NotImplementedError, match="FP8 KV cache"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + def test_rejects_model_without_attention_layers(self): + runner = _model_runner(nn.ModuleDict({})) + with pytest.raises(NotImplementedError, match="no attention layers"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + +class TestQuantSkipRejection: + """Skip-softmax cannot be combined with attention quantization.""" + + _CALIBRATED_META = { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "threshold_scale_factor": {"prefill": {"a": 7.9, "b": 8.6}}, + } + } + } + _NM_META = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4} + } + } + + def test_quantized_install_rejects_calibrated_skip(self): + attention = _bare_attention() + runner = _model_runner( + nn.ModuleDict({"attn": attention}), sparse_metadata=self._CALIBRATED_META + ) + with pytest.raises( + NotImplementedError, match="cannot be combined with attention quantization" + ): + vllm_runtime.install_vllm_nvfp4_attention(runner) + assert not isinstance(attention.impl, ModelOptSparseAttentionImpl) + + def test_sparse_only_install_rejects_skip_onto_quantized_layer(self): + """Sparse-only installs must also refuse skip onto live quantizers.""" + attention = _bare_attention() + attention.q_bmm_quantizer = SimpleNamespace(is_enabled=True) + original_impl = attention.impl + runner = _model_runner( + nn.ModuleDict({"attn": attention}), sparse_metadata=self._CALIBRATED_META + ) + with pytest.raises( + NotImplementedError, match="cannot be combined with attention quantization" + ): + vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert attention.impl is original_impl + + def test_sparse_only_install_allows_skip_on_unquantized_layer(self): + runner = _model_runner( + nn.ModuleDict({"attn": _bare_attention()}), sparse_metadata=self._CALIBRATED_META + ) + report = vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert report.installed_count == 1 + + def test_quantized_plan_allows_nm_sparsity(self, monkeypatch): + monkeypatch.setattr( + quant_plugin, + "_get_device_dtype", + lambda module: (torch.device("cpu"), torch.float16), + ) + runner = _model_runner( + nn.ModuleDict({"attn": _bare_attention()}), sparse_metadata=self._NM_META + ) + plan = vllm_runtime._plan_vllm_attention(runner, quantize=True, sparse_cfg="checkpoint") + assert len(plan.layers) == 1 + assert plan.layers[0].sparse_kw.get("sparsity_n") == 2 + + +class TestFlashInferLayoutGuard: + """HND FlashInfer caches are rejected via layout metadata, pre-measurement.""" + + def test_layout_helper_preserves_none(self, monkeypatch): + """A getter returning None must not become the truthy string 'None'.""" + fake = SimpleNamespace(get_kv_cache_layout=lambda: None) + monkeypatch.setitem(sys.modules, "vllm.v1.attention.backends.utils", fake) + assert attention_plugin._flashinfer_kv_cache_layout() is None + + fake.get_kv_cache_layout = lambda: "HND" + assert attention_plugin._flashinfer_kv_cache_layout() == "HND" + + def test_installer_rejects_hnd_layout(self, monkeypatch): + monkeypatch.setattr(attention_plugin, "_flashinfer_kv_cache_layout", lambda: "HND") + attention = _bare_attention(FlashInferImpl) + original_impl = attention.impl + runner = _model_runner(nn.ModuleDict({"attn": attention})) + with pytest.raises(NotImplementedError, match="HND"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + assert attention.impl is original_impl + + def test_forward_rejects_hnd_layout_before_cache_write(self, monkeypatch): + monkeypatch.setattr(attention_plugin, "_flashinfer_kv_cache_layout", lambda: "HND") + writes = [] + monkeypatch.setattr( + attention_plugin, + "_maybe_update_flashinfer_cache", + lambda *args, **kwargs: writes.append(1), + ) + num_heads, num_kv_heads, head_dim, page = 4, 2, 64, 16 + impl = SimpleNamespace( + num_kv_heads=num_kv_heads, + head_size=head_dim, + scale=1.0 / (head_dim**0.5), + _calibrate=True, + _calib_threshold_trials=list(TRIALS), + _calib_records=[], + ) + kv_cache = torch.zeros(3, 2, page, num_kv_heads, head_dim, dtype=torch.bfloat16) + attn_metadata = SimpleNamespace( + _modelopt_block_table=torch.zeros(1, 1, dtype=torch.int32), + _modelopt_seq_lens=torch.tensor([8], dtype=torch.int32), + _modelopt_query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + _modelopt_num_actual_tokens=1, + _modelopt_max_query_len=1, + _modelopt_max_seq_len=8, + _modelopt_causal=False, + slot_mapping=torch.zeros(1, dtype=torch.int64), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + with pytest.raises(NotImplementedError, match="HND"): + attention_plugin._flashinfer_forward( + impl, + None, + None, + q, + q[:, :num_kv_heads], + q[:, :num_kv_heads], + kv_cache, + attn_metadata, + output=torch.empty_like(q), + ) + assert not writes, "layout must be validated before the cache write" + + +class TestSparseOnlyGraphGuard: + """Commit-contract: the calibrated-decode graph guard is not quantize-gated.""" + + _CALIBRATED_META = { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "threshold_scale_factor": { + "prefill": {"a": 7.9, "b": 8.6}, + "decode": {"a": 0.12, "b": 9.8}, + }, + } + } + } + + def test_sparse_only_install_rejects_full_decode_graph(self): + attention = _bare_attention() + runner = _model_runner( + nn.ModuleDict({"attn": attention}), + sparse_metadata=self._CALIBRATED_META, + cudagraph_mode=CUDAGraphMode.FULL, + ) + with pytest.raises(NotImplementedError, match="non-FULL CUDA graph"): + vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert not isinstance(attention.impl, ModelOptSparseAttentionImpl) + + def test_sparse_only_install_allows_eager(self): + attention = _bare_attention() + runner = _model_runner( + nn.ModuleDict({"attn": attention}), + sparse_metadata=self._CALIBRATED_META, + cudagraph_mode=CUDAGraphMode.NONE, + ) + report = vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert report.installed_count == 1 + + +# --------------------------------------------------------------------------- +# Calibration forward through the FlashAttention adapter (GPU) +# --------------------------------------------------------------------------- +def _make_impl(num_heads, head_dim, num_kv_heads): + return ModelOptSparseAttentionImpl( + num_heads=num_heads, + head_size=head_dim, + scale=1.0 / (head_dim**0.5), + num_kv_heads=num_kv_heads, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + ) + + +def _paged_cache_for(seqs_kv, num_kv_heads, head_dim, page_size, device, dtype): + """Scatter per-request contiguous K/V lists into a stacked paged cache.""" + blocks_per_seq = [(kv.shape[0] + page_size - 1) // page_size for kv, _ in seqs_kv] + num_blocks = sum(blocks_per_seq) + max_blocks = max(blocks_per_seq) + k_cache = torch.zeros(num_blocks, page_size, num_kv_heads, head_dim, device=device, dtype=dtype) + v_cache = torch.zeros_like(k_cache) + block_table = torch.zeros(len(seqs_kv), max_blocks, device=device, dtype=torch.int32) + g = 0 + for b, (k, v) in enumerate(seqs_kv): + for blk in range(blocks_per_seq[b]): + block_table[b, blk] = g + ts, te = blk * page_size, min((blk + 1) * page_size, k.shape[0]) + k_cache[g, : te - ts] = k[ts:te] + v_cache[g, : te - ts] = v[ts:te] + g += 1 + return torch.stack([k_cache, v_cache], dim=0), block_table + + +def _sdpa_reference(q, k, v, is_causal): + # [tokens, heads, dim] -> [1, heads, tokens, dim] + qh, kh, vh = (t.transpose(0, 1).unsqueeze(0).float() for t in (q, k, v)) + kh = kh.repeat_interleave(q.shape[1] // k.shape[1], dim=1) + vh = vh.repeat_interleave(q.shape[1] // v.shape[1], dim=1) + if is_causal and q.shape[0] < k.shape[0]: + # Suffix-causal mask for decode/chunked prefill shapes. + mask = torch.ones(q.shape[0], k.shape[0], dtype=torch.bool, device=q.device).tril( + diagonal=k.shape[0] - q.shape[0] + ) + out = torch.nn.functional.scaled_dot_product_attention(qh, kh, vh, attn_mask=mask) + else: + out = torch.nn.functional.scaled_dot_product_attention(qh, kh, vh, is_causal=is_causal) + return out.squeeze(0).transpose(0, 1).to(q.dtype) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestCalibrationForward: + def test_mixed_batch_records_phases_and_stays_dense(self): + torch.manual_seed(0) + device, dtype = "cuda", torch.bfloat16 + num_heads, num_kv_heads, head_dim, page_size = 4, 2, 64, 16 + prefill_len, decode_ctx = 64, 48 + + k0 = torch.randn(prefill_len, num_kv_heads, head_dim, device=device, dtype=dtype) + v0 = torch.randn_like(k0) + k1 = torch.randn(decode_ctx, num_kv_heads, head_dim, device=device, dtype=dtype) + v1 = torch.randn_like(k1) + q = torch.randn(prefill_len + 1, num_heads, head_dim, device=device, dtype=dtype) + + kv_cache, block_table = _paged_cache_for( + [(k0, v0), (k1, v1)], num_kv_heads, head_dim, page_size, device, dtype + ) + attn_metadata = SimpleNamespace( + num_actual_tokens=prefill_len + 1, + max_query_len=prefill_len, + max_seq_len=max(prefill_len, decode_ctx), + query_start_loc=torch.tensor( + [0, prefill_len, prefill_len + 1], device=device, dtype=torch.int32 + ), + seq_lens=torch.tensor([prefill_len, decode_ctx], device=device, dtype=torch.int32), + block_table=block_table, + ) + + impl = _make_impl(num_heads, head_dim, num_kv_heads) + impl.sparse_kw = {} + enable_calibration([impl], TRIALS) + output = torch.empty_like(q) + out = impl.forward( + layer=None, + query=q, + key=q[:, :num_kv_heads], + value=q[:, :num_kv_heads], + kv_cache=kv_cache, + attn_metadata=attn_metadata, + output=output, + ) + + # Two records: one per request, phases decided per request. + records = impl._calib_records + assert [r["phase"] for r in records] == ["prefill", "decode"] + assert [r["sample_length"] for r in records] == [prefill_len, decode_ctx] + for record in records: + assert len(record["total_tiles"]) == len(TRIALS) + assert all(t > 0 for t in record["total_tiles"]) + assert all(0 <= s <= t for s, t in zip(record["skipped_tiles"], record["total_tiles"])) + + # Output is full dense attention (calibration never skips). + ref_prefill = _sdpa_reference(q[:prefill_len], k0, v0, is_causal=True) + ref_decode = _sdpa_reference(q[prefill_len:], k1, v1, is_causal=False) + torch.testing.assert_close(out[:prefill_len], ref_prefill, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(out[prefill_len:], ref_decode, rtol=2e-2, atol=2e-2) + + def test_collect_calibration_counts_sums_layers(self): + class FakeModel(nn.Module): + def __init__(self, impls): + super().__init__() + self._impls = impls + self.layers = nn.ModuleList([nn.Identity() for _ in impls]) + for identity, impl in zip(self.layers, impls): + identity.impl = impl + + impls = [object.__new__(ModelOptSparseAttentionImpl) for _ in range(2)] + enable_calibration(impls, TRIALS) + for idx, impl in enumerate(impls): + impl._calib_records = [ + { + "phase": "prefill", + "sample_length": 128, + "total_tiles": [4, 4, 4], + "skipped_tiles": [idx, idx + 1, idx + 2], + } + ] + model = FakeModel(impls) + assert len(list(iter_sparse_impls(model))) == 2 + disable_calibration(impls) + + counts = collect_calibration_counts(model) + assert counts["prefill"] == [ + {"sample_length": 128, "total_tiles": [8, 8, 8], "skipped_tiles": [1, 3, 5]} + ] + + def test_rejects_non_nhd_cache_layout(self): + num_heads, num_kv_heads, head_dim = 4, 2, 64 + impl = _make_impl(num_heads, head_dim, num_kv_heads) + enable_calibration([impl], TRIALS) + # HND-shaped cache: [blocks, kv_heads, page, dim] -> axis 2 != num_kv_heads. + kv_cache = torch.zeros(2, 1, num_kv_heads, 16, head_dim, dtype=torch.bfloat16) + attn_metadata = SimpleNamespace( + num_actual_tokens=1, + max_query_len=1, + max_seq_len=8, + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=torch.tensor([8], dtype=torch.int32), + block_table=torch.zeros(1, 1, dtype=torch.int32), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + with pytest.raises(NotImplementedError, match="not NHD"): + impl.forward( + layer=None, + query=q, + key=q[:, :num_kv_heads], + value=q[:, :num_kv_heads], + kv_cache=kv_cache, + attn_metadata=attn_metadata, + output=torch.empty_like(q), + ) + + def test_rejects_non_16bit_cache(self): + num_heads, num_kv_heads, head_dim = 4, 2, 64 + impl = _make_impl(num_heads, head_dim, num_kv_heads) + enable_calibration([impl], TRIALS) + kv_cache = torch.zeros(2, 1, 16, num_kv_heads, head_dim, dtype=torch.uint8) + attn_metadata = SimpleNamespace( + num_actual_tokens=1, + max_query_len=1, + max_seq_len=8, + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=torch.tensor([8], dtype=torch.int32), + block_table=torch.zeros(1, 1, dtype=torch.int32), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + with pytest.raises(NotImplementedError, match="fp16/bf16 KV cache"): + impl.forward( + layer=None, + query=q, + key=q[:, :num_kv_heads], + value=q[:, :num_kv_heads], + kv_cache=kv_cache, + attn_metadata=attn_metadata, + output=torch.empty_like(q), + ) + + +# --------------------------------------------------------------------------- +# FlashInfer adapter: cache write must precede the calibrate-kernel read +# --------------------------------------------------------------------------- +class TestFlashInferCalibrationOrdering: + def test_cache_write_happens_before_calibrate_read(self, monkeypatch): + calls = [] + monkeypatch.setattr( + attention_plugin, + "_maybe_update_flashinfer_cache", + lambda *args, **kwargs: calls.append("cache_write"), + ) + + def fake_calibrate(q, *args, **kwargs): + calls.append("calibrate") + counters = torch.zeros(len(TRIALS), 2, dtype=torch.int64) + return torch.zeros_like(q), counters + + monkeypatch.setattr(attention_plugin, "attention_calibrate", fake_calibrate) + + num_heads, num_kv_heads, head_dim, page = 4, 2, 64, 16 + impl = SimpleNamespace( + num_kv_heads=num_kv_heads, + head_size=head_dim, + scale=1.0 / (head_dim**0.5), + _calibrate=True, + _calib_threshold_trials=list(TRIALS), + _calib_records=[], + ) + kv_cache = torch.zeros(3, 2, page, num_kv_heads, head_dim, dtype=torch.bfloat16) + attn_metadata = SimpleNamespace( + _modelopt_block_table=torch.zeros(1, 1, dtype=torch.int32), + _modelopt_seq_lens=torch.tensor([8], dtype=torch.int32), + _modelopt_query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + _modelopt_num_actual_tokens=1, + _modelopt_max_query_len=1, + _modelopt_max_seq_len=8, + _modelopt_causal=False, + slot_mapping=torch.zeros(1, dtype=torch.int64), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + + out = attention_plugin._flashinfer_forward( + impl, + None, # native_forward is unused on the calibration path + None, # layer + q, + q[:, :num_kv_heads], + q[:, :num_kv_heads], + kv_cache, + attn_metadata, + output=torch.empty_like(q), + ) + + assert calls == ["cache_write", "calibrate"] + assert torch.isfinite(out).all() + assert len(impl._calib_records) == 1 + assert impl._calib_records[0]["phase"] == "decode" diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py index b3ceef30eab..9c448c1327c 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py @@ -90,7 +90,6 @@ def test_sparse_install_from_checkpoint_is_validation_atomic(): nn.ModuleDict({"valid_attn": valid, "invalid_attn": invalid}), sparse_metadata=_sparse_metadata(), ) - del runner.vllm_config with pytest.raises(NotImplementedError, match="sliding_window"): vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) @@ -107,7 +106,6 @@ def test_installer_rejects_cross_attention_layout_even_if_marked_decoder(): nn.ModuleDict({"cross_attn": attention}), sparse_metadata=_sparse_metadata(), ) - del runner.vllm_config with pytest.raises(NotImplementedError, match="layout CrossAttention"): vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) @@ -122,7 +120,6 @@ def test_sparse_install_uses_checkpoint_metadata(monkeypatch, impl_cls): nn.ModuleDict({"attn": attention}), sparse_metadata=_sparse_metadata(), ) - del runner.vllm_config monkeypatch.setattr( vllm_runtime.attention_plugin, "patch_flashinfer_metadata_builder", lambda: True ) diff --git a/tests/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index 62395ff5a7a..13ce5da55ef 100644 --- a/tests/unit/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -139,21 +139,7 @@ def test_forward_routes_every_mode_to_single_autotuner( assert kernel.kwargs["V_QDQ"] == expected_v_qdq -@pytest.mark.parametrize( - ("attention_kwargs", "expected_block_m"), - [ - ({"skip_softmax_threshold": 0.1, "measure_sparsity": True}, 128), - ( - { - "p_qdq": "nvfp4", - "skip_softmax_threshold": 0.1, - "measure_sparsity": True, - }, - 16, - ), - ], -) -def test_forward_measurement_uses_one_fixed_launch(monkeypatch, attention_kwargs, expected_block_m): +def test_forward_measurement_uses_one_fixed_launch(monkeypatch): """Counter measurement bypasses autotuning to avoid repeated atomic updates.""" pytest.importorskip("triton") @@ -173,10 +159,42 @@ def test_forward_measurement_uses_one_fixed_launch(monkeypatch, attention_kwargs starts = torch.tensor([0], dtype=torch.int32) lengths = torch.tensor([seq_len], dtype=torch.int32) - triton_fa.attention(q, k, v, starts, lengths, seq_len, **attention_kwargs) + triton_fa.attention( + q, k, v, starts, lengths, seq_len, skip_softmax_threshold=0.1, measure_sparsity=True + ) assert kernel.fn.launch_count == 1 - assert kernel.fn.kwargs["BLOCK_M"] == expected_block_m + assert kernel.fn.kwargs["BLOCK_M"] == 128 assert kernel.fn.kwargs["BLOCK_N"] == 128 assert kernel.fn.kwargs["num_stages"] == 1 assert kernel.fn.kwargs["num_warps"] == 4 + + +@pytest.mark.parametrize( + "qdq_kwargs", + [{"p_qdq": "nvfp4"}, {"p_qdq": "fp8"}, {"v_qdq": "nvfp4", "v_qdq_amax": 1.0}], +) +def test_forward_rejects_skip_softmax_with_qdq(monkeypatch, qdq_kwargs): + """Active skip-softmax rejects P/V QDQ before any kernel launch.""" + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + kernel = _ForbiddenKernel() + kernel.fn = _ForbiddenKernel() + monkeypatch.setattr(triton_fa, "_attn_fwd", kernel) + monkeypatch.setattr(triton_fa.torch.cuda, "device", lambda _device: nullcontext()) + monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) + monkeypatch.setattr(triton_fa, "_load_qdq_helpers", lambda: None) + + seq_len = 129 + q = torch.empty(seq_len, 2, 16) + k = torch.empty(seq_len, 1, 16) + v = torch.empty_like(k) + starts = torch.tensor([0], dtype=torch.int32) + lengths = torch.tensor([seq_len], dtype=torch.int32) + + with pytest.raises(ValueError, match="cannot be combined with attention quantization"): + triton_fa.attention( + q, k, v, starts, lengths, seq_len, skip_softmax_threshold=0.1, **qdq_kwargs + ) 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, + ) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py new file mode 100644 index 00000000000..7ff174952b3 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py @@ -0,0 +1,270 @@ +# 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. + +"""Unit tests for vLLM-free skip-softmax calibration helpers (no vLLM needed).""" + +import math +from types import SimpleNamespace + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration.calibrator import ( + DynamicThresholdCalibrator, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_calibration import ( + DEFAULT_THRESHOLD_TRIALS, + build_sparse_attention_config, + fit_from_counts, + merge_count_records, + merge_phase_counts, + split_records_by_phase, + stats_from_counts, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_config import ( + load_from_checkpoint_metadata, +) + + +def _record(phase, length, totals, skipped): + return { + "phase": phase, + "sample_length": length, + "total_tiles": list(totals), + "skipped_tiles": list(skipped), + } + + +class TestCountMerging: + def test_split_records_by_phase_preserves_order(self): + records = [ + _record("prefill", 100, [4], [1]), + _record("decode", 101, [2], [0]), + _record("prefill", 200, [8], [3]), + ] + split = split_records_by_phase(records) + assert [r["sample_length"] for r in split["prefill"]] == [100, 200] + assert [r["sample_length"] for r in split["decode"]] == [101] + + def test_merge_sums_counts_elementwise(self): + layer_a = [_record("prefill", 128, [10, 10], [2, 4])] + layer_b = [_record("prefill", 128, [10, 10], [1, 3])] + merged = merge_count_records([layer_a, layer_b]) + assert merged == [{"sample_length": 128, "total_tiles": [20, 20], "skipped_tiles": [3, 7]}] + + def test_merge_rejects_ragged_sources(self): + long = [_record("prefill", 1, [1], [0]), _record("prefill", 2, [1], [1])] + short = [_record("prefill", 1, [1], [1])] + with pytest.raises(ValueError, match="disagree on sample count"): + merge_count_records([long, short]) + + def test_merge_rejects_misaligned_sample_lengths(self): + with pytest.raises(ValueError, match="Misaligned calibration records"): + merge_count_records( + [[_record("prefill", 100, [1], [0])], [_record("prefill", 200, [1], [0])]] + ) + + def test_merge_rejects_threshold_width_mismatch(self): + with pytest.raises(ValueError, match="threshold-vector widths"): + merge_count_records( + [[_record("prefill", 100, [1, 2], [0, 1])], [_record("prefill", 100, [1], [0])]] + ) + + def test_merge_phase_counts_rejects_rank_phase_mismatch(self): + rank0 = {"prefill": [_record("prefill", 64, [5], [1])]} + rank1 = {"prefill": []} + with pytest.raises(ValueError, match="recorded no 'prefill' samples"): + merge_phase_counts([rank0, rank1]) + + def test_merge_phase_counts_across_ranks(self): + rank0 = {"prefill": [_record("prefill", 64, [5], [1])], "decode": []} + rank1 = {"prefill": [_record("prefill", 64, [5], [2])]} + merged = merge_phase_counts([rank0, rank1]) + assert merged["prefill"][0]["total_tiles"] == [10] + assert merged["prefill"][0]["skipped_tiles"] == [3] + assert merged["decode"] == [] + + def test_stats_from_counts_forms_ratios_after_merge(self): + stats = stats_from_counts( + [{"sample_length": 64, "total_tiles": [8, 0], "skipped_tiles": [2, 0]}] + ) + assert stats == [{"sample_length": 64, "sparsity": [0.25, 0.0]}] + + +class TestFitFromCounts: + def test_fit_recovers_synthetic_exponential(self): + a_true, b_true = 5.0, 8.0 + trials = DEFAULT_THRESHOLD_TRIALS + + def counts(length, total): + sparsity = [ + min(0.95, max(0.0, math.log(max(t * length, 1e-9) / a_true) / b_true)) + for t in trials + ] + return { + "sample_length": length, + "total_tiles": [total] * len(trials), + "skipped_tiles": [int(s * total) for s in sparsity], + } + + per_phase = {"prefill": [counts(length, 4000) for length in (2048, 4096, 8192, 16384)]} + params = fit_from_counts(per_phase, trials) + assert abs(params["prefill"]["a"] - a_true) / a_true < 0.3 + assert abs(params["prefill"]["b"] - b_true) / b_true < 0.15 + assert 0.0 <= params["prefill"]["min_observed_sparsity"] <= 1.0 + + def test_empty_phase_produces_no_fit(self): + assert fit_from_counts({"decode": []}, DEFAULT_THRESHOLD_TRIALS) == {} + + def test_fit_rejects_counter_width_vs_trials_mismatch(self): + """Consistent-but-wrong widths must not silently zip against the trials.""" + short = len(DEFAULT_THRESHOLD_TRIALS) - 1 + records = [ + {"sample_length": 4096, "total_tiles": [100] * short, "skipped_tiles": [50] * short} + ] + with pytest.raises(ValueError, match="threshold trials are configured"): + fit_from_counts({"prefill": records}, DEFAULT_THRESHOLD_TRIALS) + + +class TestCalibrateFromStats: + def _stats(self, trials): + a_true, b_true = 3.0, 9.0 + stats = [] + for length in (1024, 2048, 4096, 8192): + sparsity = [ + min(0.95, max(0.0, math.log(max(t * length, 1e-9) / a_true) / b_true)) + for t in trials + ] + stats.append({"sample_length": length, "sparsity": sparsity}) + return stats + + def test_linear_fit_reports_fit_logspace_false(self): + calibrator = DynamicThresholdCalibrator(threshold_trials=list(DEFAULT_THRESHOLD_TRIALS)) + result = calibrator.calibrate_from_stats(self._stats(DEFAULT_THRESHOLD_TRIALS), "prefill") + assert result["fit_logspace"] is False + assert "log_a" not in result + assert len(result["per_sample_sparsity"]) == 4 + + def test_logspace_fit_preserves_log_a(self): + calibrator = DynamicThresholdCalibrator( + threshold_trials=list(DEFAULT_THRESHOLD_TRIALS), fit_logspace=True + ) + result = calibrator.calibrate_from_stats(self._stats(DEFAULT_THRESHOLD_TRIALS), "prefill") + assert result["fit_logspace"] is True + assert math.isclose(math.exp(result["log_a"]), result["a"], rel_tol=1e-9) + + +class TestBuildSparseAttentionConfig: + _PARAMS = {"prefill": {"a": 7.9, "b": 8.6}, "decode": {"a": 0.12, "b": 9.8}} + + def test_canonical_schema(self): + config = build_sparse_attention_config(self._PARAMS, 0.4) + group = config["config_groups"]["group_0"] + assert group["algorithm"] == "skip_softmax" + assert group["threshold_scale_factor"]["prefill"] == {"a": 7.9, "b": 8.6} + assert group["threshold_scale_factor"]["formula"] == "a * exp(b * target_sparsity)" + assert group["target_sparsity"] == {"prefill": 0.4, "decode": 0.4} + assert config["producer"]["name"] == "modelopt" + + def test_target_sparsity_covers_only_fitted_phases(self): + """A phase without calibrated (a, b) must not claim a sparsity target.""" + config = build_sparse_attention_config({"prefill": {"a": 7.9, "b": 8.6}}, 0.4) + group = config["config_groups"]["group_0"] + assert group["target_sparsity"] == {"prefill": 0.4} + assert "decode" not in group["threshold_scale_factor"] + + def test_replaced_skip_group_keeps_layer_policy(self): + """Recalibration replaces thresholds but keeps ignore/initial_disabled_steps.""" + existing = { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "ignore": ["model.layers.0.self_attn"], + "initial_disabled_steps": 4, + "threshold_scale_factor": {"prefill": {"a": 1.0, "b": 1.0}}, + } + } + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + group = config["config_groups"]["group_0"] + assert group["ignore"] == ["model.layers.0.self_attn"] + assert group["initial_disabled_steps"] == 4 + assert group["threshold_scale_factor"]["prefill"] == {"a": 7.9, "b": 8.6} + + def test_preserves_nm_groups_and_replaces_old_skip_group(self): + existing = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4}, + "group_1": { + "algorithm": "skip_softmax", + "threshold_scale_factor": {"prefill": {"a": 1.0, "b": 1.0}}, + }, + } + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + groups = config["config_groups"] + assert len(groups) == 2 + assert groups["group_0"]["algorithm"] == "skip_softmax" + assert groups["group_0"]["threshold_scale_factor"]["prefill"]["a"] == 7.9 + assert groups["group_1"]["algorithm"] == "sparse_softmax" + assert groups["group_1"]["sparsity_n"] == 2 + + def test_round_trips_through_serving_loader(self): + config = build_sparse_attention_config(self._PARAMS, {"prefill": 0.5, "decode": 0.3}) + hf_config = SimpleNamespace(sparse_attention_config=config) + loaded = load_from_checkpoint_metadata(hf_config) + assert loaded is not None + sparse_cfg, preset = loaded + assert preset == "CHECKPOINT_CALIBRATED_SOFTMAX_SKIP" + layer_cfg = sparse_cfg["sparse_cfg"]["*attn*"] + assert layer_cfg["method"] == "triton_skip_softmax" + assert layer_cfg["threshold_scale_factor"]["decode"] == {"a": 0.12, "b": 9.8} + assert layer_cfg["target_sparse_ratio"] == {"prefill": 0.5, "decode": 0.3} + + def test_rejects_out_of_range_target_sparsity(self): + with pytest.raises(ValueError, match=r"between 0\.0 and 1\.0"): + build_sparse_attention_config(self._PARAMS, 1.5) + with pytest.raises(ValueError, match=r"between 0\.0 and 1\.0"): + build_sparse_attention_config(self._PARAMS, {"prefill": 0.5, "decode": -0.1}) + + def test_preserves_legacy_toplevel_sparse_softmax(self): + existing = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4} + }, + "sparse_softmax": {"sparsity_n": 1, "sparsity_m": 4, "dense_recent_tokens": 128}, + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + # The serving loader reads the legacy top-level key ahead of group params. + assert config["sparse_softmax"] == existing["sparse_softmax"] + loaded = load_from_checkpoint_metadata(SimpleNamespace(sparse_attention_config=config)) + assert loaded is not None + layer_cfg = loaded[0]["sparse_cfg"]["*attn*"] + assert layer_cfg["sparsity_n"] == 1 + assert layer_cfg["dense_recent_tokens"] == 128 + + def test_round_trip_with_preserved_nm_group_activates_both(self): + existing = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4} + } + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + loaded = load_from_checkpoint_metadata(SimpleNamespace(sparse_attention_config=config)) + assert loaded is not None + sparse_cfg, preset = loaded + assert preset == "CHECKPOINT_CALIBRATED_SOFTMAX_SKIP_SPARSE_SOFTMAX" + layer_cfg = sparse_cfg["sparse_cfg"]["*attn*"] + assert layer_cfg["sparsity_n"] == 2 + assert "threshold_scale_factor" in layer_cfg