From 8dcda4178a45544e4a204427655fbb82bede1202 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 24 Aug 2026 10:55:18 -0700 Subject: [PATCH] Add vLLM mask reuse capture Signed-off-by: Kai Xu --- examples/vllm_serve/collect_mask_reuse.py | 636 ++++++++ .../vllm_serve/create_fa4_source_witness.py | 57 + .../calibration/source_manifest.py | 613 ++++++++ .../sparsity/attention_sparsity/conversion.py | 6 +- .../plugins/mask_reuse_capture.py | 1274 +++++++++++++++++ .../plugins/vllm_mask_reuse_capture.py | 93 ++ .../test_collect_mask_reuse_cli.py | 392 +++++ .../test_mask_reuse_capture.py | 400 ++++++ .../test_source_manifest.py | 327 +++++ .../test_sparse_attn_calibration.py | 19 +- .../test_vllm_mask_reuse_capture_worker.py | 96 ++ 11 files changed, 3909 insertions(+), 4 deletions(-) create mode 100644 examples/vllm_serve/collect_mask_reuse.py create mode 100644 examples/vllm_serve/create_fa4_source_witness.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/calibration/source_manifest.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/plugins/vllm_mask_reuse_capture.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py diff --git a/examples/vllm_serve/collect_mask_reuse.py b/examples/vllm_serve/collect_mask_reuse.py new file mode 100644 index 00000000000..20df9b3f691 --- /dev/null +++ b/examples/vllm_serve/collect_mask_reuse.py @@ -0,0 +1,636 @@ +# 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, + merge_rank_topology_discovery_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" +TOPOLOGY_MAX_REUSE_SPAN_ENV = "MASK_REUSE_FA4_TOPOLOGY_MAX_REUSE_SPAN" +_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( + "--max-reuse-span", + type=int, + default=None, + help=( + "Maximum number of intervening attention-layer positions considered by a " + "*_topology_discovery plan" + ), + ) + 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, + max_reuse_span: int | None = None, +) -> 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" + topology_discovery = plan.endswith("_topology_discovery") + if topology_discovery: + if max_reuse_span is None or max_reuse_span <= 0: + raise CaptureContractError("a *_topology_discovery plan requires --max-reuse-span > 0") + os.environ[TOPOLOGY_MAX_REUSE_SPAN_ENV] = str(max_reuse_span) + elif max_reuse_span is not None: + raise CaptureContractError( + "--max-reuse-span is valid only with a *_topology_discovery plan" + ) + else: + os.environ.pop(TOPOLOGY_MAX_REUSE_SPAN_ENV, None) + 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, + max_reuse_span=args.max_reuse_span, + ) + # 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) + rank_captures = llm.collective_rpc("mask_reuse_capture_drain") + merge = ( + merge_rank_topology_discovery_captures + if args.plan.endswith("_topology_discovery") + else merge_rank_captures + ) + merged = merge(rank_captures, 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() + topology_discovery = args.plan.endswith("_topology_discovery") + manifest = { + "capture_manifest_schema_version": 5 if topology_discovery else 4, + "capture_protocol": ( + "modelopt_vllm_mask_reuse_topology_discovery_v1" + if topology_discovery + else "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, + "capture_count": capture_count, + "candidate_cell_count": candidate_cell_count, + "captures": capture_manifests, + } + if topology_discovery: + manifest["capture_mode"] = "topology_discovery" + manifest["max_reuse_span"] = args.max_reuse_span + manifest["topology_discovery_capture_file_sha256"] = capture_sha256 + else: + manifest["compact_capture_file_sha256"] = capture_sha256 + 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_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/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/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index b2acd137930..971e11b37c1 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py @@ -50,10 +50,14 @@ def export_threshold_scale_factor(calibration_params: dict[str, Any]) -> dict[st block: dict[str, Any] = {"formula": "a * exp(b * target_sparsity)"} for phase in ("prefill", "decode"): if phase in calibration_params: - block[phase] = { + 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]) + block[phase] = phase_params return block 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..ad6b730b5a4 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py @@ -0,0 +1,1274 @@ +# 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", + "merge_rank_topology_discovery_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", + } +) +_RANK_TOPOLOGY_CAPTURE_FIELDS = frozenset( + { + "capture_schema_version", + "capture_mode", + "rank", + "world_size", + "invocation", + "invocation_sha256", + "geometry", + "global_num_heads", + "eligible_tiles", + "attention_layers", + "max_reuse_span", + "anchor_stats_by_layer", + "consumer_candidates_by_layer", + "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"}) +_CANDIDATE_STATS_FIELDS = frozenset({"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, + } + + +def _parse_topology_candidates( + raw: object, + *, + rank: int, + global_num_heads: int, + attention_layers: Sequence[int], + max_reuse_span: int, +) -> dict[int, dict[int, tuple[int, list[list[float]]]]]: + if not isinstance(raw, Mapping): + raise CaptureContractError(f"rank {rank} consumer_candidates_by_layer must be an object") + expected: dict[int, set[int]] = {} + for position, consumer in enumerate(attention_layers[1:], start=1): + first = max(0, position - max_reuse_span) + expected[consumer] = set(attention_layers[first:position]) + result: dict[int, dict[int, tuple[int, list[list[float]]]]] = {} + for raw_consumer, raw_by_anchor in raw.items(): + if ( + not isinstance(raw_consumer, str) + or not raw_consumer.isdigit() + or raw_consumer != str(int(raw_consumer)) + ): + raise CaptureContractError("topology consumer keys must be canonical integer strings") + consumer = int(raw_consumer) + if not isinstance(raw_by_anchor, Mapping): + raise CaptureContractError( + f"consumer_candidates_by_layer[{consumer}] must be an object" + ) + by_anchor: dict[int, tuple[int, list[list[float]]]] = {} + for raw_anchor, raw_stats in raw_by_anchor.items(): + if ( + not isinstance(raw_anchor, str) + or not raw_anchor.isdigit() + or raw_anchor != str(int(raw_anchor)) + ): + raise CaptureContractError("topology anchor keys must be canonical integer strings") + anchor = int(raw_anchor) + if not isinstance(raw_stats, Mapping): + raise CaptureContractError( + f"topology candidate {anchor}->{consumer} must be an object" + ) + _exact_fields( + raw_stats, + _CANDIDATE_STATS_FIELDS, + f"topology candidate {anchor}->{consumer}", + ) + start = _integer( + raw_stats["consumer_head_start"], + f"topology candidate {anchor}->{consumer}.consumer_head_start", + ) + matrix = raw_stats["dropped_mass"] + if not isinstance(matrix, list) or not matrix: + raise CaptureContractError( + f"topology candidate {anchor}->{consumer}.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"topology candidate {anchor}->{consumer} local head " + f"{local_head} must cover every donor" + ) + rows.append( + [ + _finite_number( + value, + f"topology candidate {anchor}->{consumer}[{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} topology candidate {anchor}->{consumer} shard exceeds head count" + ) + by_anchor[anchor] = (start, rows) + result[consumer] = by_anchor + observed = {consumer: set(by_anchor) for consumer, by_anchor in result.items()} + if observed != expected: + raise CaptureContractError( + f"rank {rank} does not contain the exact topology candidate window" + ) + return result + + +@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) + + +def merge_rank_topology_discovery_captures( + values: Sequence[Mapping[str, object]], invocation: Mapping[str, object] +) -> MergedCapture: + """Merge bounded all-candidate topology-discovery shards.""" + trusted_invocation = _validate_invocation(invocation) + expected_digest = canonical_json_sha256(trusted_invocation) + world_size, ordered = _validate_rank_envelope( + values, _RANK_TOPOLOGY_CAPTURE_FIELDS, "topology capture drain" + ) + global_heads: set[int] = set() + eligible_values: set[int] = set() + geometry_values: list[dict[str, int]] = [] + attention_layer_values: list[list[int]] = [] + max_span_values: set[int] = set() + anchor_payloads: list[dict[str, dict[str, list[int] | list[float]]]] = [] + candidate_payloads: list[dict[int, dict[int, tuple[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["capture_mode"] != "topology_discovery": + raise CaptureContractError( + f"topology capture drain rank {rank} has the wrong capture_mode" + ) + if value["invocation"] != trusted_invocation: + raise CaptureContractError( + f"topology capture drain rank {rank} echoed the wrong invocation" + ) + if value["invocation_sha256"] != expected_digest: + raise CaptureContractError( + f"topology capture drain rank {rank} has the wrong invocation digest" + ) + geometry = _validate_geometry( + value["geometry"], f"topology capture drain rank {rank} geometry" + ) + if geometry != trusted_invocation["expected_geometry"]: + raise CaptureContractError( + f"topology capture drain rank {rank} measured the wrong final chunk" + ) + geometry_values.append(geometry) + global_num_heads = _integer( + value["global_num_heads"], + f"topology capture drain rank {rank} global_num_heads", + minimum=1, + ) + eligible_tiles = _integer( + value["eligible_tiles"], + f"topology capture drain rank {rank} eligible_tiles", + minimum=1, + ) + raw_layers = value["attention_layers"] + if not isinstance(raw_layers, list): + raise CaptureContractError( + f"topology capture drain rank {rank} attention_layers must be a list" + ) + attention_layers = [_integer(layer, f"rank {rank} attention layer") for layer in raw_layers] + if len(attention_layers) < 2 or attention_layers != sorted(set(attention_layers)): + raise CaptureContractError( + f"topology capture drain rank {rank} attention_layers must be sorted and unique" + ) + max_reuse_span = _integer( + value["max_reuse_span"], + f"topology capture drain rank {rank} max_reuse_span", + minimum=1, + ) + if max_reuse_span >= len(attention_layers): + raise CaptureContractError( + "topology max_reuse_span must be smaller than the attention-layer count" + ) + global_heads.add(global_num_heads) + eligible_values.add(eligible_tiles) + attention_layer_values.append(attention_layers) + max_span_values.add(max_reuse_span) + anchors = _parse_anchor_stats( + value["anchor_stats_by_layer"], + global_num_heads=global_num_heads, + eligible_tiles=eligible_tiles, + ) + if {int(layer) for layer in anchors} != set(attention_layers): + raise CaptureContractError( + f"rank {rank} anchor statistics do not cover every attention layer" + ) + anchor_payloads.append(anchors) + candidates = _parse_topology_candidates( + value["consumer_candidates_by_layer"], + rank=rank, + global_num_heads=global_num_heads, + attention_layers=attention_layers, + max_reuse_span=max_reuse_span, + ) + candidate_payloads.append(candidates) + attention_call_counts.append( + _parse_attention_call_counts(value["attention_call_counts"], rank=rank) + ) + 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 consumer, by_anchor in candidates.items(): + for anchor, (start, rows) in by_anchor.items(): + if ( + start != tp_evidence["global_head_start"] + or len(rows) != tp_evidence["local_num_heads"] + ): + raise CaptureContractError( + f"rank {rank} topology candidate {anchor}->{consumer} " + "does not match its TP head shard" + ) + + if len(global_heads) != 1 or len(eligible_values) != 1 or len(max_span_values) != 1: + raise CaptureContractError( + "topology capture ranks disagree on heads, eligible tiles, or reuse span" + ) + if any(value != geometry_values[0] for value in geometry_values[1:]): + raise CaptureContractError("topology capture ranks disagree on final-chunk geometry") + if any(value != attention_layer_values[0] for value in attention_layer_values[1:]): + raise CaptureContractError("topology capture ranks disagree on attention layers") + if any(value != anchor_payloads[0] for value in anchor_payloads[1:]): + raise CaptureContractError("topology capture ranks disagree on anchor statistics") + if any(value != attention_call_counts[0] for value in attention_call_counts[1:]): + raise CaptureContractError("topology capture ranks disagree on attention call counts") + if any(value != dense_shadow_evidence[0] for value in dense_shadow_evidence[1:]): + raise CaptureContractError("topology capture ranks disagree on dense shadow evidence") + + global_num_heads = next(iter(global_heads)) + eligible_tiles = next(iter(eligible_values)) + attention_layers = attention_layer_values[0] + max_reuse_span = next(iter(max_span_values)) + if eligible_tiles != _eligible_tiles(geometry_values[0]): + raise CaptureContractError("topology capture eligible_tiles does not match causal geometry") + 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( + "topology capture prefill attention call count does not match chunks times layers" + ) + if common_counts["decode"] != 0: + raise CaptureContractError( + "topology 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( + "topology dense shadow evidence does not cover every attention layer" + ) + + expected_edges = { + (anchor, consumer) + for position, consumer in enumerate(attention_layers[1:], start=1) + for anchor in attention_layers[max(0, position - max_reuse_span) : position] + } + merged_candidates: dict[str, dict[str, dict[str, object]]] = {} + for anchor, consumer in sorted(expected_edges): + shards = sorted( + (payload[consumer][anchor] for payload in candidate_payloads), + key=lambda value: value[0], + ) + cursor = 0 + global_rows: list[list[float]] = [] + for start, local_rows in shards: + if start != cursor: + raise CaptureContractError( + f"topology candidate {anchor}->{consumer} shards are " + f"overlapping or incomplete at {cursor}" + ) + global_rows.extend(local_rows) + cursor += len(local_rows) + if cursor != global_num_heads: + raise CaptureContractError( + f"topology candidate {anchor}->{consumer} shards cover {cursor}, " + f"expected {global_num_heads}" + ) + merged_candidates.setdefault(str(consumer), {})[str(anchor)] = {"dropped_mass": global_rows} + + capture = { + "topology_discovery_capture_schema_version": 1, + "invocation": trusted_invocation, + "geometry": geometry_values[0], + "global_num_heads": global_num_heads, + "eligible_tiles": eligible_tiles, + "attention_layers": attention_layers, + "max_reuse_span": max_reuse_span, + "anchor_stats_by_layer": anchor_payloads[0], + "consumer_candidates_by_layer": merged_candidates, + } + manifest = { + "capture_schema_version": CAPTURE_SCHEMA_VERSION, + "capture_mode": "topology_discovery", + "invocation": trusted_invocation, + "invocation_sha256": expected_digest, + "world_size": world_size, + "global_num_heads": global_num_heads, + "eligible_tiles": eligible_tiles, + "attention_layers": attention_layers, + "max_reuse_span": max_reuse_span, + "candidate_edge_count": len(expected_edges), + "candidate_cell_count": len(expected_edges) * global_num_heads * global_num_heads, + "attention_call_counts": common_counts, + "tp_head_order_evidence": tp_head_order_evidence, + "dense_shadow_evidence": common_shadow, + "topology_discovery_capture_sha256": canonical_json_sha256(capture), + } + return MergedCapture(capture, manifest) 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/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_capture.py b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py new file mode 100644 index 00000000000..32009f076d5 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py @@ -0,0 +1,400 @@ +# 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, + merge_rank_topology_discovery_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 _topology_rank_payload(invocation, rank): + payload = _rank_payload(invocation, rank) + start = rank * 2 + payload.pop("consumer_layers") + payload.update( + { + "capture_mode": "topology_discovery", + "attention_layers": [0, 1, 2, 3], + "max_reuse_span": 2, + "anchor_stats_by_layer": { + str(layer): { + "retained_tiles": [10 + layer] * 4, + "dropped_mass": [0.01 * (layer + 1)] * 4, + } + for layer in range(4) + }, + "consumer_candidates_by_layer": { + str(consumer): { + str(anchor): { + "consumer_head_start": start, + "dropped_mass": [ + [ + 0.01 * (anchor + consumer + start + local + donor + 1) + for donor in range(4) + ] + for local in range(2) + ], + } + for anchor in range(max(0, consumer - 2), consumer) + } + for consumer in range(1, 4) + }, + "dense_shadow_evidence": { + "enabled": True, + "atol_hex": (0.0).hex(), + "rtol_hex": (0.0).hex(), + "validated_layer_indices": [0, 1, 2, 3], + }, + } + ) + return payload + + +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"} + + +def test_merge_topology_discovery_captures_concatenates_every_candidate_edge(): + invocation = _invocation() + merged = merge_rank_topology_discovery_captures( + [ + _topology_rank_payload(invocation, 1), + _topology_rank_payload(invocation, 0), + ], + invocation, + ) + + assert merged.capture["attention_layers"] == [0, 1, 2, 3] + assert merged.capture["max_reuse_span"] == 2 + assert set(merged.capture["anchor_stats_by_layer"]) == {"0", "1", "2", "3"} + assert set(merged.capture["consumer_candidates_by_layer"]["3"]) == {"1", "2"} + assert merged.capture["consumer_candidates_by_layer"]["3"]["1"]["dropped_mass"][2][ + 3 + ] == pytest.approx(0.1) + assert merged.manifest["candidate_edge_count"] == 5 + assert merged.manifest["candidate_cell_count"] == 5 * 4 * 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_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 index 7ff174952b3..68d7b33b0b6 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py @@ -166,13 +166,26 @@ def test_logspace_fit_preserves_log_a(self): class TestBuildSparseAttentionConfig: - _PARAMS = {"prefill": {"a": 7.9, "b": 8.6}, "decode": {"a": 0.12, "b": 9.8}} + _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} + 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" @@ -200,7 +213,7 @@ def test_replaced_skip_group_keeps_layer_policy(self): group = config["config_groups"]["group_0"] assert group["ignore"] == ["model.layers.0.self_attn"] assert group["initial_disabled_steps"] == 4 - assert group["threshold_scale_factor"]["prefill"] == {"a": 7.9, "b": 8.6} + assert group["threshold_scale_factor"]["prefill"] == self._PARAMS["prefill"] def test_preserves_nm_groups_and_replaces_old_skip_group(self): existing = { 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",)]