calibration

Calibration framework for sparse attention methods.

Classes

`AnchorLayerStats`	Per-head BLASST mask statistics for one topology anchor layer.
`CompactMaskReuseCapture`	One prompt/target capture decoded without expanding candidate rows.
`CompactMaskReuseCaptureSource`	Re-iterable strict JSONL source used by the three-pass selector.
`DynamicThresholdCalibrator`	Dynamic threshold calibrator using Exponential model.
`MaskReuseObservation`	One prompt, target-sparsity, consumer-head, and donor-head observation.
`RulerDatasetBuilder`	Builder for RULER calibration datasets.
`StableFileSnapshot`	Exact bytes and SHA256 read from one stable no-follow descriptor.
`VerifiedCheckpointManifest`	Identity of a checkpoint whose complete file set was SHA256-verified.

Functions

`calibrate_compact_mask_reuse_policy`	Select and evaluate a fail-closed schema-v3 candidate from compact captures.
`calibrate_mask_reuse_policy`	Select and evaluate a fail-closed schema-v3 legacy candidate.
`calibrate_sparse_attention`	Calibrate sparse attention parameters for optimal sparsity.
`canonical_prefill_threshold_scale_factor`	Canonicalize ModelOpt fit parameters or exported skip-softmax metadata.
`create_checkpoint_manifest`	Create the deterministic checkpoint manifest without replacing any file.
`load_compact_mask_reuse_captures`	Return a lazy, re-iterable compact-capture source.
`load_mask_reuse_observations`	Load normalized mask-reuse observations from JSONL.
`parse_mask_reuse_observations`	Parse strict normalized observation JSONL.
`read_stable_file_snapshot`	Read and hash identical bytes from one stable regular file.
`stable_file_sha256`	Hash one stable regular file without retaining its contents.
`verify_checkpoint_manifest`	Verify the fixed manifest under `checkpoint` and every declared file.

class AnchorLayerStats

Bases: object

Per-head BLASST mask statistics for one topology anchor layer.

__init__(retained_tiles, dropped_mass)

Parameters:

retained_tiles (tuple[int, ...])
dropped_mass (tuple[float, ...])

Return type:

None

dropped_mass: tuple[float, ...]

retained_tiles: tuple[int, ...]

exception CheckpointManifestError

Bases: ValueError

Raised when a checkpoint cannot be bound to its exact file contents.

class CompactMaskReuseCapture

Bases: object

One prompt/target capture decoded without expanding candidate rows.

__init__(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, sample_length, threshold_log2, threshold_lambda, geometry, global_num_heads, eligible_tiles, anchor_stats_by_layer, consumer_layers)

Parameters:

model (str)
checkpoint_manifest_sha256 (str)
split (str)
partition (str)
inner_fold (int | None)
prompt_id (str)
source (str)
source_group_sha256 (str)
source_capture_sha256 (str)
min_kv_tokens (int)
max_kv_tokens (int | None)
target_sparsity (float)
sample_length (int)
threshold_log2 (float)
threshold_lambda (float)
geometry (Mapping[str, int])
global_num_heads (int)
eligible_tiles (int)
anchor_stats_by_layer (Mapping[int, AnchorLayerStats])
consumer_layers (Mapping[int, CompactConsumerStats])

Return type:

None

anchor_stats_by_layer: Mapping[int, AnchorLayerStats]

property bucket: tuple[int, int | None]: Return this capture’s context bounds.

checkpoint_manifest_sha256: str

consumer_layers: Mapping[int, CompactConsumerStats]

eligible_tiles: int

classmethod from_mapping(raw)

Parse one strict compact-capture object.

Parameters:: raw (Mapping[str, object])
Return type:: CompactMaskReuseCapture

geometry: Mapping[str, int]

global_num_heads: int

inner_fold: int | None

max_kv_tokens: int | None

min_kv_tokens: int

model: str

partition: str

prompt_id: str

sample_length: int

source: str

source_capture_sha256: str

source_group_sha256: str

split: str

target_sparsity: float

threshold_lambda: float

threshold_log2: float

class CompactMaskReuseCaptureSource

Bases: object

Re-iterable strict JSONL source used by the three-pass selector.

__init__(path)

Parameters:: path (Path)
Return type:: None

path: Path

sha256()

Hash exact file bytes without loading the capture bundle.

Return type:: str

class DynamicThresholdCalibrator

Bases: object

Dynamic threshold calibrator using Exponential model.

Calibration Algorithm:

For each threshold λ_j in threshold_trials: - Run ALL samples through forward_loop - For each sample i with length L_i, collect sparsity S_ij - Compute scale_factor_ij = λ_j × L_i
Fit Exponential model to ALL individual (sf_ij, S_ij) pairs: scale_factor = a * exp(b * sparsity)
Return fitted a and b parameters

At inference time (user specifies target_sparsity S*):

scale_factor = a * exp(b * S*) threshold = scale_factor / seqlen

Key insight: Using all individual data points (N_thresholds × N_samples) instead of per-threshold averages provides more accurate fitting without additional calibration time cost.

__init__(threshold_trials=None, fit_logspace=False)

Initialize dynamic threshold calibrator.

