diff --git a/CHANGELOG.rst b/CHANGELOG.rst index be05c1c441e..adfa528a7fc 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -28,6 +28,7 @@ Experimental **New Features** - Add NVFP4 and FP8 PTQ recipes with projection-output quantizers for Llama-Nemotron embedding and reranking models (``modelopt_recipes/huggingface/nemotron_llama/``) and an end-to-end HF embedding/reranking quantize-to-ONNX example (``examples/torch_onnx/hf_embedding_quant_to_onnx.py``). Quantizing the projection-Linear outputs keeps TensorRT inter-layer activations in FP4, roughly halving engine activation memory versus the plain ``nvfp4`` preset. NVFP4/MXFP8 output quantizers now export through the dynamic quantize path. ``examples/torch_onnx/torch_quant_to_onnx.py`` also gains a ``--recipe`` flag to load quantization configs from YAML recipes instead of the removed ``mtq.*_CFG`` module-constant table. +- 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. See ``examples/vllm_serve/calibrate_sparse_attn.py``. - Add the ``prepare_megatron_data_blend`` utility to prepare weighted Megatron data blends from YAML configs, including optional token-budgeted subsets for distillation workflows. See the `Megatron data preparation guide `_. - Add Learned Scale Quantization (LSQ) and Dual-LSQ support for quantization-aware distillation, including learnable ``amax`` parameters, tied-scale and pre-scale options, focused NVFP4 recipes, and scale-only training. - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 `_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change). diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 858243686d0..c5737cd91e2 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -117,6 +117,152 @@ 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. + +### Build a cross-layer mask-reuse calibration candidate + +Mask-reuse calibration is a second, offline ModelOpt stage. The current exporter +produces a **candidate only**: it deliberately cannot promote a serving policy +until grouped inner-fold stability, the preregistered 99-outer-group gate, and +the deployment-geometry gate are implemented and pass. + +First run `calibrate_sparse_attn.py` to fit the vanilla skip-softmax transfer +function and finalize the checkpoint/config. Then create the content-addressed +identity inside the exact local checkpoint that vLLM will load: + +```bash +python create_checkpoint_manifest.py --model-id Nemotron-3-Ultra +``` + +This deterministic command writes `/checkpoint_manifest.json` with +SHA256 and size for every checkpoint file except the manifest itself. It refuses +to overwrite an existing manifest, rejects symlinks, and verifies the complete +file set. Keep the checkpoint root immutable after creating it. + +Create the runtime FA4 source from the exact Git commit outside the serving +container. The generator emits a scoped `git archive` containing only +`flash_attn/` and a canonical witness for every directory and regular file: + +```bash +python create_fa4_source_witness.py /path/to/flash-attention \ + --expected-commit <40-HEX-COMMIT> \ + --archive-output fa4-runtime-source.tar \ + --manifest-output fa4-runtime-source.manifest.json +mkdir fa4-runtime-source +tar -xf fa4-runtime-source.tar -C fa4-runtime-source +``` + +Pin the printed archive and manifest hashes independently. The runtime verifier +does not invoke Git; it rejects missing, extra, changed, executable-mode, link, +reparse-point, and special-file differences from the witnessed archive. + +Then collect dense-reference mask-reuse sufficient statistics at a menu of target +sparsities under the exact runtime rule +`threshold = a * exp(b * target_sparsity) / kv_tokens`: + +```bash +python collect_mask_reuse.py \ + --model-id Nemotron-3-Ultra \ + --checkpoint-manifest-sha256 <64-HEX-CHECKPOINT-MANIFEST-SHA256> \ + --plan nemotron3_ultra_stride2 \ + --fa4-source /path/to/fa4-runtime-source \ + --fa4-source-manifest fa4-runtime-source.manifest.json \ + --fa4-source-manifest-sha256 <64-HEX-SOURCE-MANIFEST-SHA256> \ + --fa4-commit <40-HEX-COMMIT> \ + --prompts-jsonl mask_reuse_prompts.jsonl \ + --vanilla-config /config.json \ + --target-sparsities 0.5 0.6 0.7 \ + --max-model-len 1000001 \ + --tensor-parallel-size 8 \ + --validate-dense-output \ + --output mask_reuse_compact_captures.jsonl +``` + +The prompt file has one strict JSON object per capture unit with `split`, +`partition`, `inner_fold`, globally split-unique `prompt_id`, `source`, +`source_group_sha256`, `prompt`, `min_kv_tokens`, and `max_kv_tokens`. +`calibration` maps to `development` with a nonnegative fold; `heldout` maps to +`outer_test` with a null fold. One source group cannot cross partitions or +folds. The collector forces eager, batch-one, BF16, chunked-prefill +vLLM V1 execution through the installed `mask_reuse_fa4` plugin and the +explicit pinned FA4 source. It does not load a provisional reuse policy and +does not enable quantization. The verified checkpoint identity and group +assignment are included in the schema-v2 invocation and its canonical RPC +digest. The checkpoint and FA4 source are reverified after vLLM loads them and +before capture; FA4 is verified once more before evidence publication. +Each compact JSONL record includes the target +sparsity, sample length, exact binary64 +`threshold_log2` and `threshold_lambda`, rectangular query/KV geometry, every +topology anchor/head's self-mask statistics, and every consumer/donor +dropped-mass matrix once. It does not expand or repeat the matrix as millions +of candidate rows. Fixed-lambda version-4/version-5 captures are not valid +inputs to this target-sparsity path. Use the printed compact-capture SHA256 as +the `reuse_bundle_sha256` evidence value. The version-4 capture manifest binds +the independently pinned checkpoint and FA4 source witnesses, Git commit/tree, +archive digest, canonical vLLM engine arguments, rank-major TP sentinel +evidence, exact prefill/decode attention-call counts, and optional bitwise +dense-shadow coverage. `--validate-dense-output` runs a second pinned dense FA4 +call for every armed final-chunk layer and fails before publication if its BF16 +output differs bitwise from the model-trajectory output. Schema-version-3 +capture manifests are intentionally rejected because they lack the runtime +source witness; recollect them rather than migrating their provenance labels. + +Once those observations have been captured, select the per-context target, +donor-head map, and exact fallbacks and export the fail-closed candidate: + +```bash +python calibrate_mask_reuse.py \ + --checkpoint \ + --compact-captures mask_reuse_compact_captures.jsonl \ + --capture-manifest mask_reuse_compact_captures.jsonl.manifest.json \ + --vanilla-config /config.json \ + --topology mask_reuse_topology.json \ + --calibration-plan mask_reuse_calibration_plan.json \ + --family-registry mask_reuse_family_registry.json \ + --grouped-fit mask_reuse_grouped_fit.json \ + --outer-report mask_reuse_outer_report.json \ + --max-anchor-dropped-mass 0.02 \ + --max-reuse-selection-dropped-mass 0.01 \ + --max-reuse-dropped-mass 0.02 \ + --output-policy mask_reuse_candidate.json \ + --output-report mask_reuse_calibration_report.json +``` + +All six evidence hashes are computed from the exact supplied artifact files; +free-form digest claims are not accepted. The compact selector streams three +semantic passes over the captures and hashes +the exact file bytes before and after selection/evaluation. It aborts if the +capture bundle or another evidence artifact changes while calibration is +running. It minimizes the declared equal-BMM tile cost `2 * A_R + A_A`. + +Selection uses only records labeled `calibration`; the frozen policy is then +evaluated on records labeled `heldout`. Held-out violations are preserved in +the standalone report and do not silently retune the selected policy. The +candidate has `promotion_status="candidate_only"` and +`deployment_geometry_validated=false`; the serving backend rejects it. Candidate +and report publication is no-clobber and report-first with rollback, so a policy +path cannot appear without its complete report during normal process execution. +On platforms without portable directory `fsync`, directory-entry durability +across power loss is best effort. Decode is explicitly dense. + 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: @@ -131,7 +277,7 @@ 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. +- 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`. ### Compact NVFP4 attention worker @@ -148,7 +294,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..59c4b5db00b --- /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 \ + --max-reuse-dropped-mass 0.02 \ + --max-reuse-selection-dropped-mass 0.005 + +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( + "--max-reuse-dropped-mass", + type=float, + required=True, + help="Maximum allowed held-out reuse dropped mass", + ) + parser.add_argument( + "--max-reuse-selection-dropped-mass", + type=float, + default=None, + help="Optional stricter calibration-time reuse bound", + ) + 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, + max_reuse_dropped_mass=args.max_reuse_dropped_mass, + max_reuse_selection_dropped_mass=args.max_reuse_selection_dropped_mass, + 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..28f6a7eed9e --- /dev/null +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -0,0 +1,299 @@ +# 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.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. + from modelopt.torch.sparsity.attention_sparsity.calibration.ruler_dataset import ( + RulerDatasetBuilder, + ) + + 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() + 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 + config_json.write_text(json.dumps(config, indent=2)) + 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() + + # 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 + + 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) + + 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/collect_mask_reuse.py b/examples/vllm_serve/collect_mask_reuse.py new file mode 100644 index 00000000000..54d830537c8 --- /dev/null +++ b/examples/vllm_serve/collect_mask_reuse.py @@ -0,0 +1,598 @@ +# 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. + +"""Collect mask-reuse calibration observations through the vLLM V1 path. + +ModelOpt derives each trial threshold from an existing vanilla skip-softmax +``(a, b)`` fit, arms the custom backend with one exact prompt/target +invocation, and merges only raw sufficient statistics returned by every TP +rank. No promoted reuse policy is loaded during collection. + +Prompt JSONL schema (one object per line):: + + {"split":"calibration", "partition":"development", "inner_fold":0, + "prompt_id":"p0", "source":"ruler/niah", "source_group_sha256":"...", + "prompt":"...", "min_kv_tokens":8192, "max_kv_tokens":65536} + +Usage:: + + python examples/vllm_serve/collect_mask_reuse.py /path/to/checkpoint \ + --model-id Nemotron-3-Ultra \ + --checkpoint-manifest-sha256 012345... \ + --plan nemotron3_ultra_stride2 \ + --fa4-source /path/to/extracted-fa4-runtime-source \ + --fa4-source-manifest /path/to/fa4-source-manifest.json \ + --fa4-source-manifest-sha256 abcdef... \ + --fa4-commit 4c40766b... \ + --prompts-jsonl prompts.jsonl \ + --vanilla-config /path/to/config.json \ + --target-sparsities 0.5 0.6 0.7 \ + --output compact-captures.jsonl +""" + +from __future__ import annotations + +import argparse +import importlib.metadata +import json +import os +import stat +import sys +import tempfile +from hashlib import sha256 +from pathlib import Path +from typing import cast + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + read_stable_file_snapshot, + verify_checkpoint_manifest, +) +from modelopt.torch.sparsity.attention_sparsity.calibration.source_manifest import ( + SourceManifestError, + VerifiedSourceManifest, + verify_source_manifest, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.mask_reuse_capture import ( + MAX_QUERY_CHUNK_TOKENS, + CaptureContractError, + build_capture_invocation, + canonical_json_sha256, + merge_rank_captures, + parse_prompt_specs_jsonl, + parse_vanilla_prefill_fit, + validate_begin_acks, + validate_capture_statuses, +) + +CAPTURE_ENV = "MASK_REUSE_FA4_CALIBRATION_CAPTURE" +PLAN_ENV = "MASK_REUSE_FA4_PLAN" +CHECKPOINT_ENV = "MASK_REUSE_FA4_CHECKPOINT_MANIFEST_SHA256" +DENSE_SHADOW_ENV = "MASK_REUSE_FA4_CAPTURE_DENSE_SHADOW" +_POLICY_ENVS = ( + "MASK_REUSE_FA4_POLICY", + "MASK_REUSE_FA4_POLICY_SHA256", +) +_FORBIDDEN_ENGINE_KWARGS = frozenset( + { + "additional_config", + "attention_backend", + "decode_context_parallel_size", + "enable_chunked_prefill", + "enable_prefix_caching", + "enforce_eager", + "kv_cache_dtype", + "kv_transfer_config", + "max_num_batched_tokens", + "max_num_seqs", + "pipeline_parallel_size", + "quantization", + "speculative_config", + "worker_cls", + } +) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Collect ModelOpt mask-reuse sufficient statistics through vLLM V1" + ) + parser.add_argument("model", help="HF checkpoint path loaded by vLLM") + parser.add_argument( + "--model-id", + default=None, + help="Stable model name stored in observations (default: the model argument)", + ) + parser.add_argument( + "--checkpoint-manifest-sha256", + required=True, + help="Separately pinned SHA256 of the checkpoint manifest", + ) + parser.add_argument("--plan", required=True, help="Explicit mask-reuse layer-plan preset") + parser.add_argument( + "--fa4-source", + required=True, + help="FlashAttention source tree extracted from the exact witnessed git archive", + ) + parser.add_argument( + "--fa4-source-manifest", + required=True, + help="Canonical full-tree witness generated with create_fa4_source_witness.py", + ) + parser.add_argument( + "--fa4-source-manifest-sha256", + required=True, + help="Separately pinned SHA256 of --fa4-source-manifest", + ) + parser.add_argument( + "--fa4-commit", + required=True, + help="Separately pinned 40-hex FlashAttention Git commit", + ) + parser.add_argument("--prompts-jsonl", required=True, help="Strict prompt-plan JSONL") + parser.add_argument( + "--vanilla-config", + required=True, + help="ModelOpt config containing the calibrated prefill skip-softmax (a, b) fit", + ) + parser.add_argument( + "--target-sparsities", + type=float, + nargs="+", + required=True, + help="Preregistered target-sparsity menu evaluated for every prompt", + ) + parser.add_argument("--output", required=True, help="Compact normalized capture JSONL") + parser.add_argument( + "--output-manifest", + default=None, + help="Capture provenance JSON (default: .manifest.json)", + ) + parser.add_argument("--max-model-len", type=int, default=None) + parser.add_argument("--tensor-parallel-size", type=int, default=1) + parser.add_argument("--gpu-memory-utilization", type=float, default=None) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument( + "--validate-dense-output", + action="store_true", + help="Bitwise-compare every armed capture layer with a second pinned dense FA4 call", + ) + parser.add_argument( + "--engine-kwargs", + default=None, + help="JSON object of non-contract vLLM kwargs (for example hybrid-model options)", + ) + return parser + + +def _target_menu(values: list[float]) -> tuple[float, ...]: + menu = tuple(sorted(set(values))) + if not menu or any(not 0.0 < value < 1.0 for value in menu): + raise CaptureContractError("--target-sparsities values must be finite and in (0, 1)") + return menu + + +def _engine_kwargs(args: argparse.Namespace) -> dict[str, object]: + extra: dict[str, object] = {} + if args.engine_kwargs is not None: + raw = json.loads(args.engine_kwargs) + if not isinstance(raw, dict): + raise CaptureContractError("--engine-kwargs must be a JSON object") + conflicts = _FORBIDDEN_ENGINE_KWARGS & raw.keys() + if conflicts: + raise CaptureContractError( + f"--engine-kwargs cannot override capture/precision controls: {sorted(conflicts)}" + ) + extra.update(raw) + extra.update( + { + "model": args.model, + "worker_cls": ( + "modelopt.torch.sparsity.attention_sparsity.plugins." + "vllm_mask_reuse_capture.MaskReuseCaptureWorker" + ), + "attention_backend": "CUSTOM", + "dtype": "bfloat16", + "enforce_eager": True, + "enable_prefix_caching": False, + "enable_chunked_prefill": True, + "max_num_batched_tokens": MAX_QUERY_CHUNK_TOKENS, + "max_num_seqs": 1, + "disable_cascade_attn": True, + "pipeline_parallel_size": 1, + "decode_context_parallel_size": 1, + } + ) + if args.max_model_len is not None: + if args.max_model_len <= 0: + raise CaptureContractError("--max-model-len must be positive") + extra["max_model_len"] = args.max_model_len + if args.tensor_parallel_size <= 0: + raise CaptureContractError("--tensor-parallel-size must be positive") + extra["tensor_parallel_size"] = args.tensor_parallel_size + if args.gpu_memory_utilization is not None: + if not 0.0 < args.gpu_memory_utilization <= 1.0: + raise CaptureContractError("--gpu-memory-utilization must be in (0, 1]") + extra["gpu_memory_utilization"] = args.gpu_memory_utilization + if args.trust_remote_code: + extra["trust_remote_code"] = True + return extra + + +def _canonical_capture_line(capture: dict[str, object]) -> bytes: + return ( + json.dumps(capture, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + + +def _verify_fa4_source( + fa4_source: str, + fa4_source_manifest: str, + fa4_source_manifest_sha256: str, + fa4_commit: str, +) -> VerifiedSourceManifest: + try: + verified = verify_source_manifest( + fa4_source, + fa4_source_manifest, + expected_manifest_sha256=fa4_source_manifest_sha256, + expected_commit=fa4_commit, + expected_source_kind="flash-attention-4", + ) + except SourceManifestError as error: + raise CaptureContractError(f"--fa4-source verification failed: {error}") from error + required = ( + verified.source_root / "flash_attn/cute/interface.py", + verified.source_root / "flash_attn/cute/block_sparsity.py", + ) + try: + required_are_regular = all( + stat.S_ISREG(path.stat(follow_symlinks=False).st_mode) for path in required + ) + except OSError: + required_are_regular = False + if not required_are_regular: + raise CaptureContractError( + "--fa4-source must contain flash_attn/cute/interface.py and block_sparsity.py" + ) + return verified + + +def _configure_capture_environment( + plan: str, + fa4_source: str, + fa4_source_manifest: str, + fa4_source_manifest_sha256: str, + fa4_commit: str, + checkpoint_manifest_sha256: str, + *, + validate_dense_output: bool, +) -> VerifiedSourceManifest: + verified = _verify_fa4_source( + fa4_source, + fa4_source_manifest, + fa4_source_manifest_sha256, + fa4_commit, + ) + plugins = [ + entry + for entry in importlib.metadata.entry_points(group="vllm.general_plugins") + if entry.name == "mask_reuse_fa4" + ] + if len(plugins) != 1 or plugins[0].value != "mask_reuse_vllm.plugin:register": + raise CaptureContractError( + "the mask_reuse_fa4 vLLM plugin entry point is missing or ambiguous" + ) + os.environ[CAPTURE_ENV] = "1" + os.environ[PLAN_ENV] = plan + os.environ[CHECKPOINT_ENV] = checkpoint_manifest_sha256 + os.environ[DENSE_SHADOW_ENV] = "1" if validate_dense_output else "0" + os.environ["PYTHONDONTWRITEBYTECODE"] = "1" + sys.dont_write_bytecode = True + os.environ["MASK_REUSE_FA4_SOURCE"] = str(verified.source_root) + os.environ["MASK_REUSE_FA4_FORCE_DENSE"] = "0" + os.environ["VLLM_PLUGINS"] = "mask_reuse_fa4" + # Collection is deliberately policy-free. Remove inherited serving + # settings so a stale deployment policy cannot become threshold authority. + for name in _POLICY_ENVS: + os.environ.pop(name, None) + return verified + + +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 _publish_no_clobber(temporary: Path, destination: Path) -> tuple[int, int]: + """Atomically link a complete temp file only when destination is absent.""" + + observed = temporary.stat(follow_symlinks=False) + if observed.st_ino == 0: + raise RuntimeError("capture 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("capture destination changed during publication") + temporary.unlink() + _fsync_directory(destination.parent) + except BaseException: + _unlink_if_identity(destination, identity) + raise + return identity + + +def _unlink_if_identity(path: Path, identity: tuple[int, int]) -> None: + """Rollback only a file that is still the inode published by this process.""" + + if identity[1] == 0: + return + try: + stat = path.stat(follow_symlinks=False) + except FileNotFoundError: + return + if stat.st_ino != 0 and (stat.st_dev, stat.st_ino) == identity: + path.unlink() + _fsync_directory(path.parent) + + +def run(args: argparse.Namespace) -> tuple[Path, Path]: + prompt_snapshot = read_stable_file_snapshot(args.prompts_jsonl, label="prompt plan") + prompt_plan_sha256 = prompt_snapshot.sha256 + prompts = parse_prompt_specs_jsonl(prompt_snapshot.payload) + vanilla_snapshot = read_stable_file_snapshot(args.vanilla_config, label="vanilla config") + vanilla_config_sha256 = vanilla_snapshot.sha256 + threshold_scale_factor = parse_vanilla_prefill_fit(vanilla_snapshot.payload) + targets = _target_menu(args.target_sparsities) + checkpoint = verify_checkpoint_manifest(args.model, expected_model=args.model_id) + if checkpoint.sha256 != args.checkpoint_manifest_sha256: + raise CaptureContractError( + "checkpoint manifest does not match --checkpoint-manifest-sha256" + ) + model_id = checkpoint.model + output_path = Path(args.output) + manifest_path = ( + Path(args.output_manifest) + if args.output_manifest is not None + else Path(str(output_path) + ".manifest.json") + ) + if output_path.resolve() == manifest_path.resolve(): + raise CaptureContractError("observation and manifest output paths must differ") + if output_path.exists() or manifest_path.exists(): + raise CaptureContractError("capture outputs already exist; refusing to overwrite them") + output_path.parent.mkdir(parents=True, exist_ok=True) + manifest_path.parent.mkdir(parents=True, exist_ok=True) + + fa4_source = _configure_capture_environment( + args.plan, + args.fa4_source, + args.fa4_source_manifest, + args.fa4_source_manifest_sha256, + args.fa4_commit, + checkpoint.sha256, + validate_dense_output=args.validate_dense_output, + ) + # Import after setting the gate: worker subprocesses inherit the exact + # capture environment and never enter policy-backed serving mode. + from vllm import LLM, SamplingParams + + engine_kwargs = _engine_kwargs(args) + llm = LLM(**engine_kwargs) + loaded_checkpoint = verify_checkpoint_manifest(args.model, expected_model=model_id) + if loaded_checkpoint != checkpoint: + raise CaptureContractError( + "checkpoint identity changed while vLLM loaded the model; no capture was started" + ) + loaded_fa4_source = _verify_fa4_source( + args.fa4_source, + args.fa4_source_manifest, + args.fa4_source_manifest_sha256, + args.fa4_commit, + ) + if loaded_fa4_source != fa4_source: + raise CaptureContractError( + "FA4 source identity changed while vLLM loaded the model; no capture was started" + ) + statuses = llm.collective_rpc("mask_reuse_capture_status") + validate_capture_statuses(statuses) + tokenizer = llm.get_tokenizer() + sampling = SamplingParams(temperature=0.0, max_tokens=1, ignore_eos=True) + + capture_manifests: list[dict[str, object]] = [] + seen_sources: dict[str, tuple[str, str, tuple[int, int | None]]] = {} + capture_digest = sha256() + capture_count = 0 + candidate_cell_count = 0 + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=output_path.parent, + prefix=f".{output_path.name}.", + suffix=".tmp", + delete=False, + ) as temporary: + temporary_path = Path(temporary.name) + for prompt in prompts: + token_ids = tokenizer.encode(prompt.prompt, add_special_tokens=True) + if not isinstance(token_ids, list) or any( + type(token) is not int for token in token_ids + ): + raise CaptureContractError( + "tokenizer.encode must return a list of integer token IDs" + ) + if args.max_model_len is not None and len(token_ids) + 1 > args.max_model_len: + raise CaptureContractError( + f"prompt {prompt.prompt_id!r} plus one output token exceeds --max-model-len" + ) + for target in targets: + invocation = build_capture_invocation( + model=model_id, + checkpoint_manifest_sha256=checkpoint.sha256, + prompt=prompt, + prompt_token_ids=token_ids, + target_sparsity=target, + threshold_scale_factor=threshold_scale_factor, + ) + fingerprint = str(invocation["source_capture_sha256"]) + identity = (prompt.split, prompt.prompt_id, prompt.bucket) + previous = seen_sources.setdefault(fingerprint, identity) + if previous != identity: + raise CaptureContractError( + "the same tokenized source is assigned to multiple prompt captures: " + f"{previous} and {identity}" + ) + acknowledgements = llm.collective_rpc( + "mask_reuse_capture_begin", args=(invocation,) + ) + validate_begin_acks(acknowledgements, invocation) + # Passing token IDs makes sample_length and source_capture_sha256 + # identical to the request that reaches the vLLM scheduler. + llm.generate(token_ids, sampling, use_tqdm=False) + merged = merge_rank_captures( + llm.collective_rpc("mask_reuse_capture_drain"), invocation + ) + line = _canonical_capture_line(merged.capture) + temporary.write(line) + capture_digest.update(line) + capture_count += 1 + candidate_cell_count += cast("int", merged.manifest["candidate_cell_count"]) + capture_manifests.append(merged.manifest) + temporary.flush() + os.fsync(temporary.fileno()) + except BaseException: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + raise + + try: + prompt_final = read_stable_file_snapshot(args.prompts_jsonl, label="prompt plan") + vanilla_final = read_stable_file_snapshot(args.vanilla_config, label="vanilla config") + final_fa4_source = _verify_fa4_source( + args.fa4_source, + args.fa4_source_manifest, + args.fa4_source_manifest_sha256, + args.fa4_commit, + ) + except BaseException: + assert temporary_path is not None + temporary_path.unlink(missing_ok=True) + raise + if ( + prompt_final.sha256 != prompt_plan_sha256 + or vanilla_final.sha256 != vanilla_config_sha256 + or final_fa4_source != fa4_source + ): + assert temporary_path is not None + temporary_path.unlink(missing_ok=True) + raise CaptureContractError( + "prompt plan, vanilla calibration, or FA4 source changed during capture; " + "evidence was discarded" + ) + + capture_sha256 = capture_digest.hexdigest() + manifest = { + "capture_manifest_schema_version": 4, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v4", + "model": model_id, + "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": args.plan, + "fa4_source": str(fa4_source.source_root), + "fa4_source_commit": fa4_source.git_commit, + "fa4_source_git_tree": fa4_source.git_tree, + "fa4_source_git_archive_sha256": fa4_source.git_archive_sha256, + "fa4_source_manifest_path": str(fa4_source.manifest_path), + "fa4_source_manifest_sha256": fa4_source.manifest_sha256, + "fa4_source_directory_count": fa4_source.directory_count, + "fa4_source_file_count": fa4_source.file_count, + "fa4_source_total_size_bytes": fa4_source.total_size_bytes, + "engine_kwargs": engine_kwargs, + "dense_shadow_validation_requested": args.validate_dense_output, + "target_sparsity_hex": [target.hex() for target in targets], + "vanilla_threshold_scale_factor": threshold_scale_factor, + "vanilla_fit_sha256": canonical_json_sha256(threshold_scale_factor), + "vanilla_config_file_sha256": vanilla_config_sha256, + "prompt_plan_file_sha256": prompt_plan_sha256, + "compact_capture_file_sha256": capture_sha256, + "capture_count": capture_count, + "candidate_cell_count": candidate_cell_count, + "captures": capture_manifests, + } + manifest_bytes = ( + json.dumps(manifest, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + manifest_temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=manifest_path.parent, + prefix=f".{manifest_path.name}.", + suffix=".tmp", + delete=False, + ) as temporary: + manifest_temporary = Path(temporary.name) + temporary.write(manifest_bytes) + temporary.flush() + os.fsync(temporary.fileno()) + assert temporary_path is not None and manifest_temporary is not None + capture_identity = _publish_no_clobber(temporary_path, output_path) + temporary_path = None + try: + _publish_no_clobber(manifest_temporary, manifest_path) + manifest_temporary = None + except BaseException: + _unlink_if_identity(output_path, capture_identity) + raise + except FileExistsError as error: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + if manifest_temporary is not None: + manifest_temporary.unlink(missing_ok=True) + raise CaptureContractError("capture destination appeared during publication") from error + except BaseException: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + if manifest_temporary is not None: + manifest_temporary.unlink(missing_ok=True) + raise + print( + f"[ModelOpt] Wrote {capture_count} compact captures " + f"({candidate_cell_count} candidate cells) to {output_path.resolve()}" + ) + print(f"[ModelOpt] compact_capture_file_sha256={capture_sha256}") + print(f"[ModelOpt] Wrote capture manifest to {manifest_path.resolve()}") + return output_path, manifest_path + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + run(args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/vllm_serve/create_checkpoint_manifest.py b/examples/vllm_serve/create_checkpoint_manifest.py new file mode 100644 index 00000000000..8fe276f0d0f --- /dev/null +++ b/examples/vllm_serve/create_checkpoint_manifest.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Create a deterministic, no-clobber checkpoint manifest for calibration.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + create_checkpoint_manifest, +) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(allow_abbrev=False) + parser.add_argument("checkpoint", type=Path) + parser.add_argument("--model-id", required=True) + args = parser.parse_args(argv) + try: + manifest = create_checkpoint_manifest(args.checkpoint, model=args.model_id) + except (OSError, ValueError) as error: + parser.error(str(error)) + print(f"[ModelOpt] Wrote {manifest.manifest_path}") + print(f"CHECKPOINT_MANIFEST_SHA256={manifest.sha256}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/vllm_serve/create_fa4_source_witness.py b/examples/vllm_serve/create_fa4_source_witness.py new file mode 100644 index 00000000000..e32cc4f3823 --- /dev/null +++ b/examples/vllm_serve/create_fa4_source_witness.py @@ -0,0 +1,57 @@ +# 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 an exact FlashAttention source archive and its runtime witness.""" + +from __future__ import annotations + +import argparse + +from modelopt.torch.sparsity.attention_sparsity.calibration.source_manifest import ( + create_source_manifest_from_git_archive, +) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Create a Git-free runtime witness for one exact FlashAttention commit" + ) + parser.add_argument("checkout", help="Clean FlashAttention Git checkout") + parser.add_argument("--expected-commit", required=True, help="Required 40-hex HEAD commit") + parser.add_argument("--archive-output", required=True, help="New exact git-archive tar path") + parser.add_argument( + "--manifest-output", required=True, help="New canonical source-witness JSON path" + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + generated = create_source_manifest_from_git_archive( + args.checkout, + expected_commit=args.expected_commit, + source_kind="flash-attention-4", + archive_output=args.archive_output, + manifest_output=args.manifest_output, + ) + print(f"git_commit={generated.git_commit}") + print(f"git_tree={generated.git_tree}") + print(f"git_archive_sha256={generated.git_archive_sha256}") + print(f"source_manifest_sha256={generated.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..b8e9d6475c7 100644 --- a/modelopt/torch/kernels/sparsity/attention/calibrate.py +++ b/modelopt/torch/kernels/sparsity/attention/calibrate.py @@ -28,7 +28,12 @@ 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 +69,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 +144,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 +209,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 @@ -212,6 +277,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 +288,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 +315,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,6 +363,11 @@ def attention_calibrate( b_seq_len_k = b_seq_len b_start_loc_k = b_start_loc + # Paged mode: KV positions come from block_table, so the contiguous KV + # offsets are unused. Provide a dummy so Triton can compile the tl.load. + if b_start_loc_k is None: + b_start_loc_k = torch.zeros_like(b_start_loc) + num_thresholds = len(threshold_trials) # Scores already include sm_scale and LOG2E; convert lambda to log2 space only. @@ -304,6 +391,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 +437,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/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..397d61c152e --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py @@ -0,0 +1,1408 @@ +# 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 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" +_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 + + +@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, + "constraint_violation_count": self.violations, + "constraint_violation_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, + max_reuse_selection_dropped_mass: 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]] = [] + 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, + ) + choices: dict[ConsumerHead, _Choice] = {} + total_retained = 0 + fallback_count = 0 + for target in targets: + feasible: list[tuple[int, int]] = [] + for donor in range(global_num_heads): + candidate_rows = [ + index.observations[ + ( + "calibration", + prompt, + target_sparsity, + target[0], + target[1], + donor, + ) + ] + for prompt in prompts + ] + if all( + row.dropped_mass <= max_reuse_selection_dropped_mass for row in candidate_rows + ): + feasible.append((sum(row.retained_tiles for row in candidate_rows), donor)) + if feasible: + retained, donor = min(feasible) + choice = _Choice(donor, False, retained) + else: + fallback_count += 1 + retained = sum( + index.eligible[("calibration", prompt, target)] for prompt in prompts + ) + choice = _Choice(0, True, retained) + choices[target] = choice + total_retained += retained + signature = tuple(choices[target].donor_head for target in targets) + combined_tile_cost = 2 * total_retained + anchor.retained_tiles + frontier.append( + { + "target_sparsity": target_sparsity, + "anchor_safe": anchor.violations == 0, + "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 anchor.violations == 0: + rank = ( + combined_tile_cost, + fallback_count, + target_sparsity, + signature, + ) + candidates.append((rank, _Selection(target_sparsity, choices, ()))) + if not 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", + ) + _, selected = min(candidates, key=lambda item: item[0]) + return _Selection(selected.target_sparsity, selected.choices, tuple(frontier)) + + +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, + max_reuse_dropped_mass: float, + max_reuse_selection_dropped_mass: float | None = None, + source_provenance: Mapping[str, object] | None = None, +) -> dict[str, object]: + """Select and evaluate a fail-closed schema-v3 legacy candidate.""" + 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_bound = _number( + max_reuse_dropped_mass, "max_reuse_dropped_mass", minimum=0.0, maximum=1.0 + ) + selection_bound = ( + reuse_bound + if max_reuse_selection_dropped_mass is None + else _number( + max_reuse_selection_dropped_mass, + "max_reuse_selection_dropped_mass", + minimum=0.0, + maximum=1.0, + ) + ) + if selection_bound > reuse_bound: + raise MaskReuseCalibrationError( + "max_reuse_selection_dropped_mass must not exceed max_reuse_dropped_mass" + ) + 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 + 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, + max_reuse_selection_dropped_mass=selection_bound, + ) + 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_bound + ) + reuse_heldout = _evaluate_reuse(selection, index, split="heldout", maximum=reuse_bound) + 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 + 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": "sparse" + if selection.target_sparsity is not None + else selection.exact_reason, + "target_sparsity_frontier": list(selection.frontier), + "fallback_head_count": fallback_count, + "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": { + "metric": "per_prompt_candidate_reuse_dropped_mass", + "comparison": "<=", + "selection_maximum": selection_bound, + "evaluation_maximum": reuse_bound, + }, + }, + "tie_breaks": [ + "reject target sparsities exceeding the calibration anchor bound", + "minimum equal-BMM combined tile cost 2*A_R + A_A", + "fewest exact-fallback 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)), + "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, + "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..9d549a7496b --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py @@ -0,0 +1,1211 @@ +# 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. Memory scales with the current and immediately prior +capture plus ``O(targets * consumers * H^2)`` aggregate arrays rather than with +every prompt-level candidate cell. +""" + +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 modelopt + +from .checkpoint_manifest import VerifiedCheckpointManifest +from .mask_reuse import ( + _CALIBRATION_PROTOCOL, + _DEPLOYMENT_GEOMETRY_CONTRACT, + 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: + max_mass: Mapping[float, np.ndarray] + retained_by_anchor: Mapping[float, np.ndarray] + eligible_sum: dict[float, int] + anchor_evaluation: Mapping[float, _AnchorEvaluation] + + +def _selection_pass( + source: CompactMaskReuseCaptureSource, + dataset: _Dataset, + *, + max_anchor_dropped_mass: float, + max_reuse_selection_dropped_mass: 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( + max_mass={ + target: np.full( + ( + len(dataset.consumer_layers), + dataset.global_num_heads, + dataset.global_num_heads, + ), + -np.inf, + dtype=np.float64, + ) + 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(): + matrix = np.asarray(stats.dropped_mass, dtype=np.float64) + np.maximum( + accumulator.max_mass[target][consumer_index[layer]], + matrix, + out=accumulator.max_mass[target][consumer_index[layer]], + ) + + selections: dict[Bucket, _Selection] = {} + for bucket in dataset.buckets: + accumulator = accumulators[bucket] + frontier: list[Mapping[str, object]] = [] + candidates: list[tuple[tuple[object, ...], _Selection]] = [] + for target in dataset.menus[bucket]: + anchor_evaluation = accumulator.anchor_evaluation[target] + choices: dict[ConsumerHead, _Choice] = {} + total_retained = 0 + fallback_count = 0 + for layer in dataset.consumer_layers: + layer_index = consumer_index[layer] + anchor = dataset.nearest[layer] + retained = accumulator.retained_by_anchor[target][anchor_index[anchor]] + for head in range(dataset.global_num_heads): + feasible = np.flatnonzero( + accumulator.max_mass[target][layer_index, head] + <= max_reuse_selection_dropped_mass + ) + if feasible.size: + donor = min((int(retained[item]), int(item)) for item in feasible)[1] + choice = _Choice(donor, False, int(retained[donor])) + else: + fallback_count += 1 + choice = _Choice(0, True, accumulator.eligible_sum[target]) + choices[(layer, head)] = choice + total_retained += choice.retained_tiles + signature = tuple(choices[item].donor_head for item in dataset.targets) + combined_tile_cost = 2 * total_retained + anchor_evaluation.retained_tiles + frontier.append( + { + "target_sparsity": target, + "anchor_safe": anchor_evaluation.violations == 0, + "retained_reuse_tiles": total_retained, + "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 anchor_evaluation.violations == 0: + rank = ( + combined_tile_cost, + fallback_count, + target, + signature, + ) + candidates.append((rank, _Selection(target, choices, ()))) + if candidates: + _, selected = min(candidates, key=lambda item: item[0]) + selections[bucket] = _Selection( + selected.target_sparsity, + selected.choices, + tuple(frontier), + ) + 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", + ) + return selections + + +def _evaluation_pass( + source: CompactMaskReuseCaptureSource, + dataset: _Dataset, + selections: Mapping[Bucket, _Selection], + *, + max_anchor_dropped_mass: float, + max_reuse_dropped_mass: 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 > max_reuse_dropped_mass) + + 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, + max_reuse_dropped_mass: float, + max_reuse_selection_dropped_mass: float | None = None, + source_provenance: Mapping[str, object] | None = None, +) -> dict[str, object]: + """Select and evaluate a fail-closed schema-v3 candidate from compact captures.""" + 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_bound = _number( + max_reuse_dropped_mass, "max_reuse_dropped_mass", minimum=0.0, maximum=1.0 + ) + selection_bound = ( + reuse_bound + if max_reuse_selection_dropped_mass is None + else _number( + max_reuse_selection_dropped_mass, + "max_reuse_selection_dropped_mass", + minimum=0.0, + maximum=1.0, + ) + ) + if selection_bound > reuse_bound: + raise MaskReuseCalibrationError( + "max_reuse_selection_dropped_mass must not exceed max_reuse_dropped_mass" + ) + selections = _selection_pass( + source, + dataset, + max_anchor_dropped_mass=anchor_bound, + max_reuse_selection_dropped_mass=selection_bound, + ) + reuse_evaluations, anchor_evaluations = _evaluation_pass( + source, + dataset, + selections, + max_anchor_dropped_mass=anchor_bound, + max_reuse_dropped_mass=reuse_bound, + ) + 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 + 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 + 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": ( + "sparse" if selection.target_sparsity is not None else selection.exact_reason + ), + "target_sparsity_frontier": list(selection.frontier), + "fallback_head_count": fallback_count, + "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": { + "metric": "per_prompt_candidate_reuse_dropped_mass", + "comparison": "<=", + "selection_maximum": selection_bound, + "evaluation_maximum": reuse_bound, + }, + } + 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", + "minimum equal-BMM combined tile cost 2*A_R + A_A", + "fewest exact-fallback 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, + "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, + "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/calibration/source_manifest.py b/modelopt/torch/sparsity/attention_sparsity/calibration/source_manifest.py new file mode 100644 index 00000000000..837bbf43ca3 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/source_manifest.py @@ -0,0 +1,613 @@ +# 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. + +"""Git-free verification of runtime source extracted from an exact Git archive.""" + +from __future__ import annotations + +import json +import os +import shutil +import stat +import subprocess # nosec B404 - used only by the explicit outside-container generator +import tarfile +import tempfile +import unicodedata +from dataclasses import dataclass +from hashlib import sha256 +from io import BytesIO +from pathlib import Path, PurePosixPath +from typing import cast + +from .checkpoint_manifest import ( + CheckpointManifestError, + read_stable_file_snapshot, + stable_file_sha256, +) + +__all__ = [ + "GeneratedSourceManifest", + "SourceManifestError", + "VerifiedSourceManifest", + "create_source_manifest_from_git_archive", + "verify_source_manifest", +] + +_ARCHIVE_SCOPE = "flash_attn" +_MANIFEST_FIELDS = frozenset( + { + "source_manifest_schema_version", + "source_kind", + "git_commit", + "git_tree", + "git_archive_sha256", + "archive_scope", + "directories", + "files", + } +) +_FILE_FIELDS = frozenset({"path", "mode", "size_bytes", "sha256"}) +_FILE_ATTRIBUTE_REPARSE_POINT = getattr(stat, "FILE_ATTRIBUTE_REPARSE_POINT", 0x400) + + +class SourceManifestError(ValueError): + """Raised when source bytes do not match their sealed Git-archive witness.""" + + +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 SourceManifestError(f"source 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 SourceManifestError( + f"{label} fields do not match the schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + + +def _canonical_json_bytes(value: object) -> bytes: + return ( + json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + + +def _canonical_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 SourceManifestError(f"{label} must be non-empty canonical NFC text") + return value + + +def _hex_digest(value: object, length: int, label: str) -> str: + if ( + not isinstance(value, str) + or len(value) != length + or any(character not in "0123456789abcdef" for character in value) + ): + raise SourceManifestError(f"{label} must be {length} lowercase hexadecimal characters") + return value + + +def _relative_path(value: object, label: str) -> str: + text = _canonical_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 path.parts[0] == ".git" + ): + raise SourceManifestError(f"{label} must be a canonical non-.git relative POSIX path") + return text + + +def _scoped_path(value: object, label: str) -> str: + text = _relative_path(value, label) + if PurePosixPath(text).parts[0] != _ARCHIVE_SCOPE: + raise SourceManifestError(f"{label} must be within {_ARCHIVE_SCOPE!r}") + return text + + +def _is_reparse_point(value: os.stat_result) -> bool: + return bool(getattr(value, "st_file_attributes", 0) & _FILE_ATTRIBUTE_REPARSE_POINT) + + +def _file_mode(value: os.stat_result) -> str: + return "100755" if value.st_mode & 0o111 else "100644" + + +@dataclass(frozen=True, slots=True) +class _TreeSnapshot: + directories: tuple[str, ...] + files: tuple[dict[str, object], ...] + + +def _source_tree_snapshot(root: Path) -> _TreeSnapshot: + directories: set[str] = set() + files: dict[str, dict[str, object]] = {} + pending: list[tuple[Path, str]] = [(root, "")] + while pending: + directory, directory_relative = pending.pop() + try: + observed_directory = directory.stat(follow_symlinks=False) + if not stat.S_ISDIR(observed_directory.st_mode) or _is_reparse_point( + observed_directory + ): + raise SourceManifestError( + f"source directory {directory_relative or '.'!r} is not stable" + ) + with os.scandir(directory) as iterator: + children = sorted(iterator, key=lambda entry: entry.name) + except OSError as error: + raise SourceManifestError( + f"could not traverse source directory {directory_relative or '.'!r}" + ) from error + for child in children: + if not directory_relative and child.name == ".git": + continue + relative = _relative_path( + f"{directory_relative}/{child.name}".lstrip("/"), "source path" + ) + path = Path(child.path) + try: + observed = child.stat(follow_symlinks=False) + except OSError as error: + raise SourceManifestError(f"could not inspect source path {relative!r}") from error + if stat.S_ISLNK(observed.st_mode) or _is_reparse_point(observed): + raise SourceManifestError(f"source path {relative!r} must not be a link") + if stat.S_ISDIR(observed.st_mode): + directories.add(relative) + pending.append((path, relative)) + elif stat.S_ISREG(observed.st_mode): + try: + digest = stable_file_sha256(path, label=f"source file {relative!r}") + except CheckpointManifestError as error: + raise SourceManifestError( + f"could not hash stable source file {relative!r}" + ) from error + files[relative] = { + "path": relative, + "mode": _file_mode(observed), + "size_bytes": observed.st_size, + "sha256": digest, + } + else: + raise SourceManifestError(f"source path {relative!r} is not regular") + return _TreeSnapshot( + directories=tuple(sorted(directories)), + files=tuple(files[path] for path in sorted(files)), + ) + + +def _parse_manifest(payload: bytes) -> dict[str, object]: + try: + raw = json.loads(payload, object_pairs_hook=_strict_object) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise SourceManifestError("source manifest is not strict UTF-8 JSON") from error + if not isinstance(raw, dict): + raise SourceManifestError("source manifest must be a JSON object") + _exact_fields(raw, _MANIFEST_FIELDS, "source manifest") + if ( + type(raw["source_manifest_schema_version"]) is not int + or raw["source_manifest_schema_version"] != 1 + ): + raise SourceManifestError("source_manifest_schema_version must be 1") + if payload != _canonical_json_bytes(raw): + raise SourceManifestError("source manifest bytes are not canonical JSON") + _canonical_text(raw["source_kind"], "source manifest.source_kind") + _hex_digest(raw["git_commit"], 40, "source manifest.git_commit") + _hex_digest(raw["git_tree"], 40, "source manifest.git_tree") + _hex_digest(raw["git_archive_sha256"], 64, "source manifest.git_archive_sha256") + if raw["archive_scope"] != _ARCHIVE_SCOPE: + raise SourceManifestError(f"source manifest.archive_scope must be {_ARCHIVE_SCOPE!r}") + + raw_directories = raw["directories"] + if not isinstance(raw_directories, list): + raise SourceManifestError("source manifest.directories must be a list") + directories = [ + _scoped_path(value, f"source manifest.directories[{index}]") + for index, value in enumerate(raw_directories) + ] + if directories != sorted(set(directories)): + raise SourceManifestError("source manifest directories must be unique and sorted") + + raw_files = raw["files"] + if not isinstance(raw_files, list) or not raw_files: + raise SourceManifestError("source manifest.files must be a non-empty list") + paths: list[str] = [] + for index, item in enumerate(raw_files): + label = f"source manifest.files[{index}]" + if not isinstance(item, dict): + raise SourceManifestError(f"{label} must be an object") + _exact_fields(item, _FILE_FIELDS, label) + paths.append(_scoped_path(item["path"], f"{label}.path")) + if item["mode"] not in {"100644", "100755"}: + raise SourceManifestError(f"{label}.mode must be 100644 or 100755") + size = item["size_bytes"] + if isinstance(size, bool) or not isinstance(size, int) or size < 0: + raise SourceManifestError(f"{label}.size_bytes must be an integer >= 0") + _hex_digest(item["sha256"], 64, f"{label}.sha256") + if paths != sorted(set(paths)): + raise SourceManifestError("source manifest file paths must be unique and sorted") + if set(paths) & set(directories): + raise SourceManifestError("source manifest paths cannot be both files and directories") + return raw + + +@dataclass(frozen=True, slots=True) +class VerifiedSourceManifest: + """Identity of an exact, fully enumerated runtime source tree.""" + + source_root: Path + manifest_path: Path + source_kind: str + git_commit: str + git_tree: str + git_archive_sha256: str + manifest_sha256: str + directory_count: int + file_count: int + total_size_bytes: int + + +def verify_source_manifest( + source: str | Path, + manifest: str | Path, + *, + expected_manifest_sha256: str, + expected_commit: str, + expected_source_kind: str, +) -> VerifiedSourceManifest: + """Verify every non-``.git`` source path without invoking Git.""" + root = Path(source).expanduser().resolve() + if not root.is_dir(): + raise SourceManifestError("source root must be a local directory") + expected_digest = _hex_digest(expected_manifest_sha256, 64, "expected source manifest SHA256") + expected_git_commit = _hex_digest(expected_commit, 40, "expected source Git commit") + expected_kind = _canonical_text(expected_source_kind, "expected source kind") + try: + manifest_snapshot = read_stable_file_snapshot(manifest, label="source manifest") + except CheckpointManifestError as error: + raise SourceManifestError("could not read stable source manifest") from error + if manifest_snapshot.sha256 != expected_digest: + raise SourceManifestError("source manifest does not match its expected SHA256") + raw = _parse_manifest(manifest_snapshot.payload) + if raw["source_kind"] != expected_kind: + raise SourceManifestError("source manifest kind does not match the expected source kind") + if raw["git_commit"] != expected_git_commit: + raise SourceManifestError("source manifest commit does not match the expected Git commit") + + actual = _source_tree_snapshot(root) + if list(actual.directories) != raw["directories"] or list(actual.files) != raw["files"]: + raise SourceManifestError("source tree does not exactly match its sealed manifest") + if _source_tree_snapshot(root) != actual: + raise SourceManifestError("source tree changed during verification") + return VerifiedSourceManifest( + source_root=root, + manifest_path=manifest_snapshot.path.resolve(), + source_kind=expected_kind, + git_commit=expected_git_commit, + git_tree=str(raw["git_tree"]), + git_archive_sha256=str(raw["git_archive_sha256"]), + manifest_sha256=manifest_snapshot.sha256, + directory_count=len(actual.directories), + file_count=len(actual.files), + total_size_bytes=sum(cast("int", item["size_bytes"]) for item in actual.files), + ) + + +@dataclass(frozen=True, slots=True) +class GeneratedSourceManifest: + """Hashes emitted by the outside-container Git-archive generator.""" + + git_commit: str + git_tree: str + git_archive_sha256: str + manifest_sha256: str + + +def _run_git(checkout: Path, arguments: list[str]) -> bytes: + executable = shutil.which("git") + if executable is None: + raise SourceManifestError("Git is required by the outside-container generator") + try: + return subprocess.run( + [executable, "-C", str(checkout), *arguments], # nosec B603 + check=True, + capture_output=True, + ).stdout + except (OSError, subprocess.CalledProcessError) as error: + raise SourceManifestError(f"Git command failed: {' '.join(arguments)}") from error + + +def _fsync_directory(path: Path) -> None: + directory_flag = getattr(os, "O_DIRECTORY", None) + if directory_flag is None: + return + descriptor = os.open(path, os.O_RDONLY | directory_flag) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +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 stat.S_ISREG(observed.st_mode) + and not _is_reparse_point(observed) + and (observed.st_dev, observed.st_ino) == identity + ): + path.unlink() + _fsync_directory(path.parent) + + +def _temporary_payload(destination: Path, payload: bytes) -> tuple[Path, tuple[int, int], int]: + destination.parent.mkdir(parents=True, exist_ok=True) + descriptor, name = tempfile.mkstemp( + dir=destination.parent, + prefix=f".{destination.name}.", + suffix=".tmp", + ) + temporary = Path(name) + identity: tuple[int, int] | None = None + try: + with os.fdopen(os.dup(descriptor), "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + opened = os.fstat(descriptor) + observed = temporary.stat(follow_symlinks=False) + identity = (opened.st_dev, opened.st_ino) + if ( + opened.st_ino == 0 + or not stat.S_ISREG(opened.st_mode) + or _is_reparse_point(observed) + or (observed.st_dev, observed.st_ino) != identity + ): + raise SourceManifestError("source artifact temporary file has no stable identity") + return temporary, identity, descriptor + except BaseException: + if identity is None: + try: + opened = os.fstat(descriptor) + if opened.st_ino != 0 and stat.S_ISREG(opened.st_mode): + identity = (opened.st_dev, opened.st_ino) + except OSError: + pass + if os.name == "nt": + os.close(descriptor) + if identity is not None: + _unlink_if_identity(temporary, identity) + if os.name != "nt": + os.close(descriptor) + raise + + +def _publish_no_clobber( + temporary: Path, destination: Path, identity: tuple[int, int], descriptor: int +) -> None: + observed = temporary.stat(follow_symlinks=False) + opened = os.fstat(descriptor) + if ( + observed.st_ino == 0 + or (observed.st_dev, observed.st_ino) != identity + or (opened.st_dev, opened.st_ino) != identity + ): + raise SourceManifestError("source artifact temporary file changed before publication") + os.link(temporary, destination, follow_symlinks=False) + published = destination.stat(follow_symlinks=False) + opened_after = os.fstat(descriptor) + if ( + published.st_ino == 0 + or not stat.S_ISREG(published.st_mode) + or _is_reparse_point(published) + or (published.st_dev, published.st_ino) != identity + or (opened_after.st_dev, opened_after.st_ino) != identity + ): + raise SourceManifestError("source artifact destination changed during publication") + + +def _publish_source_artifacts( + archive_path: Path, + archive: bytes, + manifest_path: Path, + manifest: bytes, +) -> None: + archive_temporary_path, archive_temporary_identity, archive_descriptor = _temporary_payload( + archive_path, archive + ) + archive_temporary: Path | None = archive_temporary_path + manifest_temporary: Path | None = None + manifest_temporary_identity: tuple[int, int] | None = None + manifest_descriptor: int | None = None + archive_identity: tuple[int, int] | None = None + manifest_identity: tuple[int, int] | None = None + try: + manifest_temporary, manifest_temporary_identity, manifest_descriptor = _temporary_payload( + manifest_path, manifest + ) + assert archive_temporary is not None + archive_identity = archive_temporary_identity + _publish_no_clobber( + archive_temporary, archive_path, archive_temporary_identity, archive_descriptor + ) + if os.name == "nt": + os.close(archive_descriptor) + archive_descriptor = None + _unlink_if_identity(archive_temporary, archive_temporary_identity) + _fsync_directory(archive_path.parent) + archive_temporary = None + manifest_identity = manifest_temporary_identity + _publish_no_clobber( + manifest_temporary, + manifest_path, + manifest_temporary_identity, + manifest_descriptor, + ) + if os.name == "nt": + os.close(manifest_descriptor) + manifest_descriptor = None + _unlink_if_identity(manifest_temporary, manifest_temporary_identity) + _fsync_directory(manifest_path.parent) + manifest_temporary = None + if ( + stable_file_sha256(archive_path, label="published source archive") + != sha256(archive).hexdigest() + or stable_file_sha256(manifest_path, label="published source manifest") + != sha256(manifest).hexdigest() + ): + raise SourceManifestError("published source artifacts failed stable rehash") + except BaseException as error: + if os.name == "nt" and manifest_descriptor is not None: + os.close(manifest_descriptor) + manifest_descriptor = None + if os.name == "nt" and archive_descriptor is not None: + os.close(archive_descriptor) + archive_descriptor = None + if manifest_identity is not None: + _unlink_if_identity(manifest_path, manifest_identity) + if archive_identity is not None: + _unlink_if_identity(archive_path, archive_identity) + if manifest_temporary is not None and manifest_temporary_identity is not None: + _unlink_if_identity(manifest_temporary, manifest_temporary_identity) + if archive_temporary is not None: + _unlink_if_identity(archive_temporary, archive_temporary_identity) + if isinstance(error, FileExistsError): + raise SourceManifestError( + "source artifact destination appeared during publication" + ) from error + if isinstance(error, CheckpointManifestError): + raise SourceManifestError("could not rehash published source artifacts") from error + raise + finally: + if manifest_descriptor is not None: + os.close(manifest_descriptor) + if archive_descriptor is not None: + os.close(archive_descriptor) + + +def _manifest_from_git_archive( + archive: bytes, *, source_kind: str, git_commit: str, git_tree: str +) -> bytes: + directories: set[str] = set() + files: dict[str, dict[str, object]] = {} + try: + with tarfile.open(fileobj=BytesIO(archive), mode="r:") as handle: + for member in handle.getmembers(): + relative = _relative_path(member.name.rstrip("/"), "Git archive path") + if member.isdir(): + directories.add(relative) + elif member.isfile(): + extracted = handle.extractfile(member) + if extracted is None: + raise SourceManifestError(f"could not read Git archive file {relative!r}") + payload = extracted.read() + if len(payload) != member.size: + raise SourceManifestError(f"Git archive file {relative!r} is truncated") + files[relative] = { + "path": relative, + "mode": "100755" if member.mode & 0o111 else "100644", + "size_bytes": len(payload), + "sha256": sha256(payload).hexdigest(), + } + else: + raise SourceManifestError( + f"Git archive path {relative!r} is not a regular file or directory" + ) + except (tarfile.TarError, OSError) as error: + raise SourceManifestError("could not parse exact Git archive") from error + if not files: + raise SourceManifestError("Git archive contains no source files") + return _canonical_json_bytes( + { + "source_manifest_schema_version": 1, + "source_kind": _canonical_text(source_kind, "source kind"), + "git_commit": _hex_digest(git_commit, 40, "Git commit"), + "git_tree": _hex_digest(git_tree, 40, "Git tree"), + "git_archive_sha256": sha256(archive).hexdigest(), + "archive_scope": _ARCHIVE_SCOPE, + "directories": sorted(directories), + "files": [files[path] for path in sorted(files)], + } + ) + + +def create_source_manifest_from_git_archive( + checkout: str | Path, + *, + expected_commit: str, + source_kind: str, + archive_output: str | Path, + manifest_output: str | Path, +) -> GeneratedSourceManifest: + """Create a sealed runtime-source witness and its exact scoped Git archive.""" + root = Path(checkout).expanduser().resolve() + if not root.is_dir(): + raise SourceManifestError("Git checkout must be a local directory") + commit = _hex_digest(expected_commit, 40, "expected Git commit") + top_level = _run_git(root, ["rev-parse", "--show-toplevel"]).decode().strip() + if Path(top_level).resolve() != root: + raise SourceManifestError("Git checkout must be the repository root") + if _run_git(root, ["rev-parse", "HEAD"]).decode().strip() != commit: + raise SourceManifestError("Git checkout HEAD does not match the expected commit") + status_arguments = ["status", "--porcelain=v1", "--untracked-files=all"] + if _run_git(root, status_arguments).strip(): + raise SourceManifestError("Git checkout has tracked or untracked modifications") + tree = _hex_digest( + _run_git(root, ["rev-parse", "HEAD^{tree}"]).decode().strip(), 40, "Git tree" + ) + archive = _run_git(root, ["archive", "--format=tar", commit, "--", _ARCHIVE_SCOPE]) + manifest = _manifest_from_git_archive( + archive, source_kind=source_kind, git_commit=commit, git_tree=tree + ) + if ( + _run_git(root, ["rev-parse", "HEAD"]).decode().strip() != commit + or _run_git(root, ["rev-parse", "HEAD^{tree}"]).decode().strip() != tree + or _run_git(root, status_arguments).strip() + ): + raise SourceManifestError("Git checkout changed while its archive was generated") + + archive_path = Path(archive_output).expanduser().resolve() + manifest_path = Path(manifest_output).expanduser().resolve() + if archive_path == manifest_path: + raise SourceManifestError("archive and source-manifest outputs must differ") + _publish_source_artifacts(archive_path, archive, manifest_path, manifest) + return GeneratedSourceManifest( + git_commit=commit, + git_tree=tree, + git_archive_sha256=sha256(archive).hexdigest(), + manifest_sha256=sha256(manifest).hexdigest(), + ) 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/mask_reuse_capture.py b/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py new file mode 100644 index 00000000000..9cee1adff50 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py @@ -0,0 +1,951 @@ +# 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. + +"""Host-side contract for ModelOpt-controlled mask-reuse capture. + +The vLLM backend measures rank-local sufficient statistics. This module owns +the trusted inputs, validates every echoed invocation, merges tensor-parallel +consumer-head shards, and emits the normalized rows consumed by +``calibrate_mask_reuse_policy``. It never fabricates missing GPU statistics. +""" + +from __future__ import annotations + +import json +import math +import struct +import unicodedata +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import cast + +from modelopt.torch.sparsity.attention_sparsity.calibration.mask_reuse import ( + canonical_prefill_threshold_scale_factor, +) + +__all__ = [ + "CAPTURE_SCHEMA_VERSION", + "MAX_QUERY_CHUNK_TOKENS", + "CaptureContractError", + "MergedCapture", + "PromptSpec", + "build_capture_invocation", + "canonical_json_sha256", + "load_prompt_specs", + "load_vanilla_prefill_fit", + "merge_rank_captures", + "parse_prompt_specs_jsonl", + "parse_vanilla_prefill_fit", + "source_capture_sha256", + "validate_begin_acks", + "validate_capture_statuses", +] + +CAPTURE_SCHEMA_VERSION = 2 +MAX_QUERY_CHUNK_TOKENS = 8192 +_QUERY_START_ALIGNMENT = 128 +_SPLITS = frozenset({"calibration", "heldout"}) +_PARTITIONS = frozenset({"development", "outer_test"}) +_PROMPT_FIELDS = frozenset( + { + "split", + "partition", + "inner_fold", + "prompt_id", + "source", + "source_group_sha256", + "prompt", + "min_kv_tokens", + "max_kv_tokens", + } +) +_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"}) +_STATUS_FIELDS = frozenset({"capture_schema_version", "available", "rank", "world_size", "reason"}) +_ACK_FIELDS = frozenset( + { + "capture_schema_version", + "armed", + "rank", + "world_size", + "invocation_sha256", + } +) +_RANK_CAPTURE_FIELDS = frozenset( + { + "capture_schema_version", + "rank", + "world_size", + "invocation", + "invocation_sha256", + "geometry", + "global_num_heads", + "eligible_tiles", + "anchor_stats_by_layer", + "consumer_layers", + "attention_call_counts", + "tp_head_order_evidence", + "dense_shadow_evidence", + } +) +_ANCHOR_STATS_FIELDS = frozenset({"retained_tiles", "dropped_mass"}) +_CONSUMER_STATS_FIELDS = frozenset({"anchor_layer", "consumer_head_start", "dropped_mass"}) +_ATTENTION_CALL_COUNT_FIELDS = frozenset({"prefill", "decode"}) +_TP_HEAD_ORDER_FIELDS = frozenset( + { + "sentinel_device_type", + "gather_dim", + "local_rank", + "local_num_heads", + "gathered_rank_local_head", + } +) +_DENSE_SHADOW_FIELDS = frozenset({"enabled", "atol_hex", "rtol_hex", "validated_layer_indices"}) + + +class CaptureContractError(ValueError): + """Raised when capture inputs or backend evidence violate the contract.""" + + +def _exact_fields(value: Mapping[str, object], expected: frozenset[str], label: str) -> None: + missing = expected - value.keys() + extra = value.keys() - expected + if missing or extra: + raise CaptureContractError( + 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 CaptureContractError(f"{label} must be an integer >= {minimum}") + return value + + +def _finite_number( + value: object, + label: str, + *, + minimum: float | None = None, + maximum: float | None = None, +) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise CaptureContractError(f"{label} must be a finite number") + result = float(value) + if not math.isfinite(result): + raise CaptureContractError(f"{label} must be a finite number") + if minimum is not None and result < minimum: + raise CaptureContractError(f"{label} must be >= {minimum}") + if maximum is not None and result > maximum: + raise CaptureContractError(f"{label} must be <= {maximum}") + return result + + +def _text(value: object, label: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise CaptureContractError(f"{label} must be a non-empty string") + result = value.strip() + if unicodedata.normalize("NFC", result) != result or any( + ord(character) < 32 for character in result + ): + raise CaptureContractError(f"{label} must be NFC text without control characters") + return result + + +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 CaptureContractError(f"{label} must be a lowercase SHA256") + return value + + +def _canonical_float_hex(value: object, label: str) -> float: + if not isinstance(value, str): + raise CaptureContractError(f"{label} must be a canonical float.hex string") + try: + parsed = float.fromhex(value) + except ValueError as error: + raise CaptureContractError(f"{label} must be a canonical float.hex string") from error + if not math.isfinite(parsed) or parsed.hex() != value: + raise CaptureContractError(f"{label} must be a canonical finite float.hex string") + return parsed + + +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 CaptureContractError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + +def canonical_json_sha256(value: object) -> str: + """Hash the canonical JSON encoding used by capture RPC acknowledgements.""" + encoded = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True).encode() + return sha256(encoded).hexdigest() + + +@dataclass(frozen=True, slots=True) +class PromptSpec: + """One preregistered prompt/context capture unit.""" + + split: str + partition: str + inner_fold: int | None + prompt_id: str + source: str + source_group_sha256: str + prompt: str + min_kv_tokens: int + max_kv_tokens: int | None + + @classmethod + def from_mapping(cls, raw: Mapping[str, object]) -> PromptSpec: + """Parse one strict prompt-plan object.""" + _exact_fields(raw, _PROMPT_FIELDS, "prompt") + split = raw["split"] + if split not in _SPLITS: + raise CaptureContractError(f"prompt.split must be one of {sorted(_SPLITS)}") + partition = raw["partition"] + if partition not in _PARTITIONS: + raise CaptureContractError(f"prompt.partition must be one of {sorted(_PARTITIONS)}") + expected_split = "calibration" if partition == "development" else "heldout" + if split != expected_split: + raise CaptureContractError("prompt.split and prompt.partition disagree") + raw_fold = raw["inner_fold"] + if partition == "development": + inner_fold = _integer(raw_fold, "prompt.inner_fold") + elif raw_fold is not None: + raise CaptureContractError("outer_test prompts must have null inner_fold") + else: + inner_fold = None + prompt = raw["prompt"] + if not isinstance(prompt, str) or not prompt: + raise CaptureContractError("prompt.prompt must be a non-empty string") + minimum = _integer(raw["min_kv_tokens"], "prompt.min_kv_tokens", minimum=1) + maximum = raw["max_kv_tokens"] + if maximum is not None: + maximum = _integer(maximum, "prompt.max_kv_tokens", minimum=minimum) + return cls( + split=split, + partition=partition, + inner_fold=inner_fold, + prompt_id=_text(raw["prompt_id"], "prompt.prompt_id"), + source=_text(raw["source"], "prompt.source"), + source_group_sha256=_sha256(raw["source_group_sha256"], "prompt.source_group_sha256"), + prompt=prompt, + min_kv_tokens=minimum, + max_kv_tokens=maximum, + ) + + @property + def bucket(self) -> tuple[int, int | None]: + """Return the calibrated context-bucket bounds.""" + return self.min_kv_tokens, self.max_kv_tokens + + +def parse_prompt_specs_jsonl(payload: bytes) -> list[PromptSpec]: + """Parse exact strict-JSONL prompt bytes and reject split leakage.""" + if not isinstance(payload, bytes): + raise CaptureContractError("prompt file payload must be bytes") + try: + lines = payload.decode("utf-8").splitlines() + except UnicodeDecodeError as error: + raise CaptureContractError("prompt file must be UTF-8") from error + prompts: list[PromptSpec] = [] + 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 CaptureContractError(f"line {line_number}: invalid JSON: {error.msg}") from error + except CaptureContractError as error: + raise CaptureContractError(f"line {line_number}: {error}") from error + if not isinstance(raw, dict): + raise CaptureContractError(f"line {line_number}: prompt must be an object") + try: + prompts.append(PromptSpec.from_mapping(raw)) + except CaptureContractError as error: + raise CaptureContractError(f"line {line_number}: {error}") from error + if not prompts: + raise CaptureContractError("prompt file contains no prompts") + + by_split = {split: [prompt for prompt in prompts if prompt.split == split] for split in _SPLITS} + if any(not values for values in by_split.values()): + raise CaptureContractError("prompt file requires calibration and heldout splits") + calibration_ids = {prompt.prompt_id for prompt in by_split["calibration"]} + heldout_ids = {prompt.prompt_id for prompt in by_split["heldout"]} + for split, values in by_split.items(): + identifiers = [prompt.prompt_id for prompt in values] + if len(identifiers) != len(set(identifiers)): + raise CaptureContractError(f"prompt IDs must be unique within the {split} split") + if calibration_ids & heldout_ids: + raise CaptureContractError("prompt IDs overlap calibration and heldout splits") + group_assignments: dict[str, tuple[str, int | None]] = {} + for prompt in prompts: + assignment = (prompt.partition, prompt.inner_fold) + previous = group_assignments.setdefault(prompt.source_group_sha256, assignment) + if previous != assignment: + raise CaptureContractError( + "one source group is assigned to multiple partitions or inner folds" + ) + split_buckets = { + split: {prompt.bucket for prompt in values} for split, values in by_split.items() + } + if split_buckets["calibration"] != split_buckets["heldout"]: + raise CaptureContractError("heldout context buckets must match calibration buckets") + keys = [(prompt.split, prompt.prompt_id, prompt.bucket) for prompt in prompts] + if len(keys) != len(set(keys)): + raise CaptureContractError("prompt file repeats a split/prompt/context capture") + return sorted( + prompts, + key=lambda prompt: ( + prompt.min_kv_tokens, + math.inf if prompt.max_kv_tokens is None else prompt.max_kv_tokens, + prompt.split, + prompt.prompt_id, + ), + ) + + +def load_prompt_specs(path: str | Path) -> list[PromptSpec]: + """Load strict JSONL prompt specifications from a path.""" + return parse_prompt_specs_jsonl(Path(path).read_bytes()) + + +def _parse_strict_json_object(payload: bytes, label: str) -> Mapping[str, object]: + try: + raw = json.loads(payload, object_pairs_hook=_reject_duplicate_json_keys) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + detail = error.msg if isinstance(error, json.JSONDecodeError) else str(error) + raise CaptureContractError(f"{label} is invalid JSON: {detail}") from error + if not isinstance(raw, dict): + raise CaptureContractError(f"{label} must be a JSON object") + return raw + + +def parse_vanilla_prefill_fit(payload: bytes) -> dict[str, object]: + """Parse and canonicalize exact vanilla-calibration JSON bytes.""" + return canonical_prefill_threshold_scale_factor( + _parse_strict_json_object(payload, "vanilla calibration") + ) + + +def load_vanilla_prefill_fit(path: str | Path) -> dict[str, object]: + """Load and canonicalize an existing ModelOpt vanilla skip-softmax fit.""" + return parse_vanilla_prefill_fit(Path(path).read_bytes()) + + +def source_capture_sha256(prompt_token_ids: Sequence[int]) -> str: + """Fingerprint exact token IDs independently of prompt labels and splits.""" + if not prompt_token_ids: + raise CaptureContractError("prompt token IDs must not be empty") + digest = sha256(b"modelopt-mask-reuse-token-ids-v1\0") + digest.update(struct.pack("= 2**64: + raise CaptureContractError(f"prompt_token_ids[{index}] exceeds uint64") + digest.update(struct.pack(" dict[str, int]: + q_start = ((sample_length - 1) // MAX_QUERY_CHUNK_TOKENS) * MAX_QUERY_CHUNK_TOKENS + q_tokens = sample_length - q_start + if q_tokens <= _QUERY_START_ALIGNMENT: + raise CaptureContractError( + "the deployed q-stage-2 contract requires a final prefill chunk of at least " + "129 tokens; adjust the prompt length" + ) + return { + "q_tokens": q_tokens, + "kv_tokens": sample_length, + "q_start_tokens": q_start, + } + + +def build_capture_invocation( + *, + model: str, + checkpoint_manifest_sha256: str, + prompt: PromptSpec, + prompt_token_ids: Sequence[int], + target_sparsity: float, + threshold_scale_factor: Mapping[str, object], +) -> dict[str, object]: + """Build the exact invocation the backend must echo before evidence is trusted.""" + model = _text(model, "model") + checkpoint_identity = _sha256(checkpoint_manifest_sha256, "checkpoint_manifest_sha256") + sample_length = len(prompt_token_ids) + if sample_length < prompt.min_kv_tokens or ( + prompt.max_kv_tokens is not None and sample_length > prompt.max_kv_tokens + ): + raise CaptureContractError( + f"prompt {prompt.prompt_id!r} token length {sample_length} lies outside " + f"bucket [{prompt.min_kv_tokens}, {prompt.max_kv_tokens}]" + ) + target = _finite_number(target_sparsity, "target_sparsity", minimum=0.0, maximum=1.0) + if not 0.0 < target < 1.0: + raise CaptureContractError("target_sparsity must be in (0, 1)") + + if set(threshold_scale_factor) != {"formula", "prefill"}: + raise CaptureContractError("threshold_scale_factor must be canonical ModelOpt metadata") + params = threshold_scale_factor["prefill"] + if not isinstance(params, Mapping): + raise CaptureContractError("threshold_scale_factor.prefill must be an object") + lower = params.get("min_observed_sparsity") + upper = params.get("max_observed_sparsity") + if (lower is not None and target < float(lower)) or ( + upper is not None and target > float(upper) + ): + raise CaptureContractError( + f"target_sparsity={target} is outside the observed vanilla calibration range" + ) + a = float(params["a"]) + b = float(params["b"]) + threshold_log2 = math.log2(a) + b * target * math.log2(math.e) - math.log2(sample_length) + # ``math.exp2`` was added in Python 3.11; ModelOpt still supports 3.10. + threshold_lambda = 2.0**threshold_log2 + if not math.isfinite(threshold_log2) or not 0.0 < threshold_lambda < 1.0: + raise CaptureContractError("vanilla fit derives a threshold outside (0, 1)") + + invocation = { + "capture_schema_version": CAPTURE_SCHEMA_VERSION, + "model": model, + "checkpoint_manifest_sha256": checkpoint_identity, + "split": prompt.split, + "partition": prompt.partition, + "inner_fold": prompt.inner_fold, + "prompt_id": prompt.prompt_id, + "source": prompt.source, + "source_group_sha256": prompt.source_group_sha256, + "source_capture_sha256": source_capture_sha256(prompt_token_ids), + "min_kv_tokens": prompt.min_kv_tokens, + "max_kv_tokens": prompt.max_kv_tokens, + "target_sparsity_hex": target.hex(), + "sample_length": sample_length, + "threshold_log2_hex": threshold_log2.hex(), + "threshold_lambda_hex": threshold_lambda.hex(), + "expected_geometry": _expected_final_geometry(sample_length), + } + _validate_invocation(invocation) + return invocation + + +def _validate_geometry(raw: object, label: str) -> dict[str, int]: + if not isinstance(raw, Mapping): + raise CaptureContractError(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 > MAX_QUERY_CHUNK_TOKENS: + raise CaptureContractError(f"{label}.q_tokens exceeds {MAX_QUERY_CHUNK_TOKENS}") + 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") + if q_start % _QUERY_START_ALIGNMENT or q_start + q_tokens != kv_tokens: + raise CaptureContractError(f"{label} is not a 128-token-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 + + +def _validate_invocation(raw: object) -> dict[str, object]: + if not isinstance(raw, Mapping): + raise CaptureContractError("capture invocation must be an object") + _exact_fields(raw, _INVOCATION_FIELDS, "capture invocation") + if raw["capture_schema_version"] != CAPTURE_SCHEMA_VERSION: + raise CaptureContractError("capture invocation has an unsupported schema version") + _text(raw["model"], "capture invocation.model") + _sha256( + raw["checkpoint_manifest_sha256"], + "capture invocation.checkpoint_manifest_sha256", + ) + if raw["split"] not in _SPLITS: + raise CaptureContractError("capture invocation.split is invalid") + partition = raw["partition"] + if partition not in _PARTITIONS: + raise CaptureContractError("capture invocation.partition is invalid") + expected_split = "calibration" if partition == "development" else "heldout" + if raw["split"] != expected_split: + raise CaptureContractError("capture invocation split and partition disagree") + if partition == "development": + _integer(raw["inner_fold"], "capture invocation.inner_fold") + elif raw["inner_fold"] is not None: + raise CaptureContractError("outer_test capture invocation must have null inner_fold") + _text(raw["prompt_id"], "capture invocation.prompt_id") + _text(raw["source"], "capture invocation.source") + _sha256(raw["source_group_sha256"], "capture invocation.source_group_sha256") + _sha256(raw["source_capture_sha256"], "capture invocation.source_capture_sha256") + minimum = _integer(raw["min_kv_tokens"], "capture invocation.min_kv_tokens", minimum=1) + maximum = raw["max_kv_tokens"] + if maximum is not None: + _integer(maximum, "capture invocation.max_kv_tokens", minimum=minimum) + target = _canonical_float_hex( + raw["target_sparsity_hex"], "capture invocation.target_sparsity_hex" + ) + if not 0.0 < target < 1.0: + raise CaptureContractError("capture invocation target sparsity must be in (0, 1)") + sample_length = _integer(raw["sample_length"], "capture invocation.sample_length", minimum=1) + threshold_log2 = _canonical_float_hex( + raw["threshold_log2_hex"], "capture invocation.threshold_log2_hex" + ) + threshold_lambda = _canonical_float_hex( + raw["threshold_lambda_hex"], "capture invocation.threshold_lambda_hex" + ) + if threshold_log2 >= 0.0 or not 0.0 < threshold_lambda < 1.0: + raise CaptureContractError("capture invocation threshold must be in (0, 1)") + if (2.0**threshold_log2).hex() != threshold_lambda.hex(): + raise CaptureContractError("capture invocation threshold hex fields disagree") + geometry = _validate_geometry(raw["expected_geometry"], "expected_geometry") + if geometry["kv_tokens"] != sample_length: + raise CaptureContractError("expected geometry does not match sample_length") + return dict(raw) + + +def _validate_rank_envelope( + values: Sequence[Mapping[str, object]], expected_fields: frozenset[str], label: str +) -> tuple[int, list[Mapping[str, object]]]: + if not values: + raise CaptureContractError(f"{label} returned no rank payloads") + world_sizes: set[int] = set() + ranks: list[int] = [] + for index, value in enumerate(values): + if not isinstance(value, Mapping): + raise CaptureContractError(f"{label}[{index}] must be an object") + _exact_fields(value, expected_fields, f"{label}[{index}]") + if value["capture_schema_version"] != CAPTURE_SCHEMA_VERSION: + raise CaptureContractError(f"{label}[{index}] has an unsupported schema version") + ranks.append(_integer(value["rank"], f"{label}[{index}].rank")) + world_sizes.add(_integer(value["world_size"], f"{label}[{index}].world_size", minimum=1)) + if len(world_sizes) != 1: + raise CaptureContractError(f"{label} disagrees on world_size") + world_size = next(iter(world_sizes)) + if len(values) != world_size or sorted(ranks) != list(range(world_size)): + raise CaptureContractError(f"{label} does not exactly cover ranks [0, {world_size})") + return world_size, sorted(values, key=lambda value: cast("int", value["rank"])) + + +def validate_capture_statuses(values: Sequence[Mapping[str, object]]) -> int: + """Require the env-gated backend capture sink on every worker rank.""" + world_size, ordered = _validate_rank_envelope(values, _STATUS_FIELDS, "capture status") + failures = [] + for value in ordered: + if not isinstance(value["available"], bool): + raise CaptureContractError("capture status.available must be boolean") + reason = value["reason"] + if reason is not None and not isinstance(reason, str): + raise CaptureContractError("capture status.reason must be a string or null") + if not value["available"]: + failures.append(f"rank {value['rank']}: {reason or 'unavailable'}") + if failures: + raise CaptureContractError("mask-reuse capture is unavailable: " + "; ".join(failures)) + return world_size + + +def validate_begin_acks( + values: Sequence[Mapping[str, object]], invocation: Mapping[str, object] +) -> int: + """Validate that every rank armed the exact same invocation before generation.""" + _validate_invocation(invocation) + expected_digest = canonical_json_sha256(invocation) + world_size, ordered = _validate_rank_envelope(values, _ACK_FIELDS, "capture begin") + for value in ordered: + if value["armed"] is not True: + raise CaptureContractError(f"capture begin rank {value['rank']} did not arm") + if value["invocation_sha256"] != expected_digest: + raise CaptureContractError( + f"capture begin rank {value['rank']} acknowledged the wrong invocation" + ) + return world_size + + +def _parse_anchor_stats( + raw: object, *, global_num_heads: int, eligible_tiles: int +) -> dict[str, dict[str, list[int] | list[float]]]: + if not isinstance(raw, Mapping) or not raw: + raise CaptureContractError("anchor_stats_by_layer must be a non-empty object") + result: dict[str, dict[str, list[int] | list[float]]] = {} + for raw_layer, value in raw.items(): + if not isinstance(raw_layer, str): + raise CaptureContractError("anchor layer keys must be canonical integer strings") + layer = _integer(int(raw_layer), f"anchor layer {raw_layer}") if raw_layer.isdigit() else -1 + if layer < 0 or raw_layer != str(layer): + raise CaptureContractError("anchor layer keys must be canonical integer strings") + if not isinstance(value, Mapping): + raise CaptureContractError(f"anchor_stats_by_layer[{layer}] must be an object") + _exact_fields(value, _ANCHOR_STATS_FIELDS, f"anchor_stats_by_layer[{layer}]") + retained_raw = value["retained_tiles"] + dropped_raw = value["dropped_mass"] + if not isinstance(retained_raw, list) or not isinstance(dropped_raw, list): + raise CaptureContractError(f"anchor_stats_by_layer[{layer}] arrays must be lists") + if len(retained_raw) != global_num_heads or len(dropped_raw) != global_num_heads: + raise CaptureContractError( + f"anchor_stats_by_layer[{layer}] must cover {global_num_heads} heads" + ) + retained = [ + _integer(value, f"anchor_stats_by_layer[{layer}].retained_tiles[{head}]") + for head, value in enumerate(retained_raw) + ] + if any(value > eligible_tiles for value in retained): + raise CaptureContractError( + f"anchor_stats_by_layer[{layer}] retained tiles exceed eligible tiles" + ) + dropped = [ + _finite_number( + value, + f"anchor_stats_by_layer[{layer}].dropped_mass[{head}]", + minimum=0.0, + maximum=1.0, + ) + for head, value in enumerate(dropped_raw) + ] + result[raw_layer] = {"retained_tiles": retained, "dropped_mass": dropped} + return dict(sorted(result.items(), key=lambda item: int(item[0]))) + + +def _parse_consumer_layers( + raw: object, *, rank: int, global_num_heads: int +) -> dict[int, tuple[int, int, list[list[float]]]]: + if not isinstance(raw, Mapping) or not raw: + raise CaptureContractError(f"rank {rank} consumer_layers must be a non-empty object") + result: dict[int, tuple[int, int, list[list[float]]]] = {} + for raw_layer, value in raw.items(): + if ( + not isinstance(raw_layer, str) + or not raw_layer.isdigit() + or raw_layer != str(int(raw_layer)) + ): + raise CaptureContractError("consumer layer keys must be canonical integer strings") + layer = int(raw_layer) + if not isinstance(value, Mapping): + raise CaptureContractError(f"consumer_layers[{layer}] must be an object") + _exact_fields(value, _CONSUMER_STATS_FIELDS, f"consumer_layers[{layer}]") + anchor = _integer(value["anchor_layer"], f"consumer_layers[{layer}].anchor_layer") + if anchor >= layer: + raise CaptureContractError(f"consumer layer {layer} must follow anchor {anchor}") + start = _integer( + value["consumer_head_start"], f"consumer_layers[{layer}].consumer_head_start" + ) + matrix = value["dropped_mass"] + if not isinstance(matrix, list) or not matrix: + raise CaptureContractError(f"consumer_layers[{layer}].dropped_mass must be non-empty") + rows: list[list[float]] = [] + for local_head, raw_row in enumerate(matrix): + if not isinstance(raw_row, list) or len(raw_row) != global_num_heads: + raise CaptureContractError( + f"consumer_layers[{layer}].dropped_mass[{local_head}] must cover " + f"{global_num_heads} donor heads" + ) + rows.append( + [ + _finite_number( + value, + f"consumer_layers[{layer}].dropped_mass[{local_head}][{donor}]", + minimum=0.0, + maximum=1.0, + ) + for donor, value in enumerate(raw_row) + ] + ) + if start + len(rows) > global_num_heads: + raise CaptureContractError( + f"rank {rank} consumer layer {layer} shard exceeds head count" + ) + result[layer] = (anchor, start, rows) + return result + + +def _parse_attention_call_counts(raw: object, *, rank: int) -> dict[str, int]: + if not isinstance(raw, Mapping): + raise CaptureContractError(f"rank {rank} attention_call_counts must be an object") + _exact_fields(raw, _ATTENTION_CALL_COUNT_FIELDS, f"rank {rank} attention_call_counts") + return { + name: _integer(raw[name], f"rank {rank} attention_call_counts.{name}") + for name in ("prefill", "decode") + } + + +def _parse_tp_head_order_evidence( + raw: object, + *, + rank: int, + world_size: int, + global_num_heads: int, +) -> dict[str, object]: + if not isinstance(raw, Mapping): + raise CaptureContractError(f"rank {rank} tp_head_order_evidence must be an object") + _exact_fields(raw, _TP_HEAD_ORDER_FIELDS, f"rank {rank} tp_head_order_evidence") + if raw["sentinel_device_type"] != "cuda": + raise CaptureContractError(f"rank {rank} TP sentinel was not gathered on CUDA") + if raw["gather_dim"] != 0: + raise CaptureContractError(f"rank {rank} TP sentinel must be gathered along dim 0") + if _integer(raw["local_rank"], f"rank {rank} TP local_rank") != rank: + raise CaptureContractError(f"rank {rank} TP sentinel reports a different local rank") + local_num_heads = _integer(raw["local_num_heads"], f"rank {rank} TP local_num_heads", minimum=1) + if local_num_heads * world_size != global_num_heads: + raise CaptureContractError( + f"rank {rank} local head count does not evenly cover global heads" + ) + gathered = raw["gathered_rank_local_head"] + expected = [ + [global_head // local_num_heads, global_head % local_num_heads] + for global_head in range(global_num_heads) + ] + if gathered != expected: + raise CaptureContractError( + f"rank {rank} TP all-gather is not rank-major in global-head order" + ) + return { + "rank": rank, + "global_head_start": rank * local_num_heads, + "local_num_heads": local_num_heads, + "sentinel_device_type": "cuda", + "gather_dim": 0, + "gathered_rank_local_head": expected, + } + + +def _parse_dense_shadow_evidence(raw: object, *, rank: int) -> dict[str, object]: + if not isinstance(raw, Mapping): + raise CaptureContractError(f"rank {rank} dense_shadow_evidence must be an object") + _exact_fields(raw, _DENSE_SHADOW_FIELDS, f"rank {rank} dense_shadow_evidence") + enabled = raw["enabled"] + if not isinstance(enabled, bool): + raise CaptureContractError(f"rank {rank} dense_shadow_evidence.enabled must be boolean") + atol = _canonical_float_hex(raw["atol_hex"], f"rank {rank} dense shadow atol") + rtol = _canonical_float_hex(raw["rtol_hex"], f"rank {rank} dense shadow rtol") + if atol != 0.0 or rtol != 0.0: + raise CaptureContractError("dense shadow evidence must use bitwise zero tolerances") + raw_layers = raw["validated_layer_indices"] + if not isinstance(raw_layers, list): + raise CaptureContractError( + f"rank {rank} dense_shadow_evidence.validated_layer_indices must be a list" + ) + layers = [_integer(layer, f"rank {rank} dense shadow layer") for layer in raw_layers] + if layers != sorted(set(layers)): + raise CaptureContractError(f"rank {rank} dense shadow layers must be sorted and unique") + if not enabled and layers: + raise CaptureContractError( + f"rank {rank} disabled dense shadow cannot report validated layers" + ) + return { + "enabled": enabled, + "atol_hex": atol.hex(), + "rtol_hex": rtol.hex(), + "validated_layer_indices": layers, + } + + +@dataclass(frozen=True, slots=True) +class MergedCapture: + """One compact normalized capture plus auditable metadata.""" + + capture: dict[str, object] + manifest: dict[str, object] + + +def merge_rank_captures( + values: Sequence[Mapping[str, object]], invocation: Mapping[str, object] +) -> MergedCapture: + """Merge raw per-rank sufficient statistics without inventing absent rows.""" + trusted_invocation = _validate_invocation(invocation) + expected_digest = canonical_json_sha256(trusted_invocation) + world_size, ordered = _validate_rank_envelope(values, _RANK_CAPTURE_FIELDS, "capture drain") + global_heads: set[int] = set() + eligible_values: set[int] = set() + geometry_values: list[dict[str, int]] = [] + anchor_payloads: list[dict[str, dict[str, list[int] | list[float]]]] = [] + consumer_payloads: list[dict[int, tuple[int, int, list[list[float]]]]] = [] + attention_call_counts: list[dict[str, int]] = [] + tp_head_order_evidence: list[dict[str, object]] = [] + dense_shadow_evidence: list[dict[str, object]] = [] + for value in ordered: + rank = cast("int", value["rank"]) + if value["invocation"] != trusted_invocation: + raise CaptureContractError(f"capture drain rank {rank} echoed the wrong invocation") + if value["invocation_sha256"] != expected_digest: + raise CaptureContractError(f"capture drain rank {rank} has the wrong invocation digest") + geometry = _validate_geometry(value["geometry"], f"capture drain rank {rank} geometry") + if geometry != trusted_invocation["expected_geometry"]: + raise CaptureContractError(f"capture drain rank {rank} measured the wrong final chunk") + geometry_values.append(geometry) + global_num_heads = _integer( + value["global_num_heads"], f"capture drain rank {rank} global_num_heads", minimum=1 + ) + eligible_tiles = _integer( + value["eligible_tiles"], f"capture drain rank {rank} eligible_tiles", minimum=1 + ) + global_heads.add(global_num_heads) + eligible_values.add(eligible_tiles) + anchor_payloads.append( + _parse_anchor_stats( + value["anchor_stats_by_layer"], + global_num_heads=global_num_heads, + eligible_tiles=eligible_tiles, + ) + ) + consumers = _parse_consumer_layers( + value["consumer_layers"], rank=rank, global_num_heads=global_num_heads + ) + consumer_payloads.append(consumers) + counts = _parse_attention_call_counts(value["attention_call_counts"], rank=rank) + attention_call_counts.append(counts) + tp_evidence = _parse_tp_head_order_evidence( + value["tp_head_order_evidence"], + rank=rank, + world_size=world_size, + global_num_heads=global_num_heads, + ) + tp_head_order_evidence.append(tp_evidence) + dense_shadow_evidence.append( + _parse_dense_shadow_evidence(value["dense_shadow_evidence"], rank=rank) + ) + for layer, (_, start, rows) in consumers.items(): + if ( + start != tp_evidence["global_head_start"] + or len(rows) != tp_evidence["local_num_heads"] + ): + raise CaptureContractError( + f"rank {rank} consumer layer {layer} does not match its TP head shard" + ) + if len(global_heads) != 1 or len(eligible_values) != 1: + raise CaptureContractError("capture ranks disagree on head count or eligible tiles") + if any(value != geometry_values[0] for value in geometry_values[1:]): + raise CaptureContractError("capture ranks disagree on final-chunk geometry") + if any(value != anchor_payloads[0] for value in anchor_payloads[1:]): + raise CaptureContractError("capture ranks disagree on global anchor statistics") + if any(value != attention_call_counts[0] for value in attention_call_counts[1:]): + raise CaptureContractError("capture ranks disagree on attention call counts") + if any(value != dense_shadow_evidence[0] for value in dense_shadow_evidence[1:]): + raise CaptureContractError("capture ranks disagree on dense shadow evidence") + layer_sets = [set(value) for value in consumer_payloads] + if any(value != layer_sets[0] for value in layer_sets[1:]): + raise CaptureContractError("capture ranks disagree on consumer layer coverage") + + global_num_heads = next(iter(global_heads)) + eligible_tiles = next(iter(eligible_values)) + if eligible_tiles != _eligible_tiles(geometry_values[0]): + raise CaptureContractError( + "capture eligible_tiles does not match 128x128 bottom-right causal geometry" + ) + anchors = anchor_payloads[0] + attention_layers = sorted({int(layer) for layer in anchors} | layer_sets[0]) + expected_prefill_calls = ( + (cast("int", trusted_invocation["sample_length"]) + MAX_QUERY_CHUNK_TOKENS - 1) + // MAX_QUERY_CHUNK_TOKENS + ) * len(attention_layers) + common_counts = attention_call_counts[0] + if common_counts["prefill"] != expected_prefill_calls: + raise CaptureContractError( + "capture prefill attention call count does not match chunks times layers" + ) + if common_counts["decode"] != 0: + raise CaptureContractError("capture observed a decode attention call at max_tokens=1") + common_shadow = dense_shadow_evidence[0] + if common_shadow["enabled"] and common_shadow["validated_layer_indices"] != attention_layers: + raise CaptureContractError( + "dense shadow evidence does not cover every captured attention layer" + ) + merged_consumers: dict[str, dict[str, object]] = {} + for layer in sorted(layer_sets[0]): + shards = sorted( + (payload[layer] for payload in consumer_payloads), key=lambda value: value[1] + ) + anchor_layers = {shard[0] for shard in shards} + if len(anchor_layers) != 1: + raise CaptureContractError(f"capture ranks disagree on consumer layer {layer}'s anchor") + anchor_layer = next(iter(anchor_layers)) + anchor_stats = anchors.get(str(anchor_layer)) + if anchor_stats is None: + raise CaptureContractError( + f"consumer layer {layer} references missing anchor layer {anchor_layer}" + ) + cursor = 0 + global_rows: list[list[float]] = [] + for _, start, local_rows in shards: + if start != cursor: + raise CaptureContractError( + f"consumer layer {layer} head shards are overlapping or incomplete at {cursor}" + ) + global_rows.extend(local_rows) + cursor += len(local_rows) + if cursor != global_num_heads: + raise CaptureContractError( + f"consumer layer {layer} head shards cover {cursor}, expected {global_num_heads}" + ) + merged_consumers[str(layer)] = { + "anchor_layer": anchor_layer, + "dropped_mass": global_rows, + } + if not merged_consumers: + raise CaptureContractError("capture drain contained no consumer-head statistics") + compact_capture = { + "compact_capture_schema_version": 1, + "invocation": trusted_invocation, + "geometry": geometry_values[0], + "global_num_heads": global_num_heads, + "eligible_tiles": eligible_tiles, + "anchor_stats_by_layer": anchors, + "consumer_layers": merged_consumers, + } + manifest = { + "capture_schema_version": CAPTURE_SCHEMA_VERSION, + "invocation": trusted_invocation, + "invocation_sha256": expected_digest, + "world_size": world_size, + "global_num_heads": global_num_heads, + "eligible_tiles": eligible_tiles, + "candidate_cell_count": sum( + len(cast("Sequence[object]", value["dropped_mass"])) * global_num_heads + for value in merged_consumers.values() + ), + "attention_call_counts": common_counts, + "tp_head_order_evidence": tp_head_order_evidence, + "dense_shadow_evidence": common_shadow, + "compact_capture_sha256": canonical_json_sha256(compact_capture), + } + return MergedCapture(compact_capture, manifest) 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..2d23089ecdd --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -0,0 +1,273 @@ +# 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 + +import modelopt + +# 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 + +__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]]]) -> 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. + """ + 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)} rank(s) " + f"recorded no {phase!r} samples while others did" + ) + merged[phase] = merge_count_records([] if empty else 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. + """ + from ..calibration.calibrator import DynamicThresholdCalibrator + + 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. + """ + threshold_scale_factor: dict[str, Any] = {"formula": "a * exp(b * target_sparsity)"} + for phase in _PHASES: + if phase in calibration_params: + phase_params = { + "a": float(calibration_params[phase]["a"]), + "b": float(calibration_params[phase]["b"]), + } + for key in ("min_observed_sparsity", "max_observed_sparsity"): + if key in calibration_params[phase]: + phase_params[key] = float(calibration_params[phase][key]) + threshold_scale_factor[phase] = phase_params + + skip_group: dict[str, Any] = { + "algorithm": "skip_softmax", + "targets": ["Attention"], + "threshold_scale_factor": threshold_scale_factor, + "target_sparsity": _normalize_target_sparsity(target_sparsity), + } + + config_groups: dict[str, Any] = {"group_0": skip_group} + existing_groups = (existing_config or {}).get("config_groups") + if isinstance(existing_groups, dict): + preserved = [ + group + for group in existing_groups.values() + if isinstance(group, dict) and group.get("algorithm") != "skip_softmax" + ] + for idx, group in enumerate(preserved, start=1): + config_groups[f"group_{idx}"] = group + + result: dict[str, Any] = { + "config_groups": config_groups, + "producer": {"name": "modelopt", "version": modelopt.__version__}, + } + # 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 00411eaa4ee..1d4cfd0f4c6 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,7 @@ ) 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 def _target_sparse_ratio_for_phase(target_sparse_ratio, phase: str) -> float: @@ -239,6 +241,134 @@ 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 a decode + step (full-cache, non-causal); ``q_len > 1`` is (chunked) prefill (causal + — the kernel offsets the query into the KV span). 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] + b_start_loc = cu_seqlens_q[:batch] + b_seq_len = cu_seqlens_q[1 : batch + 1] - cu_seqlens_q[:batch] + + 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_len = int(b_seq_len[i].item()) + if q_len <= 0: + continue + q_start = int(b_start_loc[i].item()) + seq_k = int(seq_lens[i].item()) + phase = "decode" if q_len <= 1 else "prefill" + + oi, counters = attention_calibrate( + q[q_start : q_start + q_len], + k_dummy, + k_dummy, + b_start_loc=torch.zeros(1, device=q.device, dtype=torch.int32), + b_seq_len=b_seq_len[i : i + 1].to(torch.int32), + max_input_len=q_len, + is_causal=q_len > 1, + softmax_scale=impl.scale, + b_seq_len_k=seq_lens[i : i + 1].to(torch.int32), + 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 + + impl._calib_records.append( + { + "phase": phase, + "sample_length": seq_k, + "total_tiles": counters[:, 0].tolist(), + "skipped_tiles": counters[:, 1].tolist(), + } + ) + + return output + + # Resolution guards raw configured transforms; dispatch rechecks effective # sparse work after calibration and decode-only pruning. def _forward_modelopt( @@ -504,6 +634,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, @@ -686,6 +836,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, @@ -814,3 +1002,70 @@ 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. + """ + from .sparse_attn_calibration import merge_count_records, split_records_by_phase + + splits = [ + split_records_by_phase(getattr(impl, "_calib_records", [])) + for impl in iter_sparse_impls(model) + ] + phases = {phase for split in splits for phase, records in split.items() if records} + counts: dict[str, list[dict]] = {} + for phase in phases: + sources = [split.get(phase, []) for split in splits] + empty = sum(1 for source in sources if not source) + if empty: + # Every layer sees every launch, so a layer with no records for a + # phase others measured indicates a collection bug. + raise ValueError( + f"Misaligned calibration records: {empty}/{len(sources)} attention " + f"layer(s) recorded no {phase!r} samples while others did" + ) + counts[phase] = merge_count_records(sources) + return counts diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_mask_reuse_capture.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_mask_reuse_capture.py new file mode 100644 index 00000000000..0b01c3cb87d --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_mask_reuse_capture.py @@ -0,0 +1,93 @@ +# 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 worker bootstrap and RPCs for policy-free mask-reuse calibration.""" + +from __future__ import annotations + +import importlib +import os +from collections.abc import Mapping +from typing import TYPE_CHECKING + +from vllm.v1.worker.gpu_worker import Worker as BaseWorker + +if TYPE_CHECKING: + from types import ModuleType + +CAPTURE_ENV = "MASK_REUSE_FA4_CALIBRATION_CAPTURE" +PLAN_ENV = "MASK_REUSE_FA4_PLAN" +_REQUIRED_API = ( + "configure_capture_runtime", + "capture_status", + "begin_capture", + "drain_capture", +) + +__all__ = ["MaskReuseCaptureWorker"] + + +def _capture_api() -> ModuleType: + if os.environ.get(CAPTURE_ENV) != "1": + raise RuntimeError(f"{CAPTURE_ENV}=1 is required for mask-reuse calibration capture") + try: + api = importlib.import_module("mask_reuse_vllm.capture") + except ImportError as error: + raise RuntimeError( + "the custom mask-reuse backend does not expose mask_reuse_vllm.capture" + ) from error + missing = [name for name in _REQUIRED_API if not callable(getattr(api, name, None))] + if missing: + raise RuntimeError(f"mask-reuse capture API is incomplete; missing {missing}") + return api + + +def _configure_capture_before_model_load() -> ModuleType: + """Install a planner and capture provider without loading a serving policy.""" + plan_name = os.environ.get(PLAN_ENV) + if not plan_name: + raise RuntimeError(f"{PLAN_ENV} must name the explicit calibration topology preset") + api = _capture_api() + api.configure_capture_runtime(plan_name) + return api + + +class MaskReuseCaptureWorker(BaseWorker): + """Run the custom backend in env-gated, policy-free capture mode.""" + + def load_model(self, *args, **kwargs) -> None: + """Install capture runtime before vLLM constructs attention modules.""" + # The attention implementation resolves process-local runtime state + # while the model is loading. Install the policy-free capture provider + # first; a promoted v3 policy must not be required to collect its own + # calibration evidence. + api = _configure_capture_before_model_load() + super().load_model(*args, **kwargs) + status = api.capture_status() + if not isinstance(status, Mapping) or status.get("available") is not True: + reason = status.get("reason") if isinstance(status, Mapping) else None + raise RuntimeError(f"mask-reuse capture backend is unavailable: {reason or status!r}") + + def mask_reuse_capture_status(self) -> dict[str, object]: + """Return this worker rank's fail-closed capture status.""" + return dict(_capture_api().capture_status()) + + def mask_reuse_capture_begin(self, invocation: dict[str, object]) -> dict[str, object]: + """Arm exactly one prompt/target invocation on this worker rank.""" + return dict(_capture_api().begin_capture(invocation)) + + def mask_reuse_capture_drain(self) -> dict[str, object]: + """Drain one completed rank-local sufficient-stat payload.""" + return dict(_capture_api().drain_capture()) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index b4141740b64..724feb69f90 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", ] @@ -248,6 +249,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 @@ -348,11 +356,29 @@ def _plan_vllm_attention( _require_supported_vllm() errors = _global_errors(model_runner) if quantize else [] - mode = _cudagraph_mode(model_runner) if quantize else None + 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 +391,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 +521,87 @@ 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") + 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..555ef52f483 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,29 +454,29 @@ 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.""" + @pytest.mark.parametrize( + "qdq_kwargs", + [{"p_qdq": "fp8"}, {"v_qdq": "nvfp4", "v_qdq_amax": 1.0}], + ) + def test_skip_softmax_rejects_attention_quantization(self, qdq_kwargs): + """The raw contiguous kernel rejects quantization with calibrated skip.""" seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 - scale = 1.0 / (head_dim**0.5) torch.manual_seed(19) 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, + skip_softmax_threshold=1e-3, + **qdq_kwargs, + ) 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_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index b44d452be02..d1c95a1a8a6 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( @@ -1026,40 +1030,6 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): assert calls["query"].dtype == torch.float32 -def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): - """Split-local maxima must not change calibrated skip-softmax semantics.""" - impl = _clone_sparse_impl(_make_old_impl()) - impl.quant_kw = { - "p_qdq": "nvfp4", - "p_qdq_amax": 1.0, - "v_qdq": "nvfp4", - "v_qdq_amax": 6.0 * 448.0, - } - impl.sparse_kw = {"skip_softmax_threshold": 0.001} - q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) - metadata = _flash_attention_metadata(1, 16) - captured = {} - - monkeypatch.setattr(vllm_plugin, "fake_quant_v_onwrite", lambda *args, **kwargs: None) - monkeypatch.setattr( - vllm_plugin, - "triton_decode_attention", - lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("split-K kernel")), - ) - - def fake_attention(query, **kwargs): - captured.update(kwargs) - return torch.zeros_like(query) - - monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) - output = torch.empty_like(q) - assert impl.forward(None, q, q, q, kv_cache, metadata, output=output) is output - assert captured["skip_softmax_threshold"] == pytest.approx(0.001) - assert captured["p_qdq"] == "nvfp4" - assert captured["v_qdq"] == "nvfp4" - - def test_resolve_calibrated_skip_softmax_threshold_for_decode(): """Calibration conversion is phase-aware even when decode later delegates.""" sparse_kw = { 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..543fbd0a2da --- /dev/null +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -0,0 +1,544 @@ +# 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.""" + +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.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): + from modelopt.torch.quantization.plugins import vllm as quant_plugin + + 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'.""" + import sys + + 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/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index 62395ff5a7a..8e6bef1db55 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,49 @@ 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"}, {"v_qdq": "nvfp4"}]) +def test_forward_rejects_skip_softmax_with_attention_quantization(monkeypatch, qdq_kwargs): + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + 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="skip-softmax cannot be combined"): + triton_fa.attention( + q, + k, + v, + starts, + lengths, + seq_len, + skip_softmax_threshold=0.1, + measure_sparsity=True, + **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..a1ec161804e --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_calibrate_mask_reuse_cli.py @@ -0,0 +1,261 @@ +# 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", + "--max-reuse-dropped-mass", + "0.03", + "--max-reuse-selection-dropped-mass", + "0.01", + "--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 "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_collect_mask_reuse_cli.py b/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py new file mode 100644 index 00000000000..c5754b43e76 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py @@ -0,0 +1,392 @@ +# 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 launch/RPC test for the vLLM mask-reuse collection driver.""" + +import importlib.util +import json +import os +import sys +from hashlib import sha256 +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration import source_manifest +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + create_checkpoint_manifest, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.mask_reuse_capture import ( + canonical_json_sha256, +) + +_SCRIPT_PATH = Path(__file__).parents[5] / "examples/vllm_serve/collect_mask_reuse.py" +_SPEC = importlib.util.spec_from_file_location("collect_mask_reuse_cli", _SCRIPT_PATH) +assert _SPEC is not None and _SPEC.loader is not None +collect_mask_reuse_cli = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(collect_mask_reuse_cli) + +_FA4_COMMIT = "a" * 40 +_FA4_TREE = "b" * 40 +_FA4_ARCHIVE_SHA256 = "c" * 64 + + +def _write_source_witness(source: Path, destination: Path) -> str: + snapshot = source_manifest._source_tree_snapshot(source) + raw = { + "source_manifest_schema_version": 1, + "source_kind": "flash-attention-4", + "git_commit": _FA4_COMMIT, + "git_tree": _FA4_TREE, + "git_archive_sha256": _FA4_ARCHIVE_SHA256, + "archive_scope": "flash_attn", + "directories": list(snapshot.directories), + "files": list(snapshot.files), + } + payload = ( + json.dumps(raw, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + destination.write_bytes(payload) + return sha256(payload).hexdigest() + + +class _Tokenizer: + def encode(self, prompt, *, add_special_tokens): + assert add_special_tokens is True + offset = ord(prompt[0]) + return [offset + index for index in range(256)] + + +class _SamplingParams: + def __init__(self, **kwargs): + self.kwargs = kwargs + + +class _LLM: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.invocation = None + self.events = [] + self.__class__.instances.append(self) + + def get_tokenizer(self): + return _Tokenizer() + + def collective_rpc(self, method, args=()): + self.events.append(method) + if method == "mask_reuse_capture_status": + return [ + { + "capture_schema_version": 2, + "available": True, + "rank": 0, + "world_size": 1, + "reason": None, + } + ] + if method == "mask_reuse_capture_begin": + self.invocation = args[0] + return [ + { + "capture_schema_version": 2, + "armed": True, + "rank": 0, + "world_size": 1, + "invocation_sha256": canonical_json_sha256(self.invocation), + } + ] + assert method == "mask_reuse_capture_drain" + invocation = self.invocation + return [ + { + "capture_schema_version": 2, + "rank": 0, + "world_size": 1, + "invocation": invocation, + "invocation_sha256": canonical_json_sha256(invocation), + "geometry": invocation["expected_geometry"], + "global_num_heads": 2, + "eligible_tiles": 3, + "anchor_stats_by_layer": { + "0": {"retained_tiles": [2, 3], "dropped_mass": [0.01, 0.02]} + }, + "consumer_layers": { + "1": { + "anchor_layer": 0, + "consumer_head_start": 0, + "dropped_mass": [[0.01, 0.02], [0.03, 0.04]], + } + }, + "attention_call_counts": {"prefill": 2, "decode": 0}, + "tp_head_order_evidence": { + "sentinel_device_type": "cuda", + "gather_dim": 0, + "local_rank": 0, + "local_num_heads": 2, + "gathered_rank_local_head": [[0, 0], [0, 1]], + }, + "dense_shadow_evidence": { + "enabled": True, + "atol_hex": (0.0).hex(), + "rtol_hex": (0.0).hex(), + "validated_layer_indices": [0, 1], + }, + } + ] + + def generate(self, token_ids, sampling, *, use_tqdm): + assert len(token_ids) == 256 + assert sampling.kwargs == {"temperature": 0.0, "max_tokens": 1, "ignore_eos": True} + assert use_tqdm is False + self.events.append("generate") + return [] + + +def test_main_bootstraps_policy_free_backend_and_writes_normalized_evidence(tmp_path, monkeypatch): + prompts = tmp_path / "prompts.jsonl" + prompts.write_text( + "\n".join( + json.dumps( + { + "split": split, + "partition": ("development" if split == "calibration" else "outer_test"), + "inner_fold": 0 if split == "calibration" else None, + "prompt_id": prompt_id, + "source": source, + "source_group_sha256": sha256(source.encode()).hexdigest(), + "prompt": text, + "min_kv_tokens": 129, + "max_kv_tokens": 512, + } + ) + for split, prompt_id, source, text in ( + ("calibration", "cal-0", "ruler/niah", "alpha"), + ("heldout", "held-0", "longbench/qasper", "beta"), + ) + ) + + "\n", + encoding="utf-8", + ) + vanilla = tmp_path / "config.json" + vanilla.write_text( + json.dumps( + { + "threshold_scale_factor": { + "formula": "a * exp(b * target_sparsity)", + "prefill": { + "a": 1.0, + "b": 1.0, + "min_observed_sparsity": 0.5, + "max_observed_sparsity": 0.8, + }, + } + } + ), + encoding="utf-8", + ) + output = tmp_path / "observations.jsonl" + manifest = tmp_path / "manifest.json" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{}\n", encoding="utf-8") + (checkpoint / "model.safetensors").write_bytes(b"toy-weights") + checkpoint_manifest = create_checkpoint_manifest(checkpoint, model="test-model") + fa4_source = tmp_path / "flash-attention" + cute = fa4_source / "flash_attn/cute" + cute.mkdir(parents=True) + (cute / "interface.py").write_text("# pinned\n", encoding="utf-8") + (cute / "block_sparsity.py").write_text("# pinned\n", encoding="utf-8") + fa4_source_manifest = tmp_path / "fa4-source-manifest.json" + fa4_source_manifest_sha256 = _write_source_witness(fa4_source, fa4_source_manifest) + fake_vllm = SimpleNamespace(LLM=_LLM, SamplingParams=_SamplingParams) + monkeypatch.setitem(sys.modules, "vllm", fake_vllm) + monkeypatch.setattr( + collect_mask_reuse_cli.importlib.metadata, + "entry_points", + lambda **kwargs: [ + SimpleNamespace(name="mask_reuse_fa4", value="mask_reuse_vllm.plugin:register") + ], + ) + monkeypatch.setenv("MASK_REUSE_FA4_POLICY", "stale-policy.json") + monkeypatch.setenv("MASK_REUSE_FA4_POLICY_SHA256", "stale") + parsed_payloads = {} + real_parse_prompts = collect_mask_reuse_cli.parse_prompt_specs_jsonl + real_parse_vanilla = collect_mask_reuse_cli.parse_vanilla_prefill_fit + + def parse_prompts(payload): + parsed_payloads["prompts"] = payload + return real_parse_prompts(payload) + + def parse_vanilla(payload): + parsed_payloads["vanilla"] = payload + return real_parse_vanilla(payload) + + monkeypatch.setattr(collect_mask_reuse_cli, "parse_prompt_specs_jsonl", parse_prompts) + monkeypatch.setattr(collect_mask_reuse_cli, "parse_vanilla_prefill_fit", parse_vanilla) + _LLM.instances.clear() + + result = collect_mask_reuse_cli.main( + [ + str(checkpoint), + "--model-id", + "test-model", + "--checkpoint-manifest-sha256", + checkpoint_manifest.sha256, + "--plan", + "test_stride2", + "--fa4-source", + str(fa4_source), + "--fa4-source-manifest", + str(fa4_source_manifest), + "--fa4-source-manifest-sha256", + fa4_source_manifest_sha256, + "--fa4-commit", + _FA4_COMMIT, + "--prompts-jsonl", + str(prompts), + "--vanilla-config", + str(vanilla), + "--target-sparsities", + "0.7", + "--output", + str(output), + "--output-manifest", + str(manifest), + "--max-model-len", + "512", + "--validate-dense-output", + ] + ) + + assert result == 0 + engine = _LLM.instances[0] + assert engine.kwargs["attention_backend"] == "CUSTOM" + assert engine.kwargs["dtype"] == "bfloat16" + assert engine.kwargs["worker_cls"].endswith("MaskReuseCaptureWorker") + assert engine.kwargs["max_num_batched_tokens"] == 8192 + assert "quantization" not in engine.kwargs + assert engine.events == [ + "mask_reuse_capture_status", + "mask_reuse_capture_begin", + "generate", + "mask_reuse_capture_drain", + "mask_reuse_capture_begin", + "generate", + "mask_reuse_capture_drain", + ] + assert collect_mask_reuse_cli.os.environ["MASK_REUSE_FA4_CALIBRATION_CAPTURE"] == "1" + assert ( + collect_mask_reuse_cli.os.environ["MASK_REUSE_FA4_CHECKPOINT_MANIFEST_SHA256"] + == checkpoint_manifest.sha256 + ) + assert collect_mask_reuse_cli.os.environ["PYTHONDONTWRITEBYTECODE"] == "1" + assert collect_mask_reuse_cli.os.environ["MASK_REUSE_FA4_CAPTURE_DENSE_SHADOW"] == "1" + assert "MASK_REUSE_FA4_POLICY" not in collect_mask_reuse_cli.os.environ + assert "MASK_REUSE_FA4_POLICY_SHA256" not in collect_mask_reuse_cli.os.environ + + captures = [json.loads(line) for line in output.read_text().splitlines()] + assert len(captures) == 2 + assert all(capture["compact_capture_schema_version"] == 1 for capture in captures) + assert all("observations" not in capture for capture in captures) + assert {capture["invocation"]["split"] for capture in captures} == { + "calibration", + "heldout", + } + assert {capture["invocation"]["target_sparsity_hex"] for capture in captures} == {(0.7).hex()} + report = json.loads(manifest.read_text()) + assert report["capture_manifest_schema_version"] == 4 + assert report["capture_protocol"] == "modelopt_vllm_mask_reuse_target_sparsity_v4" + assert report["checkpoint_manifest_sha256"] == checkpoint_manifest.sha256 + assert report["prompt_plan_file_sha256"] == sha256(parsed_payloads["prompts"]).hexdigest() + assert report["vanilla_config_file_sha256"] == sha256(parsed_payloads["vanilla"]).hexdigest() + assert report["fa4_source_commit"] == _FA4_COMMIT + assert report["fa4_source_git_tree"] == _FA4_TREE + assert report["fa4_source_git_archive_sha256"] == _FA4_ARCHIVE_SHA256 + assert report["fa4_source_manifest_sha256"] == fa4_source_manifest_sha256 + assert report["fa4_source_file_count"] == 2 + assert report["dense_shadow_validation_requested"] is True + assert report["engine_kwargs"]["tensor_parallel_size"] == 1 + assert report["capture_count"] == len(captures) + assert "observation_count" not in report + assert len(report["captures"]) == 2 + assert {capture["invocation"]["source"] for capture in report["captures"]} == { + "ruler/niah", + "longbench/qasper", + } + + +def test_publish_no_clobber_preserves_destination_created_by_racer(tmp_path, monkeypatch): + temporary = tmp_path / "capture.tmp" + destination = tmp_path / "capture.jsonl" + temporary.write_text("ours", encoding="utf-8") + real_link = os.link + + def racing_link(source, target, **kwargs): + Path(target).write_text("racer", encoding="utf-8") + return real_link(source, target, **kwargs) + + monkeypatch.setattr(collect_mask_reuse_cli.os, "link", racing_link) + + with pytest.raises(FileExistsError): + collect_mask_reuse_cli._publish_no_clobber(temporary, destination) + + assert destination.read_text(encoding="utf-8") == "racer" + assert temporary.read_text(encoding="utf-8") == "ours" + + +def test_publish_no_clobber_rolls_back_capture_and_manifest_on_fsync_failure(tmp_path, monkeypatch): + def fail_fsync(path): + raise OSError("injected fsync failure") + + monkeypatch.setattr(collect_mask_reuse_cli, "_fsync_directory", fail_fsync) + + for name in ("capture.jsonl", "capture.manifest.json"): + temporary = tmp_path / f".{name}.tmp" + destination = tmp_path / name + temporary.write_text("complete", encoding="utf-8") + + with pytest.raises(OSError, match="injected fsync failure"): + collect_mask_reuse_cli._publish_no_clobber(temporary, destination) + + assert not destination.exists() + assert not temporary.exists() + + +def test_capture_environment_rejects_extra_fa4_source_path(tmp_path, monkeypatch): + fa4_source = tmp_path / "flash-attention" + cute = fa4_source / "flash_attn/cute" + cute.mkdir(parents=True) + (cute / "interface.py").write_text("# pinned\n", encoding="utf-8") + (cute / "block_sparsity.py").write_text("# pinned\n", encoding="utf-8") + manifest = tmp_path / "fa4-source-manifest.json" + digest = _write_source_witness(fa4_source, manifest) + (cute / "local_override.py").write_text("# extra\n", encoding="utf-8") + + with pytest.raises( + collect_mask_reuse_cli.CaptureContractError, + match="does not exactly match", + ): + collect_mask_reuse_cli._configure_capture_environment( + "test_stride2", + str(fa4_source), + str(manifest), + digest, + _FA4_COMMIT, + "0" * 64, + validate_dense_output=True, + ) 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..0531845ba23 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py @@ -0,0 +1,318 @@ +# 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, + max_reuse_dropped_mass=0.1, + max_reuse_selection_dropped_mass=0.05, + ) + + +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": [1, 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"]["constraint_violation_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 == 1 + else row + for row in _observations() + ] + + hostile = _calibrate(hostile_rows) + + assert hostile["context_policies"] == baseline["context_policies"] + assert ( + hostile["calibration_report"]["overall"]["reuse_heldout"]["constraint_violation_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_capture.py b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py new file mode 100644 index 00000000000..10064db2a91 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py @@ -0,0 +1,335 @@ +# 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 the strict ModelOpt mask-reuse capture contract.""" + +import json +import math +from dataclasses import replace +from hashlib import sha256 + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.plugins.mask_reuse_capture import ( + CaptureContractError, + PromptSpec, + build_capture_invocation, + canonical_json_sha256, + load_prompt_specs, + merge_rank_captures, + source_capture_sha256, + validate_begin_acks, + validate_capture_statuses, +) + +_CHECKPOINT = sha256(b"checkpoint-manifest").hexdigest() +_GROUP = sha256(b"source-group-0").hexdigest() + + +def _fit(): + return { + "formula": "a * exp(b * target_sparsity)", + "prefill": { + "a": 1.0, + "b": 1.0, + "min_observed_sparsity": 0.4, + "max_observed_sparsity": 0.8, + }, + } + + +def _prompt(split="calibration", prompt_id="p0"): + partition = "development" if split == "calibration" else "outer_test" + return PromptSpec( + split=split, + partition=partition, + inner_fold=0 if partition == "development" else None, + prompt_id=prompt_id, + source="ruler/niah", + source_group_sha256=_GROUP, + prompt="prompt", + min_kv_tokens=8192, + max_kv_tokens=16384, + ) + + +def _invocation(): + return build_capture_invocation( + model="test-model", + checkpoint_manifest_sha256=_CHECKPOINT, + prompt=_prompt(), + prompt_token_ids=list(range(8448)), + target_sparsity=0.7, + threshold_scale_factor=_fit(), + ) + + +def _rank_payload(invocation, rank): + start = rank * 2 + return { + "capture_schema_version": 2, + "rank": rank, + "world_size": 2, + "invocation": invocation, + "invocation_sha256": canonical_json_sha256(invocation), + "geometry": invocation["expected_geometry"], + "global_num_heads": 4, + "eligible_tiles": 131, + "anchor_stats_by_layer": { + "0": { + "retained_tiles": [10, 11, 12, 13], + "dropped_mass": [0.01, 0.02, 0.03, 0.04], + }, + "4": { + "retained_tiles": [14, 15, 16, 17], + "dropped_mass": [0.05, 0.06, 0.07, 0.08], + }, + }, + "consumer_layers": { + "1": { + "anchor_layer": 0, + "consumer_head_start": start, + "dropped_mass": [ + [0.01 * (start + local + donor + 1) for donor in range(4)] for local in range(2) + ], + }, + "5": { + "anchor_layer": 4, + "consumer_head_start": start, + "dropped_mass": [ + [0.01 * (start + local + donor + 2) for donor in range(4)] for local in range(2) + ], + }, + }, + "attention_call_counts": {"prefill": 8, "decode": 0}, + "tp_head_order_evidence": { + "sentinel_device_type": "cuda", + "gather_dim": 0, + "local_rank": rank, + "local_num_heads": 2, + "gathered_rank_local_head": [[0, 0], [0, 1], [1, 0], [1, 1]], + }, + "dense_shadow_evidence": { + "enabled": True, + "atol_hex": (0.0).hex(), + "rtol_hex": (0.0).hex(), + "validated_layer_indices": [0, 1, 4, 5], + }, + } + + +def test_build_invocation_binds_exact_threshold_source_and_final_chunk(): + invocation = _invocation() + expected_log2 = math.log2(1.0) + 0.7 * math.log2(math.e) - math.log2(8448) + + assert invocation["target_sparsity_hex"] == (0.7).hex() + assert invocation["checkpoint_manifest_sha256"] == _CHECKPOINT + assert invocation["source_group_sha256"] == _GROUP + assert invocation["partition"] == "development" + assert invocation["inner_fold"] == 0 + assert invocation["threshold_log2_hex"] == expected_log2.hex() + assert invocation["threshold_lambda_hex"] == (2.0**expected_log2).hex() + assert invocation["expected_geometry"] == { + "q_tokens": 256, + "kv_tokens": 8448, + "q_start_tokens": 8192, + } + assert invocation["source_capture_sha256"] == source_capture_sha256(list(range(8448))) + + +def test_backend_schema_v2_invocation_golden_sha_is_byte_exact(): + invocation = build_capture_invocation( + model="toy", + checkpoint_manifest_sha256=_CHECKPOINT, + prompt=PromptSpec( + split="calibration", + partition="development", + inner_fold=0, + prompt_id="p", + source="s", + source_group_sha256=_GROUP, + prompt="unused", + min_kv_tokens=1, + max_kv_tokens=None, + ), + prompt_token_ids=list(range(33_024)), + target_sparsity=0.7, + threshold_scale_factor={ + "formula": "a * exp(b * target_sparsity)", + "prefill": { + "a": 14.47, + "b": 10.91, + "min_observed_sparsity": 0.0, + "max_observed_sparsity": 1.0, + }, + }, + ) + + assert canonical_json_sha256(invocation) == ( + "d1bb38b0611b7a424f70e4567f50380ad8fef9f77ebf05df97643fea08367056" + ) + + +def test_source_fingerprint_does_not_hide_split_overlap(): + token_ids = [1, 2, 3] + assert source_capture_sha256(token_ids) == source_capture_sha256(token_ids) + assert source_capture_sha256(token_ids) != source_capture_sha256([1, 2, 4]) + + +def test_prompt_plan_requires_unique_ids_within_each_split(tmp_path): + path = tmp_path / "prompts.jsonl" + rows = [ + { + "split": split, + "partition": "development" if split == "calibration" else "outer_test", + "inner_fold": 0 if split == "calibration" else None, + "prompt_id": prompt_id, + "source": "dataset", + "source_group_sha256": sha256(f"{split}-{prompt_id}".encode()).hexdigest(), + "prompt": text, + "min_kv_tokens": minimum, + "max_kv_tokens": maximum, + } + for split, prompt_id, text, minimum, maximum in ( + ("calibration", "same", "a", 129, 512), + ("calibration", "same", "b", 513, 1024), + ("heldout", "held-0", "c", 129, 512), + ("heldout", "held-1", "d", 513, 1024), + ) + ] + path.write_text("\n".join(json.dumps(row) for row in rows) + "\n", encoding="utf-8") + + with pytest.raises(CaptureContractError, match="unique within"): + load_prompt_specs(path) + + +def test_build_invocation_rejects_unqualified_final_chunk_and_extrapolated_target(): + with pytest.raises(CaptureContractError, match="at least 129"): + build_capture_invocation( + model="test-model", + checkpoint_manifest_sha256=_CHECKPOINT, + prompt=replace(_prompt(), max_kv_tokens=9000), + prompt_token_ids=list(range(8200)), + target_sparsity=0.7, + threshold_scale_factor=_fit(), + ) + with pytest.raises(CaptureContractError, match="outside the observed"): + build_capture_invocation( + model="test-model", + checkpoint_manifest_sha256=_CHECKPOINT, + prompt=_prompt(), + prompt_token_ids=list(range(8448)), + target_sparsity=0.9, + threshold_scale_factor=_fit(), + ) + + +def test_status_and_begin_require_every_rank_and_exact_invocation(): + statuses = [ + { + "capture_schema_version": 2, + "available": True, + "rank": rank, + "world_size": 2, + "reason": None, + } + for rank in range(2) + ] + assert validate_capture_statuses(statuses) == 2 + invocation = _invocation() + digest = canonical_json_sha256(invocation) + acknowledgements = [ + { + "capture_schema_version": 2, + "armed": True, + "rank": rank, + "world_size": 2, + "invocation_sha256": digest, + } + for rank in range(2) + ] + assert validate_begin_acks(acknowledgements, invocation) == 2 + + statuses[1]["available"] = False + statuses[1]["reason"] = "sink disabled" + with pytest.raises(CaptureContractError, match="sink disabled"): + validate_capture_statuses(statuses) + with pytest.raises(CaptureContractError, match="exactly cover ranks"): + validate_begin_acks(acknowledgements[:1], invocation) + + +def test_merge_rank_captures_concatenates_consumer_shards_deterministically(): + invocation = _invocation() + merged = merge_rank_captures( + [_rank_payload(invocation, 1), _rank_payload(invocation, 0)], invocation + ) + + assert len(merged.capture["consumer_layers"]) == 2 + assert merged.manifest["world_size"] == 2 + assert merged.manifest["global_num_heads"] == 4 + assert merged.manifest["candidate_cell_count"] == 2 * 4 * 4 + assert merged.manifest["attention_call_counts"] == {"prefill": 8, "decode": 0} + assert [item["global_head_start"] for item in merged.manifest["tp_head_order_evidence"]] == [ + 0, + 2, + ] + assert merged.manifest["dense_shadow_evidence"]["validated_layer_indices"] == [0, 1, 4, 5] + consumer = merged.capture["consumer_layers"]["1"] + assert consumer["anchor_layer"] == 0 + assert consumer["dropped_mass"][2][3] == pytest.approx(0.06) + assert set(merged.capture["anchor_stats_by_layer"]) == {"0", "4"} + + +@pytest.mark.parametrize( + "corruption", + [ + "wrong_invocation", + "anchor_disagreement", + "head_gap", + "rank_permutation", + "decode_call", + "prefill_count", + "missing_dense_shadow", + ], +) +def test_merge_rank_captures_fails_closed_on_incomplete_or_disagreed_evidence(corruption): + invocation = _invocation() + rank0 = _rank_payload(invocation, 0) + rank1 = _rank_payload(invocation, 1) + if corruption == "wrong_invocation": + rank1["invocation"] = dict(invocation, prompt_id="other") + elif corruption == "anchor_disagreement": + rank1["anchor_stats_by_layer"]["0"]["retained_tiles"][0] = 9 + elif corruption == "head_gap": + rank1["consumer_layers"]["1"]["consumer_head_start"] = 3 + elif corruption == "rank_permutation": + rank1["tp_head_order_evidence"]["gathered_rank_local_head"] = [ + [1, 0], + [1, 1], + [0, 0], + [0, 1], + ] + elif corruption == "decode_call": + rank0["attention_call_counts"]["decode"] = 1 + rank1["attention_call_counts"]["decode"] = 1 + elif corruption == "prefill_count": + rank0["attention_call_counts"]["prefill"] = 7 + rank1["attention_call_counts"]["prefill"] = 7 + else: + rank0["dense_shadow_evidence"]["validated_layer_indices"] = [0, 1, 4] + rank1["dense_shadow_evidence"]["validated_layer_indices"] = [0, 1, 4] + + with pytest.raises(CaptureContractError): + merge_rank_captures([rank0, rank1], invocation) 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..347be356a6f --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py @@ -0,0 +1,394 @@ +# 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, + "max_reuse_dropped_mass": 0.025, + "max_reuse_selection_dropped_mass": 0.02, + } + + 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 == legacy_report + 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_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, + max_reuse_dropped_mass=0.025, + ) + + +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, + max_reuse_dropped_mass=0.025, + ) + + 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, + max_reuse_dropped_mass=0.025, + ) + + +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, + max_reuse_dropped_mass=0.025, + ) + + +def test_selector_minimizes_declared_equal_bmm_combined_cost(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, + max_reuse_dropped_mass=0.025, + max_reuse_selection_dropped_mass=0.02, + ) + + assert candidate["context_policies"][0]["target_sparsity"] == 0.7 + frontier = candidate["calibration_report"]["by_bucket"][0]["target_sparsity_frontier"] + assert [row["combined_tile_cost"] for row in frontier] == [14, 10] + + +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, + max_reuse_dropped_mass=0.025, + ) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py b/tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py new file mode 100644 index 00000000000..c784d32308f --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py @@ -0,0 +1,327 @@ +# 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 exact, Git-free source-tree witnesses.""" + +import json +import os +import shutil +import subprocess +import tarfile +from hashlib import sha256 +from pathlib import Path + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration import source_manifest +from modelopt.torch.sparsity.attention_sparsity.calibration.source_manifest import ( + SourceManifestError, + create_source_manifest_from_git_archive, + verify_source_manifest, +) + +_COMMIT = "a" * 40 +_TREE = "b" * 40 +_ARCHIVE_SHA256 = "c" * 64 + + +def _toy_source(root: Path) -> Path: + (root / "flash_attn/cute").mkdir(parents=True) + (root / "flash_attn/empty-dir").mkdir() + (root / "flash_attn/cute/interface.py").write_text("# interface\n", encoding="utf-8") + (root / "flash_attn/cute/block_sparsity.py").write_text("# sparse\n", encoding="utf-8") + return root + + +def _write_witness(source: Path, manifest: Path) -> str: + snapshot = source_manifest._source_tree_snapshot(source) + raw = { + "source_manifest_schema_version": 1, + "source_kind": "flash-attention-4", + "git_commit": _COMMIT, + "git_tree": _TREE, + "git_archive_sha256": _ARCHIVE_SHA256, + "archive_scope": "flash_attn", + "directories": list(snapshot.directories), + "files": list(snapshot.files), + } + payload = ( + json.dumps(raw, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + manifest.write_bytes(payload) + return sha256(payload).hexdigest() + + +def _verify(source: Path, manifest: Path, digest: str): + return verify_source_manifest( + source, + manifest, + expected_manifest_sha256=digest, + expected_commit=_COMMIT, + expected_source_kind="flash-attention-4", + ) + + +def _symlink_or_skip(path: Path, target: str) -> None: + try: + path.symlink_to(target) + except OSError as error: + pytest.skip(f"symbolic links unavailable: {error}") + + +def test_source_witness_accepts_exact_files_and_directories(tmp_path): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + digest = _write_witness(source, manifest) + + verified = _verify(source, manifest, digest) + + assert verified.git_commit == _COMMIT + assert verified.git_tree == _TREE + assert verified.manifest_sha256 == digest + assert verified.file_count == 2 + assert verified.directory_count == 3 + + +@pytest.mark.parametrize( + "change", + [ + "mutated", + "missing", + "extra", + pytest.param( + "mode", + marks=pytest.mark.skipif(os.name == "nt", reason="Windows has no POSIX execute bit"), + ), + ], +) +def test_source_witness_rejects_file_tree_changes(tmp_path, change): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + digest = _write_witness(source, manifest) + interface = source / "flash_attn/cute/interface.py" + if change == "mutated": + interface.write_text("# modified\n", encoding="utf-8") + elif change == "missing": + interface.unlink() + elif change == "extra": + (source / "flash_attn/cute/shadow.py").write_text("# extra\n", encoding="utf-8") + else: + interface.chmod(interface.stat().st_mode | 0o111) + + with pytest.raises(SourceManifestError, match="does not exactly match"): + _verify(source, manifest, digest) + + +def test_source_witness_rejects_regular_file_replaced_by_symlink(tmp_path): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + digest = _write_witness(source, manifest) + interface = source / "flash_attn/cute/interface.py" + interface.unlink() + _symlink_or_skip(interface, "block_sparsity.py") + + with pytest.raises(SourceManifestError, match="must not be a link"): + _verify(source, manifest, digest) + + +def test_source_witness_rejects_extra_symlink(tmp_path): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + digest = _write_witness(source, manifest) + _symlink_or_skip(source / "flash_attn/unexpected-link", "cute/interface.py") + with pytest.raises(SourceManifestError, match="must not be a link"): + _verify(source, manifest, digest) + + +def test_source_witness_requires_independent_manifest_and_commit_pins(tmp_path): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + digest = _write_witness(source, manifest) + + with pytest.raises(SourceManifestError, match="expected SHA256"): + _verify(source, manifest, "0" * 64) + with pytest.raises(SourceManifestError, match="expected Git commit"): + verify_source_manifest( + source, + manifest, + expected_manifest_sha256=digest, + expected_commit="0" * 40, + expected_source_kind="flash-attention-4", + ) + + +def test_source_witness_rejects_boolean_schema_version(tmp_path): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + _write_witness(source, manifest) + raw = json.loads(manifest.read_bytes()) + raw["source_manifest_schema_version"] = True + payload = ( + json.dumps(raw, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + manifest.write_bytes(payload) + + with pytest.raises(SourceManifestError, match="schema_version must be 1"): + _verify(source, manifest, sha256(payload).hexdigest()) + + +def test_source_artifact_temporary_is_cleaned_when_write_setup_fails(tmp_path, monkeypatch): + destination = tmp_path / "source.tar" + real_dup = source_manifest.os.dup + + def fail_dup(descriptor): + assert descriptor >= 0 + raise OSError("injected dup failure") + + monkeypatch.setattr(source_manifest.os, "dup", fail_dup) + + with pytest.raises(OSError, match="injected dup failure"): + source_manifest._temporary_payload(destination, b"archive") + + monkeypatch.setattr(source_manifest.os, "dup", real_dup) + assert list(tmp_path.iterdir()) == [] + + +def test_source_artifact_publication_preserves_destination_racer(tmp_path, monkeypatch): + archive = tmp_path / "source.tar" + manifest = tmp_path / "source-manifest.json" + real_link = source_manifest.os.link + link_count = 0 + + def racing_link(source, destination, **kwargs): + nonlocal link_count + link_count += 1 + if link_count == 2: + Path(destination).write_bytes(b"racer") + return real_link(source, destination, **kwargs) + + monkeypatch.setattr(source_manifest.os, "link", racing_link) + + with pytest.raises(SourceManifestError, match="destination appeared"): + source_manifest._publish_source_artifacts(archive, b"archive", manifest, b"manifest") + + assert not archive.exists() + assert manifest.read_bytes() == b"racer" + assert sorted(path.name for path in tmp_path.iterdir()) == [manifest.name] + + +def test_source_artifact_publication_preserves_post_publish_replacement(tmp_path, monkeypatch): + archive = tmp_path / "source.tar" + manifest = tmp_path / "source-manifest.json" + real_stable_hash = source_manifest.stable_file_sha256 + + def replace_archive_before_rehash(path, *, label): + if label == "published source archive": + Path(path).unlink() + Path(path).write_bytes(b"racer") + return real_stable_hash(path, label=label) + + monkeypatch.setattr(source_manifest, "stable_file_sha256", replace_archive_before_rehash) + + with pytest.raises(SourceManifestError, match="failed stable rehash"): + source_manifest._publish_source_artifacts(archive, b"archive", manifest, b"manifest") + + assert archive.read_bytes() == b"racer" + assert not manifest.exists() + assert sorted(path.name for path in tmp_path.iterdir()) == [archive.name] + + +@pytest.mark.skipif(shutil.which("git") is None, reason="Git is required by the generator") +def test_generator_witnesses_the_exact_archive_without_runtime_git(tmp_path): + checkout = tmp_path / "checkout" + checkout.mkdir() + subprocess.run(["git", "init", "-q", str(checkout)], check=True) + subprocess.run(["git", "-C", str(checkout), "config", "user.name", "Test"], check=True) + subprocess.run( + ["git", "-C", str(checkout), "config", "user.email", "test@example.com"], check=True + ) + _toy_source(checkout) + subprocess.run(["git", "-C", str(checkout), "add", "."], check=True) + subprocess.run(["git", "-C", str(checkout), "commit", "-qm", "source"], check=True) + commit = subprocess.run( + ["git", "-C", str(checkout), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + archive = tmp_path / "source.tar" + manifest = tmp_path / "source-manifest.json" + + generated = create_source_manifest_from_git_archive( + checkout, + expected_commit=commit, + source_kind="flash-attention-4", + archive_output=archive, + manifest_output=manifest, + ) + extracted = tmp_path / "extracted" + extracted.mkdir() + with tarfile.open(archive, "r:") as handle: + for member in handle.getmembers(): + path = extracted / member.name + if member.isdir(): + path.mkdir(parents=True, exist_ok=True) + else: + payload = handle.extractfile(member) + assert payload is not None + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(payload.read()) + path.chmod(member.mode) + + verified = verify_source_manifest( + extracted, + manifest, + expected_manifest_sha256=generated.manifest_sha256, + expected_commit=commit, + expected_source_kind="flash-attention-4", + ) + assert verified.git_archive_sha256 == generated.git_archive_sha256 + + +@pytest.mark.skipif(shutil.which("git") is None, reason="Git is required by the generator") +def test_generator_archive_excludes_ignored_checkout_artifacts(tmp_path): + checkout = tmp_path / "checkout" + checkout.mkdir() + subprocess.run(["git", "init", "-q", str(checkout)], check=True) + subprocess.run(["git", "-C", str(checkout), "config", "user.name", "Test"], check=True) + subprocess.run( + ["git", "-C", str(checkout), "config", "user.email", "test@example.com"], check=True + ) + (checkout / ".gitignore").write_text("*.pyc\n", encoding="utf-8") + (checkout / "flash_attn").mkdir() + (checkout / "flash_attn/source.py").write_text("# source\n", encoding="utf-8") + subprocess.run(["git", "-C", str(checkout), "add", "."], check=True) + subprocess.run(["git", "-C", str(checkout), "commit", "-qm", "source"], check=True) + (checkout / "flash_attn/shadow.pyc").write_bytes(b"ignored") + commit = subprocess.run( + ["git", "-C", str(checkout), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + archive = tmp_path / "source.tar" + manifest = tmp_path / "source-manifest.json" + generated = create_source_manifest_from_git_archive( + checkout, + expected_commit=commit, + source_kind="flash-attention-4", + archive_output=archive, + manifest_output=manifest, + ) + assert generated.manifest_sha256 == sha256(manifest.read_bytes()).hexdigest() + with tarfile.open(archive, "r:") as handle: + assert "flash_attn/shadow.pyc" not in handle.getnames() 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..00c5beec5ec --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py @@ -0,0 +1,258 @@ +# 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, + "min_observed_sparsity": 0.1, + "max_observed_sparsity": 0.9, + }, + "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, + "min_observed_sparsity": 0.1, + "max_observed_sparsity": 0.9, + } + 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_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 diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py b/tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py new file mode 100644 index 00000000000..348dacac1a9 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py @@ -0,0 +1,96 @@ +# 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 the policy-free vLLM capture worker bootstrap and RPC wrappers.""" + +import importlib +import sys +from types import ModuleType, SimpleNamespace + +import pytest + + +@pytest.fixture +def vllm_mask_reuse_capture(monkeypatch): + """Import the worker against a minimal optional-vLLM boundary.""" + + class Worker: + pass + + modules = { + "vllm": ModuleType("vllm"), + "vllm.v1": ModuleType("vllm.v1"), + "vllm.v1.worker": ModuleType("vllm.v1.worker"), + "vllm.v1.worker.gpu_worker": ModuleType("vllm.v1.worker.gpu_worker"), + } + modules["vllm.v1.worker.gpu_worker"].Worker = Worker + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + target = "modelopt.torch.sparsity.attention_sparsity.plugins.vllm_mask_reuse_capture" + sys.modules.pop(target, None) + module = importlib.import_module(target) + yield module + sys.modules.pop(target, None) + + +def _api(calls): + return SimpleNamespace( + configure_capture_runtime=lambda plan: calls.append(("configure", plan)), + capture_status=lambda: { + "capture_schema_version": 1, + "available": True, + "rank": 0, + "world_size": 1, + "reason": None, + }, + begin_capture=lambda invocation: calls.append(("begin", invocation)) or {"armed": True}, + drain_capture=lambda: calls.append(("drain",)) or {"records": []}, + ) + + +def test_bootstrap_requires_gate_and_plan(monkeypatch, vllm_mask_reuse_capture): + monkeypatch.delenv(vllm_mask_reuse_capture.CAPTURE_ENV, raising=False) + monkeypatch.setenv(vllm_mask_reuse_capture.PLAN_ENV, "qwen3_stride2") + with pytest.raises(RuntimeError, match="CALIBRATION_CAPTURE=1"): + vllm_mask_reuse_capture._configure_capture_before_model_load() + + monkeypatch.setenv(vllm_mask_reuse_capture.CAPTURE_ENV, "1") + monkeypatch.delenv(vllm_mask_reuse_capture.PLAN_ENV, raising=False) + with pytest.raises(RuntimeError, match="must name"): + vllm_mask_reuse_capture._configure_capture_before_model_load() + + +def test_bootstrap_installs_policy_free_runtime_before_load(monkeypatch, vllm_mask_reuse_capture): + calls = [] + api = _api(calls) + monkeypatch.setenv(vllm_mask_reuse_capture.CAPTURE_ENV, "1") + monkeypatch.setenv(vllm_mask_reuse_capture.PLAN_ENV, "nemotron3_ultra_stride2") + monkeypatch.setattr(vllm_mask_reuse_capture, "_capture_api", lambda: api) + + assert vllm_mask_reuse_capture._configure_capture_before_model_load() is api + assert calls == [("configure", "nemotron3_ultra_stride2")] + + +def test_worker_rpc_methods_forward_only_to_capture_api(monkeypatch, vllm_mask_reuse_capture): + calls = [] + api = _api(calls) + monkeypatch.setattr(vllm_mask_reuse_capture, "_capture_api", lambda: api) + worker = object.__new__(vllm_mask_reuse_capture.MaskReuseCaptureWorker) + invocation = {"capture_schema_version": 1} + + assert worker.mask_reuse_capture_status()["available"] is True + assert worker.mask_reuse_capture_begin(invocation) == {"armed": True} + assert worker.mask_reuse_capture_drain() == {"records": []} + assert calls == [("begin", invocation), ("drain",)]