From 1e58faa20252dc9a373fbd8a86bb4ede107db9ae Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 20:49:46 -0700 Subject: [PATCH 01/10] Fix skip-softmax kernel and serving contracts - Route active skip-softmax launches (prefill and paged decode) through the fixed 128x128 calibration tile instead of the autotuner, so realized sparsity matches the granularity thresholds were calibrated at (autotuned BLOCK_N=32 tiles skip differently than the 128x128 measurement/calibration blocks). BLOCK_M steps down on shared-memory pressure (fp32 inputs on ~100KB-smem GPUs) with a warning; BLOCK_N -- the calibrated KV skip granularity -- is never reduced. - flash_skip_softmax: exclude padded query rows from the block keep decision; fully-padded rows voted keep (block_diff == 0) and forced the last partial block row dense, under-counting sparsity. - vLLM runtime: validate the calibrated-decode CUDA-graph guard for sparse-only installs, not just quantized ones. Signed-off-by: Kai Xu --- .../kernels/common/attention/triton_fa.py | 53 +++++++++++++++---- .../methods/flash_skip_softmax.py | 10 ++++ .../plugins/vllm_runtime.py | 11 ++-- .../attention/test_triton_fa_skip_softmax.py | 7 ++- 4 files changed, 64 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 5acc9fb787c..4d69a13a4cb 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -23,6 +23,7 @@ """ import math +import warnings from typing import Any import torch @@ -1030,17 +1031,47 @@ def grid(META): # kernel dereferences the right pointers instead of triggering an # illegal memory access. with torch.cuda.device(q.device): - if do_measure: - # Runtime counters mutate global tensors, so do not run them through - # autotune candidate trials. Use one stable config for measurement. - _attn_fwd.fn[grid]( - *fwd_args, - **fwd_kwargs, - BLOCK_M=_P_QDQ_MEASURE_BLOCK_M if p_qdq_mode else _MEASURE_BLOCK_M, - BLOCK_N=_MEASURE_BLOCK_N, - num_warps=_MEASURE_NUM_WARPS, - num_stages=_MEASURE_NUM_STAGES, - ) + if do_measure or apply_skip: + # Fixed-tile launches, bypassing autotune: + # - Measurement: runtime counters mutate global tensors, so they + # must not run through autotune candidate trials. + # - Active skip-softmax: the tile-skip decision depends on the + # (BLOCK_M, BLOCK_N) geometry, and thresholds are calibrated at + # the 128x128 measurement granularity (attention_calibrate and + # flash_skip_softmax both use 128x128 blocks). Autotuned tiles + # (e.g. BLOCK_N=32) would realize a different sparsity than + # calibrated, so skip launches always use the calibration tile. + # + # If the 128-row tile exceeds the device's shared memory (fp32 + # inputs on ~100KB-smem GPUs), halve BLOCK_M and retry. BLOCK_N + # — the calibrated KV skip granularity — is never reduced. + # Smaller BLOCK_M splits the per-tile skip vote into narrower + # row groups: per-row numerics stay protected by the same + # threshold, but realized sparsity can exceed the calibrated + # target, so warn when stepping down. + block_m = _P_QDQ_MEASURE_BLOCK_M if p_qdq_mode else _MEASURE_BLOCK_M + while True: + try: + _attn_fwd.fn[grid]( + *fwd_args, + **fwd_kwargs, + BLOCK_M=block_m, + BLOCK_N=_MEASURE_BLOCK_N, + num_warps=_MEASURE_NUM_WARPS, + num_stages=_MEASURE_NUM_STAGES, + ) + break + except triton.runtime.errors.OutOfResources: + if block_m <= 16: + raise + block_m //= 2 + if apply_skip: + warnings.warn( + f"skip-softmax tile BLOCK_M reduced to {block_m} " + "(shared memory limit); realized sparsity may " + "exceed the calibrated target on this device/dtype", + stacklevel=2, + ) else: _attn_fwd[grid]( *fwd_args, diff --git a/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py b/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py index c1d6465ba66..baefbd7058d 100644 --- a/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py +++ b/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py @@ -193,6 +193,16 @@ def calc_correction_factor_and_p( dense_blocks_list = [] block_mask_0 = None block_diff = block_max - block_max_cummax + # Exclude padded query rows from the keep decision. _reshape_to_blocks + # pads the last block row with ``dtype.min``; a fully-padded row then + # has ``block_diff == 0`` (min - min), which passes ``> log_threshold`` + # and forces every block in the last partial block row to be kept + # (never skipped) — under-counting sparsity by up to one block row. + # Mask those rows to -inf so they vote "skip", matching the Triton + # kernel, which drops padding rows from its tile-skip reduction. + pad_q = padded_seq_q - seq_q + if pad_q > 0: + block_diff[:, :, -1, self.br - pad_q :, :] = float("-inf") for i, log_threshold in enumerate(log_thresholds): block_mask = (block_diff > log_threshold).any(dim=-2) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index b4141740b64..b5a578585d8 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -348,7 +348,7 @@ def _plan_vllm_attention( _require_supported_vllm() errors = _global_errors(model_runner) if quantize else [] - mode = _cudagraph_mode(model_runner) if quantize else None + mode = _cudagraph_mode(model_runner) quant_plugin: Any = _load_quant_plugin() if quantize else None plans = [] for name, module, sparse_kw in candidates: @@ -365,9 +365,12 @@ def _plan_vllm_attention( reasons.append(f"resolved dtype {dtype} must be fp16 or bf16") if capability_error := _device_capability_error(device): reasons.append(capability_error) - if quantize: - if graph_error := _sparse_graph_error(sparse_kw, mode): - reasons.append(graph_error) + # Calibrated decode skip-softmax replays through the decode kernel path, + # which a FULL decode CUDA graph would capture with a stale threshold. + # This holds for sparse-only installs exactly as for quantized ones, so + # the guard is not gated on ``quantize``. + if graph_error := _sparse_graph_error(sparse_kw, mode): + reasons.append(graph_error) new_impl, requires_flashinfer_patch, backend_error = _select_new_impl(module) if backend_error: reasons.append(backend_error) diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py index fc26c5db17c..caefb467555 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py @@ -271,5 +271,8 @@ def test_skip_softmax_via_sparsify(self, tiny_llama_dir): assert not torch.isnan(logits_skip).any(), "NaN in skip-softmax logits" assert not torch.isinf(logits_skip).any(), "Inf in skip-softmax logits" - # On short sequences (64 tokens), no tiles are skipped — output should match dense - torch.testing.assert_close(logits_skip, logits_dense, rtol=1e-3, atol=1e-3) + # On short sequences (64 tokens), no tiles are skipped — output should match + # dense up to bf16 accumulation-order noise: active skip launches run on the + # fixed 128x128 calibration tile, so their summation order differs from the + # HF dense reference (and from the autotuned dense Triton tile). + torch.testing.assert_close(logits_skip, logits_dense, rtol=1e-2, atol=8e-3) From 31721b03df85f54f42d925b580d9a917a40e2e64 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 20:56:22 -0700 Subject: [PATCH 02/10] Add paged vLLM skip-softmax calibration - attention_calibrate: read K/V through a paged cache (vLLM NHD layout) via block_table, reusing the shared paged tile loaders. Contiguous-KV behavior is unchanged; paged and contiguous launches produce identical counters and output. - vLLM adapters: calibration mode (enable_calibration / disable_calibration / iter_sparse_impls). When active, forward routes to the paged calibration kernel -- full dense attention plus multi-threshold tile-skip counting -- so generation is numerically unchanged while stats accumulate. Each scheduled request is measured independently (decode vs prefill phase per request), and raw per-threshold tile counts are recorded so tensor-parallel ranks can be aggregated by summing counts before the fit. - FlashInfer adapter writes the current K/V to the cache before the calibrate kernel reads it (releases that update the cache inside forward); fp16/bf16-only and NHD-only, both validated explicitly. Signed-off-by: Kai Xu --- .../kernels/sparsity/attention/calibrate.py | 132 +++++++++++-- .../attention_sparsity/plugins/vllm.py | 183 ++++++++++++++++++ 2 files changed, 301 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/kernels/sparsity/attention/calibrate.py b/modelopt/torch/kernels/sparsity/attention/calibrate.py index d26e781d48c..68748e7c50c 100644 --- a/modelopt/torch/kernels/sparsity/attention/calibrate.py +++ b/modelopt/torch/kernels/sparsity/attention/calibrate.py @@ -28,7 +28,12 @@ import triton import triton.language as tl -from modelopt.torch.kernels.common.attention.triton_fa import LOG2E, _apply_mask +from modelopt.torch.kernels.common.attention.triton_fa import ( + LOG2E, + _apply_mask, + _load_paged_k_tile, + _load_paged_v_tile, +) # --------------------------------------------------------------------------- @@ -64,6 +69,18 @@ def _attn_fwd_calibrate( HEAD_DIM: tl.constexpr, NUM_THRESHOLDS: tl.constexpr, PADDED_THRESHOLDS: tl.constexpr, # next_power_of_2(NUM_THRESHOLDS) for tl.arange + IS_PAGED: tl.constexpr = False, # Whether K/V are read from a paged KV cache + K_cache=None, # [num_blocks, page_size, num_kv_heads, head_dim] paged K + V_cache=None, # [num_blocks, page_size, num_kv_heads, head_dim] paged V + Block_table=None, # [batch, max_blocks_per_seq] page table + stride_kc_block=0, + stride_kc_pos=0, + stride_kc_head=0, + stride_vc_block=0, + stride_vc_pos=0, + stride_vc_head=0, + PAGE_SIZE: tl.constexpr = 16, + max_blocks_per_seq=0, ): """Forward kernel with multi-threshold sparsity measurement. @@ -126,12 +143,33 @@ def _attn_fwd_calibrate( for kv_start in range(0, kv_bound, BLOCK_N): kv_start = tl.multiple_of(kv_start, BLOCK_N) - k_offs = (kv_offset + kv_start + kv_pos[None, :]) * stride_kbs + dim_pos[:, None] - k = tl.load( - k_base + k_offs, - mask=((kv_start + kv_pos[None, :]) < seq_len_kv) & d_mask[:, None], - other=0.0, - ) + # Load K^T [BLOCK_D, BLOCK_N] from paged cache or contiguous K. + if IS_PAGED: + k = _load_paged_k_tile( + K_cache, + Block_table, + batch_idx, + kv_head_idx, + kv_start, + kv_pos, + dim_pos, + seq_len_kv, + stride_kc_block, + stride_kc_pos, + stride_kc_head, + PAGE_SIZE, + BLOCK_N, + BLOCK_D, + HEAD_DIM, + max_blocks_per_seq, + ) + else: + k_offs = (kv_offset + kv_start + kv_pos[None, :]) * stride_kbs + dim_pos[:, None] + k = tl.load( + k_base + k_offs, + mask=((kv_start + kv_pos[None, :]) < seq_len_kv) & d_mask[:, None], + other=0.0, + ) scores = tl.dot(q, k) * qk_scale scores = _apply_mask(scores, q_pos, kv_pos, seq_len_q, seq_len_kv, kv_start, IS_CAUSAL) @@ -164,12 +202,32 @@ def _attn_fwd_calibrate( row_sum = row_sum * correction + l_new acc = acc * correction[:, None] - v_offs = (kv_offset + kv_start + kv_pos[:, None]) * stride_vbs + dim_pos[None, :] - v = tl.load( - v_base + v_offs, - mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], - other=0.0, - ) + if IS_PAGED: + v = _load_paged_v_tile( + V_cache, + Block_table, + batch_idx, + kv_head_idx, + kv_start, + kv_pos, + dim_pos, + seq_len_kv, + stride_vc_block, + stride_vc_pos, + stride_vc_head, + PAGE_SIZE, + BLOCK_N, + BLOCK_D, + HEAD_DIM, + max_blocks_per_seq, + ) + else: + v_offs = (kv_offset + kv_start + kv_pos[:, None]) * stride_vbs + dim_pos[None, :] + v = tl.load( + v_base + v_offs, + mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], + other=0.0, + ) acc = tl.dot(p.to(v.dtype), v, acc) row_max = m_new @@ -212,6 +270,10 @@ def attention_calibrate( max_input_len_k: int | None = None, *, threshold_trials: list[float] | None = None, + k_cache: torch.Tensor | None = None, + v_cache: torch.Tensor | None = None, + block_table: torch.Tensor | None = None, + page_size: int = 16, ) -> tuple[torch.Tensor, torch.Tensor]: """Flash attention with multi-threshold skip-softmax sparsity measurement. @@ -219,12 +281,21 @@ def attention_calibrate( measuring how many KV tiles would be skipped at each threshold in ``threshold_trials``. No autograd — forward only. - All arguments except ``threshold_trials`` match + All positional arguments match :func:`modelopt.torch.kernels.common.attention.attention`. Args: threshold_trials: List of threshold values to measure sparsity for. Each value is converted to log2-scaled space for the kernel. + k_cache: Paged K cache ``[num_blocks, page_size, num_kv_heads, head_dim]``. + When provided, K/V are read from the paged cache via ``block_table`` + (vLLM NHD layout) instead of from the contiguous ``k``/``v`` tensors. + ``k``/``v`` are then dummies whose only meaningful dimension is + ``shape[1] == num_kv_heads`` (used to compute the GQA ratio). + v_cache: Paged V cache ``[num_blocks, page_size, num_kv_heads, head_dim]``. + block_table: Page table ``[batch, max_blocks_per_seq]`` mapping each + sequence's block indices to global page IDs. + page_size: Number of tokens per page in the KV cache. Returns: Tuple of ``(output, sparsity_counters)``: @@ -237,6 +308,10 @@ def attention_calibrate( if threshold_trials is None or len(threshold_trials) == 0: raise ValueError("threshold_trials must be a non-empty list") + is_paged = k_cache is not None + if is_paged and block_table is None: + raise ValueError("block_table is required when k_cache/v_cache are provided.") + # Calibration has only been validated with uniform-length batches (current # diffusion + RULER paths). Varlen inputs would exercise code paths in the # kernel that have not been tested — fail loudly rather than silently @@ -281,6 +356,11 @@ def attention_calibrate( b_seq_len_k = b_seq_len b_start_loc_k = b_start_loc + # Paged mode: KV positions come from block_table, so the contiguous KV + # offsets are unused. Provide a dummy so Triton can compile the tl.load. + if b_start_loc_k is None: + b_start_loc_k = torch.zeros_like(b_start_loc) + num_thresholds = len(threshold_trials) # Scores already include sm_scale and LOG2E; convert lambda to log2 space only. @@ -304,6 +384,18 @@ def attention_calibrate( num_programs * num_thresholds, dtype=torch.int32, device=q.device ) + # Paged KV cache strides (zeros when not paged; computed here so the type + # narrowing of k_cache/v_cache/block_table is explicit for the kernel call). + if is_paged: + assert k_cache is not None and v_cache is not None and block_table is not None + kc_strides = (k_cache.stride(0), k_cache.stride(1), k_cache.stride(2)) + vc_strides = (v_cache.stride(0), v_cache.stride(1), v_cache.stride(2)) + max_blocks_per_seq = block_table.shape[1] + else: + kc_strides = (0, 0, 0) + vc_strides = (0, 0, 0) + max_blocks_per_seq = 0 + # Triton launches on torch.cuda.current_device(), which is not necessarily # the device the tensors live on (e.g. under accelerate device_map="auto" # sharding). Activate the tensor's device so the kernel dereferences the @@ -338,6 +430,18 @@ def attention_calibrate( HEAD_DIM=HEAD_DIM, NUM_THRESHOLDS=num_thresholds, PADDED_THRESHOLDS=triton.next_power_of_2(num_thresholds), + IS_PAGED=is_paged, + K_cache=k_cache, + V_cache=v_cache, + Block_table=block_table, + stride_kc_block=kc_strides[0], + stride_kc_pos=kc_strides[1], + stride_kc_head=kc_strides[2], + stride_vc_block=vc_strides[0], + stride_vc_pos=vc_strides[1], + stride_vc_head=vc_strides[2], + PAGE_SIZE=page_size, + max_blocks_per_seq=max_blocks_per_seq, num_warps=4, num_stages=1, ) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 243db16a2bc..a3a5ebe1271 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -45,6 +45,7 @@ ) from modelopt.torch.kernels.common.attention.triton_fa import attention as triton_attention from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite +from modelopt.torch.kernels.sparsity.attention.calibrate import attention_calibrate @functools.cache @@ -252,6 +253,104 @@ def _resolve_forward( ) +def _calibration_active(impl) -> bool: + """Return whether skip-softmax calibration mode is enabled on an impl.""" + return bool(getattr(impl, "_calibrate", False)) and bool( + getattr(impl, "_calib_threshold_trials", None) + ) + + +def _forward_calibrate( + impl, + *, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + num_actual_tokens: int, + output: torch.Tensor, +) -> torch.Tensor: + """Measure per-request tile-skip stats via the paged Triton calibration kernel. + + Each scheduled request is calibrated independently (batch=1) so its KV + length is the per-sample length the exponential fit needs, and so the + kernel keeps the uniform-length contract it was validated against. The + kernel computes full attention, so ``output`` is written densely and the + forward pass is numerically unchanged. + + Phase and causality are decided per request: ``q_len == 1`` is a decode + step (full-cache, non-causal); ``q_len > 1`` is (chunked) prefill (causal + — the kernel offsets the query into the KV span). A mixed prefill/decode + batch therefore contributes correctly to both phase fits. + + Records raw per-threshold tile counts (not ratios) on + ``impl._calib_records`` so tensor-parallel workers can be aggregated by + summing counts before the fit. + """ + if key_cache.dtype not in (torch.float16, torch.bfloat16): + raise NotImplementedError( + f"skip-softmax calibration requires an fp16/bf16 KV cache, got {key_cache.dtype}" + ) + if key_cache.shape[2] != impl.num_kv_heads: + # NHD is the only supported paged layout: [blocks, page, kv_heads, dim]. + # An HND FlashInfer cache would put kv_heads on axis 1. + raise NotImplementedError( + f"KV cache layout is not NHD (shape {tuple(key_cache.shape)}, " + f"expected axis 2 == num_kv_heads == {impl.num_kv_heads}); " + "HND caches are unsupported for calibration" + ) + page_size = key_cache.shape[1] + trials = impl._calib_threshold_trials + batch = seq_lens.shape[0] + b_start_loc = cu_seqlens_q[:batch] + b_seq_len = cu_seqlens_q[1 : batch + 1] - cu_seqlens_q[:batch] + + q = query[:num_actual_tokens].contiguous() + # Dummy K/V: in paged mode KV is read from the cache via block_table. + # Only shape[1] (num_kv_heads) is consulted, to compute the GQA ratio. + k_dummy = torch.empty(0, impl.num_kv_heads, impl.head_size, device=q.device, dtype=q.dtype) + + for i in range(batch): + q_len = int(b_seq_len[i].item()) + if q_len <= 0: + continue + q_start = int(b_start_loc[i].item()) + seq_k = int(seq_lens[i].item()) + phase = "decode" if q_len <= 1 else "prefill" + + oi, counters = attention_calibrate( + q[q_start : q_start + q_len], + k_dummy, + k_dummy, + b_start_loc=torch.zeros(1, device=q.device, dtype=torch.int32), + b_seq_len=b_seq_len[i : i + 1].to(torch.int32), + max_input_len=q_len, + is_causal=q_len > 1, + softmax_scale=impl.scale, + b_seq_len_k=seq_lens[i : i + 1].to(torch.int32), + max_input_len_k=seq_k, + threshold_trials=trials, + k_cache=key_cache, + v_cache=value_cache, + block_table=block_table[i : i + 1], + page_size=page_size, + ) + output[q_start : q_start + q_len] = oi + + impl._calib_records.append( + { + "phase": phase, + "sample_length": seq_k, + "total_tiles": counters[:, 0].tolist(), + "skipped_tiles": counters[:, 1].tolist(), + } + ) + + return output + + # Resolution guards raw configured transforms; dispatch rechecks effective # sparse work after calibration and decode-only pruning. def _forward_modelopt( @@ -517,6 +616,26 @@ def native_forward(): ) return native_result + if _calibration_active(self): + if getattr(attn_metadata, "use_cascade", False): + # Cascade splits shared prefixes across requests, so per-request + # KV lengths are unavailable; skip measurement for this launch. + return native_forward() + # vLLM >= 0.15 writes the current K/V to the paged cache before + # impl.forward, so the calibrate kernel reads a complete cache. + key_cache, value_cache = kv_cache.unbind(0) + return _forward_calibrate( + self, + query=query, + key_cache=key_cache, + value_cache=value_cache, + block_table=attn_metadata.block_table, + seq_lens=attn_metadata.seq_lens, + cu_seqlens_q=attn_metadata.query_start_loc, + num_actual_tokens=attn_metadata.num_actual_tokens, + output=output, + ) + resolved = _resolve_forward( self, layer, @@ -705,6 +824,36 @@ def prepare_modelopt(): _maybe_update_flashinfer_cache(layer, key, value, kv_cache, attn_metadata, impl) cache_prepared = True + if _calibration_active(impl): + if getattr(attn_metadata, "use_cascade", False): + # Cascade splits shared prefixes across requests, so per-request + # KV lengths are unavailable; skip measurement for this launch. + return dense_fallback() + missing = [name for name in _FLASHINFER_METADATA_FIELDS if not hasattr(attn_metadata, name)] + if missing: + raise NotImplementedError( + "FlashInfer metadata is missing the ModelOpt calibration " + f"fields: {', '.join(missing)}" + ) + if kv_cache.ndim != 5 or kv_cache.shape[1] != 2: + raise ValueError( + "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" + ) + # Order matters: releases that update the KV cache inside forward must + # write the current K/V before the calibrate kernel reads the cache. + prepare_modelopt() + return _forward_calibrate( + impl, + query=query, + key_cache=kv_cache[:, 0], + value_cache=kv_cache[:, 1], + block_table=attn_metadata._modelopt_block_table, + seq_lens=attn_metadata._modelopt_seq_lens, + cu_seqlens_q=attn_metadata._modelopt_query_start_loc, + num_actual_tokens=attn_metadata._modelopt_num_actual_tokens, + output=output, + ) + resolved = _resolve_forward( impl, layer, @@ -833,3 +982,37 @@ def _clone_sparse_impl(old_impl, new_cls=None): new_impl = object.__new__(new_cls) new_impl.__dict__.update(old_state) return new_impl + + +def iter_sparse_impls(model): + """Yield every ModelOpt sparse attention impl reachable from a vLLM model. + + Walks ``model.named_modules()`` and returns the swapped ``impl`` of each + attention layer (FlashAttention or FlashInfer adapter). Used by the + calibration installer and RPC methods to toggle calibration mode and + harvest stats without knowing vLLM's module layout. + """ + for _, module in model.named_modules(): + impl = getattr(module, "impl", None) + if impl is None: + continue + if isinstance(impl, ModelOptSparseAttentionImpl) or ( + _FLASHINFER_IMPL_CLS is not None and isinstance(impl, _FLASHINFER_IMPL_CLS) + ): + yield impl + + +def enable_calibration(impls, threshold_trials: list[float]) -> None: + """Put a set of sparse impls into calibration mode and clear prior records.""" + if not threshold_trials: + raise ValueError("threshold_trials must be a non-empty list for calibration.") + for impl in impls: + impl._calibrate = True + impl._calib_threshold_trials = list(threshold_trials) + impl._calib_records = [] + + +def disable_calibration(impls) -> None: + """Turn off calibration mode (collected records are left intact).""" + for impl in impls: + impl._calibrate = False From e3fde5f98aa56012cda1b2eaef4dfe8718080d18 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 21:05:21 -0700 Subject: [PATCH 03/10] Add skip-softmax calibration installer, TP-count aggregation, and example - install_vllm_skip_softmax_calibration: validation-before-mutation install of calibration adapters on every attention layer -- requires eager execution, fp16/bf16 model and KV dtypes, no active attention Q/K/P/V fakequant, and a FlashAttention/FlashInfer backend; disables cascade. Measurement starts separately (enable_calibration RPC), so warmup launches are never recorded. - DynamicThresholdCalibrator.calibrate_from_stats: backend-agnostic fit stage extracted from calibrate(), preserving result fields (fit_logspace, log_a) and reporting per-sample sparsity; the HF path delegates to it unchanged. - plugins/sparse_attn_calibration (vLLM-free): raw tile-count merging across layers and tensor-parallel ranks (counts are additive; ratios are formed only after the global merge, one fit per phase -- never per rank, never by averaging fitted coefficients), plus a canonical sparse_attention_config builder matching export_sparse_attention_config so the serving loader round-trips it; existing N:M groups are preserved. - Thin example: SkipSoftmaxCalibWorker (load hook + enable/status/counts RPCs) and calibrate_sparse_attn.py driver (CLI, prompts, rank-count aggregation, fit, checkpoint-config writing). Signed-off-by: Kai Xu --- examples/vllm_serve/calibrate_sparse_attn.py | 246 ++++++++++++++++++ examples/vllm_serve/sparse_attn_worker.py | 63 ++++- .../calibration/calibrator.py | 52 +++- .../plugins/sparse_attn_calibration.py | 236 +++++++++++++++++ .../attention_sparsity/plugins/vllm.py | 22 ++ .../plugins/vllm_runtime.py | 72 +++++ .../test_sparse_attn_worker.py | 6 +- 7 files changed, 693 insertions(+), 4 deletions(-) create mode 100644 examples/vllm_serve/calibrate_sparse_attn.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py new file mode 100644 index 00000000000..6c36f8f65ec --- /dev/null +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Calibrate skip-softmax thresholds *through vLLM* and write the serving config. + +Runs calibration prompts through a vLLM ``LLM`` whose attention layers carry +the ModelOpt calibration adapters (installed by +``sparse_attn_worker.SkipSoftmaxCalibWorker`` via +``install_vllm_skip_softmax_calibration``). The paged Triton calibration +kernel measures, per candidate threshold, how many KV tiles would be skipped — +over the paged KV cache, for both prefill and decode — then this driver +aggregates the raw counts from every tensor-parallel rank and fits the +exponential model ``scale_factor = a * exp(b * sparsity)`` once per phase. + +The fitted ``(a, b)`` are written as a canonical ``sparse_attention_config`` +block (the same schema ModelOpt's HF export produces), so the serving path +(``vllm_serve_sparse_attn.py`` / ``install_vllm_sparse_attention_from_checkpoint``) +loads it without changes. Any exported N:M sparse-softmax groups already in +the checkpoint config are preserved. + +Usage: + python calibrate_sparse_attn.py \ + --prompts_file prompts.txt \ + --target_sparse_ratio 0.5 \ + --decode_tokens 32 \ + --update_checkpoint_config + +``--prompts_file`` is one prompt per line; longer, varied-length prompts give a +better fit. With no file, a tiny built-in demo set is used (fine for a smoke +test, not for a real fit). +""" + +import argparse +import json +import os +import sys +from pathlib import Path + +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_calibration import ( + DEFAULT_THRESHOLD_TRIALS, + build_sparse_attention_config, + fit_from_counts, + merge_phase_counts, +) + +_DEMO_PROMPTS = [ + "Summarize the history of computing in a few paragraphs. " * 40, + "Explain how attention works in transformer models. " * 60, + "Write a detailed essay about renewable energy sources. " * 80, +] + + +def _load_prompts(prompts_file: str | None) -> list[str]: + if prompts_file is None: + print( + "[ModelOpt] No --prompts_file given; using a tiny built-in demo set. " + "Pass real, varied-length prompts for a usable fit." + ) + return _DEMO_PROMPTS + lines = [ln.strip() for ln in Path(prompts_file).read_text().splitlines() if ln.strip()] + if not lines: + raise ValueError(f"No prompts found in {prompts_file}") + print(f"[ModelOpt] Loaded {len(lines)} calibration prompts from {prompts_file}") + return lines + + +def _existing_sparse_config(ckpt: str) -> dict | None: + """Read the checkpoint's sparse_attention_config so non-skip groups survive.""" + config_json = Path(ckpt) / "config.json" + if not config_json.is_file(): + return None + existing = json.loads(config_json.read_text()).get("sparse_attention_config") + return existing if isinstance(existing, dict) else None + + +def _write_config(ckpt: str, sparse_config: dict, update_checkpoint: bool) -> None: + """Dump the sparse_attention_config and optionally merge into config.json.""" + out_path = Path("sparse_attention_config.json") + out_path.write_text(json.dumps(sparse_config, indent=2)) + print(f"[ModelOpt] Wrote calibrated config to {out_path.resolve()}") + + if not update_checkpoint: + print( + "[ModelOpt] Re-run with --update_checkpoint_config to merge this into " + f"{ckpt}/config.json (required for vllm_serve_sparse_attn.py to pick it up)." + ) + return + + config_json = Path(ckpt) / "config.json" + config = json.loads(config_json.read_text()) + config["sparse_attention_config"] = sparse_config + config_json.write_text(json.dumps(config, indent=2)) + print(f"[ModelOpt] Merged sparse_attention_config into {config_json}") + + +def main(): + parser = argparse.ArgumentParser(description="Calibrate skip-softmax thresholds via vLLM") + parser.add_argument("model", type=str, help="Path to the HF checkpoint to calibrate") + parser.add_argument("--prompts_file", type=str, default=None, help="One prompt per line") + parser.add_argument( + "--target_sparse_ratio", + type=float, + default=0.5, + help="Target sparsity baked into the exported config (applied to both phases)", + ) + parser.add_argument( + "--decode_tokens", + type=int, + default=32, + help="Decode tokens to generate per prompt (drives decode-phase calibration)", + ) + parser.add_argument( + "--max_model_len", type=int, default=None, help="vLLM max_model_len override" + ) + parser.add_argument( + "--tensor_parallel_size", type=int, default=1, help="vLLM tensor-parallel size" + ) + parser.add_argument( + "--gpu_memory_utilization", + type=float, + default=None, + help="vLLM GPU memory utilization fraction", + ) + parser.add_argument( + "--trust_remote_code", + action="store_true", + help="Trust remote code for custom model classes (e.g. NemotronH)", + ) + parser.add_argument("--dtype", type=str, default=None, help="Model dtype, e.g. bfloat16") + parser.add_argument( + "--attention_backend", + type=str, + default=None, + help="Force the vLLM attention backend, e.g. FLASH_ATTN or FLASHINFER. " + "Default: let vLLM choose (the installer supports whichever of FlashAttention " + "/ FlashInfer is selected).", + ) + parser.add_argument( + "--engine_kwargs", + type=str, + default=None, + help="JSON dict of extra vLLM engine kwargs, e.g. " + '\'{"enable_expert_parallel": true, "mamba_cache_mode": "align"}\' ' + "for hybrid MoE/Mamba models", + ) + parser.add_argument( + "--fit_logspace", + action="store_true", + help="Fit the exponential model in log space (wide scale_factor ranges)", + ) + parser.add_argument( + "--update_checkpoint_config", + action="store_true", + help="Merge the calibrated config into /config.json in place", + ) + args = parser.parse_args() + + # Workers run in separate processes and must import the calibration worker. + repo_root = str(Path(__file__).resolve().parent) + if repo_root not in sys.path: + sys.path.insert(0, repo_root) + current = os.environ.get("PYTHONPATH") + os.environ["PYTHONPATH"] = os.pathsep.join([current, repo_root]) if current else repo_root + + from vllm import LLM, SamplingParams + + prompts = _load_prompts(args.prompts_file) + + llm_kwargs = { + "model": args.model, + "worker_cls": "sparse_attn_worker.SkipSoftmaxCalibWorker", + # The calibration installer requires eager execution: the per-request + # calibration loop cannot be CUDA-graph captured. + "enforce_eager": True, + # Shared-prefix reuse would make prefill measurements cover only the + # non-cached suffix of each prompt; the installer rejects it. + "enable_prefix_caching": False, + } + if args.max_model_len is not None: + llm_kwargs["max_model_len"] = args.max_model_len + if args.tensor_parallel_size and args.tensor_parallel_size > 1: + llm_kwargs["tensor_parallel_size"] = args.tensor_parallel_size + if args.gpu_memory_utilization is not None: + llm_kwargs["gpu_memory_utilization"] = args.gpu_memory_utilization + if args.trust_remote_code: + llm_kwargs["trust_remote_code"] = True + if args.dtype is not None: + llm_kwargs["dtype"] = args.dtype + if args.attention_backend is not None: + llm_kwargs["attention_backend"] = args.attention_backend + if args.engine_kwargs: + extra = json.loads(args.engine_kwargs) + if not isinstance(extra, dict): + raise ValueError("--engine_kwargs must be a JSON object") + llm_kwargs.update(extra) + llm = LLM(**llm_kwargs) + + trials = list(DEFAULT_THRESHOLD_TRIALS) + n_layers = llm.collective_rpc("sparse_calib_enable", args=(trials,))[0] + status = llm.collective_rpc("sparse_calib_status")[0] + print(f"[ModelOpt] Calibration enabled on {n_layers} attention layers") + print(f"[ModelOpt] Active sparse impls: {status['impl_types']}") + + # generate() drives prefill (prefill-phase stats) then decode_tokens decode + # steps (decode-phase stats). The calibration kernel computes full attention, + # so the generated text is unaffected — only tile-skip counts are recorded. + sampling = SamplingParams(temperature=0.0, max_tokens=args.decode_tokens) + llm.generate(prompts, sampling) + + # Aggregate RAW counts from every TP rank (each rank only measures its + # attention-head shard), then fit once per phase on the global counts. + rank_counts = llm.collective_rpc("sparse_calib_counts") + merged = merge_phase_counts(rank_counts) + calibration_params = fit_from_counts(merged, trials, fit_logspace=args.fit_logspace) + + if not calibration_params: + print( + "[ModelOpt] Calibration produced no valid fit — try more/longer prompts " + "so observed sparsity spans the (10%, 90%) fitting window." + ) + return + + sparse_config = build_sparse_attention_config( + calibration_params, + {"prefill": args.target_sparse_ratio, "decode": args.target_sparse_ratio}, + existing_config=_existing_sparse_config(args.model), + ) + print("[ModelOpt] Calibrated threshold_scale_factor:") + print(json.dumps(sparse_config["config_groups"]["group_0"]["threshold_scale_factor"], indent=2)) + _write_config(args.model, sparse_config, args.update_checkpoint_config) + + +if __name__ == "__main__": + main() diff --git a/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 7b5fb28bf68..27e1eebdf97 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -17,12 +17,22 @@ from vllm.v1.worker.gpu_worker import Worker as BaseWorker +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_calibration import ( + DEFAULT_THRESHOLD_TRIALS, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( + collect_calibration_counts, + disable_calibration, + enable_calibration, + iter_sparse_impls, +) from modelopt.torch.sparsity.attention_sparsity.plugins.vllm_runtime import ( install_vllm_nvfp4_attention, + install_vllm_skip_softmax_calibration, install_vllm_sparse_attention_from_checkpoint, ) -__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker"] # noqa: RUF022 +__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker", "SkipSoftmaxCalibWorker"] # noqa: RUF022 _QUANT_FORMAT_KEYS = ("q_format", "k_format", "p_format", "v_format") @@ -69,6 +79,57 @@ def load_model(self, *args, **kwargs) -> None: _print_install_report("Sparse attention", report) +class SkipSoftmaxCalibWorker(BaseWorker): + """Calibrate skip-softmax thresholds through the engine. + + Unlike :class:`SparseAttnWorker` (which serves an already-calibrated + ``sparse_attention_config``), this worker *produces* that config. The + library installer swaps calibration-capable adapters onto every attention + layer at load; measurement starts only when the driver calls + ``sparse_calib_enable`` (so warmup launches are never recorded) and raw + per-threshold tile counts are harvested with ``sparse_calib_counts`` for + the driver to aggregate across TP ranks and fit. + """ + + def load_model(self, *args, **kwargs) -> None: + """Load the model, then install calibration adapters on every layer.""" + super().load_model(*args, **kwargs) + report = install_vllm_skip_softmax_calibration(self.model_runner) + print( + f"[ModelOpt] Skip-softmax calibration installed on {report.installed_count} " + f"attention layers: {dict(report.backend_counts)}" + ) + + # -- RPC methods (invoked via LLM.collective_rpc) ---------------------- + + def sparse_calib_enable(self, threshold_trials: list[float] | None = None) -> int: + """Enter calibration mode on all installed impls; returns layer count.""" + impls = list(iter_sparse_impls(_unwrapped_model(self))) + enable_calibration(impls, list(threshold_trials or DEFAULT_THRESHOLD_TRIALS)) + return len(impls) + + def sparse_calib_status(self) -> dict: + """Report active impls and record counts, so the backend is verifiable.""" + impls = list(iter_sparse_impls(_unwrapped_model(self))) + impl_types: dict[str, int] = {} + total_records = 0 + for impl in impls: + impl_types[type(impl).__name__] = impl_types.get(type(impl).__name__, 0) + 1 + total_records += len(getattr(impl, "_calib_records", [])) + return { + "num_sparse_layers": len(impls), + "impl_types": impl_types, + "calibrating": any(getattr(impl, "_calibrate", False) for impl in impls), + "total_records": total_records, + } + + def sparse_calib_counts(self) -> dict[str, list[dict]]: + """Stop measuring and return this rank's layer-merged raw tile counts.""" + model = _unwrapped_model(self) + disable_calibration(list(iter_sparse_impls(model))) + return collect_calibration_counts(model) + + class QuantSparseAttnWorker(BaseWorker): """Install quantized attention plus optional checkpoint sparsity. diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py index aded26fefdc..a37b447a686 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py @@ -130,8 +130,6 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic # with one entry per threshold, eliminating the need for repeated forward passes. print(f"\nStage 1: Collecting {phase} sparsity data for all thresholds in one pass...") - all_data_points = [] # List of {"threshold", "length", "scale_factor", "sparsity"} - self._set_thresholds(attention_modules, self.threshold_trials) self._enable_calibration_mode(attention_modules) with torch.no_grad(): @@ -139,6 +137,29 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic per_sample_stats = self._extract_calibration_stats(attention_modules, phase=phase) self._disable_calibration_mode(attention_modules) + return self.calibrate_from_stats(per_sample_stats, phase) + + def calibrate_from_stats(self, per_sample_stats: list[dict], phase: str) -> dict[str, Any]: + """Fit the exponential model from already-collected per-sample stats. + + This is the backend-agnostic Stage 2/3 of :meth:`calibrate`. The HF and + diffusion paths reach it through :meth:`calibrate` (which runs a + ``forward_loop`` to collect the stats first); the vLLM path collects the + stats itself — one record per scheduled request — and calls this directly + so both paths share the same exponential fit. + + Args: + per_sample_stats: List of ``{"sparsity": [s_0, ..., s_n], "sample_length": L}`` + records, one per calibration sample. ``sparsity`` holds the + skipped-tile fraction at each threshold in ``threshold_trials`` + (same order, same length). + phase: Phase being calibrated ('prefill' or 'decode'). + + Returns: + Dict with calibration results including a, b, r_squared, and num_data_points. + """ + all_data_points = [] # List of {"threshold", "length", "scale_factor", "sparsity"} + for sample_stat in per_sample_stats: length = sample_stat["sample_length"] sparsity_list = sample_stat["sparsity"] @@ -153,6 +174,12 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic } ) + # Per-sample measured sparsity (one row per calibration sample: its + # skipped-tile fraction at every threshold). Printed before the fit- + # validity guard so the raw per-sample data is visible even when the fit + # bails (e.g. degenerate near-zero sparsity). + self._print_per_sample_sparsity(per_sample_stats, phase) + if len(all_data_points) < 10: warnings.warn( f"Not enough data points for {phase} calibration. " @@ -286,11 +313,32 @@ def exponential(sparsity, a, b): "fit_logspace": self.fit_logspace, "min_observed_sparsity": min_observed_sparsity, "max_observed_sparsity": max_observed_sparsity, + # Raw per-sample measured sparsity, so callers can audit the spread + # across samples (not just the fitted average). + "per_sample_sparsity": [ + { + "sample_length": s.get("sample_length", 0), + "sparsity": list(s.get("sparsity", [])), + } + for s in per_sample_stats + ], } if self.fit_logspace: result["log_a"] = float(log_a) return result + def _print_per_sample_sparsity(self, per_sample_stats: list[dict], phase: str) -> None: + """Print each sample's measured skipped-tile fraction at every threshold.""" + if not per_sample_stats: + return + print(f"\nPer-sample {phase} sparsity (skipped-tile fraction per threshold):") + header = " ".join(f"{t:>7.0e}" for t in self.threshold_trials) + print(f" {'sample':>6} {'length':>8} {header}") + for idx, stat in enumerate(per_sample_stats): + sparsity = stat.get("sparsity", []) + row = " ".join(f"{s:>7.2%}" for s in sparsity) + print(f" {idx:>6} {stat.get('sample_length', 0):>8} {row}") + def _enable_calibration_mode(self, modules: list[nn.Module]): """Enable calibration mode on sparse attention modules.""" for idx, module in enumerate(modules): diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py new file mode 100644 index 00000000000..f60bef1b867 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""vLLM-free helpers for skip-softmax calibration through a serving engine. + +The serving adapters (``plugins/vllm.py``) record **raw per-threshold tile +counts** per scheduled request. These helpers merge those counts — across the +layers of one rank and across tensor-parallel ranks — then fit the exponential +threshold model and build the canonical ``sparse_attention_config`` block. + +Counts are additive, so aggregation is a plain sum: every layer of a rank and +every TP rank observes the same launches in the same order (TP ranks each see +their head shard), which makes records align by index within a phase. Fitting +happens once, per phase, on the globally merged counts — never per rank +(head-sharded counts are incomplete) and never by averaging independently +fitted coefficients (the fit is nonlinear). + +Everything here operates on plain Python data and is unit-testable without +vLLM installed. +""" + +from typing import Any + +import modelopt + +__all__ = [ + "DEFAULT_THRESHOLD_TRIALS", + "build_sparse_attention_config", + "fit_from_counts", + "merge_count_records", + "merge_phase_counts", + "split_records_by_phase", + "stats_from_counts", +] + +# Default threshold sweep — should span sparsities from ~10% to ~95%. +DEFAULT_THRESHOLD_TRIALS = [ + 1e-4, + 1e-3, + 5e-3, + 1e-2, + 3e-2, + 5e-2, + 1e-1, + 2e-1, + 3e-1, + 5e-1, + 7e-1, + 9e-1, +] + +_PHASES = ("prefill", "decode") + + +def split_records_by_phase(records: list[dict]) -> dict[str, list[dict]]: + """Group one impl's ordered calibration records by phase, preserving order.""" + per_phase: dict[str, list[dict]] = {phase: [] for phase in _PHASES} + for record in records: + per_phase.setdefault(record["phase"], []).append(record) + return per_phase + + +def merge_count_records(sources: list[list[dict]]) -> list[dict]: + """Element-wise sum aligned raw-count records from multiple sources. + + ``sources`` is a list over sources — the layers of one rank, or the + already layer-merged records of each TP rank — where each source is an + ordered list of ``{"sample_length", "total_tiles", "skipped_tiles"}`` + records for one phase. All sources observe the same launches in the same + order, so records align by index; tile counts are additive across both + layers and head-sharded TP ranks. + """ + sources = [source for source in sources if source] + if not sources: + return [] + num_samples = min(len(source) for source in sources) + merged = [] + for i in range(num_samples): + base = sources[0][i] + total = [0] * len(base["total_tiles"]) + skipped = [0] * len(base["skipped_tiles"]) + for source in sources: + record = source[i] + if record["sample_length"] != base["sample_length"]: + raise ValueError( + "Misaligned calibration records: sample lengths differ across " + f"sources at index {i} ({record['sample_length']} vs " + f"{base['sample_length']})" + ) + total = [a + b for a, b in zip(total, record["total_tiles"])] + skipped = [a + b for a, b in zip(skipped, record["skipped_tiles"])] + merged.append( + { + "sample_length": base["sample_length"], + "total_tiles": total, + "skipped_tiles": skipped, + } + ) + return merged + + +def merge_phase_counts(rank_counts: list[dict[str, list[dict]]]) -> dict[str, list[dict]]: + """Merge per-phase raw-count records collected from every TP rank. + + ``rank_counts`` is the list of per-rank results (one + ``{"prefill": [...], "decode": [...]}`` dict per rank, as returned by + ``collect_calibration_counts``). Use ALL ranks: with tensor parallelism + each rank only measures its attention-head shard, so any single rank's + counts are incomplete. + """ + phases = {phase for rank in rank_counts for phase in rank} + return { + phase: merge_count_records([rank.get(phase, []) for rank in rank_counts]) + for phase in phases + } + + +def stats_from_counts(count_records: list[dict]) -> list[dict]: + """Convert merged raw-count records into per-sample sparsity-ratio stats. + + Returns ``{"sample_length", "sparsity"}`` records in the shape + :meth:`DynamicThresholdCalibrator.calibrate_from_stats` consumes. + """ + stats = [] + for record in count_records: + sparsity = [ + (skipped / total if total else 0.0) + for skipped, total in zip(record["skipped_tiles"], record["total_tiles"]) + ] + stats.append({"sample_length": record["sample_length"], "sparsity": sparsity}) + return stats + + +def fit_from_counts( + per_phase_counts: dict[str, list[dict]], + threshold_trials: list[float], + *, + fit_logspace: bool = False, +) -> dict[str, dict[str, float]]: + """Fit the exponential skip-softmax model from globally merged counts. + + Reuses :class:`DynamicThresholdCalibrator` so vLLM-calibrated ``(a, b)`` + are identical in form to the HF path and export unchanged via + ``threshold_scale_factor``. One fit per phase, on counts already merged + across all TP ranks and layers. + + Returns: + ``{phase: {"a", "b", "min_observed_sparsity", "max_observed_sparsity"}}`` + for each phase that produced a valid fit. + """ + from ..calibration.calibrator import DynamicThresholdCalibrator + + calibration_params: dict[str, dict[str, float]] = {} + for phase, records in per_phase_counts.items(): + if not records: + continue + calibrator = DynamicThresholdCalibrator( + threshold_trials=list(threshold_trials), fit_logspace=fit_logspace + ) + result = calibrator.calibrate_from_stats(stats_from_counts(records), phase=phase) + if "a" in result and "b" in result: + params = {"a": result["a"], "b": result["b"]} + for key in ("min_observed_sparsity", "max_observed_sparsity"): + if key in result: + params[key] = result[key] + calibration_params[phase] = params + return calibration_params + + +def _normalize_target_sparsity(target_sparsity: dict[str, float] | float) -> dict[str, float]: + if isinstance(target_sparsity, (int, float)): + return {phase: float(target_sparsity) for phase in _PHASES} + return {phase: float(target_sparsity.get(phase, 0.5)) for phase in _PHASES} + + +def build_sparse_attention_config( + calibration_params: dict[str, dict[str, float]], + target_sparsity: dict[str, float] | float = 0.5, + *, + existing_config: dict | None = None, +) -> dict[str, Any]: + """Build the canonical ``sparse_attention_config`` block for a checkpoint. + + Emits the same schema as + ``modelopt.torch.sparsity.attention_sparsity.conversion.export_sparse_attention_config`` + — a ``config_groups`` entry with ``algorithm: skip_softmax`` holding the + group-local ``threshold_scale_factor`` and ``target_sparsity`` — so + ``load_from_checkpoint_metadata`` (the serving loader) round-trips it + without changes. + + Non-skip groups from ``existing_config`` (e.g. exported N:M + ``sparse_softmax`` metadata) are preserved after the skip group; an + existing ``skip_softmax`` group is replaced by the new calibration. + """ + threshold_scale_factor: dict[str, Any] = {"formula": "a * exp(b * target_sparsity)"} + for phase in _PHASES: + if phase in calibration_params: + threshold_scale_factor[phase] = { + "a": float(calibration_params[phase]["a"]), + "b": float(calibration_params[phase]["b"]), + } + + skip_group: dict[str, Any] = { + "algorithm": "skip_softmax", + "targets": ["Attention"], + "threshold_scale_factor": threshold_scale_factor, + "target_sparsity": _normalize_target_sparsity(target_sparsity), + } + + config_groups: dict[str, Any] = {"group_0": skip_group} + existing_groups = (existing_config or {}).get("config_groups") + if isinstance(existing_groups, dict): + preserved = [ + group + for group in existing_groups.values() + if isinstance(group, dict) and group.get("algorithm") != "skip_softmax" + ] + for idx, group in enumerate(preserved, start=1): + config_groups[f"group_{idx}"] = group + + return { + "config_groups": config_groups, + "producer": {"name": "modelopt", "version": modelopt.__version__}, + } diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index a3a5ebe1271..80e30e5fc18 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -1016,3 +1016,25 @@ def disable_calibration(impls) -> None: """Turn off calibration mode (collected records are left intact).""" for impl in impls: impl._calibrate = False + + +def collect_calibration_counts(model) -> dict[str, list[dict]]: + """Harvest one rank's raw per-phase tile counts from every calibrating impl. + + Sums counts across the rank's layers per aligned sample (every layer sees + the same launches in the same order), keeping raw + ``{"sample_length", "total_tiles", "skipped_tiles"}`` records per phase. + The driver merges these across TP ranks with + :func:`~.sparse_attn_calibration.merge_phase_counts` and fits once per + phase with :func:`~.sparse_attn_calibration.fit_from_counts` — sparsity + ratios are only formed after the global merge. + """ + from .sparse_attn_calibration import merge_count_records, split_records_by_phase + + per_phase_layers: dict[str, list[list[dict]]] = {} + for impl in iter_sparse_impls(model): + split = split_records_by_phase(getattr(impl, "_calib_records", [])) + for phase, records in split.items(): + if records: + per_phase_layers.setdefault(phase, []).append(records) + return {phase: merge_count_records(layers) for phase, layers in per_phase_layers.items()} diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index b5a578585d8..04d703f46a1 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -32,6 +32,7 @@ __all__ = [ "VllmAttentionInstallReport", "install_vllm_nvfp4_attention", + "install_vllm_skip_softmax_calibration", "install_vllm_sparse_attention_from_checkpoint", ] @@ -495,6 +496,77 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor return _build_report(plan) +def _attention_quant_error(module) -> str | None: + """Reject calibration on layers with any active attention Q/K/P/V fakequant.""" + for attr in ("q_bmm_quantizer", "k_bmm_quantizer", "p_bmm_quantizer", "v_bmm_quantizer"): + if getattr(getattr(module, attr, None), "is_enabled", False): + return f"{attr} is enabled; skip-softmax calibration requires unquantized attention" + if getattr(module, "_query_quant_in_kernel", False) or getattr( + module, "_value_quant_in_kernel", False + ): + return ( + "in-kernel attention quantization flags are set; skip-softmax " + "calibration requires unquantized attention" + ) + return None + + +def install_vllm_skip_softmax_calibration(model_runner) -> VllmAttentionInstallReport: + """Install skip-softmax calibration adapters into a loaded vLLM model. + + Swaps the backend-matched ModelOpt adapter onto every attention layer and + disables cascade attention, following validation-before-mutation: every + known compatibility error — across all layers — is collected and raised + before any module is changed. Calibration itself starts separately via + :func:`~.vllm.enable_calibration` (typically over a worker RPC), so engine + warmup/profiling launches after install are served natively and never + pollute the measurement; until then the adapters delegate every forward to + the backend's native implementation. + + Requirements validated here: eager execution (``enforce_eager=True`` — + the per-request calibration loop cannot be CUDA-graph captured), fp16/bf16 + model and KV-cache dtypes, no active attention Q/K/P/V fakequant, and a + FlashAttention or FlashInfer backend per layer. + """ + from vllm.config.compilation import CUDAGraphMode + + model = _unwrapped_model(model_runner) + candidates = [ + (name, module) + for name, module in model.named_modules() + if isinstance(module, _VLLM_ATTENTION) + ] + + _require_supported_vllm() + errors = _global_errors(model_runner) + if _cudagraph_mode(model_runner) != CUDAGraphMode.NONE: + errors.append( + "skip-softmax calibration requires eager execution (enforce_eager=True); " + "the per-request calibration loop cannot be CUDA-graph captured" + ) + if not candidates: + errors.append("no attention layers were found") + + plans = [] + for name, module in candidates: + reasons = _layer_errors(module) + if quant_error := _attention_quant_error(module): + reasons.append(quant_error) + new_impl, requires_flashinfer_patch, backend_error = _select_new_impl(module) + if backend_error: + reasons.append(backend_error) + if reasons: + errors.extend(f"{name or ''}: {reason}" for reason in reasons) + continue + plans.append( + _AttentionPlan(name, module, new_impl, {}, None, None, requires_flashinfer_patch) + ) + _raise_unsupported(errors, "skip-softmax calibration") + + plan = _InstallPlan(model_runner, tuple(plans), False, "SKIP_SOFTMAX_CALIBRATION") + return _apply_vllm_attention_plans(plan) + + def install_vllm_sparse_attention_from_checkpoint( model_runner, ) -> VllmAttentionInstallReport: diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 578922db077..f87a8cbb26b 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -80,7 +80,11 @@ def guarded_import(name, *args, **kwargs): monkeypatch.setattr(builtins, "__import__", guarded_import) worker_module = _load_worker_module("sparse_attn_worker_import_test") - assert worker_module.__all__ == ["SparseAttnWorker", "QuantSparseAttnWorker"] + assert worker_module.__all__ == [ + "SparseAttnWorker", + "QuantSparseAttnWorker", + "SkipSoftmaxCalibWorker", + ] @pytest.mark.parametrize( From 97bd013652bf055d8fa579f3a1f0cf8b89eb1dc6 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 21:17:29 -0700 Subject: [PATCH 04/10] Add vLLM skip-softmax calibration tests and docs Tests: - Paged calibrate kernel: paged == contiguous (counters bit-equal, output match) across aligned and non-128-aligned lengths with shuffled block tables; decode-shaped measurement covers the full cache; high-block-ID pointer regression (int64 paged offsets, memory-gated). - Calibration/serving tile contract: an active skip-softmax launch skips exactly the tiles the calibration kernel counted at the same threshold (same 128x128 geometry, same prefix-max criterion); skip composes with P/V QDQ (sm89+). - Adapter forward: mixed prefill/decode batches record per-request phases and raw counts while output stays dense vs an SDPA reference; NHD and fp16/bf16 cache validation; layer count summing; FlashInfer cache write ordered before the calibrate-kernel read. - Installer: install-without-measure lifecycle, eager/quantizer/fp8-KV/ no-layer rejections with validation-before-mutation, and the sparse-only CUDA-graph guard for calibrated decode configs. - vLLM-free helpers: count merging/alignment, TP-rank aggregation, synthetic exponential fit recovery, calibrate_from_stats field contract, canonical config round trip through the serving loader with N:M group preservation. Docs: vLLM calibration workflow section in examples/vllm_serve/README.md. Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 14 + .../attention/test_paged_calibrate.py | 236 ++++++++++ .../test_vllm_calibration.py | 411 ++++++++++++++++++ .../test_sparse_attn_calibration.py | 203 +++++++++ 4 files changed, 864 insertions(+) create mode 100644 tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py create mode 100644 tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index fc4e8a0ebcc..424163d5371 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -178,6 +178,20 @@ Workflow: If the checkpoint has no `sparse_attention_config`, the sparse-only installer passes through and vLLM runs unchanged. Whole-model fakequant flows remain handled by `vllm_serve_fakequant.py`; the compact attention-only path is below. +### Calibrate skip-softmax thresholds through vLLM + +Instead of the HF path in step 1, thresholds can be calibrated directly through vLLM — over the paged KV cache, for both prefill and decode, with tensor parallelism: + +```bash +python calibrate_sparse_attn.py \ + --prompts_file prompts.txt --target_sparse_ratio 0.5 \ + --decode_tokens 32 --tensor_parallel_size 8 --update_checkpoint_config +``` + +`install_vllm_skip_softmax_calibration` (called by `sparse_attn_worker.SkipSoftmaxCalibWorker` at model load) swaps calibration adapters onto every attention layer after validating all of them — eager execution is required, model and KV-cache dtypes must be fp16/bf16, and no attention Q/K/P/V fakequant may be active. During `llm.generate`, the paged Triton calibration kernel computes full dense attention (generation is numerically unchanged) while counting, per candidate threshold, how many KV tiles the skip criterion would drop. The driver then collects **raw tile counts from every TP rank** (each rank only measures its head shard), merges them, fits `scale_factor = a * exp(b * sparsity)` once per phase, and writes the same canonical `sparse_attention_config` block the HF export produces — preserving any exported N:M sparse-softmax groups — so the serving workflow above picks it up unchanged. + +Calibration and serving measure skipping at the same fixed 128x128 tile geometry (active skip-softmax launches bypass the autotuner), so the sparsity realized at serve time matches the calibrated `(a, b)` model. + The reusable serving policies live in `modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py`. `install_vllm_sparse_attention_from_checkpoint` installs checkpoint-driven sparse-only attention, while `install_vllm_nvfp4_attention` installs fixed NVFP4 Q/K/P/V with optional checkpoint sparsity. Both validate every selected layer before publishing any replacement implementation and return a `VllmAttentionInstallReport` with the installed layer names and backend counts. `sparse_attn_worker.py` only invokes these APIs after vLLM loads the model. It retains `SparseAttnWorker` as the launcher's default and provides `QuantSparseAttnWorker` for the compact NVFP4 policy. Other vLLM integrations can invoke the same library APIs directly: diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py new file mode 100644 index 00000000000..9e55879dbb6 --- /dev/null +++ b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Paged-cache calibration kernel tests and the calibration/serving tile contract.""" + +import pytest +import torch +from conftest import make_qkv, make_varlen_meta + +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE + +if TRITON_KERNEL_AVAILABLE: + from modelopt.torch.kernels.common.attention import attention + from modelopt.torch.kernels.sparsity.attention.calibrate import attention_calibrate + +pytestmark = pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") + +TRIALS = [1e-3, 1e-2, 1e-1, 3e-1] + + +def _pack_paged(k, v, page_size, *, shuffle=True, num_spare_blocks=7): + """Pack one sequence's contiguous K/V into a (optionally shuffled) paged cache.""" + seq = k.shape[0] + num_blocks = (seq + page_size - 1) // page_size + order = torch.randperm(num_blocks) if shuffle else torch.arange(num_blocks) + k_cache = k.new_zeros(num_blocks + num_spare_blocks, page_size, k.shape[1], k.shape[2]) + v_cache = torch.zeros_like(k_cache) + block_table = torch.zeros(1, num_blocks, device=k.device, dtype=torch.int32) + for i in range(num_blocks): + page = int(order[i]) + num_spare_blocks # keep low pages unused + ts, te = i * page_size, min((i + 1) * page_size, seq) + k_cache[page, : te - ts] = k[ts:te] + v_cache[page, : te - ts] = v[ts:te] + block_table[0, i] = page + return k_cache, v_cache, block_table + + +class TestPagedCalibrate: + @pytest.mark.parametrize("seq_len", [256, 300, 512]) # 300: non-128-aligned padding + def test_paged_matches_contiguous_prefill(self, seq_len): + """Paged and contiguous calibration agree exactly on counters and output.""" + torch.manual_seed(0) + num_heads, num_kv_heads, head_dim, page_size = 8, 2, 64, 16 + q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + locs, lens = make_varlen_meta([seq_len]) + + out_ref, counters_ref = attention_calibrate( + q, k, v, locs, lens, seq_len, is_causal=True, threshold_trials=TRIALS + ) + + k_cache, v_cache, block_table = _pack_paged(k, v, page_size) + k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) + out_paged, counters_paged = attention_calibrate( + q, + k_dummy, + k_dummy, + locs, + lens, + seq_len, + is_causal=True, + threshold_trials=TRIALS, + b_seq_len_k=lens, + max_input_len_k=seq_len, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + page_size=page_size, + ) + + assert torch.equal(counters_ref, counters_paged) + torch.testing.assert_close(out_paged, out_ref, rtol=1e-3, atol=1e-3) + + def test_paged_decode_measures_full_cache(self): + """A one-row decode query measures every KV tile of the paged cache.""" + torch.manual_seed(1) + num_heads, num_kv_heads, head_dim, page_size = 8, 2, 64, 16 + ctx = 384 + q, k, v = make_qkv(ctx, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + k_cache, v_cache, block_table = _pack_paged(k, v, page_size) + k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) + locs = torch.zeros(1, device="cuda", dtype=torch.int32) + + _, counters = attention_calibrate( + q[:1], + k_dummy, + k_dummy, + locs, + torch.ones(1, device="cuda", dtype=torch.int32), + 1, + is_causal=False, + threshold_trials=TRIALS, + b_seq_len_k=torch.tensor([ctx], device="cuda", dtype=torch.int32), + max_input_len_k=ctx, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + page_size=page_size, + ) + + num_kv_tiles = -(-ctx // 128) + assert counters[:, 0].tolist() == [num_heads * num_kv_tiles] * len(TRIALS) + + def test_high_block_id_pointer_arithmetic(self): + """Block IDs whose int32 byte offsets would wrap still read correctly.""" + num_kv_heads, head_dim, page_size = 2, 64, 16 + block_elems = page_size * num_kv_heads * head_dim + # Smallest block ID whose element offset exceeds int32. + high_block = (2**31) // block_elems + 1 + bytes_needed = 2 * (high_block + 1) * block_elems * 2 # K and V caches, bf16 + free, _ = torch.cuda.mem_get_info() + if free < bytes_needed + (2 << 30): + pytest.skip(f"needs ~{bytes_needed / 2**30:.1f} GiB free GPU memory") + + torch.manual_seed(2) + num_heads = 4 + q, k, v = make_qkv(page_size, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + k_cache = torch.zeros( + high_block + 1, page_size, num_kv_heads, head_dim, device="cuda", dtype=torch.bfloat16 + ) + v_cache = torch.zeros_like(k_cache) + k_cache[high_block] = k + v_cache[high_block] = v + block_table = torch.tensor([[high_block]], device="cuda", dtype=torch.int32) + locs, lens = make_varlen_meta([page_size]) + + out_ref, counters_ref = attention_calibrate( + q, k, v, locs, lens, page_size, is_causal=True, threshold_trials=TRIALS + ) + k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) + out_paged, counters_paged = attention_calibrate( + q, + k_dummy, + k_dummy, + locs, + lens, + page_size, + is_causal=True, + threshold_trials=TRIALS, + b_seq_len_k=lens, + max_input_len_k=page_size, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + page_size=page_size, + ) + del k_cache, v_cache + + assert torch.equal(counters_ref, counters_paged) + torch.testing.assert_close(out_paged, out_ref, rtol=1e-3, atol=1e-3) + + +class TestCalibrationServingTileContract: + """Active skip launches and calibration must count identically (same tiles).""" + + def _contrasty_qkv(self, seq_len, num_heads, num_kv_heads, head_dim): + """K with a dominant head-of-sequence so later tiles are skippable.""" + torch.manual_seed(3) + q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + k = k * 0.05 + k[:32] = k[:32] * 600.0 # first tile dominates the running max by >> log2(threshold) + return q, k, v + + @pytest.mark.parametrize("threshold", [1e-3, 1e-2]) + def test_serve_skip_counts_equal_calibrate_counts(self, threshold): + seq_len, num_heads, num_kv_heads, head_dim = 512, 8, 2, 64 + q, k, v = self._contrasty_qkv(seq_len, num_heads, num_kv_heads, head_dim) + locs, lens = make_varlen_meta([seq_len]) + scale = 1.0 / (head_dim**0.5) + + _, counters = attention_calibrate( + q, + k, + v, + locs, + lens, + seq_len, + is_causal=True, + softmax_scale=scale, + threshold_trials=[threshold], + ) + + out = attention( + q, + k, + v, + locs, + lens, + seq_len, + is_causal=True, + softmax_scale=scale, + skip_softmax_threshold=threshold, + measure_sparsity=True, + ) + + calib_total, calib_skipped = int(counters[0, 0]), int(counters[0, 1]) + assert calib_skipped > 0, "test data must produce skippable tiles" + # Same 128x128 tile geometry and same prefix-max criterion => the serve + # kernel must skip exactly the tiles calibration predicted. + assert out._sparsity_total == calib_total + assert out._sparsity_skipped == calib_skipped + + def test_skip_composes_with_pv_qdq(self): + """Active skip at the fixed tile compiles and runs with P/V QDQ enabled.""" + if torch.cuda.get_device_capability() < (8, 9): + pytest.skip("NVFP4/FP8 QDQ needs compute capability >= 8.9 (E4M3 casts)") + seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 + q, k, v = self._contrasty_qkv(seq_len, num_heads, num_kv_heads, head_dim) + locs, lens = make_varlen_meta([seq_len]) + + out = attention( + q, + k, + v, + locs, + lens, + seq_len, + is_causal=True, + skip_softmax_threshold=1e-2, + p_qdq="nvfp4", + p_qdq_amax=1.0, + v_qdq="nvfp4", + v_qdq_amax=float(v.abs().amax()), + ) + assert torch.isfinite(out).all() diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py new file mode 100644 index 00000000000..e5192e7e7f3 --- /dev/null +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -0,0 +1,411 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for skip-softmax calibration through the vLLM adapters and installer.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl + +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as attention_plugin +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm_runtime +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( + ModelOptSparseAttentionImpl, + collect_calibration_counts, + disable_calibration, + enable_calibration, + iter_sparse_impls, +) + +TRIALS = [1e-3, 1e-1, 5e-1] + + +# --------------------------------------------------------------------------- +# Installer fakes (mirroring test_vllm_runtime.py) +# --------------------------------------------------------------------------- +def _bare_attention(impl_cls=FlashAttentionImpl): + module = object.__new__(vllm_runtime._VLLM_ATTENTION) + nn.Module.__init__(module) + module.attn_type = "decoder" + module.head_size = 64 + module.device = torch.device("cpu") + module.dtype = torch.float16 + module.impl = object.__new__(impl_cls) + module.impl.sinks = None + return module + + +def _model_runner(model, *, sparse_metadata=None, cudagraph_mode=CUDAGraphMode.NONE): + hf_config = SimpleNamespace(sparse_attention_config=sparse_metadata) + model_config = SimpleNamespace(hf_config=hf_config, dtype=torch.float16) + return SimpleNamespace( + model=model, + model_config=model_config, + cascade_attn_enabled=True, + vllm_config=SimpleNamespace( + model_config=model_config, + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + enable_dbo=False, + use_ubatching=False, + ), + cache_config=SimpleNamespace(enable_prefix_caching=False, cache_dtype="auto"), + compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), + kv_transfer_config=None, + speculative_config=None, + ), + ) + + +class TestCalibrationInstaller: + def test_installs_adapters_without_enabling_measurement(self): + first = _bare_attention() + second = _bare_attention() + runner = _model_runner(nn.ModuleDict({"a_attn": first, "b_attn": second})) + + report = vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + assert report.installed_count == 2 + assert report.sparse_algorithm == "SKIP_SOFTMAX_CALIBRATION" + assert report.cascade_disabled and runner.cascade_attn_enabled is False + for module in (first, second): + assert isinstance(module.impl, ModelOptSparseAttentionImpl) + # Measurement starts only via enable_calibration, so warmup + # launches after install are never recorded. + assert not attention_plugin._calibration_active(module.impl) + + def test_rejects_non_eager_execution(self): + runner = _model_runner( + nn.ModuleDict({"attn": _bare_attention()}), + cudagraph_mode=CUDAGraphMode.PIECEWISE, + ) + with pytest.raises(NotImplementedError, match="enforce_eager"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + def test_rejects_active_attention_quantizers_atomically(self): + quantized = _bare_attention() + quantized.q_bmm_quantizer = SimpleNamespace(is_enabled=True) + clean = _bare_attention() + clean_impl = clean.impl + runner = _model_runner(nn.ModuleDict({"q_attn": quantized, "c_attn": clean})) + + with pytest.raises(NotImplementedError, match="requires unquantized attention"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + # Validation-before-mutation: the clean layer must not be touched either. + assert clean.impl is clean_impl + assert runner.cascade_attn_enabled is True + + def test_rejects_fp8_kv_cache(self): + attention = _bare_attention() + attention.kv_cache_dtype = "fp8" + runner = _model_runner(nn.ModuleDict({"attn": attention})) + with pytest.raises(NotImplementedError, match="FP8 KV cache"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + def test_rejects_model_without_attention_layers(self): + runner = _model_runner(nn.ModuleDict({})) + with pytest.raises(NotImplementedError, match="no attention layers"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + +class TestSparseOnlyGraphGuard: + """Commit-contract: the calibrated-decode graph guard is not quantize-gated.""" + + _CALIBRATED_META = { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "threshold_scale_factor": { + "prefill": {"a": 7.9, "b": 8.6}, + "decode": {"a": 0.12, "b": 9.8}, + }, + } + } + } + + def test_sparse_only_install_rejects_full_decode_graph(self): + attention = _bare_attention() + runner = _model_runner( + nn.ModuleDict({"attn": attention}), + sparse_metadata=self._CALIBRATED_META, + cudagraph_mode=CUDAGraphMode.FULL, + ) + with pytest.raises(NotImplementedError, match="non-FULL CUDA graph"): + vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert not isinstance(attention.impl, ModelOptSparseAttentionImpl) + + def test_sparse_only_install_allows_eager(self): + attention = _bare_attention() + runner = _model_runner( + nn.ModuleDict({"attn": attention}), + sparse_metadata=self._CALIBRATED_META, + cudagraph_mode=CUDAGraphMode.NONE, + ) + report = vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert report.installed_count == 1 + + +# --------------------------------------------------------------------------- +# Calibration forward through the FlashAttention adapter (GPU) +# --------------------------------------------------------------------------- +def _make_impl(num_heads, head_dim, num_kv_heads): + return ModelOptSparseAttentionImpl( + num_heads=num_heads, + head_size=head_dim, + scale=1.0 / (head_dim**0.5), + num_kv_heads=num_kv_heads, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + ) + + +def _paged_cache_for(seqs_kv, num_kv_heads, head_dim, page_size, device, dtype): + """Scatter per-request contiguous K/V lists into a stacked paged cache.""" + blocks_per_seq = [(kv.shape[0] + page_size - 1) // page_size for kv, _ in seqs_kv] + num_blocks = sum(blocks_per_seq) + max_blocks = max(blocks_per_seq) + k_cache = torch.zeros(num_blocks, page_size, num_kv_heads, head_dim, device=device, dtype=dtype) + v_cache = torch.zeros_like(k_cache) + block_table = torch.zeros(len(seqs_kv), max_blocks, device=device, dtype=torch.int32) + g = 0 + for b, (k, v) in enumerate(seqs_kv): + for blk in range(blocks_per_seq[b]): + block_table[b, blk] = g + ts, te = blk * page_size, min((blk + 1) * page_size, k.shape[0]) + k_cache[g, : te - ts] = k[ts:te] + v_cache[g, : te - ts] = v[ts:te] + g += 1 + return torch.stack([k_cache, v_cache], dim=0), block_table + + +def _sdpa_reference(q, k, v, is_causal): + # [tokens, heads, dim] -> [1, heads, tokens, dim] + qh, kh, vh = (t.transpose(0, 1).unsqueeze(0).float() for t in (q, k, v)) + kh = kh.repeat_interleave(q.shape[1] // k.shape[1], dim=1) + vh = vh.repeat_interleave(q.shape[1] // v.shape[1], dim=1) + if is_causal and q.shape[0] < k.shape[0]: + # Suffix-causal mask for decode/chunked prefill shapes. + mask = torch.ones(q.shape[0], k.shape[0], dtype=torch.bool, device=q.device).tril( + diagonal=k.shape[0] - q.shape[0] + ) + out = torch.nn.functional.scaled_dot_product_attention(qh, kh, vh, attn_mask=mask) + else: + out = torch.nn.functional.scaled_dot_product_attention(qh, kh, vh, is_causal=is_causal) + return out.squeeze(0).transpose(0, 1).to(q.dtype) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestCalibrationForward: + def test_mixed_batch_records_phases_and_stays_dense(self): + torch.manual_seed(0) + device, dtype = "cuda", torch.bfloat16 + num_heads, num_kv_heads, head_dim, page_size = 4, 2, 64, 16 + prefill_len, decode_ctx = 64, 48 + + k0 = torch.randn(prefill_len, num_kv_heads, head_dim, device=device, dtype=dtype) + v0 = torch.randn_like(k0) + k1 = torch.randn(decode_ctx, num_kv_heads, head_dim, device=device, dtype=dtype) + v1 = torch.randn_like(k1) + q = torch.randn(prefill_len + 1, num_heads, head_dim, device=device, dtype=dtype) + + kv_cache, block_table = _paged_cache_for( + [(k0, v0), (k1, v1)], num_kv_heads, head_dim, page_size, device, dtype + ) + attn_metadata = SimpleNamespace( + num_actual_tokens=prefill_len + 1, + max_query_len=prefill_len, + max_seq_len=max(prefill_len, decode_ctx), + query_start_loc=torch.tensor( + [0, prefill_len, prefill_len + 1], device=device, dtype=torch.int32 + ), + seq_lens=torch.tensor([prefill_len, decode_ctx], device=device, dtype=torch.int32), + block_table=block_table, + ) + + impl = _make_impl(num_heads, head_dim, num_kv_heads) + impl.sparse_kw = {} + enable_calibration([impl], TRIALS) + output = torch.empty_like(q) + out = impl.forward( + layer=None, + query=q, + key=q[:, :num_kv_heads], + value=q[:, :num_kv_heads], + kv_cache=kv_cache, + attn_metadata=attn_metadata, + output=output, + ) + + # Two records: one per request, phases decided per request. + records = impl._calib_records + assert [r["phase"] for r in records] == ["prefill", "decode"] + assert [r["sample_length"] for r in records] == [prefill_len, decode_ctx] + for record in records: + assert len(record["total_tiles"]) == len(TRIALS) + assert all(t > 0 for t in record["total_tiles"]) + assert all(0 <= s <= t for s, t in zip(record["skipped_tiles"], record["total_tiles"])) + + # Output is full dense attention (calibration never skips). + ref_prefill = _sdpa_reference(q[:prefill_len], k0, v0, is_causal=True) + ref_decode = _sdpa_reference(q[prefill_len:], k1, v1, is_causal=False) + torch.testing.assert_close(out[:prefill_len], ref_prefill, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(out[prefill_len:], ref_decode, rtol=2e-2, atol=2e-2) + + def test_collect_calibration_counts_sums_layers(self): + class FakeModel(nn.Module): + def __init__(self, impls): + super().__init__() + self._impls = impls + self.layers = nn.ModuleList([nn.Identity() for _ in impls]) + for identity, impl in zip(self.layers, impls): + identity.impl = impl + + impls = [object.__new__(ModelOptSparseAttentionImpl) for _ in range(2)] + enable_calibration(impls, TRIALS) + for idx, impl in enumerate(impls): + impl._calib_records = [ + { + "phase": "prefill", + "sample_length": 128, + "total_tiles": [4, 4, 4], + "skipped_tiles": [idx, idx + 1, idx + 2], + } + ] + model = FakeModel(impls) + assert len(list(iter_sparse_impls(model))) == 2 + disable_calibration(impls) + + counts = collect_calibration_counts(model) + assert counts["prefill"] == [ + {"sample_length": 128, "total_tiles": [8, 8, 8], "skipped_tiles": [1, 3, 5]} + ] + + def test_rejects_non_nhd_cache_layout(self): + num_heads, num_kv_heads, head_dim = 4, 2, 64 + impl = _make_impl(num_heads, head_dim, num_kv_heads) + enable_calibration([impl], TRIALS) + # HND-shaped cache: [blocks, kv_heads, page, dim] -> axis 2 != num_kv_heads. + kv_cache = torch.zeros(2, 1, num_kv_heads, 16, head_dim, dtype=torch.bfloat16) + attn_metadata = SimpleNamespace( + num_actual_tokens=1, + max_query_len=1, + max_seq_len=8, + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=torch.tensor([8], dtype=torch.int32), + block_table=torch.zeros(1, 1, dtype=torch.int32), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + with pytest.raises(NotImplementedError, match="not NHD"): + impl.forward( + layer=None, + query=q, + key=q[:, :num_kv_heads], + value=q[:, :num_kv_heads], + kv_cache=kv_cache, + attn_metadata=attn_metadata, + output=torch.empty_like(q), + ) + + def test_rejects_non_16bit_cache(self): + num_heads, num_kv_heads, head_dim = 4, 2, 64 + impl = _make_impl(num_heads, head_dim, num_kv_heads) + enable_calibration([impl], TRIALS) + kv_cache = torch.zeros(2, 1, 16, num_kv_heads, head_dim, dtype=torch.uint8) + attn_metadata = SimpleNamespace( + num_actual_tokens=1, + max_query_len=1, + max_seq_len=8, + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=torch.tensor([8], dtype=torch.int32), + block_table=torch.zeros(1, 1, dtype=torch.int32), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + with pytest.raises(NotImplementedError, match="fp16/bf16 KV cache"): + impl.forward( + layer=None, + query=q, + key=q[:, :num_kv_heads], + value=q[:, :num_kv_heads], + kv_cache=kv_cache, + attn_metadata=attn_metadata, + output=torch.empty_like(q), + ) + + +# --------------------------------------------------------------------------- +# FlashInfer adapter: cache write must precede the calibrate-kernel read +# --------------------------------------------------------------------------- +class TestFlashInferCalibrationOrdering: + def test_cache_write_happens_before_calibrate_read(self, monkeypatch): + calls = [] + monkeypatch.setattr( + attention_plugin, + "_maybe_update_flashinfer_cache", + lambda *args, **kwargs: calls.append("cache_write"), + ) + + def fake_calibrate(q, *args, **kwargs): + calls.append("calibrate") + counters = torch.zeros(len(TRIALS), 2, dtype=torch.int64) + return torch.zeros_like(q), counters + + monkeypatch.setattr(attention_plugin, "attention_calibrate", fake_calibrate) + + num_heads, num_kv_heads, head_dim, page = 4, 2, 64, 16 + impl = SimpleNamespace( + num_kv_heads=num_kv_heads, + head_size=head_dim, + scale=1.0 / (head_dim**0.5), + _calibrate=True, + _calib_threshold_trials=list(TRIALS), + _calib_records=[], + ) + kv_cache = torch.zeros(3, 2, page, num_kv_heads, head_dim, dtype=torch.bfloat16) + attn_metadata = SimpleNamespace( + _modelopt_block_table=torch.zeros(1, 1, dtype=torch.int32), + _modelopt_seq_lens=torch.tensor([8], dtype=torch.int32), + _modelopt_query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + _modelopt_num_actual_tokens=1, + _modelopt_max_query_len=1, + _modelopt_max_seq_len=8, + _modelopt_causal=False, + slot_mapping=torch.zeros(1, dtype=torch.int64), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + + out = attention_plugin._flashinfer_forward( + impl, + None, # native_forward is unused on the calibration path + None, # layer + q, + q[:, :num_kv_heads], + q[:, :num_kv_heads], + kv_cache, + attn_metadata, + output=torch.empty_like(q), + ) + + assert calls == ["cache_write", "calibrate"] + assert torch.isfinite(out).all() + assert len(impl._calib_records) == 1 + assert impl._calib_records[0]["phase"] == "decode" diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py new file mode 100644 index 00000000000..77f836c9205 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for vLLM-free skip-softmax calibration helpers (no vLLM needed).""" + +import math +from types import SimpleNamespace + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration.calibrator import ( + DynamicThresholdCalibrator, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_calibration import ( + DEFAULT_THRESHOLD_TRIALS, + build_sparse_attention_config, + fit_from_counts, + merge_count_records, + merge_phase_counts, + split_records_by_phase, + stats_from_counts, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_config import ( + load_from_checkpoint_metadata, +) + + +def _record(phase, length, totals, skipped): + return { + "phase": phase, + "sample_length": length, + "total_tiles": list(totals), + "skipped_tiles": list(skipped), + } + + +class TestCountMerging: + def test_split_records_by_phase_preserves_order(self): + records = [ + _record("prefill", 100, [4], [1]), + _record("decode", 101, [2], [0]), + _record("prefill", 200, [8], [3]), + ] + split = split_records_by_phase(records) + assert [r["sample_length"] for r in split["prefill"]] == [100, 200] + assert [r["sample_length"] for r in split["decode"]] == [101] + + def test_merge_sums_counts_elementwise(self): + layer_a = [_record("prefill", 128, [10, 10], [2, 4])] + layer_b = [_record("prefill", 128, [10, 10], [1, 3])] + merged = merge_count_records([layer_a, layer_b]) + assert merged == [{"sample_length": 128, "total_tiles": [20, 20], "skipped_tiles": [3, 7]}] + + def test_merge_truncates_to_shortest_source(self): + long = [_record("prefill", 1, [1], [0]), _record("prefill", 2, [1], [1])] + short = [_record("prefill", 1, [1], [1])] + merged = merge_count_records([long, short]) + assert len(merged) == 1 + assert merged[0]["skipped_tiles"] == [1] + + def test_merge_rejects_misaligned_sample_lengths(self): + with pytest.raises(ValueError, match="Misaligned calibration records"): + merge_count_records( + [[_record("prefill", 100, [1], [0])], [_record("prefill", 200, [1], [0])]] + ) + + def test_merge_phase_counts_across_ranks(self): + rank0 = {"prefill": [_record("prefill", 64, [5], [1])], "decode": []} + rank1 = {"prefill": [_record("prefill", 64, [5], [2])]} + merged = merge_phase_counts([rank0, rank1]) + assert merged["prefill"][0]["total_tiles"] == [10] + assert merged["prefill"][0]["skipped_tiles"] == [3] + assert merged["decode"] == [] + + def test_stats_from_counts_forms_ratios_after_merge(self): + stats = stats_from_counts( + [{"sample_length": 64, "total_tiles": [8, 0], "skipped_tiles": [2, 0]}] + ) + assert stats == [{"sample_length": 64, "sparsity": [0.25, 0.0]}] + + +class TestFitFromCounts: + def test_fit_recovers_synthetic_exponential(self): + a_true, b_true = 5.0, 8.0 + trials = DEFAULT_THRESHOLD_TRIALS + + def counts(length, total): + sparsity = [ + min(0.95, max(0.0, math.log(max(t * length, 1e-9) / a_true) / b_true)) + for t in trials + ] + return { + "sample_length": length, + "total_tiles": [total] * len(trials), + "skipped_tiles": [int(s * total) for s in sparsity], + } + + per_phase = {"prefill": [counts(length, 4000) for length in (2048, 4096, 8192, 16384)]} + params = fit_from_counts(per_phase, trials) + assert abs(params["prefill"]["a"] - a_true) / a_true < 0.3 + assert abs(params["prefill"]["b"] - b_true) / b_true < 0.15 + assert 0.0 <= params["prefill"]["min_observed_sparsity"] <= 1.0 + + def test_empty_phase_produces_no_fit(self): + assert fit_from_counts({"decode": []}, DEFAULT_THRESHOLD_TRIALS) == {} + + +class TestCalibrateFromStats: + def _stats(self, trials): + a_true, b_true = 3.0, 9.0 + stats = [] + for length in (1024, 2048, 4096, 8192): + sparsity = [ + min(0.95, max(0.0, math.log(max(t * length, 1e-9) / a_true) / b_true)) + for t in trials + ] + stats.append({"sample_length": length, "sparsity": sparsity}) + return stats + + def test_linear_fit_reports_fit_logspace_false(self): + calibrator = DynamicThresholdCalibrator(threshold_trials=list(DEFAULT_THRESHOLD_TRIALS)) + result = calibrator.calibrate_from_stats(self._stats(DEFAULT_THRESHOLD_TRIALS), "prefill") + assert result["fit_logspace"] is False + assert "log_a" not in result + assert len(result["per_sample_sparsity"]) == 4 + + def test_logspace_fit_preserves_log_a(self): + calibrator = DynamicThresholdCalibrator( + threshold_trials=list(DEFAULT_THRESHOLD_TRIALS), fit_logspace=True + ) + result = calibrator.calibrate_from_stats(self._stats(DEFAULT_THRESHOLD_TRIALS), "prefill") + assert result["fit_logspace"] is True + assert math.isclose(math.exp(result["log_a"]), result["a"], rel_tol=1e-9) + + +class TestBuildSparseAttentionConfig: + _PARAMS = {"prefill": {"a": 7.9, "b": 8.6}, "decode": {"a": 0.12, "b": 9.8}} + + def test_canonical_schema(self): + config = build_sparse_attention_config(self._PARAMS, 0.4) + group = config["config_groups"]["group_0"] + assert group["algorithm"] == "skip_softmax" + assert group["threshold_scale_factor"]["prefill"] == {"a": 7.9, "b": 8.6} + assert group["threshold_scale_factor"]["formula"] == "a * exp(b * target_sparsity)" + assert group["target_sparsity"] == {"prefill": 0.4, "decode": 0.4} + assert config["producer"]["name"] == "modelopt" + + def test_preserves_nm_groups_and_replaces_old_skip_group(self): + existing = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4}, + "group_1": { + "algorithm": "skip_softmax", + "threshold_scale_factor": {"prefill": {"a": 1.0, "b": 1.0}}, + }, + } + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + groups = config["config_groups"] + assert len(groups) == 2 + assert groups["group_0"]["algorithm"] == "skip_softmax" + assert groups["group_0"]["threshold_scale_factor"]["prefill"]["a"] == 7.9 + assert groups["group_1"]["algorithm"] == "sparse_softmax" + assert groups["group_1"]["sparsity_n"] == 2 + + def test_round_trips_through_serving_loader(self): + config = build_sparse_attention_config(self._PARAMS, {"prefill": 0.5, "decode": 0.3}) + hf_config = SimpleNamespace(sparse_attention_config=config) + loaded = load_from_checkpoint_metadata(hf_config) + assert loaded is not None + sparse_cfg, preset = loaded + assert preset == "CHECKPOINT_CALIBRATED_SOFTMAX_SKIP" + layer_cfg = sparse_cfg["sparse_cfg"]["*attn*"] + assert layer_cfg["method"] == "triton_skip_softmax" + assert layer_cfg["threshold_scale_factor"]["decode"] == {"a": 0.12, "b": 9.8} + assert layer_cfg["target_sparse_ratio"] == {"prefill": 0.5, "decode": 0.3} + + def test_round_trip_with_preserved_nm_group_activates_both(self): + existing = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4} + } + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + loaded = load_from_checkpoint_metadata(SimpleNamespace(sparse_attention_config=config)) + assert loaded is not None + sparse_cfg, preset = loaded + assert preset == "CHECKPOINT_CALIBRATED_SOFTMAX_SKIP_SPARSE_SOFTMAX" + layer_cfg = sparse_cfg["sparse_cfg"]["*attn*"] + assert layer_cfg["sparsity_n"] == 2 + assert "threshold_scale_factor" in layer_cfg From aca90e40c3561ca1ddf4b32355d80e4689c84b41 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 18 Jul 2026 21:27:13 -0700 Subject: [PATCH 05/10] Add changelog entry for vLLM skip-softmax calibration Signed-off-by: Kai Xu --- CHANGELOG.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6687ebd31ea..a8a1e8b2720 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,10 @@ Changelog **New Features** +*Sparsity* + +- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. See ``examples/vllm_serve/calibrate_sparse_attn.py``. + *Quantization* - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported. From a0461394cb130b6d72691ac406ba9a11e42866d5 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 18 Jul 2026 21:27:17 -0700 Subject: [PATCH 06/10] Default vLLM calibration prompts to the RULER dataset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Build calibration prompts with RulerDatasetBuilder — the same dataset the PyTorch (HF) calibration path uses — so vLLM- and HF-calibrated thresholds are fit on identical data (defaults mirror the HF path: 24 samples, max_seqlen 32768; NIAH essay haystack via download_ruler_data.sh and --calib_data_dir). The built-in demo prompt fallback is removed; --prompts_file remains as an explicit override for custom calibration data. Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 8 +- examples/vllm_serve/calibrate_sparse_attn.py | 91 ++++++++++++++------ 2 files changed, 74 insertions(+), 25 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 424163d5371..43628fbb7a3 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -183,11 +183,17 @@ If the checkpoint has no `sparse_attention_config`, the sparse-only installer pa Instead of the HF path in step 1, thresholds can be calibrated directly through vLLM — over the paged KV cache, for both prefill and decode, with tensor parallelism: ```bash +# One-time: fetch the RULER essay haystack +bash ../llm_sparsity/attention_sparsity/download_ruler_data.sh + python calibrate_sparse_attn.py \ - --prompts_file prompts.txt --target_sparse_ratio 0.5 \ + --calib_data_dir ../llm_sparsity/attention_sparsity/data \ + --target_sparse_ratio 0.5 \ --decode_tokens 32 --tensor_parallel_size 8 --update_checkpoint_config ``` +Calibration prompts default to the **RULER dataset** via the same `RulerDatasetBuilder` the HF calibration path uses (`--calib_samples` / `--calib_max_seqlen` mirror the HF defaults of 24 / 32768), so vLLM- and PyTorch-calibrated thresholds are fit on identical data. `--prompts_file` (one prompt per line) substitutes custom calibration data. + `install_vllm_skip_softmax_calibration` (called by `sparse_attn_worker.SkipSoftmaxCalibWorker` at model load) swaps calibration adapters onto every attention layer after validating all of them — eager execution is required, model and KV-cache dtypes must be fp16/bf16, and no attention Q/K/P/V fakequant may be active. During `llm.generate`, the paged Triton calibration kernel computes full dense attention (generation is numerically unchanged) while counting, per candidate threshold, how many KV tiles the skip criterion would drop. The driver then collects **raw tile counts from every TP rank** (each rank only measures its head shard), merges them, fits `scale_factor = a * exp(b * sparsity)` once per phase, and writes the same canonical `sparse_attention_config` block the HF export produces — preserving any exported N:M sparse-softmax groups — so the serving workflow above picks it up unchanged. Calibration and serving measure skipping at the same fixed 128x128 tile geometry (active skip-softmax launches bypass the autotuner), so the sparsity realized at serve time matches the calibrated `(a, b)` model. diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py index 6c36f8f65ec..65091485981 100644 --- a/examples/vllm_serve/calibrate_sparse_attn.py +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -32,14 +32,16 @@ Usage: python calibrate_sparse_attn.py \ - --prompts_file prompts.txt \ --target_sparse_ratio 0.5 \ --decode_tokens 32 \ --update_checkpoint_config -``--prompts_file`` is one prompt per line; longer, varied-length prompts give a -better fit. With no file, a tiny built-in demo set is used (fine for a smoke -test, not for a real fit). +Calibration prompts default to the RULER dataset — the same +``RulerDatasetBuilder`` the PyTorch (HF) calibration path uses — so both paths +calibrate on identical data. NIAH tasks need the essay haystack downloaded by +``examples/llm_sparsity/attention_sparsity/download_ruler_data.sh`` (point +``--calib_data_dir`` at its ``data`` directory). ``--prompts_file`` (one prompt +per line) overrides the RULER set with custom calibration data. """ import argparse @@ -55,25 +57,39 @@ merge_phase_counts, ) -_DEMO_PROMPTS = [ - "Summarize the history of computing in a few paragraphs. " * 40, - "Explain how attention works in transformer models. " * 60, - "Write a detailed essay about renewable energy sources. " * 80, -] +def _load_prompts(llm, args) -> list[str]: + """Load override prompts from a file, or build the default RULER set.""" + if args.prompts_file is not None: + lines = [ + ln.strip() for ln in Path(args.prompts_file).read_text().splitlines() if ln.strip() + ] + if not lines: + raise ValueError(f"No prompts found in {args.prompts_file}") + print(f"[ModelOpt] Loaded {len(lines)} calibration prompts from {args.prompts_file}") + return lines -def _load_prompts(prompts_file: str | None) -> list[str]: - if prompts_file is None: - print( - "[ModelOpt] No --prompts_file given; using a tiny built-in demo set. " - "Pass real, varied-length prompts for a usable fit." - ) - return _DEMO_PROMPTS - lines = [ln.strip() for ln in Path(prompts_file).read_text().splitlines() if ln.strip()] - if not lines: - raise ValueError(f"No prompts found in {prompts_file}") - print(f"[ModelOpt] Loaded {len(lines)} calibration prompts from {prompts_file}") - return lines + # Same dataset as the HF calibration path (calibration/calibrate.py), so the + # vLLM- and PyTorch-calibrated thresholds are fit on identical data. + from modelopt.torch.sparsity.attention_sparsity.calibration.ruler_dataset import ( + RulerDatasetBuilder, + ) + + builder = RulerDatasetBuilder( + samples=args.calib_samples, + max_seqlen=args.calib_max_seqlen, + tokenizer_name_or_path=llm.get_tokenizer(), + max_length_filter=int(args.calib_max_seqlen * 1.5), + data_dir=args.calib_data_dir, + ) + samples = builder.build_calibration_dataset() + prompts = [sample["input"] for sample in samples] + lengths = sorted(sample["length"] for sample in samples) + print( + f"[ModelOpt] Built {len(prompts)} RULER calibration prompts " + f"(token lengths {lengths[0]}..{lengths[-1]})" + ) + return prompts def _existing_sparse_config(ckpt: str) -> dict | None: @@ -108,7 +124,33 @@ def _write_config(ckpt: str, sparse_config: dict, update_checkpoint: bool) -> No def main(): parser = argparse.ArgumentParser(description="Calibrate skip-softmax thresholds via vLLM") parser.add_argument("model", type=str, help="Path to the HF checkpoint to calibrate") - parser.add_argument("--prompts_file", type=str, default=None, help="One prompt per line") + parser.add_argument( + "--prompts_file", + type=str, + default=None, + help="Optional custom calibration prompts (one per line), overriding the " + "default RULER dataset", + ) + parser.add_argument( + "--calib_samples", + type=int, + default=24, + help="Total RULER samples, distributed across length bins (HF-path default: 24)", + ) + parser.add_argument( + "--calib_max_seqlen", + type=int, + default=32768, + help="Maximum RULER sequence length; length bins descend in powers of 2. " + "Must fit within --max_model_len together with --decode_tokens.", + ) + parser.add_argument( + "--calib_data_dir", + type=str, + default=None, + help="RULER data directory containing the 'essays' haystack (populated by " + "examples/llm_sparsity/attention_sparsity/download_ruler_data.sh)", + ) parser.add_argument( "--target_sparse_ratio", type=float, @@ -176,8 +218,6 @@ def main(): from vllm import LLM, SamplingParams - prompts = _load_prompts(args.prompts_file) - llm_kwargs = { "model": args.model, "worker_cls": "sparse_attn_worker.SkipSoftmaxCalibWorker", @@ -207,6 +247,9 @@ def main(): llm_kwargs.update(extra) llm = LLM(**llm_kwargs) + # Built after engine init so the RULER builder reuses the engine's tokenizer. + prompts = _load_prompts(llm, args) + trials = list(DEFAULT_THRESHOLD_TRIALS) n_layers = llm.collective_rpc("sparse_calib_enable", args=(trials,))[0] status = llm.collective_rpc("sparse_calib_status")[0] From 27b749f78242cdfc9a256db080e5d0bbc00afa5d Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 18 Jul 2026 23:11:04 -0700 Subject: [PATCH 07/10] Address review: tile-contract, quant-composition, and calibration data-flow hardening - Reject skip-softmax combined with ANY attention quantization at plan time: quantized Q/K/P change the score distribution the thresholds were calibrated on (N:M sparse softmax still composes with quantization). - The 128x128 skip tile is now a hard contract: configurations that cannot compile it (fp32 on ~100KB-smem GPUs) are rejected with a clear error instead of re-tiled (the BLOCK_M step-down changed realized sparsity). - FlashInfer layout: query vLLM layout metadata and reject HND at install and before the calibration cache write; the NHD shape check remains as a fallback (it is ambiguous when page_size equals the per-rank KV heads). - One canonical threshold sweep shared by the HF and vLLM calibration paths (DEFAULT_THRESHOLD_TRIALS hoisted from DynamicThresholdCalibrator); ignore_eos forces the full decode length; the driver exits nonzero unless every requested phase produced a valid fit (no silent partial export). - Count merging treats alignment as a contract: mismatched sample counts, lengths, threshold widths, or per-rank/per-layer phase coverage raise instead of silently truncating. - Preserve legacy top-level sparse_softmax metadata through recalibration. - Docs: no-sparsification wording (kernel numerics differ from the native backend); clarify prefix-caching support (sparse-only serving supports suffix attention; quantized installs and calibration reject it). - High-block-ID regression test halves its allocation (V aliases K storage). Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 +- examples/vllm_serve/README.md | 6 +- examples/vllm_serve/calibrate_sparse_attn.py | 18 ++-- .../kernels/common/attention/triton_fa.py | 52 ++++----- .../calibration/calibrator.py | 50 +++++---- .../plugins/sparse_attn_calibration.py | 77 ++++++++----- .../attention_sparsity/plugins/vllm.py | 65 +++++++++-- .../plugins/vllm_runtime.py | 23 ++++ .../attention/test_paged_calibrate.py | 12 +-- .../attention/test_triton_fa_calibrate.py | 9 +- .../attention/test_triton_fa_skip_softmax.py | 31 ++++-- .../test_vllm_calibration.py | 101 ++++++++++++++++++ .../test_sparse_attn_calibration.py | 35 +++++- 13 files changed, 357 insertions(+), 124 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a8a1e8b2720..6d9d2152752 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog *Sparsity* -- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. See ``examples/vllm_serve/calibrate_sparse_attn.py``. +- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. Configurations that cannot compile the 128x128 tile are rejected rather than re-tiled, and skip-softmax cannot be combined with attention quantization. See ``examples/vllm_serve/calibrate_sparse_attn.py``. *Quantization* diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 43628fbb7a3..34a802c85e2 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -194,7 +194,7 @@ python calibrate_sparse_attn.py \ Calibration prompts default to the **RULER dataset** via the same `RulerDatasetBuilder` the HF calibration path uses (`--calib_samples` / `--calib_max_seqlen` mirror the HF defaults of 24 / 32768), so vLLM- and PyTorch-calibrated thresholds are fit on identical data. `--prompts_file` (one prompt per line) substitutes custom calibration data. -`install_vllm_skip_softmax_calibration` (called by `sparse_attn_worker.SkipSoftmaxCalibWorker` at model load) swaps calibration adapters onto every attention layer after validating all of them — eager execution is required, model and KV-cache dtypes must be fp16/bf16, and no attention Q/K/P/V fakequant may be active. During `llm.generate`, the paged Triton calibration kernel computes full dense attention (generation is numerically unchanged) while counting, per candidate threshold, how many KV tiles the skip criterion would drop. The driver then collects **raw tile counts from every TP rank** (each rank only measures its head shard), merges them, fits `scale_factor = a * exp(b * sparsity)` once per phase, and writes the same canonical `sparse_attention_config` block the HF export produces — preserving any exported N:M sparse-softmax groups — so the serving workflow above picks it up unchanged. +`install_vllm_skip_softmax_calibration` (called by `sparse_attn_worker.SkipSoftmaxCalibWorker` at model load) swaps calibration adapters onto every attention layer after validating all of them — eager execution is required, model and KV-cache dtypes must be fp16/bf16, and no attention Q/K/P/V fakequant may be active. During `llm.generate`, the paged Triton calibration kernel computes full dense attention — no sparsification is applied to generation, though the dense kernel's numerics differ slightly from the native backend's — while counting, per candidate threshold, how many KV tiles the skip criterion would drop. The driver then collects **raw tile counts from every TP rank** (each rank only measures its head shard), merges them, fits `scale_factor = a * exp(b * sparsity)` once per phase, and writes the same canonical `sparse_attention_config` block the HF export produces — preserving any exported N:M sparse-softmax groups — so the serving workflow above picks it up unchanged. Calibration and serving measure skipping at the same fixed 128x128 tile geometry (active skip-softmax launches bypass the autotuner), so the sparsity realized at serve time matches the calibrated `(a, b)` model. @@ -212,7 +212,7 @@ report = install_vllm_nvfp4_attention(model_runner, sparse_cfg="checkpoint") Limitations: -- vLLM V1 chunked prefill and prefix-cache suffix attention are supported by offsetting query positions into the longer KV span. +- vLLM V1 chunked prefill and prefix-cache suffix attention are supported by offsetting query positions into the longer KV span. This applies to sparse-only serving; quantized attention installs and skip-softmax calibration reject `enable_prefix_caching` (quantize-on-write and per-request measurement both require uncached prefills). - `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`. ### Compact NVFP4 attention worker @@ -229,7 +229,7 @@ python vllm_serve_sparse_attn.py -tp 8 \ The installer supports both FlashInfer and FlashAttention, and the worker prints the installed adapter counts. Pass `--attention-backend FLASHINFER` or `--attention-backend FLASH_ATTN` only when an explicit override is needed. -This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant format to Q/K/P/V. Q is dynamic; missing K/V scales default to global scale 1.0, and P defaults to amax 1.0. Existing scalar attention amax values are preserved, but this path does not calibrate or restore them itself. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored. +This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant format to Q/K/P/V. Q is dynamic; missing K/V scales default to global scale 1.0, and P defaults to amax 1.0. Existing scalar attention amax values are preserved, but this path does not calibrate or restore them itself. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored for N:M sparse softmax; calibrated skip-softmax groups are rejected in combination with attention quantization, because quantized Q/K/P change the score distribution the skip thresholds were calibrated on. Decode uses a fixed 32-split, 128-key-tile schedule. P QDQ consumes split-local, unnormalized online-softmax probabilities, so changing that schedule can change diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py index 65091485981..bc5ed8281c9 100644 --- a/examples/vllm_serve/calibrate_sparse_attn.py +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -257,9 +257,11 @@ def main(): print(f"[ModelOpt] Active sparse impls: {status['impl_types']}") # generate() drives prefill (prefill-phase stats) then decode_tokens decode - # steps (decode-phase stats). The calibration kernel computes full attention, - # so the generated text is unaffected — only tile-skip counts are recorded. - sampling = SamplingParams(temperature=0.0, max_tokens=args.decode_tokens) + # steps (decode-phase stats). No sparsification is applied during + # calibration — the kernel computes full dense attention while recording + # tile-skip counts. ignore_eos forces the full decode length so early EOS + # cannot thin the decode-phase statistics. + sampling = SamplingParams(temperature=0.0, max_tokens=args.decode_tokens, ignore_eos=True) llm.generate(prompts, sampling) # Aggregate RAW counts from every TP rank (each rank only measures its @@ -268,12 +270,16 @@ def main(): merged = merge_phase_counts(rank_counts) calibration_params = fit_from_counts(merged, trials, fit_logspace=args.fit_logspace) - if not calibration_params: + requested_phases = ["prefill"] + (["decode"] if args.decode_tokens > 0 else []) + missing = [phase for phase in requested_phases if phase not in calibration_params] + if missing: print( - "[ModelOpt] Calibration produced no valid fit — try more/longer prompts " + f"[ModelOpt] Calibration FAILED: no valid fit for phase(s) {', '.join(missing)}. " + "No config was written — a partially calibrated export would silently serve " + "the missing phase dense. Try more/longer prompts (and more decode tokens) " "so observed sparsity spans the (10%, 90%) fitting window." ) - return + sys.exit(1) sparse_config = build_sparse_attention_config( calibration_params, diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 4d69a13a4cb..145de3a51e2 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -23,7 +23,6 @@ """ import math -import warnings from typing import Any import torch @@ -1042,36 +1041,27 @@ def grid(META): # (e.g. BLOCK_N=32) would realize a different sparsity than # calibrated, so skip launches always use the calibration tile. # - # If the 128-row tile exceeds the device's shared memory (fp32 - # inputs on ~100KB-smem GPUs), halve BLOCK_M and retry. BLOCK_N - # — the calibrated KV skip granularity — is never reduced. - # Smaller BLOCK_M splits the per-tile skip vote into narrower - # row groups: per-row numerics stay protected by the same - # threshold, but realized sparsity can exceed the calibrated - # target, so warn when stepping down. - block_m = _P_QDQ_MEASURE_BLOCK_M if p_qdq_mode else _MEASURE_BLOCK_M - while True: - try: - _attn_fwd.fn[grid]( - *fwd_args, - **fwd_kwargs, - BLOCK_M=block_m, - BLOCK_N=_MEASURE_BLOCK_N, - num_warps=_MEASURE_NUM_WARPS, - num_stages=_MEASURE_NUM_STAGES, - ) - break - except triton.runtime.errors.OutOfResources: - if block_m <= 16: - raise - block_m //= 2 - if apply_skip: - warnings.warn( - f"skip-softmax tile BLOCK_M reduced to {block_m} " - "(shared memory limit); realized sparsity may " - "exceed the calibrated target on this device/dtype", - stacklevel=2, - ) + # The tile is a contract, not a preference: configurations that + # cannot compile it (e.g. fp32 inputs on ~100KB-shared-memory + # GPUs) are rejected rather than re-tiled, because a different + # tile realizes a different sparsity than was calibrated. + try: + _attn_fwd.fn[grid]( + *fwd_args, + **fwd_kwargs, + BLOCK_M=_P_QDQ_MEASURE_BLOCK_M if p_qdq_mode else _MEASURE_BLOCK_M, + BLOCK_N=_MEASURE_BLOCK_N, + num_warps=_MEASURE_NUM_WARPS, + num_stages=_MEASURE_NUM_STAGES, + ) + except triton.runtime.errors.OutOfResources as err: + raise RuntimeError( + "skip-softmax requires the fixed 128x128 calibration tile, " + f"which exceeds this GPU's shared memory for {q.dtype} " + f"inputs ({err}). Use fp16/bf16 inputs or a device with " + "more shared memory; re-tiling would change the " + "calibrated sparsity contract." + ) from err else: _attn_fwd[grid]( *fwd_args, diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py index a37b447a686..5e604c039ca 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py @@ -28,6 +28,33 @@ from ..stats_manager import SparseAttentionStatsManager from ..utils import get_sparse_attention_modules +# Canonical skip-softmax threshold sweep — should span sparsities from ~10% to +# ~95%. Shared by the HF calibration path (this class's default) and the vLLM +# calibration path (``plugins/sparse_attn_calibration.py``), so both fit on the +# same trial grid. +DEFAULT_THRESHOLD_TRIALS = [ + 1e-6, + 5e-6, + 1e-5, + 5e-5, + 1e-4, + 5e-4, + 1e-3, + 5e-3, + 1e-2, + 2e-2, + 5e-2, + 1e-1, + 2e-1, + 3e-1, + 5e-1, + 7e-1, + 8e-1, + 9e-1, + 9.5e-1, + 9.9e-1, +] + class DynamicThresholdCalibrator: """Dynamic threshold calibrator using Exponential model. @@ -67,28 +94,7 @@ def __init__( where scale_factors span many orders of magnitude. """ # Default threshold trials if not provided - self.threshold_trials = threshold_trials or [ - 1e-6, - 5e-6, - 1e-5, - 5e-5, - 1e-4, - 5e-4, - 1e-3, - 5e-3, - 1e-2, - 2e-2, - 5e-2, - 1e-1, - 2e-1, - 3e-1, - 5e-1, - 7e-1, - 8e-1, - 9e-1, - 9.5e-1, - 9.9e-1, - ] + self.threshold_trials = threshold_trials or list(DEFAULT_THRESHOLD_TRIALS) self.fit_logspace = fit_logspace def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dict[str, Any]: diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py index f60bef1b867..b66a5c27436 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -35,6 +35,10 @@ import modelopt +# One canonical sweep for both calibration paths: re-exported from the HF-path +# calibrator so the vLLM path fits on the identical trial grid. +from ..calibration.calibrator import DEFAULT_THRESHOLD_TRIALS + __all__ = [ "DEFAULT_THRESHOLD_TRIALS", "build_sparse_attention_config", @@ -45,22 +49,6 @@ "stats_from_counts", ] -# Default threshold sweep — should span sparsities from ~10% to ~95%. -DEFAULT_THRESHOLD_TRIALS = [ - 1e-4, - 1e-3, - 5e-3, - 1e-2, - 3e-2, - 5e-2, - 1e-1, - 2e-1, - 3e-1, - 5e-1, - 7e-1, - 9e-1, -] - _PHASES = ("prefill", "decode") @@ -81,24 +69,40 @@ def merge_count_records(sources: list[list[dict]]) -> list[dict]: records for one phase. All sources observe the same launches in the same order, so records align by index; tile counts are additive across both layers and head-sharded TP ranks. + + Alignment is a contract: every source must report the same number of + samples, the same per-sample lengths, and the same threshold-vector width. + Any mismatch indicates a collection bug and raises rather than silently + dropping records. """ - sources = [source for source in sources if source] if not sources: return [] - num_samples = min(len(source) for source in sources) + sample_counts = {len(source) for source in sources} + if len(sample_counts) != 1: + raise ValueError( + "Misaligned calibration records: sources disagree on sample count " + f"({sorted(sample_counts)})" + ) + num_samples = sample_counts.pop() merged = [] for i in range(num_samples): base = sources[0][i] - total = [0] * len(base["total_tiles"]) - skipped = [0] * len(base["skipped_tiles"]) - for source in sources: - record = source[i] + width = len(base["total_tiles"]) + total = [0] * width + skipped = [0] * width + for record in (source[i] for source in sources): if record["sample_length"] != base["sample_length"]: raise ValueError( "Misaligned calibration records: sample lengths differ across " f"sources at index {i} ({record['sample_length']} vs " f"{base['sample_length']})" ) + if len(record["total_tiles"]) != width or len(record["skipped_tiles"]) != width: + raise ValueError( + "Misaligned calibration records: threshold-vector widths differ " + f"across sources at index {i} " + f"({len(record['total_tiles'])}/{len(record['skipped_tiles'])} vs {width})" + ) total = [a + b for a, b in zip(total, record["total_tiles"])] skipped = [a + b for a, b in zip(skipped, record["skipped_tiles"])] merged.append( @@ -118,13 +122,21 @@ def merge_phase_counts(rank_counts: list[dict[str, list[dict]]]) -> dict[str, li ``{"prefill": [...], "decode": [...]}`` dict per rank, as returned by ``collect_calibration_counts``). Use ALL ranks: with tensor parallelism each rank only measures its attention-head shard, so any single rank's - counts are incomplete. + counts are incomplete. A phase recorded by some ranks but not others + indicates a collection bug and raises. """ phases = {phase for rank in rank_counts for phase in rank} - return { - phase: merge_count_records([rank.get(phase, []) for rank in rank_counts]) - for phase in phases - } + merged: dict[str, list[dict]] = {} + for phase in phases: + sources = [rank.get(phase, []) for rank in rank_counts] + empty = sum(1 for source in sources if not source) + if empty and empty != len(sources): + raise ValueError( + f"Misaligned calibration records: {empty}/{len(sources)} rank(s) " + f"recorded no {phase!r} samples while others did" + ) + merged[phase] = merge_count_records([] if empty else sources) + return merged def stats_from_counts(count_records: list[dict]) -> list[dict]: @@ -180,7 +192,7 @@ def fit_from_counts( def _normalize_target_sparsity(target_sparsity: dict[str, float] | float) -> dict[str, float]: - if isinstance(target_sparsity, (int, float)): + if isinstance(target_sparsity, int | float): return {phase: float(target_sparsity) for phase in _PHASES} return {phase: float(target_sparsity.get(phase, 0.5)) for phase in _PHASES} @@ -230,7 +242,14 @@ def build_sparse_attention_config( for idx, group in enumerate(preserved, start=1): config_groups[f"group_{idx}"] = group - return { + result: dict[str, Any] = { "config_groups": config_groups, "producer": {"name": "modelopt", "version": modelopt.__version__}, } + # Legacy checkpoints carry N:M parameters as a top-level ``sparse_softmax`` + # dict (read by the serving loader ahead of group params) — preserve it so + # recalibration does not silently reset N:M settings to defaults. + legacy_sparse_softmax = (existing_config or {}).get("sparse_softmax") + if isinstance(legacy_sparse_softmax, dict): + result["sparse_softmax"] = legacy_sparse_softmax + return result diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 80e30e5fc18..2ebf2aa39ac 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -28,6 +28,7 @@ """ import functools +import importlib import inspect import math import warnings @@ -260,6 +261,32 @@ def _calibration_active(impl) -> bool: ) +def _flashinfer_kv_cache_layout() -> str | None: + """Best-effort query of vLLM's configured FlashInfer KV-cache layout. + + Returns ``"NHD"`` / ``"HND"`` when the running vLLM exposes the layout + (newer releases select HND for some Blackwell FlashInfer paths), or + ``None`` when it cannot be determined — callers then fall back to the + shape-based check in :func:`_forward_calibrate`. + """ + for module_name in ( + "vllm.v1.attention.backends.utils", + "vllm.attention.backends.utils", + ): + try: + module = importlib.import_module(module_name) + except ImportError: + continue + getter = getattr(module, "get_kv_cache_layout", None) + if getter is None: + continue + try: + return str(getter()) + except Exception: + return None + return None + + def _forward_calibrate( impl, *, @@ -277,8 +304,9 @@ def _forward_calibrate( Each scheduled request is calibrated independently (batch=1) so its KV length is the per-sample length the exponential fit needs, and so the kernel keeps the uniform-length contract it was validated against. The - kernel computes full attention, so ``output`` is written densely and the - forward pass is numerically unchanged. + kernel computes full attention, so ``output`` is written densely — no + sparsification is applied to generation (the dense Triton kernel's + numerics may differ slightly from the native backend's). Phase and causality are decided per request: ``q_len == 1`` is a decode step (full-cache, non-causal); ``q_len > 1`` is (chunked) prefill (causal @@ -839,6 +867,14 @@ def prepare_modelopt(): raise ValueError( "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" ) + layout = _flashinfer_kv_cache_layout() + if layout is not None and layout.upper() != "NHD": + # Authoritative layout metadata beats the shape heuristic (which is + # ambiguous when page_size equals the per-rank KV-head count). + raise NotImplementedError( + f"FlashInfer KV-cache layout {layout!r} is unsupported for " + "skip-softmax calibration; only NHD is supported" + ) # Order matters: releases that update the KV cache inside forward must # write the current K/V before the calibrate kernel reads the cache. prepare_modelopt() @@ -1031,10 +1067,21 @@ def collect_calibration_counts(model) -> dict[str, list[dict]]: """ from .sparse_attn_calibration import merge_count_records, split_records_by_phase - per_phase_layers: dict[str, list[list[dict]]] = {} - for impl in iter_sparse_impls(model): - split = split_records_by_phase(getattr(impl, "_calib_records", [])) - for phase, records in split.items(): - if records: - per_phase_layers.setdefault(phase, []).append(records) - return {phase: merge_count_records(layers) for phase, layers in per_phase_layers.items()} + splits = [ + split_records_by_phase(getattr(impl, "_calib_records", [])) + for impl in iter_sparse_impls(model) + ] + phases = {phase for split in splits for phase, records in split.items() if records} + counts: dict[str, list[dict]] = {} + for phase in phases: + sources = [split.get(phase, []) for split in splits] + empty = sum(1 for source in sources if not source) + if empty: + # Every layer sees every launch, so a layer with no records for a + # phase others measured indicates a collection bug. + raise ValueError( + f"Misaligned calibration records: {empty}/{len(sources)} attention " + f"layer(s) recorded no {phase!r} samples while others did" + ) + counts[phase] = merge_count_records(sources) + return counts diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index 04d703f46a1..a00a7ca012e 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -249,6 +249,13 @@ def _device_capability_error(device) -> str | None: return None +def _skip_softmax_active(sparse_kw: dict[str, Any]) -> bool: + """Return whether a layer's sparse config contains skip-softmax work.""" + return bool(sparse_kw) and ( + "skip_softmax_threshold" in sparse_kw or "threshold_scale_factor" in sparse_kw + ) + + def _sparse_graph_error(sparse_kw: dict[str, Any], mode) -> str | None: from vllm.config.compilation import CUDAGraphMode @@ -354,6 +361,16 @@ def _plan_vllm_attention( plans = [] for name, module, sparse_kw in candidates: reasons = _layer_errors(module) + if quantize and _skip_softmax_active(sparse_kw): + # Quantized Q/K/P change the attention-score distribution the skip + # thresholds were calibrated on, so the calibrated sparsity contract + # no longer holds. N:M sparse softmax has no calibrated threshold + # and composes with quantization. + reasons.append( + "skip-softmax cannot be combined with attention quantization; " + "serve skip-softmax unquantized or drop the skip_softmax group " + "(N:M sparse softmax composes with quantization)" + ) device = dtype = None if quantize: device, dtype = quant_plugin._get_device_dtype(module) @@ -561,6 +578,12 @@ def install_vllm_skip_softmax_calibration(model_runner) -> VllmAttentionInstallR plans.append( _AttentionPlan(name, module, new_impl, {}, None, None, requires_flashinfer_patch) ) + if any(plan.requires_flashinfer_patch for plan in plans): + layout = attention_plugin._flashinfer_kv_cache_layout() + if layout is not None and layout.upper() != "NHD": + errors.append( + f"FlashInfer KV-cache layout {layout!r} is unsupported for calibration (NHD only)" + ) _raise_unsupported(errors, "skip-softmax calibration") plan = _InstallPlan(model_runner, tuple(plans), False, "SKIP_SOFTMAX_CALIBRATION") diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py index 9e55879dbb6..7020f90d7c5 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py @@ -116,27 +116,27 @@ def test_high_block_id_pointer_arithmetic(self): """Block IDs whose int32 byte offsets would wrap still read correctly.""" num_kv_heads, head_dim, page_size = 2, 64, 16 block_elems = page_size * num_kv_heads * head_dim - # Smallest block ID whose element offset exceeds int32. + # Smallest block ID whose element offset exceeds int32. V aliases the K + # cache storage (same values on both operands), halving the allocation. high_block = (2**31) // block_elems + 1 - bytes_needed = 2 * (high_block + 1) * block_elems * 2 # K and V caches, bf16 + bytes_needed = (high_block + 1) * block_elems * 2 # one shared K/V cache, bf16 free, _ = torch.cuda.mem_get_info() if free < bytes_needed + (2 << 30): pytest.skip(f"needs ~{bytes_needed / 2**30:.1f} GiB free GPU memory") torch.manual_seed(2) num_heads = 4 - q, k, v = make_qkv(page_size, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + q, k, _ = make_qkv(page_size, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) k_cache = torch.zeros( high_block + 1, page_size, num_kv_heads, head_dim, device="cuda", dtype=torch.bfloat16 ) - v_cache = torch.zeros_like(k_cache) + v_cache = k_cache # alias: V reads the same storage (and the same values) k_cache[high_block] = k - v_cache[high_block] = v block_table = torch.tensor([[high_block]], device="cuda", dtype=torch.int32) locs, lens = make_varlen_meta([page_size]) out_ref, counters_ref = attention_calibrate( - q, k, v, locs, lens, page_size, is_causal=True, threshold_trials=TRIALS + q, k, k, locs, lens, page_size, is_causal=True, threshold_trials=TRIALS ) k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) out_paged, counters_paged = attention_calibrate( diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py index 2806e39b4a5..598f97a3fda 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py @@ -356,11 +356,16 @@ class TestBackwardWithSparsity: """Backward pass with skip-softmax (covers _attn_bwd_dq / _attn_bwd_dkdv).""" def test_backward_with_skip_softmax(self): - """Backward pass runs without error when skip-softmax is active.""" + """Backward pass runs without error when skip-softmax is active. + + fp16 rather than fp32: active skip launches require the fixed 128x128 + calibration tile, which fp32 inputs cannot compile on ~100KB-shared- + memory GPUs (such configurations are rejected by design). + """ seq_len, num_heads, head_dim = 128, 4, 64 scale = 1.0 / (head_dim**0.5) torch.manual_seed(7) - q, k, v = make_qkv(seq_len, num_heads, num_heads, head_dim, dtype=torch.float32) + q, k, v = make_qkv(seq_len, num_heads, num_heads, head_dim, dtype=torch.float16) q.requires_grad_(True) k.requires_grad_(True) v.requires_grad_(True) diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py index caefb467555..033952d34eb 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py @@ -217,17 +217,26 @@ def test_triton_matches_pytorch_reference(self): locs = torch.arange(batch, device="cuda", dtype=torch.int32) * seq_len lens = torch.full((batch,), seq_len, device="cuda", dtype=torch.int32) - triton_out = attention( - q_flat, - k_flat, - v_flat, - locs, - lens, - seq_len, - is_causal=True, - softmax_scale=scale, - skip_softmax_threshold=threshold, - ) + try: + triton_out = attention( + q_flat, + k_flat, + v_flat, + locs, + lens, + seq_len, + is_causal=True, + softmax_scale=scale, + skip_softmax_threshold=threshold, + ) + except RuntimeError as err: + if "shared memory" in str(err): + # Active skip requires the fixed 128x128 calibration tile; fp32 + # inputs cannot compile it on ~100KB-shared-memory GPUs and the + # configuration is rejected by design. fp32 is kept here for a + # tight reference comparison on GPUs that support it. + pytest.skip("fp32 skip tile exceeds this GPU's shared memory") + raise triton_out_4d = triton_out.view(batch, seq_len, num_heads, head_dim).permute(0, 2, 1, 3) # Both outputs should be close — same algorithm, different implementations. diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py index e5192e7e7f3..62a75b92d8f 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -22,6 +22,7 @@ from torch import nn from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.attention.backends.flashinfer import FlashInferImpl from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as attention_plugin @@ -126,6 +127,106 @@ def test_rejects_model_without_attention_layers(self): vllm_runtime.install_vllm_skip_softmax_calibration(runner) +class TestQuantSkipRejection: + """Skip-softmax cannot be combined with attention quantization.""" + + _CALIBRATED_META = { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "threshold_scale_factor": {"prefill": {"a": 7.9, "b": 8.6}}, + } + } + } + _NM_META = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4} + } + } + + def test_quantized_install_rejects_calibrated_skip(self): + attention = _bare_attention() + runner = _model_runner( + nn.ModuleDict({"attn": attention}), sparse_metadata=self._CALIBRATED_META + ) + with pytest.raises( + NotImplementedError, match="cannot be combined with attention quantization" + ): + vllm_runtime.install_vllm_nvfp4_attention(runner) + assert not isinstance(attention.impl, ModelOptSparseAttentionImpl) + + def test_quantized_plan_allows_nm_sparsity(self, monkeypatch): + from modelopt.torch.quantization.plugins import vllm as quant_plugin + + monkeypatch.setattr( + quant_plugin, + "_get_device_dtype", + lambda module: (torch.device("cpu"), torch.float16), + ) + runner = _model_runner( + nn.ModuleDict({"attn": _bare_attention()}), sparse_metadata=self._NM_META + ) + plan = vllm_runtime._plan_vllm_attention(runner, quantize=True, sparse_cfg="checkpoint") + assert len(plan.layers) == 1 + assert plan.layers[0].sparse_kw.get("sparsity_n") == 2 + + +class TestFlashInferLayoutGuard: + """HND FlashInfer caches are rejected via layout metadata, pre-measurement.""" + + def test_installer_rejects_hnd_layout(self, monkeypatch): + monkeypatch.setattr(attention_plugin, "_flashinfer_kv_cache_layout", lambda: "HND") + attention = _bare_attention(FlashInferImpl) + original_impl = attention.impl + runner = _model_runner(nn.ModuleDict({"attn": attention})) + with pytest.raises(NotImplementedError, match="HND"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + assert attention.impl is original_impl + + def test_forward_rejects_hnd_layout_before_cache_write(self, monkeypatch): + monkeypatch.setattr(attention_plugin, "_flashinfer_kv_cache_layout", lambda: "HND") + writes = [] + monkeypatch.setattr( + attention_plugin, + "_maybe_update_flashinfer_cache", + lambda *args, **kwargs: writes.append(1), + ) + num_heads, num_kv_heads, head_dim, page = 4, 2, 64, 16 + impl = SimpleNamespace( + num_kv_heads=num_kv_heads, + head_size=head_dim, + scale=1.0 / (head_dim**0.5), + _calibrate=True, + _calib_threshold_trials=list(TRIALS), + _calib_records=[], + ) + kv_cache = torch.zeros(3, 2, page, num_kv_heads, head_dim, dtype=torch.bfloat16) + attn_metadata = SimpleNamespace( + _modelopt_block_table=torch.zeros(1, 1, dtype=torch.int32), + _modelopt_seq_lens=torch.tensor([8], dtype=torch.int32), + _modelopt_query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + _modelopt_num_actual_tokens=1, + _modelopt_max_query_len=1, + _modelopt_max_seq_len=8, + _modelopt_causal=False, + slot_mapping=torch.zeros(1, dtype=torch.int64), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + with pytest.raises(NotImplementedError, match="HND"): + attention_plugin._flashinfer_forward( + impl, + None, + None, + q, + q[:, :num_kv_heads], + q[:, :num_kv_heads], + kv_cache, + attn_metadata, + output=torch.empty_like(q), + ) + assert not writes, "layout must be validated before the cache write" + + class TestSparseOnlyGraphGuard: """Commit-contract: the calibrated-decode graph guard is not quantize-gated.""" 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 77f836c9205..85e459dd650 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 @@ -63,12 +63,11 @@ def test_merge_sums_counts_elementwise(self): merged = merge_count_records([layer_a, layer_b]) assert merged == [{"sample_length": 128, "total_tiles": [20, 20], "skipped_tiles": [3, 7]}] - def test_merge_truncates_to_shortest_source(self): + def test_merge_rejects_ragged_sources(self): long = [_record("prefill", 1, [1], [0]), _record("prefill", 2, [1], [1])] short = [_record("prefill", 1, [1], [1])] - merged = merge_count_records([long, short]) - assert len(merged) == 1 - assert merged[0]["skipped_tiles"] == [1] + with pytest.raises(ValueError, match="disagree on sample count"): + merge_count_records([long, short]) def test_merge_rejects_misaligned_sample_lengths(self): with pytest.raises(ValueError, match="Misaligned calibration records"): @@ -76,6 +75,18 @@ def test_merge_rejects_misaligned_sample_lengths(self): [[_record("prefill", 100, [1], [0])], [_record("prefill", 200, [1], [0])]] ) + def test_merge_rejects_threshold_width_mismatch(self): + with pytest.raises(ValueError, match="threshold-vector widths"): + merge_count_records( + [[_record("prefill", 100, [1, 2], [0, 1])], [_record("prefill", 100, [1], [0])]] + ) + + def test_merge_phase_counts_rejects_rank_phase_mismatch(self): + rank0 = {"prefill": [_record("prefill", 64, [5], [1])]} + rank1 = {"prefill": []} + with pytest.raises(ValueError, match="recorded no 'prefill' samples"): + merge_phase_counts([rank0, rank1]) + def test_merge_phase_counts_across_ranks(self): rank0 = {"prefill": [_record("prefill", 64, [5], [1])], "decode": []} rank1 = {"prefill": [_record("prefill", 64, [5], [2])]} @@ -187,6 +198,22 @@ def test_round_trips_through_serving_loader(self): assert layer_cfg["threshold_scale_factor"]["decode"] == {"a": 0.12, "b": 9.8} assert layer_cfg["target_sparse_ratio"] == {"prefill": 0.5, "decode": 0.3} + def test_preserves_legacy_toplevel_sparse_softmax(self): + existing = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4} + }, + "sparse_softmax": {"sparsity_n": 1, "sparsity_m": 4, "dense_recent_tokens": 128}, + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + # The serving loader reads the legacy top-level key ahead of group params. + assert config["sparse_softmax"] == existing["sparse_softmax"] + loaded = load_from_checkpoint_metadata(SimpleNamespace(sparse_attention_config=config)) + assert loaded is not None + layer_cfg = loaded[0]["sparse_cfg"]["*attn*"] + assert layer_cfg["sparsity_n"] == 1 + assert layer_cfg["dense_recent_tokens"] == 128 + def test_round_trip_with_preserved_nm_group_activates_both(self): existing = { "config_groups": { From 907af6f0aa15a327287e0c61662a1c96d467a41a Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sun, 19 Jul 2026 15:42:38 -0700 Subject: [PATCH 08/10] Address review round 2: close remaining skip+quant routes and edge cases - Skip-softmax + attention quantization is now unreachable from every direction: sparse-only installs reject calibrated skip onto layers with active attention quantizers (not just quantized installs adding skip), and the raw kernel API itself rejects P/V QDQ with an active skip threshold -- which makes the P-QDQ 16-row measurement tile dead code, so the fixed skip tile is unconditionally 128x128. - FlashInfer layout helper preserves a genuine None from get_kv_cache_layout (str(None) became the truthy string None and both guards hard-rejected valid configurations instead of using the shape fallback). - Counter vectors are validated against len(threshold_trials) in both calibrate_from_stats and fit_from_counts: consistently short vectors previously zipped silently and could misattribute sparsities. - Driver decode semantics match vLLM: the first output token comes from the prefill forward, so --decode_tokens now means decode-attention steps and generation runs decode_tokens + 1 output tokens. - Calibrate kernel mirrors the serving kernel IEEE fp32 QK dot so raw fp32 calibration and serving round near-threshold scores identically. - target_sparsity validated to [0, 1] (same range as the HF config). Signed-off-by: Kai Xu --- examples/vllm_serve/calibrate_sparse_attn.py | 18 +++++---- .../kernels/common/attention/triton_fa.py | 15 ++++++- .../kernels/sparsity/attention/calibrate.py | 10 ++++- .../calibration/calibrator.py | 6 +++ .../plugins/sparse_attn_calibration.py | 18 ++++++++- .../attention_sparsity/plugins/vllm.py | 5 ++- .../plugins/vllm_runtime.py | 40 ++++++++++++------- .../attention/test_paged_calibrate.py | 34 +++++++--------- .../test_vllm_calibration.py | 32 +++++++++++++++ .../test_sparse_attn_calibration.py | 15 +++++++ 10 files changed, 147 insertions(+), 46 deletions(-) diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py index bc5ed8281c9..28f6a7eed9e 100644 --- a/examples/vllm_serve/calibrate_sparse_attn.py +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -161,7 +161,9 @@ def main(): "--decode_tokens", type=int, default=32, - help="Decode tokens to generate per prompt (drives decode-phase calibration)", + help="Decode attention steps per prompt (drives decode-phase calibration). " + "Generation runs decode_tokens + 1 output tokens: the first output token " + "comes from the prefill forward and performs no decode attention.", ) parser.add_argument( "--max_model_len", type=int, default=None, help="vLLM max_model_len override" @@ -256,12 +258,14 @@ def main(): print(f"[ModelOpt] Calibration enabled on {n_layers} attention layers") print(f"[ModelOpt] Active sparse impls: {status['impl_types']}") - # generate() drives prefill (prefill-phase stats) then decode_tokens decode - # steps (decode-phase stats). No sparsification is applied during - # calibration — the kernel computes full dense attention while recording - # tile-skip counts. ignore_eos forces the full decode length so early EOS - # cannot thin the decode-phase statistics. - sampling = SamplingParams(temperature=0.0, max_tokens=args.decode_tokens, ignore_eos=True) + # generate() drives prefill (prefill-phase stats) then decode steps + # (decode-phase stats). No sparsification is applied during calibration — + # the kernel computes full dense attention while recording tile-skip + # counts. ignore_eos forces the full decode length so early EOS cannot + # thin the decode-phase statistics. max_tokens is decode_tokens + 1: the + # first output token comes from the prefill forward, so decode_tokens + # decode-attention steps need one extra output token. + sampling = SamplingParams(temperature=0.0, max_tokens=args.decode_tokens + 1, ignore_eos=True) llm.generate(prompts, sampling) # Aggregate RAW counts from every TP rank (each rank only measures its diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 145de3a51e2..4a223d4c2ab 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -87,7 +87,6 @@ def _load_qdq_helpers() -> None: ] _MEASURE_BLOCK_M = 128 -_P_QDQ_MEASURE_BLOCK_M = 16 # 128 so the kernel sparsity-measurement block matches the PyTorch # calibration/reference granularity. This is deliberately independent of the # autotuned compute tile. @@ -944,6 +943,16 @@ def forward( else: apply_skip = False skip_threshold_log2 = 0.0 + if apply_skip and (p_qdq_mode or v_qdq_mode): + # Quantized operands change what the calibrated skip thresholds mean, + # and P-QDQ additionally uses a different measurement tile geometry. + # The vLLM installers reject this composition at plan time; the raw + # kernel API rejects it here so no path can serve it. + raise ValueError( + "skip-softmax cannot be combined with attention quantization " + "(P/V QDQ): the calibrated tile-skip contract does not hold " + "under quantized operands" + ) o = torch.empty_like(q) lse = torch.empty(q.shape[0], num_q_heads, device=q.device, dtype=torch.float32) @@ -1046,10 +1055,12 @@ def grid(META): # GPUs) are rejected rather than re-tiled, because a different # tile realizes a different sparsity than was calibrated. try: + # P/V QDQ is rejected above when skip is active, so the tile + # here is unconditionally the 128x128 calibration geometry. _attn_fwd.fn[grid]( *fwd_args, **fwd_kwargs, - BLOCK_M=_P_QDQ_MEASURE_BLOCK_M if p_qdq_mode else _MEASURE_BLOCK_M, + BLOCK_M=_MEASURE_BLOCK_M, BLOCK_N=_MEASURE_BLOCK_N, num_warps=_MEASURE_NUM_WARPS, num_stages=_MEASURE_NUM_STAGES, diff --git a/modelopt/torch/kernels/sparsity/attention/calibrate.py b/modelopt/torch/kernels/sparsity/attention/calibrate.py index 68748e7c50c..b8e9d6475c7 100644 --- a/modelopt/torch/kernels/sparsity/attention/calibrate.py +++ b/modelopt/torch/kernels/sparsity/attention/calibrate.py @@ -69,6 +69,7 @@ def _attn_fwd_calibrate( HEAD_DIM: tl.constexpr, NUM_THRESHOLDS: tl.constexpr, PADDED_THRESHOLDS: tl.constexpr, # next_power_of_2(NUM_THRESHOLDS) for tl.arange + Q_IS_FP32: tl.constexpr = False, # match the serving kernel's IEEE fp32 QK dot IS_PAGED: tl.constexpr = False, # Whether K/V are read from a paged KV cache K_cache=None, # [num_blocks, page_size, num_kv_heads, head_dim] paged K V_cache=None, # [num_blocks, page_size, num_kv_heads, head_dim] paged V @@ -171,7 +172,13 @@ def _attn_fwd_calibrate( other=0.0, ) - scores = tl.dot(q, k) * qk_scale + # Match the serving kernel's QK precision: fp32 Q uses the IEEE dot + # (default tl.dot is TF32 for fp32 inputs), so near-threshold scores + # round to the same skip decisions in calibration and serving. + if Q_IS_FP32: + scores = tl.dot(q, k.to(tl.float32), input_precision="ieee") * qk_scale + else: + scores = tl.dot(q, k) * qk_scale scores = _apply_mask(scores, q_pos, kv_pos, seq_len_q, seq_len_kv, kv_start, IS_CAUSAL) tile_row_max = tl.max(scores, 1) @@ -430,6 +437,7 @@ def attention_calibrate( HEAD_DIM=HEAD_DIM, NUM_THRESHOLDS=num_thresholds, PADDED_THRESHOLDS=triton.next_power_of_2(num_thresholds), + Q_IS_FP32=q.dtype == torch.float32, IS_PAGED=is_paged, K_cache=k_cache, V_cache=v_cache, diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py index 5e604c039ca..ef46758ead0 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py @@ -169,6 +169,12 @@ def calibrate_from_stats(self, per_sample_stats: list[dict], phase: str) -> dict for sample_stat in per_sample_stats: length = sample_stat["sample_length"] sparsity_list = sample_stat["sparsity"] + if len(sparsity_list) != len(self.threshold_trials): + # A silent zip would misattribute sparsities to thresholds. + raise ValueError( + f"per-sample sparsity has {len(sparsity_list)} entries but " + f"{len(self.threshold_trials)} threshold trials are configured" + ) for threshold, sparsity in zip(self.threshold_trials, sparsity_list): scale_factor = threshold * length all_data_points.append( diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py index b66a5c27436..32fb5d67160 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -178,6 +178,12 @@ def fit_from_counts( for phase, records in per_phase_counts.items(): if not records: continue + for record in records: + if len(record["total_tiles"]) != len(threshold_trials): + raise ValueError( + f"{phase} record has {len(record['total_tiles'])} counters but " + f"{len(threshold_trials)} threshold trials are configured" + ) calibrator = DynamicThresholdCalibrator( threshold_trials=list(threshold_trials), fit_logspace=fit_logspace ) @@ -193,8 +199,16 @@ def fit_from_counts( def _normalize_target_sparsity(target_sparsity: dict[str, float] | float) -> dict[str, float]: if isinstance(target_sparsity, int | float): - return {phase: float(target_sparsity) for phase in _PHASES} - return {phase: float(target_sparsity.get(phase, 0.5)) for phase in _PHASES} + values = {phase: float(target_sparsity) for phase in _PHASES} + else: + values = {phase: float(target_sparsity.get(phase, 0.5)) for phase in _PHASES} + for phase, value in values.items(): + # Same range the HF calibration config enforces. + if not 0.0 <= value <= 1.0: + raise ValueError( + f"target_sparsity for phase {phase!r} must be between 0.0 and 1.0, got {value}" + ) + return values def build_sparse_attention_config( diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 2ebf2aa39ac..8bede9ed665 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -281,9 +281,12 @@ def _flashinfer_kv_cache_layout() -> str | None: if getter is None: continue try: - return str(getter()) + value = getter() except Exception: return None + # Preserve a genuine None (layout unset) so the shape fallback runs; + # str(None) would become the truthy string "None" and hard-reject. + return None if value is None else str(value) return None diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index a00a7ca012e..724feb69f90 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -361,16 +361,24 @@ def _plan_vllm_attention( plans = [] for name, module, sparse_kw in candidates: reasons = _layer_errors(module) - if quantize and _skip_softmax_active(sparse_kw): + if _skip_softmax_active(sparse_kw): # Quantized Q/K/P change the attention-score distribution the skip # thresholds were calibrated on, so the calibrated sparsity contract - # no longer holds. N:M sparse softmax has no calibrated threshold - # and composes with quantization. - reasons.append( - "skip-softmax cannot be combined with attention quantization; " - "serve skip-softmax unquantized or drop the skip_softmax group " - "(N:M sparse softmax composes with quantization)" + # no longer holds. This guards both installation directions: quantized + # installs adding skip, and sparse-only installs onto layers that + # already carry active attention quantizers. N:M sparse softmax has + # no calibrated threshold and composes with quantization. + active = ( + "attention quantization is being installed" + if quantize + else _active_attention_quantization(module) ) + if active: + reasons.append( + f"skip-softmax cannot be combined with attention quantization ({active}); " + "serve skip-softmax unquantized or drop the skip_softmax group " + "(N:M sparse softmax composes with quantization)" + ) device = dtype = None if quantize: device, dtype = quant_plugin._get_device_dtype(module) @@ -513,18 +521,22 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor return _build_report(plan) -def _attention_quant_error(module) -> str | None: - """Reject calibration on layers with any active attention Q/K/P/V fakequant.""" +def _active_attention_quantization(module) -> str | None: + """Describe any active attention Q/K/P/V quantization on a layer, or None.""" for attr in ("q_bmm_quantizer", "k_bmm_quantizer", "p_bmm_quantizer", "v_bmm_quantizer"): if getattr(getattr(module, attr, None), "is_enabled", False): - return f"{attr} is enabled; skip-softmax calibration requires unquantized attention" + return f"{attr} is enabled" if getattr(module, "_query_quant_in_kernel", False) or getattr( module, "_value_quant_in_kernel", False ): - return ( - "in-kernel attention quantization flags are set; skip-softmax " - "calibration requires unquantized attention" - ) + return "in-kernel attention quantization flags are set" + return None + + +def _attention_quant_error(module) -> str | None: + """Reject calibration on layers with any active attention Q/K/P/V fakequant.""" + if active := _active_attention_quantization(module): + return f"{active}; skip-softmax calibration requires unquantized attention" return None diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py index 7020f90d7c5..c7390382098 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py @@ -211,26 +211,22 @@ def test_serve_skip_counts_equal_calibrate_counts(self, threshold): assert out._sparsity_total == calib_total assert out._sparsity_skipped == calib_skipped - def test_skip_composes_with_pv_qdq(self): - """Active skip at the fixed tile compiles and runs with P/V QDQ enabled.""" - if torch.cuda.get_device_capability() < (8, 9): - pytest.skip("NVFP4/FP8 QDQ needs compute capability >= 8.9 (E4M3 casts)") + @pytest.mark.parametrize("qdq_kw", [{"p_qdq": "nvfp4"}, {"v_qdq": "nvfp4", "v_qdq_amax": 1.0}]) + def test_skip_rejects_pv_qdq(self, qdq_kw): + """Active skip rejects P/V QDQ: quantized operands break the calibrated contract.""" seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 q, k, v = self._contrasty_qkv(seq_len, num_heads, num_kv_heads, head_dim) locs, lens = make_varlen_meta([seq_len]) - out = attention( - q, - k, - v, - locs, - lens, - seq_len, - is_causal=True, - skip_softmax_threshold=1e-2, - p_qdq="nvfp4", - p_qdq_amax=1.0, - v_qdq="nvfp4", - v_qdq_amax=float(v.abs().amax()), - ) - assert torch.isfinite(out).all() + with pytest.raises(ValueError, match="cannot be combined with attention quantization"): + attention( + q, + k, + v, + locs, + lens, + seq_len, + is_causal=True, + skip_softmax_threshold=1e-2, + **qdq_kw, + ) diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py index 62a75b92d8f..543fbd0a2da 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -155,6 +155,27 @@ def test_quantized_install_rejects_calibrated_skip(self): vllm_runtime.install_vllm_nvfp4_attention(runner) assert not isinstance(attention.impl, ModelOptSparseAttentionImpl) + def test_sparse_only_install_rejects_skip_onto_quantized_layer(self): + """Sparse-only installs must also refuse skip onto live quantizers.""" + attention = _bare_attention() + attention.q_bmm_quantizer = SimpleNamespace(is_enabled=True) + original_impl = attention.impl + runner = _model_runner( + nn.ModuleDict({"attn": attention}), sparse_metadata=self._CALIBRATED_META + ) + with pytest.raises( + NotImplementedError, match="cannot be combined with attention quantization" + ): + vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert attention.impl is original_impl + + def test_sparse_only_install_allows_skip_on_unquantized_layer(self): + runner = _model_runner( + nn.ModuleDict({"attn": _bare_attention()}), sparse_metadata=self._CALIBRATED_META + ) + report = vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert report.installed_count == 1 + def test_quantized_plan_allows_nm_sparsity(self, monkeypatch): from modelopt.torch.quantization.plugins import vllm as quant_plugin @@ -174,6 +195,17 @@ def test_quantized_plan_allows_nm_sparsity(self, monkeypatch): class TestFlashInferLayoutGuard: """HND FlashInfer caches are rejected via layout metadata, pre-measurement.""" + def test_layout_helper_preserves_none(self, monkeypatch): + """A getter returning None must not become the truthy string 'None'.""" + import sys + + fake = SimpleNamespace(get_kv_cache_layout=lambda: None) + monkeypatch.setitem(sys.modules, "vllm.v1.attention.backends.utils", fake) + assert attention_plugin._flashinfer_kv_cache_layout() is None + + fake.get_kv_cache_layout = lambda: "HND" + assert attention_plugin._flashinfer_kv_cache_layout() == "HND" + def test_installer_rejects_hnd_layout(self, monkeypatch): monkeypatch.setattr(attention_plugin, "_flashinfer_kv_cache_layout", lambda: "HND") attention = _bare_attention(FlashInferImpl) 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 85e459dd650..e85ad3ab77e 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 @@ -127,6 +127,15 @@ def counts(length, total): def test_empty_phase_produces_no_fit(self): assert fit_from_counts({"decode": []}, DEFAULT_THRESHOLD_TRIALS) == {} + def test_fit_rejects_counter_width_vs_trials_mismatch(self): + """Consistent-but-wrong widths must not silently zip against the trials.""" + short = len(DEFAULT_THRESHOLD_TRIALS) - 1 + records = [ + {"sample_length": 4096, "total_tiles": [100] * short, "skipped_tiles": [50] * short} + ] + with pytest.raises(ValueError, match="threshold trials are configured"): + fit_from_counts({"prefill": records}, DEFAULT_THRESHOLD_TRIALS) + class TestCalibrateFromStats: def _stats(self, trials): @@ -198,6 +207,12 @@ def test_round_trips_through_serving_loader(self): assert layer_cfg["threshold_scale_factor"]["decode"] == {"a": 0.12, "b": 9.8} assert layer_cfg["target_sparse_ratio"] == {"prefill": 0.5, "decode": 0.3} + def test_rejects_out_of_range_target_sparsity(self): + with pytest.raises(ValueError, match=r"between 0\.0 and 1\.0"): + build_sparse_attention_config(self._PARAMS, 1.5) + with pytest.raises(ValueError, match=r"between 0\.0 and 1\.0"): + build_sparse_attention_config(self._PARAMS, {"prefill": 0.5, "decode": -0.1}) + def test_preserves_legacy_toplevel_sparse_softmax(self): existing = { "config_groups": { From f361582cd68e20492f9d7f21c6cb237e53824927 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 24 Aug 2026 21:12:38 -0700 Subject: [PATCH 09/10] Address skip-softmax calibration review findings Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 +- examples/vllm_serve/README.md | 3 +- examples/vllm_serve/calibrate_sparse_attn.py | 33 ++++- .../kernels/sparsity/attention/calibrate.py | 38 +++-- modelopt/torch/quantization/tensor_quant.py | 5 +- .../sparsity/attention_sparsity/conversion.py | 39 +++-- .../plugins/sparse_attn_calibration.py | 67 +++++---- .../attention_sparsity/plugins/vllm.py | 135 +++++++++++++----- .../plugins/vllm_runtime.py | 33 ++++- .../common/attention/test_triton_fa_p_qdq.py | 36 ++--- .../quantization/test_tensor_quant_cuda.py | 33 +++++ .../test_vllm_calibration.py | 6 +- .../attention_sparsity/test_vllm_runtime.py | 3 - .../common/attention/test_triton_fa.py | 52 ++++--- .../test_sparse_attn_calibration.py | 25 ++++ 15 files changed, 372 insertions(+), 138 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6d9d2152752..a2019000088 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog *Sparsity* -- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. Configurations that cannot compile the 128x128 tile are rejected rather than re-tiled, and skip-softmax cannot be combined with attention quantization. See ``examples/vllm_serve/calibrate_sparse_attn.py``. +- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. Configurations that cannot compile the 128x128 tile are rejected rather than re-tiled, and skip-softmax cannot be combined with attention quantization. **Behavior change for sparse-only serving installs:** ``install_vllm_sparse_attention_from_checkpoint`` now runs the same engine-level compatibility validation as quantized installs (decode context parallelism, DBO, speculative decoding, and FULL mixed-batch CUDA graphs are rejected; prefix caching remains supported), and checkpoints with a calibrated ``decode`` ``threshold_scale_factor`` are rejected under a FULL decode CUDA graph mode — including vLLM's default ``FULL_AND_PIECEWISE`` — because the captured graph would replay one request's stale threshold. Previously such deployments started and silently served miscalibrated sparsity; pass ``--enforce-eager`` (or select a non-FULL decode graph mode). See ``examples/vllm_serve/calibrate_sparse_attn.py``. *Quantization* diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 34a802c85e2..c5596003cd3 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -213,7 +213,8 @@ report = install_vllm_nvfp4_attention(model_runner, sparse_cfg="checkpoint") Limitations: - vLLM V1 chunked prefill and prefix-cache suffix attention are supported by offsetting query positions into the longer KV span. This applies to sparse-only serving; quantized attention installs and skip-softmax calibration reject `enable_prefix_caching` (quantize-on-write and per-request measurement both require uncached prefills). -- `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`. +- `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`. Checkpoints with a calibrated `decode` `threshold_scale_factor` are rejected at install under a FULL decode CUDA graph mode (including vLLM's default `FULL_AND_PIECEWISE`): the captured graph would replay one request's stale threshold. +- Sparse-only installs validate engine-level compatibility like quantized installs do: decode context parallelism, DBO, speculative decoding, and FULL mixed-batch CUDA graphs are rejected (prefix caching remains supported, per the bullet above). ### Compact NVFP4 attention worker diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py index 28f6a7eed9e..42c4d36a972 100644 --- a/examples/vllm_serve/calibrate_sparse_attn.py +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -50,6 +50,7 @@ import sys from pathlib import Path +from modelopt.torch.sparsity.attention_sparsity.calibration.ruler_dataset import RulerDatasetBuilder from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_calibration import ( DEFAULT_THRESHOLD_TRIALS, build_sparse_attention_config, @@ -71,10 +72,6 @@ def _load_prompts(llm, args) -> list[str]: # Same dataset as the HF calibration path (calibration/calibrate.py), so the # vLLM- and PyTorch-calibrated thresholds are fit on identical data. - from modelopt.torch.sparsity.attention_sparsity.calibration.ruler_dataset import ( - RulerDatasetBuilder, - ) - builder = RulerDatasetBuilder( samples=args.calib_samples, max_seqlen=args.calib_max_seqlen, @@ -83,6 +80,12 @@ def _load_prompts(llm, args) -> list[str]: data_dir=args.calib_data_dir, ) samples = builder.build_calibration_dataset() + if not samples: + raise ValueError( + "RULER produced no calibration samples (all candidates exceeded " + f"max_length_filter={int(args.calib_max_seqlen * 1.5)} tokens). " + "Adjust --calib_max_seqlen / --calib_samples, or pass --prompts_file." + ) prompts = [sample["input"] for sample in samples] lengths = sorted(sample["length"] for sample in samples) print( @@ -117,7 +120,11 @@ def _write_config(ckpt: str, sparse_config: dict, update_checkpoint: bool) -> No config_json = Path(ckpt) / "config.json" config = json.loads(config_json.read_text()) config["sparse_attention_config"] = sparse_config - config_json.write_text(json.dumps(config, indent=2)) + # Atomic replace: a crash mid-write must not truncate the checkpoint's + # config.json (write_text would rewrite it in place). + tmp_path = config_json.with_name(config_json.name + ".tmp") + tmp_path.write_text(json.dumps(config, indent=2)) + os.replace(tmp_path, config_json) print(f"[ModelOpt] Merged sparse_attention_config into {config_json}") @@ -211,6 +218,14 @@ def main(): ) args = parser.parse_args() + if args.update_checkpoint_config and not (Path(args.model) / "config.json").is_file(): + # Fail before the (expensive, multi-GPU) calibration run, not after: + # merging requires a local checkpoint directory, not a HF hub ID. + parser.error( + f"--update_checkpoint_config requires a local checkpoint directory " + f"containing config.json; {args.model!r} has none" + ) + # Workers run in separate processes and must import the calibration worker. repo_root = str(Path(__file__).resolve().parent) if repo_root not in sys.path: @@ -218,6 +233,8 @@ def main(): current = os.environ.get("PYTHONPATH") os.environ["PYTHONPATH"] = os.pathsep.join([current, repo_root]) if current else repo_root + # Deferred heavy import: keep argparse/--help (and arg errors) fast, and + # only import vLLM after the PYTHONPATH setup above. from vllm import LLM, SamplingParams llm_kwargs = { @@ -284,6 +301,12 @@ def main(): "so observed sparsity spans the (10%, 90%) fitting window." ) sys.exit(1) + # Export only requested phases: a stray record (e.g. a scheduling corner + # case classified into an unrequested phase) must not bake an + # uncalibrated-by-intent phase into the config. + calibration_params = { + phase: params for phase, params in calibration_params.items() if phase in requested_phases + } sparse_config = build_sparse_attention_config( calibration_params, diff --git a/modelopt/torch/kernels/sparsity/attention/calibrate.py b/modelopt/torch/kernels/sparsity/attention/calibrate.py index b8e9d6475c7..3655c4cbfcd 100644 --- a/modelopt/torch/kernels/sparsity/attention/calibrate.py +++ b/modelopt/torch/kernels/sparsity/attention/calibrate.py @@ -22,6 +22,7 @@ ``modelopt.torch.sparsity.attention_sparsity`` to fit a skip threshold. """ +import functools import math import torch @@ -263,6 +264,22 @@ def _attn_fwd_calibrate( tl.store(Out + o_ptrs, acc, mask=(q_pos[:, None] < seq_len_q) & d_mask[None, :]) +@functools.lru_cache(maxsize=64) +def _log2_threshold_tensor( + threshold_trials: tuple[float, ...], device: torch.device +) -> torch.Tensor: + """Build the log2-space threshold tensor, cached per (trials, device). + + Scores already include sm_scale and LOG2E; convert lambda to log2 space + only. Trials are constant for a whole calibration run, and the vLLM path + calls :func:`attention_calibrate` once per request per layer per step, so + rebuilding (and re-uploading) the tensor per call would be pure waste. + """ + return torch.tensor( + [math.log2(t) for t in threshold_trials], dtype=torch.float32, device=device + ) + + def attention_calibrate( q: torch.Tensor, k: torch.Tensor, @@ -363,19 +380,22 @@ def attention_calibrate( b_seq_len_k = b_seq_len b_start_loc_k = b_start_loc - # Paged mode: KV positions come from block_table, so the contiguous KV - # offsets are unused. Provide a dummy so Triton can compile the tl.load. if b_start_loc_k is None: - b_start_loc_k = torch.zeros_like(b_start_loc) + if not is_paged: + # A zeros dummy here would silently read every sequence's K/V from + # offset 0 — fail loudly instead (contiguous K/V needs real offsets). + raise ValueError( + "b_start_loc_k is required when b_seq_len_k is provided for " + "contiguous (non-paged) K/V" + ) + # Paged mode: KV positions come from block_table, so the contiguous KV + # offsets are unused. Alias b_start_loc (same shape/dtype/device) so + # Triton can compile the tl.load without allocating a dummy per call. + b_start_loc_k = b_start_loc num_thresholds = len(threshold_trials) - # Scores already include sm_scale and LOG2E; convert lambda to log2 space only. - threshold_tensor = torch.tensor( - [math.log2(t) for t in threshold_trials], - dtype=torch.float32, - device=q.device, - ) + threshold_tensor = _log2_threshold_tensor(tuple(threshold_trials), q.device) o = torch.empty_like(q) diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..9e5b6c186ec 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -180,7 +180,10 @@ def _dynamic_block_quantize_impl( and not DISABLE_TRITON_KERNEL and amax is not None ): - return triton_kernel.fp4_fake_quant_block(inputs, amax) + # Forward the configured block size: the kernel defaults to 16, and + # silently ignoring a non-16 NVFP4 block config here would diverge + # from the cuda_ext fallback below (which honors block_size). + return triton_kernel.fp4_fake_quant_block(inputs, amax, block_size=block_size) cuda_ext_mx = get_cuda_ext_mx(raise_if_failed=True) return cuda_ext_mx.fused_amax_convert( inputs, diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index 8c41b895a8a..b2acd137930 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py @@ -39,6 +39,29 @@ ) +def export_threshold_scale_factor(calibration_params: dict[str, Any]) -> dict[str, Any]: + """Build the canonical per-phase ``threshold_scale_factor`` export block. + + Single source of the exported skip-softmax schema fragment, shared by the + HF exporter (:func:`export_sparse_attention_config`) and the vLLM + calibration path (``plugins.sparse_attn_calibration``), so the serving + loader sees one format regardless of which path produced the checkpoint. + """ + block: dict[str, Any] = {"formula": "a * exp(b * target_sparsity)"} + for phase in ("prefill", "decode"): + if phase in calibration_params: + block[phase] = { + "a": float(calibration_params[phase]["a"]), + "b": float(calibration_params[phase]["b"]), + } + return block + + +def export_config_producer() -> dict[str, str]: + """Build the canonical ``producer`` block of ``sparse_attention_config``.""" + return {"name": "modelopt", "version": mo_version} + + def _set_attn_implementation(model: nn.Module, config: SparseAttentionConfig) -> None: """Set the correct attn_implementation based on the sparse attention method/backend. @@ -469,16 +492,7 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: skip_group["initial_disabled_steps"] = initial_disabled_steps # threshold_scale_factor (a * exp(b * target_sparsity)) and target_sparsity are # skip-softmax-specific, so they live in this group. - threshold_scale_factor: dict[str, Any] = { - "formula": "a * exp(b * target_sparsity)", - } - for phase in ["prefill", "decode"]: - if phase in calibration_params: - threshold_scale_factor[phase] = { - "a": calibration_params[phase]["a"], - "b": calibration_params[phase]["b"], - } - skip_group["threshold_scale_factor"] = threshold_scale_factor + skip_group["threshold_scale_factor"] = export_threshold_scale_factor(calibration_params) if target_sparse_ratio is not None: skip_group["target_sparsity"] = target_sparse_ratio config_groups[f"group_{group_idx}"] = skip_group @@ -495,10 +509,7 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: return { "config_groups": config_groups, - "producer": { - "name": "modelopt", - "version": mo_version, - }, + "producer": export_config_producer(), } diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py index 32fb5d67160..e0bd2250d87 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -33,11 +33,13 @@ from typing import Any -import modelopt - # One canonical sweep for both calibration paths: re-exported from the HF-path # calibrator so the vLLM path fits on the identical trial grid. -from ..calibration.calibrator import DEFAULT_THRESHOLD_TRIALS +from ..calibration.calibrator import DEFAULT_THRESHOLD_TRIALS, DynamicThresholdCalibrator + +# One canonical schema for both calibration paths: the exported skip-softmax +# blocks come from the HF exporter's helpers so the formats cannot drift. +from ..conversion import export_config_producer, export_threshold_scale_factor __all__ = [ "DEFAULT_THRESHOLD_TRIALS", @@ -115,7 +117,9 @@ def merge_count_records(sources: list[list[dict]]) -> list[dict]: return merged -def merge_phase_counts(rank_counts: list[dict[str, list[dict]]]) -> dict[str, list[dict]]: +def merge_phase_counts( + rank_counts: list[dict[str, list[dict]]], *, source_desc: str = "rank" +) -> dict[str, list[dict]]: """Merge per-phase raw-count records collected from every TP rank. ``rank_counts`` is the list of per-rank results (one @@ -123,7 +127,9 @@ def merge_phase_counts(rank_counts: list[dict[str, list[dict]]]) -> dict[str, li ``collect_calibration_counts``). Use ALL ranks: with tensor parallelism each rank only measures its attention-head shard, so any single rank's counts are incomplete. A phase recorded by some ranks but not others - indicates a collection bug and raises. + indicates a collection bug and raises. The same merge also sums one + rank's per-layer splits (``collect_calibration_counts`` delegates here + with ``source_desc="attention layer"``). """ phases = {phase for rank in rank_counts for phase in rank} merged: dict[str, list[dict]] = {} @@ -132,10 +138,11 @@ def merge_phase_counts(rank_counts: list[dict[str, list[dict]]]) -> dict[str, li empty = sum(1 for source in sources if not source) if empty and empty != len(sources): raise ValueError( - f"Misaligned calibration records: {empty}/{len(sources)} rank(s) " + f"Misaligned calibration records: {empty}/{len(sources)} {source_desc}(s) " f"recorded no {phase!r} samples while others did" ) - merged[phase] = merge_count_records([] if empty else sources) + # All-empty sources merge to [] (merge_count_records sums zero samples). + merged[phase] = merge_count_records(sources) return merged @@ -172,8 +179,6 @@ def fit_from_counts( ``{phase: {"a", "b", "min_observed_sparsity", "max_observed_sparsity"}}`` for each phase that produced a valid fit. """ - from ..calibration.calibrator import DynamicThresholdCalibrator - calibration_params: dict[str, dict[str, float]] = {} for phase, records in per_phase_counts.items(): if not records: @@ -228,37 +233,49 @@ def build_sparse_attention_config( Non-skip groups from ``existing_config`` (e.g. exported N:M ``sparse_softmax`` metadata) are preserved after the skip group; an - existing ``skip_softmax`` group is replaced by the new calibration. + existing ``skip_softmax`` group is replaced by the new calibration, + carrying over its layer policy (``ignore`` — layers deliberately kept + dense — and ``initial_disabled_steps``): recalibration replaces the + fitted thresholds, not which layers the export sparsifies. """ - threshold_scale_factor: dict[str, Any] = {"formula": "a * exp(b * target_sparsity)"} - for phase in _PHASES: - if phase in calibration_params: - threshold_scale_factor[phase] = { - "a": float(calibration_params[phase]["a"]), - "b": float(calibration_params[phase]["b"]), - } + # target_sparsity covers only fitted phases: claiming a target for a phase + # without calibrated (a, b) would advertise sparsity the serving path + # silently serves dense (it needs the per-phase scale factors). + target_sparsity_by_phase = { + phase: value + for phase, value in _normalize_target_sparsity(target_sparsity).items() + if phase in calibration_params + } skip_group: dict[str, Any] = { "algorithm": "skip_softmax", "targets": ["Attention"], - "threshold_scale_factor": threshold_scale_factor, - "target_sparsity": _normalize_target_sparsity(target_sparsity), + "threshold_scale_factor": export_threshold_scale_factor(calibration_params), + "target_sparsity": target_sparsity_by_phase, } config_groups: dict[str, Any] = {"group_0": skip_group} existing_groups = (existing_config or {}).get("config_groups") if isinstance(existing_groups, dict): - preserved = [ - group - for group in existing_groups.values() - if isinstance(group, dict) and group.get("algorithm") != "skip_softmax" - ] + preserved = [] + for group in existing_groups.values(): + if not isinstance(group, dict): + continue + if group.get("algorithm") == "skip_softmax": + # Keep the replaced group's layer policy: dropping ``ignore`` + # would sparsify layers the original export deliberately kept + # dense (e.g. first/last blocks). + for key in ("ignore", "initial_disabled_steps"): + if key in group and key not in skip_group: + skip_group[key] = group[key] + else: + preserved.append(group) for idx, group in enumerate(preserved, start=1): config_groups[f"group_{idx}"] = group result: dict[str, Any] = { "config_groups": config_groups, - "producer": {"name": "modelopt", "version": modelopt.__version__}, + "producer": export_config_producer(), } # Legacy checkpoints carry N:M parameters as a top-level ``sparse_softmax`` # dict (read by the serving loader ahead of group params) — preserve it so diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 8bede9ed665..a6d3abffa4e 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -48,6 +48,17 @@ from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite from modelopt.torch.kernels.sparsity.attention.calibrate import attention_calibrate +from .sparse_attn_calibration import merge_phase_counts, split_records_by_phase + +__all__ = [ + "ModelOptSparseAttentionBackend", + "ModelOptSparseAttentionImpl", + "collect_calibration_counts", + "disable_calibration", + "enable_calibration", + "iter_sparse_impls", +] + @functools.cache def _flash_attention_kv_cache_layout() -> str: @@ -311,10 +322,17 @@ def _forward_calibrate( sparsification is applied to generation (the dense Triton kernel's numerics may differ slightly from the native backend's). - Phase and causality are decided per request: ``q_len == 1`` is a decode - step (full-cache, non-causal); ``q_len > 1`` is (chunked) prefill (causal - — the kernel offsets the query into the KV span). A mixed prefill/decode - batch therefore contributes correctly to both phase fits. + Phase and causality are decided per request: ``q_len > 1`` is (chunked) + prefill (causal — the kernel offsets the query into the KV span). A + ``q_len == 1`` row is a decode step (full-cache, non-causal) only when + its KV span exceeds the request's prompt (at least one generated token); + a 1-token row still inside the prompt is the final chunk of a chunked + prefill and is recorded as prefill. Prompt lengths come from the runner's + input batch (the installer attaches the runner as ``_calib_model_runner``; + the input batch is resolved per forward because vLLM can rebuild it after + install, same request order as the metadata rows); without it, + ``q_len == 1`` falls back to decode. A mixed prefill/decode batch + therefore contributes correctly to both phase fits. Records raw per-threshold tile counts (not ratios) on ``impl._calib_records`` so tensor-parallel workers can be aggregated by @@ -335,8 +353,20 @@ def _forward_calibrate( page_size = key_cache.shape[1] trials = impl._calib_threshold_trials batch = seq_lens.shape[0] - b_start_loc = cu_seqlens_q[:batch] - b_seq_len = cu_seqlens_q[1 : batch + 1] - cu_seqlens_q[:batch] + # Hoist per-request tensors out of the loop: kernel args are sliced views + # of these, so the loop performs no allocations or casts. + b_seq_len_i32 = (cu_seqlens_q[1 : batch + 1] - cu_seqlens_q[:batch]).to(torch.int32) + seq_lens_i32 = seq_lens[:batch].to(torch.int32) + b_start_loc_zero = torch.zeros(1, device=query.device, dtype=torch.int32) + # vLLM v1 metadata tensors are GPU-resident: take one host copy per launch + # instead of stalling the stream with per-request ``.item()`` syncs. + cu_seqlens_q_cpu = cu_seqlens_q[: batch + 1].cpu() + seq_lens_cpu = seq_lens[:batch].cpu() + # Per-request prompt lengths (same request order as the metadata rows) + # distinguish decode steps from 1-token final chunks of a chunked + # prefill. Resolved from the runner per forward: vLLM can replace + # input_batch after install (KV-cache init for hybrid models). + input_batch = getattr(getattr(impl, "_calib_model_runner", None), "input_batch", None) q = query[:num_actual_tokens].contiguous() # Dummy K/V: in paged mode KV is read from the cache via block_table. @@ -344,23 +374,29 @@ def _forward_calibrate( k_dummy = torch.empty(0, impl.num_kv_heads, impl.head_size, device=q.device, dtype=q.dtype) for i in range(batch): - q_len = int(b_seq_len[i].item()) + q_start = int(cu_seqlens_q_cpu[i]) + q_len = int(cu_seqlens_q_cpu[i + 1]) - q_start if q_len <= 0: continue - q_start = int(b_start_loc[i].item()) - seq_k = int(seq_lens[i].item()) - phase = "decode" if q_len <= 1 else "prefill" + seq_k = int(seq_lens_cpu[i]) + if q_len > 1: + phase = "prefill" + elif input_batch is not None and seq_k <= int(input_batch.num_prompt_tokens[i]): + # 1-token final chunk of a chunked prefill: still inside the prompt. + phase = "prefill" + else: + phase = "decode" oi, counters = attention_calibrate( q[q_start : q_start + q_len], k_dummy, k_dummy, - b_start_loc=torch.zeros(1, device=q.device, dtype=torch.int32), - b_seq_len=b_seq_len[i : i + 1].to(torch.int32), + b_start_loc=b_start_loc_zero, + b_seq_len=b_seq_len_i32[i : i + 1], max_input_len=q_len, is_causal=q_len > 1, softmax_scale=impl.scale, - b_seq_len_k=seq_lens[i : i + 1].to(torch.int32), + b_seq_len_k=seq_lens_i32[i : i + 1], max_input_len_k=seq_k, threshold_trials=trials, k_cache=key_cache, @@ -370,12 +406,14 @@ def _forward_calibrate( ) output[q_start : q_start + q_len] = oi + # One host transfer for both counter columns (counters is GPU-resident). + counters_cpu = counters.cpu() impl._calib_records.append( { "phase": phase, "sample_length": seq_k, - "total_tiles": counters[:, 0].tolist(), - "skipped_tiles": counters[:, 1].tolist(), + "total_tiles": counters_cpu[:, 0].tolist(), + "skipped_tiles": counters_cpu[:, 1].tolist(), } ) @@ -518,6 +556,8 @@ def _dispatch_modelopt( num_prefills: int, num_decode_tokens: int, num_prefill_tokens: int, + max_seq_len_decode: int | None = None, + max_seq_len_prefill: int | None = None, **common_kw, ) -> torch.Tensor: """Run the ModelOpt path, splitting mixed decode+prefill batches by phase. @@ -527,8 +567,14 @@ def _dispatch_modelopt( ``q_len==1`` decode rows with ``q_len>1`` (chunked-)prefill rows, ``max_query_len > 1`` and the whole batch would otherwise take the prefill skip-softmax path. Split so each phase runs its own schedule -- decode rows - always take the fixed decode path. Both the FlashAttention and FlashInfer - adapters share this dispatch. + always take the fixed decode path. + + Both adapters route through this dispatch, but the split is live only on + FlashInfer: its metadata carries the ``num_decodes``/``num_prefills`` + counts (and vLLM reorders those batches decode-first). vLLM's + FlashAttention metadata has no phase counts, so FA mixed batches fall + through to the whole-batch path and are classified by ``max_query_len`` + alone (decode rows then follow the prefill contract for that launch). """ if not (num_decodes and num_prefills): return _forward_modelopt( @@ -554,6 +600,18 @@ def _dispatch_modelopt( if not common_kw.get("quant_active", False): common_kw["dense_fallback"]() + # Each phase derives its skip threshold from its own KV maximum: reusing + # the batch-global max_seq_len (e.g. a co-scheduled 32k prefill next to 2k + # decodes) would shrink the decode threshold far below — much denser than + # — the calibrated target. Fall back to the batch-global value only when + # the builder did not provide per-phase maxima. + decode_kw = dict(common_kw) + if max_seq_len_decode is not None: + decode_kw["max_seq_len"] = max_seq_len_decode + prefill_kw = dict(common_kw) + if max_seq_len_prefill is not None: + prefill_kw["max_seq_len"] = max_seq_len_prefill + _forward_modelopt( impl, query=query[:num_decode_tokens], @@ -563,7 +621,7 @@ def _dispatch_modelopt( num_actual_tokens=num_decode_tokens, max_query_len=num_decode_tokens // num_decodes, output=output[:num_decode_tokens], - **common_kw, + **decode_kw, ) prefill_start = num_decode_tokens prefill_cu_seqlens_q = cu_seqlens_q[num_decodes:] - cu_seqlens_q[num_decodes] @@ -576,7 +634,7 @@ def _dispatch_modelopt( num_actual_tokens=num_prefill_tokens, max_query_len=max_query_len, output=output[prefill_start : prefill_start + num_prefill_tokens], - **common_kw, + **prefill_kw, ) return output @@ -778,6 +836,23 @@ def build(*args, **kwargs): common = build_sig.bind(*args, **kwargs).arguments["common_attn_metadata"] for target, source in _FLASHINFER_METADATA_FIELDS.items(): setattr(metadata, target, getattr(common, source)) + # Per-phase KV maxima for the mixed-batch split (batch is reordered + # decode-first): computed once per build — not per layer forward — so + # the split's threshold derivation neither reuses the batch-global max + # nor syncs the stream inside every layer. + num_decodes = getattr(metadata, "num_decodes", 0) + num_prefills = getattr(metadata, "num_prefills", 0) + max_seq_len_decode = max_seq_len_prefill = None + if num_decodes and num_prefills: + # Prefer the host-resident copy the runner may already carry; + # fall back to one device->host copy per mixed-batch build. + seq_lens_cpu = getattr(common, "_seq_lens_cpu", None) + if seq_lens_cpu is None: + seq_lens_cpu = common.seq_lens.cpu() + max_seq_len_decode = int(seq_lens_cpu[:num_decodes].max()) + max_seq_len_prefill = int(seq_lens_cpu[num_decodes : num_decodes + num_prefills].max()) + metadata._modelopt_max_seq_len_decode = max_seq_len_decode + metadata._modelopt_max_seq_len_prefill = max_seq_len_prefill return metadata setattr(build, "_modelopt_sparse_metadata_patch", True) @@ -940,6 +1015,8 @@ def prepare_modelopt(): num_prefills=getattr(attn_metadata, "num_prefills", 0), num_decode_tokens=getattr(attn_metadata, "num_decode_tokens", 0), num_prefill_tokens=getattr(attn_metadata, "num_prefill_tokens", 0), + max_seq_len_decode=getattr(attn_metadata, "_modelopt_max_seq_len_decode", None), + max_seq_len_prefill=getattr(attn_metadata, "_modelopt_max_seq_len_prefill", None), **common_kw, ) @@ -1068,23 +1145,11 @@ def collect_calibration_counts(model) -> dict[str, list[dict]]: phase with :func:`~.sparse_attn_calibration.fit_from_counts` — sparsity ratios are only formed after the global merge. """ - from .sparse_attn_calibration import merge_count_records, split_records_by_phase - splits = [ split_records_by_phase(getattr(impl, "_calib_records", [])) for impl in iter_sparse_impls(model) ] - phases = {phase for split in splits for phase, records in split.items() if records} - counts: dict[str, list[dict]] = {} - for phase in phases: - sources = [split.get(phase, []) for split in splits] - empty = sum(1 for source in sources if not source) - if empty: - # Every layer sees every launch, so a layer with no records for a - # phase others measured indicates a collection bug. - raise ValueError( - f"Misaligned calibration records: {empty}/{len(sources)} attention " - f"layer(s) recorded no {phase!r} samples while others did" - ) - counts[phase] = merge_count_records(sources) - return counts + # Same merge as the cross-rank aggregation: every layer sees every launch, + # so a layer with no records for a phase others measured indicates a + # collection bug (merge_phase_counts raises). + return merge_phase_counts(splits, source_desc="attention layer") diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index 724feb69f90..6e2ad911288 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -156,7 +156,7 @@ def _cudagraph_mode(model_runner): return mode if mode is not None else CUDAGraphMode.NONE -def _global_errors(model_runner) -> list[str]: +def _global_errors(model_runner, *, sparse_only: bool = False) -> list[str]: config = getattr(model_runner, "vllm_config", None) if config is None: return ["model_runner.vllm_config is required"] @@ -174,10 +174,16 @@ def _global_errors(model_runner) -> list[str]: errors.append("decode_context_parallel_size must be 1") if getattr(parallel, "enable_dbo", False) or getattr(parallel, "use_ubatching", False): errors.append("DBO/ubatching is unsupported") - if getattr(cache_config, "enable_prefix_caching", False): - errors.append("prefix caching is unsupported") - if getattr(config, "kv_transfer_config", None) is not None: - errors.append("KV transfer is unsupported") + if not sparse_only: + # Prefix caching and KV transfer break only flows that quantize the + # cache on write or measure per-request prefills (quantized installs, + # skip-softmax calibration). Sparse-only serving reads the cache + # unmodified and supports prefix-cache suffix attention by offsetting + # query positions (see the vllm_serve README limitations). + if getattr(cache_config, "enable_prefix_caching", False): + errors.append("prefix caching is unsupported") + if getattr(config, "kv_transfer_config", None) is not None: + errors.append("KV transfer is unsupported") if getattr(config, "speculative_config", None) is not None: errors.append("speculative decoding is unsupported") if _cudagraph_mode(model_runner).mixed_mode() == CUDAGraphMode.FULL: @@ -355,7 +361,12 @@ def _plan_vllm_attention( ) _require_supported_vllm() - errors = _global_errors(model_runner) if quantize else [] + # Engine-level checks apply to sparse-only installs too: the ModelOpt + # kernel path silently ignores decode context parallelism, DBO, and + # speculative decoding, and FULL mixed-batch graphs would capture stale + # per-launch thresholds (same rationale as the decode graph guard below). + # Sparse-only installs skip only the cache-mutation checks. + errors = _global_errors(model_runner, sparse_only=not quantize) mode = _cudagraph_mode(model_runner) quant_plugin: Any = _load_quant_plugin() if quantize else None plans = [] @@ -599,6 +610,16 @@ def install_vllm_skip_softmax_calibration(model_runner) -> VllmAttentionInstallR _raise_unsupported(errors, "skip-softmax calibration") plan = _InstallPlan(model_runner, tuple(plans), False, "SKIP_SOFTMAX_CALIBRATION") + # Per-request prompt lengths (same request order as the attention-metadata + # rows, kept aligned by vLLM's in-place batch reorder) let the adapter + # classify q_len == 1 rows: a decode step's KV span exceeds the prompt, + # while a 1-token final chunk of a chunked prefill is still inside it. + # The runner — not its input_batch — is attached: vLLM can rebuild + # input_batch after load_model (may_reinitialize_input_batch during KV + # cache init, e.g. hybrid Mamba/attention models), so the adapter must + # resolve the live object per forward. + for attention_plan in plans: + attention_plan.new_impl._calib_model_runner = model_runner return _apply_vllm_attention_plans(plan) diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index 60523d44457..99e8a3a83a4 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -20,7 +20,7 @@ import pytest import torch -from conftest import make_qkv, make_varlen_meta, sdpa_reference +from conftest import make_qkv, make_varlen_meta from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor, e2m1_values @@ -454,9 +454,12 @@ def test_invalid_amax_raises(self): with pytest.raises(ValueError, match="p_qdq_amax"): attention(q, k, v, locs, lens, 8, p_qdq="fp8", p_qdq_amax=0.0) - @requires_native_e4m3 - def test_composes_with_skip_softmax(self): - """p_qdq composes with the skip-softmax feature in one launch.""" + def test_rejects_skip_softmax(self): + """p_qdq cannot combine with active skip-softmax in one launch. + + Quantized P changes the score distribution the skip thresholds were + calibrated on, so the kernel rejects the composition (pre-launch). + """ seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 scale = 1.0 / (head_dim**0.5) @@ -464,19 +467,18 @@ def test_composes_with_skip_softmax(self): q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.float16) locs, lens = make_varlen_meta([seq_len]) - o = attention( - q, - k, - v, - locs, - lens, - seq_len, - softmax_scale=scale, - p_qdq="fp8", - skip_softmax_threshold=1e-3, - ) - ref = sdpa_reference(q, k, v, locs, lens) - torch.testing.assert_close(o, ref, rtol=5e-2, atol=5e-2) + with pytest.raises(ValueError, match="cannot be combined with attention quantization"): + attention( + q, + k, + v, + locs, + lens, + seq_len, + softmax_scale=scale, + p_qdq="fp8", + skip_softmax_threshold=1e-3, + ) def test_invalid_mode_raises(self): q, k, v = make_qkv(8, 2, 2, 32, dtype=torch.float16) diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index cf802abaf40..2e2a0105846 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -354,3 +354,36 @@ def test_static_vs_dynamic_fp4_kernels(self, set_torch_dtype, block_size, num_bl f"Mean abs diff: {(output_static - output_dynamic).abs().mean()}\n" f"Max relative diff: {((output_static - output_dynamic).abs() / (output_dynamic.abs() + 1e-8)).max()}" ) + + +def test_dynamic_block_quantize_forwards_block_size(monkeypatch): + """NVFP4 dynamic dispatch must pass the configured block size to the Triton kernel. + + The Triton kernel defaults to block_size=16; dropping the argument would + silently quantize non-16 NVFP4 configs with 16-element blocks while the + cuda_ext fallback honors the configured size (device-dependent numerics). + """ + recorded = {} + + def fake_fp4_fake_quant_block(inputs, amax, block_size=16): + recorded["block_size"] = block_size + return inputs.clone() + + monkeypatch.setattr(tensor_quant, "DISABLE_TRITON_KERNEL", False) + monkeypatch.setattr(triton_kernel, "IS_AVAILABLE", True, raising=False) + monkeypatch.setattr( + triton_kernel, "fp4_fake_quant_block", fake_fp4_fake_quant_block, raising=False + ) + + inputs = torch.randn(4, 64, device="cuda") + amax = inputs.abs().amax() + tensor_quant._dynamic_block_quantize_impl( + inputs, + 32, # non-default NVFP4 block size + amax, + 4, # num_bits total (E2M1) + 2, # exponent_bits -> (2, 1) + 8, # scale_num_bits -> (4, 3) + 4, # scale_exponent_bits + ) + assert recorded["block_size"] == 32 diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py index 543fbd0a2da..d3563d226cf 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -15,6 +15,7 @@ """Tests for skip-softmax calibration through the vLLM adapters and installer.""" +import sys from types import SimpleNamespace import pytest @@ -25,6 +26,7 @@ from vllm.v1.attention.backends.flashinfer import FlashInferImpl from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.quantization.plugins import vllm as quant_plugin from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as attention_plugin from modelopt.torch.sparsity.attention_sparsity.plugins import vllm_runtime from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( @@ -177,8 +179,6 @@ def test_sparse_only_install_allows_skip_on_unquantized_layer(self): assert report.installed_count == 1 def test_quantized_plan_allows_nm_sparsity(self, monkeypatch): - from modelopt.torch.quantization.plugins import vllm as quant_plugin - monkeypatch.setattr( quant_plugin, "_get_device_dtype", @@ -197,8 +197,6 @@ class TestFlashInferLayoutGuard: def test_layout_helper_preserves_none(self, monkeypatch): """A getter returning None must not become the truthy string 'None'.""" - import sys - fake = SimpleNamespace(get_kv_cache_layout=lambda: None) monkeypatch.setitem(sys.modules, "vllm.v1.attention.backends.utils", fake) assert attention_plugin._flashinfer_kv_cache_layout() is None diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py index b3ceef30eab..9c448c1327c 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py @@ -90,7 +90,6 @@ def test_sparse_install_from_checkpoint_is_validation_atomic(): nn.ModuleDict({"valid_attn": valid, "invalid_attn": invalid}), sparse_metadata=_sparse_metadata(), ) - del runner.vllm_config with pytest.raises(NotImplementedError, match="sliding_window"): vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) @@ -107,7 +106,6 @@ def test_installer_rejects_cross_attention_layout_even_if_marked_decoder(): nn.ModuleDict({"cross_attn": attention}), sparse_metadata=_sparse_metadata(), ) - del runner.vllm_config with pytest.raises(NotImplementedError, match="layout CrossAttention"): vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) @@ -122,7 +120,6 @@ def test_sparse_install_uses_checkpoint_metadata(monkeypatch, impl_cls): nn.ModuleDict({"attn": attention}), sparse_metadata=_sparse_metadata(), ) - del runner.vllm_config monkeypatch.setattr( vllm_runtime.attention_plugin, "patch_flashinfer_metadata_builder", lambda: True ) diff --git a/tests/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index 62395ff5a7a..13ce5da55ef 100644 --- a/tests/unit/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -139,21 +139,7 @@ def test_forward_routes_every_mode_to_single_autotuner( assert kernel.kwargs["V_QDQ"] == expected_v_qdq -@pytest.mark.parametrize( - ("attention_kwargs", "expected_block_m"), - [ - ({"skip_softmax_threshold": 0.1, "measure_sparsity": True}, 128), - ( - { - "p_qdq": "nvfp4", - "skip_softmax_threshold": 0.1, - "measure_sparsity": True, - }, - 16, - ), - ], -) -def test_forward_measurement_uses_one_fixed_launch(monkeypatch, attention_kwargs, expected_block_m): +def test_forward_measurement_uses_one_fixed_launch(monkeypatch): """Counter measurement bypasses autotuning to avoid repeated atomic updates.""" pytest.importorskip("triton") @@ -173,10 +159,42 @@ def test_forward_measurement_uses_one_fixed_launch(monkeypatch, attention_kwargs starts = torch.tensor([0], dtype=torch.int32) lengths = torch.tensor([seq_len], dtype=torch.int32) - triton_fa.attention(q, k, v, starts, lengths, seq_len, **attention_kwargs) + triton_fa.attention( + q, k, v, starts, lengths, seq_len, skip_softmax_threshold=0.1, measure_sparsity=True + ) assert kernel.fn.launch_count == 1 - assert kernel.fn.kwargs["BLOCK_M"] == expected_block_m + assert kernel.fn.kwargs["BLOCK_M"] == 128 assert kernel.fn.kwargs["BLOCK_N"] == 128 assert kernel.fn.kwargs["num_stages"] == 1 assert kernel.fn.kwargs["num_warps"] == 4 + + +@pytest.mark.parametrize( + "qdq_kwargs", + [{"p_qdq": "nvfp4"}, {"p_qdq": "fp8"}, {"v_qdq": "nvfp4", "v_qdq_amax": 1.0}], +) +def test_forward_rejects_skip_softmax_with_qdq(monkeypatch, qdq_kwargs): + """Active skip-softmax rejects P/V QDQ before any kernel launch.""" + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + kernel = _ForbiddenKernel() + kernel.fn = _ForbiddenKernel() + monkeypatch.setattr(triton_fa, "_attn_fwd", kernel) + monkeypatch.setattr(triton_fa.torch.cuda, "device", lambda _device: nullcontext()) + monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) + monkeypatch.setattr(triton_fa, "_load_qdq_helpers", lambda: None) + + seq_len = 129 + q = torch.empty(seq_len, 2, 16) + k = torch.empty(seq_len, 1, 16) + v = torch.empty_like(k) + starts = torch.tensor([0], dtype=torch.int32) + lengths = torch.tensor([seq_len], dtype=torch.int32) + + with pytest.raises(ValueError, match="cannot be combined with attention quantization"): + triton_fa.attention( + q, k, v, starts, lengths, seq_len, skip_softmax_threshold=0.1, **qdq_kwargs + ) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py index e85ad3ab77e..7ff174952b3 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 @@ -177,6 +177,31 @@ def test_canonical_schema(self): assert group["target_sparsity"] == {"prefill": 0.4, "decode": 0.4} assert config["producer"]["name"] == "modelopt" + def test_target_sparsity_covers_only_fitted_phases(self): + """A phase without calibrated (a, b) must not claim a sparsity target.""" + config = build_sparse_attention_config({"prefill": {"a": 7.9, "b": 8.6}}, 0.4) + group = config["config_groups"]["group_0"] + assert group["target_sparsity"] == {"prefill": 0.4} + assert "decode" not in group["threshold_scale_factor"] + + def test_replaced_skip_group_keeps_layer_policy(self): + """Recalibration replaces thresholds but keeps ignore/initial_disabled_steps.""" + existing = { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "ignore": ["model.layers.0.self_attn"], + "initial_disabled_steps": 4, + "threshold_scale_factor": {"prefill": {"a": 1.0, "b": 1.0}}, + } + } + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + group = config["config_groups"]["group_0"] + assert group["ignore"] == ["model.layers.0.self_attn"] + assert group["initial_disabled_steps"] == 4 + assert group["threshold_scale_factor"]["prefill"] == {"a": 7.9, "b": 8.6} + def test_preserves_nm_groups_and_replaces_old_skip_group(self): existing = { "config_groups": { From 2229a8c62c9f8d8c675a59d0f39c46442da7cc1a Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 24 Aug 2026 10:54:18 -0700 Subject: [PATCH 10/10] Add target-first offline mask reuse calibration Signed-off-by: Kai Xu --- examples/vllm_serve/calibrate_mask_reuse.py | 486 +++++ .../vllm_serve/create_checkpoint_manifest.py | 43 + .../calibration/__init__.py | 44 + .../calibration/checkpoint_manifest.py | 439 +++++ .../calibration/mask_reuse.py | 1702 +++++++++++++++++ .../calibration/mask_reuse_compact.py | 1545 +++++++++++++++ .../test_calibrate_mask_reuse_cli.py | 263 +++ .../test_checkpoint_manifest.py | 177 ++ .../test_mask_reuse_calibration.py | 322 ++++ .../test_mask_reuse_compact_calibration.py | 408 ++++ 10 files changed, 5429 insertions(+) create mode 100644 examples/vllm_serve/calibrate_mask_reuse.py create mode 100644 examples/vllm_serve/create_checkpoint_manifest.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/calibration/checkpoint_manifest.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_calibrate_mask_reuse_cli.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py diff --git a/examples/vllm_serve/calibrate_mask_reuse.py b/examples/vllm_serve/calibrate_mask_reuse.py new file mode 100644 index 00000000000..c4535082140 --- /dev/null +++ b/examples/vllm_serve/calibrate_mask_reuse.py @@ -0,0 +1,486 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Build a fail-closed schema-v3 mask-reuse candidate from compact captures. + +The input is compact capture JSONL emitted by ``collect_mask_reuse.py``. Its +streaming selector never expands the full consumer-head by donor-head matrix +into repeated row objects. This command cannot promote a serving policy until +the grouped inner/outer protocol is implemented and its preregistered gates pass. + +Example:: + + python examples/vllm_serve/calibrate_mask_reuse.py \ + --checkpoint /path/to/checkpoint \ + --compact-captures compact-captures.jsonl \ + --capture-manifest compact-captures.jsonl.manifest.json \ + --vanilla-config sparse_attention_config.json \ + --topology topology.json \ + --calibration-plan calibration-plan.json \ + --family-registry family-registry.json \ + --grouped-fit grouped-fit.json \ + --outer-report outer-report.json \ + --max-anchor-dropped-mass 0.02 \ + --reuse-dropped-mass-report-threshold 0.02 \ + --target-bmm1-skip-ratio 0.10 + +The final output is candidate-only and must be rejected by serving. +""" + +from __future__ import annotations + +import argparse +import json +import os +import tempfile +from collections.abc import Mapping +from hashlib import sha256 +from pathlib import Path + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + StableFileSnapshot, + read_stable_file_snapshot, + stable_file_sha256, + verify_checkpoint_manifest, +) +from modelopt.torch.sparsity.attention_sparsity.calibration.mask_reuse_compact import ( + calibrate_compact_mask_reuse_policy, + load_compact_mask_reuse_captures, +) + +_EVIDENCE_ARTIFACTS = { + "calibration_plan_sha256": "calibration_plan", + "family_registry_sha256": "family_registry", + "grouped_fit_sha256": "grouped_fit", + "outer_report_sha256": "outer_report", +} + +_CAPTURE_MANIFEST_FIELDS = frozenset( + { + "capture_manifest_schema_version", + "capture_protocol", + "model", + "checkpoint_manifest_sha256", + "checkpoint_manifest_path", + "checkpoint_file_count", + "checkpoint_total_size_bytes", + "plan", + "fa4_source", + "fa4_source_commit", + "fa4_source_git_tree", + "fa4_source_git_archive_sha256", + "fa4_source_manifest_path", + "fa4_source_manifest_sha256", + "fa4_source_directory_count", + "fa4_source_file_count", + "fa4_source_total_size_bytes", + "engine_kwargs", + "dense_shadow_validation_requested", + "target_sparsity_hex", + "vanilla_threshold_scale_factor", + "vanilla_fit_sha256", + "vanilla_config_file_sha256", + "prompt_plan_file_sha256", + "compact_capture_file_sha256", + "capture_count", + "candidate_cell_count", + "captures", + } +) + + +def _reject_duplicate_json_keys(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + +def _parse_json_object(payload: bytes, *, path: Path, label: str) -> dict[str, object]: + try: + raw = json.loads( + payload, + object_pairs_hook=_reject_duplicate_json_keys, + ) + except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as error: + raise ValueError(f"could not load {label} from {path}: {error}") from error + if not isinstance(raw, dict): + raise ValueError(f"{label} must contain a JSON object") + return raw + + +def _load_json_snapshot(path: Path, *, label: str) -> tuple[dict[str, object], StableFileSnapshot]: + snapshot = read_stable_file_snapshot(path, label=label) + return _parse_json_object(snapshot.payload, path=path, label=label), snapshot + + +def _load_json_object(path: Path, *, label: str) -> dict[str, object]: + """Load strict JSON from one stable no-follow byte snapshot.""" + return _load_json_snapshot(path, label=label)[0] + + +def _stable_file_sha256(path: Path, *, label: str) -> str: + return stable_file_sha256(path, label=label) + + +def _evidence_artifacts( + args: argparse.Namespace, *, vanilla_fit_sha256: str +) -> tuple[dict[str, str], dict[str, Path]]: + paths = { + field: Path(getattr(args, attribute)) for field, attribute in _EVIDENCE_ARTIFACTS.items() + } + paths["vanilla_fit_sha256"] = args.vanilla_config + paths["reuse_bundle_sha256"] = args.compact_captures + evidence = { + field: ( + vanilla_fit_sha256 + if field == "vanilla_fit_sha256" + else _stable_file_sha256(path, label=field) + ) + for field, path in paths.items() + } + return evidence, paths + + +def _canonical_json_bytes(value: object) -> bytes: + return ( + json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode("utf-8") + + +def _fsync_directory(path: Path) -> None: + directory_flag = getattr(os, "O_DIRECTORY", None) + if directory_flag is None: + # Windows lacks portable directory fsync; no-clobber linking remains + # atomic, while crash durability of the directory entry is best effort. + return + descriptor = os.open(path, os.O_RDONLY | directory_flag) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def _temporary_payload(path: Path, payload: bytes) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + mode="wb", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as handle: + temporary = Path(handle.name) + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + return temporary + + +def _unlink_if_identity(path: Path, identity: tuple[int, int]) -> None: + if identity[1] == 0: + return + try: + observed = path.stat(follow_symlinks=False) + except FileNotFoundError: + return + if observed.st_ino != 0 and (observed.st_dev, observed.st_ino) == identity: + path.unlink() + _fsync_directory(path.parent) + + +def _publish_no_clobber(temporary: Path, destination: Path) -> tuple[int, int]: + observed = temporary.stat(follow_symlinks=False) + if observed.st_ino == 0: + raise RuntimeError("candidate temporary file has no stable identity") + identity = observed.st_dev, observed.st_ino + os.link(temporary, destination, follow_symlinks=False) + try: + published = destination.stat(follow_symlinks=False) + if published.st_ino == 0 or (published.st_dev, published.st_ino) != identity: + raise RuntimeError("candidate destination changed during publication") + temporary.unlink() + _fsync_directory(destination.parent) + except BaseException: + _unlink_if_identity(destination, identity) + raise + return identity + + +def _publish_candidate_outputs( + policy_path: Path, policy_payload: bytes, report_path: Path, report_payload: bytes +) -> None: + if policy_path.exists() or report_path.exists(): + raise FileExistsError("candidate outputs already exist; refusing to overwrite them") + policy_temporary: Path | None = None + report_temporary: Path | None = None + report_identity: tuple[int, int] | None = None + try: + policy_temporary = _temporary_payload(policy_path, policy_payload) + report_temporary = _temporary_payload(report_path, report_payload) + report_identity = _publish_no_clobber(report_temporary, report_path) + report_temporary = None + _publish_no_clobber(policy_temporary, policy_path) + policy_temporary = None + except BaseException: + if report_identity is not None: + _unlink_if_identity(report_path, report_identity) + if policy_temporary is not None: + policy_temporary.unlink(missing_ok=True) + if report_temporary is not None: + report_temporary.unlink(missing_ok=True) + raise + + +def _validate_capture_manifest( + raw: Mapping[str, object], + *, + checkpoint_sha256: str, + model: str, + compact_capture_sha256: str, + vanilla_config_sha256: str, +) -> None: + missing = _CAPTURE_MANIFEST_FIELDS - raw.keys() + extra = raw.keys() - _CAPTURE_MANIFEST_FIELDS + if missing or extra: + raise ValueError( + "capture manifest fields do not match schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + expected = { + "capture_manifest_schema_version": 4, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v4", + "model": model, + "checkpoint_manifest_sha256": checkpoint_sha256, + "compact_capture_file_sha256": compact_capture_sha256, + "vanilla_config_file_sha256": vanilla_config_sha256, + } + for field, value in expected.items(): + if raw[field] != value: + raise ValueError(f"capture manifest {field} does not match its verified input") + for field, length in { + "fa4_source_commit": 40, + "fa4_source_git_tree": 40, + "fa4_source_git_archive_sha256": 64, + "fa4_source_manifest_sha256": 64, + }.items(): + value = raw[field] + if ( + not isinstance(value, str) + or len(value) != length + or any(character not in "0123456789abcdef" for character in value) + ): + raise ValueError(f"capture manifest {field} is not canonical hexadecimal evidence") + for field in ("fa4_source", "fa4_source_manifest_path"): + if not isinstance(raw[field], str) or not raw[field]: + raise ValueError(f"capture manifest {field} must be a non-empty path") + for field in ( + "fa4_source_directory_count", + "fa4_source_file_count", + "fa4_source_total_size_bytes", + ): + value = raw[field] + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"capture manifest {field} must be an integer >= 0") + if raw["fa4_source_file_count"] == 0: + raise ValueError("capture manifest must bind at least one FA4 source file") + if not isinstance(raw["engine_kwargs"], Mapping): + raise ValueError("capture manifest engine_kwargs must be an object") + if not isinstance(raw["dense_shadow_validation_requested"], bool): + raise ValueError("capture manifest dense_shadow_validation_requested must be boolean") + if isinstance(raw["capture_count"], bool) or not isinstance(raw["capture_count"], int): + raise ValueError("capture manifest capture_count must be an integer") + if raw["capture_count"] <= 0: + raise ValueError("capture manifest must contain at least one capture") + captures = raw["captures"] + if not isinstance(captures, list) or len(captures) != raw["capture_count"]: + raise ValueError("capture manifest captures do not match capture_count") + candidate_cell_count = raw["candidate_cell_count"] + if ( + isinstance(candidate_cell_count, bool) + or not isinstance(candidate_cell_count, int) + or candidate_cell_count <= 0 + ): + raise ValueError("capture manifest candidate_cell_count must be positive") + observed_cells = 0 + for index, capture in enumerate(captures): + if not isinstance(capture, Mapping) or not isinstance( + capture.get("candidate_cell_count"), int + ): + raise ValueError( + f"capture manifest captures[{index}].candidate_cell_count must be an integer" + ) + observed_cells += int(capture["candidate_cell_count"]) + if observed_cells != candidate_cell_count: + raise ValueError("capture manifest candidate-cell total is inconsistent") + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Build a fail-closed schema-v3 candidate from schema-v4 mask-reuse capture evidence" + ), + allow_abbrev=False, + ) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument( + "--compact-captures", + type=Path, + required=True, + help="Compact capture JSONL emitted by collect_mask_reuse.py (recommended)", + ) + parser.add_argument("--capture-manifest", type=Path, required=True) + parser.add_argument( + "--vanilla-config", + type=Path, + required=True, + help="ModelOpt sparse_attention_config JSON or checkpoint config.json", + ) + parser.add_argument( + "--topology", + type=Path, + required=True, + help="JSON object containing anchors and nearest layer mappings", + ) + parser.add_argument("--calibration-plan", type=Path, required=True) + parser.add_argument("--family-registry", type=Path, required=True) + parser.add_argument("--grouped-fit", type=Path, required=True) + parser.add_argument("--outer-report", type=Path, required=True) + parser.add_argument( + "--max-anchor-dropped-mass", + type=float, + required=True, + help="Maximum allowed anchor dropped mass", + ) + parser.add_argument( + "--reuse-dropped-mass-report-threshold", + type=float, + required=True, + help="Diagnostic reuse threshold; does not affect selection", + ) + parser.add_argument( + "--target-bmm1-skip-ratio", + type=float, + required=True, + help="Minimum model-wide BMM1 tile skip ratio required in every context bucket", + ) + parser.add_argument( + "--output-policy", + type=Path, + default=Path("mask_reuse_candidate.json"), + help="Fail-closed schema-v3 candidate output path", + ) + parser.add_argument( + "--output-report", + type=Path, + default=Path("mask_reuse_calibration_report.json"), + help="Standalone calibration-report output path", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + if args.output_policy.resolve() == args.output_report.resolve(): + parser.error("--output-policy and --output-report must be different paths") + + try: + checkpoint = verify_checkpoint_manifest(args.checkpoint) + vanilla_config, vanilla_snapshot = _load_json_snapshot( + args.vanilla_config, label="vanilla config" + ) + topology, topology_snapshot = _load_json_snapshot(args.topology, label="topology") + capture_manifest, capture_manifest_snapshot = _load_json_snapshot( + args.capture_manifest, label="capture manifest" + ) + evidence, evidence_paths = _evidence_artifacts( + args, vanilla_fit_sha256=vanilla_snapshot.sha256 + ) + capture_manifest_sha256 = capture_manifest_snapshot.sha256 + topology_sha256 = topology_snapshot.sha256 + _validate_capture_manifest( + capture_manifest, + checkpoint_sha256=checkpoint.sha256, + model=checkpoint.model, + compact_capture_sha256=evidence["reuse_bundle_sha256"], + vanilla_config_sha256=evidence["vanilla_fit_sha256"], + ) + artifact = calibrate_compact_mask_reuse_policy( + load_compact_mask_reuse_captures(args.compact_captures), + vanilla_calibration=vanilla_config, + topology=topology, + checkpoint_manifest=checkpoint, + evidence=evidence, + max_anchor_dropped_mass=args.max_anchor_dropped_mass, + reuse_dropped_mass_report_threshold=(args.reuse_dropped_mass_report_threshold), + target_bmm1_skip_ratio=args.target_bmm1_skip_ratio, + source_provenance={ + "capture_manifest_sha256": capture_manifest_sha256, + "topology_file_sha256": topology_sha256, + }, + ) + provenance = artifact.get("provenance") + if not isinstance(provenance, Mapping): + raise ValueError("calibrator returned no provenance object") + if provenance.get("input_capture_count") != capture_manifest["capture_count"]: + raise ValueError("calibrator capture count does not match capture manifest") + if provenance.get("candidate_cell_count") != capture_manifest["candidate_cell_count"]: + raise ValueError("calibrator candidate-cell count does not match capture manifest") + for field, path in evidence_paths.items(): + if _stable_file_sha256(path, label=field) != evidence[field]: + raise ValueError(f"{field} artifact changed during calibration") + if ( + _stable_file_sha256(args.capture_manifest, label="capture manifest") + != capture_manifest_sha256 + ): + raise ValueError("capture manifest changed during calibration") + if _stable_file_sha256(args.topology, label="topology") != topology_sha256: + raise ValueError("topology changed during calibration") + if verify_checkpoint_manifest(args.checkpoint) != checkpoint: + raise ValueError("checkpoint changed during calibration") + except (OSError, ValueError) as error: + parser.error(str(error)) + + if ( + artifact.get("promotion_status") != "candidate_only" + or artifact.get("deployment_geometry_validated") is not False + ): + parser.error("calibrator did not return a fail-closed candidate-only artifact") + + report = artifact.get("calibration_report") + if not isinstance(report, Mapping): + parser.error("calibrator returned no calibration_report object") + + try: + policy_payload = _canonical_json_bytes(artifact) + report_payload = _canonical_json_bytes(report) + _publish_candidate_outputs( + args.output_policy, policy_payload, args.output_report, report_payload + ) + except (OSError, FileExistsError) as error: + parser.error(f"could not write calibration outputs: {error}") + + policy_digest = sha256(policy_payload).hexdigest() + print(f"[ModelOpt] Wrote fail-closed mask-reuse candidate to {args.output_policy.resolve()}") + print(f"[ModelOpt] Wrote calibration report to {args.output_report.resolve()}") + print(f"MASK_REUSE_FA4_CANDIDATE_SHA256={policy_digest}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/vllm_serve/create_checkpoint_manifest.py b/examples/vllm_serve/create_checkpoint_manifest.py new file mode 100644 index 00000000000..8fe276f0d0f --- /dev/null +++ b/examples/vllm_serve/create_checkpoint_manifest.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Create a deterministic, no-clobber checkpoint manifest for calibration.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + create_checkpoint_manifest, +) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(allow_abbrev=False) + parser.add_argument("checkpoint", type=Path) + parser.add_argument("--model-id", required=True) + args = parser.parse_args(argv) + try: + manifest = create_checkpoint_manifest(args.checkpoint, model=args.model_id) + except (OSError, ValueError) as error: + parser.error(str(error)) + print(f"[ModelOpt] Wrote {manifest.manifest_path}") + print(f"CHECKPOINT_MANIFEST_SHA256={manifest.sha256}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py b/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py index 87088f805bd..c8b55c03a3a 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py @@ -17,10 +17,54 @@ from .calibrate import calibrate_sparse_attention from .calibrator import DynamicThresholdCalibrator +from .checkpoint_manifest import ( + CHECKPOINT_MANIFEST_NAME, + CheckpointManifestError, + StableFileSnapshot, + VerifiedCheckpointManifest, + create_checkpoint_manifest, + read_stable_file_snapshot, + stable_file_sha256, + verify_checkpoint_manifest, +) +from .mask_reuse import ( + AnchorLayerStats, + MaskReuseCalibrationError, + MaskReuseObservation, + calibrate_mask_reuse_policy, + canonical_prefill_threshold_scale_factor, + load_mask_reuse_observations, + parse_mask_reuse_observations, +) +from .mask_reuse_compact import ( + CompactMaskReuseCapture, + CompactMaskReuseCaptureSource, + calibrate_compact_mask_reuse_policy, + load_compact_mask_reuse_captures, +) from .ruler_dataset import RulerDatasetBuilder __all__ = [ + "CHECKPOINT_MANIFEST_NAME", + "AnchorLayerStats", + "CheckpointManifestError", + "CompactMaskReuseCapture", + "CompactMaskReuseCaptureSource", "DynamicThresholdCalibrator", + "MaskReuseCalibrationError", + "MaskReuseObservation", "RulerDatasetBuilder", + "StableFileSnapshot", + "VerifiedCheckpointManifest", + "calibrate_compact_mask_reuse_policy", + "calibrate_mask_reuse_policy", "calibrate_sparse_attention", + "canonical_prefill_threshold_scale_factor", + "create_checkpoint_manifest", + "load_compact_mask_reuse_captures", + "load_mask_reuse_observations", + "parse_mask_reuse_observations", + "read_stable_file_snapshot", + "stable_file_sha256", + "verify_checkpoint_manifest", ] diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/checkpoint_manifest.py b/modelopt/torch/sparsity/attention_sparsity/calibration/checkpoint_manifest.py new file mode 100644 index 00000000000..dd79737fb25 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/checkpoint_manifest.py @@ -0,0 +1,439 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Strict, content-addressed checkpoint identity for mask-reuse calibration.""" + +from __future__ import annotations + +import json +import os +import stat +import tempfile +import unicodedata +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path, PurePosixPath + +__all__ = [ + "CHECKPOINT_MANIFEST_NAME", + "CheckpointManifestError", + "StableFileSnapshot", + "VerifiedCheckpointManifest", + "create_checkpoint_manifest", + "read_stable_file_snapshot", + "stable_file_sha256", + "verify_checkpoint_manifest", +] + +CHECKPOINT_MANIFEST_NAME = "checkpoint_manifest.json" +_MANIFEST_FIELDS = frozenset({"checkpoint_manifest_schema_version", "model", "files"}) +_FILE_FIELDS = frozenset({"path", "size_bytes", "sha256"}) +_WEIGHT_SUFFIXES = frozenset({".bin", ".pt", ".safetensors"}) +_FILE_ATTRIBUTE_REPARSE_POINT = getattr(stat, "FILE_ATTRIBUTE_REPARSE_POINT", 0x400) + + +class CheckpointManifestError(ValueError): + """Raised when a checkpoint cannot be bound to its exact file contents.""" + + +def _strict_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise CheckpointManifestError(f"checkpoint manifest repeats JSON key {key!r}") + result[key] = value + return result + + +def _exact_fields(raw: dict[str, object], expected: frozenset[str], label: str) -> None: + missing = expected - raw.keys() + extra = raw.keys() - expected + if missing or extra: + raise CheckpointManifestError( + f"{label} fields do not match the schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + + +def _text(value: object, label: str) -> str: + if ( + not isinstance(value, str) + or not value + or value != value.strip() + or unicodedata.normalize("NFC", value) != value + or any(ord(character) < 32 for character in value) + ): + raise CheckpointManifestError(f"{label} must be non-empty canonical NFC text") + return value + + +def _sha256(value: object, label: str) -> str: + if ( + not isinstance(value, str) + or len(value) != 64 + or any(character not in "0123456789abcdef" for character in value) + ): + raise CheckpointManifestError(f"{label} must be a lowercase SHA256") + return value + + +def _canonical_json_bytes(value: object) -> bytes: + return ( + json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + + +def _is_link_like(value: os.stat_result) -> bool: + """Return whether a no-follow stat identifies a symlink or Windows reparse point.""" + return stat.S_ISLNK(value.st_mode) or bool( + getattr(value, "st_file_attributes", 0) & _FILE_ATTRIBUTE_REPARSE_POINT + ) + + +def _same_file(left: os.stat_result, right: os.stat_result) -> bool: + """Compare file identities, failing closed when an inode is unavailable.""" + return left.st_ino != 0 and right.st_ino != 0 and os.path.samestat(left, right) + + +def _open_stable_regular(path: Path, label: str) -> tuple[int, os.stat_result]: + try: + named_before = path.stat(follow_symlinks=False) + except OSError as error: + raise CheckpointManifestError(f"could not inspect {label}") from error + if _is_link_like(named_before): + raise CheckpointManifestError(f"could not open {label} without following symlinks") + if not stat.S_ISREG(named_before.st_mode): + raise CheckpointManifestError(f"{label} must be one stable regular file, not a symlink") + try: + # Windows has no O_NOFOLLOW. The no-follow pre/post stats and handle + # identity checks keep that fallback fail-closed before any bytes are read. + descriptor = os.open( + path, + os.O_RDONLY | getattr(os, "O_BINARY", 0) | getattr(os, "O_NOFOLLOW", 0), + ) + except OSError as error: + raise CheckpointManifestError( + f"could not open {label} without following symlinks" + ) from error + try: + opened = os.fstat(descriptor) + named = path.stat(follow_symlinks=False) + except OSError: + os.close(descriptor) + raise + if _is_link_like(named): + os.close(descriptor) + raise CheckpointManifestError(f"could not open {label} without following symlinks") + if ( + not stat.S_ISREG(opened.st_mode) + or not stat.S_ISREG(named.st_mode) + or not _same_file(named_before, opened) + or not _same_file(opened, named) + ): + os.close(descriptor) + raise CheckpointManifestError(f"{label} must be one stable regular file, not a symlink") + return descriptor, opened + + +def _hash_stable_regular( + path: Path, label: str, *, capture_payload: bool = False +) -> tuple[int, str, bytes | None]: + descriptor, before = _open_stable_regular(path, label) + digest = sha256() + payload = bytearray() if capture_payload else None + observed_size = 0 + try: + for chunk in iter(lambda: os.read(descriptor, 1024 * 1024), b""): + observed_size += len(chunk) + digest.update(chunk) + if payload is not None: + payload.extend(chunk) + after = os.fstat(descriptor) + named_after = path.stat(follow_symlinks=False) + except OSError as error: + raise CheckpointManifestError(f"could not hash stable {label}") from error + finally: + os.close(descriptor) + if ( + _is_link_like(named_after) + or not stat.S_ISREG(named_after.st_mode) + or not _same_file(before, after) + or not _same_file(after, named_after) + or (before.st_size, before.st_mtime_ns) != (after.st_size, after.st_mtime_ns) + or (after.st_size, after.st_mtime_ns) != (named_after.st_size, named_after.st_mtime_ns) + ): + raise CheckpointManifestError(f"{label} changed while it was being hashed") + return observed_size, digest.hexdigest(), None if payload is None else bytes(payload) + + +@dataclass(frozen=True, slots=True) +class StableFileSnapshot: + """Exact bytes and SHA256 read from one stable no-follow descriptor.""" + + path: Path + payload: bytes + sha256: str + + +def read_stable_file_snapshot(path: str | Path, *, label: str) -> StableFileSnapshot: + """Read and hash identical bytes from one stable regular file.""" + source = Path(path) + _, digest, payload = _hash_stable_regular(source, label, capture_payload=True) + assert payload is not None + return StableFileSnapshot(source, payload, digest) + + +def stable_file_sha256(path: str | Path, *, label: str) -> str: + """Hash one stable regular file without retaining its contents.""" + _, digest, _ = _hash_stable_regular(Path(path), label) + return digest + + +def _checkpoint_files(root: Path, manifest_path: Path) -> set[str]: + files: set[str] = set() + pending = [root] + while pending: + directory = pending.pop() + try: + with os.scandir(directory) as iterator: + entries = sorted(iterator, key=lambda entry: entry.name) + except OSError as error: + raise CheckpointManifestError( + f"could not traverse checkpoint directory {directory}" + ) from error + for entry in entries: + path = Path(entry.path) + relative = path.relative_to(root).as_posix() + try: + observed = entry.stat(follow_symlinks=False) + except OSError as error: + raise CheckpointManifestError( + f"could not inspect checkpoint path {relative!r}" + ) from error + if _is_link_like(observed): + raise CheckpointManifestError( + f"checkpoint contains forbidden symlink or reparse point {relative!r}" + ) + if stat.S_ISDIR(observed.st_mode): + pending.append(path) + elif stat.S_ISREG(observed.st_mode): + if path != manifest_path: + files.add(relative) + else: + raise CheckpointManifestError(f"checkpoint contains non-regular path {relative!r}") + return files + + +def _fsync_directory(path: Path) -> None: + directory_flag = getattr(os, "O_DIRECTORY", None) + if directory_flag is None: + # Python on Windows cannot portably open and fsync a directory. The + # complete temporary file is still fsynced before its no-clobber link. + return + descriptor = os.open(path, os.O_RDONLY | directory_flag) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def create_checkpoint_manifest(checkpoint: str | Path, *, model: str) -> VerifiedCheckpointManifest: + """Create the deterministic checkpoint manifest without replacing any file.""" + root = Path(checkpoint).expanduser().resolve() + if not root.is_dir(): + raise CheckpointManifestError("checkpoint must be a local directory") + manifest_path = root / CHECKPOINT_MANIFEST_NAME + if os.path.lexists(manifest_path): + raise CheckpointManifestError( + f"{CHECKPOINT_MANIFEST_NAME} already exists; refusing to overwrite it" + ) + files = _checkpoint_files(root, manifest_path) + entries = [] + for relative in sorted(files): + size, digest, _ = _hash_stable_regular(root / relative, f"checkpoint file {relative!r}") + entries.append({"path": relative, "size_bytes": size, "sha256": digest}) + if _checkpoint_files(root, manifest_path) != files: + raise CheckpointManifestError("checkpoint file set changed while building manifest") + payload = _canonical_json_bytes( + { + "checkpoint_manifest_schema_version": 1, + "model": _text(model, "model"), + "files": entries, + } + ) + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=root, + prefix=f".{CHECKPOINT_MANIFEST_NAME}.", + suffix=".tmp", + delete=False, + ) as handle: + temporary = Path(handle.name) + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + observed = temporary.stat(follow_symlinks=False) + if observed.st_ino == 0: + raise CheckpointManifestError( + "checkpoint manifest temporary file has no stable identity" + ) + identity = observed.st_dev, observed.st_ino + os.link(temporary, manifest_path, follow_symlinks=False) + try: + published = manifest_path.stat(follow_symlinks=False) + if ( + _is_link_like(published) + or not stat.S_ISREG(published.st_mode) + or not _same_file(observed, published) + ): + raise CheckpointManifestError( + "checkpoint manifest destination changed during publication" + ) + temporary.unlink() + temporary = None + _fsync_directory(root) + except BaseException: + try: + published = manifest_path.stat(follow_symlinks=False) + if published.st_ino != 0 and (published.st_dev, published.st_ino) == identity: + manifest_path.unlink() + _fsync_directory(root) + finally: + raise + except FileExistsError as error: + raise CheckpointManifestError( + f"{CHECKPOINT_MANIFEST_NAME} appeared during publication; refusing to overwrite it" + ) from error + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + return verify_checkpoint_manifest(root, expected_model=model) + + +def _relative_path(value: object, label: str) -> str: + text = _text(value, label) + path = PurePosixPath(text) + if ( + path.is_absolute() + or text != path.as_posix() + or "\\" in text + or any(part in {"", ".", ".."} for part in path.parts) + or text == CHECKPOINT_MANIFEST_NAME + ): + raise CheckpointManifestError(f"{label} must be a canonical relative POSIX path") + return text + + +@dataclass(frozen=True, slots=True) +class VerifiedCheckpointManifest: + """Identity of a checkpoint whose complete file set was SHA256-verified.""" + + checkpoint_root: Path + manifest_path: Path + model: str + sha256: str + file_count: int + total_size_bytes: int + + +def verify_checkpoint_manifest( + checkpoint: str | Path, *, expected_model: str | None = None +) -> VerifiedCheckpointManifest: + """Verify the fixed manifest under ``checkpoint`` and every declared file. + + The manifest must enumerate every regular file below the loaded checkpoint + directory except itself. This prevents a manifest that binds only a subset + of weights or remote-code/tokenizer inputs from naming the checkpoint. + """ + root = Path(checkpoint).expanduser().resolve() + if not root.is_dir(): + raise CheckpointManifestError("checkpoint must be a local directory") + manifest_path = root / CHECKPOINT_MANIFEST_NAME + _, manifest_digest, payload = _hash_stable_regular( + manifest_path, "checkpoint manifest", capture_payload=True + ) + assert payload is not None + try: + raw = json.loads(payload, object_pairs_hook=_strict_object) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise CheckpointManifestError("checkpoint manifest is not strict UTF-8 JSON") from error + if not isinstance(raw, dict): + raise CheckpointManifestError("checkpoint manifest must be a JSON object") + _exact_fields(raw, _MANIFEST_FIELDS, "checkpoint manifest") + if raw["checkpoint_manifest_schema_version"] != 1: + raise CheckpointManifestError("checkpoint_manifest_schema_version must be 1") + if payload != _canonical_json_bytes(raw): + raise CheckpointManifestError("checkpoint manifest bytes are not canonical JSON") + model = _text(raw["model"], "checkpoint manifest.model") + if expected_model is not None and model != expected_model: + raise CheckpointManifestError( + f"checkpoint manifest model {model!r} does not match requested model {expected_model!r}" + ) + raw_files = raw["files"] + if not isinstance(raw_files, list) or not raw_files: + raise CheckpointManifestError("checkpoint manifest.files must be a non-empty list") + + declared: dict[str, tuple[int, str]] = {} + for index, item in enumerate(raw_files): + label = f"checkpoint manifest.files[{index}]" + if not isinstance(item, dict): + raise CheckpointManifestError(f"{label} must be an object") + _exact_fields(item, _FILE_FIELDS, label) + relative = _relative_path(item["path"], f"{label}.path") + size = item["size_bytes"] + if isinstance(size, bool) or not isinstance(size, int) or size < 0: + raise CheckpointManifestError(f"{label}.size_bytes must be an integer >= 0") + digest = _sha256(item["sha256"], f"{label}.sha256") + if relative in declared: + raise CheckpointManifestError(f"checkpoint manifest repeats file {relative!r}") + declared[relative] = (size, digest) + if list(declared) != sorted(declared): + raise CheckpointManifestError("checkpoint manifest files must be sorted by path") + if "config.json" not in declared or not any( + Path(relative).suffix in _WEIGHT_SUFFIXES for relative in declared + ): + raise CheckpointManifestError( + "checkpoint manifest must bind config.json and at least one model weight file" + ) + + actual = _checkpoint_files(root, manifest_path) + if actual != set(declared): + raise CheckpointManifestError( + "checkpoint manifest does not exactly cover checkpoint files; " + f"missing={sorted(actual - set(declared))}, extra={sorted(set(declared) - actual)}" + ) + total_size = 0 + for relative, (expected_size, expected_digest) in declared.items(): + path = root / relative + observed_size, observed_digest, _ = _hash_stable_regular( + path, f"checkpoint file {relative!r}" + ) + if observed_size != expected_size or observed_digest != expected_digest: + raise CheckpointManifestError( + f"checkpoint file {relative!r} does not match its size/SHA256 manifest entry" + ) + total_size += observed_size + if _checkpoint_files(root, manifest_path) != actual: + raise CheckpointManifestError("checkpoint file set changed during verification") + return VerifiedCheckpointManifest( + checkpoint_root=root, + manifest_path=manifest_path, + model=model, + sha256=manifest_digest, + file_count=len(declared), + total_size_bytes=total_size, + ) diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py new file mode 100644 index 00000000000..e7b9f6bbce3 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py @@ -0,0 +1,1702 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline calibration and schema-v3 export for cross-layer mask reuse. + +The selector consumes prompt-level observations measured at target sparsities +derived from an existing ModelOpt skip-softmax fit. Calibration observations +alone select one target sparsity per context bucket and one donor head (or an +exact fallback) per consumer head. Held-out observations only evaluate the +frozen policy. + +This module intentionally has no serving-backend dependency. It exports the +JSON-safe schema consumed by the mask-reuse attention backend. +""" + +from __future__ import annotations + +import json +import math +from collections import defaultdict +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import cast + +import pulp + +import modelopt + +from .checkpoint_manifest import VerifiedCheckpointManifest + +__all__ = [ + "AnchorLayerStats", + "MaskReuseCalibrationError", + "MaskReuseObservation", + "calibrate_mask_reuse_policy", + "canonical_prefill_threshold_scale_factor", + "load_mask_reuse_observations", + "parse_mask_reuse_observations", +] + + +_FORMULA = "a * exp(b * target_sparsity)" +_SPLITS = frozenset({"calibration", "heldout"}) +_OBSERVATION_FIELDS = frozenset( + { + "model", + "min_kv_tokens", + "max_kv_tokens", + "target_sparsity", + "sample_length", + "threshold_lambda", + "threshold_log2", + "q_tokens", + "kv_tokens", + "q_start_tokens", + "split", + "prompt_id", + "source_capture_sha256", + "anchor_layer", + "consumer_layer", + "consumer_head", + "donor_head", + "retained_tiles", + "eligible_tiles", + "anchor_dropped_mass", + "anchor_stats_by_layer", + "dropped_mass", + } +) +_CALIBRATION_PROTOCOL = "modelopt_mask_reuse_target_sparsity_v1" +_SOLVER_LEXICOGRAPHIC_ATOL = 1e-7 +_EVIDENCE_FIELDS = frozenset( + { + "calibration_plan_sha256", + "family_registry_sha256", + "vanilla_fit_sha256", + "reuse_bundle_sha256", + "grouped_fit_sha256", + "outer_report_sha256", + } +) +_DEPLOYMENT_GEOMETRY_CONTRACT: dict[str, object] = { + "schema_version": 1, + "batch_size": 1, + "max_query_chunk_tokens": 8192, + "query_block_tokens": 128, + "key_block_tokens": 128, + "qstage2_query_pair_tokens": 256, + "kv_page_tokens": 16, + "head_dim": 128, + "causal": True, + "bottom_right_aligned": True, + "query_chunk_start_alignment_tokens": 128, + "attention_dtype": "bfloat16", + "kv_cache_dtype": "bfloat16", + "common_prefix": False, + "cascade_attention": False, + "context_parallel_size": 1, + "pipeline_parallel_size": 1, +} + +Bucket = tuple[int, int | None] +ConsumerHead = tuple[int, int] +ObservationKey = tuple[str, str, float, int, int, int] +AnchorKey = tuple[str, str, float, int, int] + + +class MaskReuseCalibrationError(ValueError): + """Raised when observations cannot produce a trustworthy reuse policy.""" + + +@dataclass(frozen=True, slots=True) +class AnchorLayerStats: + """Per-head BLASST mask statistics for one topology anchor layer.""" + + retained_tiles: tuple[int, ...] + dropped_mass: tuple[float, ...] + + +def _integer(value: object, name: str, *, minimum: int) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise MaskReuseCalibrationError(f"{name} must be an integer >= {minimum}") + return value + + +def _number( + value: object, + name: str, + *, + minimum: float, + maximum: float | None = None, + minimum_inclusive: bool = True, +) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise MaskReuseCalibrationError(f"{name} must be a finite number") + result = float(value) + below = result < minimum if minimum_inclusive else result <= minimum + if not math.isfinite(result) or below or (maximum is not None and result > maximum): + raise MaskReuseCalibrationError(f"{name} is outside its valid range") + return result + + +def _sha256(value: object, name: str) -> str: + if not isinstance(value, str): + raise MaskReuseCalibrationError(f"{name} must be a lowercase SHA256") + normalized = value.strip().lower() + if len(normalized) != 64 or any( + character not in "0123456789abcdef" for character in normalized + ): + raise MaskReuseCalibrationError(f"{name} must be a lowercase SHA256") + return normalized + + +def _parse_anchor_stats_by_layer( + value: object, + *, + require_canonical_string_keys: bool, +) -> dict[int, AnchorLayerStats]: + if not isinstance(value, Mapping) or not value: + raise MaskReuseCalibrationError("anchor_stats_by_layer must be a non-empty object") + parsed: dict[int, AnchorLayerStats] = {} + for raw_layer, raw_stats in value.items(): + if isinstance(raw_layer, bool): + raise MaskReuseCalibrationError("anchor_stats_by_layer has a non-integer layer key") + try: + layer = int(raw_layer) + except (TypeError, ValueError) as error: + raise MaskReuseCalibrationError( + "anchor_stats_by_layer has a non-integer layer key" + ) from error + if layer < 0 or ( + require_canonical_string_keys + and (not isinstance(raw_layer, str) or raw_layer != str(layer)) + ): + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer layer key {raw_layer!r} is not canonical" + ) + if layer in parsed: + raise MaskReuseCalibrationError(f"anchor_stats_by_layer repeats layer {layer}") + if isinstance(raw_stats, AnchorLayerStats): + raw_retained = raw_stats.retained_tiles + raw_dropped = raw_stats.dropped_mass + elif isinstance(raw_stats, Mapping): + missing = {"retained_tiles", "dropped_mass"} - raw_stats.keys() + extra = raw_stats.keys() - {"retained_tiles", "dropped_mass"} + if missing or extra: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}] requires exactly retained_tiles " + f"and dropped_mass; missing={sorted(missing)}, extra={sorted(extra)}" + ) + raw_retained = raw_stats["retained_tiles"] + raw_dropped = raw_stats["dropped_mass"] + else: + raise MaskReuseCalibrationError(f"anchor_stats_by_layer[{layer}] must be an object") + if not isinstance(raw_retained, list | tuple) or not raw_retained: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}].retained_tiles must be a non-empty list" + ) + if not isinstance(raw_dropped, list | tuple) or not raw_dropped: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}].dropped_mass must be a non-empty list" + ) + retained = tuple( + _integer( + item, + f"anchor_stats_by_layer[{layer}].retained_tiles[{head}]", + minimum=0, + ) + for head, item in enumerate(raw_retained) + ) + dropped = tuple( + _number( + item, + f"anchor_stats_by_layer[{layer}].dropped_mass[{head}]", + minimum=0.0, + maximum=1.0, + ) + for head, item in enumerate(raw_dropped) + ) + if len(retained) != len(dropped): + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}] head arrays differ in width" + ) + parsed[layer] = AnchorLayerStats(retained, dropped) + return dict(sorted(parsed.items())) + + +@dataclass(frozen=True, slots=True) +class MaskReuseObservation: + """One prompt, target-sparsity, consumer-head, and donor-head observation.""" + + model: str + min_kv_tokens: int + max_kv_tokens: int | None + target_sparsity: float + sample_length: int + threshold_lambda: float + threshold_log2: float + q_tokens: int + kv_tokens: int + q_start_tokens: int + split: str + prompt_id: str + source_capture_sha256: str + anchor_layer: int + consumer_layer: int + consumer_head: int + donor_head: int + retained_tiles: int + eligible_tiles: int + anchor_dropped_mass: float + anchor_stats_by_layer: Mapping[int, AnchorLayerStats] + dropped_mass: float + + def __post_init__(self) -> None: + if not isinstance(self.model, str) or not self.model.strip(): + raise MaskReuseCalibrationError("model must be a non-empty string") + if not isinstance(self.prompt_id, str) or not self.prompt_id.strip(): + raise MaskReuseCalibrationError("prompt_id must be a non-empty string") + if self.split not in _SPLITS: + raise MaskReuseCalibrationError(f"split must be one of {sorted(_SPLITS)}") + minimum = _integer(self.min_kv_tokens, "min_kv_tokens", minimum=1) + maximum = self.max_kv_tokens + if maximum is not None: + maximum = _integer(maximum, "max_kv_tokens", minimum=minimum) + sample_length = _integer(self.sample_length, "sample_length", minimum=1) + if sample_length < minimum or (maximum is not None and sample_length > maximum): + raise MaskReuseCalibrationError("sample_length lies outside its context bucket") + q_tokens = _integer(self.q_tokens, "q_tokens", minimum=129) + if q_tokens > int(cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["max_query_chunk_tokens"])): + raise MaskReuseCalibrationError("q_tokens exceeds the deployment geometry limit") + kv_tokens = _integer(self.kv_tokens, "kv_tokens", minimum=1) + q_start_tokens = _integer(self.q_start_tokens, "q_start_tokens", minimum=0) + if sample_length != kv_tokens: + raise MaskReuseCalibrationError("sample_length must equal kv_tokens") + if q_start_tokens + q_tokens != kv_tokens: + raise MaskReuseCalibrationError("q_start_tokens + q_tokens must equal kv_tokens") + alignment = int( + cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["query_chunk_start_alignment_tokens"]) + ) + if q_start_tokens % alignment: + raise MaskReuseCalibrationError("q_start_tokens must be 128-token aligned") + anchor_layer = _integer(self.anchor_layer, "anchor_layer", minimum=0) + consumer_layer = _integer(self.consumer_layer, "consumer_layer", minimum=0) + if anchor_layer >= consumer_layer: + raise MaskReuseCalibrationError("anchor_layer must precede consumer_layer") + _integer(self.consumer_head, "consumer_head", minimum=0) + _integer(self.donor_head, "donor_head", minimum=0) + retained = _integer(self.retained_tiles, "retained_tiles", minimum=0) + eligible = _integer(self.eligible_tiles, "eligible_tiles", minimum=1) + if retained > eligible: + raise MaskReuseCalibrationError("retained_tiles must not exceed eligible_tiles") + + object.__setattr__(self, "model", self.model.strip()) + object.__setattr__(self, "prompt_id", self.prompt_id.strip()) + object.__setattr__( + self, + "source_capture_sha256", + _sha256(self.source_capture_sha256, "source_capture_sha256"), + ) + object.__setattr__( + self, + "target_sparsity", + _number( + self.target_sparsity, + "target_sparsity", + minimum=0.0, + maximum=1.0, + minimum_inclusive=False, + ), + ) + if self.target_sparsity >= 1.0: + raise MaskReuseCalibrationError("target_sparsity must be in (0, 1)") + object.__setattr__( + self, + "threshold_lambda", + _number( + self.threshold_lambda, + "threshold_lambda", + minimum=0.0, + maximum=1.0, + minimum_inclusive=False, + ), + ) + if self.threshold_lambda >= 1.0: + raise MaskReuseCalibrationError("threshold_lambda must be in (0, 1)") + object.__setattr__( + self, + "threshold_log2", + _number(self.threshold_log2, "threshold_log2", minimum=-math.inf, maximum=0.0), + ) + object.__setattr__( + self, + "anchor_dropped_mass", + _number(self.anchor_dropped_mass, "anchor_dropped_mass", minimum=0.0, maximum=1.0), + ) + object.__setattr__( + self, + "anchor_stats_by_layer", + _parse_anchor_stats_by_layer( + self.anchor_stats_by_layer, + require_canonical_string_keys=False, + ), + ) + object.__setattr__( + self, + "dropped_mass", + _number(self.dropped_mass, "dropped_mass", minimum=0.0, maximum=1.0), + ) + + @classmethod + def from_mapping(cls, raw: Mapping[str, object]) -> MaskReuseObservation: + """Build a validated observation from normalized JSON.""" + missing = _OBSERVATION_FIELDS - raw.keys() + extra = raw.keys() - _OBSERVATION_FIELDS + if missing or extra: + raise MaskReuseCalibrationError( + f"observation fields do not match the schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + values = dict(raw) + values["anchor_stats_by_layer"] = _parse_anchor_stats_by_layer( + raw["anchor_stats_by_layer"], + require_canonical_string_keys=True, + ) + return cls(**values) # type: ignore[arg-type] + + def to_mapping(self) -> dict[str, object]: + """Return the normalized JSON representation.""" + result = { + field: getattr(self, field) + for field in _OBSERVATION_FIELDS + if field != "anchor_stats_by_layer" + } + result["anchor_stats_by_layer"] = { + str(layer): { + "retained_tiles": list(stats.retained_tiles), + "dropped_mass": list(stats.dropped_mass), + } + for layer, stats in self.anchor_stats_by_layer.items() + } + return result + + +def _reject_duplicate_json_keys(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise MaskReuseCalibrationError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + +def parse_mask_reuse_observations(lines: Iterable[str]) -> list[MaskReuseObservation]: + """Parse strict normalized observation JSONL.""" + observations: list[MaskReuseObservation] = [] + for line_number, line in enumerate(lines, start=1): + if not line.strip(): + continue + try: + raw = json.loads(line, object_pairs_hook=_reject_duplicate_json_keys) + except json.JSONDecodeError as error: + raise MaskReuseCalibrationError( + f"line {line_number}: invalid JSON: {error.msg}" + ) from error + except MaskReuseCalibrationError as error: + raise MaskReuseCalibrationError(f"line {line_number}: {error}") from error + if not isinstance(raw, dict): + raise MaskReuseCalibrationError(f"line {line_number}: observation must be an object") + try: + observations.append(MaskReuseObservation.from_mapping(raw)) + except MaskReuseCalibrationError as error: + raise MaskReuseCalibrationError(f"line {line_number}: {error}") from error + if not observations: + raise MaskReuseCalibrationError("input contains no mask-reuse observations") + return observations + + +def load_mask_reuse_observations(path: str | Path) -> list[MaskReuseObservation]: + """Load normalized mask-reuse observations from JSONL.""" + with Path(path).open(encoding="utf-8") as handle: + return parse_mask_reuse_observations(handle) + + +def _find_skip_softmax_group(raw: Mapping[str, object]) -> Mapping[str, object]: + current: object = raw + if "sparse_attention_config" in raw: + current = raw["sparse_attention_config"] + if not isinstance(current, Mapping): + raise MaskReuseCalibrationError("sparse_attention_config must be an object") + if "config_groups" in current: + groups = current["config_groups"] + if not isinstance(groups, Mapping): + raise MaskReuseCalibrationError("config_groups must be an object") + matches = [ + group + for group in groups.values() + if isinstance(group, Mapping) + and ( + group.get("algorithm") == "skip_softmax" + or group.get("sparse_algo") == "softmax_skip" + ) + ] + if len(matches) != 1: + raise MaskReuseCalibrationError( + "vanilla config must contain exactly one skip_softmax config group" + ) + selected = matches[0] + if selected.get("sparse_algo") == "softmax_skip" and "threshold_scale_factor" in current: + # Older ModelOpt serving calibration stored the fit beside a + # ``sparse_algo: softmax_skip`` group instead of inside it. + current = current["threshold_scale_factor"] + else: + current = selected + if isinstance(current, Mapping) and "threshold_scale_factor" in current: + current = current["threshold_scale_factor"] + if not isinstance(current, Mapping): + raise MaskReuseCalibrationError("threshold_scale_factor must be an object") + return current + + +def canonical_prefill_threshold_scale_factor( + vanilla_calibration: Mapping[str, object], +) -> dict[str, object]: + """Canonicalize ModelOpt fit parameters or exported skip-softmax metadata.""" + raw = _find_skip_softmax_group(vanilla_calibration) + if "calibration_params" in raw: + raw = _find_skip_softmax_group(raw["calibration_params"]) # type: ignore[arg-type] + formula = raw.get("formula", _FORMULA) + if formula != _FORMULA: + raise MaskReuseCalibrationError("vanilla calibration uses an unsupported formula") + params = raw.get("prefill") + if not isinstance(params, Mapping): + raise MaskReuseCalibrationError("vanilla calibration requires prefill fit parameters") + unknown = params.keys() - { + "a", + "b", + "min_observed_sparsity", + "max_observed_sparsity", + } + if unknown or not {"a", "b"} <= params.keys(): + raise MaskReuseCalibrationError( + f"prefill fit requires a and b and contains unknown fields {sorted(unknown)}" + ) + prefill: dict[str, float] = { + "a": _number(params["a"], "prefill.a", minimum=0.0, minimum_inclusive=False), + "b": _number(params["b"], "prefill.b", minimum=0.0, maximum=20.0), + } + bounds = {"min_observed_sparsity", "max_observed_sparsity"} & params.keys() + if bounds and len(bounds) != 2: + raise MaskReuseCalibrationError("observed sparsity bounds must appear together") + if bounds: + lower = _number( + params["min_observed_sparsity"], + "prefill.min_observed_sparsity", + minimum=0.0, + maximum=1.0, + ) + upper = _number( + params["max_observed_sparsity"], + "prefill.max_observed_sparsity", + minimum=0.0, + maximum=1.0, + ) + if lower > upper: + raise MaskReuseCalibrationError("observed sparsity range is reversed") + prefill.update(min_observed_sparsity=lower, max_observed_sparsity=upper) + return {"formula": _FORMULA, "prefill": prefill} + + +def _normalize_topology(raw: Mapping[str, object]) -> tuple[tuple[int, ...], dict[int, int]]: + if set(raw) != {"anchors", "nearest"}: + raise MaskReuseCalibrationError("topology must contain exactly anchors and nearest") + raw_anchors = raw["anchors"] + if not isinstance(raw_anchors, list) or not raw_anchors: + raise MaskReuseCalibrationError("topology anchors must be a non-empty list") + anchors = tuple(sorted(_integer(value, "topology anchor", minimum=0) for value in raw_anchors)) + if len(anchors) != len(set(anchors)): + raise MaskReuseCalibrationError("topology anchors must be unique") + raw_nearest = raw["nearest"] + if not isinstance(raw_nearest, Mapping): + raise MaskReuseCalibrationError("topology nearest must be an object") + nearest: dict[int, int] = {} + for raw_layer, raw_anchor in raw_nearest.items(): + try: + layer = int(raw_layer) + except (TypeError, ValueError) as error: + raise MaskReuseCalibrationError("topology nearest has a non-integer key") from error + if str(layer) != str(raw_layer) or layer in nearest: + raise MaskReuseCalibrationError("topology nearest keys must be canonical and unique") + nearest[layer] = _integer(raw_anchor, f"topology nearest[{layer}]", minimum=0) + anchor_set = set(anchors) + if not anchor_set <= nearest.keys(): + raise MaskReuseCalibrationError("topology nearest must include every anchor") + for layer, anchor in nearest.items(): + if anchor not in anchor_set or (layer != anchor and anchor >= layer): + raise MaskReuseCalibrationError(f"topology layer {layer} has invalid anchor {anchor}") + if layer in anchor_set and anchor != layer: + raise MaskReuseCalibrationError(f"topology anchor {layer} must map to itself") + if not any(layer != anchor for layer, anchor in nearest.items()): + raise MaskReuseCalibrationError("topology must contain at least one reuse layer") + return anchors, dict(sorted(nearest.items())) + + +def _bucket_key(bucket: Bucket) -> tuple[int, float]: + return bucket[0], math.inf if bucket[1] is None else float(bucket[1]) + + +@dataclass(frozen=True, slots=True) +class _Choice: + donor_head: int + fallback: bool + retained_tiles: int + + +@dataclass(frozen=True, slots=True) +class _Selection: + target_sparsity: float | None + choices: Mapping[ConsumerHead, _Choice] + frontier: tuple[Mapping[str, object], ...] + exact_reason: str | None = None + bmm1_eligible_tiles: int = 0 + bmm1_skipped_tiles: int = 0 + target_bmm1_skip_ratio_met: bool = False + worst_prompt_reuse_dropped_mass: float = 0.0 + mean_prompt_reuse_dropped_mass: float = 0.0 + worst_individual_reuse_dropped_mass: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class _DonorOption: + choice: _Choice + bmm1_skipped_tiles: int + risk_by_prompt: tuple[float, ...] + + +@dataclass(frozen=True, slots=True) +class _BucketIndex: + observations: Mapping[ObservationKey, MaskReuseObservation] + prompts: Mapping[str, tuple[str, ...]] + target_menus: Mapping[tuple[str, str], frozenset[float]] + donor_menus: Mapping[tuple[str, str, float, ConsumerHead], frozenset[int]] + eligible: Mapping[tuple[str, str, ConsumerHead], int] + anchor_masks: Mapping[AnchorKey, tuple[int, int, float]] + anchors: tuple[int, ...] + + +@dataclass(slots=True) +class _ReuseEvaluation: + eligible_tiles: int = 0 + retained_tiles: int = 0 + sparse_observations: int = 0 + violations: int = 0 + dropped_mass_sum: float = 0.0 + worst_dropped_mass: float = 0.0 + + def add(self, other: _ReuseEvaluation) -> None: + self.eligible_tiles += other.eligible_tiles + self.retained_tiles += other.retained_tiles + self.sparse_observations += other.sparse_observations + self.violations += other.violations + self.dropped_mass_sum += other.dropped_mass_sum + self.worst_dropped_mass = max(self.worst_dropped_mass, other.worst_dropped_mass) + + def to_mapping(self) -> dict[str, object]: + return { + "eligible_tiles": self.eligible_tiles, + "retained_tiles": self.retained_tiles, + "bmm1_tile_savings_fraction": ( + 1.0 - self.retained_tiles / self.eligible_tiles if self.eligible_tiles else 0.0 + ), + "sparse_head_prompt_observations": self.sparse_observations, + "report_threshold_exceedance_count": self.violations, + "report_threshold_exceedance_rate": ( + self.violations / self.sparse_observations if self.sparse_observations else 0.0 + ), + "mean_dropped_mass": ( + self.dropped_mass_sum / self.sparse_observations + if self.sparse_observations + else 0.0 + ), + "worst_dropped_mass": self.worst_dropped_mass, + } + + +@dataclass(slots=True) +class _AnchorEvaluation: + eligible_tiles: int = 0 + retained_tiles: int = 0 + prompt_count: int = 0 + violations: int = 0 + prompt_mean_sum: float = 0.0 + worst_prompt_mean: float = 0.0 + + def add(self, other: _AnchorEvaluation) -> None: + self.eligible_tiles += other.eligible_tiles + self.retained_tiles += other.retained_tiles + self.prompt_count += other.prompt_count + self.violations += other.violations + self.prompt_mean_sum += other.prompt_mean_sum + self.worst_prompt_mean = max(self.worst_prompt_mean, other.worst_prompt_mean) + + def to_mapping(self, *, exact: bool = False) -> dict[str, object]: + return { + "policy_exact": exact, + "constraint_statistic": "worst_prompt_mean_anchor_dropped_mass", + "eligible_tiles": self.eligible_tiles, + "retained_tiles": self.retained_tiles, + "bmm2_tile_savings_fraction": ( + 1.0 - self.retained_tiles / self.eligible_tiles if self.eligible_tiles else 0.0 + ), + "evaluated_prompt_count": self.prompt_count, + "constraint_violation_count": self.violations, + "constraint_violation_rate": ( + self.violations / self.prompt_count if self.prompt_count else 0.0 + ), + "mean_prompt_mean_anchor_dropped_mass": ( + self.prompt_mean_sum / self.prompt_count if self.prompt_count else 0.0 + ), + "worst_prompt_mean_anchor_dropped_mass": self.worst_prompt_mean, + } + + +def _normalize_observations( + values: Sequence[MaskReuseObservation | Mapping[str, object]], +) -> list[MaskReuseObservation]: + if not values: + raise MaskReuseCalibrationError("at least one observation is required") + return [ + value + if isinstance(value, MaskReuseObservation) + else MaskReuseObservation.from_mapping(value) + for value in values + ] + + +def _validate_thresholds( + observations: Sequence[MaskReuseObservation], threshold_scale_factor: Mapping[str, object] +) -> None: + params = threshold_scale_factor["prefill"] + assert isinstance(params, Mapping) + a = float(params["a"]) + b = float(params["b"]) + lower = params.get("min_observed_sparsity") + upper = params.get("max_observed_sparsity") + for observation in observations: + target = observation.target_sparsity + if (lower is not None and target < float(lower)) or ( + upper is not None and target > float(upper) + ): + raise MaskReuseCalibrationError( + f"target_sparsity={target} is outside the observed vanilla calibration range" + ) + expected_log2 = ( + math.log2(a) + b * target * math.log2(math.e) - math.log2(observation.sample_length) + ) + expected_lambda = 2.0**expected_log2 + if not 0.0 < expected_lambda < 1.0: + raise MaskReuseCalibrationError( + "vanilla calibration derives a threshold outside (0, 1) for " + f"prompt={observation.prompt_id!r}, target_sparsity={target}" + ) + if observation.threshold_log2.hex() != expected_log2.hex(): + raise MaskReuseCalibrationError( + "threshold_log2 does not match the log-domain vanilla fit for " + f"prompt={observation.prompt_id!r}, target_sparsity={target}: " + f"observed={observation.threshold_log2.hex()}, " + f"expected={expected_log2.hex()}" + ) + if observation.threshold_lambda.hex() != expected_lambda.hex(): + raise MaskReuseCalibrationError( + "threshold_lambda does not match " + "exp2(log2(a) + b * target_sparsity * log2(e) - log2(sample_length)) for " + f"prompt={observation.prompt_id!r}, target_sparsity={target}: " + f"observed={observation.threshold_lambda.hex()}, " + f"expected={expected_lambda.hex()}" + ) + + +def _validate_dataset( + observations: Sequence[MaskReuseObservation], + *, + nearest: Mapping[int, int], +) -> tuple[ + str, int, tuple[Bucket, ...], tuple[ConsumerHead, ...], dict[Bucket, list[MaskReuseObservation]] +]: + by_split = { + split: [observation for observation in observations if observation.split == split] + for split in _SPLITS + } + if any(not rows for rows in by_split.values()): + raise MaskReuseCalibrationError("observations require calibration and heldout splits") + models = {observation.model for observation in observations} + if len(models) != 1: + raise MaskReuseCalibrationError("observations must contain exactly one model") + calibration_prompts = {row.prompt_id for row in by_split["calibration"]} + heldout_prompts = {row.prompt_id for row in by_split["heldout"]} + if calibration_prompts & heldout_prompts: + raise MaskReuseCalibrationError("prompt IDs overlap calibration and heldout splits") + calibration_sources = {row.source_capture_sha256 for row in by_split["calibration"]} + heldout_sources = {row.source_capture_sha256 for row in by_split["heldout"]} + if calibration_sources & heldout_sources: + raise MaskReuseCalibrationError("source captures overlap calibration and heldout splits") + + consumer_to_anchor = {layer: anchor for layer, anchor in nearest.items() if layer != anchor} + max_head = max(max(row.consumer_head, row.donor_head) for row in observations) + global_num_heads = max_head + 1 + targets = tuple( + (layer, head) for layer in sorted(consumer_to_anchor) for head in range(global_num_heads) + ) + expected_targets = set(targets) + seen: set[tuple[object, ...]] = set() + capture_sources: dict[tuple[str, str, Bucket], str] = {} + by_bucket: dict[Bucket, list[MaskReuseObservation]] = defaultdict(list) + for row in observations: + expected_anchor = consumer_to_anchor.get(row.consumer_layer) + if expected_anchor != row.anchor_layer: + raise MaskReuseCalibrationError( + f"consumer layer {row.consumer_layer} does not match the explicit topology" + ) + bucket = (row.min_kv_tokens, row.max_kv_tokens) + by_bucket[bucket].append(row) + identity = ( + row.model, + bucket, + row.split, + row.prompt_id, + row.target_sparsity, + row.consumer_layer, + row.consumer_head, + row.donor_head, + ) + if identity in seen: + raise MaskReuseCalibrationError("observations contain a duplicate candidate row") + seen.add(identity) + capture = (row.split, row.prompt_id, bucket) + previous_source = capture_sources.setdefault(capture, row.source_capture_sha256) + if previous_source != row.source_capture_sha256: + raise MaskReuseCalibrationError("one prompt/context capture has multiple fingerprints") + + split_buckets = { + split: {(row.min_kv_tokens, row.max_kv_tokens) for row in rows} + for split, rows in by_split.items() + } + if split_buckets["calibration"] != split_buckets["heldout"]: + raise MaskReuseCalibrationError("heldout context buckets must match calibration buckets") + buckets = tuple(sorted(split_buckets["calibration"], key=_bucket_key)) + previous_max: int | None = 0 + for minimum, maximum in buckets: + if previous_max is None or minimum <= previous_max: + raise MaskReuseCalibrationError("context buckets must be ordered and non-overlapping") + previous_max = maximum + for bucket in buckets: + for split in _SPLITS: + actual = { + (row.consumer_layer, row.consumer_head) + for row in by_bucket[bucket] + if row.split == split + } + if actual != expected_targets: + raise MaskReuseCalibrationError( + f"{split} bucket {bucket} does not cover every consumer head" + ) + return next(iter(models)), global_num_heads, buckets, targets, dict(by_bucket) + + +def _index_bucket( + rows: Sequence[MaskReuseObservation], + *, + anchors: tuple[int, ...], + global_num_heads: int, + targets: Sequence[ConsumerHead], +) -> _BucketIndex: + observations: dict[ObservationKey, MaskReuseObservation] = {} + prompts: dict[str, set[str]] = defaultdict(set) + target_menus: dict[tuple[str, str], set[float]] = defaultdict(set) + donor_menus: dict[tuple[str, str, float, ConsumerHead], set[int]] = defaultdict(set) + eligible: dict[tuple[str, str, ConsumerHead], int] = {} + anchor_masks: dict[AnchorKey, tuple[int, int, float]] = {} + prompt_targets: dict[tuple[str, str, float], tuple[int, float, float]] = {} + anchor_payloads: dict[tuple[str, str, float], Mapping[int, AnchorLayerStats]] = {} + payload_eligible: dict[tuple[str, str, float], int] = {} + for row in rows: + target = (row.consumer_layer, row.consumer_head) + key: ObservationKey = ( + row.split, + row.prompt_id, + row.target_sparsity, + row.consumer_layer, + row.consumer_head, + row.donor_head, + ) + observations[key] = row + prompts[row.split].add(row.prompt_id) + target_menus[(row.split, row.prompt_id)].add(row.target_sparsity) + donor_menus[(row.split, row.prompt_id, row.target_sparsity, target)].add(row.donor_head) + eligible_key = (row.split, row.prompt_id, target) + previous_eligible = eligible.setdefault(eligible_key, row.eligible_tiles) + if previous_eligible != row.eligible_tiles: + raise MaskReuseCalibrationError("eligible_tiles differs across candidates") + prompt_target = (row.split, row.prompt_id, row.target_sparsity) + sample_and_threshold = ( + row.sample_length, + row.threshold_lambda, + row.threshold_log2, + ) + if prompt_targets.setdefault(prompt_target, sample_and_threshold) != sample_and_threshold: + raise MaskReuseCalibrationError( + "sample_length, threshold_lambda, or threshold_log2 differs within a capture" + ) + previous_payload = anchor_payloads.setdefault(prompt_target, row.anchor_stats_by_layer) + if previous_payload != row.anchor_stats_by_layer: + raise MaskReuseCalibrationError( + "anchor_stats_by_layer differs across repeated candidate rows" + ) + previous_payload_eligible = payload_eligible.setdefault(prompt_target, row.eligible_tiles) + if previous_payload_eligible != row.eligible_tiles: + raise MaskReuseCalibrationError( + "eligible_tiles differs within one capture and target_sparsity" + ) + + expected_anchors = set(anchors) + for prompt_target, payload in anchor_payloads.items(): + actual_anchors = set(payload) + if actual_anchors != expected_anchors: + raise MaskReuseCalibrationError( + "anchor_stats_by_layer does not exactly cover topology anchors; " + f"missing={sorted(expected_anchors - actual_anchors)}, " + f"extra={sorted(actual_anchors - expected_anchors)}" + ) + eligible_tiles = payload_eligible[prompt_target] + for anchor, stats in payload.items(): + if len(stats.retained_tiles) != global_num_heads: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{anchor}] has head width " + f"{len(stats.retained_tiles)}, expected {global_num_heads}" + ) + for head, (retained, dropped) in enumerate( + zip(stats.retained_tiles, stats.dropped_mass, strict=True) + ): + if retained > eligible_tiles: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{anchor}].retained_tiles[{head}] " + "exceeds eligible_tiles" + ) + split, prompt, target_sparsity = prompt_target + anchor_masks[(split, prompt, target_sparsity, anchor, head)] = ( + retained, + eligible_tiles, + dropped, + ) + + for row in rows: + stats = row.anchor_stats_by_layer[row.anchor_layer] + payload_candidate = ( + stats.retained_tiles[row.donor_head], + stats.dropped_mass[row.donor_head], + ) + if payload_candidate != (row.retained_tiles, row.anchor_dropped_mass): + raise MaskReuseCalibrationError( + "candidate retained_tiles/anchor_dropped_mass does not match anchor_stats_by_layer" + ) + + first_prompt = min(prompts["calibration"]) + expected_targets = target_menus[("calibration", first_prompt)] + full_donors = set(range(global_num_heads)) + for split in _SPLITS: + for prompt in prompts[split]: + if target_menus[(split, prompt)] != expected_targets: + raise MaskReuseCalibrationError("target_sparsity menu differs across captures") + for target_sparsity in expected_targets: + for target in targets: + if donor_menus[(split, prompt, target_sparsity, target)] != full_donors: + raise MaskReuseCalibrationError("candidate donor menu is incomplete") + for anchor in anchors: + for head in range(global_num_heads): + if (split, prompt, target_sparsity, anchor, head) not in anchor_masks: + raise MaskReuseCalibrationError( + "anchor/head mask observations are incomplete" + ) + return _BucketIndex( + observations=observations, + prompts={split: tuple(sorted(values)) for split, values in prompts.items()}, + target_menus={key: frozenset(values) for key, values in target_menus.items()}, + donor_menus={key: frozenset(values) for key, values in donor_menus.items()}, + eligible=eligible, + anchor_masks=anchor_masks, + anchors=anchors, + ) + + +def _evaluate_anchor( + index: _BucketIndex, + *, + split: str, + target_sparsity: float | None, + global_num_heads: int, + maximum: float, +) -> _AnchorEvaluation: + result = _AnchorEvaluation() + for prompt in index.prompts[split]: + selected_target = ( + min(index.target_menus[(split, prompt)]) if target_sparsity is None else target_sparsity + ) + values: list[float] = [] + retained_tiles = 0 + eligible_tiles = 0 + for anchor in index.anchors: + for head in range(global_num_heads): + retained, eligible, dropped = index.anchor_masks[ + (split, prompt, selected_target, anchor, head) + ] + values.append(0.0 if target_sparsity is None else dropped) + retained_tiles += eligible if target_sparsity is None else retained + eligible_tiles += eligible + prompt_mean = sum(values) / len(values) + result.eligible_tiles += eligible_tiles + result.retained_tiles += retained_tiles + result.prompt_count += 1 + result.prompt_mean_sum += prompt_mean + result.worst_prompt_mean = max(result.worst_prompt_mean, prompt_mean) + result.violations += int(prompt_mean > maximum) + return result + + +def _select_bucket( + index: _BucketIndex, + *, + targets: Sequence[ConsumerHead], + global_num_heads: int, + max_anchor_dropped_mass: float, + target_bmm1_skip_ratio: float, +) -> _Selection: + prompts = index.prompts["calibration"] + target_menu = tuple(sorted(index.target_menus[("calibration", prompts[0])])) + frontier: list[Mapping[str, object]] = [] + candidates: list[tuple[tuple[object, ...], _Selection]] = [] + maximum_candidates: list[tuple[tuple[object, ...], _Selection]] = [] + + def donor_options(target_sparsity: float, target: ConsumerHead) -> tuple[_DonorOption, ...]: + eligible = sum(index.eligible[("calibration", prompt, target)] for prompt in prompts) + options = [ + _DonorOption( + _Choice(0, True, eligible), + 0, + tuple(0.0 for _ in prompts), + ) + ] + for donor in range(global_num_heads): + rows = [ + index.observations[ + ( + "calibration", + prompt, + target_sparsity, + target[0], + target[1], + donor, + ) + ] + for prompt in prompts + ] + retained = sum(row.retained_tiles for row in rows) + options.append( + _DonorOption( + _Choice(donor, False, retained), + eligible - retained, + tuple(row.dropped_mass for row in rows), + ) + ) + pareto = [] + for candidate_index, candidate in enumerate(options): + dominated = False + for other_index, other in enumerate(options): + if candidate_index == other_index: + continue + no_worse = other.bmm1_skipped_tiles >= candidate.bmm1_skipped_tiles and all( + other_risk <= candidate_risk + for other_risk, candidate_risk in zip( + other.risk_by_prompt, + candidate.risk_by_prompt, + strict=True, + ) + ) + strictly_better = other.bmm1_skipped_tiles > candidate.bmm1_skipped_tiles or any( + other_risk < candidate_risk + for other_risk, candidate_risk in zip( + other.risk_by_prompt, + candidate.risk_by_prompt, + strict=True, + ) + ) + canonical_tie = not strictly_better and ( + (other.choice.fallback and not candidate.choice.fallback) + or ( + other.choice.fallback == candidate.choice.fallback + and other.choice.donor_head < candidate.choice.donor_head + ) + ) + if no_worse and (strictly_better or canonical_tie): + dominated = True + break + if not dominated: + pareto.append(candidate) + return tuple(pareto) + + def solve_target( + target_sparsity: float, + *, + minimum_bmm1_skipped_tiles: int | None, + maximize_bmm1_skipped_tiles: bool, + target_met: bool, + ) -> _Selection | None: + problem = pulp.LpProblem("legacy_mask_reuse", pulp.LpMinimize) + variables: dict[tuple[ConsumerHead, int], pulp.LpVariable] = {} + options: dict[tuple[ConsumerHead, int], _DonorOption] = {} + prompt_risk_terms: dict[str, list[object]] = defaultdict(list) + for target_index, target in enumerate(targets): + menu = donor_options(target_sparsity, target) + choice_variables = [] + for option_index, option in enumerate(menu): + key = (target, option_index) + variable = pulp.LpVariable( + f"choice_{target_index}_{option_index}", + lowBound=0, + upBound=1, + cat="Binary", + ) + variables[key] = variable + options[key] = option + choice_variables.append(variable) + for prompt, risk in zip(prompts, option.risk_by_prompt, strict=True): + prompt_risk_terms[prompt].append(risk * variable) + problem += pulp.lpSum(choice_variables) == 1, f"choose_{target_index}" + + bmm1_skipped = pulp.lpSum( + options[key].bmm1_skipped_tiles * variable for key, variable in variables.items() + ) + retained_reuse = pulp.lpSum( + options[key].choice.retained_tiles * variable for key, variable in variables.items() + ) + reuse_count = pulp.lpSum( + int(not options[key].choice.fallback) * variable for key, variable in variables.items() + ) + all_layer_head_count = len(index.anchors) * global_num_heads + len(targets) + worst_prompt_risk = pulp.LpVariable("worst_prompt_reuse_dropped_mass", lowBound=0.0) + for prompt_index, prompt in enumerate(prompts): + problem += ( + pulp.lpSum(prompt_risk_terms[prompt]) <= all_layer_head_count * worst_prompt_risk, + f"reuse_risk_{prompt_index}", + ) + if minimum_bmm1_skipped_tiles is not None: + problem += bmm1_skipped >= minimum_bmm1_skipped_tiles, "minimum_bmm1_skips" + solver = pulp.PULP_CBC_CMD(msg=False, threads=1, options=["randomSeed 0"]) + warm_solver = pulp.PULP_CBC_CMD( + msg=False, + threads=1, + options=["randomSeed 0"], + warmStart=True, + ) + has_incumbent = False + + def minimize(expression: object) -> bool: + nonlocal has_incumbent + problem.setObjective(expression) + status = problem.solve(warm_solver if has_incumbent else solver) + has_incumbent = status == pulp.LpStatusOptimal + return has_incumbent + + def minimize_and_fix(expression: object, name: str, *, integral: bool) -> bool: + nonlocal problem + if not minimize(expression): + return False + raw_value = pulp.value(expression) + value = 0.0 if raw_value is None else float(raw_value) + if integral: + problem += expression == round(value), name + else: + problem += expression <= value + _SOLVER_LEXICOGRAPHIC_ATOL, name + return True + + if maximize_bmm1_skipped_tiles and not minimize_and_fix( + -bmm1_skipped, "fix_maximum_bmm1_skips", integral=True + ): + return None + if not minimize_and_fix(worst_prompt_risk, "fix_worst_prompt_risk", integral=False): + return None + if not minimize_and_fix(retained_reuse, "fix_retained_reuse", integral=True): + raise MaskReuseCalibrationError( + "legacy selector lost feasibility after fixing worst-prompt reuse risk" + ) + donor_signature = pulp.lpSum( + (options[key].choice.donor_head + 1) * variable + for key, variable in variables.items() + if not options[key].choice.fallback + ) + donor_base = global_num_heads * len(targets) + 1 + if not minimize(reuse_count * donor_base + donor_signature): + raise MaskReuseCalibrationError( + "legacy selector lost deterministic tie-break feasibility" + ) + + choices: dict[ConsumerHead, _Choice] = {} + prompt_totals = dict.fromkeys(prompts, 0.0) + worst_individual = 0.0 + skipped_tiles = 0 + for key, variable in variables.items(): + if variable.value() <= 0.5: + continue + target, _ = key + option = options[key] + choices[target] = option.choice + skipped_tiles += option.bmm1_skipped_tiles + if not option.choice.fallback: + for prompt, risk in zip(prompts, option.risk_by_prompt, strict=True): + prompt_totals[prompt] += risk + worst_individual = max(worst_individual, risk) + eligible_per_head = sum( + index.eligible[("calibration", prompt, targets[0])] for prompt in prompts + ) + eligible_tiles = eligible_per_head * all_layer_head_count + prompt_risks = [prompt_totals[prompt] / all_layer_head_count for prompt in prompts] + return _Selection( + target_sparsity, + choices, + (), + bmm1_eligible_tiles=eligible_tiles, + bmm1_skipped_tiles=skipped_tiles, + target_bmm1_skip_ratio_met=target_met, + worst_prompt_reuse_dropped_mass=max(prompt_risks, default=0.0), + mean_prompt_reuse_dropped_mass=( + sum(prompt_risks) / len(prompt_risks) if prompt_risks else 0.0 + ), + worst_individual_reuse_dropped_mass=worst_individual, + ) + + for target_sparsity in target_menu: + anchor = _evaluate_anchor( + index, + split="calibration", + target_sparsity=target_sparsity, + global_num_heads=global_num_heads, + maximum=max_anchor_dropped_mass, + ) + eligible_per_head = sum( + index.eligible[("calibration", prompt, targets[0])] for prompt in prompts + ) + eligible_tiles = eligible_per_head * (len(index.anchors) * global_num_heads + len(targets)) + required_tiles = math.ceil(target_bmm1_skip_ratio * eligible_tiles) + selected = None + if anchor.violations == 0: + selected = solve_target( + target_sparsity, + minimum_bmm1_skipped_tiles=required_tiles, + maximize_bmm1_skipped_tiles=False, + target_met=True, + ) + total_retained = ( + None + if selected is None + else sum(choice.retained_tiles for choice in selected.choices.values()) + ) + fallback_count = ( + None + if selected is None + else sum(choice.fallback for choice in selected.choices.values()) + ) + combined_tile_cost = ( + None if total_retained is None else 2 * total_retained + anchor.retained_tiles + ) + frontier.append( + { + "target_sparsity": target_sparsity, + "anchor_safe": anchor.violations == 0, + "target_bmm1_skip_ratio": target_bmm1_skip_ratio, + "target_bmm1_skip_ratio_feasible": selected is not None, + "retained_reuse_tiles": total_retained, + "retained_anchor_tiles": anchor.retained_tiles, + "combined_tile_cost": combined_tile_cost, + "fallback_head_count": fallback_count, + "anchor_calibration": anchor.to_mapping(), + } + ) + if selected is not None: + rank = ( + selected.worst_prompt_reuse_dropped_mass, + combined_tile_cost, + target_sparsity, + ) + candidates.append((rank, selected)) + elif anchor.violations == 0: + maximum = solve_target( + target_sparsity, + minimum_bmm1_skipped_tiles=None, + maximize_bmm1_skipped_tiles=True, + target_met=False, + ) + if maximum is not None: + maximum_candidates.append( + ( + ( + -maximum.bmm1_skipped_tiles, + maximum.worst_prompt_reuse_dropped_mass, + target_sparsity, + ), + maximum, + ) + ) + if not candidates and not maximum_candidates: + dense_choices = { + target: _Choice( + 0, True, sum(index.eligible[("calibration", prompt, target)] for prompt in prompts) + ) + for target in targets + } + return _Selection( + None, + dense_choices, + tuple(frontier), + "no_target_sparsity_satisfied_anchor_calibration_constraint", + bmm1_eligible_tiles=( + sum(index.eligible[("calibration", prompt, targets[0])] for prompt in prompts) + * (len(index.anchors) * global_num_heads + len(targets)) + ), + ) + if candidates: + _, selected = min(candidates, key=lambda item: item[0]) + else: + _, selected = min(maximum_candidates, key=lambda item: item[0]) + return _Selection( + selected.target_sparsity, + selected.choices, + tuple(frontier), + bmm1_eligible_tiles=selected.bmm1_eligible_tiles, + bmm1_skipped_tiles=selected.bmm1_skipped_tiles, + target_bmm1_skip_ratio_met=selected.target_bmm1_skip_ratio_met, + worst_prompt_reuse_dropped_mass=selected.worst_prompt_reuse_dropped_mass, + mean_prompt_reuse_dropped_mass=selected.mean_prompt_reuse_dropped_mass, + worst_individual_reuse_dropped_mass=selected.worst_individual_reuse_dropped_mass, + ) + + +def _evaluate_reuse( + selection: _Selection, + index: _BucketIndex, + *, + split: str, + maximum: float, +) -> _ReuseEvaluation: + result = _ReuseEvaluation() + for prompt in index.prompts[split]: + for target, choice in sorted(selection.choices.items()): + eligible = index.eligible[(split, prompt, target)] + result.eligible_tiles += eligible + if choice.fallback: + result.retained_tiles += eligible + continue + assert selection.target_sparsity is not None + row = index.observations[ + ( + split, + prompt, + selection.target_sparsity, + target[0], + target[1], + choice.donor_head, + ) + ] + result.retained_tiles += row.retained_tiles + result.sparse_observations += 1 + result.dropped_mass_sum += row.dropped_mass + result.worst_dropped_mass = max(result.worst_dropped_mass, row.dropped_mass) + result.violations += int(row.dropped_mass > maximum) + return result + + +def _canonical_digest(observations: Sequence[MaskReuseObservation]) -> str: + rows = [ + json.dumps(row.to_mapping(), sort_keys=True, separators=(",", ":")).encode() + for row in observations + ] + digest = sha256() + for row in sorted(rows): + digest.update(sha256(row).digest()) + return digest.hexdigest() + + +def _deployment_geometry( + observations: Sequence[MaskReuseObservation], +) -> dict[str, object]: + by_capture: dict[tuple[object, ...], tuple[int, int, int]] = {} + for row in observations: + identity = ( + row.split, + row.prompt_id, + row.source_capture_sha256, + row.min_kv_tokens, + row.max_kv_tokens, + ) + geometry = (row.q_tokens, row.kv_tokens, row.q_start_tokens) + if by_capture.setdefault(identity, geometry) != geometry: + raise MaskReuseCalibrationError( + "q_tokens, kv_tokens, or q_start_tokens differs within one capture" + ) + geometry_rows = [ + { + "split": identity[0], + "prompt_id": identity[1], + "source_capture_sha256": identity[2], + "min_kv_tokens": identity[3], + "max_kv_tokens": identity[4], + "q_tokens": geometry[0], + "kv_tokens": geometry[1], + "q_start_tokens": geometry[2], + } + for identity, geometry in sorted(by_capture.items(), key=lambda item: str(item[0])) + ] + return { + "contract": dict(_DEPLOYMENT_GEOMETRY_CONTRACT), + "observations": geometry_rows, + } + + +def _canonical_evidence(raw: Mapping[str, object]) -> dict[str, str]: + missing = _EVIDENCE_FIELDS - raw.keys() + extra = raw.keys() - _EVIDENCE_FIELDS + if missing or extra: + raise MaskReuseCalibrationError( + f"evidence fields do not match schema v3; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + return {field: _sha256(raw[field], f"evidence.{field}") for field in sorted(raw)} + + +def calibrate_mask_reuse_policy( + observations: Sequence[MaskReuseObservation | Mapping[str, object]], + *, + vanilla_calibration: Mapping[str, object], + topology: Mapping[str, object], + checkpoint_manifest: VerifiedCheckpointManifest, + evidence: Mapping[str, object], + max_anchor_dropped_mass: float, + reuse_dropped_mass_report_threshold: float, + target_bmm1_skip_ratio: float, + source_provenance: Mapping[str, object] | None = None, +) -> dict[str, object]: + """Select a minimum-risk legacy candidate under a per-bucket BMM1 target.""" + if not isinstance(checkpoint_manifest, VerifiedCheckpointManifest): + raise MaskReuseCalibrationError( + "checkpoint_manifest must be returned by verify_checkpoint_manifest" + ) + rows = _normalize_observations(observations) + threshold_scale_factor = canonical_prefill_threshold_scale_factor(vanilla_calibration) + prefill_params = threshold_scale_factor["prefill"] + assert isinstance(prefill_params, Mapping) + if not {"min_observed_sparsity", "max_observed_sparsity"} <= prefill_params.keys(): + raise MaskReuseCalibrationError( + "schema-v3 export requires observed prefill sparsity bounds" + ) + _validate_thresholds(rows, threshold_scale_factor) + anchors, nearest = _normalize_topology(topology) + checkpoint_identity = checkpoint_manifest.sha256 + geometry = _deployment_geometry(rows) + canonical_evidence = _canonical_evidence(evidence) + anchor_bound = _number( + max_anchor_dropped_mass, "max_anchor_dropped_mass", minimum=0.0, maximum=1.0 + ) + reuse_report_threshold = _number( + reuse_dropped_mass_report_threshold, + "reuse_dropped_mass_report_threshold", + minimum=0.0, + maximum=1.0, + ) + bmm1_target = _number( + target_bmm1_skip_ratio, + "target_bmm1_skip_ratio", + minimum=0.0, + maximum=1.0, + ) + model, global_num_heads, buckets, targets, by_bucket = _validate_dataset(rows, nearest=nearest) + if model != checkpoint_manifest.model: + raise MaskReuseCalibrationError( + "observation model does not match the verified checkpoint manifest" + ) + + context_policies: list[dict[str, object]] = [] + bucket_reports: list[dict[str, object]] = [] + target_menus: list[dict[str, object]] = [] + overall_reuse_calibration = _ReuseEvaluation() + overall_reuse_heldout = _ReuseEvaluation() + overall_anchor_calibration = _AnchorEvaluation() + overall_anchor_heldout = _AnchorEvaluation() + total_fallback = 0 + total_bmm1_eligible = 0 + total_bmm1_skipped = 0 + all_bmm1_targets_met = True + for bounds in buckets: + index = _index_bucket( + by_bucket[bounds], + anchors=anchors, + global_num_heads=global_num_heads, + targets=targets, + ) + selection = _select_bucket( + index, + targets=targets, + global_num_heads=global_num_heads, + max_anchor_dropped_mass=anchor_bound, + target_bmm1_skip_ratio=bmm1_target, + ) + if selection.target_sparsity is not None and bounds[1] is None: + raise MaskReuseCalibrationError( + "a deployment-qualified sparse context bucket requires a finite maximum" + ) + reuse_calibration = _evaluate_reuse( + selection, index, split="calibration", maximum=reuse_report_threshold + ) + reuse_heldout = _evaluate_reuse( + selection, index, split="heldout", maximum=reuse_report_threshold + ) + anchor_calibration = _evaluate_anchor( + index, + split="calibration", + target_sparsity=selection.target_sparsity, + global_num_heads=global_num_heads, + maximum=anchor_bound, + ) + anchor_heldout = _evaluate_anchor( + index, + split="heldout", + target_sparsity=selection.target_sparsity, + global_num_heads=global_num_heads, + maximum=anchor_bound, + ) + overall_reuse_calibration.add(reuse_calibration) + overall_reuse_heldout.add(reuse_heldout) + overall_anchor_calibration.add(anchor_calibration) + overall_anchor_heldout.add(anchor_heldout) + + policy: dict[str, object] + if selection.target_sparsity is None: + policy = { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "exact": True, + } + headmaps: dict[str, list[int]] = {} + fallback_heads: dict[str, list[int]] = {} + else: + headmaps = { + str(layer): [ + selection.choices[(layer, head)].donor_head for head in range(global_num_heads) + ] + for layer in sorted({target[0] for target in targets}) + } + fallback_heads = { + str(layer): [ + head + for head in range(global_num_heads) + if selection.choices[(layer, head)].fallback + ] + for layer in sorted({target[0] for target in targets}) + } + policy = { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "target_sparsity": selection.target_sparsity, + "headmaps": headmaps, + "fallback_heads": fallback_heads, + } + context_policies.append(policy) + fallback_count = sum(len(heads) for heads in fallback_heads.values()) + total_fallback += fallback_count + total_bmm1_eligible += selection.bmm1_eligible_tiles + total_bmm1_skipped += selection.bmm1_skipped_tiles + all_bmm1_targets_met &= selection.target_bmm1_skip_ratio_met + menu = [row["target_sparsity"] for row in selection.frontier] + target_menus.append( + { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "target_sparsities": menu, + } + ) + bucket_reports.append( + { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "selected_target_sparsity": selection.target_sparsity, + "selection_status": ( + "target_bmm1_skip_ratio_met" + if selection.target_sparsity is not None + and selection.target_bmm1_skip_ratio_met + else ( + "target_bmm1_skip_ratio_unmet_maximum_feasible" + if selection.target_sparsity is not None + else selection.exact_reason + ) + ), + "target_sparsity_frontier": list(selection.frontier), + "fallback_head_count": fallback_count, + "bmm1_skip_objective": { + "target": bmm1_target, + "target_met": selection.target_bmm1_skip_ratio_met, + "eligible_tiles": selection.bmm1_eligible_tiles, + "skipped_tiles": selection.bmm1_skipped_tiles, + "achieved": ( + selection.bmm1_skipped_tiles / selection.bmm1_eligible_tiles + if selection.bmm1_eligible_tiles + else 0.0 + ), + }, + "reuse_selection_objective": { + "hard_maximum": None, + "worst_development_prompt_model_wide_dropped_mass": ( + selection.worst_prompt_reuse_dropped_mass + ), + "mean_development_prompt_model_wide_dropped_mass": ( + selection.mean_prompt_reuse_dropped_mass + ), + "worst_individual_dropped_mass": ( + selection.worst_individual_reuse_dropped_mass + ), + }, + "reuse_calibration": reuse_calibration.to_mapping(), + "reuse_heldout": reuse_heldout.to_mapping(), + "anchor_calibration": anchor_calibration.to_mapping( + exact=selection.target_sparsity is None + ), + "anchor_heldout": anchor_heldout.to_mapping( + exact=selection.target_sparsity is None + ), + } + ) + + provenance: dict[str, object] = { + "calibrator": "modelopt.mask_reuse", + "observation_schema_version": 1, + "input_observation_count": len(rows), + "canonical_input_sha256": _canonical_digest(rows), + "calibration_prompt_ids": sorted( + {row.prompt_id for row in rows if row.split == "calibration"} + ), + "heldout_prompt_ids": sorted({row.prompt_id for row in rows if row.split == "heldout"}), + "selection_split": "calibration", + "evaluation_split": "heldout", + "threshold_semantics": "a * exp(b * target_sparsity) / sample_length", + "threshold_implementation": { + "threshold_log2": ("log2(a) + b * target_sparsity * log2(e) - log2(sample_length)"), + "threshold_lambda": "exp2(threshold_log2)", + "validation": "exact IEEE-754 binary64 hex equality", + }, + "checkpoint_manifest_sha256": checkpoint_identity, + "target_sparsity_menus": target_menus, + "constraints": { + "anchor": { + "metric": "worst_prompt_mean_anchor_dropped_mass", + "comparison": "<=", + "maximum": anchor_bound, + }, + "reuse": { + "selection_metric": ( + "per_prompt_mean_across_all_attention_layers_and_heads_reuse_dropped_mass" + ), + "selection_hard_maximum": None, + "report_metric": "per_prompt_candidate_reuse_dropped_mass", + "report_threshold": reuse_report_threshold, + "report_threshold_affects_selection": False, + }, + "bmm1_skip_ratio": { + "metric": "per_context_bucket_model_wide_eligible_bmm1_tile_skip_ratio", + "comparison": ">=", + "target": bmm1_target, + }, + }, + "tie_breaks": [ + "reject target sparsities exceeding the calibration anchor bound", + "meet the requested BMM1 skip ratio in every context bucket", + "minimum worst-prompt model-wide reuse dropped mass", + "minimum equal-BMM combined tile cost 2*A_R + A_A", + "fewest reused consumer heads", + "smallest target sparsity", + "lexicographically smallest donor-head map", + ], + "selection_cost": { + "formula": "2 * retained_reuse_tiles + retained_anchor_tiles", + "bmm1_weight": 1.0, + "bmm2_weight": 1.0, + }, + } + if source_provenance is not None: + provenance["source"] = dict(source_provenance) + + report = { + "model": model, + "checkpoint_manifest_sha256": checkpoint_identity, + "constraints": provenance["constraints"], + "selection_unit": "context_bucket", + "by_bucket": bucket_reports, + "overall": { + "consumer_head_bucket_count": len(targets) * len(buckets), + "fallback_head_bucket_count": total_fallback, + "fallback_fraction": total_fallback / (len(targets) * len(buckets)), + "bmm1_skip_objective": { + "target_per_context_bucket": bmm1_target, + "all_context_buckets_met": all_bmm1_targets_met, + "eligible_tiles": total_bmm1_eligible, + "skipped_tiles": total_bmm1_skipped, + "achieved_model_wide": ( + total_bmm1_skipped / total_bmm1_eligible if total_bmm1_eligible else 0.0 + ), + }, + "reuse_calibration": overall_reuse_calibration.to_mapping(), + "reuse_heldout": overall_reuse_heldout.to_mapping(), + "anchor_calibration": overall_anchor_calibration.to_mapping(), + "anchor_heldout": overall_anchor_heldout.to_mapping(), + }, + "promotion": { + "status": "candidate_only", + "eligible": False, + "reasons": [ + "legacy observations are not capture-schema-v2 checkpoint-bound records", + "grouped inner-fold and preregistered outer gates were not evaluated", + "deployment rectangular-geometry promotion gate was not evaluated", + ], + }, + } + return { + "version": 3, + "promotion_status": "candidate_only", + "phase": "prefill", + "decode": {"mode": "dense"}, + "calibration_protocol": _CALIBRATION_PROTOCOL, + "producer": {"name": "modelopt", "version": modelopt.__version__}, + "evidence": canonical_evidence, + "threshold_scale_factor": threshold_scale_factor, + "model": model, + "checkpoint_manifest_sha256": checkpoint_identity, + "global_num_heads": global_num_heads, + "target_bmm1_skip_ratio": bmm1_target, + "anchors": list(anchors), + "nearest": {str(layer): anchor for layer, anchor in nearest.items()}, + "deployment_geometry_validated": False, + "deployment_geometry": geometry, + "context_policies": context_policies, + "provenance": provenance, + "calibration_report": report, + } diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py new file mode 100644 index 00000000000..5c7c38a51fd --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py @@ -0,0 +1,1545 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Streaming selector for compact mask-reuse capture JSONL. + +One compact record stores the anchor vectors and consumer ``[H, H]`` risk +matrices for a single prompt and target sparsity. The selector makes three +semantic passes (validation, calibration selection, frozen-policy evaluation) +and hashes the exact file bytes before and after selection/evaluation to detect +concurrent mutation. Selection retains development risk matrices long enough to +optimize the worst prompt under the requested BMM1 target; held-out captures +remain evaluation-only. +""" + +from __future__ import annotations + +import json +import math +import unicodedata +from collections import defaultdict +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import cast + +import numpy as np +import pulp + +import modelopt + +from .checkpoint_manifest import VerifiedCheckpointManifest +from .mask_reuse import ( + _CALIBRATION_PROTOCOL, + _DEPLOYMENT_GEOMETRY_CONTRACT, + _SOLVER_LEXICOGRAPHIC_ATOL, + AnchorLayerStats, + Bucket, + ConsumerHead, + MaskReuseCalibrationError, + _AnchorEvaluation, + _bucket_key, + _canonical_evidence, + _Choice, + _normalize_topology, + _number, + _ReuseEvaluation, + _Selection, + _sha256, + canonical_prefill_threshold_scale_factor, +) + +__all__ = [ + "CompactMaskReuseCapture", + "CompactMaskReuseCaptureSource", + "calibrate_compact_mask_reuse_policy", + "load_compact_mask_reuse_captures", +] + +_CAPTURE_FIELDS = frozenset( + { + "compact_capture_schema_version", + "invocation", + "geometry", + "global_num_heads", + "eligible_tiles", + "anchor_stats_by_layer", + "consumer_layers", + } +) +_INVOCATION_FIELDS = frozenset( + { + "capture_schema_version", + "model", + "checkpoint_manifest_sha256", + "split", + "partition", + "inner_fold", + "prompt_id", + "source", + "source_group_sha256", + "source_capture_sha256", + "min_kv_tokens", + "max_kv_tokens", + "target_sparsity_hex", + "sample_length", + "threshold_log2_hex", + "threshold_lambda_hex", + "expected_geometry", + } +) +_GEOMETRY_FIELDS = frozenset({"q_tokens", "kv_tokens", "q_start_tokens"}) +_ANCHOR_FIELDS = frozenset({"retained_tiles", "dropped_mass"}) +_CONSUMER_FIELDS = frozenset({"anchor_layer", "dropped_mass"}) +_SPLITS = ("calibration", "heldout") +_PARTITIONS = ("development", "outer_test") +_MONOTONIC_ATOL = 1e-7 + + +def _exact_fields(raw: Mapping[str, object], expected: frozenset[str], label: str) -> None: + missing = expected - raw.keys() + extra = raw.keys() - expected + if missing or extra: + raise MaskReuseCalibrationError( + f"{label} fields do not match the schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + + +def _integer(value: object, label: str, *, minimum: int = 0) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise MaskReuseCalibrationError(f"{label} must be an integer >= {minimum}") + return value + + +def _text(value: object, label: str) -> str: + if not isinstance(value, str) or not value: + raise MaskReuseCalibrationError(f"{label} must be a non-empty string") + if unicodedata.normalize("NFC", value) != value or any( + ord(character) < 32 for character in value + ): + raise MaskReuseCalibrationError(f"{label} must be NFC text without control characters") + return value + + +def _float_hex(value: object, label: str) -> float: + if not isinstance(value, str): + raise MaskReuseCalibrationError(f"{label} must be a canonical float.hex string") + try: + parsed = float.fromhex(value) + except ValueError as error: + raise MaskReuseCalibrationError(f"{label} must be a canonical float.hex string") from error + if not math.isfinite(parsed) or parsed.hex() != value: + raise MaskReuseCalibrationError(f"{label} must be a canonical finite float.hex string") + return parsed + + +def _geometry(raw: object, label: str) -> dict[str, int]: + if not isinstance(raw, Mapping): + raise MaskReuseCalibrationError(f"{label} must be an object") + _exact_fields(raw, _GEOMETRY_FIELDS, label) + q_tokens = _integer(raw["q_tokens"], f"{label}.q_tokens", minimum=129) + if q_tokens > int(cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["max_query_chunk_tokens"])): + raise MaskReuseCalibrationError(f"{label}.q_tokens exceeds the deployment maximum") + kv_tokens = _integer(raw["kv_tokens"], f"{label}.kv_tokens", minimum=q_tokens) + q_start = _integer(raw["q_start_tokens"], f"{label}.q_start_tokens") + alignment = int( + cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["query_chunk_start_alignment_tokens"]) + ) + if q_start % alignment or q_start + q_tokens != kv_tokens: + raise MaskReuseCalibrationError(f"{label} is not a bottom-right aligned final chunk") + return {"q_tokens": q_tokens, "kv_tokens": kv_tokens, "q_start_tokens": q_start} + + +def _eligible_tiles(geometry: Mapping[str, int]) -> int: + q_blocks = (geometry["q_tokens"] + 127) // 128 + first_eligible = geometry["q_start_tokens"] // 128 + 1 + return q_blocks * (2 * first_eligible + q_blocks - 1) // 2 + + +@dataclass(frozen=True, slots=True) +class CompactConsumerStats: + """One consumer layer's global consumer-by-donor dropped-mass matrix.""" + + anchor_layer: int + dropped_mass: Sequence[Sequence[float]] + + +@dataclass(frozen=True, slots=True) +class CompactMaskReuseCapture: + """One prompt/target capture decoded without expanding candidate rows.""" + + model: str + checkpoint_manifest_sha256: str + split: str + partition: str + inner_fold: int | None + prompt_id: str + source: str + source_group_sha256: str + source_capture_sha256: str + min_kv_tokens: int + max_kv_tokens: int | None + target_sparsity: float + sample_length: int + threshold_log2: float + threshold_lambda: float + geometry: Mapping[str, int] + global_num_heads: int + eligible_tiles: int + anchor_stats_by_layer: Mapping[int, AnchorLayerStats] + consumer_layers: Mapping[int, CompactConsumerStats] + + @property + def bucket(self) -> Bucket: + """Return this capture's context bounds.""" + return self.min_kv_tokens, self.max_kv_tokens + + @classmethod + def from_mapping(cls, raw: Mapping[str, object]) -> CompactMaskReuseCapture: + """Parse one strict compact-capture object.""" + _exact_fields(raw, _CAPTURE_FIELDS, "compact capture") + if raw["compact_capture_schema_version"] != 1: + raise MaskReuseCalibrationError("compact_capture_schema_version must be 1") + invocation = raw["invocation"] + if not isinstance(invocation, Mapping): + raise MaskReuseCalibrationError("compact capture invocation must be an object") + _exact_fields(invocation, _INVOCATION_FIELDS, "compact capture invocation") + if invocation["capture_schema_version"] != 2: + raise MaskReuseCalibrationError("capture_schema_version must be 2") + split = invocation["split"] + if split not in _SPLITS: + raise MaskReuseCalibrationError("capture split must be calibration or heldout") + partition = invocation["partition"] + if partition not in _PARTITIONS: + raise MaskReuseCalibrationError("capture partition must be development or outer_test") + expected_split = "calibration" if partition == "development" else "heldout" + if split != expected_split: + raise MaskReuseCalibrationError("capture split and partition disagree") + raw_fold = invocation["inner_fold"] + if partition == "development": + inner_fold = _integer(raw_fold, "inner_fold") + elif raw_fold is not None: + raise MaskReuseCalibrationError("outer_test capture must have null inner_fold") + else: + inner_fold = None + minimum = _integer(invocation["min_kv_tokens"], "min_kv_tokens", minimum=1) + maximum = invocation["max_kv_tokens"] + if maximum is not None: + maximum = _integer(maximum, "max_kv_tokens", minimum=minimum) + sample_length = _integer(invocation["sample_length"], "sample_length", minimum=1) + if sample_length < minimum or (maximum is not None and sample_length > maximum): + raise MaskReuseCalibrationError("sample_length lies outside its context bucket") + target = _float_hex(invocation["target_sparsity_hex"], "target_sparsity_hex") + threshold_log2 = _float_hex(invocation["threshold_log2_hex"], "threshold_log2_hex") + threshold_lambda = _float_hex(invocation["threshold_lambda_hex"], "threshold_lambda_hex") + if not 0.0 < target < 1.0: + raise MaskReuseCalibrationError("target_sparsity must be in (0, 1)") + if threshold_log2 >= 0.0 or not 0.0 < threshold_lambda < 1.0: + raise MaskReuseCalibrationError("threshold must be in (0, 1)") + if (2.0**threshold_log2).hex() != threshold_lambda.hex(): + raise MaskReuseCalibrationError("threshold lambda and log2 fields disagree") + expected_geometry = _geometry(invocation["expected_geometry"], "expected_geometry") + observed_geometry = _geometry(raw["geometry"], "geometry") + if ( + expected_geometry != observed_geometry + or observed_geometry["kv_tokens"] != sample_length + ): + raise MaskReuseCalibrationError("capture geometry does not match its invocation") + global_num_heads = _integer(raw["global_num_heads"], "global_num_heads", minimum=1) + eligible_tiles = _integer(raw["eligible_tiles"], "eligible_tiles", minimum=1) + if eligible_tiles != _eligible_tiles(observed_geometry): + raise MaskReuseCalibrationError( + "eligible_tiles does not match 128x128 bottom-right causal geometry" + ) + + raw_anchors = raw["anchor_stats_by_layer"] + if not isinstance(raw_anchors, Mapping) or not raw_anchors: + raise MaskReuseCalibrationError("anchor_stats_by_layer must be non-empty") + anchors: dict[int, AnchorLayerStats] = {} + for raw_layer, raw_stats in raw_anchors.items(): + if ( + not isinstance(raw_layer, str) + or not raw_layer.isdigit() + or raw_layer != str(int(raw_layer)) + ): + raise MaskReuseCalibrationError("anchor layer keys must be canonical integers") + layer = int(raw_layer) + if not isinstance(raw_stats, Mapping): + raise MaskReuseCalibrationError(f"anchor_stats_by_layer[{layer}] must be an object") + _exact_fields(raw_stats, _ANCHOR_FIELDS, f"anchor_stats_by_layer[{layer}]") + retained_raw = raw_stats["retained_tiles"] + dropped_raw = raw_stats["dropped_mass"] + if not isinstance(retained_raw, list) or not isinstance(dropped_raw, list): + raise MaskReuseCalibrationError(f"anchor_stats_by_layer[{layer}] must use arrays") + if len(retained_raw) != global_num_heads or len(dropped_raw) != global_num_heads: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}] does not cover all global heads" + ) + retained = tuple( + _integer(value, f"anchor {layer} retained[{head}]") + for head, value in enumerate(retained_raw) + ) + if any(value > eligible_tiles for value in retained): + raise MaskReuseCalibrationError(f"anchor {layer} retained tiles exceed eligible") + dropped = tuple( + _number(value, f"anchor {layer} dropped[{head}]", minimum=0.0, maximum=1.0) + for head, value in enumerate(dropped_raw) + ) + anchors[layer] = AnchorLayerStats(retained, dropped) + + raw_consumers = raw["consumer_layers"] + if not isinstance(raw_consumers, Mapping) or not raw_consumers: + raise MaskReuseCalibrationError("consumer_layers must be non-empty") + consumers: dict[int, CompactConsumerStats] = {} + for raw_layer, raw_stats in raw_consumers.items(): + if ( + not isinstance(raw_layer, str) + or not raw_layer.isdigit() + or raw_layer != str(int(raw_layer)) + ): + raise MaskReuseCalibrationError("consumer layer keys must be canonical integers") + layer = int(raw_layer) + if not isinstance(raw_stats, Mapping): + raise MaskReuseCalibrationError(f"consumer_layers[{layer}] must be an object") + _exact_fields(raw_stats, _CONSUMER_FIELDS, f"consumer_layers[{layer}]") + anchor = _integer(raw_stats["anchor_layer"], f"consumer_layers[{layer}].anchor") + matrix = raw_stats["dropped_mass"] + if not isinstance(matrix, list) or len(matrix) != global_num_heads: + raise MaskReuseCalibrationError( + f"consumer_layers[{layer}] must have {global_num_heads} consumer rows" + ) + for consumer_head, row in enumerate(matrix): + if not isinstance(row, list) or len(row) != global_num_heads: + raise MaskReuseCalibrationError( + f"consumer_layers[{layer}][{consumer_head}] must cover all donors" + ) + for donor_head, value in enumerate(row): + _number( + value, + f"consumer_layers[{layer}][{consumer_head}][{donor_head}]", + minimum=0.0, + maximum=1.0, + ) + consumers[layer] = CompactConsumerStats(anchor, matrix) + return cls( + model=_text(invocation["model"], "model"), + checkpoint_manifest_sha256=_sha256( + invocation["checkpoint_manifest_sha256"], "checkpoint_manifest_sha256" + ), + split=split, + partition=partition, + inner_fold=inner_fold, + prompt_id=_text(invocation["prompt_id"], "prompt_id"), + source=_text(invocation["source"], "source"), + source_group_sha256=_sha256(invocation["source_group_sha256"], "source_group_sha256"), + source_capture_sha256=_sha256( + invocation["source_capture_sha256"], "source_capture_sha256" + ), + min_kv_tokens=minimum, + max_kv_tokens=maximum, + target_sparsity=target, + sample_length=sample_length, + threshold_log2=threshold_log2, + threshold_lambda=threshold_lambda, + geometry=observed_geometry, + global_num_heads=global_num_heads, + eligible_tiles=eligible_tiles, + anchor_stats_by_layer=dict(sorted(anchors.items())), + consumer_layers=dict(sorted(consumers.items())), + ) + + +def _reject_duplicate_json_keys(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise MaskReuseCalibrationError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + +@dataclass(frozen=True, slots=True) +class CompactMaskReuseCaptureSource: + """Re-iterable strict JSONL source used by the three-pass selector.""" + + path: Path + + def __iter__(self) -> Iterator[CompactMaskReuseCapture]: + with self.path.open(encoding="utf-8") as handle: + seen = False + for line_number, line in enumerate(handle, start=1): + if not line.strip(): + continue + seen = True + try: + raw = json.loads(line, object_pairs_hook=_reject_duplicate_json_keys) + except json.JSONDecodeError as error: + raise MaskReuseCalibrationError( + f"line {line_number}: invalid JSON: {error.msg}" + ) from error + except MaskReuseCalibrationError as error: + raise MaskReuseCalibrationError(f"line {line_number}: {error}") from error + if not isinstance(raw, dict): + raise MaskReuseCalibrationError( + f"line {line_number}: compact capture must be an object" + ) + try: + yield CompactMaskReuseCapture.from_mapping(raw) + except MaskReuseCalibrationError as error: + raise MaskReuseCalibrationError(f"line {line_number}: {error}") from error + if not seen: + raise MaskReuseCalibrationError("compact capture input is empty") + + def sha256(self) -> str: + """Hash exact file bytes without loading the capture bundle.""" + digest = sha256() + with self.path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def load_compact_mask_reuse_captures(path: str | Path) -> CompactMaskReuseCaptureSource: + """Return a lazy, re-iterable compact-capture source.""" + source = CompactMaskReuseCaptureSource(Path(path)) + if not source.path.is_file(): + raise MaskReuseCalibrationError(f"compact capture file does not exist: {source.path}") + return source + + +@dataclass(frozen=True, slots=True) +class _Dataset: + model: str + checkpoint_manifest_sha256: str + global_num_heads: int + buckets: tuple[Bucket, ...] + anchors: tuple[int, ...] + nearest: Mapping[int, int] + consumer_layers: tuple[int, ...] + targets: tuple[ConsumerHead, ...] + prompts: Mapping[tuple[Bucket, str], tuple[str, ...]] + menus: Mapping[Bucket, tuple[float, ...]] + deployment_geometry: Mapping[str, object] + prompt_sources: tuple[Mapping[str, object], ...] + calibration_prompt_ids: tuple[str, ...] + heldout_prompt_ids: tuple[str, ...] + capture_count: int + + +def _validate_threshold(capture: CompactMaskReuseCapture, fit: Mapping[str, object]) -> None: + params = fit["prefill"] + assert isinstance(params, Mapping) + lower = params.get("min_observed_sparsity") + upper = params.get("max_observed_sparsity") + if (lower is not None and capture.target_sparsity < float(lower)) or ( + upper is not None and capture.target_sparsity > float(upper) + ): + raise MaskReuseCalibrationError( + f"target_sparsity={capture.target_sparsity} is outside the vanilla fit range" + ) + expected_log2 = ( + math.log2(float(params["a"])) + + float(params["b"]) * capture.target_sparsity * math.log2(math.e) + - math.log2(capture.sample_length) + ) + expected_lambda = 2.0**expected_log2 + if capture.threshold_log2.hex() != expected_log2.hex(): + raise MaskReuseCalibrationError("compact capture threshold_log2 differs from vanilla fit") + if capture.threshold_lambda.hex() != expected_lambda.hex(): + raise MaskReuseCalibrationError("compact capture threshold_lambda differs from vanilla fit") + + +def _capture_order_key(capture: CompactMaskReuseCapture) -> tuple[object, ...]: + return ( + capture.min_kv_tokens, + math.inf if capture.max_kv_tokens is None else capture.max_kv_tokens, + _SPLITS.index(capture.split), + capture.prompt_id, + capture.target_sparsity, + ) + + +def _validate_monotonic_pair( + previous: CompactMaskReuseCapture, current: CompactMaskReuseCapture +) -> None: + identity = (current.bucket, current.split, current.prompt_id) + if identity != (previous.bucket, previous.split, previous.prompt_id): + return + if current.target_sparsity <= previous.target_sparsity: + raise MaskReuseCalibrationError("target sparsities must be strictly increasing per prompt") + for layer, current_stats in current.anchor_stats_by_layer.items(): + previous_stats = previous.anchor_stats_by_layer[layer] + if any( + current_value > previous_value + for current_value, previous_value in zip( + current_stats.retained_tiles, + previous_stats.retained_tiles, + strict=True, + ) + ): + raise MaskReuseCalibrationError( + f"anchor {layer} retained counts increase with target sparsity" + ) + if any( + current_value + _MONOTONIC_ATOL < previous_value + for current_value, previous_value in zip( + current_stats.dropped_mass, + previous_stats.dropped_mass, + strict=True, + ) + ): + raise MaskReuseCalibrationError( + f"anchor {layer} dropped mass decreases with target sparsity" + ) + for layer, current_consumer_stats in current.consumer_layers.items(): + previous_consumer_stats = previous.consumer_layers[layer] + for consumer_head, (current_row, previous_row) in enumerate( + zip( + current_consumer_stats.dropped_mass, + previous_consumer_stats.dropped_mass, + strict=True, + ) + ): + if any( + current_value + _MONOTONIC_ATOL < previous_value + for current_value, previous_value in zip(current_row, previous_row, strict=True) + ): + raise MaskReuseCalibrationError( + f"consumer {layer} head {consumer_head} dropped mass decreases " + "with target sparsity" + ) + + +def _validate_dataset( + source: CompactMaskReuseCaptureSource, + *, + fit: Mapping[str, object], + anchors: tuple[int, ...], + nearest: Mapping[int, int], +) -> _Dataset: + models: set[str] = set() + checkpoint_identities: set[str] = set() + head_counts: set[int] = set() + split_counts: dict[str, int] = defaultdict(int) + split_buckets: dict[str, set[Bucket]] = defaultdict(set) + prompts: dict[tuple[Bucket, str], set[str]] = defaultdict(set) + menus: dict[tuple[Bucket, str, str], set[float]] = defaultdict(set) + capture_sources: dict[tuple[Bucket, str, str], tuple[object, ...]] = {} + fingerprint_owner: dict[tuple[str, str], tuple[Bucket, str]] = {} + seen: set[tuple[Bucket, str, str, float]] = set() + geometry_by_prompt: dict[tuple[Bucket, str, str], tuple[object, ...]] = {} + prompt_source_rows: dict[tuple[Bucket, str, str], Mapping[str, object]] = {} + split_prompt_ids: dict[str, set[str]] = defaultdict(set) + prompt_buckets: dict[tuple[str, str], Bucket] = {} + split_fingerprints: dict[str, set[str]] = defaultdict(set) + group_assignments: dict[str, tuple[str, int | None]] = {} + fingerprint_groups: dict[str, str] = {} + expected_anchor_set = set(anchors) + consumer_to_anchor = {layer: anchor for layer, anchor in nearest.items() if layer != anchor} + expected_consumer_set = set(consumer_to_anchor) + capture_count = 0 + previous_order: tuple[object, ...] | None = None + previous_capture: CompactMaskReuseCapture | None = None + for capture in source: + capture_count += 1 + order = _capture_order_key(capture) + if previous_order is not None and order <= previous_order: + raise MaskReuseCalibrationError("compact capture records are not in canonical order") + if previous_capture is not None: + _validate_monotonic_pair(previous_capture, capture) + previous_order = order + previous_capture = capture + _validate_threshold(capture, fit) + if set(capture.anchor_stats_by_layer) != expected_anchor_set: + raise MaskReuseCalibrationError("compact capture does not cover every topology anchor") + if set(capture.consumer_layers) != expected_consumer_set: + raise MaskReuseCalibrationError("compact capture does not cover every reuse layer") + for layer, stats in capture.consumer_layers.items(): + if stats.anchor_layer != consumer_to_anchor[layer]: + raise MaskReuseCalibrationError( + f"consumer layer {layer} does not match the explicit topology" + ) + models.add(capture.model) + checkpoint_identities.add(capture.checkpoint_manifest_sha256) + head_counts.add(capture.global_num_heads) + split_counts[capture.split] += 1 + split_buckets[capture.split].add(capture.bucket) + prompts[(capture.bucket, capture.split)].add(capture.prompt_id) + menus[(capture.bucket, capture.split, capture.prompt_id)].add(capture.target_sparsity) + split_prompt_ids[capture.split].add(capture.prompt_id) + prompt_identity = (capture.split, capture.prompt_id) + if prompt_buckets.setdefault(prompt_identity, capture.bucket) != capture.bucket: + raise MaskReuseCalibrationError("a prompt ID is assigned to multiple context buckets") + split_fingerprints[capture.split].add(capture.source_capture_sha256) + assignment = (capture.partition, capture.inner_fold) + previous_assignment = group_assignments.setdefault(capture.source_group_sha256, assignment) + if previous_assignment != assignment: + raise MaskReuseCalibrationError( + "one source group is assigned to multiple partitions or inner folds" + ) + previous_group = fingerprint_groups.setdefault( + capture.source_capture_sha256, capture.source_group_sha256 + ) + if previous_group != capture.source_group_sha256: + raise MaskReuseCalibrationError( + "one rendered source capture is assigned to multiple source groups" + ) + key = (capture.bucket, capture.split, capture.prompt_id, capture.target_sparsity) + if key in seen: + raise MaskReuseCalibrationError("compact captures contain a duplicate prompt target") + seen.add(key) + prompt_key = (capture.bucket, capture.split, capture.prompt_id) + source_identity = ( + capture.source, + capture.source_group_sha256, + capture.partition, + capture.inner_fold, + capture.source_capture_sha256, + capture.sample_length, + tuple(sorted(capture.geometry.items())), + capture.eligible_tiles, + ) + if capture_sources.setdefault(prompt_key, source_identity) != source_identity: + raise MaskReuseCalibrationError("prompt metadata differs across target sparsities") + owner_key = (capture.split, capture.source_capture_sha256) + owner = (capture.bucket, capture.prompt_id) + if fingerprint_owner.setdefault(owner_key, owner) != owner: + raise MaskReuseCalibrationError("one source fingerprint names multiple prompt captures") + geometry_identity = ( + capture.source_capture_sha256, + capture.geometry["q_tokens"], + capture.geometry["kv_tokens"], + capture.geometry["q_start_tokens"], + ) + if geometry_by_prompt.setdefault(prompt_key, geometry_identity) != geometry_identity: + raise MaskReuseCalibrationError("prompt geometry differs across target sparsities") + prompt_source_rows[prompt_key] = { + "min_kv_tokens": capture.min_kv_tokens, + "max_kv_tokens": capture.max_kv_tokens, + "split": capture.split, + "prompt_id": capture.prompt_id, + "source": capture.source, + "source_group_sha256": capture.source_group_sha256, + "partition": capture.partition, + "inner_fold": capture.inner_fold, + "source_capture_sha256": capture.source_capture_sha256, + } + if capture_count == 0 or any(not split_counts[split] for split in _SPLITS): + raise MaskReuseCalibrationError("compact captures require calibration and heldout splits") + if len(models) != 1 or len(head_counts) != 1 or len(checkpoint_identities) != 1: + raise MaskReuseCalibrationError( + "compact captures must use one model, checkpoint, and head count" + ) + if split_prompt_ids["calibration"] & split_prompt_ids["heldout"]: + raise MaskReuseCalibrationError("prompt IDs overlap calibration and heldout splits") + if split_fingerprints["calibration"] & split_fingerprints["heldout"]: + raise MaskReuseCalibrationError("source captures overlap calibration and heldout splits") + if split_buckets["calibration"] != split_buckets["heldout"]: + raise MaskReuseCalibrationError("heldout context buckets must match calibration buckets") + buckets = tuple(sorted(split_buckets["calibration"], key=_bucket_key)) + previous_max: int | None = 0 + for minimum, maximum in buckets: + if previous_max is None or minimum <= previous_max: + raise MaskReuseCalibrationError("context buckets must be ordered and non-overlapping") + previous_max = maximum + canonical_menus: dict[Bucket, tuple[float, ...]] = {} + for bucket in buckets: + calibration_prompts = sorted(prompts[(bucket, "calibration")]) + if not calibration_prompts: + raise MaskReuseCalibrationError(f"bucket {bucket} has no calibration prompts") + menu = menus[(bucket, "calibration", calibration_prompts[0])] + if not menu: + raise MaskReuseCalibrationError(f"bucket {bucket} has no target menu") + for split in _SPLITS: + for prompt in prompts[(bucket, split)]: + if menus[(bucket, split, prompt)] != menu: + raise MaskReuseCalibrationError("target-sparsity menu differs across prompts") + canonical_menus[bucket] = tuple(sorted(menu)) + + geometry_rows = [] + for (bucket, split, prompt), identity in sorted( + geometry_by_prompt.items(), key=lambda item: str(item[0]) + ): + geometry_rows.append( + { + "split": split, + "prompt_id": prompt, + "source_capture_sha256": identity[0], + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "q_tokens": identity[1], + "kv_tokens": identity[2], + "q_start_tokens": identity[3], + } + ) + global_num_heads = next(iter(head_counts)) + consumer_layers = tuple(sorted(expected_consumer_set)) + targets = tuple((layer, head) for layer in consumer_layers for head in range(global_num_heads)) + return _Dataset( + model=next(iter(models)), + checkpoint_manifest_sha256=next(iter(checkpoint_identities)), + global_num_heads=global_num_heads, + buckets=buckets, + anchors=anchors, + nearest=nearest, + consumer_layers=consumer_layers, + targets=targets, + prompts={key: tuple(sorted(value)) for key, value in prompts.items()}, + menus=canonical_menus, + deployment_geometry={ + "contract": dict(_DEPLOYMENT_GEOMETRY_CONTRACT), + "observations": geometry_rows, + }, + prompt_sources=tuple( + prompt_source_rows[key] for key in sorted(prompt_source_rows, key=str) + ), + calibration_prompt_ids=tuple(sorted(split_prompt_ids["calibration"])), + heldout_prompt_ids=tuple(sorted(split_prompt_ids["heldout"])), + capture_count=capture_count, + ) + + +@dataclass(slots=True) +class _SelectionAccumulator: + prompt_mass: Mapping[float, dict[str, np.ndarray]] + retained_by_anchor: Mapping[float, np.ndarray] + eligible_sum: dict[float, int] + anchor_evaluation: Mapping[float, _AnchorEvaluation] + + +@dataclass(frozen=True, slots=True) +class _DonorOption: + choice: _Choice + bmm1_skipped_tiles: int + risk_by_prompt: tuple[float, ...] + + +def _selection_pass( + source: CompactMaskReuseCaptureSource, + dataset: _Dataset, + *, + max_anchor_dropped_mass: float, + target_bmm1_skip_ratio: float, +) -> dict[Bucket, _Selection]: + consumer_index = {layer: index for index, layer in enumerate(dataset.consumer_layers)} + anchor_index = {layer: index for index, layer in enumerate(dataset.anchors)} + accumulators: dict[Bucket, _SelectionAccumulator] = {} + for bucket in dataset.buckets: + menus = dataset.menus[bucket] + accumulators[bucket] = _SelectionAccumulator( + prompt_mass={target: {} for target in menus}, + retained_by_anchor={ + target: np.zeros((len(dataset.anchors), dataset.global_num_heads), dtype=np.int64) + for target in menus + }, + eligible_sum=dict.fromkeys(menus, 0), + anchor_evaluation={target: _AnchorEvaluation() for target in menus}, + ) + for capture in source: + if capture.split != "calibration": + continue + accumulator = accumulators[capture.bucket] + target = capture.target_sparsity + accumulator.eligible_sum[target] += capture.eligible_tiles + for layer, stats in capture.anchor_stats_by_layer.items(): + accumulator.retained_by_anchor[target][anchor_index[layer]] += np.asarray( + stats.retained_tiles, dtype=np.int64 + ) + evaluation = accumulator.anchor_evaluation[target] + dropped = [ + value + for stats in capture.anchor_stats_by_layer.values() + for value in stats.dropped_mass + ] + prompt_mean = sum(dropped) / len(dropped) + evaluation.eligible_tiles += ( + capture.eligible_tiles * len(dataset.anchors) * dataset.global_num_heads + ) + evaluation.retained_tiles += sum( + sum(stats.retained_tiles) for stats in capture.anchor_stats_by_layer.values() + ) + evaluation.prompt_count += 1 + evaluation.prompt_mean_sum += prompt_mean + evaluation.worst_prompt_mean = max(evaluation.worst_prompt_mean, prompt_mean) + evaluation.violations += int(prompt_mean > max_anchor_dropped_mass) + for layer, stats in capture.consumer_layers.items(): + if capture.prompt_id not in accumulator.prompt_mass[target]: + accumulator.prompt_mass[target][capture.prompt_id] = np.empty( + ( + len(dataset.consumer_layers), + dataset.global_num_heads, + dataset.global_num_heads, + ), + dtype=np.float64, + ) + accumulator.prompt_mass[target][capture.prompt_id][consumer_index[layer]] = np.asarray( + stats.dropped_mass, dtype=np.float64 + ) + + def donor_options( + accumulator: _SelectionAccumulator, + *, + target: float, + layer: int, + head: int, + retained: np.ndarray, + prompts: tuple[str, ...], + ) -> tuple[_DonorOption, ...]: + options = [ + _DonorOption( + _Choice(0, True, accumulator.eligible_sum[target]), + 0, + tuple(0.0 for _ in prompts), + ) + ] + layer_row = consumer_index[layer] + for donor in range(dataset.global_num_heads): + retained_tiles = int(retained[donor]) + options.append( + _DonorOption( + _Choice(donor, False, retained_tiles), + accumulator.eligible_sum[target] - retained_tiles, + tuple( + float(accumulator.prompt_mass[target][prompt][layer_row, head, donor]) + for prompt in prompts + ), + ) + ) + pareto = [] + for candidate_index, candidate in enumerate(options): + dominated = False + for other_index, other in enumerate(options): + if candidate_index == other_index: + continue + no_worse = other.bmm1_skipped_tiles >= candidate.bmm1_skipped_tiles and all( + other_risk <= candidate_risk + for other_risk, candidate_risk in zip( + other.risk_by_prompt, + candidate.risk_by_prompt, + strict=True, + ) + ) + strictly_better = other.bmm1_skipped_tiles > candidate.bmm1_skipped_tiles or any( + other_risk < candidate_risk + for other_risk, candidate_risk in zip( + other.risk_by_prompt, + candidate.risk_by_prompt, + strict=True, + ) + ) + canonical_tie = not strictly_better and ( + (other.choice.fallback and not candidate.choice.fallback) + or ( + other.choice.fallback == candidate.choice.fallback + and other.choice.donor_head < candidate.choice.donor_head + ) + ) + if no_worse and (strictly_better or canonical_tie): + dominated = True + break + if not dominated: + pareto.append(candidate) + return tuple(pareto) + + def solve_target( + bucket: Bucket, + target: float, + *, + minimum_bmm1_skipped_tiles: int | None, + maximize_bmm1_skipped_tiles: bool, + target_met: bool, + ) -> _Selection | None: + accumulator = accumulators[bucket] + prompts = dataset.prompts[(bucket, "calibration")] + problem = pulp.LpProblem("compact_mask_reuse", pulp.LpMinimize) + variables: dict[tuple[int, int, int], pulp.LpVariable] = {} + options: dict[tuple[int, int, int], _DonorOption] = {} + prompt_risk_terms: dict[str, list[object]] = defaultdict(list) + for layer in dataset.consumer_layers: + anchor = dataset.nearest[layer] + retained = accumulator.retained_by_anchor[target][anchor_index[anchor]] + for head in range(dataset.global_num_heads): + menu = donor_options( + accumulator, + target=target, + layer=layer, + head=head, + retained=retained, + prompts=prompts, + ) + choice_variables = [] + for option_index, option in enumerate(menu): + key = (layer, head, option_index) + variable = pulp.LpVariable( + f"choice_{layer}_{head}_{option_index}", + lowBound=0, + upBound=1, + cat="Binary", + ) + variables[key] = variable + options[key] = option + choice_variables.append(variable) + for prompt, risk in zip(prompts, option.risk_by_prompt, strict=True): + prompt_risk_terms[prompt].append(risk * variable) + problem += pulp.lpSum(choice_variables) == 1, f"choose_{layer}_{head}" + + bmm1_skipped = pulp.lpSum( + options[key].bmm1_skipped_tiles * variable for key, variable in variables.items() + ) + retained_reuse = pulp.lpSum( + options[key].choice.retained_tiles * variable for key, variable in variables.items() + ) + reuse_count = pulp.lpSum( + int(not options[key].choice.fallback) * variable for key, variable in variables.items() + ) + normalizer = len(dataset.nearest) * dataset.global_num_heads + worst_prompt_risk = pulp.LpVariable("worst_prompt_reuse_dropped_mass", lowBound=0.0) + for prompt_index, prompt in enumerate(prompts): + problem += ( + pulp.lpSum(prompt_risk_terms[prompt]) <= normalizer * worst_prompt_risk, + f"reuse_risk_{prompt_index}", + ) + if minimum_bmm1_skipped_tiles is not None: + problem += bmm1_skipped >= minimum_bmm1_skipped_tiles, "minimum_bmm1_skips" + solver = pulp.PULP_CBC_CMD(msg=False, threads=1, options=["randomSeed 0"]) + warm_solver = pulp.PULP_CBC_CMD( + msg=False, + threads=1, + options=["randomSeed 0"], + warmStart=True, + ) + has_incumbent = False + + def minimize(expression: object) -> bool: + nonlocal has_incumbent + problem.setObjective(expression) + status = problem.solve(warm_solver if has_incumbent else solver) + has_incumbent = status == pulp.LpStatusOptimal + return has_incumbent + + def minimize_and_fix(expression: object, name: str, *, integral: bool) -> bool: + nonlocal problem + if not minimize(expression): + return False + raw_value = pulp.value(expression) + value = 0.0 if raw_value is None else float(raw_value) + if integral: + problem += expression == round(value), name + else: + problem += expression <= value + _SOLVER_LEXICOGRAPHIC_ATOL, name + return True + + if maximize_bmm1_skipped_tiles and not minimize_and_fix( + -bmm1_skipped, "fix_maximum_bmm1_skips", integral=True + ): + return None + if not minimize_and_fix(worst_prompt_risk, "fix_worst_prompt_risk", integral=False): + return None + if not minimize_and_fix(retained_reuse, "fix_retained_reuse", integral=True): + raise MaskReuseCalibrationError( + "compact selector lost feasibility after fixing worst-prompt reuse risk" + ) + donor_signature = pulp.lpSum( + (options[key].choice.donor_head + 1) * variable + for key, variable in variables.items() + if not options[key].choice.fallback + ) + donor_base = dataset.global_num_heads * len(dataset.targets) + 1 + if not minimize(reuse_count * donor_base + donor_signature): + raise MaskReuseCalibrationError( + "compact selector lost deterministic tie-break feasibility" + ) + + choices: dict[ConsumerHead, _Choice] = {} + prompt_totals = dict.fromkeys(prompts, 0.0) + worst_individual = 0.0 + skipped_tiles = 0 + retained_tiles = 0 + for key, variable in variables.items(): + if variable.value() <= 0.5: + continue + layer, head, _ = key + option = options[key] + choices[(layer, head)] = option.choice + skipped_tiles += option.bmm1_skipped_tiles + retained_tiles += option.choice.retained_tiles + if not option.choice.fallback: + for prompt, risk in zip(prompts, option.risk_by_prompt, strict=True): + prompt_totals[prompt] += risk + worst_individual = max(worst_individual, risk) + prompt_risks = [prompt_totals[prompt] / normalizer for prompt in prompts] + eligible_tiles = ( + accumulator.eligible_sum[target] * dataset.global_num_heads * len(dataset.nearest) + ) + return _Selection( + target, + choices, + (), + bmm1_eligible_tiles=eligible_tiles, + bmm1_skipped_tiles=skipped_tiles, + target_bmm1_skip_ratio_met=target_met, + worst_prompt_reuse_dropped_mass=max(prompt_risks, default=0.0), + mean_prompt_reuse_dropped_mass=( + sum(prompt_risks) / len(prompt_risks) if prompt_risks else 0.0 + ), + worst_individual_reuse_dropped_mass=worst_individual, + ) + + selections: dict[Bucket, _Selection] = {} + for bucket in dataset.buckets: + accumulator = accumulators[bucket] + frontier: list[dict[str, object]] = [] + candidates: list[tuple[tuple[object, ...], _Selection]] = [] + maximum_candidates: list[tuple[tuple[object, ...], _Selection]] = [] + for target in dataset.menus[bucket]: + anchor_evaluation = accumulator.anchor_evaluation[target] + eligible_tiles = ( + accumulator.eligible_sum[target] * dataset.global_num_heads * len(dataset.nearest) + ) + required_tiles = math.ceil(target_bmm1_skip_ratio * eligible_tiles) + selected = None + if anchor_evaluation.violations == 0: + selected = solve_target( + bucket, + target, + minimum_bmm1_skipped_tiles=required_tiles, + maximize_bmm1_skipped_tiles=False, + target_met=True, + ) + fallback_count = ( + None + if selected is None + else sum(choice.fallback for choice in selected.choices.values()) + ) + retained_reuse = ( + None + if selected is None + else sum(choice.retained_tiles for choice in selected.choices.values()) + ) + combined_tile_cost = ( + None + if retained_reuse is None + else 2 * retained_reuse + anchor_evaluation.retained_tiles + ) + frontier.append( + { + "target_sparsity": target, + "anchor_safe": anchor_evaluation.violations == 0, + "target_bmm1_skip_ratio": target_bmm1_skip_ratio, + "target_bmm1_skip_ratio_feasible": selected is not None, + "retained_reuse_tiles": retained_reuse, + "retained_anchor_tiles": anchor_evaluation.retained_tiles, + "combined_tile_cost": combined_tile_cost, + "fallback_head_count": fallback_count, + "anchor_calibration": anchor_evaluation.to_mapping(), + } + ) + if selected is not None: + rank = ( + selected.worst_prompt_reuse_dropped_mass, + combined_tile_cost, + target, + ) + candidates.append((rank, selected)) + frontier[-1].update( + { + "bmm1_skipped_tiles": selected.bmm1_skipped_tiles, + "achieved_bmm1_skip_ratio": ( + selected.bmm1_skipped_tiles / selected.bmm1_eligible_tiles + ), + "worst_prompt_model_wide_reuse_dropped_mass": ( + selected.worst_prompt_reuse_dropped_mass + ), + } + ) + elif anchor_evaluation.violations == 0: + maximum = solve_target( + bucket, + target, + minimum_bmm1_skipped_tiles=None, + maximize_bmm1_skipped_tiles=True, + target_met=False, + ) + if maximum is not None: + maximum_candidates.append( + ( + ( + -maximum.bmm1_skipped_tiles, + maximum.worst_prompt_reuse_dropped_mass, + target, + ), + maximum, + ) + ) + frontier[-1].update( + { + "maximum_feasible_bmm1_skipped_tiles": (maximum.bmm1_skipped_tiles), + "maximum_feasible_bmm1_skip_ratio": ( + maximum.bmm1_skipped_tiles / maximum.bmm1_eligible_tiles + ), + } + ) + if candidates: + _, selected = min(candidates, key=lambda item: item[0]) + selections[bucket] = _Selection( + selected.target_sparsity, + selected.choices, + tuple(frontier), + bmm1_eligible_tiles=selected.bmm1_eligible_tiles, + bmm1_skipped_tiles=selected.bmm1_skipped_tiles, + target_bmm1_skip_ratio_met=True, + worst_prompt_reuse_dropped_mass=selected.worst_prompt_reuse_dropped_mass, + mean_prompt_reuse_dropped_mass=selected.mean_prompt_reuse_dropped_mass, + worst_individual_reuse_dropped_mass=(selected.worst_individual_reuse_dropped_mass), + ) + elif maximum_candidates: + _, selected = min(maximum_candidates, key=lambda item: item[0]) + selections[bucket] = _Selection( + selected.target_sparsity, + selected.choices, + tuple(frontier), + bmm1_eligible_tiles=selected.bmm1_eligible_tiles, + bmm1_skipped_tiles=selected.bmm1_skipped_tiles, + target_bmm1_skip_ratio_met=False, + worst_prompt_reuse_dropped_mass=selected.worst_prompt_reuse_dropped_mass, + mean_prompt_reuse_dropped_mass=selected.mean_prompt_reuse_dropped_mass, + worst_individual_reuse_dropped_mass=(selected.worst_individual_reuse_dropped_mass), + ) + else: + target = dataset.menus[bucket][0] + dense_choices = { + item: _Choice(0, True, accumulators[bucket].eligible_sum[target]) + for item in dataset.targets + } + selections[bucket] = _Selection( + None, + dense_choices, + tuple(frontier), + "no_target_sparsity_satisfied_anchor_calibration_constraint", + bmm1_eligible_tiles=( + accumulator.eligible_sum[target] + * dataset.global_num_heads + * len(dataset.nearest) + ), + ) + return selections + + +def _evaluation_pass( + source: CompactMaskReuseCaptureSource, + dataset: _Dataset, + selections: Mapping[Bucket, _Selection], + *, + max_anchor_dropped_mass: float, + reuse_dropped_mass_report_threshold: float, +) -> tuple[ + dict[tuple[Bucket, str], _ReuseEvaluation], + dict[tuple[Bucket, str], _AnchorEvaluation], +]: + reuse = {(bucket, split): _ReuseEvaluation() for bucket in dataset.buckets for split in _SPLITS} + anchor = { + (bucket, split): _AnchorEvaluation() for bucket in dataset.buckets for split in _SPLITS + } + for capture in source: + selection = selections[capture.bucket] + evaluation_target = ( + dataset.menus[capture.bucket][0] + if selection.target_sparsity is None + else selection.target_sparsity + ) + if capture.target_sparsity != evaluation_target: + continue + reuse_result = reuse[(capture.bucket, capture.split)] + for layer, stats in capture.consumer_layers.items(): + anchor_stats = capture.anchor_stats_by_layer[stats.anchor_layer] + for head in range(dataset.global_num_heads): + choice = selection.choices[(layer, head)] + reuse_result.eligible_tiles += capture.eligible_tiles + if choice.fallback: + reuse_result.retained_tiles += capture.eligible_tiles + continue + donor = choice.donor_head + dropped = float(stats.dropped_mass[head][donor]) + reuse_result.retained_tiles += anchor_stats.retained_tiles[donor] + reuse_result.sparse_observations += 1 + reuse_result.dropped_mass_sum += dropped + reuse_result.worst_dropped_mass = max(reuse_result.worst_dropped_mass, dropped) + reuse_result.violations += int(dropped > reuse_dropped_mass_report_threshold) + + anchor_result = anchor[(capture.bucket, capture.split)] + dropped_values: list[float] = [] + retained_tiles = 0 + eligible_tiles = 0 + for anchor_stats in capture.anchor_stats_by_layer.values(): + eligible_tiles += capture.eligible_tiles * dataset.global_num_heads + if selection.target_sparsity is None: + retained_tiles += capture.eligible_tiles * dataset.global_num_heads + dropped_values.extend([0.0] * dataset.global_num_heads) + else: + retained_tiles += sum(anchor_stats.retained_tiles) + dropped_values.extend(anchor_stats.dropped_mass) + prompt_mean = sum(dropped_values) / len(dropped_values) + anchor_result.eligible_tiles += eligible_tiles + anchor_result.retained_tiles += retained_tiles + anchor_result.prompt_count += 1 + anchor_result.prompt_mean_sum += prompt_mean + anchor_result.worst_prompt_mean = max(anchor_result.worst_prompt_mean, prompt_mean) + anchor_result.violations += int(prompt_mean > max_anchor_dropped_mass) + for bucket in dataset.buckets: + for split in _SPLITS: + expected = len(dataset.prompts[(bucket, split)]) + if anchor[(bucket, split)].prompt_count != expected: + raise MaskReuseCalibrationError( + f"evaluation pass did not cover every {split} prompt in bucket {bucket}" + ) + return reuse, anchor + + +def calibrate_compact_mask_reuse_policy( + captures: CompactMaskReuseCaptureSource | str | Path, + *, + vanilla_calibration: Mapping[str, object], + topology: Mapping[str, object], + checkpoint_manifest: VerifiedCheckpointManifest, + evidence: Mapping[str, object], + max_anchor_dropped_mass: float, + reuse_dropped_mass_report_threshold: float, + target_bmm1_skip_ratio: float, + source_provenance: Mapping[str, object] | None = None, +) -> dict[str, object]: + """Select a minimum-risk schema-v3 candidate under a per-bucket BMM1 target.""" + if not isinstance(checkpoint_manifest, VerifiedCheckpointManifest): + raise MaskReuseCalibrationError( + "checkpoint_manifest must be returned by verify_checkpoint_manifest" + ) + source = ( + captures + if isinstance(captures, CompactMaskReuseCaptureSource) + else load_compact_mask_reuse_captures(captures) + ) + threshold_scale_factor = canonical_prefill_threshold_scale_factor(vanilla_calibration) + params = threshold_scale_factor["prefill"] + assert isinstance(params, Mapping) + if not {"min_observed_sparsity", "max_observed_sparsity"} <= params.keys(): + raise MaskReuseCalibrationError("compact schema-v3 export requires vanilla fit bounds") + anchors, nearest = _normalize_topology(topology) + dataset = _validate_dataset( + source, + fit=threshold_scale_factor, + anchors=anchors, + nearest=nearest, + ) + checkpoint_identity = checkpoint_manifest.sha256 + if dataset.checkpoint_manifest_sha256 != checkpoint_identity: + raise MaskReuseCalibrationError( + "compact captures do not match the verified checkpoint manifest" + ) + if dataset.model != checkpoint_manifest.model: + raise MaskReuseCalibrationError( + "compact capture model does not match the verified checkpoint manifest" + ) + canonical_evidence = _canonical_evidence(evidence) + input_sha256 = source.sha256() + if canonical_evidence["reuse_bundle_sha256"] != input_sha256: + raise MaskReuseCalibrationError( + "evidence.reuse_bundle_sha256 does not match the compact capture file" + ) + anchor_bound = _number( + max_anchor_dropped_mass, "max_anchor_dropped_mass", minimum=0.0, maximum=1.0 + ) + reuse_report_threshold = _number( + reuse_dropped_mass_report_threshold, + "reuse_dropped_mass_report_threshold", + minimum=0.0, + maximum=1.0, + ) + bmm1_target = _number( + target_bmm1_skip_ratio, + "target_bmm1_skip_ratio", + minimum=0.0, + maximum=1.0, + ) + selections = _selection_pass( + source, + dataset, + max_anchor_dropped_mass=anchor_bound, + target_bmm1_skip_ratio=bmm1_target, + ) + reuse_evaluations, anchor_evaluations = _evaluation_pass( + source, + dataset, + selections, + max_anchor_dropped_mass=anchor_bound, + reuse_dropped_mass_report_threshold=reuse_report_threshold, + ) + if source.sha256() != input_sha256: + raise MaskReuseCalibrationError( + "compact capture file changed during calibration; discard this result" + ) + + context_policies = [] + bucket_reports = [] + target_menus = [] + overall_reuse_calibration = _ReuseEvaluation() + overall_reuse_heldout = _ReuseEvaluation() + overall_anchor_calibration = _AnchorEvaluation() + overall_anchor_heldout = _AnchorEvaluation() + total_fallback = 0 + total_bmm1_eligible = 0 + total_bmm1_skipped = 0 + for bucket in dataset.buckets: + selection = selections[bucket] + if selection.target_sparsity is not None and bucket[1] is None: + raise MaskReuseCalibrationError( + "a deployment-qualified sparse context bucket requires a finite maximum" + ) + reuse_calibration = reuse_evaluations[(bucket, "calibration")] + reuse_heldout = reuse_evaluations[(bucket, "heldout")] + anchor_calibration = anchor_evaluations[(bucket, "calibration")] + anchor_heldout = anchor_evaluations[(bucket, "heldout")] + overall_reuse_calibration.add(reuse_calibration) + overall_reuse_heldout.add(reuse_heldout) + overall_anchor_calibration.add(anchor_calibration) + overall_anchor_heldout.add(anchor_heldout) + policy: dict[str, object] + if selection.target_sparsity is None: + policy = {"min_kv_tokens": bucket[0], "max_kv_tokens": bucket[1], "exact": True} + headmaps: dict[str, list[int]] = {} + fallback_heads: dict[str, list[int]] = {} + else: + headmaps = { + str(layer): [ + selection.choices[(layer, head)].donor_head + for head in range(dataset.global_num_heads) + ] + for layer in dataset.consumer_layers + } + fallback_heads = { + str(layer): [ + head + for head in range(dataset.global_num_heads) + if selection.choices[(layer, head)].fallback + ] + for layer in dataset.consumer_layers + } + policy = { + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "target_sparsity": selection.target_sparsity, + "headmaps": headmaps, + "fallback_heads": fallback_heads, + } + context_policies.append(policy) + fallback_count = sum(len(heads) for heads in fallback_heads.values()) + total_fallback += fallback_count + total_bmm1_eligible += selection.bmm1_eligible_tiles + total_bmm1_skipped += selection.bmm1_skipped_tiles + target_menus.append( + { + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "target_sparsities": [row["target_sparsity"] for row in selection.frontier], + } + ) + bucket_reports.append( + { + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "selected_target_sparsity": selection.target_sparsity, + "selection_status": ( + "target_bmm1_skip_ratio_met" + if selection.target_sparsity is not None + and selection.target_bmm1_skip_ratio_met + else ( + "target_bmm1_skip_ratio_unmet_maximum_feasible" + if selection.target_sparsity is not None + else selection.exact_reason + ) + ), + "target_sparsity_frontier": list(selection.frontier), + "fallback_head_count": fallback_count, + "bmm1_skip_objective": { + "target": bmm1_target, + "target_met": selection.target_bmm1_skip_ratio_met, + "eligible_tiles": selection.bmm1_eligible_tiles, + "skipped_tiles": selection.bmm1_skipped_tiles, + "achieved": ( + selection.bmm1_skipped_tiles / selection.bmm1_eligible_tiles + if selection.bmm1_eligible_tiles + else 0.0 + ), + }, + "reuse_selection_objective": { + "hard_maximum": None, + "worst_development_prompt_model_wide_dropped_mass": ( + selection.worst_prompt_reuse_dropped_mass + ), + "mean_development_prompt_model_wide_dropped_mass": ( + selection.mean_prompt_reuse_dropped_mass + ), + "worst_individual_dropped_mass": ( + selection.worst_individual_reuse_dropped_mass + ), + }, + "reuse_calibration": reuse_calibration.to_mapping(), + "reuse_heldout": reuse_heldout.to_mapping(), + "anchor_calibration": anchor_calibration.to_mapping( + exact=selection.target_sparsity is None + ), + "anchor_heldout": anchor_heldout.to_mapping( + exact=selection.target_sparsity is None + ), + } + ) + + candidate_cell_count = ( + dataset.capture_count + * len(dataset.consumer_layers) + * dataset.global_num_heads + * dataset.global_num_heads + ) + constraints = { + "anchor": { + "metric": "worst_prompt_mean_anchor_dropped_mass", + "comparison": "<=", + "maximum": anchor_bound, + }, + "reuse": { + "selection_metric": ( + "per_prompt_mean_across_all_attention_layers_and_heads_reuse_dropped_mass" + ), + "selection_hard_maximum": None, + "report_metric": "per_prompt_candidate_reuse_dropped_mass", + "report_threshold": reuse_report_threshold, + "report_threshold_affects_selection": False, + }, + "bmm1_skip_ratio": { + "metric": "per_context_bucket_model_wide_eligible_bmm1_tile_skip_ratio", + "comparison": ">=", + "target": bmm1_target, + }, + } + provenance: dict[str, object] = { + "calibrator": "modelopt.mask_reuse.compact_streaming", + "compact_capture_schema_version": 1, + "input_capture_count": dataset.capture_count, + "candidate_cell_count": candidate_cell_count, + "canonical_input_sha256": input_sha256, + "calibration_prompt_ids": list(dataset.calibration_prompt_ids), + "heldout_prompt_ids": list(dataset.heldout_prompt_ids), + "prompt_sources": list(dataset.prompt_sources), + "development_source_group_sha256": sorted( + { + str(row["source_group_sha256"]) + for row in dataset.prompt_sources + if row["partition"] == "development" + } + ), + "outer_test_source_group_sha256": sorted( + { + str(row["source_group_sha256"]) + for row in dataset.prompt_sources + if row["partition"] == "outer_test" + } + ), + "selection_split": "calibration", + "evaluation_split": "heldout", + "streaming_passes": ["validation", "calibration_selection", "frozen_evaluation"], + "threshold_semantics": "a * exp(b * target_sparsity) / sample_length", + "threshold_implementation": { + "threshold_log2": "log2(a) + b * target_sparsity * log2(e) - log2(sample_length)", + "threshold_lambda": "exp2(threshold_log2)", + "validation": "exact IEEE-754 binary64 hex equality", + }, + "checkpoint_manifest_sha256": checkpoint_identity, + "target_sparsity_menus": target_menus, + "constraints": constraints, + "tie_breaks": [ + "reject target sparsities exceeding the calibration anchor bound", + "meet the requested BMM1 skip ratio in every context bucket", + "minimum worst-prompt model-wide reuse dropped mass", + "minimum equal-BMM combined tile cost 2*A_R + A_A", + "fewest reused consumer heads", + "smallest target sparsity", + "lexicographically smallest donor-head map", + ], + "selection_cost": { + "formula": "2 * retained_reuse_tiles + retained_anchor_tiles", + "bmm1_weight": 1.0, + "bmm2_weight": 1.0, + }, + } + if source_provenance is not None: + provenance["source"] = dict(source_provenance) + denominator = len(dataset.targets) * len(dataset.buckets) + report = { + "model": dataset.model, + "checkpoint_manifest_sha256": checkpoint_identity, + "constraints": constraints, + "selection_unit": "context_bucket", + "by_bucket": bucket_reports, + "overall": { + "consumer_head_bucket_count": denominator, + "fallback_head_bucket_count": total_fallback, + "fallback_fraction": total_fallback / denominator, + "bmm1_skip_objective": { + "target_per_context_bucket": bmm1_target, + "all_context_buckets_met": all( + selection.target_bmm1_skip_ratio_met for selection in selections.values() + ), + "eligible_tiles": total_bmm1_eligible, + "skipped_tiles": total_bmm1_skipped, + "achieved_model_wide": ( + total_bmm1_skipped / total_bmm1_eligible if total_bmm1_eligible else 0.0 + ), + }, + "reuse_calibration": overall_reuse_calibration.to_mapping(), + "reuse_heldout": overall_reuse_heldout.to_mapping(), + "anchor_calibration": overall_anchor_calibration.to_mapping(), + "anchor_heldout": overall_anchor_heldout.to_mapping(), + }, + "promotion": { + "status": "candidate_only", + "eligible": False, + "reasons": [ + "grouped inner-fold modal and safety stability not evaluated", + "preregistered outer gate with at least 99 independent groups per family cell not evaluated", + "deployment rectangular-geometry promotion gate not evaluated", + ], + }, + } + return { + "version": 3, + "promotion_status": "candidate_only", + "phase": "prefill", + "decode": {"mode": "dense"}, + "calibration_protocol": _CALIBRATION_PROTOCOL, + "producer": {"name": "modelopt", "version": modelopt.__version__}, + "evidence": canonical_evidence, + "threshold_scale_factor": threshold_scale_factor, + "model": dataset.model, + "checkpoint_manifest_sha256": checkpoint_identity, + "global_num_heads": dataset.global_num_heads, + "target_bmm1_skip_ratio": bmm1_target, + "anchors": list(dataset.anchors), + "nearest": {str(layer): anchor for layer, anchor in dataset.nearest.items()}, + "deployment_geometry_validated": False, + "deployment_geometry": dataset.deployment_geometry, + "context_policies": context_policies, + "provenance": provenance, + "calibration_report": report, + } diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_calibrate_mask_reuse_cli.py b/tests/unit/torch/sparsity/attention_sparsity/test_calibrate_mask_reuse_cli.py new file mode 100644 index 00000000000..2c5244b256c --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_calibrate_mask_reuse_cli.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Focused tests for fail-closed mask-reuse candidate publication.""" + +import importlib.util +import json +import os +from hashlib import sha256 +from pathlib import Path + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + create_checkpoint_manifest, +) + +_SCRIPT_PATH = Path(__file__).parents[5] / "examples/vllm_serve/calibrate_mask_reuse.py" +_SPEC = importlib.util.spec_from_file_location("calibrate_mask_reuse_cli", _SCRIPT_PATH) +assert _SPEC is not None and _SPEC.loader is not None +calibrate_mask_reuse_cli = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(calibrate_mask_reuse_cli) + + +def _canonical(value): + return ( + json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + + +def _base_args(tmp_path: Path): + checkpoint_root = tmp_path / "checkpoint" + checkpoint_root.mkdir() + (checkpoint_root / "config.json").write_text("{}\n", encoding="utf-8") + (checkpoint_root / "model.safetensors").write_bytes(b"weights") + checkpoint = create_checkpoint_manifest(checkpoint_root, model="test-model") + + captures = tmp_path / "captures.jsonl" + captures.write_text("unused\n", encoding="utf-8") + vanilla = tmp_path / "vanilla.json" + vanilla.write_text('{"vanilla":true}\n', encoding="utf-8") + topology = tmp_path / "topology.json" + topology.write_text('{"anchors":[0],"nearest":{"0":0}}\n', encoding="utf-8") + artifacts = {} + for name in ("calibration_plan", "family_registry", "grouped_fit", "outer_report"): + path = tmp_path / f"{name}.json" + path.write_text(f'{{"artifact":"{name}"}}\n', encoding="utf-8") + artifacts[name] = path + capture_manifest = tmp_path / "capture-manifest.json" + capture_manifest.write_bytes( + _canonical( + { + "capture_manifest_schema_version": 4, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v4", + "model": "test-model", + "checkpoint_manifest_sha256": checkpoint.sha256, + "checkpoint_manifest_path": str(checkpoint.manifest_path), + "checkpoint_file_count": checkpoint.file_count, + "checkpoint_total_size_bytes": checkpoint.total_size_bytes, + "plan": "test_stride2", + "fa4_source": "/source", + "fa4_source_commit": "a" * 40, + "fa4_source_git_tree": "b" * 40, + "fa4_source_git_archive_sha256": "c" * 64, + "fa4_source_manifest_path": "/source-manifest.json", + "fa4_source_manifest_sha256": "d" * 64, + "fa4_source_directory_count": 4, + "fa4_source_file_count": 2, + "fa4_source_total_size_bytes": 42, + "engine_kwargs": {"tensor_parallel_size": 2}, + "dense_shadow_validation_requested": True, + "target_sparsity_hex": [(0.7).hex()], + "vanilla_threshold_scale_factor": {"formula": "unused"}, + "vanilla_fit_sha256": sha256(b"normalized-fit").hexdigest(), + "vanilla_config_file_sha256": sha256(vanilla.read_bytes()).hexdigest(), + "prompt_plan_file_sha256": sha256(b"prompts").hexdigest(), + "compact_capture_file_sha256": sha256(captures.read_bytes()).hexdigest(), + "capture_count": 1, + "candidate_cell_count": 4, + "captures": [{"candidate_cell_count": 4}], + } + ) + ) + policy = tmp_path / "candidate.json" + report = tmp_path / "report.json" + args = [ + "--checkpoint", + str(checkpoint_root), + "--compact-captures", + str(captures), + "--capture-manifest", + str(capture_manifest), + "--vanilla-config", + str(vanilla), + "--topology", + str(topology), + "--calibration-plan", + str(artifacts["calibration_plan"]), + "--family-registry", + str(artifacts["family_registry"]), + "--grouped-fit", + str(artifacts["grouped_fit"]), + "--outer-report", + str(artifacts["outer_report"]), + "--max-anchor-dropped-mass", + "0.02", + "--reuse-dropped-mass-report-threshold", + "0.03", + "--target-bmm1-skip-ratio", + "0.10", + "--output-policy", + str(policy), + "--output-report", + str(report), + ] + return args, policy, report, captures, vanilla, artifacts + + +def test_main_verifies_artifacts_and_atomically_writes_candidate(tmp_path, monkeypatch, capsys): + args, policy_path, report_path, captures, vanilla, artifacts = _base_args(tmp_path) + source = object() + monkeypatch.setattr( + calibrate_mask_reuse_cli, "load_compact_mask_reuse_captures", lambda path: source + ) + captured = {} + artifact = { + "version": 3, + "promotion_status": "candidate_only", + "deployment_geometry_validated": False, + "provenance": {"input_capture_count": 1, "candidate_cell_count": 4}, + "calibration_report": {"promotion": {"eligible": False}}, + } + + def fake_calibrate(compact_source, **kwargs): + captured["source"] = compact_source + captured.update(kwargs) + return artifact + + monkeypatch.setattr( + calibrate_mask_reuse_cli, "calibrate_compact_mask_reuse_policy", fake_calibrate + ) + + assert calibrate_mask_reuse_cli.main(args) == 0 + + expected_policy = _canonical(artifact) + assert policy_path.read_bytes() == expected_policy + assert report_path.read_bytes() == _canonical(artifact["calibration_report"]) + assert captured["source"] is source + assert captured["evidence"] == { + "calibration_plan_sha256": sha256(artifacts["calibration_plan"].read_bytes()).hexdigest(), + "family_registry_sha256": sha256(artifacts["family_registry"].read_bytes()).hexdigest(), + "grouped_fit_sha256": sha256(artifacts["grouped_fit"].read_bytes()).hexdigest(), + "outer_report_sha256": sha256(artifacts["outer_report"].read_bytes()).hexdigest(), + "vanilla_fit_sha256": sha256(vanilla.read_bytes()).hexdigest(), + "reuse_bundle_sha256": sha256(captures.read_bytes()).hexdigest(), + } + assert captured["reuse_dropped_mass_report_threshold"] == 0.03 + assert captured["target_bmm1_skip_ratio"] == 0.10 + assert "MASK_REUSE_FA4_CANDIDATE_SHA256=" in capsys.readouterr().out + + +def test_vanilla_mutation_after_semantic_snapshot_cannot_publish(tmp_path, monkeypatch, capsys): + args, policy_path, report_path, _, vanilla, _ = _base_args(tmp_path) + original_payload = vanilla.read_bytes() + real_evidence_artifacts = calibrate_mask_reuse_cli._evidence_artifacts + + def mutate_after_snapshot(namespace, *, vanilla_fit_sha256): + vanilla.write_text('{"vanilla":false}\n', encoding="utf-8") + return real_evidence_artifacts(namespace, vanilla_fit_sha256=vanilla_fit_sha256) + + monkeypatch.setattr(calibrate_mask_reuse_cli, "_evidence_artifacts", mutate_after_snapshot) + monkeypatch.setattr( + calibrate_mask_reuse_cli, + "load_compact_mask_reuse_captures", + lambda path: object(), + ) + + def fake_calibrate(compact_source, **kwargs): + assert kwargs["vanilla_calibration"] == {"vanilla": True} + assert kwargs["evidence"]["vanilla_fit_sha256"] == sha256(original_payload).hexdigest() + return { + "promotion_status": "candidate_only", + "deployment_geometry_validated": False, + "provenance": {"input_capture_count": 1, "candidate_cell_count": 4}, + "calibration_report": {"promotion": {"eligible": False}}, + } + + monkeypatch.setattr( + calibrate_mask_reuse_cli, + "calibrate_compact_mask_reuse_policy", + fake_calibrate, + ) + + with pytest.raises(SystemExit, match="2"): + calibrate_mask_reuse_cli.main(args) + + assert "vanilla_fit_sha256 artifact changed during calibration" in capsys.readouterr().err + assert not policy_path.exists() + assert not report_path.exists() + + +def test_main_rejects_capture_manifest_hash_mismatch(tmp_path, capsys): + args, _, _, captures, _, _ = _base_args(tmp_path) + captures.write_text("changed\n", encoding="utf-8") + + with pytest.raises(SystemExit, match="2"): + calibrate_mask_reuse_cli.main(args) + + assert "compact_capture_file_sha256" in capsys.readouterr().err + + +def test_load_json_object_rejects_duplicate_keys(tmp_path): + path = tmp_path / "duplicate.json" + path.write_text('{"anchors":[0],"anchors":[1]}', encoding="utf-8") + + with pytest.raises(ValueError, match="duplicate JSON key 'anchors'"): + calibrate_mask_reuse_cli._load_json_object(path, label="topology") + + +def test_candidate_publication_rolls_back_report_when_policy_race_wins(tmp_path, monkeypatch): + policy = tmp_path / "candidate.json" + report = tmp_path / "report.json" + real_link = os.link + call_count = 0 + + def racing_link(source, target, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 2: + Path(target).write_bytes(b"racer") + return real_link(source, target, **kwargs) + + monkeypatch.setattr(calibrate_mask_reuse_cli.os, "link", racing_link) + + with pytest.raises(FileExistsError): + calibrate_mask_reuse_cli._publish_candidate_outputs(policy, b"ours", report, b"report") + + assert policy.read_bytes() == b"racer" + assert not report.exists() + + +def test_candidate_publication_refuses_existing_output(tmp_path): + policy = tmp_path / "candidate.json" + report = tmp_path / "report.json" + policy.write_bytes(b"existing") + + with pytest.raises(FileExistsError, match="refusing to overwrite"): + calibrate_mask_reuse_cli._publish_candidate_outputs(policy, b"ours", report, b"report") + + assert policy.read_bytes() == b"existing" + assert not report.exists() diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py b/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py new file mode 100644 index 00000000000..8e90f4feb97 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for content-addressed checkpoint identity.""" + +from hashlib import sha256 +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration import checkpoint_manifest +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + CHECKPOINT_MANIFEST_NAME, + CheckpointManifestError, + create_checkpoint_manifest, + read_stable_file_snapshot, + verify_checkpoint_manifest, +) + + +def _toy_checkpoint(path: Path) -> Path: + path.mkdir() + (path / "config.json").write_text('{"model_type":"toy"}\n', encoding="utf-8") + (path / "model.safetensors").write_bytes(b"toy-weights") + (path / "tokenizer.json").write_text("{}\n", encoding="utf-8") + return path + + +def test_manifest_generation_is_deterministic_verified_and_no_clobber(tmp_path): + first_root = _toy_checkpoint(tmp_path / "first") + second_root = _toy_checkpoint(tmp_path / "second") + + first = create_checkpoint_manifest(first_root, model="toy-model") + second = create_checkpoint_manifest(second_root, model="toy-model") + + assert first.sha256 == second.sha256 + assert first.manifest_path.read_bytes() == second.manifest_path.read_bytes() + assert verify_checkpoint_manifest(first_root, expected_model="toy-model") == first + with pytest.raises(CheckpointManifestError, match="refusing to overwrite"): + create_checkpoint_manifest(first_root, model="toy-model") + + +def test_verifier_rejects_content_mutation_and_symlinks(tmp_path): + root = _toy_checkpoint(tmp_path / "checkpoint") + create_checkpoint_manifest(root, model="toy-model") + (root / "model.safetensors").write_bytes(b"mutated") + with pytest.raises(CheckpointManifestError, match="does not match"): + verify_checkpoint_manifest(root) + + other = _toy_checkpoint(tmp_path / "symlinked") + (other / "alias.bin").symlink_to(other / "model.safetensors") + with pytest.raises(CheckpointManifestError, match="forbidden symlink"): + create_checkpoint_manifest(other, model="toy-model") + + +def test_verifier_rejects_symlink_manifest(tmp_path): + root = _toy_checkpoint(tmp_path / "checkpoint") + target = tmp_path / "outside.json" + target.write_text("{}\n", encoding="utf-8") + (root / CHECKPOINT_MANIFEST_NAME).symlink_to(target) + + with pytest.raises(CheckpointManifestError, match="without following symlinks"): + verify_checkpoint_manifest(root) + + +def test_portable_open_fallback_still_rejects_symlink(tmp_path, monkeypatch): + target = tmp_path / "target.json" + target.write_text("{}\n", encoding="utf-8") + alias = tmp_path / "alias.json" + alias.symlink_to(target) + monkeypatch.delattr(checkpoint_manifest.os, "O_NOFOLLOW", raising=False) + + with pytest.raises(CheckpointManifestError, match="without following symlinks"): + read_stable_file_snapshot(alias, label="fallback input") + + +def test_portable_open_fallback_preserves_exact_binary_bytes(tmp_path, monkeypatch): + path = tmp_path / "binary.dat" + payload = b"line-1\r\nline-2\x00\n" + path.write_bytes(payload) + monkeypatch.delattr(checkpoint_manifest.os, "O_NOFOLLOW", raising=False) + + snapshot = read_stable_file_snapshot(path, label="binary input") + + assert snapshot.payload == payload + assert snapshot.sha256 == sha256(payload).hexdigest() + + +def test_portable_open_uses_binary_flag_when_available(tmp_path, monkeypatch): + path = tmp_path / "binary.dat" + path.write_bytes(b"payload") + binary_flag = 1 << 29 + observed_flags = [] + real_open = checkpoint_manifest.os.open + monkeypatch.setattr(checkpoint_manifest.os, "O_BINARY", binary_flag, raising=False) + + def recording_open(source, flags, *args, **kwargs): + observed_flags.append(flags) + return real_open(source, flags & ~binary_flag, *args, **kwargs) + + monkeypatch.setattr(checkpoint_manifest.os, "open", recording_open) + + snapshot = read_stable_file_snapshot(path, label="binary input") + + assert snapshot.payload == b"payload" + assert observed_flags and observed_flags[0] & binary_flag + + +def test_portable_open_rejects_path_swap_before_read(tmp_path, monkeypatch): + path = tmp_path / "input.dat" + replacement = tmp_path / "replacement.dat" + path.write_bytes(b"expected") + replacement.write_bytes(b"attacker") + real_open = checkpoint_manifest.os.open + swapped = False + monkeypatch.delattr(checkpoint_manifest.os, "O_NOFOLLOW", raising=False) + + def swapping_open(source, flags, *args, **kwargs): + nonlocal swapped + if Path(source) == path and not swapped: + swapped = True + path.unlink() + replacement.rename(path) + return real_open(source, flags, *args, **kwargs) + + monkeypatch.setattr(checkpoint_manifest.os, "open", swapping_open) + + with pytest.raises(CheckpointManifestError, match="stable regular file"): + read_stable_file_snapshot(path, label="swapped input") + + +def test_windows_reparse_attribute_is_link_like(): + observed = SimpleNamespace( + st_mode=0, + st_file_attributes=checkpoint_manifest._FILE_ATTRIBUTE_REPARSE_POINT, + ) + + assert checkpoint_manifest._is_link_like(observed) + + +def test_manifest_creation_without_directory_fsync_support(tmp_path, monkeypatch): + root = _toy_checkpoint(tmp_path / "checkpoint") + monkeypatch.delattr(checkpoint_manifest.os, "O_DIRECTORY", raising=False) + + created = create_checkpoint_manifest(root, model="toy-model") + + assert verify_checkpoint_manifest(root, expected_model="toy-model") == created + + +def test_manifest_publication_preserves_destination_created_by_racer(tmp_path, monkeypatch): + root = _toy_checkpoint(tmp_path / "checkpoint") + manifest = root / CHECKPOINT_MANIFEST_NAME + real_link = checkpoint_manifest.os.link + + def racing_link(source, target, **kwargs): + Path(target).write_text("racer\n", encoding="utf-8") + return real_link(source, target, **kwargs) + + monkeypatch.setattr(checkpoint_manifest.os, "link", racing_link) + + with pytest.raises(CheckpointManifestError, match="appeared during publication"): + create_checkpoint_manifest(root, model="toy-model") + + assert manifest.read_text(encoding="utf-8") == "racer\n" diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py new file mode 100644 index 00000000000..8f2b95623a5 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py @@ -0,0 +1,322 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Deterministic tests for ModelOpt-owned mask-reuse calibration.""" + +import json +import math +from dataclasses import replace +from hashlib import sha256 +from pathlib import Path + +import pytest + +import modelopt +from modelopt.torch.sparsity.attention_sparsity.calibration import ( + AnchorLayerStats, + MaskReuseCalibrationError, + MaskReuseObservation, + calibrate_mask_reuse_policy, + canonical_prefill_threshold_scale_factor, + parse_mask_reuse_observations, +) +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + VerifiedCheckpointManifest, +) + +A = 14.47 +B = 10.91 +VANILLA_FIT = { + "prefill": { + "a": A, + "b": B, + "min_observed_sparsity": 0.4, + "max_observed_sparsity": 0.8, + } +} +TOPOLOGY = {"anchors": [0, 2], "nearest": {"0": 0, "1": 0, "2": 2}} +CHECKPOINT = sha256(b"checkpoint").hexdigest() +VERIFIED_CHECKPOINT = VerifiedCheckpointManifest( + checkpoint_root=Path("/verified-checkpoint"), + manifest_path=Path("/verified-checkpoint/checkpoint_manifest.json"), + model="toy", + sha256=CHECKPOINT, + file_count=2, + total_size_bytes=1, +) +EVIDENCE = { + field: sha256(field.encode()).hexdigest() + for field in ( + "calibration_plan_sha256", + "family_registry_sha256", + "vanilla_fit_sha256", + "reuse_bundle_sha256", + "grouped_fit_sha256", + "outer_report_sha256", + ) +} + + +def _source(prompt: str) -> str: + return sha256(prompt.encode()).hexdigest() + + +def _thresholds(target_sparsity: float, sample_length: int) -> tuple[float, float]: + threshold_log2 = ( + math.log2(A) + B * target_sparsity * math.log2(math.e) - math.log2(sample_length) + ) + return 2.0**threshold_log2, threshold_log2 + + +def _observations() -> list[MaskReuseObservation]: + rows = [] + prompts = { + "calibration": (("cal-0", 65_536), ("cal-1", 98_304)), + "heldout": (("held-0", 65_536), ("held-1", 98_304)), + } + for split, samples in prompts.items(): + for prompt, sample_length in samples: + for target_sparsity in (0.5, 0.7): + threshold, threshold_log2 = _thresholds(target_sparsity, sample_length) + anchor_retained = { + 0.5: (80, 90), + 0.7: (40, 20), + }[target_sparsity] + anchor_dropped = (0.03, 0.03) if target_sparsity == 0.5 else (0.07, 0.08) + anchor_stats = { + 0: AnchorLayerStats(anchor_retained, anchor_dropped), + 2: AnchorLayerStats( + (60, 60) if target_sparsity == 0.5 else (30, 30), + (0.02, 0.02), + ), + } + for consumer_head in range(2): + for donor_head in range(2): + retained = anchor_retained[donor_head] + if target_sparsity == 0.7 and consumer_head == 1: + dropped_mass = 0.07 + 0.01 * donor_head + else: + dropped_mass = 0.02 + 0.01 * donor_head + rows.append( + MaskReuseObservation( + model="toy", + min_kv_tokens=65_536, + max_kv_tokens=131_072, + target_sparsity=target_sparsity, + sample_length=sample_length, + threshold_lambda=threshold, + threshold_log2=threshold_log2, + q_tokens=8192, + kv_tokens=sample_length, + q_start_tokens=sample_length - 8192, + split=split, + prompt_id=prompt, + source_capture_sha256=_source(prompt), + anchor_layer=0, + consumer_layer=1, + consumer_head=consumer_head, + donor_head=donor_head, + retained_tiles=retained, + eligible_tiles=100, + anchor_dropped_mass=( + 0.03 if target_sparsity == 0.5 else 0.07 + 0.01 * donor_head + ), + anchor_stats_by_layer=anchor_stats, + dropped_mass=dropped_mass, + ) + ) + return rows + + +def _calibrate(rows): + return calibrate_mask_reuse_policy( + rows, + vanilla_calibration=VANILLA_FIT, + topology=TOPOLOGY, + checkpoint_manifest=VERIFIED_CHECKPOINT, + evidence=EVIDENCE, + max_anchor_dropped_mass=0.1, + reuse_dropped_mass_report_threshold=0.1, + target_bmm1_skip_ratio=0.1, + ) + + +def test_selects_target_sparsity_and_exports_backend_v3(): + artifact = _calibrate(_observations()) + + assert artifact["version"] == 3 + assert artifact["phase"] == "prefill" + assert artifact["decode"] == {"mode": "dense"} + assert artifact["calibration_protocol"] == "modelopt_mask_reuse_target_sparsity_v1" + assert artifact["producer"] == {"name": "modelopt", "version": modelopt.__version__} + assert artifact["evidence"] == EVIDENCE + assert artifact["threshold_scale_factor"] == { + "formula": "a * exp(b * target_sparsity)", + "prefill": VANILLA_FIT["prefill"], + } + assert artifact["context_policies"] == [ + { + "min_kv_tokens": 65_536, + "max_kv_tokens": 131_072, + "target_sparsity": 0.7, + "headmaps": {"1": [0, 0]}, + "fallback_heads": {"1": [1]}, + } + ] + assert artifact["promotion_status"] == "candidate_only" + assert artifact["deployment_geometry_validated"] is False + assert artifact["deployment_geometry"]["contract"]["kv_page_tokens"] == 16 + assert len(artifact["deployment_geometry"]["observations"]) == 4 + assert ( + artifact["calibration_report"]["overall"]["reuse_heldout"][ + "report_threshold_exceedance_rate" + ] + == 0.0 + ) + json.dumps(artifact) + + +def test_heldout_values_evaluate_but_cannot_change_selection(): + baseline = _calibrate(_observations()) + hostile_rows = [ + replace(row, dropped_mass=0.9) + if row.split == "heldout" + and row.target_sparsity == 0.7 + and row.consumer_head == 0 + and row.donor_head == 0 + else row + for row in _observations() + ] + + hostile = _calibrate(hostile_rows) + + assert hostile["context_policies"] == baseline["context_policies"] + assert ( + hostile["calibration_report"]["overall"]["reuse_heldout"][ + "report_threshold_exceedance_rate" + ] + == 1.0 + ) + assert hostile["calibration_report"]["overall"]["reuse_heldout"]["worst_dropped_mass"] == 0.9 + + +def test_rejects_relabelled_fixed_lambda_observation(): + rows = _observations() + rows[0] = replace( + rows[0], + threshold_lambda=math.nextafter(rows[0].threshold_lambda, math.inf), + ) + + with pytest.raises(MaskReuseCalibrationError, match="threshold_lambda does not match"): + _calibrate(rows) + + +def test_rejects_inexact_log2_launch_argument(): + rows = _observations() + rows[0] = replace( + rows[0], + threshold_log2=math.nextafter(rows[0].threshold_log2, math.inf), + ) + + with pytest.raises(MaskReuseCalibrationError, match="threshold_log2 does not match"): + _calibrate(rows) + + +def test_anchor_gate_includes_anchor_without_consumer_layer(): + rows = [ + replace( + row, + anchor_stats_by_layer={ + **row.anchor_stats_by_layer, + 2: AnchorLayerStats((30, 30), (0.2, 0.2)), + }, + ) + if row.target_sparsity == 0.7 + else row + for row in _observations() + ] + + artifact = _calibrate(rows) + + assert artifact["context_policies"][0]["target_sparsity"] == 0.5 + + +def test_rejects_inconsistent_repeated_anchor_payload(): + rows = _observations() + rows[0] = replace( + rows[0], + anchor_stats_by_layer={ + **rows[0].anchor_stats_by_layer, + 2: AnchorLayerStats((59, 60), (0.02, 0.02)), + }, + ) + + with pytest.raises(MaskReuseCalibrationError, match="anchor_stats_by_layer differs"): + _calibrate(rows) + + +def test_jsonl_rejects_duplicate_keys(): + with pytest.raises(MaskReuseCalibrationError, match="duplicate JSON key 'model'"): + parse_mask_reuse_observations(['{"model":"first","model":"second"}']) + + +def test_prefill_fit_rejects_b_above_backend_limit(): + invalid = {"prefill": {**VANILLA_FIT["prefill"], "b": 20.000_001}} + + with pytest.raises(MaskReuseCalibrationError, match=r"prefill\.b"): + canonical_prefill_threshold_scale_factor(invalid) + + +def test_canonicalizes_existing_sparse_attention_config(): + exported = { + "sparse_attention_config": { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "threshold_scale_factor": { + "formula": "a * exp(b * target_sparsity)", + "prefill": VANILLA_FIT["prefill"], + }, + } + } + } + } + + assert canonical_prefill_threshold_scale_factor(exported) == { + "formula": "a * exp(b * target_sparsity)", + "prefill": VANILLA_FIT["prefill"], + } + + +def test_canonicalizes_legacy_modelopt_serving_calibration(): + exported = { + "config_groups": { + "group_0": { + "sparse_algo": "softmax_skip", + "targets": ["Attention"], + } + }, + "threshold_scale_factor": { + "formula": "a * exp(b * target_sparsity)", + "prefill": {"a": 1.6771257955393728, "b": 8.894668875002724}, + "decode": {"a": 0.006180090552526715, "b": 10.23399476354776}, + }, + "target_sparse_ratio": {"prefill": 0.5, "decode": 0.5}, + } + + assert canonical_prefill_threshold_scale_factor(exported) == { + "formula": "a * exp(b * target_sparsity)", + "prefill": {"a": 1.6771257955393728, "b": 8.894668875002724}, + } diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py new file mode 100644 index 00000000000..67ab4214107 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py @@ -0,0 +1,408 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for streaming compact-capture policy selection.""" + +import inspect +import json +import math +from hashlib import sha256 + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration import mask_reuse_compact +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + create_checkpoint_manifest, +) +from modelopt.torch.sparsity.attention_sparsity.calibration.mask_reuse import ( + MaskReuseCalibrationError, + calibrate_mask_reuse_policy, +) +from modelopt.torch.sparsity.attention_sparsity.calibration.mask_reuse_compact import ( + calibrate_compact_mask_reuse_policy, + load_compact_mask_reuse_captures, +) + +_FIT = { + "threshold_scale_factor": { + "formula": "a * exp(b * target_sparsity)", + "prefill": { + "a": 1.0, + "b": 1.0, + "min_observed_sparsity": 0.4, + "max_observed_sparsity": 0.8, + }, + } +} +_TOPOLOGY = {"anchors": [0], "nearest": {"0": 0, "1": 0}} + + +def _checkpoint(tmp_path): + root = tmp_path / "checkpoint" + root.mkdir() + (root / "config.json").write_text("{}\n", encoding="utf-8") + (root / "model.safetensors").write_bytes(b"weights") + return create_checkpoint_manifest(root, model="test-model") + + +def _capture(split, prompt_id, target, checkpoint_sha256): + threshold_log2 = math.log2(1.0) + target * math.log2(math.e) - math.log2(256) + if target == 0.5: + retained = [2, 3] + anchor_dropped = [0.005, 0.005] + matrix = [[0.001, 0.001], [0.001, 0.001]] + else: + retained = [1, 2] + anchor_dropped = [0.02, 0.02] + matrix = ( + [[0.01, 0.03], [0.04, 0.01]] + if split == "calibration" + else [[0.015, 0.05], [0.05, 0.015]] + ) + invocation = { + "capture_schema_version": 2, + "model": "test-model", + "checkpoint_manifest_sha256": checkpoint_sha256, + "split": split, + "partition": "development" if split == "calibration" else "outer_test", + "inner_fold": 0 if split == "calibration" else None, + "prompt_id": prompt_id, + "source": f"dataset/{split}", + "source_group_sha256": sha256(f"group/{prompt_id}".encode()).hexdigest(), + "source_capture_sha256": sha256(prompt_id.encode()).hexdigest(), + "min_kv_tokens": 129, + "max_kv_tokens": 512, + "target_sparsity_hex": target.hex(), + "sample_length": 256, + "threshold_log2_hex": threshold_log2.hex(), + "threshold_lambda_hex": (2.0**threshold_log2).hex(), + "expected_geometry": {"q_tokens": 256, "kv_tokens": 256, "q_start_tokens": 0}, + } + return { + "compact_capture_schema_version": 1, + "invocation": invocation, + "geometry": invocation["expected_geometry"], + "global_num_heads": 2, + "eligible_tiles": 3, + "anchor_stats_by_layer": { + "0": {"retained_tiles": retained, "dropped_mass": anchor_dropped} + }, + "consumer_layers": {"1": {"anchor_layer": 0, "dropped_mass": matrix}}, + } + + +def _write_captures(path, checkpoint_sha256): + captures = [ + _capture(split, prompt, target, checkpoint_sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + payload = b"".join( + ( + json.dumps(capture, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + for capture in captures + ) + path.write_bytes(payload) + return captures, sha256(payload).hexdigest() + + +def _write_payload(path, captures): + payload = b"".join( + ( + json.dumps(capture, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + for capture in captures + ) + path.write_bytes(payload) + return sha256(payload).hexdigest() + + +def _evidence(reuse_bundle_sha256): + fields = ( + "calibration_plan_sha256", + "family_registry_sha256", + "vanilla_fit_sha256", + "reuse_bundle_sha256", + "grouped_fit_sha256", + "outer_report_sha256", + ) + return { + field: reuse_bundle_sha256 + if field == "reuse_bundle_sha256" + else sha256(field.encode()).hexdigest() + for field in fields + } + + +def _expanded_rows(captures): + rows = [] + for capture in captures: + invocation = capture["invocation"] + anchors = capture["anchor_stats_by_layer"] + anchor = anchors["0"] + matrix = capture["consumer_layers"]["1"]["dropped_mass"] + for consumer_head in range(2): + rows.extend( + { + "model": invocation["model"], + "min_kv_tokens": invocation["min_kv_tokens"], + "max_kv_tokens": invocation["max_kv_tokens"], + "target_sparsity": float.fromhex(invocation["target_sparsity_hex"]), + "sample_length": invocation["sample_length"], + "threshold_lambda": float.fromhex(invocation["threshold_lambda_hex"]), + "threshold_log2": float.fromhex(invocation["threshold_log2_hex"]), + "q_tokens": 256, + "kv_tokens": 256, + "q_start_tokens": 0, + "split": invocation["split"], + "prompt_id": invocation["prompt_id"], + "source_capture_sha256": invocation["source_capture_sha256"], + "anchor_layer": 0, + "consumer_layer": 1, + "consumer_head": consumer_head, + "donor_head": donor_head, + "retained_tiles": anchor["retained_tiles"][donor_head], + "eligible_tiles": 3, + "anchor_dropped_mass": anchor["dropped_mass"][donor_head], + "anchor_stats_by_layer": anchors, + "dropped_mass": matrix[consumer_head][donor_head], + } + for donor_head in range(2) + ) + return rows + + +def test_streaming_compact_selector_matches_legacy_row_policy(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures, digest = _write_captures(path, checkpoint.sha256) + kwargs = { + "vanilla_calibration": _FIT, + "topology": _TOPOLOGY, + "checkpoint_manifest": checkpoint, + "evidence": _evidence(digest), + "max_anchor_dropped_mass": 0.03, + "reuse_dropped_mass_report_threshold": 0.025, + "target_bmm1_skip_ratio": 0.25, + } + + compact = calibrate_compact_mask_reuse_policy(load_compact_mask_reuse_captures(path), **kwargs) + legacy = calibrate_mask_reuse_policy(_expanded_rows(captures), **kwargs) + + assert compact["context_policies"] == legacy["context_policies"] + assert compact["context_policies"][0]["target_sparsity"] == 0.7 + assert compact["context_policies"][0]["headmaps"] == {"1": [0, 1]} + compact_report = dict(compact["calibration_report"]) + legacy_report = dict(legacy["calibration_report"]) + compact_report.pop("promotion") + legacy_report.pop("promotion") + assert compact_report["constraints"] == legacy_report["constraints"] + assert compact_report["overall"] == legacy_report["overall"] + assert compact["provenance"]["input_capture_count"] == 4 + assert compact["provenance"]["candidate_cell_count"] == 16 + assert compact["provenance"]["streaming_passes"] == [ + "validation", + "calibration_selection", + "frozen_evaluation", + ] + assert compact["promotion_status"] == "candidate_only" + assert compact["deployment_geometry_validated"] is False + + +def test_compact_public_interface_has_no_hard_reuse_risk_gate(): + parameters = inspect.signature(calibrate_compact_mask_reuse_policy).parameters + + assert parameters["target_bmm1_skip_ratio"].default is inspect.Parameter.empty + assert "max_reuse_selection_dropped_mass" not in parameters + + +def test_compact_selector_binds_reuse_bundle_sha(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + _, digest = _write_captures(path, checkpoint.sha256) + evidence = _evidence(digest) + evidence["reuse_bundle_sha256"] = sha256(b"wrong").hexdigest() + + with pytest.raises(MaskReuseCalibrationError, match="does not match"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=evidence, + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + ) + + +def test_compact_selector_binds_verified_checkpoint_and_disjoint_groups(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures = [ + _capture(split, prompt, target, checkpoint.sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + captures[0]["invocation"]["checkpoint_manifest_sha256"] = "0" * 64 + digest = _write_payload(path, captures) + with pytest.raises(MaskReuseCalibrationError, match="one model, checkpoint"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + ) + + shared_group = captures[1]["invocation"]["source_group_sha256"] + for capture in captures: + capture["invocation"]["checkpoint_manifest_sha256"] = checkpoint.sha256 + capture["invocation"]["source_group_sha256"] = shared_group + digest = _write_payload(path, captures) + with pytest.raises(MaskReuseCalibrationError, match="multiple partitions"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + ) + + +def test_compact_selector_rejects_file_changed_during_evaluation(tmp_path, monkeypatch): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + _, digest = _write_captures(path, checkpoint.sha256) + real_evaluation_pass = mask_reuse_compact._evaluation_pass + + def mutate_after_evaluation(*args, **kwargs): + result = real_evaluation_pass(*args, **kwargs) + with path.open("ab") as handle: + handle.write(b"\n") + return result + + monkeypatch.setattr(mask_reuse_compact, "_evaluation_pass", mutate_after_evaluation) + + with pytest.raises(MaskReuseCalibrationError, match="changed during calibration"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + ) + + +def test_selector_meets_bmm1_target_before_minimizing_reuse_risk(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures = [ + _capture(split, prompt, target, checkpoint.sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + for capture in captures: + target = float.fromhex(capture["invocation"]["target_sparsity_hex"]) + if target == 0.5: + capture["anchor_stats_by_layer"]["0"] = { + "retained_tiles": [1, 3], + "dropped_mass": [0.001, 0.001], + } + capture["anchor_stats_by_layer"]["2"] = { + "retained_tiles": [3, 3], + "dropped_mass": [0.001, 0.001], + } + capture["consumer_layers"]["1"]["dropped_mass"] = [ + [0.001, 0.001], + [0.001, 0.001], + ] + else: + capture["anchor_stats_by_layer"]["0"] = { + "retained_tiles": [0, 2], + "dropped_mass": [0.002, 0.002], + } + capture["anchor_stats_by_layer"]["2"] = { + "retained_tiles": [0, 0], + "dropped_mass": [0.002, 0.002], + } + capture["consumer_layers"]["1"]["dropped_mass"] = [ + [0.03, 0.01], + [0.03, 0.01], + ] + digest = _write_payload(path, captures) + + candidate = calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology={"anchors": [0, 2], "nearest": {"0": 0, "1": 0, "2": 2}}, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + ) + + assert candidate["context_policies"][0]["target_sparsity"] == 0.7 + frontier = candidate["calibration_report"]["by_bucket"][0]["target_sparsity_frontier"] + assert [row["target_bmm1_skip_ratio_feasible"] for row in frontier] == [False, True] + assert frontier[1]["combined_tile_cost"] == 2 + + +def test_validation_pass_does_not_retain_capture_objects(): + implementation = inspect.getsource(mask_reuse_compact._validate_dataset) + + assert ".append(capture)" not in implementation + assert "list[CompactMaskReuseCapture]" not in implementation + + +@pytest.mark.parametrize("corruption", ["eligible", "retained_monotonic", "mass_monotonic"]) +def test_compact_validation_rejects_impossible_geometry_or_sparsity_trend(tmp_path, corruption): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures = [ + _capture(split, prompt, target, checkpoint.sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + if corruption == "eligible": + captures[0]["eligible_tiles"] = 4 + elif corruption == "retained_monotonic": + captures[1]["anchor_stats_by_layer"]["0"]["retained_tiles"][0] = 3 + else: + captures[1]["consumer_layers"]["1"]["dropped_mass"][0][0] = 0.0 + digest = _write_payload(path, captures) + + with pytest.raises(MaskReuseCalibrationError): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + reuse_dropped_mass_report_threshold=0.025, + target_bmm1_skip_ratio=0.25, + )