Parameters:

threshold_trials (list[float] | None) – List of thresholds to try during calibration. Should span a range that achieves sparsities from ~10% to ~95%.
fit_logspace (bool) – If True, fit the exponential model in log space (minimizes relative error). Recommended for diffusion models where scale_factors span many orders of magnitude.

calibrate(model, forward_loop, phase)

Calibrate a and b parameters for Exponential model.

Algorithm:

Set thresholds = threshold_trials on all modules, run ONE forward pass. Each module returns a sparsity list (one entry per threshold) per sample. Unpack to get (scale_factor_ij = λ_j × L_i, sparsity_ij) pairs.
Fit Exponential model to ALL (sf_ij, S_ij) pairs: scale_factor = a * exp(b * sparsity)
Return fitted a and b parameters

At inference time (user specifies target_sparsity S*):

scale_factor = a * exp(b * S*) threshold = scale_factor / seqlen

Parameters:

model (Module) – The model with sparse attention modules
forward_loop (Callable) – Callable that takes model and forwards calibration data
phase (str) – Phase to calibrate (‘prefill’ or ‘decode’)

Returns:

Dict with calibration results including a, b, r_squared, and num_data_points

Return type:

dict[str, Any]

calibrate_from_stats(per_sample_stats, phase)

Fit the exponential model from already-collected per-sample stats.

This is the backend-agnostic Stage 2/3 of calibrate(). The HF and diffusion paths reach it through calibrate() (which runs a forward_loop to collect the stats first); the vLLM path collects the stats itself — one record per scheduled request — and calls this directly so both paths share the same exponential fit.

Parameters:

per_sample_stats (list[dict]) – List of {"sparsity": [s_0, ..., s_n], "sample_length": L} records, one per calibration sample. sparsity holds the skipped-tile fraction at each threshold in threshold_trials (same order, same length).
phase (str) – Phase being calibrated (‘prefill’ or ‘decode’).

Returns:

Dict with calibration results including a, b, r_squared, and num_data_points.

Return type:

dict[str, Any]

exception MaskReuseCalibrationError

Bases: ValueError

Raised when observations cannot produce a trustworthy reuse policy.

class MaskReuseObservation

Bases: object

One prompt, target-sparsity, consumer-head, and donor-head observation.

__init__(model, min_kv_tokens, max_kv_tokens, target_sparsity, sample_length, threshold_lambda, threshold_log2, q_tokens, kv_tokens, q_start_tokens, split, prompt_id, source_capture_sha256, anchor_layer, consumer_layer, consumer_head, donor_head, retained_tiles, eligible_tiles, anchor_dropped_mass, anchor_stats_by_layer, dropped_mass)

Parameters:

model (str)
min_kv_tokens (int)
max_kv_tokens (int | None)
target_sparsity (float)
sample_length (int)
threshold_lambda (float)
threshold_log2 (float)
q_tokens (int)
kv_tokens (int)
q_start_tokens (int)
split (str)
prompt_id (str)
source_capture_sha256 (str)
anchor_layer (int)
consumer_layer (int)
consumer_head (int)
donor_head (int)
retained_tiles (int)
eligible_tiles (int)
anchor_dropped_mass (float)
anchor_stats_by_layer (Mapping[int, AnchorLayerStats])
dropped_mass (float)

Return type:

None

anchor_dropped_mass: float

anchor_layer: int

anchor_stats_by_layer: Mapping[int, AnchorLayerStats]

consumer_head: int

consumer_layer: int

donor_head: int

dropped_mass: float

eligible_tiles: int

classmethod from_mapping(raw)

Build a validated observation from normalized JSON.

Parameters:: raw (Mapping[str, object])
Return type:: MaskReuseObservation

kv_tokens: int

max_kv_tokens: int | None

min_kv_tokens: int

model: str

prompt_id: str

q_start_tokens: int

q_tokens: int

retained_tiles: int

sample_length: int

source_capture_sha256: str

split: str

target_sparsity: float

threshold_lambda: float

threshold_log2: float

to_mapping()

Return the normalized JSON representation.

Return type:: dict[str, object]

class RulerDatasetBuilder

Bases: object

Builder for RULER calibration datasets.

__init__(samples, max_seqlen, tokenizer_name_or_path, num_length_bins=4, max_length_filter=65536, seed=42, cache_dir=None, data_dir=None)

Initialize RULER dataset builder.

Parameters:

samples (int) – Total number of samples to generate (distributed evenly across length bins)
max_seqlen (int) – Maximum sequence length (length bins auto-generated as powers of 2)
tokenizer_name_or_path (str | object) – HuggingFace tokenizer path or tokenizer object
seed (int) – Random seed for reproducibility
num_length_bins (int) – Number of length bins to generate (default: 4)
max_length_filter (int) – Maximum sequence length to keep (default: 65536)
cache_dir (str | None) – Optional cache directory. If None, uses ~/.cache/modelopt/data/
data_dir (str | Path | None) – Optional path to RULER data directory (contains ‘essays’ subdir). Required for NIAH tasks with essay haystack when not using pip default layout.

Note

