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.
`TopologyDiscoveryCapture`	One prompt/target all-earlier-layer topology-discovery capture.
`TopologyDiscoveryCaptureSource`	Re-iterable strict JSONL source for topology discovery.
`VerifiedCheckpointManifest`	Identity of a checkpoint whose complete file set was SHA256-verified.

Functions

`calibrate_compact_mask_reuse_policy`	Select a minimum-risk schema-v3 candidate under a per-bucket BMM1 target.
`calibrate_mask_reuse_policy`	Select a minimum-risk legacy candidate under a per-bucket BMM1 target.
`calibrate_mask_reuse_topology`	Select a global candidate topology and evaluate it without held-out tuning.
`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.
`load_topology_discovery_captures`	Return a lazy topology-discovery capture source.
`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 TopologyDiscoveryCapture

Bases: object

One prompt/target all-earlier-layer topology-discovery capture.

__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, attention_layers, max_reuse_span, anchor_stats_by_layer, consumer_candidates_by_layer)

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)
attention_layers (tuple[int, ...])
max_reuse_span (int)
anchor_stats_by_layer (Mapping[int, AnchorLayerStats])
consumer_candidates_by_layer (Mapping[int, Mapping[int, tuple[tuple[float, ...], ...]]])

Return type:

None

anchor_stats_by_layer: Mapping[int, AnchorLayerStats]

attention_layers: tuple[int, ...]

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

checkpoint_manifest_sha256: str

consumer_candidates_by_layer: Mapping[int, Mapping[int, tuple[tuple[float, ...], ...]]]

eligible_tiles: int

classmethod from_mapping(raw)

Parse one strict topology-discovery capture object.

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

geometry: Mapping[str, int]

global_num_heads: int

inner_fold: int | None

max_kv_tokens: int | None

max_reuse_span: int

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 TopologyDiscoveryCaptureSource

Bases: object

Re-iterable strict JSONL source for topology discovery.

__init__(path)

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

path: Path

sha256()

Hash exact file bytes.

Return type:: 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, reuse_dropped_mass_report_threshold, target_bmm1_skip_ratio, source_provenance=None)

Select a minimum-risk schema-v3 candidate under a per-bucket BMM1 target.

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)
reuse_dropped_mass_report_threshold (float)
target_bmm1_skip_ratio (float)
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, reuse_dropped_mass_report_threshold, target_bmm1_skip_ratio, source_provenance=None)

Select a minimum-risk legacy candidate under a per-bucket BMM1 target.

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)
reuse_dropped_mass_report_threshold (float)
target_bmm1_skip_ratio (float)
source_provenance (Mapping[str, object] | None)

Return type:

dict[str, object]

calibrate_mask_reuse_topology(captures, *, vanilla_calibration, checkpoint_manifest, evidence, max_anchor_dropped_mass, reuse_dropped_mass_report_threshold, target_bmm1_skip_ratio)

Select a global candidate topology and evaluate it without held-out tuning.

target_bmm1_skip_ratio is the minimum fraction of eligible QK tiles skipped by non-fallback reuse consumers across all attention layers in the calibration split. Among policies that meet the target, selection minimizes development reuse dropped mass. If the target is structurally infeasible, the returned candidate maximizes BMM1 skips and explicitly reports the gap.

reuse_dropped_mass_report_threshold only counts diagnostic violations after selection. It never accepts, rejects, or retunes a policy.

Parameters:

captures (TopologyDiscoveryCaptureSource | str | Path)
vanilla_calibration (Mapping[str, object])
checkpoint_manifest (VerifiedCheckpointManifest)
evidence (Mapping[str, object])
max_anchor_dropped_mass (float)
reuse_dropped_mass_report_threshold (float)
target_bmm1_skip_ratio (float)

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]

load_topology_discovery_captures(path)

Return a lazy topology-discovery capture source.

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

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