Length bins are auto-generated as descending powers of 2: [max_seqlen, max_seqlen/2, max_seqlen/4, …] Generation stops when num_length_bins is reached or length < 1024. Subtasks are set to all the difficult tasks defined in RULER_TASKS.

build_calibration_dataset()

Build the complete calibration dataset.

If cache_dir was set, checks cache first and returns cached data if present. Otherwise generates the dataset, saves to cache (if cache_dir set), and returns.

Returns:: List of calibration samples with ‘input’ and ‘length’ fields
Return type:: list[dict[str, Any]]

class StableFileSnapshot

Bases: object

Exact bytes and SHA256 read from one stable no-follow descriptor.

__init__(path, payload, sha256)

Parameters:

path (Path)
payload (bytes)
sha256 (str)

Return type:

None

path: Path

payload: bytes

sha256: str

class VerifiedCheckpointManifest

Bases: object

Identity of a checkpoint whose complete file set was SHA256-verified.

__init__(checkpoint_root, manifest_path, model, sha256, file_count, total_size_bytes)

Parameters:

checkpoint_root (Path)
manifest_path (Path)
model (str)
sha256 (str)
file_count (int)
total_size_bytes (int)

Return type:

None

checkpoint_root: Path

file_count: int

manifest_path: Path

model: str

sha256: str

total_size_bytes: int

calibrate_compact_mask_reuse_policy(captures, *, vanilla_calibration, topology, checkpoint_manifest, evidence, max_anchor_dropped_mass, max_reuse_dropped_mass, max_reuse_selection_dropped_mass=None, source_provenance=None)

Select and evaluate a fail-closed schema-v3 candidate from compact captures.

Parameters:

captures (CompactMaskReuseCaptureSource | str | Path)
vanilla_calibration (Mapping[str, object])
topology (Mapping[str, object])
checkpoint_manifest (VerifiedCheckpointManifest)
evidence (Mapping[str, object])
max_anchor_dropped_mass (float)
max_reuse_dropped_mass (float)
max_reuse_selection_dropped_mass (float | None)
source_provenance (Mapping[str, object] | None)

Return type:

dict[str, object]

calibrate_mask_reuse_policy(observations, *, vanilla_calibration, topology, checkpoint_manifest, evidence, max_anchor_dropped_mass, max_reuse_dropped_mass, max_reuse_selection_dropped_mass=None, source_provenance=None)

Select and evaluate a fail-closed schema-v3 legacy candidate.

Parameters:

observations (Sequence[MaskReuseObservation | Mapping[str, object]])
vanilla_calibration (Mapping[str, object])
topology (Mapping[str, object])
checkpoint_manifest (VerifiedCheckpointManifest)
evidence (Mapping[str, object])
max_anchor_dropped_mass (float)
max_reuse_dropped_mass (float)
max_reuse_selection_dropped_mass (float | None)
source_provenance (Mapping[str, object] | None)

Return type:

dict[str, object]

calibrate_sparse_attention(model, config, forward_loop=None)

Calibrate sparse attention parameters for optimal sparsity.

Supports both prefill and decode phase calibration with per-phase target sparsity.

Parameters:

model (Module) – Model with sparse attention modules
config (dict[str, Any]) – Sparse attention configuration dict
forward_loop (Callable | None) – Callable that forwards calibration data through model. If None, auto-generates RULER dataset. Only used for prefill.

Returns:

Dictionary with calibration results for each phase

Return type:

dict[str, Any]

canonical_prefill_threshold_scale_factor(vanilla_calibration)

Canonicalize ModelOpt fit parameters or exported skip-softmax metadata.

Parameters:: vanilla_calibration (Mapping[str, object])
Return type:: dict[str, object]

create_checkpoint_manifest(checkpoint, *, model)

Create the deterministic checkpoint manifest without replacing any file.

Parameters:

checkpoint (str | Path)
model (str)

Return type:

VerifiedCheckpointManifest

load_compact_mask_reuse_captures(path)

Return a lazy, re-iterable compact-capture source.

Parameters:: path (str | Path)
Return type:: CompactMaskReuseCaptureSource

load_mask_reuse_observations(path)

Load normalized mask-reuse observations from JSONL.

Parameters:: path (str | Path)
Return type:: list[MaskReuseObservation]

parse_mask_reuse_observations(lines)

Parse strict normalized observation JSONL.

Parameters:: lines (Iterable[str])
Return type:: list[MaskReuseObservation]

read_stable_file_snapshot(path, *, label)

Read and hash identical bytes from one stable regular file.

Parameters:

path (str | Path)
label (str)

Return type:

StableFileSnapshot

stable_file_sha256(path, *, label)

Hash one stable regular file without retaining its contents.

Parameters:

path (str | Path)
label (str)

Return type:

str

verify_checkpoint_manifest(checkpoint, *, expected_model=None)

Verify the fixed manifest under checkpoint and every declared file.

The manifest must enumerate every regular file below the loaded checkpoint directory except itself. This prevents a manifest that binds only a subset of weights or remote-code/tokenizer inputs from naming the checkpoint.

Parameters:

checkpoint (str | Path)
expected_model (str | None)

Return type:

VerifiedCheckpointManifest