From ab76d9988b320298f181f697f45e0f0e4d2dc4d4 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 20:49:46 -0700 Subject: [PATCH 01/16] 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 b47be80907c3e64ead4c7bec5c4ed979259be860 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 20:56:22 -0700 Subject: [PATCH 02/16] 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 00411eaa4ee..0fb67a911d6 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 def _target_sparse_ratio_for_phase(target_sparse_ratio, phase: str) -> float: @@ -239,6 +240,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( @@ -504,6 +603,26 @@ def native_forward(): ) return native_result + if _calibration_active(self): + if getattr(attn_metadata, "use_cascade", False): + # Cascade splits shared prefixes across requests, so per-request + # KV lengths are unavailable; skip measurement for this launch. + return native_forward() + # vLLM >= 0.15 writes the current K/V to the paged cache before + # impl.forward, so the calibrate kernel reads a complete cache. + key_cache, value_cache = kv_cache.unbind(0) + return _forward_calibrate( + self, + query=query, + key_cache=key_cache, + value_cache=value_cache, + block_table=attn_metadata.block_table, + seq_lens=attn_metadata.seq_lens, + cu_seqlens_q=attn_metadata.query_start_loc, + num_actual_tokens=attn_metadata.num_actual_tokens, + output=output, + ) + resolved = _resolve_forward( self, layer, @@ -686,6 +805,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, @@ -814,3 +963,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 5661eee7728d9afb622c91a14eb069d239a12592 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 21:05:21 -0700 Subject: [PATCH 03/16] 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 0fb67a911d6..c8219fae592 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -997,3 +997,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 b44d452be02..57a3a9549ca 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 a732783fec12d6e6b8b11b77ee332fcab0c0aff1 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 21:17:29 -0700 Subject: [PATCH 04/16] 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 858243686d0..29d8965d854 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -117,6 +117,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 bfcb0c838b4a26cb5b4f6302a7dc96d104dbf3d4 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 18 Jul 2026 21:27:13 -0700 Subject: [PATCH 05/16] Add changelog entry for vLLM skip-softmax calibration Signed-off-by: Kai Xu --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index be05c1c441e..5a7554c2763 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -28,6 +28,7 @@ Experimental **New Features** - Add NVFP4 and FP8 PTQ recipes with projection-output quantizers for Llama-Nemotron embedding and reranking models (``modelopt_recipes/huggingface/nemotron_llama/``) and an end-to-end HF embedding/reranking quantize-to-ONNX example (``examples/torch_onnx/hf_embedding_quant_to_onnx.py``). Quantizing the projection-Linear outputs keeps TensorRT inter-layer activations in FP4, roughly halving engine activation memory versus the plain ``nvfp4`` preset. NVFP4/MXFP8 output quantizers now export through the dynamic quantize path. ``examples/torch_onnx/torch_quant_to_onnx.py`` also gains a ``--recipe`` flag to load quantization configs from YAML recipes instead of the removed ``mtq.*_CFG`` module-constant table. +- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. See ``examples/vllm_serve/calibrate_sparse_attn.py``. - Add the ``prepare_megatron_data_blend`` utility to prepare weighted Megatron data blends from YAML configs, including optional token-budgeted subsets for distillation workflows. See the `Megatron data preparation guide `_. - Add Learned Scale Quantization (LSQ) and Dual-LSQ support for quantization-aware distillation, including learnable ``amax`` parameters, tied-scale and pre-scale options, focused NVFP4 recipes, and scale-only training. - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 `_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change). From 6f3f737dfa66f6cb72993ac3e3112a95b4e48f40 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 18 Jul 2026 21:27:17 -0700 Subject: [PATCH 06/16] 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 29d8965d854..fe2b0e80cb8 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -122,11 +122,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 f564882b2bb551b7ed97e0571997a09b50e5d8e1 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 18 Jul 2026 23:11:04 -0700 Subject: [PATCH 07/16] 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 5a7554c2763..adfa528a7fc 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -28,7 +28,7 @@ Experimental **New Features** - Add NVFP4 and FP8 PTQ recipes with projection-output quantizers for Llama-Nemotron embedding and reranking models (``modelopt_recipes/huggingface/nemotron_llama/``) and an end-to-end HF embedding/reranking quantize-to-ONNX example (``examples/torch_onnx/hf_embedding_quant_to_onnx.py``). Quantizing the projection-Linear outputs keeps TensorRT inter-layer activations in FP4, roughly halving engine activation memory versus the plain ``nvfp4`` preset. NVFP4/MXFP8 output quantizers now export through the dynamic quantize path. ``examples/torch_onnx/torch_quant_to_onnx.py`` also gains a ``--recipe`` flag to load quantization configs from YAML recipes instead of the removed ``mtq.*_CFG`` module-constant table. -- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. 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``. - Add the ``prepare_megatron_data_blend`` utility to prepare weighted Megatron data blends from YAML configs, including optional token-budgeted subsets for distillation workflows. See the `Megatron data preparation guide `_. - Add Learned Scale Quantization (LSQ) and Dual-LSQ support for quantization-aware distillation, including learnable ``amax`` parameters, tied-scale and pre-scale options, focused NVFP4 recipes, and scale-only training. - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 `_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change). diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index fe2b0e80cb8..6cb7b8f0a0f 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -133,7 +133,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. @@ -151,7 +151,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 @@ -168,7 +168,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 c8219fae592..027b07eb6c1 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 @@ -247,6 +248,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, *, @@ -264,8 +291,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 @@ -820,6 +848,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() @@ -1012,10 +1048,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 6bc0de7dc7b565e84841f03217b7fa16fd222d6c Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sun, 19 Jul 2026 15:42:38 -0700 Subject: [PATCH 08/16] 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 027b07eb6c1..1d4cfd0f4c6 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -268,9 +268,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 cd405520445b358c56b2c6b34898704597e181c6 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 1 Aug 2026 05:27:51 -0700 Subject: [PATCH 09/16] Add fail-closed mask reuse calibration Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 94 ++ examples/vllm_serve/calibrate_mask_reuse.py | 455 ++++++ examples/vllm_serve/collect_mask_reuse.py | 487 ++++++ .../vllm_serve/create_checkpoint_manifest.py | 43 + .../calibration/__init__.py | 42 + .../calibration/checkpoint_manifest.py | 361 +++++ .../calibration/mask_reuse.py | 1398 +++++++++++++++++ .../calibration/mask_reuse_compact.py | 1211 ++++++++++++++ .../plugins/mask_reuse_capture.py | 809 ++++++++++ .../plugins/sparse_attn_calibration.py | 6 +- .../plugins/vllm_mask_reuse_capture.py | 93 ++ .../test_calibrate_mask_reuse_cli.py | 252 +++ .../test_checkpoint_manifest.py | 72 + .../test_collect_mask_reuse_cli.py | 344 ++++ .../test_mask_reuse_calibration.py | 296 ++++ .../test_mask_reuse_capture.py | 288 ++++ .../test_mask_reuse_compact_calibration.py | 394 +++++ .../test_sparse_attn_calibration.py | 17 +- .../test_vllm_mask_reuse_capture_worker.py | 73 + 19 files changed, 6732 insertions(+), 3 deletions(-) create mode 100644 examples/vllm_serve/calibrate_mask_reuse.py create mode 100644 examples/vllm_serve/collect_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 modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/plugins/vllm_mask_reuse_capture.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_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_collect_mask_reuse_cli.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_capture.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 6cb7b8f0a0f..ba7c22ee968 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -137,6 +137,100 @@ Calibration prompts default to the **RULER dataset** via the same `RulerDatasetB Calibration and serving measure skipping at the same fixed 128x128 tile geometry (active skip-softmax launches bypass the autotuner), so the sparsity realized at serve time matches the calibrated `(a, b)` model. +### Build a cross-layer mask-reuse calibration candidate + +Mask-reuse calibration is a second, offline ModelOpt stage. The current exporter +produces a **candidate only**: it deliberately cannot promote a serving policy +until grouped inner-fold stability, the preregistered 99-outer-group gate, and +the deployment-geometry gate are implemented and pass. + +First run `calibrate_sparse_attn.py` to fit the vanilla skip-softmax transfer +function and finalize the checkpoint/config. Then create the content-addressed +identity inside the exact local checkpoint that vLLM will load: + +```bash +python create_checkpoint_manifest.py --model-id Nemotron-3-Ultra +``` + +This deterministic command writes `/checkpoint_manifest.json` with +SHA256 and size for every checkpoint file except the manifest itself. It refuses +to overwrite an existing manifest, rejects symlinks, and verifies the complete +file set. Keep the checkpoint root immutable after creating it. + +Then collect dense-reference mask-reuse sufficient statistics at a menu of target +sparsities under the exact runtime rule +`threshold = a * exp(b * target_sparsity) / kv_tokens`: + +```bash +python collect_mask_reuse.py \ + --model-id Nemotron-3-Ultra \ + --plan nemotron3_ultra_stride2 \ + --fa4-source /path/to/flash-attention-at-4c40766b \ + --prompts-jsonl mask_reuse_prompts.jsonl \ + --vanilla-config /config.json \ + --target-sparsities 0.5 0.6 0.7 \ + --max-model-len 1000001 \ + --tensor-parallel-size 8 \ + --output mask_reuse_compact_captures.jsonl +``` + +The prompt file has one strict JSON object per capture unit with `split`, +`partition`, `inner_fold`, globally split-unique `prompt_id`, `source`, +`source_group_sha256`, `prompt`, `min_kv_tokens`, and `max_kv_tokens`. +`calibration` maps to `development` with a nonnegative fold; `heldout` maps to +`outer_test` with a null fold. One source group cannot cross partitions or +folds. The collector forces eager, batch-one, BF16, chunked-prefill +vLLM V1 execution through the installed `mask_reuse_fa4` plugin and the +explicit pinned FA4 source. It does not load a provisional reuse policy and +does not enable quantization. The verified checkpoint identity and group +assignment are included in the schema-v2 invocation and its canonical RPC +digest. The checkpoint is reverified after vLLM loads it and before capture. +Each compact JSONL record includes the target +sparsity, sample length, exact binary64 +`threshold_log2` and `threshold_lambda`, rectangular query/KV geometry, every +topology anchor/head's self-mask statistics, and every consumer/donor +dropped-mass matrix once. It does not expand or repeat the matrix as millions +of candidate rows. Fixed-lambda version-4/version-5 captures are not valid +inputs to this target-sparsity path. Use the printed compact-capture SHA256 as +the `reuse_bundle_sha256` evidence value. The FA4 source must be a clean Git +checkout; its exact commit is recorded in the capture manifest. + +Once those observations have been captured, select the per-context target, +donor-head map, and exact fallbacks and export the fail-closed candidate: + +```bash +python calibrate_mask_reuse.py \ + --checkpoint \ + --compact-captures mask_reuse_compact_captures.jsonl \ + --capture-manifest mask_reuse_compact_captures.jsonl.manifest.json \ + --vanilla-config /config.json \ + --topology mask_reuse_topology.json \ + --calibration-plan mask_reuse_calibration_plan.json \ + --family-registry mask_reuse_family_registry.json \ + --grouped-fit mask_reuse_grouped_fit.json \ + --outer-report mask_reuse_outer_report.json \ + --max-anchor-dropped-mass 0.02 \ + --max-reuse-selection-dropped-mass 0.01 \ + --max-reuse-dropped-mass 0.02 \ + --output-policy mask_reuse_candidate.json \ + --output-report mask_reuse_calibration_report.json +``` + +All six evidence hashes are computed from the exact supplied artifact files; +free-form digest claims are not accepted. The compact selector streams three +semantic passes over the captures and hashes +the exact file bytes before and after selection/evaluation. It aborts if the +capture bundle or another evidence artifact changes while calibration is +running. It minimizes the declared equal-BMM tile cost `2 * A_R + A_A`. + +Selection uses only records labeled `calibration`; the frozen policy is then +evaluated on records labeled `heldout`. Held-out violations are preserved in +the standalone report and do not silently retune the selected policy. The +candidate has `promotion_status="candidate_only"` and +`deployment_geometry_validated=false`; the serving backend rejects it. Candidate +and report publication is no-clobber and report-first with rollback, so a policy +path cannot appear without its complete report. Decode is explicitly dense. + The reusable serving policies live in `modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py`. `install_vllm_sparse_attention_from_checkpoint` installs checkpoint-driven sparse-only attention, while `install_vllm_nvfp4_attention` installs fixed NVFP4 Q/K/P/V with optional checkpoint sparsity. Both validate every selected layer before publishing any replacement implementation and return a `VllmAttentionInstallReport` with the installed layer names and backend counts. `sparse_attn_worker.py` only invokes these APIs after vLLM loads the model. It retains `SparseAttnWorker` as the launcher's default and provides `QuantSparseAttnWorker` for the compact NVFP4 policy. Other vLLM integrations can invoke the same library APIs directly: diff --git a/examples/vllm_serve/calibrate_mask_reuse.py b/examples/vllm_serve/calibrate_mask_reuse.py new file mode 100644 index 00000000000..65bdffe84c0 --- /dev/null +++ b/examples/vllm_serve/calibrate_mask_reuse.py @@ -0,0 +1,455 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Build a fail-closed schema-v3 mask-reuse candidate from compact captures. + +The input is compact capture JSONL emitted by ``collect_mask_reuse.py``. Its +streaming selector never expands the full consumer-head by donor-head matrix +into repeated row objects. This command cannot promote a serving policy until +the grouped inner/outer protocol is implemented and its preregistered gates pass. + +Example:: + + python examples/vllm_serve/calibrate_mask_reuse.py \ + --checkpoint /path/to/checkpoint \ + --compact-captures compact-captures.jsonl \ + --capture-manifest compact-captures.jsonl.manifest.json \ + --vanilla-config sparse_attention_config.json \ + --topology topology.json \ + --calibration-plan calibration-plan.json \ + --family-registry family-registry.json \ + --grouped-fit grouped-fit.json \ + --outer-report outer-report.json \ + --max-anchor-dropped-mass 0.02 \ + --max-reuse-dropped-mass 0.02 \ + --max-reuse-selection-dropped-mass 0.005 + +The final output is candidate-only and must be rejected by serving. +""" + +from __future__ import annotations + +import argparse +import json +import os +import stat +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, + 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", + "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: + if path.is_symlink(): + raise ValueError(f"{label} must not be a symlink") + try: + descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW) + except OSError as error: + raise ValueError(f"could not open {label} at {path}") from error + before = os.fstat(descriptor) + digest = sha256() + try: + with os.fdopen(descriptor, "rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + after = os.fstat(handle.fileno()) + named = path.stat(follow_symlinks=False) + except OSError as error: + raise ValueError(f"could not hash stable {label} at {path}") from error + identities = { + (value.st_dev, value.st_ino, value.st_size, value.st_mtime_ns) + for value in (before, after, named) + } + if len(identities) != 1 or not stat.S_ISREG(named.st_mode): + raise ValueError(f"{label} changed while it was being hashed") + return digest.hexdigest() + + +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: + descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) + 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: + try: + observed = path.stat(follow_symlinks=False) + except FileNotFoundError: + return + if (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) + identity = observed.st_dev, observed.st_ino + os.link(temporary, destination, follow_symlinks=False) + try: + 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": 2, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v2", + "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") + 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 compact mask-reuse captures", + allow_abbrev=False, + ) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument( + "--compact-captures", + type=Path, + required=True, + help="Compact capture JSONL emitted by collect_mask_reuse.py (recommended)", + ) + parser.add_argument("--capture-manifest", type=Path, required=True) + parser.add_argument( + "--vanilla-config", + type=Path, + required=True, + help="ModelOpt sparse_attention_config JSON or checkpoint config.json", + ) + parser.add_argument( + "--topology", + type=Path, + required=True, + help="JSON object containing anchors and nearest layer mappings", + ) + parser.add_argument("--calibration-plan", type=Path, required=True) + parser.add_argument("--family-registry", type=Path, required=True) + parser.add_argument("--grouped-fit", type=Path, required=True) + parser.add_argument("--outer-report", type=Path, required=True) + parser.add_argument( + "--max-anchor-dropped-mass", + type=float, + required=True, + help="Maximum allowed anchor dropped mass", + ) + parser.add_argument( + "--max-reuse-dropped-mass", + type=float, + required=True, + help="Maximum allowed held-out reuse dropped mass", + ) + parser.add_argument( + "--max-reuse-selection-dropped-mass", + type=float, + default=None, + help="Optional stricter calibration-time reuse bound", + ) + parser.add_argument( + "--output-policy", + type=Path, + default=Path("mask_reuse_candidate.json"), + help="Fail-closed schema-v3 candidate output path", + ) + parser.add_argument( + "--output-report", + type=Path, + default=Path("mask_reuse_calibration_report.json"), + help="Standalone calibration-report output path", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + if args.output_policy.resolve() == args.output_report.resolve(): + parser.error("--output-policy and --output-report must be different paths") + + try: + checkpoint = verify_checkpoint_manifest(args.checkpoint) + vanilla_config, vanilla_snapshot = _load_json_snapshot( + args.vanilla_config, label="vanilla config" + ) + topology, topology_snapshot = _load_json_snapshot(args.topology, label="topology") + capture_manifest, capture_manifest_snapshot = _load_json_snapshot( + args.capture_manifest, label="capture manifest" + ) + evidence, evidence_paths = _evidence_artifacts( + args, vanilla_fit_sha256=vanilla_snapshot.sha256 + ) + capture_manifest_sha256 = capture_manifest_snapshot.sha256 + topology_sha256 = topology_snapshot.sha256 + _validate_capture_manifest( + capture_manifest, + checkpoint_sha256=checkpoint.sha256, + model=checkpoint.model, + compact_capture_sha256=evidence["reuse_bundle_sha256"], + vanilla_config_sha256=evidence["vanilla_fit_sha256"], + ) + artifact = calibrate_compact_mask_reuse_policy( + load_compact_mask_reuse_captures(args.compact_captures), + vanilla_calibration=vanilla_config, + topology=topology, + checkpoint_manifest=checkpoint, + evidence=evidence, + max_anchor_dropped_mass=args.max_anchor_dropped_mass, + max_reuse_dropped_mass=args.max_reuse_dropped_mass, + max_reuse_selection_dropped_mass=args.max_reuse_selection_dropped_mass, + source_provenance={ + "capture_manifest_sha256": capture_manifest_sha256, + "topology_file_sha256": topology_sha256, + }, + ) + provenance = artifact.get("provenance") + if not isinstance(provenance, Mapping): + raise ValueError("calibrator returned no provenance object") + if provenance.get("input_capture_count") != capture_manifest["capture_count"]: + raise ValueError("calibrator capture count does not match capture manifest") + if provenance.get("candidate_cell_count") != capture_manifest["candidate_cell_count"]: + raise ValueError("calibrator candidate-cell count does not match capture manifest") + for field, path in evidence_paths.items(): + if _stable_file_sha256(path, label=field) != evidence[field]: + raise ValueError(f"{field} artifact changed during calibration") + if ( + _stable_file_sha256(args.capture_manifest, label="capture manifest") + != capture_manifest_sha256 + ): + raise ValueError("capture manifest changed during calibration") + if _stable_file_sha256(args.topology, label="topology") != topology_sha256: + raise ValueError("topology changed during calibration") + if verify_checkpoint_manifest(args.checkpoint) != checkpoint: + raise ValueError("checkpoint changed during calibration") + except (OSError, ValueError) as error: + parser.error(str(error)) + + if ( + artifact.get("promotion_status") != "candidate_only" + or artifact.get("deployment_geometry_validated") is not False + ): + parser.error("calibrator did not return a fail-closed candidate-only artifact") + + report = artifact.get("calibration_report") + if not isinstance(report, Mapping): + parser.error("calibrator returned no calibration_report object") + + try: + policy_payload = _canonical_json_bytes(artifact) + report_payload = _canonical_json_bytes(report) + _publish_candidate_outputs( + args.output_policy, policy_payload, args.output_report, report_payload + ) + except (OSError, FileExistsError) as error: + parser.error(f"could not write calibration outputs: {error}") + + policy_digest = sha256(policy_payload).hexdigest() + print(f"[ModelOpt] Wrote fail-closed mask-reuse candidate to {args.output_policy.resolve()}") + print(f"[ModelOpt] Wrote calibration report to {args.output_report.resolve()}") + print(f"MASK_REUSE_FA4_CANDIDATE_SHA256={policy_digest}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/vllm_serve/collect_mask_reuse.py b/examples/vllm_serve/collect_mask_reuse.py new file mode 100644 index 00000000000..a86390d93a7 --- /dev/null +++ b/examples/vllm_serve/collect_mask_reuse.py @@ -0,0 +1,487 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Collect mask-reuse calibration observations through the vLLM V1 path. + +ModelOpt derives each trial threshold from an existing vanilla skip-softmax +``(a, b)`` fit, arms the custom backend with one exact prompt/target +invocation, and merges only raw sufficient statistics returned by every TP +rank. No promoted reuse policy is loaded during collection. + +Prompt JSONL schema (one object per line):: + + {"split":"calibration", "partition":"development", "inner_fold":0, + "prompt_id":"p0", "source":"ruler/niah", "source_group_sha256":"...", + "prompt":"...", "min_kv_tokens":8192, "max_kv_tokens":65536} + +Usage:: + + python examples/vllm_serve/collect_mask_reuse.py /path/to/checkpoint \ + --model-id Nemotron-3-Ultra \ + --plan nemotron3_ultra_stride2 \ + --fa4-source /path/to/flash-attention-at-4c40766b \ + --prompts-jsonl prompts.jsonl \ + --vanilla-config /path/to/config.json \ + --target-sparsities 0.5 0.6 0.7 \ + --output compact-captures.jsonl +""" + +from __future__ import annotations + +import argparse +import importlib.metadata +import json +import os +import subprocess +import sys +import tempfile +from hashlib import sha256 +from pathlib import Path +from typing import cast + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + read_stable_file_snapshot, + verify_checkpoint_manifest, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.mask_reuse_capture import ( + MAX_QUERY_CHUNK_TOKENS, + CaptureContractError, + build_capture_invocation, + canonical_json_sha256, + merge_rank_captures, + parse_prompt_specs_jsonl, + parse_vanilla_prefill_fit, + validate_begin_acks, + validate_capture_statuses, +) + +CAPTURE_ENV = "MASK_REUSE_FA4_CALIBRATION_CAPTURE" +PLAN_ENV = "MASK_REUSE_FA4_PLAN" +CHECKPOINT_ENV = "MASK_REUSE_FA4_CHECKPOINT_MANIFEST_SHA256" +_POLICY_ENVS = ( + "MASK_REUSE_FA4_POLICY", + "MASK_REUSE_FA4_POLICY_SHA256", +) +_FORBIDDEN_ENGINE_KWARGS = frozenset( + { + "additional_config", + "attention_backend", + "decode_context_parallel_size", + "enable_chunked_prefill", + "enable_prefix_caching", + "enforce_eager", + "kv_cache_dtype", + "kv_transfer_config", + "max_num_batched_tokens", + "max_num_seqs", + "pipeline_parallel_size", + "quantization", + "speculative_config", + "worker_cls", + } +) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Collect ModelOpt mask-reuse sufficient statistics through vLLM V1" + ) + parser.add_argument("model", help="HF checkpoint path loaded by vLLM") + parser.add_argument( + "--model-id", + default=None, + help="Stable model name stored in observations (default: the model argument)", + ) + parser.add_argument("--plan", required=True, help="Explicit mask-reuse layer-plan preset") + parser.add_argument( + "--fa4-source", + required=True, + help="Pinned FlashAttention source checkout containing flash_attn/cute", + ) + parser.add_argument("--prompts-jsonl", required=True, help="Strict prompt-plan JSONL") + parser.add_argument( + "--vanilla-config", + required=True, + help="ModelOpt config containing the calibrated prefill skip-softmax (a, b) fit", + ) + parser.add_argument( + "--target-sparsities", + type=float, + nargs="+", + required=True, + help="Preregistered target-sparsity menu evaluated for every prompt", + ) + parser.add_argument("--output", required=True, help="Compact normalized capture JSONL") + parser.add_argument( + "--output-manifest", + default=None, + help="Capture provenance JSON (default: .manifest.json)", + ) + parser.add_argument("--max-model-len", type=int, default=None) + parser.add_argument("--tensor-parallel-size", type=int, default=1) + parser.add_argument("--gpu-memory-utilization", type=float, default=None) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument( + "--engine-kwargs", + default=None, + help="JSON object of non-contract vLLM kwargs (for example hybrid-model options)", + ) + return parser + + +def _target_menu(values: list[float]) -> tuple[float, ...]: + menu = tuple(sorted(set(values))) + if not menu or any(not 0.0 < value < 1.0 for value in menu): + raise CaptureContractError("--target-sparsities values must be finite and in (0, 1)") + return menu + + +def _engine_kwargs(args: argparse.Namespace) -> dict[str, object]: + extra: dict[str, object] = {} + if args.engine_kwargs is not None: + raw = json.loads(args.engine_kwargs) + if not isinstance(raw, dict): + raise CaptureContractError("--engine-kwargs must be a JSON object") + conflicts = _FORBIDDEN_ENGINE_KWARGS & raw.keys() + if conflicts: + raise CaptureContractError( + f"--engine-kwargs cannot override capture/precision controls: {sorted(conflicts)}" + ) + extra.update(raw) + extra.update( + { + "model": args.model, + "worker_cls": ( + "modelopt.torch.sparsity.attention_sparsity.plugins." + "vllm_mask_reuse_capture.MaskReuseCaptureWorker" + ), + "attention_backend": "CUSTOM", + "dtype": "bfloat16", + "enforce_eager": True, + "enable_prefix_caching": False, + "enable_chunked_prefill": True, + "max_num_batched_tokens": MAX_QUERY_CHUNK_TOKENS, + "max_num_seqs": 1, + "disable_cascade_attn": True, + "pipeline_parallel_size": 1, + "decode_context_parallel_size": 1, + } + ) + if args.max_model_len is not None: + if args.max_model_len <= 0: + raise CaptureContractError("--max-model-len must be positive") + extra["max_model_len"] = args.max_model_len + if args.tensor_parallel_size <= 0: + raise CaptureContractError("--tensor-parallel-size must be positive") + extra["tensor_parallel_size"] = args.tensor_parallel_size + if args.gpu_memory_utilization is not None: + if not 0.0 < args.gpu_memory_utilization <= 1.0: + raise CaptureContractError("--gpu-memory-utilization must be in (0, 1]") + extra["gpu_memory_utilization"] = args.gpu_memory_utilization + if args.trust_remote_code: + extra["trust_remote_code"] = True + return extra + + +def _canonical_capture_line(capture: dict[str, object]) -> bytes: + return ( + json.dumps(capture, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + + +def _configure_capture_environment( + plan: str, fa4_source: str, checkpoint_manifest_sha256: str +) -> tuple[Path, str]: + source = Path(fa4_source).expanduser().resolve() + required = ( + source / "flash_attn/cute/interface.py", + source / "flash_attn/cute/block_sparsity.py", + ) + if not source.is_dir() or any(not path.is_file() for path in required): + raise CaptureContractError( + "--fa4-source must contain flash_attn/cute/interface.py and block_sparsity.py" + ) + try: + commit = subprocess.run( + ["git", "-C", str(source), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + dirty = subprocess.run( + ["git", "-C", str(source), "status", "--porcelain", "--untracked-files=all"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + except (OSError, subprocess.CalledProcessError) as error: + raise CaptureContractError("--fa4-source must be a readable Git checkout") from error + if len(commit) != 40 or any(character not in "0123456789abcdef" for character in commit): + raise CaptureContractError("--fa4-source HEAD is not a canonical Git commit") + if dirty: + raise CaptureContractError( + "--fa4-source has tracked or untracked modifications; pin a clean commit" + ) + plugins = [ + entry + for entry in importlib.metadata.entry_points(group="vllm.general_plugins") + if entry.name == "mask_reuse_fa4" + ] + if len(plugins) != 1 or plugins[0].value != "mask_reuse_vllm.plugin:register": + raise CaptureContractError( + "the mask_reuse_fa4 vLLM plugin entry point is missing or ambiguous" + ) + os.environ[CAPTURE_ENV] = "1" + os.environ[PLAN_ENV] = plan + os.environ[CHECKPOINT_ENV] = checkpoint_manifest_sha256 + os.environ["PYTHONDONTWRITEBYTECODE"] = "1" + sys.dont_write_bytecode = True + os.environ["MASK_REUSE_FA4_SOURCE"] = str(source) + os.environ["MASK_REUSE_FA4_FORCE_DENSE"] = "0" + os.environ["VLLM_PLUGINS"] = "mask_reuse_fa4" + # Collection is deliberately policy-free. Remove inherited serving + # settings so a stale deployment policy cannot become threshold authority. + for name in _POLICY_ENVS: + os.environ.pop(name, None) + return source, commit + + +def _fsync_directory(path: Path) -> None: + descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def _publish_no_clobber(temporary: Path, destination: Path) -> tuple[int, int]: + """Atomically link a complete temp file only when destination is absent.""" + + identity = (temporary.stat().st_dev, temporary.stat().st_ino) + os.link(temporary, destination) + try: + temporary.unlink() + _fsync_directory(destination.parent) + except BaseException: + _unlink_if_identity(destination, identity) + raise + return identity + + +def _unlink_if_identity(path: Path, identity: tuple[int, int]) -> None: + """Rollback only a file that is still the inode published by this process.""" + + try: + stat = path.stat() + except FileNotFoundError: + return + if (stat.st_dev, stat.st_ino) == identity: + path.unlink() + _fsync_directory(path.parent) + + +def run(args: argparse.Namespace) -> tuple[Path, Path]: + prompt_snapshot = read_stable_file_snapshot(args.prompts_jsonl, label="prompt plan") + prompt_plan_sha256 = prompt_snapshot.sha256 + prompts = parse_prompt_specs_jsonl(prompt_snapshot.payload) + vanilla_snapshot = read_stable_file_snapshot(args.vanilla_config, label="vanilla config") + vanilla_config_sha256 = vanilla_snapshot.sha256 + threshold_scale_factor = parse_vanilla_prefill_fit(vanilla_snapshot.payload) + targets = _target_menu(args.target_sparsities) + checkpoint = verify_checkpoint_manifest(args.model, expected_model=args.model_id) + model_id = checkpoint.model + output_path = Path(args.output) + manifest_path = ( + Path(args.output_manifest) + if args.output_manifest is not None + else Path(str(output_path) + ".manifest.json") + ) + if output_path.resolve() == manifest_path.resolve(): + raise CaptureContractError("observation and manifest output paths must differ") + if output_path.exists() or manifest_path.exists(): + raise CaptureContractError("capture outputs already exist; refusing to overwrite them") + output_path.parent.mkdir(parents=True, exist_ok=True) + manifest_path.parent.mkdir(parents=True, exist_ok=True) + + fa4_source, fa4_source_commit = _configure_capture_environment( + args.plan, args.fa4_source, checkpoint.sha256 + ) + # Import after setting the gate: worker subprocesses inherit the exact + # capture environment and never enter policy-backed serving mode. + from vllm import LLM, SamplingParams + + llm = LLM(**_engine_kwargs(args)) + loaded_checkpoint = verify_checkpoint_manifest(args.model, expected_model=model_id) + if loaded_checkpoint != checkpoint: + raise CaptureContractError( + "checkpoint identity changed while vLLM loaded the model; no capture was started" + ) + statuses = llm.collective_rpc("mask_reuse_capture_status") + validate_capture_statuses(statuses) + tokenizer = llm.get_tokenizer() + sampling = SamplingParams(temperature=0.0, max_tokens=1, ignore_eos=True) + + capture_manifests: list[dict[str, object]] = [] + seen_sources: dict[str, tuple[str, str, tuple[int, int | None]]] = {} + capture_digest = sha256() + capture_count = 0 + candidate_cell_count = 0 + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=output_path.parent, + prefix=f".{output_path.name}.", + suffix=".tmp", + delete=False, + ) as temporary: + temporary_path = Path(temporary.name) + for prompt in prompts: + token_ids = tokenizer.encode(prompt.prompt, add_special_tokens=True) + if not isinstance(token_ids, list) or any( + type(token) is not int for token in token_ids + ): + raise CaptureContractError( + "tokenizer.encode must return a list of integer token IDs" + ) + if args.max_model_len is not None and len(token_ids) + 1 > args.max_model_len: + raise CaptureContractError( + f"prompt {prompt.prompt_id!r} plus one output token exceeds --max-model-len" + ) + for target in targets: + invocation = build_capture_invocation( + model=model_id, + checkpoint_manifest_sha256=checkpoint.sha256, + prompt=prompt, + prompt_token_ids=token_ids, + target_sparsity=target, + threshold_scale_factor=threshold_scale_factor, + ) + fingerprint = str(invocation["source_capture_sha256"]) + identity = (prompt.split, prompt.prompt_id, prompt.bucket) + previous = seen_sources.setdefault(fingerprint, identity) + if previous != identity: + raise CaptureContractError( + "the same tokenized source is assigned to multiple prompt captures: " + f"{previous} and {identity}" + ) + acknowledgements = llm.collective_rpc( + "mask_reuse_capture_begin", args=(invocation,) + ) + validate_begin_acks(acknowledgements, invocation) + # Passing token IDs makes sample_length and source_capture_sha256 + # identical to the request that reaches the vLLM scheduler. + llm.generate(token_ids, sampling, use_tqdm=False) + merged = merge_rank_captures( + llm.collective_rpc("mask_reuse_capture_drain"), invocation + ) + line = _canonical_capture_line(merged.capture) + temporary.write(line) + capture_digest.update(line) + capture_count += 1 + candidate_cell_count += cast("int", merged.manifest["candidate_cell_count"]) + capture_manifests.append(merged.manifest) + temporary.flush() + os.fsync(temporary.fileno()) + except BaseException: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + raise + + prompt_final = read_stable_file_snapshot(args.prompts_jsonl, label="prompt plan") + vanilla_final = read_stable_file_snapshot(args.vanilla_config, label="vanilla config") + if prompt_final.sha256 != prompt_plan_sha256 or vanilla_final.sha256 != vanilla_config_sha256: + assert temporary_path is not None + temporary_path.unlink(missing_ok=True) + raise CaptureContractError( + "prompt plan or vanilla calibration changed during capture; evidence was discarded" + ) + + capture_sha256 = capture_digest.hexdigest() + manifest = { + "capture_manifest_schema_version": 2, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v2", + "model": model_id, + "checkpoint_manifest_sha256": checkpoint.sha256, + "checkpoint_manifest_path": str(checkpoint.manifest_path), + "checkpoint_file_count": checkpoint.file_count, + "checkpoint_total_size_bytes": checkpoint.total_size_bytes, + "plan": args.plan, + "fa4_source": str(fa4_source), + "fa4_source_commit": fa4_source_commit, + "target_sparsity_hex": [target.hex() for target in targets], + "vanilla_threshold_scale_factor": threshold_scale_factor, + "vanilla_fit_sha256": canonical_json_sha256(threshold_scale_factor), + "vanilla_config_file_sha256": vanilla_config_sha256, + "prompt_plan_file_sha256": prompt_plan_sha256, + "compact_capture_file_sha256": capture_sha256, + "capture_count": capture_count, + "candidate_cell_count": candidate_cell_count, + "captures": capture_manifests, + } + manifest_bytes = ( + json.dumps(manifest, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + manifest_temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=manifest_path.parent, + prefix=f".{manifest_path.name}.", + suffix=".tmp", + delete=False, + ) as temporary: + manifest_temporary = Path(temporary.name) + temporary.write(manifest_bytes) + temporary.flush() + os.fsync(temporary.fileno()) + assert temporary_path is not None and manifest_temporary is not None + capture_identity = _publish_no_clobber(temporary_path, output_path) + temporary_path = None + try: + _publish_no_clobber(manifest_temporary, manifest_path) + manifest_temporary = None + except BaseException: + _unlink_if_identity(output_path, capture_identity) + raise + except FileExistsError as error: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + if manifest_temporary is not None: + manifest_temporary.unlink(missing_ok=True) + raise CaptureContractError("capture destination appeared during publication") from error + except BaseException: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + if manifest_temporary is not None: + manifest_temporary.unlink(missing_ok=True) + raise + print( + f"[ModelOpt] Wrote {capture_count} compact captures " + f"({candidate_cell_count} candidate cells) to {output_path.resolve()}" + ) + print(f"[ModelOpt] compact_capture_file_sha256={capture_sha256}") + print(f"[ModelOpt] Wrote capture manifest to {manifest_path.resolve()}") + return output_path, manifest_path + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + run(args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/vllm_serve/create_checkpoint_manifest.py b/examples/vllm_serve/create_checkpoint_manifest.py new file mode 100644 index 00000000000..8fe276f0d0f --- /dev/null +++ b/examples/vllm_serve/create_checkpoint_manifest.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Create a deterministic, no-clobber checkpoint manifest for calibration.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + create_checkpoint_manifest, +) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(allow_abbrev=False) + parser.add_argument("checkpoint", type=Path) + parser.add_argument("--model-id", required=True) + args = parser.parse_args(argv) + try: + manifest = create_checkpoint_manifest(args.checkpoint, model=args.model_id) + except (OSError, ValueError) as error: + parser.error(str(error)) + print(f"[ModelOpt] Wrote {manifest.manifest_path}") + print(f"CHECKPOINT_MANIFEST_SHA256={manifest.sha256}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py b/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py index 87088f805bd..950a411182f 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py @@ -17,10 +17,52 @@ 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, + 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", + "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..6ce3cadf274 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/checkpoint_manifest.py @@ -0,0 +1,361 @@ +# 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", + "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"}) + + +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 _open_stable_regular(path: Path, label: str) -> tuple[int, os.stat_result]: + try: + descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW) + except OSError as error: + raise CheckpointManifestError( + f"could not open {label} without following symlinks" + ) from error + opened = os.fstat(descriptor) + try: + named = path.stat(follow_symlinks=False) + except OSError: + os.close(descriptor) + raise + if ( + not stat.S_ISREG(opened.st_mode) + or stat.S_ISLNK(named.st_mode) + or (opened.st_dev, opened.st_ino) != (named.st_dev, named.st_ino) + ): + 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: + with os.fdopen(descriptor, "rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + observed_size += len(chunk) + digest.update(chunk) + if payload is not None: + payload.extend(chunk) + after = os.fstat(handle.fileno()) + named_after = path.stat(follow_symlinks=False) + except OSError as error: + raise CheckpointManifestError(f"could not hash stable {label}") from error + identity_before = (before.st_dev, before.st_ino, before.st_size, before.st_mtime_ns) + identity_after = (after.st_dev, after.st_ino, after.st_size, after.st_mtime_ns) + identity_named = ( + named_after.st_dev, + named_after.st_ino, + named_after.st_size, + named_after.st_mtime_ns, + ) + if identity_before != identity_after or identity_after != identity_named: + 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 _checkpoint_files(root: Path, manifest_path: Path) -> set[str]: + files: set[str] = set() + for path in root.rglob("*"): + if path.is_symlink(): + raise CheckpointManifestError( + f"checkpoint contains forbidden symlink {path.relative_to(root).as_posix()!r}" + ) + if path.is_file() and path != manifest_path: + files.add(path.relative_to(root).as_posix()) + return files + + +def _fsync_directory(path: Path) -> None: + descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) + 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) + identity = observed.st_dev, observed.st_ino + os.link(temporary, manifest_path, follow_symlinks=False) + try: + temporary.unlink() + temporary = None + _fsync_directory(root) + except BaseException: + try: + published = manifest_path.stat(follow_symlinks=False) + if (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..d99f785055b --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py @@ -0,0 +1,1398 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline calibration and schema-v3 export for cross-layer mask reuse. + +The selector consumes prompt-level observations measured at target sparsities +derived from an existing ModelOpt skip-softmax fit. Calibration observations +alone select one target sparsity per context bucket and one donor head (or an +exact fallback) per consumer head. Held-out observations only evaluate the +frozen policy. + +This module intentionally has no serving-backend dependency. It exports the +JSON-safe schema consumed by the mask-reuse attention backend. +""" + +from __future__ import annotations + +import json +import math +from collections import defaultdict +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import cast + +import modelopt + +from .checkpoint_manifest import VerifiedCheckpointManifest + +__all__ = [ + "AnchorLayerStats", + "MaskReuseCalibrationError", + "MaskReuseObservation", + "calibrate_mask_reuse_policy", + "canonical_prefill_threshold_scale_factor", + "load_mask_reuse_observations", + "parse_mask_reuse_observations", +] + + +_FORMULA = "a * exp(b * target_sparsity)" +_SPLITS = frozenset({"calibration", "heldout"}) +_OBSERVATION_FIELDS = frozenset( + { + "model", + "min_kv_tokens", + "max_kv_tokens", + "target_sparsity", + "sample_length", + "threshold_lambda", + "threshold_log2", + "q_tokens", + "kv_tokens", + "q_start_tokens", + "split", + "prompt_id", + "source_capture_sha256", + "anchor_layer", + "consumer_layer", + "consumer_head", + "donor_head", + "retained_tiles", + "eligible_tiles", + "anchor_dropped_mass", + "anchor_stats_by_layer", + "dropped_mass", + } +) +_CALIBRATION_PROTOCOL = "modelopt_mask_reuse_target_sparsity_v1" +_EVIDENCE_FIELDS = frozenset( + { + "calibration_plan_sha256", + "family_registry_sha256", + "vanilla_fit_sha256", + "reuse_bundle_sha256", + "grouped_fit_sha256", + "outer_report_sha256", + } +) +_DEPLOYMENT_GEOMETRY_CONTRACT: dict[str, object] = { + "schema_version": 1, + "batch_size": 1, + "max_query_chunk_tokens": 8192, + "query_block_tokens": 128, + "key_block_tokens": 128, + "qstage2_query_pair_tokens": 256, + "kv_page_tokens": 16, + "head_dim": 128, + "causal": True, + "bottom_right_aligned": True, + "query_chunk_start_alignment_tokens": 128, + "attention_dtype": "bfloat16", + "kv_cache_dtype": "bfloat16", + "common_prefix": False, + "cascade_attention": False, + "context_parallel_size": 1, + "pipeline_parallel_size": 1, +} + +Bucket = tuple[int, int | None] +ConsumerHead = tuple[int, int] +ObservationKey = tuple[str, str, float, int, int, int] +AnchorKey = tuple[str, str, float, int, int] + + +class MaskReuseCalibrationError(ValueError): + """Raised when observations cannot produce a trustworthy reuse policy.""" + + +@dataclass(frozen=True, slots=True) +class AnchorLayerStats: + """Per-head BLASST mask statistics for one topology anchor layer.""" + + retained_tiles: tuple[int, ...] + dropped_mass: tuple[float, ...] + + +def _integer(value: object, name: str, *, minimum: int) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise MaskReuseCalibrationError(f"{name} must be an integer >= {minimum}") + return value + + +def _number( + value: object, + name: str, + *, + minimum: float, + maximum: float | None = None, + minimum_inclusive: bool = True, +) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise MaskReuseCalibrationError(f"{name} must be a finite number") + result = float(value) + below = result < minimum if minimum_inclusive else result <= minimum + if not math.isfinite(result) or below or (maximum is not None and result > maximum): + raise MaskReuseCalibrationError(f"{name} is outside its valid range") + return result + + +def _sha256(value: object, name: str) -> str: + if not isinstance(value, str): + raise MaskReuseCalibrationError(f"{name} must be a lowercase SHA256") + normalized = value.strip().lower() + if len(normalized) != 64 or any( + character not in "0123456789abcdef" for character in normalized + ): + raise MaskReuseCalibrationError(f"{name} must be a lowercase SHA256") + return normalized + + +def _parse_anchor_stats_by_layer( + value: object, + *, + require_canonical_string_keys: bool, +) -> dict[int, AnchorLayerStats]: + if not isinstance(value, Mapping) or not value: + raise MaskReuseCalibrationError("anchor_stats_by_layer must be a non-empty object") + parsed: dict[int, AnchorLayerStats] = {} + for raw_layer, raw_stats in value.items(): + if isinstance(raw_layer, bool): + raise MaskReuseCalibrationError("anchor_stats_by_layer has a non-integer layer key") + try: + layer = int(raw_layer) + except (TypeError, ValueError) as error: + raise MaskReuseCalibrationError( + "anchor_stats_by_layer has a non-integer layer key" + ) from error + if layer < 0 or ( + require_canonical_string_keys + and (not isinstance(raw_layer, str) or raw_layer != str(layer)) + ): + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer layer key {raw_layer!r} is not canonical" + ) + if layer in parsed: + raise MaskReuseCalibrationError(f"anchor_stats_by_layer repeats layer {layer}") + if isinstance(raw_stats, AnchorLayerStats): + raw_retained = raw_stats.retained_tiles + raw_dropped = raw_stats.dropped_mass + elif isinstance(raw_stats, Mapping): + missing = {"retained_tiles", "dropped_mass"} - raw_stats.keys() + extra = raw_stats.keys() - {"retained_tiles", "dropped_mass"} + if missing or extra: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}] requires exactly retained_tiles " + f"and dropped_mass; missing={sorted(missing)}, extra={sorted(extra)}" + ) + raw_retained = raw_stats["retained_tiles"] + raw_dropped = raw_stats["dropped_mass"] + else: + raise MaskReuseCalibrationError(f"anchor_stats_by_layer[{layer}] must be an object") + if not isinstance(raw_retained, list | tuple) or not raw_retained: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}].retained_tiles must be a non-empty list" + ) + if not isinstance(raw_dropped, list | tuple) or not raw_dropped: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}].dropped_mass must be a non-empty list" + ) + retained = tuple( + _integer( + item, + f"anchor_stats_by_layer[{layer}].retained_tiles[{head}]", + minimum=0, + ) + for head, item in enumerate(raw_retained) + ) + dropped = tuple( + _number( + item, + f"anchor_stats_by_layer[{layer}].dropped_mass[{head}]", + minimum=0.0, + maximum=1.0, + ) + for head, item in enumerate(raw_dropped) + ) + if len(retained) != len(dropped): + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{layer}] head arrays differ in width" + ) + parsed[layer] = AnchorLayerStats(retained, dropped) + return dict(sorted(parsed.items())) + + +@dataclass(frozen=True, slots=True) +class MaskReuseObservation: + """One prompt, target-sparsity, consumer-head, and donor-head observation.""" + + model: str + min_kv_tokens: int + max_kv_tokens: int | None + target_sparsity: float + sample_length: int + threshold_lambda: float + threshold_log2: float + q_tokens: int + kv_tokens: int + q_start_tokens: int + split: str + prompt_id: str + source_capture_sha256: str + anchor_layer: int + consumer_layer: int + consumer_head: int + donor_head: int + retained_tiles: int + eligible_tiles: int + anchor_dropped_mass: float + anchor_stats_by_layer: Mapping[int, AnchorLayerStats] + dropped_mass: float + + def __post_init__(self) -> None: + if not isinstance(self.model, str) or not self.model.strip(): + raise MaskReuseCalibrationError("model must be a non-empty string") + if not isinstance(self.prompt_id, str) or not self.prompt_id.strip(): + raise MaskReuseCalibrationError("prompt_id must be a non-empty string") + if self.split not in _SPLITS: + raise MaskReuseCalibrationError(f"split must be one of {sorted(_SPLITS)}") + minimum = _integer(self.min_kv_tokens, "min_kv_tokens", minimum=1) + maximum = self.max_kv_tokens + if maximum is not None: + maximum = _integer(maximum, "max_kv_tokens", minimum=minimum) + sample_length = _integer(self.sample_length, "sample_length", minimum=1) + if sample_length < minimum or (maximum is not None and sample_length > maximum): + raise MaskReuseCalibrationError("sample_length lies outside its context bucket") + q_tokens = _integer(self.q_tokens, "q_tokens", minimum=129) + if q_tokens > int(cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["max_query_chunk_tokens"])): + raise MaskReuseCalibrationError("q_tokens exceeds the deployment geometry limit") + kv_tokens = _integer(self.kv_tokens, "kv_tokens", minimum=1) + q_start_tokens = _integer(self.q_start_tokens, "q_start_tokens", minimum=0) + if sample_length != kv_tokens: + raise MaskReuseCalibrationError("sample_length must equal kv_tokens") + if q_start_tokens + q_tokens != kv_tokens: + raise MaskReuseCalibrationError("q_start_tokens + q_tokens must equal kv_tokens") + alignment = int( + cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["query_chunk_start_alignment_tokens"]) + ) + if q_start_tokens % alignment: + raise MaskReuseCalibrationError("q_start_tokens must be 128-token aligned") + anchor_layer = _integer(self.anchor_layer, "anchor_layer", minimum=0) + consumer_layer = _integer(self.consumer_layer, "consumer_layer", minimum=0) + if anchor_layer >= consumer_layer: + raise MaskReuseCalibrationError("anchor_layer must precede consumer_layer") + _integer(self.consumer_head, "consumer_head", minimum=0) + _integer(self.donor_head, "donor_head", minimum=0) + retained = _integer(self.retained_tiles, "retained_tiles", minimum=0) + eligible = _integer(self.eligible_tiles, "eligible_tiles", minimum=1) + if retained > eligible: + raise MaskReuseCalibrationError("retained_tiles must not exceed eligible_tiles") + + object.__setattr__(self, "model", self.model.strip()) + object.__setattr__(self, "prompt_id", self.prompt_id.strip()) + object.__setattr__( + self, + "source_capture_sha256", + _sha256(self.source_capture_sha256, "source_capture_sha256"), + ) + object.__setattr__( + self, + "target_sparsity", + _number( + self.target_sparsity, + "target_sparsity", + minimum=0.0, + maximum=1.0, + minimum_inclusive=False, + ), + ) + if self.target_sparsity >= 1.0: + raise MaskReuseCalibrationError("target_sparsity must be in (0, 1)") + object.__setattr__( + self, + "threshold_lambda", + _number( + self.threshold_lambda, + "threshold_lambda", + minimum=0.0, + maximum=1.0, + minimum_inclusive=False, + ), + ) + if self.threshold_lambda >= 1.0: + raise MaskReuseCalibrationError("threshold_lambda must be in (0, 1)") + object.__setattr__( + self, + "threshold_log2", + _number(self.threshold_log2, "threshold_log2", minimum=-math.inf, maximum=0.0), + ) + object.__setattr__( + self, + "anchor_dropped_mass", + _number(self.anchor_dropped_mass, "anchor_dropped_mass", minimum=0.0, maximum=1.0), + ) + object.__setattr__( + self, + "anchor_stats_by_layer", + _parse_anchor_stats_by_layer( + self.anchor_stats_by_layer, + require_canonical_string_keys=False, + ), + ) + object.__setattr__( + self, + "dropped_mass", + _number(self.dropped_mass, "dropped_mass", minimum=0.0, maximum=1.0), + ) + + @classmethod + def from_mapping(cls, raw: Mapping[str, object]) -> MaskReuseObservation: + """Build a validated observation from normalized JSON.""" + missing = _OBSERVATION_FIELDS - raw.keys() + extra = raw.keys() - _OBSERVATION_FIELDS + if missing or extra: + raise MaskReuseCalibrationError( + f"observation fields do not match the schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + values = dict(raw) + values["anchor_stats_by_layer"] = _parse_anchor_stats_by_layer( + raw["anchor_stats_by_layer"], + require_canonical_string_keys=True, + ) + return cls(**values) # type: ignore[arg-type] + + def to_mapping(self) -> dict[str, object]: + """Return the normalized JSON representation.""" + result = { + field: getattr(self, field) + for field in _OBSERVATION_FIELDS + if field != "anchor_stats_by_layer" + } + result["anchor_stats_by_layer"] = { + str(layer): { + "retained_tiles": list(stats.retained_tiles), + "dropped_mass": list(stats.dropped_mass), + } + for layer, stats in self.anchor_stats_by_layer.items() + } + return result + + +def _reject_duplicate_json_keys(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise MaskReuseCalibrationError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + +def parse_mask_reuse_observations(lines: Iterable[str]) -> list[MaskReuseObservation]: + """Parse strict normalized observation JSONL.""" + observations: list[MaskReuseObservation] = [] + for line_number, line in enumerate(lines, start=1): + if not line.strip(): + continue + try: + raw = json.loads(line, object_pairs_hook=_reject_duplicate_json_keys) + except json.JSONDecodeError as error: + raise MaskReuseCalibrationError( + f"line {line_number}: invalid JSON: {error.msg}" + ) from error + except MaskReuseCalibrationError as error: + raise MaskReuseCalibrationError(f"line {line_number}: {error}") from error + if not isinstance(raw, dict): + raise MaskReuseCalibrationError(f"line {line_number}: observation must be an object") + try: + observations.append(MaskReuseObservation.from_mapping(raw)) + except MaskReuseCalibrationError as error: + raise MaskReuseCalibrationError(f"line {line_number}: {error}") from error + if not observations: + raise MaskReuseCalibrationError("input contains no mask-reuse observations") + return observations + + +def load_mask_reuse_observations(path: str | Path) -> list[MaskReuseObservation]: + """Load normalized mask-reuse observations from JSONL.""" + with Path(path).open(encoding="utf-8") as handle: + return parse_mask_reuse_observations(handle) + + +def _find_skip_softmax_group(raw: Mapping[str, object]) -> Mapping[str, object]: + current: object = raw + if "sparse_attention_config" in raw: + current = raw["sparse_attention_config"] + if not isinstance(current, Mapping): + raise MaskReuseCalibrationError("sparse_attention_config must be an object") + if "config_groups" in current: + groups = current["config_groups"] + if not isinstance(groups, Mapping): + raise MaskReuseCalibrationError("config_groups must be an object") + matches = [ + group + for group in groups.values() + if isinstance(group, Mapping) and group.get("algorithm") == "skip_softmax" + ] + if len(matches) != 1: + raise MaskReuseCalibrationError( + "vanilla config must contain exactly one skip_softmax config group" + ) + current = matches[0] + if isinstance(current, Mapping) and "threshold_scale_factor" in current: + current = current["threshold_scale_factor"] + if not isinstance(current, Mapping): + raise MaskReuseCalibrationError("threshold_scale_factor must be an object") + return current + + +def canonical_prefill_threshold_scale_factor( + vanilla_calibration: Mapping[str, object], +) -> dict[str, object]: + """Canonicalize ModelOpt fit parameters or exported skip-softmax metadata.""" + raw = _find_skip_softmax_group(vanilla_calibration) + if "calibration_params" in raw: + raw = _find_skip_softmax_group(raw["calibration_params"]) # type: ignore[arg-type] + formula = raw.get("formula", _FORMULA) + if formula != _FORMULA: + raise MaskReuseCalibrationError("vanilla calibration uses an unsupported formula") + params = raw.get("prefill") + if not isinstance(params, Mapping): + raise MaskReuseCalibrationError("vanilla calibration requires prefill fit parameters") + unknown = params.keys() - { + "a", + "b", + "min_observed_sparsity", + "max_observed_sparsity", + } + if unknown or not {"a", "b"} <= params.keys(): + raise MaskReuseCalibrationError( + f"prefill fit requires a and b and contains unknown fields {sorted(unknown)}" + ) + prefill: dict[str, float] = { + "a": _number(params["a"], "prefill.a", minimum=0.0, minimum_inclusive=False), + "b": _number(params["b"], "prefill.b", minimum=0.0, maximum=20.0), + } + bounds = {"min_observed_sparsity", "max_observed_sparsity"} & params.keys() + if bounds and len(bounds) != 2: + raise MaskReuseCalibrationError("observed sparsity bounds must appear together") + if bounds: + lower = _number( + params["min_observed_sparsity"], + "prefill.min_observed_sparsity", + minimum=0.0, + maximum=1.0, + ) + upper = _number( + params["max_observed_sparsity"], + "prefill.max_observed_sparsity", + minimum=0.0, + maximum=1.0, + ) + if lower > upper: + raise MaskReuseCalibrationError("observed sparsity range is reversed") + prefill.update(min_observed_sparsity=lower, max_observed_sparsity=upper) + return {"formula": _FORMULA, "prefill": prefill} + + +def _normalize_topology(raw: Mapping[str, object]) -> tuple[tuple[int, ...], dict[int, int]]: + if set(raw) != {"anchors", "nearest"}: + raise MaskReuseCalibrationError("topology must contain exactly anchors and nearest") + raw_anchors = raw["anchors"] + if not isinstance(raw_anchors, list) or not raw_anchors: + raise MaskReuseCalibrationError("topology anchors must be a non-empty list") + anchors = tuple(sorted(_integer(value, "topology anchor", minimum=0) for value in raw_anchors)) + if len(anchors) != len(set(anchors)): + raise MaskReuseCalibrationError("topology anchors must be unique") + raw_nearest = raw["nearest"] + if not isinstance(raw_nearest, Mapping): + raise MaskReuseCalibrationError("topology nearest must be an object") + nearest: dict[int, int] = {} + for raw_layer, raw_anchor in raw_nearest.items(): + try: + layer = int(raw_layer) + except (TypeError, ValueError) as error: + raise MaskReuseCalibrationError("topology nearest has a non-integer key") from error + if str(layer) != str(raw_layer) or layer in nearest: + raise MaskReuseCalibrationError("topology nearest keys must be canonical and unique") + nearest[layer] = _integer(raw_anchor, f"topology nearest[{layer}]", minimum=0) + anchor_set = set(anchors) + if not anchor_set <= nearest.keys(): + raise MaskReuseCalibrationError("topology nearest must include every anchor") + for layer, anchor in nearest.items(): + if anchor not in anchor_set or (layer != anchor and anchor >= layer): + raise MaskReuseCalibrationError(f"topology layer {layer} has invalid anchor {anchor}") + if layer in anchor_set and anchor != layer: + raise MaskReuseCalibrationError(f"topology anchor {layer} must map to itself") + if not any(layer != anchor for layer, anchor in nearest.items()): + raise MaskReuseCalibrationError("topology must contain at least one reuse layer") + return anchors, dict(sorted(nearest.items())) + + +def _bucket_key(bucket: Bucket) -> tuple[int, float]: + return bucket[0], math.inf if bucket[1] is None else float(bucket[1]) + + +@dataclass(frozen=True, slots=True) +class _Choice: + donor_head: int + fallback: bool + retained_tiles: int + + +@dataclass(frozen=True, slots=True) +class _Selection: + target_sparsity: float | None + choices: Mapping[ConsumerHead, _Choice] + frontier: tuple[Mapping[str, object], ...] + exact_reason: str | None = None + + +@dataclass(frozen=True, slots=True) +class _BucketIndex: + observations: Mapping[ObservationKey, MaskReuseObservation] + prompts: Mapping[str, tuple[str, ...]] + target_menus: Mapping[tuple[str, str], frozenset[float]] + donor_menus: Mapping[tuple[str, str, float, ConsumerHead], frozenset[int]] + eligible: Mapping[tuple[str, str, ConsumerHead], int] + anchor_masks: Mapping[AnchorKey, tuple[int, int, float]] + anchors: tuple[int, ...] + + +@dataclass(slots=True) +class _ReuseEvaluation: + eligible_tiles: int = 0 + retained_tiles: int = 0 + sparse_observations: int = 0 + violations: int = 0 + dropped_mass_sum: float = 0.0 + worst_dropped_mass: float = 0.0 + + def add(self, other: _ReuseEvaluation) -> None: + self.eligible_tiles += other.eligible_tiles + self.retained_tiles += other.retained_tiles + self.sparse_observations += other.sparse_observations + self.violations += other.violations + self.dropped_mass_sum += other.dropped_mass_sum + self.worst_dropped_mass = max(self.worst_dropped_mass, other.worst_dropped_mass) + + def to_mapping(self) -> dict[str, object]: + return { + "eligible_tiles": self.eligible_tiles, + "retained_tiles": self.retained_tiles, + "bmm1_tile_savings_fraction": ( + 1.0 - self.retained_tiles / self.eligible_tiles if self.eligible_tiles else 0.0 + ), + "sparse_head_prompt_observations": self.sparse_observations, + "constraint_violation_count": self.violations, + "constraint_violation_rate": ( + self.violations / self.sparse_observations if self.sparse_observations else 0.0 + ), + "mean_dropped_mass": ( + self.dropped_mass_sum / self.sparse_observations + if self.sparse_observations + else 0.0 + ), + "worst_dropped_mass": self.worst_dropped_mass, + } + + +@dataclass(slots=True) +class _AnchorEvaluation: + eligible_tiles: int = 0 + retained_tiles: int = 0 + prompt_count: int = 0 + violations: int = 0 + prompt_mean_sum: float = 0.0 + worst_prompt_mean: float = 0.0 + + def add(self, other: _AnchorEvaluation) -> None: + self.eligible_tiles += other.eligible_tiles + self.retained_tiles += other.retained_tiles + self.prompt_count += other.prompt_count + self.violations += other.violations + self.prompt_mean_sum += other.prompt_mean_sum + self.worst_prompt_mean = max(self.worst_prompt_mean, other.worst_prompt_mean) + + def to_mapping(self, *, exact: bool = False) -> dict[str, object]: + return { + "policy_exact": exact, + "constraint_statistic": "worst_prompt_mean_anchor_dropped_mass", + "eligible_tiles": self.eligible_tiles, + "retained_tiles": self.retained_tiles, + "bmm2_tile_savings_fraction": ( + 1.0 - self.retained_tiles / self.eligible_tiles if self.eligible_tiles else 0.0 + ), + "evaluated_prompt_count": self.prompt_count, + "constraint_violation_count": self.violations, + "constraint_violation_rate": ( + self.violations / self.prompt_count if self.prompt_count else 0.0 + ), + "mean_prompt_mean_anchor_dropped_mass": ( + self.prompt_mean_sum / self.prompt_count if self.prompt_count else 0.0 + ), + "worst_prompt_mean_anchor_dropped_mass": self.worst_prompt_mean, + } + + +def _normalize_observations( + values: Sequence[MaskReuseObservation | Mapping[str, object]], +) -> list[MaskReuseObservation]: + if not values: + raise MaskReuseCalibrationError("at least one observation is required") + return [ + value + if isinstance(value, MaskReuseObservation) + else MaskReuseObservation.from_mapping(value) + for value in values + ] + + +def _validate_thresholds( + observations: Sequence[MaskReuseObservation], threshold_scale_factor: Mapping[str, object] +) -> None: + params = threshold_scale_factor["prefill"] + assert isinstance(params, Mapping) + a = float(params["a"]) + b = float(params["b"]) + lower = params.get("min_observed_sparsity") + upper = params.get("max_observed_sparsity") + for observation in observations: + target = observation.target_sparsity + if (lower is not None and target < float(lower)) or ( + upper is not None and target > float(upper) + ): + raise MaskReuseCalibrationError( + f"target_sparsity={target} is outside the observed vanilla calibration range" + ) + expected_log2 = ( + math.log2(a) + b * target * math.log2(math.e) - math.log2(observation.sample_length) + ) + expected_lambda = float(getattr(math, "exp2")(expected_log2)) + if not 0.0 < expected_lambda < 1.0: + raise MaskReuseCalibrationError( + "vanilla calibration derives a threshold outside (0, 1) for " + f"prompt={observation.prompt_id!r}, target_sparsity={target}" + ) + if observation.threshold_log2.hex() != expected_log2.hex(): + raise MaskReuseCalibrationError( + "threshold_log2 does not match the log-domain vanilla fit for " + f"prompt={observation.prompt_id!r}, target_sparsity={target}: " + f"observed={observation.threshold_log2.hex()}, " + f"expected={expected_log2.hex()}" + ) + if observation.threshold_lambda.hex() != expected_lambda.hex(): + raise MaskReuseCalibrationError( + "threshold_lambda does not match " + "exp2(log2(a) + b * target_sparsity * log2(e) - log2(sample_length)) for " + f"prompt={observation.prompt_id!r}, target_sparsity={target}: " + f"observed={observation.threshold_lambda.hex()}, " + f"expected={expected_lambda.hex()}" + ) + + +def _validate_dataset( + observations: Sequence[MaskReuseObservation], + *, + nearest: Mapping[int, int], +) -> tuple[ + str, int, tuple[Bucket, ...], tuple[ConsumerHead, ...], dict[Bucket, list[MaskReuseObservation]] +]: + by_split = { + split: [observation for observation in observations if observation.split == split] + for split in _SPLITS + } + if any(not rows for rows in by_split.values()): + raise MaskReuseCalibrationError("observations require calibration and heldout splits") + models = {observation.model for observation in observations} + if len(models) != 1: + raise MaskReuseCalibrationError("observations must contain exactly one model") + calibration_prompts = {row.prompt_id for row in by_split["calibration"]} + heldout_prompts = {row.prompt_id for row in by_split["heldout"]} + if calibration_prompts & heldout_prompts: + raise MaskReuseCalibrationError("prompt IDs overlap calibration and heldout splits") + calibration_sources = {row.source_capture_sha256 for row in by_split["calibration"]} + heldout_sources = {row.source_capture_sha256 for row in by_split["heldout"]} + if calibration_sources & heldout_sources: + raise MaskReuseCalibrationError("source captures overlap calibration and heldout splits") + + consumer_to_anchor = {layer: anchor for layer, anchor in nearest.items() if layer != anchor} + max_head = max(max(row.consumer_head, row.donor_head) for row in observations) + global_num_heads = max_head + 1 + targets = tuple( + (layer, head) for layer in sorted(consumer_to_anchor) for head in range(global_num_heads) + ) + expected_targets = set(targets) + seen: set[tuple[object, ...]] = set() + capture_sources: dict[tuple[str, str, Bucket], str] = {} + by_bucket: dict[Bucket, list[MaskReuseObservation]] = defaultdict(list) + for row in observations: + expected_anchor = consumer_to_anchor.get(row.consumer_layer) + if expected_anchor != row.anchor_layer: + raise MaskReuseCalibrationError( + f"consumer layer {row.consumer_layer} does not match the explicit topology" + ) + bucket = (row.min_kv_tokens, row.max_kv_tokens) + by_bucket[bucket].append(row) + identity = ( + row.model, + bucket, + row.split, + row.prompt_id, + row.target_sparsity, + row.consumer_layer, + row.consumer_head, + row.donor_head, + ) + if identity in seen: + raise MaskReuseCalibrationError("observations contain a duplicate candidate row") + seen.add(identity) + capture = (row.split, row.prompt_id, bucket) + previous_source = capture_sources.setdefault(capture, row.source_capture_sha256) + if previous_source != row.source_capture_sha256: + raise MaskReuseCalibrationError("one prompt/context capture has multiple fingerprints") + + split_buckets = { + split: {(row.min_kv_tokens, row.max_kv_tokens) for row in rows} + for split, rows in by_split.items() + } + if split_buckets["calibration"] != split_buckets["heldout"]: + raise MaskReuseCalibrationError("heldout context buckets must match calibration buckets") + buckets = tuple(sorted(split_buckets["calibration"], key=_bucket_key)) + previous_max: int | None = 0 + for minimum, maximum in buckets: + if previous_max is None or minimum <= previous_max: + raise MaskReuseCalibrationError("context buckets must be ordered and non-overlapping") + previous_max = maximum + for bucket in buckets: + for split in _SPLITS: + actual = { + (row.consumer_layer, row.consumer_head) + for row in by_bucket[bucket] + if row.split == split + } + if actual != expected_targets: + raise MaskReuseCalibrationError( + f"{split} bucket {bucket} does not cover every consumer head" + ) + return next(iter(models)), global_num_heads, buckets, targets, dict(by_bucket) + + +def _index_bucket( + rows: Sequence[MaskReuseObservation], + *, + anchors: tuple[int, ...], + global_num_heads: int, + targets: Sequence[ConsumerHead], +) -> _BucketIndex: + observations: dict[ObservationKey, MaskReuseObservation] = {} + prompts: dict[str, set[str]] = defaultdict(set) + target_menus: dict[tuple[str, str], set[float]] = defaultdict(set) + donor_menus: dict[tuple[str, str, float, ConsumerHead], set[int]] = defaultdict(set) + eligible: dict[tuple[str, str, ConsumerHead], int] = {} + anchor_masks: dict[AnchorKey, tuple[int, int, float]] = {} + prompt_targets: dict[tuple[str, str, float], tuple[int, float, float]] = {} + anchor_payloads: dict[tuple[str, str, float], Mapping[int, AnchorLayerStats]] = {} + payload_eligible: dict[tuple[str, str, float], int] = {} + for row in rows: + target = (row.consumer_layer, row.consumer_head) + key: ObservationKey = ( + row.split, + row.prompt_id, + row.target_sparsity, + row.consumer_layer, + row.consumer_head, + row.donor_head, + ) + observations[key] = row + prompts[row.split].add(row.prompt_id) + target_menus[(row.split, row.prompt_id)].add(row.target_sparsity) + donor_menus[(row.split, row.prompt_id, row.target_sparsity, target)].add(row.donor_head) + eligible_key = (row.split, row.prompt_id, target) + previous_eligible = eligible.setdefault(eligible_key, row.eligible_tiles) + if previous_eligible != row.eligible_tiles: + raise MaskReuseCalibrationError("eligible_tiles differs across candidates") + prompt_target = (row.split, row.prompt_id, row.target_sparsity) + sample_and_threshold = ( + row.sample_length, + row.threshold_lambda, + row.threshold_log2, + ) + if prompt_targets.setdefault(prompt_target, sample_and_threshold) != sample_and_threshold: + raise MaskReuseCalibrationError( + "sample_length, threshold_lambda, or threshold_log2 differs within a capture" + ) + previous_payload = anchor_payloads.setdefault(prompt_target, row.anchor_stats_by_layer) + if previous_payload != row.anchor_stats_by_layer: + raise MaskReuseCalibrationError( + "anchor_stats_by_layer differs across repeated candidate rows" + ) + previous_payload_eligible = payload_eligible.setdefault(prompt_target, row.eligible_tiles) + if previous_payload_eligible != row.eligible_tiles: + raise MaskReuseCalibrationError( + "eligible_tiles differs within one capture and target_sparsity" + ) + + expected_anchors = set(anchors) + for prompt_target, payload in anchor_payloads.items(): + actual_anchors = set(payload) + if actual_anchors != expected_anchors: + raise MaskReuseCalibrationError( + "anchor_stats_by_layer does not exactly cover topology anchors; " + f"missing={sorted(expected_anchors - actual_anchors)}, " + f"extra={sorted(actual_anchors - expected_anchors)}" + ) + eligible_tiles = payload_eligible[prompt_target] + for anchor, stats in payload.items(): + if len(stats.retained_tiles) != global_num_heads: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{anchor}] has head width " + f"{len(stats.retained_tiles)}, expected {global_num_heads}" + ) + for head, (retained, dropped) in enumerate( + zip(stats.retained_tiles, stats.dropped_mass, strict=True) + ): + if retained > eligible_tiles: + raise MaskReuseCalibrationError( + f"anchor_stats_by_layer[{anchor}].retained_tiles[{head}] " + "exceeds eligible_tiles" + ) + split, prompt, target_sparsity = prompt_target + anchor_masks[(split, prompt, target_sparsity, anchor, head)] = ( + retained, + eligible_tiles, + dropped, + ) + + for row in rows: + stats = row.anchor_stats_by_layer[row.anchor_layer] + payload_candidate = ( + stats.retained_tiles[row.donor_head], + stats.dropped_mass[row.donor_head], + ) + if payload_candidate != (row.retained_tiles, row.anchor_dropped_mass): + raise MaskReuseCalibrationError( + "candidate retained_tiles/anchor_dropped_mass does not match anchor_stats_by_layer" + ) + + first_prompt = min(prompts["calibration"]) + expected_targets = target_menus[("calibration", first_prompt)] + full_donors = set(range(global_num_heads)) + for split in _SPLITS: + for prompt in prompts[split]: + if target_menus[(split, prompt)] != expected_targets: + raise MaskReuseCalibrationError("target_sparsity menu differs across captures") + for target_sparsity in expected_targets: + for target in targets: + if donor_menus[(split, prompt, target_sparsity, target)] != full_donors: + raise MaskReuseCalibrationError("candidate donor menu is incomplete") + for anchor in anchors: + for head in range(global_num_heads): + if (split, prompt, target_sparsity, anchor, head) not in anchor_masks: + raise MaskReuseCalibrationError( + "anchor/head mask observations are incomplete" + ) + return _BucketIndex( + observations=observations, + prompts={split: tuple(sorted(values)) for split, values in prompts.items()}, + target_menus={key: frozenset(values) for key, values in target_menus.items()}, + donor_menus={key: frozenset(values) for key, values in donor_menus.items()}, + eligible=eligible, + anchor_masks=anchor_masks, + anchors=anchors, + ) + + +def _evaluate_anchor( + index: _BucketIndex, + *, + split: str, + target_sparsity: float | None, + global_num_heads: int, + maximum: float, +) -> _AnchorEvaluation: + result = _AnchorEvaluation() + for prompt in index.prompts[split]: + selected_target = ( + min(index.target_menus[(split, prompt)]) if target_sparsity is None else target_sparsity + ) + values: list[float] = [] + retained_tiles = 0 + eligible_tiles = 0 + for anchor in index.anchors: + for head in range(global_num_heads): + retained, eligible, dropped = index.anchor_masks[ + (split, prompt, selected_target, anchor, head) + ] + values.append(0.0 if target_sparsity is None else dropped) + retained_tiles += eligible if target_sparsity is None else retained + eligible_tiles += eligible + prompt_mean = sum(values) / len(values) + result.eligible_tiles += eligible_tiles + result.retained_tiles += retained_tiles + result.prompt_count += 1 + result.prompt_mean_sum += prompt_mean + result.worst_prompt_mean = max(result.worst_prompt_mean, prompt_mean) + result.violations += int(prompt_mean > maximum) + return result + + +def _select_bucket( + index: _BucketIndex, + *, + targets: Sequence[ConsumerHead], + global_num_heads: int, + max_anchor_dropped_mass: float, + max_reuse_selection_dropped_mass: float, +) -> _Selection: + prompts = index.prompts["calibration"] + target_menu = tuple(sorted(index.target_menus[("calibration", prompts[0])])) + frontier: list[Mapping[str, object]] = [] + candidates: list[tuple[tuple[object, ...], _Selection]] = [] + for target_sparsity in target_menu: + anchor = _evaluate_anchor( + index, + split="calibration", + target_sparsity=target_sparsity, + global_num_heads=global_num_heads, + maximum=max_anchor_dropped_mass, + ) + choices: dict[ConsumerHead, _Choice] = {} + total_retained = 0 + fallback_count = 0 + for target in targets: + feasible: list[tuple[int, int]] = [] + for donor in range(global_num_heads): + candidate_rows = [ + index.observations[ + ( + "calibration", + prompt, + target_sparsity, + target[0], + target[1], + donor, + ) + ] + for prompt in prompts + ] + if all( + row.dropped_mass <= max_reuse_selection_dropped_mass for row in candidate_rows + ): + feasible.append((sum(row.retained_tiles for row in candidate_rows), donor)) + if feasible: + retained, donor = min(feasible) + choice = _Choice(donor, False, retained) + else: + fallback_count += 1 + retained = sum( + index.eligible[("calibration", prompt, target)] for prompt in prompts + ) + choice = _Choice(0, True, retained) + choices[target] = choice + total_retained += retained + signature = tuple(choices[target].donor_head for target in targets) + combined_tile_cost = 2 * total_retained + anchor.retained_tiles + frontier.append( + { + "target_sparsity": target_sparsity, + "anchor_safe": anchor.violations == 0, + "retained_reuse_tiles": total_retained, + "retained_anchor_tiles": anchor.retained_tiles, + "combined_tile_cost": combined_tile_cost, + "fallback_head_count": fallback_count, + "anchor_calibration": anchor.to_mapping(), + } + ) + if anchor.violations == 0: + rank = ( + combined_tile_cost, + fallback_count, + target_sparsity, + signature, + ) + candidates.append((rank, _Selection(target_sparsity, choices, ()))) + if not candidates: + dense_choices = { + target: _Choice( + 0, True, sum(index.eligible[("calibration", prompt, target)] for prompt in prompts) + ) + for target in targets + } + return _Selection( + None, + dense_choices, + tuple(frontier), + "no_target_sparsity_satisfied_anchor_calibration_constraint", + ) + _, selected = min(candidates, key=lambda item: item[0]) + return _Selection(selected.target_sparsity, selected.choices, tuple(frontier)) + + +def _evaluate_reuse( + selection: _Selection, + index: _BucketIndex, + *, + split: str, + maximum: float, +) -> _ReuseEvaluation: + result = _ReuseEvaluation() + for prompt in index.prompts[split]: + for target, choice in sorted(selection.choices.items()): + eligible = index.eligible[(split, prompt, target)] + result.eligible_tiles += eligible + if choice.fallback: + result.retained_tiles += eligible + continue + assert selection.target_sparsity is not None + row = index.observations[ + ( + split, + prompt, + selection.target_sparsity, + target[0], + target[1], + choice.donor_head, + ) + ] + result.retained_tiles += row.retained_tiles + result.sparse_observations += 1 + result.dropped_mass_sum += row.dropped_mass + result.worst_dropped_mass = max(result.worst_dropped_mass, row.dropped_mass) + result.violations += int(row.dropped_mass > maximum) + return result + + +def _canonical_digest(observations: Sequence[MaskReuseObservation]) -> str: + rows = [ + json.dumps(row.to_mapping(), sort_keys=True, separators=(",", ":")).encode() + for row in observations + ] + digest = sha256() + for row in sorted(rows): + digest.update(sha256(row).digest()) + return digest.hexdigest() + + +def _deployment_geometry( + observations: Sequence[MaskReuseObservation], +) -> dict[str, object]: + by_capture: dict[tuple[object, ...], tuple[int, int, int]] = {} + for row in observations: + identity = ( + row.split, + row.prompt_id, + row.source_capture_sha256, + row.min_kv_tokens, + row.max_kv_tokens, + ) + geometry = (row.q_tokens, row.kv_tokens, row.q_start_tokens) + if by_capture.setdefault(identity, geometry) != geometry: + raise MaskReuseCalibrationError( + "q_tokens, kv_tokens, or q_start_tokens differs within one capture" + ) + geometry_rows = [ + { + "split": identity[0], + "prompt_id": identity[1], + "source_capture_sha256": identity[2], + "min_kv_tokens": identity[3], + "max_kv_tokens": identity[4], + "q_tokens": geometry[0], + "kv_tokens": geometry[1], + "q_start_tokens": geometry[2], + } + for identity, geometry in sorted(by_capture.items(), key=lambda item: str(item[0])) + ] + return { + "contract": dict(_DEPLOYMENT_GEOMETRY_CONTRACT), + "observations": geometry_rows, + } + + +def _canonical_evidence(raw: Mapping[str, object]) -> dict[str, str]: + missing = _EVIDENCE_FIELDS - raw.keys() + extra = raw.keys() - _EVIDENCE_FIELDS + if missing or extra: + raise MaskReuseCalibrationError( + f"evidence fields do not match schema v3; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + return {field: _sha256(raw[field], f"evidence.{field}") for field in sorted(raw)} + + +def calibrate_mask_reuse_policy( + observations: Sequence[MaskReuseObservation | Mapping[str, object]], + *, + vanilla_calibration: Mapping[str, object], + topology: Mapping[str, object], + checkpoint_manifest: VerifiedCheckpointManifest, + evidence: Mapping[str, object], + max_anchor_dropped_mass: float, + max_reuse_dropped_mass: float, + max_reuse_selection_dropped_mass: float | None = None, + source_provenance: Mapping[str, object] | None = None, +) -> dict[str, object]: + """Select and evaluate a fail-closed schema-v3 legacy candidate.""" + if not isinstance(checkpoint_manifest, VerifiedCheckpointManifest): + raise MaskReuseCalibrationError( + "checkpoint_manifest must be returned by verify_checkpoint_manifest" + ) + rows = _normalize_observations(observations) + threshold_scale_factor = canonical_prefill_threshold_scale_factor(vanilla_calibration) + prefill_params = threshold_scale_factor["prefill"] + assert isinstance(prefill_params, Mapping) + if not {"min_observed_sparsity", "max_observed_sparsity"} <= prefill_params.keys(): + raise MaskReuseCalibrationError( + "schema-v3 export requires observed prefill sparsity bounds" + ) + _validate_thresholds(rows, threshold_scale_factor) + anchors, nearest = _normalize_topology(topology) + checkpoint_identity = checkpoint_manifest.sha256 + geometry = _deployment_geometry(rows) + canonical_evidence = _canonical_evidence(evidence) + anchor_bound = _number( + max_anchor_dropped_mass, "max_anchor_dropped_mass", minimum=0.0, maximum=1.0 + ) + reuse_bound = _number( + max_reuse_dropped_mass, "max_reuse_dropped_mass", minimum=0.0, maximum=1.0 + ) + selection_bound = ( + reuse_bound + if max_reuse_selection_dropped_mass is None + else _number( + max_reuse_selection_dropped_mass, + "max_reuse_selection_dropped_mass", + minimum=0.0, + maximum=1.0, + ) + ) + if selection_bound > reuse_bound: + raise MaskReuseCalibrationError( + "max_reuse_selection_dropped_mass must not exceed max_reuse_dropped_mass" + ) + model, global_num_heads, buckets, targets, by_bucket = _validate_dataset(rows, nearest=nearest) + if model != checkpoint_manifest.model: + raise MaskReuseCalibrationError( + "observation model does not match the verified checkpoint manifest" + ) + + context_policies: list[dict[str, object]] = [] + bucket_reports: list[dict[str, object]] = [] + target_menus: list[dict[str, object]] = [] + overall_reuse_calibration = _ReuseEvaluation() + overall_reuse_heldout = _ReuseEvaluation() + overall_anchor_calibration = _AnchorEvaluation() + overall_anchor_heldout = _AnchorEvaluation() + total_fallback = 0 + for bounds in buckets: + index = _index_bucket( + by_bucket[bounds], + anchors=anchors, + global_num_heads=global_num_heads, + targets=targets, + ) + selection = _select_bucket( + index, + targets=targets, + global_num_heads=global_num_heads, + max_anchor_dropped_mass=anchor_bound, + max_reuse_selection_dropped_mass=selection_bound, + ) + if selection.target_sparsity is not None and bounds[1] is None: + raise MaskReuseCalibrationError( + "a deployment-qualified sparse context bucket requires a finite maximum" + ) + reuse_calibration = _evaluate_reuse( + selection, index, split="calibration", maximum=reuse_bound + ) + reuse_heldout = _evaluate_reuse(selection, index, split="heldout", maximum=reuse_bound) + anchor_calibration = _evaluate_anchor( + index, + split="calibration", + target_sparsity=selection.target_sparsity, + global_num_heads=global_num_heads, + maximum=anchor_bound, + ) + anchor_heldout = _evaluate_anchor( + index, + split="heldout", + target_sparsity=selection.target_sparsity, + global_num_heads=global_num_heads, + maximum=anchor_bound, + ) + overall_reuse_calibration.add(reuse_calibration) + overall_reuse_heldout.add(reuse_heldout) + overall_anchor_calibration.add(anchor_calibration) + overall_anchor_heldout.add(anchor_heldout) + + policy: dict[str, object] + if selection.target_sparsity is None: + policy = { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "exact": True, + } + headmaps: dict[str, list[int]] = {} + fallback_heads: dict[str, list[int]] = {} + else: + headmaps = { + str(layer): [ + selection.choices[(layer, head)].donor_head for head in range(global_num_heads) + ] + for layer in sorted({target[0] for target in targets}) + } + fallback_heads = { + str(layer): [ + head + for head in range(global_num_heads) + if selection.choices[(layer, head)].fallback + ] + for layer in sorted({target[0] for target in targets}) + } + policy = { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "target_sparsity": selection.target_sparsity, + "headmaps": headmaps, + "fallback_heads": fallback_heads, + } + context_policies.append(policy) + fallback_count = sum(len(heads) for heads in fallback_heads.values()) + total_fallback += fallback_count + menu = [row["target_sparsity"] for row in selection.frontier] + target_menus.append( + { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "target_sparsities": menu, + } + ) + bucket_reports.append( + { + "min_kv_tokens": bounds[0], + "max_kv_tokens": bounds[1], + "selected_target_sparsity": selection.target_sparsity, + "selection_status": "sparse" + if selection.target_sparsity is not None + else selection.exact_reason, + "target_sparsity_frontier": list(selection.frontier), + "fallback_head_count": fallback_count, + "reuse_calibration": reuse_calibration.to_mapping(), + "reuse_heldout": reuse_heldout.to_mapping(), + "anchor_calibration": anchor_calibration.to_mapping( + exact=selection.target_sparsity is None + ), + "anchor_heldout": anchor_heldout.to_mapping( + exact=selection.target_sparsity is None + ), + } + ) + + provenance: dict[str, object] = { + "calibrator": "modelopt.mask_reuse", + "observation_schema_version": 1, + "input_observation_count": len(rows), + "canonical_input_sha256": _canonical_digest(rows), + "calibration_prompt_ids": sorted( + {row.prompt_id for row in rows if row.split == "calibration"} + ), + "heldout_prompt_ids": sorted({row.prompt_id for row in rows if row.split == "heldout"}), + "selection_split": "calibration", + "evaluation_split": "heldout", + "threshold_semantics": "a * exp(b * target_sparsity) / sample_length", + "threshold_implementation": { + "threshold_log2": ("log2(a) + b * target_sparsity * log2(e) - log2(sample_length)"), + "threshold_lambda": "exp2(threshold_log2)", + "validation": "exact IEEE-754 binary64 hex equality", + }, + "checkpoint_manifest_sha256": checkpoint_identity, + "target_sparsity_menus": target_menus, + "constraints": { + "anchor": { + "metric": "worst_prompt_mean_anchor_dropped_mass", + "comparison": "<=", + "maximum": anchor_bound, + }, + "reuse": { + "metric": "per_prompt_candidate_reuse_dropped_mass", + "comparison": "<=", + "selection_maximum": selection_bound, + "evaluation_maximum": reuse_bound, + }, + }, + "tie_breaks": [ + "reject target sparsities exceeding the calibration anchor bound", + "minimum equal-BMM combined tile cost 2*A_R + A_A", + "fewest exact-fallback consumer heads", + "smallest target sparsity", + "lexicographically smallest donor-head map", + ], + "selection_cost": { + "formula": "2 * retained_reuse_tiles + retained_anchor_tiles", + "bmm1_weight": 1.0, + "bmm2_weight": 1.0, + }, + } + if source_provenance is not None: + provenance["source"] = dict(source_provenance) + + report = { + "model": model, + "checkpoint_manifest_sha256": checkpoint_identity, + "constraints": provenance["constraints"], + "selection_unit": "context_bucket", + "by_bucket": bucket_reports, + "overall": { + "consumer_head_bucket_count": len(targets) * len(buckets), + "fallback_head_bucket_count": total_fallback, + "fallback_fraction": total_fallback / (len(targets) * len(buckets)), + "reuse_calibration": overall_reuse_calibration.to_mapping(), + "reuse_heldout": overall_reuse_heldout.to_mapping(), + "anchor_calibration": overall_anchor_calibration.to_mapping(), + "anchor_heldout": overall_anchor_heldout.to_mapping(), + }, + "promotion": { + "status": "candidate_only", + "eligible": False, + "reasons": [ + "legacy observations are not capture-schema-v2 checkpoint-bound records", + "grouped inner-fold and preregistered outer gates were not evaluated", + "deployment rectangular-geometry promotion gate was not evaluated", + ], + }, + } + return { + "version": 3, + "promotion_status": "candidate_only", + "phase": "prefill", + "decode": {"mode": "dense"}, + "calibration_protocol": _CALIBRATION_PROTOCOL, + "producer": {"name": "modelopt", "version": modelopt.__version__}, + "evidence": canonical_evidence, + "threshold_scale_factor": threshold_scale_factor, + "model": model, + "checkpoint_manifest_sha256": checkpoint_identity, + "global_num_heads": global_num_heads, + "anchors": list(anchors), + "nearest": {str(layer): anchor for layer, anchor in nearest.items()}, + "deployment_geometry_validated": False, + "deployment_geometry": geometry, + "context_policies": context_policies, + "provenance": provenance, + "calibration_report": report, + } diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py new file mode 100644 index 00000000000..6977134bb73 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py @@ -0,0 +1,1211 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Streaming selector for compact mask-reuse capture JSONL. + +One compact record stores the anchor vectors and consumer ``[H, H]`` risk +matrices for a single prompt and target sparsity. The selector makes three +semantic passes (validation, calibration selection, frozen-policy evaluation) +and hashes the exact file bytes before and after selection/evaluation to detect +concurrent mutation. Memory scales with the current and immediately prior +capture plus ``O(targets * consumers * H^2)`` aggregate arrays rather than with +every prompt-level candidate cell. +""" + +from __future__ import annotations + +import json +import math +import unicodedata +from collections import defaultdict +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import cast + +import numpy as np + +import modelopt + +from .checkpoint_manifest import VerifiedCheckpointManifest +from .mask_reuse import ( + _CALIBRATION_PROTOCOL, + _DEPLOYMENT_GEOMETRY_CONTRACT, + AnchorLayerStats, + Bucket, + ConsumerHead, + MaskReuseCalibrationError, + _AnchorEvaluation, + _bucket_key, + _canonical_evidence, + _Choice, + _normalize_topology, + _number, + _ReuseEvaluation, + _Selection, + _sha256, + canonical_prefill_threshold_scale_factor, +) + +__all__ = [ + "CompactMaskReuseCapture", + "CompactMaskReuseCaptureSource", + "calibrate_compact_mask_reuse_policy", + "load_compact_mask_reuse_captures", +] + +_CAPTURE_FIELDS = frozenset( + { + "compact_capture_schema_version", + "invocation", + "geometry", + "global_num_heads", + "eligible_tiles", + "anchor_stats_by_layer", + "consumer_layers", + } +) +_INVOCATION_FIELDS = frozenset( + { + "capture_schema_version", + "model", + "checkpoint_manifest_sha256", + "split", + "partition", + "inner_fold", + "prompt_id", + "source", + "source_group_sha256", + "source_capture_sha256", + "min_kv_tokens", + "max_kv_tokens", + "target_sparsity_hex", + "sample_length", + "threshold_log2_hex", + "threshold_lambda_hex", + "expected_geometry", + } +) +_GEOMETRY_FIELDS = frozenset({"q_tokens", "kv_tokens", "q_start_tokens"}) +_ANCHOR_FIELDS = frozenset({"retained_tiles", "dropped_mass"}) +_CONSUMER_FIELDS = frozenset({"anchor_layer", "dropped_mass"}) +_SPLITS = ("calibration", "heldout") +_PARTITIONS = ("development", "outer_test") +_MONOTONIC_ATOL = 1e-7 + + +def _exact_fields(raw: Mapping[str, object], expected: frozenset[str], label: str) -> None: + missing = expected - raw.keys() + extra = raw.keys() - expected + if missing or extra: + raise MaskReuseCalibrationError( + f"{label} fields do not match the schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + + +def _integer(value: object, label: str, *, minimum: int = 0) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise MaskReuseCalibrationError(f"{label} must be an integer >= {minimum}") + return value + + +def _text(value: object, label: str) -> str: + if not isinstance(value, str) or not value: + raise MaskReuseCalibrationError(f"{label} must be a non-empty string") + if unicodedata.normalize("NFC", value) != value or any( + ord(character) < 32 for character in value + ): + raise MaskReuseCalibrationError(f"{label} must be NFC text without control characters") + return value + + +def _float_hex(value: object, label: str) -> float: + if not isinstance(value, str): + raise MaskReuseCalibrationError(f"{label} must be a canonical float.hex string") + try: + parsed = float.fromhex(value) + except ValueError as error: + raise MaskReuseCalibrationError(f"{label} must be a canonical float.hex string") from error + if not math.isfinite(parsed) or parsed.hex() != value: + raise MaskReuseCalibrationError(f"{label} must be a canonical finite float.hex string") + return parsed + + +def _geometry(raw: object, label: str) -> dict[str, int]: + if not isinstance(raw, Mapping): + raise MaskReuseCalibrationError(f"{label} must be an object") + _exact_fields(raw, _GEOMETRY_FIELDS, label) + q_tokens = _integer(raw["q_tokens"], f"{label}.q_tokens", minimum=129) + if q_tokens > int(cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["max_query_chunk_tokens"])): + raise MaskReuseCalibrationError(f"{label}.q_tokens exceeds the deployment maximum") + kv_tokens = _integer(raw["kv_tokens"], f"{label}.kv_tokens", minimum=q_tokens) + q_start = _integer(raw["q_start_tokens"], f"{label}.q_start_tokens") + alignment = int( + cast("int", _DEPLOYMENT_GEOMETRY_CONTRACT["query_chunk_start_alignment_tokens"]) + ) + if q_start % alignment or q_start + q_tokens != kv_tokens: + raise MaskReuseCalibrationError(f"{label} is not a bottom-right aligned final chunk") + return {"q_tokens": q_tokens, "kv_tokens": kv_tokens, "q_start_tokens": q_start} + + +def _eligible_tiles(geometry: Mapping[str, int]) -> int: + q_blocks = (geometry["q_tokens"] + 127) // 128 + first_eligible = geometry["q_start_tokens"] // 128 + 1 + return q_blocks * (2 * first_eligible + q_blocks - 1) // 2 + + +@dataclass(frozen=True, slots=True) +class CompactConsumerStats: + """One consumer layer's global consumer-by-donor dropped-mass matrix.""" + + anchor_layer: int + dropped_mass: Sequence[Sequence[float]] + + +@dataclass(frozen=True, slots=True) +class CompactMaskReuseCapture: + """One prompt/target capture decoded without expanding candidate rows.""" + + model: str + checkpoint_manifest_sha256: str + split: str + partition: str + inner_fold: int | None + prompt_id: str + source: str + source_group_sha256: str + source_capture_sha256: str + min_kv_tokens: int + max_kv_tokens: int | None + target_sparsity: float + sample_length: int + threshold_log2: float + threshold_lambda: float + geometry: Mapping[str, int] + global_num_heads: int + eligible_tiles: int + anchor_stats_by_layer: Mapping[int, AnchorLayerStats] + consumer_layers: Mapping[int, CompactConsumerStats] + + @property + def bucket(self) -> Bucket: + """Return this capture's context bounds.""" + return self.min_kv_tokens, self.max_kv_tokens + + @classmethod + def from_mapping(cls, raw: Mapping[str, object]) -> CompactMaskReuseCapture: + """Parse one strict compact-capture object.""" + _exact_fields(raw, _CAPTURE_FIELDS, "compact capture") + if raw["compact_capture_schema_version"] != 1: + raise MaskReuseCalibrationError("compact_capture_schema_version must be 1") + invocation = raw["invocation"] + if not isinstance(invocation, Mapping): + raise MaskReuseCalibrationError("compact capture invocation must be an object") + _exact_fields(invocation, _INVOCATION_FIELDS, "compact capture invocation") + if invocation["capture_schema_version"] != 2: + raise MaskReuseCalibrationError("capture_schema_version must be 2") + split = invocation["split"] + if split not in _SPLITS: + raise MaskReuseCalibrationError("capture split must be calibration or heldout") + partition = invocation["partition"] + if partition not in _PARTITIONS: + raise MaskReuseCalibrationError("capture partition must be development or outer_test") + expected_split = "calibration" if partition == "development" else "heldout" + if split != expected_split: + raise MaskReuseCalibrationError("capture split and partition disagree") + raw_fold = invocation["inner_fold"] + if partition == "development": + inner_fold = _integer(raw_fold, "inner_fold") + elif raw_fold is not None: + raise MaskReuseCalibrationError("outer_test capture must have null inner_fold") + else: + inner_fold = None + minimum = _integer(invocation["min_kv_tokens"], "min_kv_tokens", minimum=1) + maximum = invocation["max_kv_tokens"] + if maximum is not None: + maximum = _integer(maximum, "max_kv_tokens", minimum=minimum) + sample_length = _integer(invocation["sample_length"], "sample_length", minimum=1) + if sample_length < minimum or (maximum is not None and sample_length > maximum): + raise MaskReuseCalibrationError("sample_length lies outside its context bucket") + target = _float_hex(invocation["target_sparsity_hex"], "target_sparsity_hex") + threshold_log2 = _float_hex(invocation["threshold_log2_hex"], "threshold_log2_hex") + threshold_lambda = _float_hex(invocation["threshold_lambda_hex"], "threshold_lambda_hex") + if not 0.0 < target < 1.0: + raise MaskReuseCalibrationError("target_sparsity must be in (0, 1)") + if threshold_log2 >= 0.0 or not 0.0 < threshold_lambda < 1.0: + raise MaskReuseCalibrationError("threshold must be in (0, 1)") + if float(getattr(math, "exp2")(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 = float(getattr(math, "exp2")(expected_log2)) + if capture.threshold_log2.hex() != expected_log2.hex(): + raise MaskReuseCalibrationError("compact capture threshold_log2 differs from vanilla fit") + if capture.threshold_lambda.hex() != expected_lambda.hex(): + raise MaskReuseCalibrationError("compact capture threshold_lambda differs from vanilla fit") + + +def _capture_order_key(capture: CompactMaskReuseCapture) -> tuple[object, ...]: + return ( + capture.min_kv_tokens, + math.inf if capture.max_kv_tokens is None else capture.max_kv_tokens, + _SPLITS.index(capture.split), + capture.prompt_id, + capture.target_sparsity, + ) + + +def _validate_monotonic_pair( + previous: CompactMaskReuseCapture, current: CompactMaskReuseCapture +) -> None: + identity = (current.bucket, current.split, current.prompt_id) + if identity != (previous.bucket, previous.split, previous.prompt_id): + return + if current.target_sparsity <= previous.target_sparsity: + raise MaskReuseCalibrationError("target sparsities must be strictly increasing per prompt") + for layer, current_stats in current.anchor_stats_by_layer.items(): + previous_stats = previous.anchor_stats_by_layer[layer] + if any( + current_value > previous_value + for current_value, previous_value in zip( + current_stats.retained_tiles, + previous_stats.retained_tiles, + strict=True, + ) + ): + raise MaskReuseCalibrationError( + f"anchor {layer} retained counts increase with target sparsity" + ) + if any( + current_value + _MONOTONIC_ATOL < previous_value + for current_value, previous_value in zip( + current_stats.dropped_mass, + previous_stats.dropped_mass, + strict=True, + ) + ): + raise MaskReuseCalibrationError( + f"anchor {layer} dropped mass decreases with target sparsity" + ) + for layer, current_consumer_stats in current.consumer_layers.items(): + previous_consumer_stats = previous.consumer_layers[layer] + for consumer_head, (current_row, previous_row) in enumerate( + zip( + current_consumer_stats.dropped_mass, + previous_consumer_stats.dropped_mass, + strict=True, + ) + ): + if any( + current_value + _MONOTONIC_ATOL < previous_value + for current_value, previous_value in zip(current_row, previous_row, strict=True) + ): + raise MaskReuseCalibrationError( + f"consumer {layer} head {consumer_head} dropped mass decreases " + "with target sparsity" + ) + + +def _validate_dataset( + source: CompactMaskReuseCaptureSource, + *, + fit: Mapping[str, object], + anchors: tuple[int, ...], + nearest: Mapping[int, int], +) -> _Dataset: + models: set[str] = set() + checkpoint_identities: set[str] = set() + head_counts: set[int] = set() + split_counts: dict[str, int] = defaultdict(int) + split_buckets: dict[str, set[Bucket]] = defaultdict(set) + prompts: dict[tuple[Bucket, str], set[str]] = defaultdict(set) + menus: dict[tuple[Bucket, str, str], set[float]] = defaultdict(set) + capture_sources: dict[tuple[Bucket, str, str], tuple[object, ...]] = {} + fingerprint_owner: dict[tuple[str, str], tuple[Bucket, str]] = {} + seen: set[tuple[Bucket, str, str, float]] = set() + geometry_by_prompt: dict[tuple[Bucket, str, str], tuple[object, ...]] = {} + prompt_source_rows: dict[tuple[Bucket, str, str], Mapping[str, object]] = {} + split_prompt_ids: dict[str, set[str]] = defaultdict(set) + prompt_buckets: dict[tuple[str, str], Bucket] = {} + split_fingerprints: dict[str, set[str]] = defaultdict(set) + group_assignments: dict[str, tuple[str, int | None]] = {} + fingerprint_groups: dict[str, str] = {} + expected_anchor_set = set(anchors) + consumer_to_anchor = {layer: anchor for layer, anchor in nearest.items() if layer != anchor} + expected_consumer_set = set(consumer_to_anchor) + capture_count = 0 + previous_order: tuple[object, ...] | None = None + previous_capture: CompactMaskReuseCapture | None = None + for capture in source: + capture_count += 1 + order = _capture_order_key(capture) + if previous_order is not None and order <= previous_order: + raise MaskReuseCalibrationError("compact capture records are not in canonical order") + if previous_capture is not None: + _validate_monotonic_pair(previous_capture, capture) + previous_order = order + previous_capture = capture + _validate_threshold(capture, fit) + if set(capture.anchor_stats_by_layer) != expected_anchor_set: + raise MaskReuseCalibrationError("compact capture does not cover every topology anchor") + if set(capture.consumer_layers) != expected_consumer_set: + raise MaskReuseCalibrationError("compact capture does not cover every reuse layer") + for layer, stats in capture.consumer_layers.items(): + if stats.anchor_layer != consumer_to_anchor[layer]: + raise MaskReuseCalibrationError( + f"consumer layer {layer} does not match the explicit topology" + ) + models.add(capture.model) + checkpoint_identities.add(capture.checkpoint_manifest_sha256) + head_counts.add(capture.global_num_heads) + split_counts[capture.split] += 1 + split_buckets[capture.split].add(capture.bucket) + prompts[(capture.bucket, capture.split)].add(capture.prompt_id) + menus[(capture.bucket, capture.split, capture.prompt_id)].add(capture.target_sparsity) + split_prompt_ids[capture.split].add(capture.prompt_id) + prompt_identity = (capture.split, capture.prompt_id) + if prompt_buckets.setdefault(prompt_identity, capture.bucket) != capture.bucket: + raise MaskReuseCalibrationError("a prompt ID is assigned to multiple context buckets") + split_fingerprints[capture.split].add(capture.source_capture_sha256) + assignment = (capture.partition, capture.inner_fold) + previous_assignment = group_assignments.setdefault(capture.source_group_sha256, assignment) + if previous_assignment != assignment: + raise MaskReuseCalibrationError( + "one source group is assigned to multiple partitions or inner folds" + ) + previous_group = fingerprint_groups.setdefault( + capture.source_capture_sha256, capture.source_group_sha256 + ) + if previous_group != capture.source_group_sha256: + raise MaskReuseCalibrationError( + "one rendered source capture is assigned to multiple source groups" + ) + key = (capture.bucket, capture.split, capture.prompt_id, capture.target_sparsity) + if key in seen: + raise MaskReuseCalibrationError("compact captures contain a duplicate prompt target") + seen.add(key) + prompt_key = (capture.bucket, capture.split, capture.prompt_id) + source_identity = ( + capture.source, + capture.source_group_sha256, + capture.partition, + capture.inner_fold, + capture.source_capture_sha256, + capture.sample_length, + tuple(sorted(capture.geometry.items())), + capture.eligible_tiles, + ) + if capture_sources.setdefault(prompt_key, source_identity) != source_identity: + raise MaskReuseCalibrationError("prompt metadata differs across target sparsities") + owner_key = (capture.split, capture.source_capture_sha256) + owner = (capture.bucket, capture.prompt_id) + if fingerprint_owner.setdefault(owner_key, owner) != owner: + raise MaskReuseCalibrationError("one source fingerprint names multiple prompt captures") + geometry_identity = ( + capture.source_capture_sha256, + capture.geometry["q_tokens"], + capture.geometry["kv_tokens"], + capture.geometry["q_start_tokens"], + ) + if geometry_by_prompt.setdefault(prompt_key, geometry_identity) != geometry_identity: + raise MaskReuseCalibrationError("prompt geometry differs across target sparsities") + prompt_source_rows[prompt_key] = { + "min_kv_tokens": capture.min_kv_tokens, + "max_kv_tokens": capture.max_kv_tokens, + "split": capture.split, + "prompt_id": capture.prompt_id, + "source": capture.source, + "source_group_sha256": capture.source_group_sha256, + "partition": capture.partition, + "inner_fold": capture.inner_fold, + "source_capture_sha256": capture.source_capture_sha256, + } + if capture_count == 0 or any(not split_counts[split] for split in _SPLITS): + raise MaskReuseCalibrationError("compact captures require calibration and heldout splits") + if len(models) != 1 or len(head_counts) != 1 or len(checkpoint_identities) != 1: + raise MaskReuseCalibrationError( + "compact captures must use one model, checkpoint, and head count" + ) + if split_prompt_ids["calibration"] & split_prompt_ids["heldout"]: + raise MaskReuseCalibrationError("prompt IDs overlap calibration and heldout splits") + if split_fingerprints["calibration"] & split_fingerprints["heldout"]: + raise MaskReuseCalibrationError("source captures overlap calibration and heldout splits") + if split_buckets["calibration"] != split_buckets["heldout"]: + raise MaskReuseCalibrationError("heldout context buckets must match calibration buckets") + buckets = tuple(sorted(split_buckets["calibration"], key=_bucket_key)) + previous_max: int | None = 0 + for minimum, maximum in buckets: + if previous_max is None or minimum <= previous_max: + raise MaskReuseCalibrationError("context buckets must be ordered and non-overlapping") + previous_max = maximum + canonical_menus: dict[Bucket, tuple[float, ...]] = {} + for bucket in buckets: + calibration_prompts = sorted(prompts[(bucket, "calibration")]) + if not calibration_prompts: + raise MaskReuseCalibrationError(f"bucket {bucket} has no calibration prompts") + menu = menus[(bucket, "calibration", calibration_prompts[0])] + if not menu: + raise MaskReuseCalibrationError(f"bucket {bucket} has no target menu") + for split in _SPLITS: + for prompt in prompts[(bucket, split)]: + if menus[(bucket, split, prompt)] != menu: + raise MaskReuseCalibrationError("target-sparsity menu differs across prompts") + canonical_menus[bucket] = tuple(sorted(menu)) + + geometry_rows = [] + for (bucket, split, prompt), identity in sorted( + geometry_by_prompt.items(), key=lambda item: str(item[0]) + ): + geometry_rows.append( + { + "split": split, + "prompt_id": prompt, + "source_capture_sha256": identity[0], + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "q_tokens": identity[1], + "kv_tokens": identity[2], + "q_start_tokens": identity[3], + } + ) + global_num_heads = next(iter(head_counts)) + consumer_layers = tuple(sorted(expected_consumer_set)) + targets = tuple((layer, head) for layer in consumer_layers for head in range(global_num_heads)) + return _Dataset( + model=next(iter(models)), + checkpoint_manifest_sha256=next(iter(checkpoint_identities)), + global_num_heads=global_num_heads, + buckets=buckets, + anchors=anchors, + nearest=nearest, + consumer_layers=consumer_layers, + targets=targets, + prompts={key: tuple(sorted(value)) for key, value in prompts.items()}, + menus=canonical_menus, + deployment_geometry={ + "contract": dict(_DEPLOYMENT_GEOMETRY_CONTRACT), + "observations": geometry_rows, + }, + prompt_sources=tuple( + prompt_source_rows[key] for key in sorted(prompt_source_rows, key=str) + ), + calibration_prompt_ids=tuple(sorted(split_prompt_ids["calibration"])), + heldout_prompt_ids=tuple(sorted(split_prompt_ids["heldout"])), + capture_count=capture_count, + ) + + +@dataclass(slots=True) +class _SelectionAccumulator: + max_mass: Mapping[float, np.ndarray] + retained_by_anchor: Mapping[float, np.ndarray] + eligible_sum: dict[float, int] + anchor_evaluation: Mapping[float, _AnchorEvaluation] + + +def _selection_pass( + source: CompactMaskReuseCaptureSource, + dataset: _Dataset, + *, + max_anchor_dropped_mass: float, + max_reuse_selection_dropped_mass: float, +) -> dict[Bucket, _Selection]: + consumer_index = {layer: index for index, layer in enumerate(dataset.consumer_layers)} + anchor_index = {layer: index for index, layer in enumerate(dataset.anchors)} + accumulators: dict[Bucket, _SelectionAccumulator] = {} + for bucket in dataset.buckets: + menus = dataset.menus[bucket] + accumulators[bucket] = _SelectionAccumulator( + max_mass={ + target: np.full( + ( + len(dataset.consumer_layers), + dataset.global_num_heads, + dataset.global_num_heads, + ), + -np.inf, + dtype=np.float64, + ) + for target in menus + }, + retained_by_anchor={ + target: np.zeros((len(dataset.anchors), dataset.global_num_heads), dtype=np.int64) + for target in menus + }, + eligible_sum=dict.fromkeys(menus, 0), + anchor_evaluation={target: _AnchorEvaluation() for target in menus}, + ) + for capture in source: + if capture.split != "calibration": + continue + accumulator = accumulators[capture.bucket] + target = capture.target_sparsity + accumulator.eligible_sum[target] += capture.eligible_tiles + for layer, stats in capture.anchor_stats_by_layer.items(): + accumulator.retained_by_anchor[target][anchor_index[layer]] += np.asarray( + stats.retained_tiles, dtype=np.int64 + ) + evaluation = accumulator.anchor_evaluation[target] + dropped = [ + value + for stats in capture.anchor_stats_by_layer.values() + for value in stats.dropped_mass + ] + prompt_mean = sum(dropped) / len(dropped) + evaluation.eligible_tiles += ( + capture.eligible_tiles * len(dataset.anchors) * dataset.global_num_heads + ) + evaluation.retained_tiles += sum( + sum(stats.retained_tiles) for stats in capture.anchor_stats_by_layer.values() + ) + evaluation.prompt_count += 1 + evaluation.prompt_mean_sum += prompt_mean + evaluation.worst_prompt_mean = max(evaluation.worst_prompt_mean, prompt_mean) + evaluation.violations += int(prompt_mean > max_anchor_dropped_mass) + for layer, stats in capture.consumer_layers.items(): + matrix = np.asarray(stats.dropped_mass, dtype=np.float64) + np.maximum( + accumulator.max_mass[target][consumer_index[layer]], + matrix, + out=accumulator.max_mass[target][consumer_index[layer]], + ) + + selections: dict[Bucket, _Selection] = {} + for bucket in dataset.buckets: + accumulator = accumulators[bucket] + frontier: list[Mapping[str, object]] = [] + candidates: list[tuple[tuple[object, ...], _Selection]] = [] + for target in dataset.menus[bucket]: + anchor_evaluation = accumulator.anchor_evaluation[target] + choices: dict[ConsumerHead, _Choice] = {} + total_retained = 0 + fallback_count = 0 + for layer in dataset.consumer_layers: + layer_index = consumer_index[layer] + anchor = dataset.nearest[layer] + retained = accumulator.retained_by_anchor[target][anchor_index[anchor]] + for head in range(dataset.global_num_heads): + feasible = np.flatnonzero( + accumulator.max_mass[target][layer_index, head] + <= max_reuse_selection_dropped_mass + ) + if feasible.size: + donor = min((int(retained[item]), int(item)) for item in feasible)[1] + choice = _Choice(donor, False, int(retained[donor])) + else: + fallback_count += 1 + choice = _Choice(0, True, accumulator.eligible_sum[target]) + choices[(layer, head)] = choice + total_retained += choice.retained_tiles + signature = tuple(choices[item].donor_head for item in dataset.targets) + combined_tile_cost = 2 * total_retained + anchor_evaluation.retained_tiles + frontier.append( + { + "target_sparsity": target, + "anchor_safe": anchor_evaluation.violations == 0, + "retained_reuse_tiles": total_retained, + "retained_anchor_tiles": anchor_evaluation.retained_tiles, + "combined_tile_cost": combined_tile_cost, + "fallback_head_count": fallback_count, + "anchor_calibration": anchor_evaluation.to_mapping(), + } + ) + if anchor_evaluation.violations == 0: + rank = ( + combined_tile_cost, + fallback_count, + target, + signature, + ) + candidates.append((rank, _Selection(target, choices, ()))) + if candidates: + _, selected = min(candidates, key=lambda item: item[0]) + selections[bucket] = _Selection( + selected.target_sparsity, + selected.choices, + tuple(frontier), + ) + else: + target = dataset.menus[bucket][0] + dense_choices = { + item: _Choice(0, True, accumulators[bucket].eligible_sum[target]) + for item in dataset.targets + } + selections[bucket] = _Selection( + None, + dense_choices, + tuple(frontier), + "no_target_sparsity_satisfied_anchor_calibration_constraint", + ) + return selections + + +def _evaluation_pass( + source: CompactMaskReuseCaptureSource, + dataset: _Dataset, + selections: Mapping[Bucket, _Selection], + *, + max_anchor_dropped_mass: float, + max_reuse_dropped_mass: float, +) -> tuple[ + dict[tuple[Bucket, str], _ReuseEvaluation], + dict[tuple[Bucket, str], _AnchorEvaluation], +]: + reuse = {(bucket, split): _ReuseEvaluation() for bucket in dataset.buckets for split in _SPLITS} + anchor = { + (bucket, split): _AnchorEvaluation() for bucket in dataset.buckets for split in _SPLITS + } + for capture in source: + selection = selections[capture.bucket] + evaluation_target = ( + dataset.menus[capture.bucket][0] + if selection.target_sparsity is None + else selection.target_sparsity + ) + if capture.target_sparsity != evaluation_target: + continue + reuse_result = reuse[(capture.bucket, capture.split)] + for layer, stats in capture.consumer_layers.items(): + anchor_stats = capture.anchor_stats_by_layer[stats.anchor_layer] + for head in range(dataset.global_num_heads): + choice = selection.choices[(layer, head)] + reuse_result.eligible_tiles += capture.eligible_tiles + if choice.fallback: + reuse_result.retained_tiles += capture.eligible_tiles + continue + donor = choice.donor_head + dropped = float(stats.dropped_mass[head][donor]) + reuse_result.retained_tiles += anchor_stats.retained_tiles[donor] + reuse_result.sparse_observations += 1 + reuse_result.dropped_mass_sum += dropped + reuse_result.worst_dropped_mass = max(reuse_result.worst_dropped_mass, dropped) + reuse_result.violations += int(dropped > max_reuse_dropped_mass) + + anchor_result = anchor[(capture.bucket, capture.split)] + dropped_values: list[float] = [] + retained_tiles = 0 + eligible_tiles = 0 + for anchor_stats in capture.anchor_stats_by_layer.values(): + eligible_tiles += capture.eligible_tiles * dataset.global_num_heads + if selection.target_sparsity is None: + retained_tiles += capture.eligible_tiles * dataset.global_num_heads + dropped_values.extend([0.0] * dataset.global_num_heads) + else: + retained_tiles += sum(anchor_stats.retained_tiles) + dropped_values.extend(anchor_stats.dropped_mass) + prompt_mean = sum(dropped_values) / len(dropped_values) + anchor_result.eligible_tiles += eligible_tiles + anchor_result.retained_tiles += retained_tiles + anchor_result.prompt_count += 1 + anchor_result.prompt_mean_sum += prompt_mean + anchor_result.worst_prompt_mean = max(anchor_result.worst_prompt_mean, prompt_mean) + anchor_result.violations += int(prompt_mean > max_anchor_dropped_mass) + for bucket in dataset.buckets: + for split in _SPLITS: + expected = len(dataset.prompts[(bucket, split)]) + if anchor[(bucket, split)].prompt_count != expected: + raise MaskReuseCalibrationError( + f"evaluation pass did not cover every {split} prompt in bucket {bucket}" + ) + return reuse, anchor + + +def calibrate_compact_mask_reuse_policy( + captures: CompactMaskReuseCaptureSource | str | Path, + *, + vanilla_calibration: Mapping[str, object], + topology: Mapping[str, object], + checkpoint_manifest: VerifiedCheckpointManifest, + evidence: Mapping[str, object], + max_anchor_dropped_mass: float, + max_reuse_dropped_mass: float, + max_reuse_selection_dropped_mass: float | None = None, + source_provenance: Mapping[str, object] | None = None, +) -> dict[str, object]: + """Select and evaluate a fail-closed schema-v3 candidate from compact captures.""" + if not isinstance(checkpoint_manifest, VerifiedCheckpointManifest): + raise MaskReuseCalibrationError( + "checkpoint_manifest must be returned by verify_checkpoint_manifest" + ) + source = ( + captures + if isinstance(captures, CompactMaskReuseCaptureSource) + else load_compact_mask_reuse_captures(captures) + ) + threshold_scale_factor = canonical_prefill_threshold_scale_factor(vanilla_calibration) + params = threshold_scale_factor["prefill"] + assert isinstance(params, Mapping) + if not {"min_observed_sparsity", "max_observed_sparsity"} <= params.keys(): + raise MaskReuseCalibrationError("compact schema-v3 export requires vanilla fit bounds") + anchors, nearest = _normalize_topology(topology) + dataset = _validate_dataset( + source, + fit=threshold_scale_factor, + anchors=anchors, + nearest=nearest, + ) + checkpoint_identity = checkpoint_manifest.sha256 + if dataset.checkpoint_manifest_sha256 != checkpoint_identity: + raise MaskReuseCalibrationError( + "compact captures do not match the verified checkpoint manifest" + ) + if dataset.model != checkpoint_manifest.model: + raise MaskReuseCalibrationError( + "compact capture model does not match the verified checkpoint manifest" + ) + canonical_evidence = _canonical_evidence(evidence) + input_sha256 = source.sha256() + if canonical_evidence["reuse_bundle_sha256"] != input_sha256: + raise MaskReuseCalibrationError( + "evidence.reuse_bundle_sha256 does not match the compact capture file" + ) + anchor_bound = _number( + max_anchor_dropped_mass, "max_anchor_dropped_mass", minimum=0.0, maximum=1.0 + ) + reuse_bound = _number( + max_reuse_dropped_mass, "max_reuse_dropped_mass", minimum=0.0, maximum=1.0 + ) + selection_bound = ( + reuse_bound + if max_reuse_selection_dropped_mass is None + else _number( + max_reuse_selection_dropped_mass, + "max_reuse_selection_dropped_mass", + minimum=0.0, + maximum=1.0, + ) + ) + if selection_bound > reuse_bound: + raise MaskReuseCalibrationError( + "max_reuse_selection_dropped_mass must not exceed max_reuse_dropped_mass" + ) + selections = _selection_pass( + source, + dataset, + max_anchor_dropped_mass=anchor_bound, + max_reuse_selection_dropped_mass=selection_bound, + ) + reuse_evaluations, anchor_evaluations = _evaluation_pass( + source, + dataset, + selections, + max_anchor_dropped_mass=anchor_bound, + max_reuse_dropped_mass=reuse_bound, + ) + if source.sha256() != input_sha256: + raise MaskReuseCalibrationError( + "compact capture file changed during calibration; discard this result" + ) + + context_policies = [] + bucket_reports = [] + target_menus = [] + overall_reuse_calibration = _ReuseEvaluation() + overall_reuse_heldout = _ReuseEvaluation() + overall_anchor_calibration = _AnchorEvaluation() + overall_anchor_heldout = _AnchorEvaluation() + total_fallback = 0 + for bucket in dataset.buckets: + selection = selections[bucket] + if selection.target_sparsity is not None and bucket[1] is None: + raise MaskReuseCalibrationError( + "a deployment-qualified sparse context bucket requires a finite maximum" + ) + reuse_calibration = reuse_evaluations[(bucket, "calibration")] + reuse_heldout = reuse_evaluations[(bucket, "heldout")] + anchor_calibration = anchor_evaluations[(bucket, "calibration")] + anchor_heldout = anchor_evaluations[(bucket, "heldout")] + overall_reuse_calibration.add(reuse_calibration) + overall_reuse_heldout.add(reuse_heldout) + overall_anchor_calibration.add(anchor_calibration) + overall_anchor_heldout.add(anchor_heldout) + policy: dict[str, object] + if selection.target_sparsity is None: + policy = {"min_kv_tokens": bucket[0], "max_kv_tokens": bucket[1], "exact": True} + headmaps: dict[str, list[int]] = {} + fallback_heads: dict[str, list[int]] = {} + else: + headmaps = { + str(layer): [ + selection.choices[(layer, head)].donor_head + for head in range(dataset.global_num_heads) + ] + for layer in dataset.consumer_layers + } + fallback_heads = { + str(layer): [ + head + for head in range(dataset.global_num_heads) + if selection.choices[(layer, head)].fallback + ] + for layer in dataset.consumer_layers + } + policy = { + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "target_sparsity": selection.target_sparsity, + "headmaps": headmaps, + "fallback_heads": fallback_heads, + } + context_policies.append(policy) + fallback_count = sum(len(heads) for heads in fallback_heads.values()) + total_fallback += fallback_count + target_menus.append( + { + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "target_sparsities": [row["target_sparsity"] for row in selection.frontier], + } + ) + bucket_reports.append( + { + "min_kv_tokens": bucket[0], + "max_kv_tokens": bucket[1], + "selected_target_sparsity": selection.target_sparsity, + "selection_status": ( + "sparse" if selection.target_sparsity is not None else selection.exact_reason + ), + "target_sparsity_frontier": list(selection.frontier), + "fallback_head_count": fallback_count, + "reuse_calibration": reuse_calibration.to_mapping(), + "reuse_heldout": reuse_heldout.to_mapping(), + "anchor_calibration": anchor_calibration.to_mapping( + exact=selection.target_sparsity is None + ), + "anchor_heldout": anchor_heldout.to_mapping( + exact=selection.target_sparsity is None + ), + } + ) + + candidate_cell_count = ( + dataset.capture_count + * len(dataset.consumer_layers) + * dataset.global_num_heads + * dataset.global_num_heads + ) + constraints = { + "anchor": { + "metric": "worst_prompt_mean_anchor_dropped_mass", + "comparison": "<=", + "maximum": anchor_bound, + }, + "reuse": { + "metric": "per_prompt_candidate_reuse_dropped_mass", + "comparison": "<=", + "selection_maximum": selection_bound, + "evaluation_maximum": reuse_bound, + }, + } + provenance: dict[str, object] = { + "calibrator": "modelopt.mask_reuse.compact_streaming", + "compact_capture_schema_version": 1, + "input_capture_count": dataset.capture_count, + "candidate_cell_count": candidate_cell_count, + "canonical_input_sha256": input_sha256, + "calibration_prompt_ids": list(dataset.calibration_prompt_ids), + "heldout_prompt_ids": list(dataset.heldout_prompt_ids), + "prompt_sources": list(dataset.prompt_sources), + "development_source_group_sha256": sorted( + { + str(row["source_group_sha256"]) + for row in dataset.prompt_sources + if row["partition"] == "development" + } + ), + "outer_test_source_group_sha256": sorted( + { + str(row["source_group_sha256"]) + for row in dataset.prompt_sources + if row["partition"] == "outer_test" + } + ), + "selection_split": "calibration", + "evaluation_split": "heldout", + "streaming_passes": ["validation", "calibration_selection", "frozen_evaluation"], + "threshold_semantics": "a * exp(b * target_sparsity) / sample_length", + "threshold_implementation": { + "threshold_log2": "log2(a) + b * target_sparsity * log2(e) - log2(sample_length)", + "threshold_lambda": "exp2(threshold_log2)", + "validation": "exact IEEE-754 binary64 hex equality", + }, + "checkpoint_manifest_sha256": checkpoint_identity, + "target_sparsity_menus": target_menus, + "constraints": constraints, + "tie_breaks": [ + "reject target sparsities exceeding the calibration anchor bound", + "minimum equal-BMM combined tile cost 2*A_R + A_A", + "fewest exact-fallback consumer heads", + "smallest target sparsity", + "lexicographically smallest donor-head map", + ], + "selection_cost": { + "formula": "2 * retained_reuse_tiles + retained_anchor_tiles", + "bmm1_weight": 1.0, + "bmm2_weight": 1.0, + }, + } + if source_provenance is not None: + provenance["source"] = dict(source_provenance) + denominator = len(dataset.targets) * len(dataset.buckets) + report = { + "model": dataset.model, + "checkpoint_manifest_sha256": checkpoint_identity, + "constraints": constraints, + "selection_unit": "context_bucket", + "by_bucket": bucket_reports, + "overall": { + "consumer_head_bucket_count": denominator, + "fallback_head_bucket_count": total_fallback, + "fallback_fraction": total_fallback / denominator, + "reuse_calibration": overall_reuse_calibration.to_mapping(), + "reuse_heldout": overall_reuse_heldout.to_mapping(), + "anchor_calibration": overall_anchor_calibration.to_mapping(), + "anchor_heldout": overall_anchor_heldout.to_mapping(), + }, + "promotion": { + "status": "candidate_only", + "eligible": False, + "reasons": [ + "grouped inner-fold modal and safety stability not evaluated", + "preregistered outer gate with at least 99 independent groups per family cell not evaluated", + "deployment rectangular-geometry promotion gate not evaluated", + ], + }, + } + return { + "version": 3, + "promotion_status": "candidate_only", + "phase": "prefill", + "decode": {"mode": "dense"}, + "calibration_protocol": _CALIBRATION_PROTOCOL, + "producer": {"name": "modelopt", "version": modelopt.__version__}, + "evidence": canonical_evidence, + "threshold_scale_factor": threshold_scale_factor, + "model": dataset.model, + "checkpoint_manifest_sha256": checkpoint_identity, + "global_num_heads": dataset.global_num_heads, + "anchors": list(dataset.anchors), + "nearest": {str(layer): anchor for layer, anchor in dataset.nearest.items()}, + "deployment_geometry_validated": False, + "deployment_geometry": dataset.deployment_geometry, + "context_policies": context_policies, + "provenance": provenance, + "calibration_report": report, + } diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py b/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py new file mode 100644 index 00000000000..1a1160af22c --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py @@ -0,0 +1,809 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Host-side contract for ModelOpt-controlled mask-reuse capture. + +The vLLM backend measures rank-local sufficient statistics. This module owns +the trusted inputs, validates every echoed invocation, merges tensor-parallel +consumer-head shards, and emits the normalized rows consumed by +``calibrate_mask_reuse_policy``. It never fabricates missing GPU statistics. +""" + +from __future__ import annotations + +import json +import math +import struct +import unicodedata +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import cast + +from modelopt.torch.sparsity.attention_sparsity.calibration.mask_reuse import ( + canonical_prefill_threshold_scale_factor, +) + +__all__ = [ + "CAPTURE_SCHEMA_VERSION", + "MAX_QUERY_CHUNK_TOKENS", + "CaptureContractError", + "MergedCapture", + "PromptSpec", + "build_capture_invocation", + "canonical_json_sha256", + "load_prompt_specs", + "load_vanilla_prefill_fit", + "merge_rank_captures", + "parse_prompt_specs_jsonl", + "parse_vanilla_prefill_fit", + "source_capture_sha256", + "validate_begin_acks", + "validate_capture_statuses", +] + +CAPTURE_SCHEMA_VERSION = 2 +MAX_QUERY_CHUNK_TOKENS = 8192 +_QUERY_START_ALIGNMENT = 128 +_SPLITS = frozenset({"calibration", "heldout"}) +_PARTITIONS = frozenset({"development", "outer_test"}) +_PROMPT_FIELDS = frozenset( + { + "split", + "partition", + "inner_fold", + "prompt_id", + "source", + "source_group_sha256", + "prompt", + "min_kv_tokens", + "max_kv_tokens", + } +) +_INVOCATION_FIELDS = frozenset( + { + "capture_schema_version", + "model", + "checkpoint_manifest_sha256", + "split", + "partition", + "inner_fold", + "prompt_id", + "source", + "source_group_sha256", + "source_capture_sha256", + "min_kv_tokens", + "max_kv_tokens", + "target_sparsity_hex", + "sample_length", + "threshold_log2_hex", + "threshold_lambda_hex", + "expected_geometry", + } +) +_GEOMETRY_FIELDS = frozenset({"q_tokens", "kv_tokens", "q_start_tokens"}) +_STATUS_FIELDS = frozenset({"capture_schema_version", "available", "rank", "world_size", "reason"}) +_ACK_FIELDS = frozenset( + { + "capture_schema_version", + "armed", + "rank", + "world_size", + "invocation_sha256", + } +) +_RANK_CAPTURE_FIELDS = frozenset( + { + "capture_schema_version", + "rank", + "world_size", + "invocation", + "invocation_sha256", + "geometry", + "global_num_heads", + "eligible_tiles", + "anchor_stats_by_layer", + "consumer_layers", + } +) +_ANCHOR_STATS_FIELDS = frozenset({"retained_tiles", "dropped_mass"}) +_CONSUMER_STATS_FIELDS = frozenset({"anchor_layer", "consumer_head_start", "dropped_mass"}) + + +class CaptureContractError(ValueError): + """Raised when capture inputs or backend evidence violate the contract.""" + + +def _exact_fields(value: Mapping[str, object], expected: frozenset[str], label: str) -> None: + missing = expected - value.keys() + extra = value.keys() - expected + if missing or extra: + raise CaptureContractError( + f"{label} fields do not match the schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + + +def _integer(value: object, label: str, *, minimum: int = 0) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise CaptureContractError(f"{label} must be an integer >= {minimum}") + return value + + +def _finite_number( + value: object, + label: str, + *, + minimum: float | None = None, + maximum: float | None = None, +) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise CaptureContractError(f"{label} must be a finite number") + result = float(value) + if not math.isfinite(result): + raise CaptureContractError(f"{label} must be a finite number") + if minimum is not None and result < minimum: + raise CaptureContractError(f"{label} must be >= {minimum}") + if maximum is not None and result > maximum: + raise CaptureContractError(f"{label} must be <= {maximum}") + return result + + +def _text(value: object, label: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise CaptureContractError(f"{label} must be a non-empty string") + result = value.strip() + if unicodedata.normalize("NFC", result) != result or any( + ord(character) < 32 for character in result + ): + raise CaptureContractError(f"{label} must be NFC text without control characters") + return result + + +def _sha256(value: object, label: str) -> str: + if ( + not isinstance(value, str) + or len(value) != 64 + or any(character not in "0123456789abcdef" for character in value) + ): + raise CaptureContractError(f"{label} must be a lowercase SHA256") + return value + + +def _canonical_float_hex(value: object, label: str) -> float: + if not isinstance(value, str): + raise CaptureContractError(f"{label} must be a canonical float.hex string") + try: + parsed = float.fromhex(value) + except ValueError as error: + raise CaptureContractError(f"{label} must be a canonical float.hex string") from error + if not math.isfinite(parsed) or parsed.hex() != value: + raise CaptureContractError(f"{label} must be a canonical finite float.hex string") + return parsed + + +def _reject_duplicate_json_keys(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise CaptureContractError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + +def canonical_json_sha256(value: object) -> str: + """Hash the canonical JSON encoding used by capture RPC acknowledgements.""" + encoded = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True).encode() + return sha256(encoded).hexdigest() + + +@dataclass(frozen=True, slots=True) +class PromptSpec: + """One preregistered prompt/context capture unit.""" + + split: str + partition: str + inner_fold: int | None + prompt_id: str + source: str + source_group_sha256: str + prompt: str + min_kv_tokens: int + max_kv_tokens: int | None + + @classmethod + def from_mapping(cls, raw: Mapping[str, object]) -> PromptSpec: + """Parse one strict prompt-plan object.""" + _exact_fields(raw, _PROMPT_FIELDS, "prompt") + split = raw["split"] + if split not in _SPLITS: + raise CaptureContractError(f"prompt.split must be one of {sorted(_SPLITS)}") + partition = raw["partition"] + if partition not in _PARTITIONS: + raise CaptureContractError(f"prompt.partition must be one of {sorted(_PARTITIONS)}") + expected_split = "calibration" if partition == "development" else "heldout" + if split != expected_split: + raise CaptureContractError("prompt.split and prompt.partition disagree") + raw_fold = raw["inner_fold"] + if partition == "development": + inner_fold = _integer(raw_fold, "prompt.inner_fold") + elif raw_fold is not None: + raise CaptureContractError("outer_test prompts must have null inner_fold") + else: + inner_fold = None + prompt = raw["prompt"] + if not isinstance(prompt, str) or not prompt: + raise CaptureContractError("prompt.prompt must be a non-empty string") + minimum = _integer(raw["min_kv_tokens"], "prompt.min_kv_tokens", minimum=1) + maximum = raw["max_kv_tokens"] + if maximum is not None: + maximum = _integer(maximum, "prompt.max_kv_tokens", minimum=minimum) + return cls( + split=split, + partition=partition, + inner_fold=inner_fold, + prompt_id=_text(raw["prompt_id"], "prompt.prompt_id"), + source=_text(raw["source"], "prompt.source"), + source_group_sha256=_sha256(raw["source_group_sha256"], "prompt.source_group_sha256"), + prompt=prompt, + min_kv_tokens=minimum, + max_kv_tokens=maximum, + ) + + @property + def bucket(self) -> tuple[int, int | None]: + """Return the calibrated context-bucket bounds.""" + return self.min_kv_tokens, self.max_kv_tokens + + +def parse_prompt_specs_jsonl(payload: bytes) -> list[PromptSpec]: + """Parse exact strict-JSONL prompt bytes and reject split leakage.""" + if not isinstance(payload, bytes): + raise CaptureContractError("prompt file payload must be bytes") + try: + lines = payload.decode("utf-8").splitlines() + except UnicodeDecodeError as error: + raise CaptureContractError("prompt file must be UTF-8") from error + prompts: list[PromptSpec] = [] + for line_number, line in enumerate(lines, start=1): + if not line.strip(): + continue + try: + raw = json.loads(line, object_pairs_hook=_reject_duplicate_json_keys) + except json.JSONDecodeError as error: + raise CaptureContractError(f"line {line_number}: invalid JSON: {error.msg}") from error + except CaptureContractError as error: + raise CaptureContractError(f"line {line_number}: {error}") from error + if not isinstance(raw, dict): + raise CaptureContractError(f"line {line_number}: prompt must be an object") + try: + prompts.append(PromptSpec.from_mapping(raw)) + except CaptureContractError as error: + raise CaptureContractError(f"line {line_number}: {error}") from error + if not prompts: + raise CaptureContractError("prompt file contains no prompts") + + by_split = {split: [prompt for prompt in prompts if prompt.split == split] for split in _SPLITS} + if any(not values for values in by_split.values()): + raise CaptureContractError("prompt file requires calibration and heldout splits") + calibration_ids = {prompt.prompt_id for prompt in by_split["calibration"]} + heldout_ids = {prompt.prompt_id for prompt in by_split["heldout"]} + for split, values in by_split.items(): + identifiers = [prompt.prompt_id for prompt in values] + if len(identifiers) != len(set(identifiers)): + raise CaptureContractError(f"prompt IDs must be unique within the {split} split") + if calibration_ids & heldout_ids: + raise CaptureContractError("prompt IDs overlap calibration and heldout splits") + group_assignments: dict[str, tuple[str, int | None]] = {} + for prompt in prompts: + assignment = (prompt.partition, prompt.inner_fold) + previous = group_assignments.setdefault(prompt.source_group_sha256, assignment) + if previous != assignment: + raise CaptureContractError( + "one source group is assigned to multiple partitions or inner folds" + ) + split_buckets = { + split: {prompt.bucket for prompt in values} for split, values in by_split.items() + } + if split_buckets["calibration"] != split_buckets["heldout"]: + raise CaptureContractError("heldout context buckets must match calibration buckets") + keys = [(prompt.split, prompt.prompt_id, prompt.bucket) for prompt in prompts] + if len(keys) != len(set(keys)): + raise CaptureContractError("prompt file repeats a split/prompt/context capture") + return sorted( + prompts, + key=lambda prompt: ( + prompt.min_kv_tokens, + math.inf if prompt.max_kv_tokens is None else prompt.max_kv_tokens, + prompt.split, + prompt.prompt_id, + ), + ) + + +def load_prompt_specs(path: str | Path) -> list[PromptSpec]: + """Load strict JSONL prompt specifications from a path.""" + return parse_prompt_specs_jsonl(Path(path).read_bytes()) + + +def _parse_strict_json_object(payload: bytes, label: str) -> Mapping[str, object]: + try: + raw = json.loads(payload, object_pairs_hook=_reject_duplicate_json_keys) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + detail = error.msg if isinstance(error, json.JSONDecodeError) else str(error) + raise CaptureContractError(f"{label} is invalid JSON: {detail}") from error + if not isinstance(raw, dict): + raise CaptureContractError(f"{label} must be a JSON object") + return raw + + +def parse_vanilla_prefill_fit(payload: bytes) -> dict[str, object]: + """Parse and canonicalize exact vanilla-calibration JSON bytes.""" + return canonical_prefill_threshold_scale_factor( + _parse_strict_json_object(payload, "vanilla calibration") + ) + + +def load_vanilla_prefill_fit(path: str | Path) -> dict[str, object]: + """Load and canonicalize an existing ModelOpt vanilla skip-softmax fit.""" + return parse_vanilla_prefill_fit(Path(path).read_bytes()) + + +def source_capture_sha256(prompt_token_ids: Sequence[int]) -> str: + """Fingerprint exact token IDs independently of prompt labels and splits.""" + if not prompt_token_ids: + raise CaptureContractError("prompt token IDs must not be empty") + digest = sha256(b"modelopt-mask-reuse-token-ids-v1\0") + digest.update(struct.pack("= 2**64: + raise CaptureContractError(f"prompt_token_ids[{index}] exceeds uint64") + digest.update(struct.pack(" dict[str, int]: + q_start = ((sample_length - 1) // MAX_QUERY_CHUNK_TOKENS) * MAX_QUERY_CHUNK_TOKENS + q_tokens = sample_length - q_start + if q_tokens <= _QUERY_START_ALIGNMENT: + raise CaptureContractError( + "the deployed q-stage-2 contract requires a final prefill chunk of at least " + "129 tokens; adjust the prompt length" + ) + return { + "q_tokens": q_tokens, + "kv_tokens": sample_length, + "q_start_tokens": q_start, + } + + +def build_capture_invocation( + *, + model: str, + checkpoint_manifest_sha256: str, + prompt: PromptSpec, + prompt_token_ids: Sequence[int], + target_sparsity: float, + threshold_scale_factor: Mapping[str, object], +) -> dict[str, object]: + """Build the exact invocation the backend must echo before evidence is trusted.""" + model = _text(model, "model") + checkpoint_identity = _sha256(checkpoint_manifest_sha256, "checkpoint_manifest_sha256") + sample_length = len(prompt_token_ids) + if sample_length < prompt.min_kv_tokens or ( + prompt.max_kv_tokens is not None and sample_length > prompt.max_kv_tokens + ): + raise CaptureContractError( + f"prompt {prompt.prompt_id!r} token length {sample_length} lies outside " + f"bucket [{prompt.min_kv_tokens}, {prompt.max_kv_tokens}]" + ) + target = _finite_number(target_sparsity, "target_sparsity", minimum=0.0, maximum=1.0) + if not 0.0 < target < 1.0: + raise CaptureContractError("target_sparsity must be in (0, 1)") + + if set(threshold_scale_factor) != {"formula", "prefill"}: + raise CaptureContractError("threshold_scale_factor must be canonical ModelOpt metadata") + params = threshold_scale_factor["prefill"] + if not isinstance(params, Mapping): + raise CaptureContractError("threshold_scale_factor.prefill must be an object") + lower = params.get("min_observed_sparsity") + upper = params.get("max_observed_sparsity") + if (lower is not None and target < float(lower)) or ( + upper is not None and target > float(upper) + ): + raise CaptureContractError( + f"target_sparsity={target} is outside the observed vanilla calibration range" + ) + a = float(params["a"]) + b = float(params["b"]) + threshold_log2 = math.log2(a) + b * target * math.log2(math.e) - math.log2(sample_length) + threshold_lambda = math.exp2(threshold_log2) # type: ignore[attr-defined] + if not math.isfinite(threshold_log2) or not 0.0 < threshold_lambda < 1.0: + raise CaptureContractError("vanilla fit derives a threshold outside (0, 1)") + + invocation = { + "capture_schema_version": CAPTURE_SCHEMA_VERSION, + "model": model, + "checkpoint_manifest_sha256": checkpoint_identity, + "split": prompt.split, + "partition": prompt.partition, + "inner_fold": prompt.inner_fold, + "prompt_id": prompt.prompt_id, + "source": prompt.source, + "source_group_sha256": prompt.source_group_sha256, + "source_capture_sha256": source_capture_sha256(prompt_token_ids), + "min_kv_tokens": prompt.min_kv_tokens, + "max_kv_tokens": prompt.max_kv_tokens, + "target_sparsity_hex": target.hex(), + "sample_length": sample_length, + "threshold_log2_hex": threshold_log2.hex(), + "threshold_lambda_hex": threshold_lambda.hex(), + "expected_geometry": _expected_final_geometry(sample_length), + } + _validate_invocation(invocation) + return invocation + + +def _validate_geometry(raw: object, label: str) -> dict[str, int]: + if not isinstance(raw, Mapping): + raise CaptureContractError(f"{label} must be an object") + _exact_fields(raw, _GEOMETRY_FIELDS, label) + q_tokens = _integer(raw["q_tokens"], f"{label}.q_tokens", minimum=129) + if q_tokens > MAX_QUERY_CHUNK_TOKENS: + raise CaptureContractError(f"{label}.q_tokens exceeds {MAX_QUERY_CHUNK_TOKENS}") + kv_tokens = _integer(raw["kv_tokens"], f"{label}.kv_tokens", minimum=q_tokens) + q_start = _integer(raw["q_start_tokens"], f"{label}.q_start_tokens") + if q_start % _QUERY_START_ALIGNMENT or q_start + q_tokens != kv_tokens: + raise CaptureContractError(f"{label} is not a 128-token-aligned final chunk") + return {"q_tokens": q_tokens, "kv_tokens": kv_tokens, "q_start_tokens": q_start} + + +def _eligible_tiles(geometry: Mapping[str, int]) -> int: + q_blocks = (geometry["q_tokens"] + 127) // 128 + first_eligible = geometry["q_start_tokens"] // 128 + 1 + return q_blocks * (2 * first_eligible + q_blocks - 1) // 2 + + +def _validate_invocation(raw: object) -> dict[str, object]: + if not isinstance(raw, Mapping): + raise CaptureContractError("capture invocation must be an object") + _exact_fields(raw, _INVOCATION_FIELDS, "capture invocation") + if raw["capture_schema_version"] != CAPTURE_SCHEMA_VERSION: + raise CaptureContractError("capture invocation has an unsupported schema version") + _text(raw["model"], "capture invocation.model") + _sha256( + raw["checkpoint_manifest_sha256"], + "capture invocation.checkpoint_manifest_sha256", + ) + if raw["split"] not in _SPLITS: + raise CaptureContractError("capture invocation.split is invalid") + partition = raw["partition"] + if partition not in _PARTITIONS: + raise CaptureContractError("capture invocation.partition is invalid") + expected_split = "calibration" if partition == "development" else "heldout" + if raw["split"] != expected_split: + raise CaptureContractError("capture invocation split and partition disagree") + if partition == "development": + _integer(raw["inner_fold"], "capture invocation.inner_fold") + elif raw["inner_fold"] is not None: + raise CaptureContractError("outer_test capture invocation must have null inner_fold") + _text(raw["prompt_id"], "capture invocation.prompt_id") + _text(raw["source"], "capture invocation.source") + _sha256(raw["source_group_sha256"], "capture invocation.source_group_sha256") + _sha256(raw["source_capture_sha256"], "capture invocation.source_capture_sha256") + minimum = _integer(raw["min_kv_tokens"], "capture invocation.min_kv_tokens", minimum=1) + maximum = raw["max_kv_tokens"] + if maximum is not None: + _integer(maximum, "capture invocation.max_kv_tokens", minimum=minimum) + target = _canonical_float_hex( + raw["target_sparsity_hex"], "capture invocation.target_sparsity_hex" + ) + if not 0.0 < target < 1.0: + raise CaptureContractError("capture invocation target sparsity must be in (0, 1)") + sample_length = _integer(raw["sample_length"], "capture invocation.sample_length", minimum=1) + threshold_log2 = _canonical_float_hex( + raw["threshold_log2_hex"], "capture invocation.threshold_log2_hex" + ) + threshold_lambda = _canonical_float_hex( + raw["threshold_lambda_hex"], "capture invocation.threshold_lambda_hex" + ) + if threshold_log2 >= 0.0 or not 0.0 < threshold_lambda < 1.0: + raise CaptureContractError("capture invocation threshold must be in (0, 1)") + if math.exp2(threshold_log2).hex() != threshold_lambda.hex(): # type: ignore[attr-defined] + raise CaptureContractError("capture invocation threshold hex fields disagree") + geometry = _validate_geometry(raw["expected_geometry"], "expected_geometry") + if geometry["kv_tokens"] != sample_length: + raise CaptureContractError("expected geometry does not match sample_length") + return dict(raw) + + +def _validate_rank_envelope( + values: Sequence[Mapping[str, object]], expected_fields: frozenset[str], label: str +) -> tuple[int, list[Mapping[str, object]]]: + if not values: + raise CaptureContractError(f"{label} returned no rank payloads") + world_sizes: set[int] = set() + ranks: list[int] = [] + for index, value in enumerate(values): + if not isinstance(value, Mapping): + raise CaptureContractError(f"{label}[{index}] must be an object") + _exact_fields(value, expected_fields, f"{label}[{index}]") + if value["capture_schema_version"] != CAPTURE_SCHEMA_VERSION: + raise CaptureContractError(f"{label}[{index}] has an unsupported schema version") + ranks.append(_integer(value["rank"], f"{label}[{index}].rank")) + world_sizes.add(_integer(value["world_size"], f"{label}[{index}].world_size", minimum=1)) + if len(world_sizes) != 1: + raise CaptureContractError(f"{label} disagrees on world_size") + world_size = next(iter(world_sizes)) + if len(values) != world_size or sorted(ranks) != list(range(world_size)): + raise CaptureContractError(f"{label} does not exactly cover ranks [0, {world_size})") + return world_size, sorted(values, key=lambda value: cast("int", value["rank"])) + + +def validate_capture_statuses(values: Sequence[Mapping[str, object]]) -> int: + """Require the env-gated backend capture sink on every worker rank.""" + world_size, ordered = _validate_rank_envelope(values, _STATUS_FIELDS, "capture status") + failures = [] + for value in ordered: + if not isinstance(value["available"], bool): + raise CaptureContractError("capture status.available must be boolean") + reason = value["reason"] + if reason is not None and not isinstance(reason, str): + raise CaptureContractError("capture status.reason must be a string or null") + if not value["available"]: + failures.append(f"rank {value['rank']}: {reason or 'unavailable'}") + if failures: + raise CaptureContractError("mask-reuse capture is unavailable: " + "; ".join(failures)) + return world_size + + +def validate_begin_acks( + values: Sequence[Mapping[str, object]], invocation: Mapping[str, object] +) -> int: + """Validate that every rank armed the exact same invocation before generation.""" + _validate_invocation(invocation) + expected_digest = canonical_json_sha256(invocation) + world_size, ordered = _validate_rank_envelope(values, _ACK_FIELDS, "capture begin") + for value in ordered: + if value["armed"] is not True: + raise CaptureContractError(f"capture begin rank {value['rank']} did not arm") + if value["invocation_sha256"] != expected_digest: + raise CaptureContractError( + f"capture begin rank {value['rank']} acknowledged the wrong invocation" + ) + return world_size + + +def _parse_anchor_stats( + raw: object, *, global_num_heads: int, eligible_tiles: int +) -> dict[str, dict[str, list[int] | list[float]]]: + if not isinstance(raw, Mapping) or not raw: + raise CaptureContractError("anchor_stats_by_layer must be a non-empty object") + result: dict[str, dict[str, list[int] | list[float]]] = {} + for raw_layer, value in raw.items(): + if not isinstance(raw_layer, str): + raise CaptureContractError("anchor layer keys must be canonical integer strings") + layer = _integer(int(raw_layer), f"anchor layer {raw_layer}") if raw_layer.isdigit() else -1 + if layer < 0 or raw_layer != str(layer): + raise CaptureContractError("anchor layer keys must be canonical integer strings") + if not isinstance(value, Mapping): + raise CaptureContractError(f"anchor_stats_by_layer[{layer}] must be an object") + _exact_fields(value, _ANCHOR_STATS_FIELDS, f"anchor_stats_by_layer[{layer}]") + retained_raw = value["retained_tiles"] + dropped_raw = value["dropped_mass"] + if not isinstance(retained_raw, list) or not isinstance(dropped_raw, list): + raise CaptureContractError(f"anchor_stats_by_layer[{layer}] arrays must be lists") + if len(retained_raw) != global_num_heads or len(dropped_raw) != global_num_heads: + raise CaptureContractError( + f"anchor_stats_by_layer[{layer}] must cover {global_num_heads} heads" + ) + retained = [ + _integer(value, f"anchor_stats_by_layer[{layer}].retained_tiles[{head}]") + for head, value in enumerate(retained_raw) + ] + if any(value > eligible_tiles for value in retained): + raise CaptureContractError( + f"anchor_stats_by_layer[{layer}] retained tiles exceed eligible tiles" + ) + dropped = [ + _finite_number( + value, + f"anchor_stats_by_layer[{layer}].dropped_mass[{head}]", + minimum=0.0, + maximum=1.0, + ) + for head, value in enumerate(dropped_raw) + ] + result[raw_layer] = {"retained_tiles": retained, "dropped_mass": dropped} + return dict(sorted(result.items(), key=lambda item: int(item[0]))) + + +def _parse_consumer_layers( + raw: object, *, rank: int, global_num_heads: int +) -> dict[int, tuple[int, int, list[list[float]]]]: + if not isinstance(raw, Mapping) or not raw: + raise CaptureContractError(f"rank {rank} consumer_layers must be a non-empty object") + result: dict[int, tuple[int, int, list[list[float]]]] = {} + for raw_layer, value in raw.items(): + if ( + not isinstance(raw_layer, str) + or not raw_layer.isdigit() + or raw_layer != str(int(raw_layer)) + ): + raise CaptureContractError("consumer layer keys must be canonical integer strings") + layer = int(raw_layer) + if not isinstance(value, Mapping): + raise CaptureContractError(f"consumer_layers[{layer}] must be an object") + _exact_fields(value, _CONSUMER_STATS_FIELDS, f"consumer_layers[{layer}]") + anchor = _integer(value["anchor_layer"], f"consumer_layers[{layer}].anchor_layer") + if anchor >= layer: + raise CaptureContractError(f"consumer layer {layer} must follow anchor {anchor}") + start = _integer( + value["consumer_head_start"], f"consumer_layers[{layer}].consumer_head_start" + ) + matrix = value["dropped_mass"] + if not isinstance(matrix, list) or not matrix: + raise CaptureContractError(f"consumer_layers[{layer}].dropped_mass must be non-empty") + rows: list[list[float]] = [] + for local_head, raw_row in enumerate(matrix): + if not isinstance(raw_row, list) or len(raw_row) != global_num_heads: + raise CaptureContractError( + f"consumer_layers[{layer}].dropped_mass[{local_head}] must cover " + f"{global_num_heads} donor heads" + ) + rows.append( + [ + _finite_number( + value, + f"consumer_layers[{layer}].dropped_mass[{local_head}][{donor}]", + minimum=0.0, + maximum=1.0, + ) + for donor, value in enumerate(raw_row) + ] + ) + if start + len(rows) > global_num_heads: + raise CaptureContractError( + f"rank {rank} consumer layer {layer} shard exceeds head count" + ) + result[layer] = (anchor, start, rows) + return result + + +@dataclass(frozen=True, slots=True) +class MergedCapture: + """One compact normalized capture plus auditable metadata.""" + + capture: dict[str, object] + manifest: dict[str, object] + + +def merge_rank_captures( + values: Sequence[Mapping[str, object]], invocation: Mapping[str, object] +) -> MergedCapture: + """Merge raw per-rank sufficient statistics without inventing absent rows.""" + trusted_invocation = _validate_invocation(invocation) + expected_digest = canonical_json_sha256(trusted_invocation) + world_size, ordered = _validate_rank_envelope(values, _RANK_CAPTURE_FIELDS, "capture drain") + global_heads: set[int] = set() + eligible_values: set[int] = set() + geometry_values: list[dict[str, int]] = [] + anchor_payloads: list[dict[str, dict[str, list[int] | list[float]]]] = [] + consumer_payloads: list[dict[int, tuple[int, int, list[list[float]]]]] = [] + for value in ordered: + rank = cast("int", value["rank"]) + if value["invocation"] != trusted_invocation: + raise CaptureContractError(f"capture drain rank {rank} echoed the wrong invocation") + if value["invocation_sha256"] != expected_digest: + raise CaptureContractError(f"capture drain rank {rank} has the wrong invocation digest") + geometry = _validate_geometry(value["geometry"], f"capture drain rank {rank} geometry") + if geometry != trusted_invocation["expected_geometry"]: + raise CaptureContractError(f"capture drain rank {rank} measured the wrong final chunk") + geometry_values.append(geometry) + global_num_heads = _integer( + value["global_num_heads"], f"capture drain rank {rank} global_num_heads", minimum=1 + ) + eligible_tiles = _integer( + value["eligible_tiles"], f"capture drain rank {rank} eligible_tiles", minimum=1 + ) + global_heads.add(global_num_heads) + eligible_values.add(eligible_tiles) + anchor_payloads.append( + _parse_anchor_stats( + value["anchor_stats_by_layer"], + global_num_heads=global_num_heads, + eligible_tiles=eligible_tiles, + ) + ) + consumer_payloads.append( + _parse_consumer_layers( + value["consumer_layers"], rank=rank, global_num_heads=global_num_heads + ) + ) + if len(global_heads) != 1 or len(eligible_values) != 1: + raise CaptureContractError("capture ranks disagree on head count or eligible tiles") + if any(value != geometry_values[0] for value in geometry_values[1:]): + raise CaptureContractError("capture ranks disagree on final-chunk geometry") + if any(value != anchor_payloads[0] for value in anchor_payloads[1:]): + raise CaptureContractError("capture ranks disagree on global anchor statistics") + layer_sets = [set(value) for value in consumer_payloads] + if any(value != layer_sets[0] for value in layer_sets[1:]): + raise CaptureContractError("capture ranks disagree on consumer layer coverage") + + global_num_heads = next(iter(global_heads)) + eligible_tiles = next(iter(eligible_values)) + if eligible_tiles != _eligible_tiles(geometry_values[0]): + raise CaptureContractError( + "capture eligible_tiles does not match 128x128 bottom-right causal geometry" + ) + anchors = anchor_payloads[0] + merged_consumers: dict[str, dict[str, object]] = {} + for layer in sorted(layer_sets[0]): + shards = sorted( + (payload[layer] for payload in consumer_payloads), key=lambda value: value[1] + ) + anchor_layers = {shard[0] for shard in shards} + if len(anchor_layers) != 1: + raise CaptureContractError(f"capture ranks disagree on consumer layer {layer}'s anchor") + anchor_layer = next(iter(anchor_layers)) + anchor_stats = anchors.get(str(anchor_layer)) + if anchor_stats is None: + raise CaptureContractError( + f"consumer layer {layer} references missing anchor layer {anchor_layer}" + ) + cursor = 0 + global_rows: list[list[float]] = [] + for _, start, local_rows in shards: + if start != cursor: + raise CaptureContractError( + f"consumer layer {layer} head shards are overlapping or incomplete at {cursor}" + ) + global_rows.extend(local_rows) + cursor += len(local_rows) + if cursor != global_num_heads: + raise CaptureContractError( + f"consumer layer {layer} head shards cover {cursor}, expected {global_num_heads}" + ) + merged_consumers[str(layer)] = { + "anchor_layer": anchor_layer, + "dropped_mass": global_rows, + } + if not merged_consumers: + raise CaptureContractError("capture drain contained no consumer-head statistics") + compact_capture = { + "compact_capture_schema_version": 1, + "invocation": trusted_invocation, + "geometry": geometry_values[0], + "global_num_heads": global_num_heads, + "eligible_tiles": eligible_tiles, + "anchor_stats_by_layer": anchors, + "consumer_layers": merged_consumers, + } + manifest = { + "capture_schema_version": CAPTURE_SCHEMA_VERSION, + "invocation": trusted_invocation, + "invocation_sha256": expected_digest, + "world_size": world_size, + "global_num_heads": global_num_heads, + "eligible_tiles": eligible_tiles, + "candidate_cell_count": sum( + len(cast("Sequence[object]", value["dropped_mass"])) * global_num_heads + for value in merged_consumers.values() + ), + "compact_capture_sha256": canonical_json_sha256(compact_capture), + } + return MergedCapture(compact_capture, manifest) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py index 32fb5d67160..2d23089ecdd 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -233,10 +233,14 @@ def build_sparse_attention_config( 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] = { + phase_params = { "a": float(calibration_params[phase]["a"]), "b": float(calibration_params[phase]["b"]), } + for key in ("min_observed_sparsity", "max_observed_sparsity"): + if key in calibration_params[phase]: + phase_params[key] = float(calibration_params[phase][key]) + threshold_scale_factor[phase] = phase_params skip_group: dict[str, Any] = { "algorithm": "skip_softmax", diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_mask_reuse_capture.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_mask_reuse_capture.py new file mode 100644 index 00000000000..0b01c3cb87d --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_mask_reuse_capture.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""vLLM worker bootstrap and RPCs for policy-free mask-reuse calibration.""" + +from __future__ import annotations + +import importlib +import os +from collections.abc import Mapping +from typing import TYPE_CHECKING + +from vllm.v1.worker.gpu_worker import Worker as BaseWorker + +if TYPE_CHECKING: + from types import ModuleType + +CAPTURE_ENV = "MASK_REUSE_FA4_CALIBRATION_CAPTURE" +PLAN_ENV = "MASK_REUSE_FA4_PLAN" +_REQUIRED_API = ( + "configure_capture_runtime", + "capture_status", + "begin_capture", + "drain_capture", +) + +__all__ = ["MaskReuseCaptureWorker"] + + +def _capture_api() -> ModuleType: + if os.environ.get(CAPTURE_ENV) != "1": + raise RuntimeError(f"{CAPTURE_ENV}=1 is required for mask-reuse calibration capture") + try: + api = importlib.import_module("mask_reuse_vllm.capture") + except ImportError as error: + raise RuntimeError( + "the custom mask-reuse backend does not expose mask_reuse_vllm.capture" + ) from error + missing = [name for name in _REQUIRED_API if not callable(getattr(api, name, None))] + if missing: + raise RuntimeError(f"mask-reuse capture API is incomplete; missing {missing}") + return api + + +def _configure_capture_before_model_load() -> ModuleType: + """Install a planner and capture provider without loading a serving policy.""" + plan_name = os.environ.get(PLAN_ENV) + if not plan_name: + raise RuntimeError(f"{PLAN_ENV} must name the explicit calibration topology preset") + api = _capture_api() + api.configure_capture_runtime(plan_name) + return api + + +class MaskReuseCaptureWorker(BaseWorker): + """Run the custom backend in env-gated, policy-free capture mode.""" + + def load_model(self, *args, **kwargs) -> None: + """Install capture runtime before vLLM constructs attention modules.""" + # The attention implementation resolves process-local runtime state + # while the model is loading. Install the policy-free capture provider + # first; a promoted v3 policy must not be required to collect its own + # calibration evidence. + api = _configure_capture_before_model_load() + super().load_model(*args, **kwargs) + status = api.capture_status() + if not isinstance(status, Mapping) or status.get("available") is not True: + reason = status.get("reason") if isinstance(status, Mapping) else None + raise RuntimeError(f"mask-reuse capture backend is unavailable: {reason or status!r}") + + def mask_reuse_capture_status(self) -> dict[str, object]: + """Return this worker rank's fail-closed capture status.""" + return dict(_capture_api().capture_status()) + + def mask_reuse_capture_begin(self, invocation: dict[str, object]) -> dict[str, object]: + """Arm exactly one prompt/target invocation on this worker rank.""" + return dict(_capture_api().begin_capture(invocation)) + + def mask_reuse_capture_drain(self) -> dict[str, object]: + """Drain one completed rank-local sufficient-stat payload.""" + return dict(_capture_api().drain_capture()) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_calibrate_mask_reuse_cli.py b/tests/unit/torch/sparsity/attention_sparsity/test_calibrate_mask_reuse_cli.py new file mode 100644 index 00000000000..a65405fa1c2 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_calibrate_mask_reuse_cli.py @@ -0,0 +1,252 @@ +# 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": 2, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v2", + "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, + "target_sparsity_hex": [(0.7).hex()], + "vanilla_threshold_scale_factor": {"formula": "unused"}, + "vanilla_fit_sha256": sha256(b"normalized-fit").hexdigest(), + "vanilla_config_file_sha256": sha256(vanilla.read_bytes()).hexdigest(), + "prompt_plan_file_sha256": sha256(b"prompts").hexdigest(), + "compact_capture_file_sha256": sha256(captures.read_bytes()).hexdigest(), + "capture_count": 1, + "candidate_cell_count": 4, + "captures": [{"candidate_cell_count": 4}], + } + ) + ) + policy = tmp_path / "candidate.json" + report = tmp_path / "report.json" + args = [ + "--checkpoint", + str(checkpoint_root), + "--compact-captures", + str(captures), + "--capture-manifest", + str(capture_manifest), + "--vanilla-config", + str(vanilla), + "--topology", + str(topology), + "--calibration-plan", + str(artifacts["calibration_plan"]), + "--family-registry", + str(artifacts["family_registry"]), + "--grouped-fit", + str(artifacts["grouped_fit"]), + "--outer-report", + str(artifacts["outer_report"]), + "--max-anchor-dropped-mass", + "0.02", + "--max-reuse-dropped-mass", + "0.03", + "--max-reuse-selection-dropped-mass", + "0.01", + "--output-policy", + str(policy), + "--output-report", + str(report), + ] + return args, policy, report, captures, vanilla, artifacts + + +def test_main_verifies_artifacts_and_atomically_writes_candidate(tmp_path, monkeypatch, capsys): + args, policy_path, report_path, captures, vanilla, artifacts = _base_args(tmp_path) + source = object() + monkeypatch.setattr( + calibrate_mask_reuse_cli, "load_compact_mask_reuse_captures", lambda path: source + ) + captured = {} + artifact = { + "version": 3, + "promotion_status": "candidate_only", + "deployment_geometry_validated": False, + "provenance": {"input_capture_count": 1, "candidate_cell_count": 4}, + "calibration_report": {"promotion": {"eligible": False}}, + } + + def fake_calibrate(compact_source, **kwargs): + captured["source"] = compact_source + captured.update(kwargs) + return artifact + + monkeypatch.setattr( + calibrate_mask_reuse_cli, "calibrate_compact_mask_reuse_policy", fake_calibrate + ) + + assert calibrate_mask_reuse_cli.main(args) == 0 + + expected_policy = _canonical(artifact) + assert policy_path.read_bytes() == expected_policy + assert report_path.read_bytes() == _canonical(artifact["calibration_report"]) + assert captured["source"] is source + assert captured["evidence"] == { + "calibration_plan_sha256": sha256(artifacts["calibration_plan"].read_bytes()).hexdigest(), + "family_registry_sha256": sha256(artifacts["family_registry"].read_bytes()).hexdigest(), + "grouped_fit_sha256": sha256(artifacts["grouped_fit"].read_bytes()).hexdigest(), + "outer_report_sha256": sha256(artifacts["outer_report"].read_bytes()).hexdigest(), + "vanilla_fit_sha256": sha256(vanilla.read_bytes()).hexdigest(), + "reuse_bundle_sha256": sha256(captures.read_bytes()).hexdigest(), + } + assert "MASK_REUSE_FA4_CANDIDATE_SHA256=" in capsys.readouterr().out + + +def test_vanilla_mutation_after_semantic_snapshot_cannot_publish(tmp_path, monkeypatch, capsys): + args, policy_path, report_path, _, vanilla, _ = _base_args(tmp_path) + original_payload = vanilla.read_bytes() + real_evidence_artifacts = calibrate_mask_reuse_cli._evidence_artifacts + + def mutate_after_snapshot(namespace, *, vanilla_fit_sha256): + vanilla.write_text('{"vanilla":false}\n', encoding="utf-8") + return real_evidence_artifacts(namespace, vanilla_fit_sha256=vanilla_fit_sha256) + + monkeypatch.setattr(calibrate_mask_reuse_cli, "_evidence_artifacts", mutate_after_snapshot) + monkeypatch.setattr( + calibrate_mask_reuse_cli, + "load_compact_mask_reuse_captures", + lambda path: object(), + ) + + def fake_calibrate(compact_source, **kwargs): + assert kwargs["vanilla_calibration"] == {"vanilla": True} + assert kwargs["evidence"]["vanilla_fit_sha256"] == sha256(original_payload).hexdigest() + return { + "promotion_status": "candidate_only", + "deployment_geometry_validated": False, + "provenance": {"input_capture_count": 1, "candidate_cell_count": 4}, + "calibration_report": {"promotion": {"eligible": False}}, + } + + monkeypatch.setattr( + calibrate_mask_reuse_cli, + "calibrate_compact_mask_reuse_policy", + fake_calibrate, + ) + + with pytest.raises(SystemExit, match="2"): + calibrate_mask_reuse_cli.main(args) + + assert "vanilla_fit_sha256 artifact changed during calibration" in capsys.readouterr().err + assert not policy_path.exists() + assert not report_path.exists() + + +def test_main_rejects_capture_manifest_hash_mismatch(tmp_path, capsys): + args, _, _, captures, _, _ = _base_args(tmp_path) + captures.write_text("changed\n", encoding="utf-8") + + with pytest.raises(SystemExit, match="2"): + calibrate_mask_reuse_cli.main(args) + + assert "compact_capture_file_sha256" in capsys.readouterr().err + + +def test_load_json_object_rejects_duplicate_keys(tmp_path): + path = tmp_path / "duplicate.json" + path.write_text('{"anchors":[0],"anchors":[1]}', encoding="utf-8") + + with pytest.raises(ValueError, match="duplicate JSON key 'anchors'"): + calibrate_mask_reuse_cli._load_json_object(path, label="topology") + + +def test_candidate_publication_rolls_back_report_when_policy_race_wins(tmp_path, monkeypatch): + policy = tmp_path / "candidate.json" + report = tmp_path / "report.json" + real_link = os.link + call_count = 0 + + def racing_link(source, target, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 2: + Path(target).write_bytes(b"racer") + return real_link(source, target, **kwargs) + + monkeypatch.setattr(calibrate_mask_reuse_cli.os, "link", racing_link) + + with pytest.raises(FileExistsError): + calibrate_mask_reuse_cli._publish_candidate_outputs(policy, b"ours", report, b"report") + + assert policy.read_bytes() == b"racer" + assert not report.exists() + + +def test_candidate_publication_refuses_existing_output(tmp_path): + policy = tmp_path / "candidate.json" + report = tmp_path / "report.json" + policy.write_bytes(b"existing") + + with pytest.raises(FileExistsError, match="refusing to overwrite"): + calibrate_mask_reuse_cli._publish_candidate_outputs(policy, b"ours", report, b"report") + + assert policy.read_bytes() == b"existing" + assert not report.exists() diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py b/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py new file mode 100644 index 00000000000..d1090296e0b --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py @@ -0,0 +1,72 @@ +# 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 pathlib import Path + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + CHECKPOINT_MANIFEST_NAME, + CheckpointManifestError, + create_checkpoint_manifest, + 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) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py b/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py new file mode 100644 index 00000000000..c2f10b9d66e --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py @@ -0,0 +1,344 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Focused launch/RPC test for the vLLM mask-reuse collection driver.""" + +import importlib.util +import json +import os +import sys +from hashlib import sha256 +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + create_checkpoint_manifest, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.mask_reuse_capture import ( + canonical_json_sha256, +) + +_SCRIPT_PATH = Path(__file__).parents[5] / "examples/vllm_serve/collect_mask_reuse.py" +_SPEC = importlib.util.spec_from_file_location("collect_mask_reuse_cli", _SCRIPT_PATH) +assert _SPEC is not None and _SPEC.loader is not None +collect_mask_reuse_cli = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(collect_mask_reuse_cli) + + +class _Tokenizer: + def encode(self, prompt, *, add_special_tokens): + assert add_special_tokens is True + offset = ord(prompt[0]) + return [offset + index for index in range(256)] + + +class _SamplingParams: + def __init__(self, **kwargs): + self.kwargs = kwargs + + +class _LLM: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.invocation = None + self.events = [] + self.__class__.instances.append(self) + + def get_tokenizer(self): + return _Tokenizer() + + def collective_rpc(self, method, args=()): + self.events.append(method) + if method == "mask_reuse_capture_status": + return [ + { + "capture_schema_version": 2, + "available": True, + "rank": 0, + "world_size": 1, + "reason": None, + } + ] + if method == "mask_reuse_capture_begin": + self.invocation = args[0] + return [ + { + "capture_schema_version": 2, + "armed": True, + "rank": 0, + "world_size": 1, + "invocation_sha256": canonical_json_sha256(self.invocation), + } + ] + assert method == "mask_reuse_capture_drain" + invocation = self.invocation + return [ + { + "capture_schema_version": 2, + "rank": 0, + "world_size": 1, + "invocation": invocation, + "invocation_sha256": canonical_json_sha256(invocation), + "geometry": invocation["expected_geometry"], + "global_num_heads": 2, + "eligible_tiles": 3, + "anchor_stats_by_layer": { + "0": {"retained_tiles": [2, 3], "dropped_mass": [0.01, 0.02]} + }, + "consumer_layers": { + "1": { + "anchor_layer": 0, + "consumer_head_start": 0, + "dropped_mass": [[0.01, 0.02], [0.03, 0.04]], + } + }, + } + ] + + def generate(self, token_ids, sampling, *, use_tqdm): + assert len(token_ids) == 256 + assert sampling.kwargs == {"temperature": 0.0, "max_tokens": 1, "ignore_eos": True} + assert use_tqdm is False + self.events.append("generate") + return [] + + +def test_main_bootstraps_policy_free_backend_and_writes_normalized_evidence(tmp_path, monkeypatch): + prompts = tmp_path / "prompts.jsonl" + prompts.write_text( + "\n".join( + json.dumps( + { + "split": split, + "partition": ("development" if split == "calibration" else "outer_test"), + "inner_fold": 0 if split == "calibration" else None, + "prompt_id": prompt_id, + "source": source, + "source_group_sha256": sha256(source.encode()).hexdigest(), + "prompt": text, + "min_kv_tokens": 129, + "max_kv_tokens": 512, + } + ) + for split, prompt_id, source, text in ( + ("calibration", "cal-0", "ruler/niah", "alpha"), + ("heldout", "held-0", "longbench/qasper", "beta"), + ) + ) + + "\n", + encoding="utf-8", + ) + vanilla = tmp_path / "config.json" + vanilla.write_text( + json.dumps( + { + "threshold_scale_factor": { + "formula": "a * exp(b * target_sparsity)", + "prefill": { + "a": 1.0, + "b": 1.0, + "min_observed_sparsity": 0.5, + "max_observed_sparsity": 0.8, + }, + } + } + ), + encoding="utf-8", + ) + output = tmp_path / "observations.jsonl" + manifest = tmp_path / "manifest.json" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{}\n", encoding="utf-8") + (checkpoint / "model.safetensors").write_bytes(b"toy-weights") + checkpoint_manifest = create_checkpoint_manifest(checkpoint, model="test-model") + fa4_source = tmp_path / "flash-attention" + cute = fa4_source / "flash_attn/cute" + cute.mkdir(parents=True) + (cute / "interface.py").write_text("# pinned\n", encoding="utf-8") + (cute / "block_sparsity.py").write_text("# pinned\n", encoding="utf-8") + fake_vllm = SimpleNamespace(LLM=_LLM, SamplingParams=_SamplingParams) + monkeypatch.setitem(sys.modules, "vllm", fake_vllm) + monkeypatch.setattr( + collect_mask_reuse_cli.importlib.metadata, + "entry_points", + lambda **kwargs: [ + SimpleNamespace(name="mask_reuse_fa4", value="mask_reuse_vllm.plugin:register") + ], + ) + monkeypatch.setattr( + collect_mask_reuse_cli.subprocess, + "run", + lambda command, **kwargs: SimpleNamespace( + stdout="a" * 40 + "\n" if command[-2:] == ["rev-parse", "HEAD"] else "" + ), + ) + monkeypatch.setenv("MASK_REUSE_FA4_POLICY", "stale-policy.json") + monkeypatch.setenv("MASK_REUSE_FA4_POLICY_SHA256", "stale") + parsed_payloads = {} + real_parse_prompts = collect_mask_reuse_cli.parse_prompt_specs_jsonl + real_parse_vanilla = collect_mask_reuse_cli.parse_vanilla_prefill_fit + + def parse_prompts(payload): + parsed_payloads["prompts"] = payload + return real_parse_prompts(payload) + + def parse_vanilla(payload): + parsed_payloads["vanilla"] = payload + return real_parse_vanilla(payload) + + monkeypatch.setattr(collect_mask_reuse_cli, "parse_prompt_specs_jsonl", parse_prompts) + monkeypatch.setattr(collect_mask_reuse_cli, "parse_vanilla_prefill_fit", parse_vanilla) + _LLM.instances.clear() + + result = collect_mask_reuse_cli.main( + [ + str(checkpoint), + "--model-id", + "test-model", + "--plan", + "test_stride2", + "--fa4-source", + str(fa4_source), + "--prompts-jsonl", + str(prompts), + "--vanilla-config", + str(vanilla), + "--target-sparsities", + "0.7", + "--output", + str(output), + "--output-manifest", + str(manifest), + "--max-model-len", + "512", + ] + ) + + assert result == 0 + engine = _LLM.instances[0] + assert engine.kwargs["attention_backend"] == "CUSTOM" + assert engine.kwargs["dtype"] == "bfloat16" + assert engine.kwargs["worker_cls"].endswith("MaskReuseCaptureWorker") + assert engine.kwargs["max_num_batched_tokens"] == 8192 + assert "quantization" not in engine.kwargs + assert engine.events == [ + "mask_reuse_capture_status", + "mask_reuse_capture_begin", + "generate", + "mask_reuse_capture_drain", + "mask_reuse_capture_begin", + "generate", + "mask_reuse_capture_drain", + ] + assert collect_mask_reuse_cli.os.environ["MASK_REUSE_FA4_CALIBRATION_CAPTURE"] == "1" + assert ( + collect_mask_reuse_cli.os.environ["MASK_REUSE_FA4_CHECKPOINT_MANIFEST_SHA256"] + == checkpoint_manifest.sha256 + ) + assert collect_mask_reuse_cli.os.environ["PYTHONDONTWRITEBYTECODE"] == "1" + assert "MASK_REUSE_FA4_POLICY" not in collect_mask_reuse_cli.os.environ + assert "MASK_REUSE_FA4_POLICY_SHA256" not in collect_mask_reuse_cli.os.environ + + captures = [json.loads(line) for line in output.read_text().splitlines()] + assert len(captures) == 2 + assert all(capture["compact_capture_schema_version"] == 1 for capture in captures) + assert all("observations" not in capture for capture in captures) + assert {capture["invocation"]["split"] for capture in captures} == { + "calibration", + "heldout", + } + assert {capture["invocation"]["target_sparsity_hex"] for capture in captures} == {(0.7).hex()} + report = json.loads(manifest.read_text()) + assert report["capture_protocol"] == "modelopt_vllm_mask_reuse_target_sparsity_v2" + assert report["checkpoint_manifest_sha256"] == checkpoint_manifest.sha256 + assert report["prompt_plan_file_sha256"] == sha256(parsed_payloads["prompts"]).hexdigest() + assert report["vanilla_config_file_sha256"] == sha256(parsed_payloads["vanilla"]).hexdigest() + assert report["fa4_source_commit"] == "a" * 40 + assert report["capture_count"] == len(captures) + assert "observation_count" not in report + assert len(report["captures"]) == 2 + assert {capture["invocation"]["source"] for capture in report["captures"]} == { + "ruler/niah", + "longbench/qasper", + } + + +def test_publish_no_clobber_preserves_destination_created_by_racer(tmp_path, monkeypatch): + temporary = tmp_path / "capture.tmp" + destination = tmp_path / "capture.jsonl" + temporary.write_text("ours", encoding="utf-8") + real_link = os.link + + def racing_link(source, target): + Path(target).write_text("racer", encoding="utf-8") + return real_link(source, target) + + monkeypatch.setattr(collect_mask_reuse_cli.os, "link", racing_link) + + with pytest.raises(FileExistsError): + collect_mask_reuse_cli._publish_no_clobber(temporary, destination) + + assert destination.read_text(encoding="utf-8") == "racer" + assert temporary.read_text(encoding="utf-8") == "ours" + + +def test_publish_no_clobber_rolls_back_capture_and_manifest_on_fsync_failure(tmp_path, monkeypatch): + def fail_fsync(path): + raise OSError("injected fsync failure") + + monkeypatch.setattr(collect_mask_reuse_cli, "_fsync_directory", fail_fsync) + + for name in ("capture.jsonl", "capture.manifest.json"): + temporary = tmp_path / f".{name}.tmp" + destination = tmp_path / name + temporary.write_text("complete", encoding="utf-8") + + with pytest.raises(OSError, match="injected fsync failure"): + collect_mask_reuse_cli._publish_no_clobber(temporary, destination) + + assert not destination.exists() + assert not temporary.exists() + + +def test_capture_environment_rejects_untracked_fa4_source(tmp_path, monkeypatch): + fa4_source = tmp_path / "flash-attention" + cute = fa4_source / "flash_attn/cute" + cute.mkdir(parents=True) + (cute / "interface.py").write_text("# pinned\n", encoding="utf-8") + (cute / "block_sparsity.py").write_text("# pinned\n", encoding="utf-8") + monkeypatch.setattr( + collect_mask_reuse_cli.subprocess, + "run", + lambda command, **kwargs: SimpleNamespace( + stdout=( + "a" * 40 + "\n" + if command[-2:] == ["rev-parse", "HEAD"] + else "?? flash_attn/cute/local_override.py\n" + ) + ), + ) + + with pytest.raises( + collect_mask_reuse_cli.CaptureContractError, + match="tracked or untracked modifications", + ): + collect_mask_reuse_cli._configure_capture_environment( + "test_stride2", str(fa4_source), "0" * 64 + ) 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..da8cde69c5b --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py @@ -0,0 +1,296 @@ +# 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 math.exp2(threshold_log2), threshold_log2 + + +def _observations() -> list[MaskReuseObservation]: + rows = [] + prompts = { + "calibration": (("cal-0", 65_536), ("cal-1", 98_304)), + "heldout": (("held-0", 65_536), ("held-1", 98_304)), + } + for split, samples in prompts.items(): + for prompt, sample_length in samples: + for target_sparsity in (0.5, 0.7): + threshold, threshold_log2 = _thresholds(target_sparsity, sample_length) + anchor_retained = { + 0.5: (80, 90), + 0.7: (40, 20), + }[target_sparsity] + anchor_dropped = (0.03, 0.03) if target_sparsity == 0.5 else (0.07, 0.08) + anchor_stats = { + 0: AnchorLayerStats(anchor_retained, anchor_dropped), + 2: AnchorLayerStats( + (60, 60) if target_sparsity == 0.5 else (30, 30), + (0.02, 0.02), + ), + } + for consumer_head in range(2): + for donor_head in range(2): + retained = anchor_retained[donor_head] + if target_sparsity == 0.7 and consumer_head == 1: + dropped_mass = 0.07 + 0.01 * donor_head + else: + dropped_mass = 0.02 + 0.01 * donor_head + rows.append( + MaskReuseObservation( + model="toy", + min_kv_tokens=65_536, + max_kv_tokens=131_072, + target_sparsity=target_sparsity, + sample_length=sample_length, + threshold_lambda=threshold, + threshold_log2=threshold_log2, + q_tokens=8192, + kv_tokens=sample_length, + q_start_tokens=sample_length - 8192, + split=split, + prompt_id=prompt, + source_capture_sha256=_source(prompt), + anchor_layer=0, + consumer_layer=1, + consumer_head=consumer_head, + donor_head=donor_head, + retained_tiles=retained, + eligible_tiles=100, + anchor_dropped_mass=( + 0.03 if target_sparsity == 0.5 else 0.07 + 0.01 * donor_head + ), + anchor_stats_by_layer=anchor_stats, + dropped_mass=dropped_mass, + ) + ) + return rows + + +def _calibrate(rows): + return calibrate_mask_reuse_policy( + rows, + vanilla_calibration=VANILLA_FIT, + topology=TOPOLOGY, + checkpoint_manifest=VERIFIED_CHECKPOINT, + evidence=EVIDENCE, + max_anchor_dropped_mass=0.1, + max_reuse_dropped_mass=0.1, + max_reuse_selection_dropped_mass=0.05, + ) + + +def test_selects_target_sparsity_and_exports_backend_v3(): + artifact = _calibrate(_observations()) + + assert artifact["version"] == 3 + assert artifact["phase"] == "prefill" + assert artifact["decode"] == {"mode": "dense"} + assert artifact["calibration_protocol"] == "modelopt_mask_reuse_target_sparsity_v1" + assert artifact["producer"] == {"name": "modelopt", "version": modelopt.__version__} + assert artifact["evidence"] == EVIDENCE + assert artifact["threshold_scale_factor"] == { + "formula": "a * exp(b * target_sparsity)", + "prefill": VANILLA_FIT["prefill"], + } + assert artifact["context_policies"] == [ + { + "min_kv_tokens": 65_536, + "max_kv_tokens": 131_072, + "target_sparsity": 0.7, + "headmaps": {"1": [1, 0]}, + "fallback_heads": {"1": [1]}, + } + ] + assert artifact["promotion_status"] == "candidate_only" + assert artifact["deployment_geometry_validated"] is False + assert artifact["deployment_geometry"]["contract"]["kv_page_tokens"] == 16 + assert len(artifact["deployment_geometry"]["observations"]) == 4 + assert ( + artifact["calibration_report"]["overall"]["reuse_heldout"]["constraint_violation_rate"] + == 0.0 + ) + json.dumps(artifact) + + +def test_heldout_values_evaluate_but_cannot_change_selection(): + baseline = _calibrate(_observations()) + hostile_rows = [ + replace(row, dropped_mass=0.9) + if row.split == "heldout" + and row.target_sparsity == 0.7 + and row.consumer_head == 0 + and row.donor_head == 1 + else row + for row in _observations() + ] + + hostile = _calibrate(hostile_rows) + + assert hostile["context_policies"] == baseline["context_policies"] + assert ( + hostile["calibration_report"]["overall"]["reuse_heldout"]["constraint_violation_rate"] + == 1.0 + ) + assert hostile["calibration_report"]["overall"]["reuse_heldout"]["worst_dropped_mass"] == 0.9 + + +def test_rejects_relabelled_fixed_lambda_observation(): + rows = _observations() + rows[0] = replace( + rows[0], + threshold_lambda=math.nextafter(rows[0].threshold_lambda, math.inf), + ) + + with pytest.raises(MaskReuseCalibrationError, match="threshold_lambda does not match"): + _calibrate(rows) + + +def test_rejects_inexact_log2_launch_argument(): + rows = _observations() + rows[0] = replace( + rows[0], + threshold_log2=math.nextafter(rows[0].threshold_log2, math.inf), + ) + + with pytest.raises(MaskReuseCalibrationError, match="threshold_log2 does not match"): + _calibrate(rows) + + +def test_anchor_gate_includes_anchor_without_consumer_layer(): + rows = [ + replace( + row, + anchor_stats_by_layer={ + **row.anchor_stats_by_layer, + 2: AnchorLayerStats((30, 30), (0.2, 0.2)), + }, + ) + if row.target_sparsity == 0.7 + else row + for row in _observations() + ] + + artifact = _calibrate(rows) + + assert artifact["context_policies"][0]["target_sparsity"] == 0.5 + + +def test_rejects_inconsistent_repeated_anchor_payload(): + rows = _observations() + rows[0] = replace( + rows[0], + anchor_stats_by_layer={ + **rows[0].anchor_stats_by_layer, + 2: AnchorLayerStats((59, 60), (0.02, 0.02)), + }, + ) + + with pytest.raises(MaskReuseCalibrationError, match="anchor_stats_by_layer differs"): + _calibrate(rows) + + +def test_jsonl_rejects_duplicate_keys(): + with pytest.raises(MaskReuseCalibrationError, match="duplicate JSON key 'model'"): + parse_mask_reuse_observations(['{"model":"first","model":"second"}']) + + +def test_prefill_fit_rejects_b_above_backend_limit(): + invalid = {"prefill": {**VANILLA_FIT["prefill"], "b": 20.000_001}} + + with pytest.raises(MaskReuseCalibrationError, match=r"prefill\.b"): + canonical_prefill_threshold_scale_factor(invalid) + + +def test_canonicalizes_existing_sparse_attention_config(): + exported = { + "sparse_attention_config": { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "threshold_scale_factor": { + "formula": "a * exp(b * target_sparsity)", + "prefill": VANILLA_FIT["prefill"], + }, + } + } + } + } + + assert canonical_prefill_threshold_scale_factor(exported) == { + "formula": "a * exp(b * target_sparsity)", + "prefill": VANILLA_FIT["prefill"], + } diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py new file mode 100644 index 00000000000..c22e0d2c1c2 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the strict ModelOpt mask-reuse capture contract.""" + +import json +import math +from dataclasses import replace +from hashlib import sha256 + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.plugins.mask_reuse_capture import ( + CaptureContractError, + PromptSpec, + build_capture_invocation, + canonical_json_sha256, + load_prompt_specs, + merge_rank_captures, + source_capture_sha256, + validate_begin_acks, + validate_capture_statuses, +) + +_CHECKPOINT = sha256(b"checkpoint-manifest").hexdigest() +_GROUP = sha256(b"source-group-0").hexdigest() + + +def _fit(): + return { + "formula": "a * exp(b * target_sparsity)", + "prefill": { + "a": 1.0, + "b": 1.0, + "min_observed_sparsity": 0.4, + "max_observed_sparsity": 0.8, + }, + } + + +def _prompt(split="calibration", prompt_id="p0"): + partition = "development" if split == "calibration" else "outer_test" + return PromptSpec( + split=split, + partition=partition, + inner_fold=0 if partition == "development" else None, + prompt_id=prompt_id, + source="ruler/niah", + source_group_sha256=_GROUP, + prompt="prompt", + min_kv_tokens=8192, + max_kv_tokens=16384, + ) + + +def _invocation(): + return build_capture_invocation( + model="test-model", + checkpoint_manifest_sha256=_CHECKPOINT, + prompt=_prompt(), + prompt_token_ids=list(range(8448)), + target_sparsity=0.7, + threshold_scale_factor=_fit(), + ) + + +def _rank_payload(invocation, rank): + start = rank * 2 + return { + "capture_schema_version": 2, + "rank": rank, + "world_size": 2, + "invocation": invocation, + "invocation_sha256": canonical_json_sha256(invocation), + "geometry": invocation["expected_geometry"], + "global_num_heads": 4, + "eligible_tiles": 131, + "anchor_stats_by_layer": { + "0": { + "retained_tiles": [10, 11, 12, 13], + "dropped_mass": [0.01, 0.02, 0.03, 0.04], + }, + "4": { + "retained_tiles": [14, 15, 16, 17], + "dropped_mass": [0.05, 0.06, 0.07, 0.08], + }, + }, + "consumer_layers": { + "1": { + "anchor_layer": 0, + "consumer_head_start": start, + "dropped_mass": [ + [0.01 * (start + local + donor + 1) for donor in range(4)] for local in range(2) + ], + }, + "5": { + "anchor_layer": 4, + "consumer_head_start": start, + "dropped_mass": [ + [0.01 * (start + local + donor + 2) for donor in range(4)] for local in range(2) + ], + }, + }, + } + + +def test_build_invocation_binds_exact_threshold_source_and_final_chunk(): + invocation = _invocation() + expected_log2 = math.log2(1.0) + 0.7 * math.log2(math.e) - math.log2(8448) + + assert invocation["target_sparsity_hex"] == (0.7).hex() + assert invocation["checkpoint_manifest_sha256"] == _CHECKPOINT + assert invocation["source_group_sha256"] == _GROUP + assert invocation["partition"] == "development" + assert invocation["inner_fold"] == 0 + assert invocation["threshold_log2_hex"] == expected_log2.hex() + assert invocation["threshold_lambda_hex"] == math.exp2(expected_log2).hex() + assert invocation["expected_geometry"] == { + "q_tokens": 256, + "kv_tokens": 8448, + "q_start_tokens": 8192, + } + assert invocation["source_capture_sha256"] == source_capture_sha256(list(range(8448))) + + +def test_backend_schema_v2_invocation_golden_sha_is_byte_exact(): + invocation = build_capture_invocation( + model="toy", + checkpoint_manifest_sha256=_CHECKPOINT, + prompt=PromptSpec( + split="calibration", + partition="development", + inner_fold=0, + prompt_id="p", + source="s", + source_group_sha256=_GROUP, + prompt="unused", + min_kv_tokens=1, + max_kv_tokens=None, + ), + prompt_token_ids=list(range(33_024)), + target_sparsity=0.7, + threshold_scale_factor={ + "formula": "a * exp(b * target_sparsity)", + "prefill": { + "a": 14.47, + "b": 10.91, + "min_observed_sparsity": 0.0, + "max_observed_sparsity": 1.0, + }, + }, + ) + + assert canonical_json_sha256(invocation) == ( + "d1bb38b0611b7a424f70e4567f50380ad8fef9f77ebf05df97643fea08367056" + ) + + +def test_source_fingerprint_does_not_hide_split_overlap(): + token_ids = [1, 2, 3] + assert source_capture_sha256(token_ids) == source_capture_sha256(token_ids) + assert source_capture_sha256(token_ids) != source_capture_sha256([1, 2, 4]) + + +def test_prompt_plan_requires_unique_ids_within_each_split(tmp_path): + path = tmp_path / "prompts.jsonl" + rows = [ + { + "split": split, + "partition": "development" if split == "calibration" else "outer_test", + "inner_fold": 0 if split == "calibration" else None, + "prompt_id": prompt_id, + "source": "dataset", + "source_group_sha256": sha256(f"{split}-{prompt_id}".encode()).hexdigest(), + "prompt": text, + "min_kv_tokens": minimum, + "max_kv_tokens": maximum, + } + for split, prompt_id, text, minimum, maximum in ( + ("calibration", "same", "a", 129, 512), + ("calibration", "same", "b", 513, 1024), + ("heldout", "held-0", "c", 129, 512), + ("heldout", "held-1", "d", 513, 1024), + ) + ] + path.write_text("\n".join(json.dumps(row) for row in rows) + "\n", encoding="utf-8") + + with pytest.raises(CaptureContractError, match="unique within"): + load_prompt_specs(path) + + +def test_build_invocation_rejects_unqualified_final_chunk_and_extrapolated_target(): + with pytest.raises(CaptureContractError, match="at least 129"): + build_capture_invocation( + model="test-model", + checkpoint_manifest_sha256=_CHECKPOINT, + prompt=replace(_prompt(), max_kv_tokens=9000), + prompt_token_ids=list(range(8200)), + target_sparsity=0.7, + threshold_scale_factor=_fit(), + ) + with pytest.raises(CaptureContractError, match="outside the observed"): + build_capture_invocation( + model="test-model", + checkpoint_manifest_sha256=_CHECKPOINT, + prompt=_prompt(), + prompt_token_ids=list(range(8448)), + target_sparsity=0.9, + threshold_scale_factor=_fit(), + ) + + +def test_status_and_begin_require_every_rank_and_exact_invocation(): + statuses = [ + { + "capture_schema_version": 2, + "available": True, + "rank": rank, + "world_size": 2, + "reason": None, + } + for rank in range(2) + ] + assert validate_capture_statuses(statuses) == 2 + invocation = _invocation() + digest = canonical_json_sha256(invocation) + acknowledgements = [ + { + "capture_schema_version": 2, + "armed": True, + "rank": rank, + "world_size": 2, + "invocation_sha256": digest, + } + for rank in range(2) + ] + assert validate_begin_acks(acknowledgements, invocation) == 2 + + statuses[1]["available"] = False + statuses[1]["reason"] = "sink disabled" + with pytest.raises(CaptureContractError, match="sink disabled"): + validate_capture_statuses(statuses) + with pytest.raises(CaptureContractError, match="exactly cover ranks"): + validate_begin_acks(acknowledgements[:1], invocation) + + +def test_merge_rank_captures_concatenates_consumer_shards_deterministically(): + invocation = _invocation() + merged = merge_rank_captures( + [_rank_payload(invocation, 1), _rank_payload(invocation, 0)], invocation + ) + + assert len(merged.capture["consumer_layers"]) == 2 + assert merged.manifest["world_size"] == 2 + assert merged.manifest["global_num_heads"] == 4 + assert merged.manifest["candidate_cell_count"] == 2 * 4 * 4 + consumer = merged.capture["consumer_layers"]["1"] + assert consumer["anchor_layer"] == 0 + assert consumer["dropped_mass"][2][3] == pytest.approx(0.06) + assert set(merged.capture["anchor_stats_by_layer"]) == {"0", "4"} + + +@pytest.mark.parametrize("corruption", ["wrong_invocation", "anchor_disagreement", "head_gap"]) +def test_merge_rank_captures_fails_closed_on_incomplete_or_disagreed_evidence(corruption): + invocation = _invocation() + rank0 = _rank_payload(invocation, 0) + rank1 = _rank_payload(invocation, 1) + if corruption == "wrong_invocation": + rank1["invocation"] = dict(invocation, prompt_id="other") + elif corruption == "anchor_disagreement": + rank1["anchor_stats_by_layer"]["0"]["retained_tiles"][0] = 9 + else: + rank1["consumer_layers"]["1"]["consumer_head_start"] = 3 + + with pytest.raises(CaptureContractError): + merge_rank_captures([rank0, rank1], invocation) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py new file mode 100644 index 00000000000..627e5e5e0a4 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_compact_calibration.py @@ -0,0 +1,394 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for streaming compact-capture policy selection.""" + +import inspect +import json +import math +from hashlib import sha256 + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration import mask_reuse_compact +from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( + create_checkpoint_manifest, +) +from modelopt.torch.sparsity.attention_sparsity.calibration.mask_reuse import ( + MaskReuseCalibrationError, + calibrate_mask_reuse_policy, +) +from modelopt.torch.sparsity.attention_sparsity.calibration.mask_reuse_compact import ( + calibrate_compact_mask_reuse_policy, + load_compact_mask_reuse_captures, +) + +_FIT = { + "threshold_scale_factor": { + "formula": "a * exp(b * target_sparsity)", + "prefill": { + "a": 1.0, + "b": 1.0, + "min_observed_sparsity": 0.4, + "max_observed_sparsity": 0.8, + }, + } +} +_TOPOLOGY = {"anchors": [0], "nearest": {"0": 0, "1": 0}} + + +def _checkpoint(tmp_path): + root = tmp_path / "checkpoint" + root.mkdir() + (root / "config.json").write_text("{}\n", encoding="utf-8") + (root / "model.safetensors").write_bytes(b"weights") + return create_checkpoint_manifest(root, model="test-model") + + +def _capture(split, prompt_id, target, checkpoint_sha256): + threshold_log2 = math.log2(1.0) + target * math.log2(math.e) - math.log2(256) + if target == 0.5: + retained = [2, 3] + anchor_dropped = [0.005, 0.005] + matrix = [[0.001, 0.001], [0.001, 0.001]] + else: + retained = [1, 2] + anchor_dropped = [0.02, 0.02] + matrix = ( + [[0.01, 0.03], [0.04, 0.01]] + if split == "calibration" + else [[0.015, 0.05], [0.05, 0.015]] + ) + invocation = { + "capture_schema_version": 2, + "model": "test-model", + "checkpoint_manifest_sha256": checkpoint_sha256, + "split": split, + "partition": "development" if split == "calibration" else "outer_test", + "inner_fold": 0 if split == "calibration" else None, + "prompt_id": prompt_id, + "source": f"dataset/{split}", + "source_group_sha256": sha256(f"group/{prompt_id}".encode()).hexdigest(), + "source_capture_sha256": sha256(prompt_id.encode()).hexdigest(), + "min_kv_tokens": 129, + "max_kv_tokens": 512, + "target_sparsity_hex": target.hex(), + "sample_length": 256, + "threshold_log2_hex": threshold_log2.hex(), + "threshold_lambda_hex": math.exp2(threshold_log2).hex(), + "expected_geometry": {"q_tokens": 256, "kv_tokens": 256, "q_start_tokens": 0}, + } + return { + "compact_capture_schema_version": 1, + "invocation": invocation, + "geometry": invocation["expected_geometry"], + "global_num_heads": 2, + "eligible_tiles": 3, + "anchor_stats_by_layer": { + "0": {"retained_tiles": retained, "dropped_mass": anchor_dropped} + }, + "consumer_layers": {"1": {"anchor_layer": 0, "dropped_mass": matrix}}, + } + + +def _write_captures(path, checkpoint_sha256): + captures = [ + _capture(split, prompt, target, checkpoint_sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + payload = b"".join( + ( + json.dumps(capture, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + for capture in captures + ) + path.write_bytes(payload) + return captures, sha256(payload).hexdigest() + + +def _write_payload(path, captures): + payload = b"".join( + ( + json.dumps(capture, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + for capture in captures + ) + path.write_bytes(payload) + return sha256(payload).hexdigest() + + +def _evidence(reuse_bundle_sha256): + fields = ( + "calibration_plan_sha256", + "family_registry_sha256", + "vanilla_fit_sha256", + "reuse_bundle_sha256", + "grouped_fit_sha256", + "outer_report_sha256", + ) + return { + field: reuse_bundle_sha256 + if field == "reuse_bundle_sha256" + else sha256(field.encode()).hexdigest() + for field in fields + } + + +def _expanded_rows(captures): + rows = [] + for capture in captures: + invocation = capture["invocation"] + anchors = capture["anchor_stats_by_layer"] + anchor = anchors["0"] + matrix = capture["consumer_layers"]["1"]["dropped_mass"] + for consumer_head in range(2): + rows.extend( + { + "model": invocation["model"], + "min_kv_tokens": invocation["min_kv_tokens"], + "max_kv_tokens": invocation["max_kv_tokens"], + "target_sparsity": float.fromhex(invocation["target_sparsity_hex"]), + "sample_length": invocation["sample_length"], + "threshold_lambda": float.fromhex(invocation["threshold_lambda_hex"]), + "threshold_log2": float.fromhex(invocation["threshold_log2_hex"]), + "q_tokens": 256, + "kv_tokens": 256, + "q_start_tokens": 0, + "split": invocation["split"], + "prompt_id": invocation["prompt_id"], + "source_capture_sha256": invocation["source_capture_sha256"], + "anchor_layer": 0, + "consumer_layer": 1, + "consumer_head": consumer_head, + "donor_head": donor_head, + "retained_tiles": anchor["retained_tiles"][donor_head], + "eligible_tiles": 3, + "anchor_dropped_mass": anchor["dropped_mass"][donor_head], + "anchor_stats_by_layer": anchors, + "dropped_mass": matrix[consumer_head][donor_head], + } + for donor_head in range(2) + ) + return rows + + +def test_streaming_compact_selector_matches_legacy_row_policy(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures, digest = _write_captures(path, checkpoint.sha256) + kwargs = { + "vanilla_calibration": _FIT, + "topology": _TOPOLOGY, + "checkpoint_manifest": checkpoint, + "evidence": _evidence(digest), + "max_anchor_dropped_mass": 0.03, + "max_reuse_dropped_mass": 0.025, + "max_reuse_selection_dropped_mass": 0.02, + } + + compact = calibrate_compact_mask_reuse_policy(load_compact_mask_reuse_captures(path), **kwargs) + legacy = calibrate_mask_reuse_policy(_expanded_rows(captures), **kwargs) + + assert compact["context_policies"] == legacy["context_policies"] + assert compact["context_policies"][0]["target_sparsity"] == 0.7 + assert compact["context_policies"][0]["headmaps"] == {"1": [0, 1]} + compact_report = dict(compact["calibration_report"]) + legacy_report = dict(legacy["calibration_report"]) + compact_report.pop("promotion") + legacy_report.pop("promotion") + assert compact_report == legacy_report + assert compact["provenance"]["input_capture_count"] == 4 + assert compact["provenance"]["candidate_cell_count"] == 16 + assert compact["provenance"]["streaming_passes"] == [ + "validation", + "calibration_selection", + "frozen_evaluation", + ] + assert compact["promotion_status"] == "candidate_only" + assert compact["deployment_geometry_validated"] is False + + +def test_compact_selector_binds_reuse_bundle_sha(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + _, digest = _write_captures(path, checkpoint.sha256) + evidence = _evidence(digest) + evidence["reuse_bundle_sha256"] = sha256(b"wrong").hexdigest() + + with pytest.raises(MaskReuseCalibrationError, match="does not match"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=evidence, + max_anchor_dropped_mass=0.03, + max_reuse_dropped_mass=0.025, + ) + + +def test_compact_selector_binds_verified_checkpoint_and_disjoint_groups(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures = [ + _capture(split, prompt, target, checkpoint.sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + captures[0]["invocation"]["checkpoint_manifest_sha256"] = "0" * 64 + digest = _write_payload(path, captures) + with pytest.raises(MaskReuseCalibrationError, match="one model, checkpoint"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + max_reuse_dropped_mass=0.025, + ) + + shared_group = captures[1]["invocation"]["source_group_sha256"] + for capture in captures: + capture["invocation"]["checkpoint_manifest_sha256"] = checkpoint.sha256 + capture["invocation"]["source_group_sha256"] = shared_group + digest = _write_payload(path, captures) + with pytest.raises(MaskReuseCalibrationError, match="multiple partitions"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + max_reuse_dropped_mass=0.025, + ) + + +def test_compact_selector_rejects_file_changed_during_evaluation(tmp_path, monkeypatch): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + _, digest = _write_captures(path, checkpoint.sha256) + real_evaluation_pass = mask_reuse_compact._evaluation_pass + + def mutate_after_evaluation(*args, **kwargs): + result = real_evaluation_pass(*args, **kwargs) + with path.open("ab") as handle: + handle.write(b"\n") + return result + + monkeypatch.setattr(mask_reuse_compact, "_evaluation_pass", mutate_after_evaluation) + + with pytest.raises(MaskReuseCalibrationError, match="changed during calibration"): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + max_reuse_dropped_mass=0.025, + ) + + +def test_selector_minimizes_declared_equal_bmm_combined_cost(tmp_path): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures = [ + _capture(split, prompt, target, checkpoint.sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + for capture in captures: + target = float.fromhex(capture["invocation"]["target_sparsity_hex"]) + if target == 0.5: + capture["anchor_stats_by_layer"]["0"] = { + "retained_tiles": [1, 3], + "dropped_mass": [0.001, 0.001], + } + capture["anchor_stats_by_layer"]["2"] = { + "retained_tiles": [3, 3], + "dropped_mass": [0.001, 0.001], + } + capture["consumer_layers"]["1"]["dropped_mass"] = [ + [0.001, 0.001], + [0.001, 0.001], + ] + else: + capture["anchor_stats_by_layer"]["0"] = { + "retained_tiles": [0, 2], + "dropped_mass": [0.002, 0.002], + } + capture["anchor_stats_by_layer"]["2"] = { + "retained_tiles": [0, 0], + "dropped_mass": [0.002, 0.002], + } + capture["consumer_layers"]["1"]["dropped_mass"] = [ + [0.03, 0.01], + [0.03, 0.01], + ] + digest = _write_payload(path, captures) + + candidate = calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology={"anchors": [0, 2], "nearest": {"0": 0, "1": 0, "2": 2}}, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + max_reuse_dropped_mass=0.025, + max_reuse_selection_dropped_mass=0.02, + ) + + assert candidate["context_policies"][0]["target_sparsity"] == 0.7 + frontier = candidate["calibration_report"]["by_bucket"][0]["target_sparsity_frontier"] + assert [row["combined_tile_cost"] for row in frontier] == [14, 10] + + +def test_validation_pass_does_not_retain_capture_objects(): + implementation = inspect.getsource(mask_reuse_compact._validate_dataset) + + assert ".append(capture)" not in implementation + assert "list[CompactMaskReuseCapture]" not in implementation + + +@pytest.mark.parametrize("corruption", ["eligible", "retained_monotonic", "mass_monotonic"]) +def test_compact_validation_rejects_impossible_geometry_or_sparsity_trend(tmp_path, corruption): + path = tmp_path / "compact.jsonl" + checkpoint = _checkpoint(tmp_path) + captures = [ + _capture(split, prompt, target, checkpoint.sha256) + for split, prompt in (("calibration", "cal-0"), ("heldout", "held-0")) + for target in (0.5, 0.7) + ] + if corruption == "eligible": + captures[0]["eligible_tiles"] = 4 + elif corruption == "retained_monotonic": + captures[1]["anchor_stats_by_layer"]["0"]["retained_tiles"][0] = 3 + else: + captures[1]["consumer_layers"]["1"]["dropped_mass"][0][0] = 0.0 + digest = _write_payload(path, captures) + + with pytest.raises(MaskReuseCalibrationError): + calibrate_compact_mask_reuse_policy( + path, + vanilla_calibration=_FIT, + topology=_TOPOLOGY, + checkpoint_manifest=checkpoint, + evidence=_evidence(digest), + max_anchor_dropped_mass=0.03, + max_reuse_dropped_mass=0.025, + ) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py index e85ad3ab77e..00c5beec5ec 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py @@ -166,13 +166,26 @@ def test_logspace_fit_preserves_log_a(self): class TestBuildSparseAttentionConfig: - _PARAMS = {"prefill": {"a": 7.9, "b": 8.6}, "decode": {"a": 0.12, "b": 9.8}} + _PARAMS = { + "prefill": { + "a": 7.9, + "b": 8.6, + "min_observed_sparsity": 0.1, + "max_observed_sparsity": 0.9, + }, + "decode": {"a": 0.12, "b": 9.8}, + } def test_canonical_schema(self): config = build_sparse_attention_config(self._PARAMS, 0.4) group = config["config_groups"]["group_0"] assert group["algorithm"] == "skip_softmax" - assert group["threshold_scale_factor"]["prefill"] == {"a": 7.9, "b": 8.6} + assert group["threshold_scale_factor"]["prefill"] == { + "a": 7.9, + "b": 8.6, + "min_observed_sparsity": 0.1, + "max_observed_sparsity": 0.9, + } assert group["threshold_scale_factor"]["formula"] == "a * exp(b * target_sparsity)" assert group["target_sparsity"] == {"prefill": 0.4, "decode": 0.4} assert config["producer"]["name"] == "modelopt" diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py b/tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py new file mode 100644 index 00000000000..4f57005f752 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the policy-free vLLM capture worker bootstrap and RPC wrappers.""" + +from types import SimpleNamespace + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm_mask_reuse_capture + + +def _api(calls): + return SimpleNamespace( + configure_capture_runtime=lambda plan: calls.append(("configure", plan)), + capture_status=lambda: { + "capture_schema_version": 1, + "available": True, + "rank": 0, + "world_size": 1, + "reason": None, + }, + begin_capture=lambda invocation: calls.append(("begin", invocation)) or {"armed": True}, + drain_capture=lambda: calls.append(("drain",)) or {"records": []}, + ) + + +def test_bootstrap_requires_gate_and_plan(monkeypatch): + monkeypatch.delenv(vllm_mask_reuse_capture.CAPTURE_ENV, raising=False) + monkeypatch.setenv(vllm_mask_reuse_capture.PLAN_ENV, "qwen3_stride2") + with pytest.raises(RuntimeError, match="CALIBRATION_CAPTURE=1"): + vllm_mask_reuse_capture._configure_capture_before_model_load() + + monkeypatch.setenv(vllm_mask_reuse_capture.CAPTURE_ENV, "1") + monkeypatch.delenv(vllm_mask_reuse_capture.PLAN_ENV, raising=False) + with pytest.raises(RuntimeError, match="must name"): + vllm_mask_reuse_capture._configure_capture_before_model_load() + + +def test_bootstrap_installs_policy_free_runtime_before_load(monkeypatch): + calls = [] + api = _api(calls) + monkeypatch.setenv(vllm_mask_reuse_capture.CAPTURE_ENV, "1") + monkeypatch.setenv(vllm_mask_reuse_capture.PLAN_ENV, "nemotron3_ultra_stride2") + monkeypatch.setattr(vllm_mask_reuse_capture, "_capture_api", lambda: api) + + assert vllm_mask_reuse_capture._configure_capture_before_model_load() is api + assert calls == [("configure", "nemotron3_ultra_stride2")] + + +def test_worker_rpc_methods_forward_only_to_capture_api(monkeypatch): + calls = [] + api = _api(calls) + monkeypatch.setattr(vllm_mask_reuse_capture, "_capture_api", lambda: api) + worker = object.__new__(vllm_mask_reuse_capture.MaskReuseCaptureWorker) + invocation = {"capture_schema_version": 1} + + assert worker.mask_reuse_capture_status()["available"] is True + assert worker.mask_reuse_capture_begin(invocation) == {"armed": True} + assert worker.mask_reuse_capture_drain() == {"records": []} + assert calls == [("begin", invocation), ("drain",)] From 7ec4ad978feb44846e10c8f1a47e0a1c71321750 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 1 Aug 2026 07:47:03 -0700 Subject: [PATCH 10/16] Add live mask reuse capture evidence Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 10 +- examples/vllm_serve/calibrate_mask_reuse.py | 10 +- examples/vllm_serve/collect_mask_reuse.py | 27 +++- .../plugins/mask_reuse_capture.py | 149 +++++++++++++++++- .../test_calibrate_mask_reuse_cli.py | 6 +- .../test_collect_mask_reuse_cli.py | 26 ++- .../test_mask_reuse_capture.py | 51 +++++- 7 files changed, 260 insertions(+), 19 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index ba7c22ee968..3830da94362 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -165,12 +165,13 @@ sparsities under the exact runtime rule python collect_mask_reuse.py \ --model-id Nemotron-3-Ultra \ --plan nemotron3_ultra_stride2 \ - --fa4-source /path/to/flash-attention-at-4c40766b \ + --fa4-source /path/to/pinned-calibrated-threshold-flash-attention \ --prompts-jsonl mask_reuse_prompts.jsonl \ --vanilla-config /config.json \ --target-sparsities 0.5 0.6 0.7 \ --max-model-len 1000001 \ --tensor-parallel-size 8 \ + --validate-dense-output \ --output mask_reuse_compact_captures.jsonl ``` @@ -193,7 +194,12 @@ dropped-mass matrix once. It does not expand or repeat the matrix as millions of candidate rows. Fixed-lambda version-4/version-5 captures are not valid inputs to this target-sparsity path. Use the printed compact-capture SHA256 as the `reuse_bundle_sha256` evidence value. The FA4 source must be a clean Git -checkout; its exact commit is recorded in the capture manifest. +checkout; its exact commit is recorded in the capture manifest. The version-3 +capture manifest also retains the canonical vLLM engine arguments, rank-major +TP sentinel evidence, exact prefill/decode attention-call counts, and optional +bitwise dense-shadow coverage. `--validate-dense-output` runs a second pinned +dense FA4 call for every armed final-chunk layer and fails before publication +if its BF16 output differs bitwise from the model-trajectory output. Once those observations have been captured, select the per-context target, donor-head map, and exact fallbacks and export the fail-closed candidate: diff --git a/examples/vllm_serve/calibrate_mask_reuse.py b/examples/vllm_serve/calibrate_mask_reuse.py index 65bdffe84c0..72480908ffe 100644 --- a/examples/vllm_serve/calibrate_mask_reuse.py +++ b/examples/vllm_serve/calibrate_mask_reuse.py @@ -79,6 +79,8 @@ "plan", "fa4_source", "fa4_source_commit", + "engine_kwargs", + "dense_shadow_validation_requested", "target_sparsity_hex", "vanilla_threshold_scale_factor", "vanilla_fit_sha256", @@ -263,8 +265,8 @@ def _validate_capture_manifest( f"missing={sorted(missing)}, extra={sorted(extra)}" ) expected = { - "capture_manifest_schema_version": 2, - "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v2", + "capture_manifest_schema_version": 3, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v3", "model": model, "checkpoint_manifest_sha256": checkpoint_sha256, "compact_capture_file_sha256": compact_capture_sha256, @@ -273,6 +275,10 @@ def _validate_capture_manifest( for field, value in expected.items(): if raw[field] != value: raise ValueError(f"capture manifest {field} does not match its verified input") + 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: diff --git a/examples/vllm_serve/collect_mask_reuse.py b/examples/vllm_serve/collect_mask_reuse.py index a86390d93a7..76aaf76ef7a 100644 --- a/examples/vllm_serve/collect_mask_reuse.py +++ b/examples/vllm_serve/collect_mask_reuse.py @@ -70,6 +70,7 @@ CAPTURE_ENV = "MASK_REUSE_FA4_CALIBRATION_CAPTURE" PLAN_ENV = "MASK_REUSE_FA4_PLAN" CHECKPOINT_ENV = "MASK_REUSE_FA4_CHECKPOINT_MANIFEST_SHA256" +DENSE_SHADOW_ENV = "MASK_REUSE_FA4_CAPTURE_DENSE_SHADOW" _POLICY_ENVS = ( "MASK_REUSE_FA4_POLICY", "MASK_REUSE_FA4_POLICY_SHA256", @@ -133,6 +134,11 @@ def _parser() -> argparse.ArgumentParser: parser.add_argument("--tensor-parallel-size", type=int, default=1) parser.add_argument("--gpu-memory-utilization", type=float, default=None) parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument( + "--validate-dense-output", + action="store_true", + help="Bitwise-compare every armed capture layer with a second pinned dense FA4 call", + ) parser.add_argument( "--engine-kwargs", default=None, @@ -202,7 +208,11 @@ def _canonical_capture_line(capture: dict[str, object]) -> bytes: def _configure_capture_environment( - plan: str, fa4_source: str, checkpoint_manifest_sha256: str + plan: str, + fa4_source: str, + checkpoint_manifest_sha256: str, + *, + validate_dense_output: bool, ) -> tuple[Path, str]: source = Path(fa4_source).expanduser().resolve() required = ( @@ -246,6 +256,7 @@ def _configure_capture_environment( os.environ[CAPTURE_ENV] = "1" os.environ[PLAN_ENV] = plan os.environ[CHECKPOINT_ENV] = checkpoint_manifest_sha256 + os.environ[DENSE_SHADOW_ENV] = "1" if validate_dense_output else "0" os.environ["PYTHONDONTWRITEBYTECODE"] = "1" sys.dont_write_bytecode = True os.environ["MASK_REUSE_FA4_SOURCE"] = str(source) @@ -316,13 +327,17 @@ def run(args: argparse.Namespace) -> tuple[Path, Path]: manifest_path.parent.mkdir(parents=True, exist_ok=True) fa4_source, fa4_source_commit = _configure_capture_environment( - args.plan, args.fa4_source, checkpoint.sha256 + args.plan, + args.fa4_source, + checkpoint.sha256, + validate_dense_output=args.validate_dense_output, ) # Import after setting the gate: worker subprocesses inherit the exact # capture environment and never enter policy-backed serving mode. from vllm import LLM, SamplingParams - llm = LLM(**_engine_kwargs(args)) + engine_kwargs = _engine_kwargs(args) + llm = LLM(**engine_kwargs) loaded_checkpoint = verify_checkpoint_manifest(args.model, expected_model=model_id) if loaded_checkpoint != checkpoint: raise CaptureContractError( @@ -411,8 +426,8 @@ def run(args: argparse.Namespace) -> tuple[Path, Path]: capture_sha256 = capture_digest.hexdigest() manifest = { - "capture_manifest_schema_version": 2, - "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v2", + "capture_manifest_schema_version": 3, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v3", "model": model_id, "checkpoint_manifest_sha256": checkpoint.sha256, "checkpoint_manifest_path": str(checkpoint.manifest_path), @@ -421,6 +436,8 @@ def run(args: argparse.Namespace) -> tuple[Path, Path]: "plan": args.plan, "fa4_source": str(fa4_source), "fa4_source_commit": fa4_source_commit, + "engine_kwargs": engine_kwargs, + "dense_shadow_validation_requested": args.validate_dense_output, "target_sparsity_hex": [target.hex() for target in targets], "vanilla_threshold_scale_factor": threshold_scale_factor, "vanilla_fit_sha256": canonical_json_sha256(threshold_scale_factor), diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py b/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py index 1a1160af22c..6dd5bcdd134 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py @@ -117,10 +117,24 @@ "eligible_tiles", "anchor_stats_by_layer", "consumer_layers", + "attention_call_counts", + "tp_head_order_evidence", + "dense_shadow_evidence", } ) _ANCHOR_STATS_FIELDS = frozenset({"retained_tiles", "dropped_mass"}) _CONSUMER_STATS_FIELDS = frozenset({"anchor_layer", "consumer_head_start", "dropped_mass"}) +_ATTENTION_CALL_COUNT_FIELDS = frozenset({"prefill", "decode"}) +_TP_HEAD_ORDER_FIELDS = frozenset( + { + "sentinel_device_type", + "gather_dim", + "local_rank", + "local_num_heads", + "gathered_rank_local_head", + } +) +_DENSE_SHADOW_FIELDS = frozenset({"enabled", "atol_hex", "rtol_hex", "validated_layer_indices"}) class CaptureContractError(ValueError): @@ -684,6 +698,87 @@ def _parse_consumer_layers( return result +def _parse_attention_call_counts(raw: object, *, rank: int) -> dict[str, int]: + if not isinstance(raw, Mapping): + raise CaptureContractError(f"rank {rank} attention_call_counts must be an object") + _exact_fields(raw, _ATTENTION_CALL_COUNT_FIELDS, f"rank {rank} attention_call_counts") + return { + name: _integer(raw[name], f"rank {rank} attention_call_counts.{name}") + for name in ("prefill", "decode") + } + + +def _parse_tp_head_order_evidence( + raw: object, + *, + rank: int, + world_size: int, + global_num_heads: int, +) -> dict[str, object]: + if not isinstance(raw, Mapping): + raise CaptureContractError(f"rank {rank} tp_head_order_evidence must be an object") + _exact_fields(raw, _TP_HEAD_ORDER_FIELDS, f"rank {rank} tp_head_order_evidence") + if raw["sentinel_device_type"] != "cuda": + raise CaptureContractError(f"rank {rank} TP sentinel was not gathered on CUDA") + if raw["gather_dim"] != 0: + raise CaptureContractError(f"rank {rank} TP sentinel must be gathered along dim 0") + if _integer(raw["local_rank"], f"rank {rank} TP local_rank") != rank: + raise CaptureContractError(f"rank {rank} TP sentinel reports a different local rank") + local_num_heads = _integer(raw["local_num_heads"], f"rank {rank} TP local_num_heads", minimum=1) + if local_num_heads * world_size != global_num_heads: + raise CaptureContractError( + f"rank {rank} local head count does not evenly cover global heads" + ) + gathered = raw["gathered_rank_local_head"] + expected = [ + [global_head // local_num_heads, global_head % local_num_heads] + for global_head in range(global_num_heads) + ] + if gathered != expected: + raise CaptureContractError( + f"rank {rank} TP all-gather is not rank-major in global-head order" + ) + return { + "rank": rank, + "global_head_start": rank * local_num_heads, + "local_num_heads": local_num_heads, + "sentinel_device_type": "cuda", + "gather_dim": 0, + "gathered_rank_local_head": expected, + } + + +def _parse_dense_shadow_evidence(raw: object, *, rank: int) -> dict[str, object]: + if not isinstance(raw, Mapping): + raise CaptureContractError(f"rank {rank} dense_shadow_evidence must be an object") + _exact_fields(raw, _DENSE_SHADOW_FIELDS, f"rank {rank} dense_shadow_evidence") + enabled = raw["enabled"] + if not isinstance(enabled, bool): + raise CaptureContractError(f"rank {rank} dense_shadow_evidence.enabled must be boolean") + atol = _canonical_float_hex(raw["atol_hex"], f"rank {rank} dense shadow atol") + rtol = _canonical_float_hex(raw["rtol_hex"], f"rank {rank} dense shadow rtol") + if atol != 0.0 or rtol != 0.0: + raise CaptureContractError("dense shadow evidence must use bitwise zero tolerances") + raw_layers = raw["validated_layer_indices"] + if not isinstance(raw_layers, list): + raise CaptureContractError( + f"rank {rank} dense_shadow_evidence.validated_layer_indices must be a list" + ) + layers = [_integer(layer, f"rank {rank} dense shadow layer") for layer in raw_layers] + if layers != sorted(set(layers)): + raise CaptureContractError(f"rank {rank} dense shadow layers must be sorted and unique") + if not enabled and layers: + raise CaptureContractError( + f"rank {rank} disabled dense shadow cannot report validated layers" + ) + return { + "enabled": enabled, + "atol_hex": atol.hex(), + "rtol_hex": rtol.hex(), + "validated_layer_indices": layers, + } + + @dataclass(frozen=True, slots=True) class MergedCapture: """One compact normalized capture plus auditable metadata.""" @@ -704,6 +799,9 @@ def merge_rank_captures( geometry_values: list[dict[str, int]] = [] anchor_payloads: list[dict[str, dict[str, list[int] | list[float]]]] = [] consumer_payloads: list[dict[int, tuple[int, int, list[list[float]]]]] = [] + attention_call_counts: list[dict[str, int]] = [] + tp_head_order_evidence: list[dict[str, object]] = [] + dense_shadow_evidence: list[dict[str, object]] = [] for value in ordered: rank = cast("int", value["rank"]) if value["invocation"] != trusted_invocation: @@ -729,17 +827,40 @@ def merge_rank_captures( eligible_tiles=eligible_tiles, ) ) - consumer_payloads.append( - _parse_consumer_layers( - value["consumer_layers"], rank=rank, global_num_heads=global_num_heads - ) + consumers = _parse_consumer_layers( + value["consumer_layers"], rank=rank, global_num_heads=global_num_heads + ) + consumer_payloads.append(consumers) + counts = _parse_attention_call_counts(value["attention_call_counts"], rank=rank) + attention_call_counts.append(counts) + tp_evidence = _parse_tp_head_order_evidence( + value["tp_head_order_evidence"], + rank=rank, + world_size=world_size, + global_num_heads=global_num_heads, + ) + tp_head_order_evidence.append(tp_evidence) + dense_shadow_evidence.append( + _parse_dense_shadow_evidence(value["dense_shadow_evidence"], rank=rank) ) + for layer, (_, start, rows) in consumers.items(): + if ( + start != tp_evidence["global_head_start"] + or len(rows) != tp_evidence["local_num_heads"] + ): + raise CaptureContractError( + f"rank {rank} consumer layer {layer} does not match its TP head shard" + ) if len(global_heads) != 1 or len(eligible_values) != 1: raise CaptureContractError("capture ranks disagree on head count or eligible tiles") if any(value != geometry_values[0] for value in geometry_values[1:]): raise CaptureContractError("capture ranks disagree on final-chunk geometry") if any(value != anchor_payloads[0] for value in anchor_payloads[1:]): raise CaptureContractError("capture ranks disagree on global anchor statistics") + if any(value != attention_call_counts[0] for value in attention_call_counts[1:]): + raise CaptureContractError("capture ranks disagree on attention call counts") + if any(value != dense_shadow_evidence[0] for value in dense_shadow_evidence[1:]): + raise CaptureContractError("capture ranks disagree on dense shadow evidence") layer_sets = [set(value) for value in consumer_payloads] if any(value != layer_sets[0] for value in layer_sets[1:]): raise CaptureContractError("capture ranks disagree on consumer layer coverage") @@ -751,6 +872,23 @@ def merge_rank_captures( "capture eligible_tiles does not match 128x128 bottom-right causal geometry" ) anchors = anchor_payloads[0] + attention_layers = sorted({int(layer) for layer in anchors} | layer_sets[0]) + expected_prefill_calls = ( + (cast("int", trusted_invocation["sample_length"]) + MAX_QUERY_CHUNK_TOKENS - 1) + // MAX_QUERY_CHUNK_TOKENS + ) * len(attention_layers) + common_counts = attention_call_counts[0] + if common_counts["prefill"] != expected_prefill_calls: + raise CaptureContractError( + "capture prefill attention call count does not match chunks times layers" + ) + if common_counts["decode"] != 0: + raise CaptureContractError("capture observed a decode attention call at max_tokens=1") + common_shadow = dense_shadow_evidence[0] + if common_shadow["enabled"] and common_shadow["validated_layer_indices"] != attention_layers: + raise CaptureContractError( + "dense shadow evidence does not cover every captured attention layer" + ) merged_consumers: dict[str, dict[str, object]] = {} for layer in sorted(layer_sets[0]): shards = sorted( @@ -804,6 +942,9 @@ def merge_rank_captures( len(cast("Sequence[object]", value["dropped_mass"])) * global_num_heads for value in merged_consumers.values() ), + "attention_call_counts": common_counts, + "tp_head_order_evidence": tp_head_order_evidence, + "dense_shadow_evidence": common_shadow, "compact_capture_sha256": canonical_json_sha256(compact_capture), } return MergedCapture(compact_capture, manifest) diff --git a/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 index a65405fa1c2..feacf8fda53 100644 --- 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 @@ -62,8 +62,8 @@ def _base_args(tmp_path: Path): capture_manifest.write_bytes( _canonical( { - "capture_manifest_schema_version": 2, - "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v2", + "capture_manifest_schema_version": 3, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v3", "model": "test-model", "checkpoint_manifest_sha256": checkpoint.sha256, "checkpoint_manifest_path": str(checkpoint.manifest_path), @@ -72,6 +72,8 @@ def _base_args(tmp_path: Path): "plan": "test_stride2", "fa4_source": "/source", "fa4_source_commit": "a" * 40, + "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(), diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py b/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py index c2f10b9d66e..6a005873f5f 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py @@ -108,6 +108,20 @@ def collective_rpc(self, method, args=()): "dropped_mass": [[0.01, 0.02], [0.03, 0.04]], } }, + "attention_call_counts": {"prefill": 2, "decode": 0}, + "tp_head_order_evidence": { + "sentinel_device_type": "cuda", + "gather_dim": 0, + "local_rank": 0, + "local_num_heads": 2, + "gathered_rank_local_head": [[0, 0], [0, 1]], + }, + "dense_shadow_evidence": { + "enabled": True, + "atol_hex": (0.0).hex(), + "rtol_hex": (0.0).hex(), + "validated_layer_indices": [0, 1], + }, } ] @@ -228,6 +242,7 @@ def parse_vanilla(payload): str(manifest), "--max-model-len", "512", + "--validate-dense-output", ] ) @@ -253,6 +268,7 @@ def parse_vanilla(payload): == checkpoint_manifest.sha256 ) assert collect_mask_reuse_cli.os.environ["PYTHONDONTWRITEBYTECODE"] == "1" + assert collect_mask_reuse_cli.os.environ["MASK_REUSE_FA4_CAPTURE_DENSE_SHADOW"] == "1" assert "MASK_REUSE_FA4_POLICY" not in collect_mask_reuse_cli.os.environ assert "MASK_REUSE_FA4_POLICY_SHA256" not in collect_mask_reuse_cli.os.environ @@ -266,11 +282,14 @@ def parse_vanilla(payload): } assert {capture["invocation"]["target_sparsity_hex"] for capture in captures} == {(0.7).hex()} report = json.loads(manifest.read_text()) - assert report["capture_protocol"] == "modelopt_vllm_mask_reuse_target_sparsity_v2" + assert report["capture_manifest_schema_version"] == 3 + assert report["capture_protocol"] == "modelopt_vllm_mask_reuse_target_sparsity_v3" assert report["checkpoint_manifest_sha256"] == checkpoint_manifest.sha256 assert report["prompt_plan_file_sha256"] == sha256(parsed_payloads["prompts"]).hexdigest() assert report["vanilla_config_file_sha256"] == sha256(parsed_payloads["vanilla"]).hexdigest() assert report["fa4_source_commit"] == "a" * 40 + assert report["dense_shadow_validation_requested"] is True + assert report["engine_kwargs"]["tensor_parallel_size"] == 1 assert report["capture_count"] == len(captures) assert "observation_count" not in report assert len(report["captures"]) == 2 @@ -340,5 +359,8 @@ def test_capture_environment_rejects_untracked_fa4_source(tmp_path, monkeypatch) match="tracked or untracked modifications", ): collect_mask_reuse_cli._configure_capture_environment( - "test_stride2", str(fa4_source), "0" * 64 + "test_stride2", + str(fa4_source), + "0" * 64, + validate_dense_output=True, ) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py index c22e0d2c1c2..ff6c904daf8 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py @@ -113,6 +113,20 @@ def _rank_payload(invocation, rank): ], }, }, + "attention_call_counts": {"prefill": 8, "decode": 0}, + "tp_head_order_evidence": { + "sentinel_device_type": "cuda", + "gather_dim": 0, + "local_rank": rank, + "local_num_heads": 2, + "gathered_rank_local_head": [[0, 0], [0, 1], [1, 0], [1, 1]], + }, + "dense_shadow_evidence": { + "enabled": True, + "atol_hex": (0.0).hex(), + "rtol_hex": (0.0).hex(), + "validated_layer_indices": [0, 1, 4, 5], + }, } @@ -266,13 +280,30 @@ def test_merge_rank_captures_concatenates_consumer_shards_deterministically(): assert merged.manifest["world_size"] == 2 assert merged.manifest["global_num_heads"] == 4 assert merged.manifest["candidate_cell_count"] == 2 * 4 * 4 + assert merged.manifest["attention_call_counts"] == {"prefill": 8, "decode": 0} + assert [item["global_head_start"] for item in merged.manifest["tp_head_order_evidence"]] == [ + 0, + 2, + ] + assert merged.manifest["dense_shadow_evidence"]["validated_layer_indices"] == [0, 1, 4, 5] consumer = merged.capture["consumer_layers"]["1"] assert consumer["anchor_layer"] == 0 assert consumer["dropped_mass"][2][3] == pytest.approx(0.06) assert set(merged.capture["anchor_stats_by_layer"]) == {"0", "4"} -@pytest.mark.parametrize("corruption", ["wrong_invocation", "anchor_disagreement", "head_gap"]) +@pytest.mark.parametrize( + "corruption", + [ + "wrong_invocation", + "anchor_disagreement", + "head_gap", + "rank_permutation", + "decode_call", + "prefill_count", + "missing_dense_shadow", + ], +) def test_merge_rank_captures_fails_closed_on_incomplete_or_disagreed_evidence(corruption): invocation = _invocation() rank0 = _rank_payload(invocation, 0) @@ -281,8 +312,24 @@ def test_merge_rank_captures_fails_closed_on_incomplete_or_disagreed_evidence(co rank1["invocation"] = dict(invocation, prompt_id="other") elif corruption == "anchor_disagreement": rank1["anchor_stats_by_layer"]["0"]["retained_tiles"][0] = 9 - else: + elif corruption == "head_gap": rank1["consumer_layers"]["1"]["consumer_head_start"] = 3 + elif corruption == "rank_permutation": + rank1["tp_head_order_evidence"]["gathered_rank_local_head"] = [ + [1, 0], + [1, 1], + [0, 0], + [0, 1], + ] + elif corruption == "decode_call": + rank0["attention_call_counts"]["decode"] = 1 + rank1["attention_call_counts"]["decode"] = 1 + elif corruption == "prefill_count": + rank0["attention_call_counts"]["prefill"] = 7 + rank1["attention_call_counts"]["prefill"] = 7 + else: + rank0["dense_shadow_evidence"]["validated_layer_indices"] = [0, 1, 4] + rank1["dense_shadow_evidence"]["validated_layer_indices"] = [0, 1, 4] with pytest.raises(CaptureContractError): merge_rank_captures([rank0, rank1], invocation) From cd39ff485f104d7380463415f5eb80e026b327bc Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 1 Aug 2026 08:29:00 -0700 Subject: [PATCH 11/16] Isolate vLLM worker unit test dependency Signed-off-by: Kai Xu --- .../test_vllm_mask_reuse_capture_worker.py | 33 ++++++++++++++++--- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py b/tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py index 4f57005f752..348dacac1a9 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_vllm_mask_reuse_capture_worker.py @@ -15,11 +15,34 @@ """Tests for the policy-free vLLM capture worker bootstrap and RPC wrappers.""" -from types import SimpleNamespace +import importlib +import sys +from types import ModuleType, SimpleNamespace import pytest -from modelopt.torch.sparsity.attention_sparsity.plugins import vllm_mask_reuse_capture + +@pytest.fixture +def vllm_mask_reuse_capture(monkeypatch): + """Import the worker against a minimal optional-vLLM boundary.""" + + class Worker: + pass + + modules = { + "vllm": ModuleType("vllm"), + "vllm.v1": ModuleType("vllm.v1"), + "vllm.v1.worker": ModuleType("vllm.v1.worker"), + "vllm.v1.worker.gpu_worker": ModuleType("vllm.v1.worker.gpu_worker"), + } + modules["vllm.v1.worker.gpu_worker"].Worker = Worker + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + target = "modelopt.torch.sparsity.attention_sparsity.plugins.vllm_mask_reuse_capture" + sys.modules.pop(target, None) + module = importlib.import_module(target) + yield module + sys.modules.pop(target, None) def _api(calls): @@ -37,7 +60,7 @@ def _api(calls): ) -def test_bootstrap_requires_gate_and_plan(monkeypatch): +def test_bootstrap_requires_gate_and_plan(monkeypatch, vllm_mask_reuse_capture): monkeypatch.delenv(vllm_mask_reuse_capture.CAPTURE_ENV, raising=False) monkeypatch.setenv(vllm_mask_reuse_capture.PLAN_ENV, "qwen3_stride2") with pytest.raises(RuntimeError, match="CALIBRATION_CAPTURE=1"): @@ -49,7 +72,7 @@ def test_bootstrap_requires_gate_and_plan(monkeypatch): vllm_mask_reuse_capture._configure_capture_before_model_load() -def test_bootstrap_installs_policy_free_runtime_before_load(monkeypatch): +def test_bootstrap_installs_policy_free_runtime_before_load(monkeypatch, vllm_mask_reuse_capture): calls = [] api = _api(calls) monkeypatch.setenv(vllm_mask_reuse_capture.CAPTURE_ENV, "1") @@ -60,7 +83,7 @@ def test_bootstrap_installs_policy_free_runtime_before_load(monkeypatch): assert calls == [("configure", "nemotron3_ultra_stride2")] -def test_worker_rpc_methods_forward_only_to_capture_api(monkeypatch): +def test_worker_rpc_methods_forward_only_to_capture_api(monkeypatch, vllm_mask_reuse_capture): calls = [] api = _api(calls) monkeypatch.setattr(vllm_mask_reuse_capture, "_capture_api", lambda: api) From 32ec2b8762b9f5d2286546bb83a3ba9a33078ac0 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 1 Aug 2026 08:49:07 -0700 Subject: [PATCH 12/16] Align skip-softmax tests with quantization guard Signed-off-by: Kai Xu --- .../common/attention/test_triton_fa_p_qdq.py | 36 +++++------ .../test_sparse_attn_worker.py | 34 ----------- .../common/attention/test_triton_fa.py | 59 +++++++++++++------ 3 files changed, 60 insertions(+), 69 deletions(-) diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index 60523d44457..555ef52f483 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -20,7 +20,7 @@ import pytest import torch -from conftest import make_qkv, make_varlen_meta, sdpa_reference +from conftest import make_qkv, make_varlen_meta from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor, e2m1_values @@ -454,29 +454,29 @@ def test_invalid_amax_raises(self): with pytest.raises(ValueError, match="p_qdq_amax"): attention(q, k, v, locs, lens, 8, p_qdq="fp8", p_qdq_amax=0.0) - @requires_native_e4m3 - def test_composes_with_skip_softmax(self): - """p_qdq composes with the skip-softmax feature in one launch.""" + @pytest.mark.parametrize( + "qdq_kwargs", + [{"p_qdq": "fp8"}, {"v_qdq": "nvfp4", "v_qdq_amax": 1.0}], + ) + def test_skip_softmax_rejects_attention_quantization(self, qdq_kwargs): + """The raw contiguous kernel rejects quantization with calibrated skip.""" seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 - scale = 1.0 / (head_dim**0.5) torch.manual_seed(19) q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.float16) locs, lens = make_varlen_meta([seq_len]) - o = attention( - q, - k, - v, - locs, - lens, - seq_len, - softmax_scale=scale, - p_qdq="fp8", - skip_softmax_threshold=1e-3, - ) - ref = sdpa_reference(q, k, v, locs, lens) - torch.testing.assert_close(o, ref, rtol=5e-2, atol=5e-2) + with pytest.raises(ValueError, match="cannot be combined with attention quantization"): + attention( + q, + k, + v, + locs, + lens, + seq_len, + skip_softmax_threshold=1e-3, + **qdq_kwargs, + ) def test_invalid_mode_raises(self): q, k, v = make_qkv(8, 2, 2, 32, dtype=torch.float16) diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 57a3a9549ca..d1c95a1a8a6 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -1030,40 +1030,6 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): assert calls["query"].dtype == torch.float32 -def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): - """Split-local maxima must not change calibrated skip-softmax semantics.""" - impl = _clone_sparse_impl(_make_old_impl()) - impl.quant_kw = { - "p_qdq": "nvfp4", - "p_qdq_amax": 1.0, - "v_qdq": "nvfp4", - "v_qdq_amax": 6.0 * 448.0, - } - impl.sparse_kw = {"skip_softmax_threshold": 0.001} - q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) - metadata = _flash_attention_metadata(1, 16) - captured = {} - - monkeypatch.setattr(vllm_plugin, "fake_quant_v_onwrite", lambda *args, **kwargs: None) - monkeypatch.setattr( - vllm_plugin, - "triton_decode_attention", - lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("split-K kernel")), - ) - - def fake_attention(query, **kwargs): - captured.update(kwargs) - return torch.zeros_like(query) - - monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) - output = torch.empty_like(q) - assert impl.forward(None, q, q, q, kv_cache, metadata, output=output) is output - assert captured["skip_softmax_threshold"] == pytest.approx(0.001) - assert captured["p_qdq"] == "nvfp4" - assert captured["v_qdq"] == "nvfp4" - - def test_resolve_calibrated_skip_softmax_threshold_for_decode(): """Calibration conversion is phase-aware even when decode later delegates.""" sparse_kw = { diff --git a/tests/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index 62395ff5a7a..8e6bef1db55 100644 --- a/tests/unit/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -139,21 +139,7 @@ def test_forward_routes_every_mode_to_single_autotuner( assert kernel.kwargs["V_QDQ"] == expected_v_qdq -@pytest.mark.parametrize( - ("attention_kwargs", "expected_block_m"), - [ - ({"skip_softmax_threshold": 0.1, "measure_sparsity": True}, 128), - ( - { - "p_qdq": "nvfp4", - "skip_softmax_threshold": 0.1, - "measure_sparsity": True, - }, - 16, - ), - ], -) -def test_forward_measurement_uses_one_fixed_launch(monkeypatch, attention_kwargs, expected_block_m): +def test_forward_measurement_uses_one_fixed_launch(monkeypatch): """Counter measurement bypasses autotuning to avoid repeated atomic updates.""" pytest.importorskip("triton") @@ -173,10 +159,49 @@ def test_forward_measurement_uses_one_fixed_launch(monkeypatch, attention_kwargs starts = torch.tensor([0], dtype=torch.int32) lengths = torch.tensor([seq_len], dtype=torch.int32) - triton_fa.attention(q, k, v, starts, lengths, seq_len, **attention_kwargs) + triton_fa.attention( + q, + k, + v, + starts, + lengths, + seq_len, + skip_softmax_threshold=0.1, + measure_sparsity=True, + ) assert kernel.fn.launch_count == 1 - assert kernel.fn.kwargs["BLOCK_M"] == expected_block_m + assert kernel.fn.kwargs["BLOCK_M"] == 128 assert kernel.fn.kwargs["BLOCK_N"] == 128 assert kernel.fn.kwargs["num_stages"] == 1 assert kernel.fn.kwargs["num_warps"] == 4 + + +@pytest.mark.parametrize("qdq_kwargs", [{"p_qdq": "nvfp4"}, {"v_qdq": "nvfp4"}]) +def test_forward_rejects_skip_softmax_with_attention_quantization(monkeypatch, qdq_kwargs): + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) + monkeypatch.setattr(triton_fa, "_load_qdq_helpers", lambda: None) + + seq_len = 129 + q = torch.empty(seq_len, 2, 16) + k = torch.empty(seq_len, 1, 16) + v = torch.empty_like(k) + starts = torch.tensor([0], dtype=torch.int32) + lengths = torch.tensor([seq_len], dtype=torch.int32) + + with pytest.raises(ValueError, match="skip-softmax cannot be combined"): + triton_fa.attention( + q, + k, + v, + starts, + lengths, + seq_len, + skip_softmax_threshold=0.1, + measure_sparsity=True, + **qdq_kwargs, + ) From f11025d8daff7ae250cdcc6abb045386bb4e52a3 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 1 Aug 2026 09:05:59 -0700 Subject: [PATCH 13/16] Support mask reuse capture on Python 3.10 Signed-off-by: Kai Xu --- .../attention_sparsity/plugins/mask_reuse_capture.py | 5 +++-- .../attention_sparsity/test_mask_reuse_calibration.py | 2 +- .../sparsity/attention_sparsity/test_mask_reuse_capture.py | 2 +- .../test_mask_reuse_compact_calibration.py | 2 +- 4 files changed, 6 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py b/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py index 6dd5bcdd134..9cee1adff50 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/mask_reuse_capture.py @@ -445,7 +445,8 @@ def build_capture_invocation( a = float(params["a"]) b = float(params["b"]) threshold_log2 = math.log2(a) + b * target * math.log2(math.e) - math.log2(sample_length) - threshold_lambda = math.exp2(threshold_log2) # type: ignore[attr-defined] + # ``math.exp2`` was added in Python 3.11; ModelOpt still supports 3.10. + threshold_lambda = 2.0**threshold_log2 if not math.isfinite(threshold_log2) or not 0.0 < threshold_lambda < 1.0: raise CaptureContractError("vanilla fit derives a threshold outside (0, 1)") @@ -537,7 +538,7 @@ def _validate_invocation(raw: object) -> dict[str, object]: ) if threshold_log2 >= 0.0 or not 0.0 < threshold_lambda < 1.0: raise CaptureContractError("capture invocation threshold must be in (0, 1)") - if math.exp2(threshold_log2).hex() != threshold_lambda.hex(): # type: ignore[attr-defined] + if (2.0**threshold_log2).hex() != threshold_lambda.hex(): raise CaptureContractError("capture invocation threshold hex fields disagree") geometry = _validate_geometry(raw["expected_geometry"], "expected_geometry") if geometry["kv_tokens"] != sample_length: 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 index da8cde69c5b..e43173440c3 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py @@ -77,7 +77,7 @@ def _thresholds(target_sparsity: float, sample_length: int) -> tuple[float, floa threshold_log2 = ( math.log2(A) + B * target_sparsity * math.log2(math.e) - math.log2(sample_length) ) - return math.exp2(threshold_log2), threshold_log2 + return 2.0**threshold_log2, threshold_log2 def _observations() -> list[MaskReuseObservation]: diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py index ff6c904daf8..10064db2a91 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_capture.py @@ -140,7 +140,7 @@ def test_build_invocation_binds_exact_threshold_source_and_final_chunk(): assert invocation["partition"] == "development" assert invocation["inner_fold"] == 0 assert invocation["threshold_log2_hex"] == expected_log2.hex() - assert invocation["threshold_lambda_hex"] == math.exp2(expected_log2).hex() + assert invocation["threshold_lambda_hex"] == (2.0**expected_log2).hex() assert invocation["expected_geometry"] == { "q_tokens": 256, "kv_tokens": 8448, 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 index 627e5e5e0a4..347be356a6f 100644 --- 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 @@ -87,7 +87,7 @@ def _capture(split, prompt_id, target, checkpoint_sha256): "target_sparsity_hex": target.hex(), "sample_length": 256, "threshold_log2_hex": threshold_log2.hex(), - "threshold_lambda_hex": math.exp2(threshold_log2).hex(), + "threshold_lambda_hex": (2.0**threshold_log2).hex(), "expected_geometry": {"q_tokens": 256, "kv_tokens": 256, "q_start_tokens": 0}, } return { From f3de07342268d7317f2b3b8da73aa34965121933 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 1 Aug 2026 09:32:08 -0700 Subject: [PATCH 14/16] Harden mask reuse capture portability Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 4 +- examples/vllm_serve/calibrate_mask_reuse.py | 42 +++--- examples/vllm_serve/collect_mask_reuse.py | 23 ++- .../calibration/__init__.py | 2 + .../calibration/checkpoint_manifest.py | 134 ++++++++++++++---- .../calibration/mask_reuse.py | 14 +- .../test_checkpoint_manifest.py | 105 ++++++++++++++ .../test_collect_mask_reuse_cli.py | 4 +- .../test_mask_reuse_calibration.py | 22 +++ 9 files changed, 286 insertions(+), 64 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 3830da94362..bec64b271ef 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -235,7 +235,9 @@ the standalone report and do not silently retune the selected policy. The candidate has `promotion_status="candidate_only"` and `deployment_geometry_validated=false`; the serving backend rejects it. Candidate and report publication is no-clobber and report-first with rollback, so a policy -path cannot appear without its complete report. Decode is explicitly dense. +path cannot appear without its complete report during normal process execution. +On platforms without portable directory `fsync`, directory-entry durability +across power loss is best effort. Decode is explicitly dense. The reusable serving policies live in `modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py`. `install_vllm_sparse_attention_from_checkpoint` installs checkpoint-driven sparse-only attention, while `install_vllm_nvfp4_attention` installs fixed NVFP4 Q/K/P/V with optional checkpoint sparsity. Both validate every selected layer before publishing any replacement implementation and return a `VllmAttentionInstallReport` with the installed layer names and backend counts. diff --git a/examples/vllm_serve/calibrate_mask_reuse.py b/examples/vllm_serve/calibrate_mask_reuse.py index 72480908ffe..cb2d95184d9 100644 --- a/examples/vllm_serve/calibrate_mask_reuse.py +++ b/examples/vllm_serve/calibrate_mask_reuse.py @@ -44,7 +44,6 @@ import argparse import json import os -import stat import tempfile from collections.abc import Mapping from hashlib import sha256 @@ -53,6 +52,7 @@ 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 ( @@ -127,29 +127,7 @@ def _load_json_object(path: Path, *, label: str) -> dict[str, object]: def _stable_file_sha256(path: Path, *, label: str) -> str: - if path.is_symlink(): - raise ValueError(f"{label} must not be a symlink") - try: - descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW) - except OSError as error: - raise ValueError(f"could not open {label} at {path}") from error - before = os.fstat(descriptor) - digest = sha256() - try: - with os.fdopen(descriptor, "rb") as handle: - for chunk in iter(lambda: handle.read(1024 * 1024), b""): - digest.update(chunk) - after = os.fstat(handle.fileno()) - named = path.stat(follow_symlinks=False) - except OSError as error: - raise ValueError(f"could not hash stable {label} at {path}") from error - identities = { - (value.st_dev, value.st_ino, value.st_size, value.st_mtime_ns) - for value in (before, after, named) - } - if len(identities) != 1 or not stat.S_ISREG(named.st_mode): - raise ValueError(f"{label} changed while it was being hashed") - return digest.hexdigest() + return stable_file_sha256(path, label=label) def _evidence_artifacts( @@ -178,7 +156,12 @@ def _canonical_json_bytes(value: object) -> bytes: def _fsync_directory(path: Path) -> None: - descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) + 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: @@ -202,20 +185,27 @@ def _temporary_payload(path: Path, payload: bytes) -> Path: 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_dev, observed.st_ino) == identity: + 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: diff --git a/examples/vllm_serve/collect_mask_reuse.py b/examples/vllm_serve/collect_mask_reuse.py index 76aaf76ef7a..61406fa9718 100644 --- a/examples/vllm_serve/collect_mask_reuse.py +++ b/examples/vllm_serve/collect_mask_reuse.py @@ -270,7 +270,12 @@ def _configure_capture_environment( def _fsync_directory(path: Path) -> None: - descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) + 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: @@ -280,9 +285,15 @@ def _fsync_directory(path: Path) -> None: def _publish_no_clobber(temporary: Path, destination: Path) -> tuple[int, int]: """Atomically link a complete temp file only when destination is absent.""" - identity = (temporary.stat().st_dev, temporary.stat().st_ino) - os.link(temporary, destination) + observed = temporary.stat(follow_symlinks=False) + if observed.st_ino == 0: + raise RuntimeError("capture temporary file has no stable identity") + identity = observed.st_dev, observed.st_ino + os.link(temporary, destination, follow_symlinks=False) try: + published = destination.stat(follow_symlinks=False) + if published.st_ino == 0 or (published.st_dev, published.st_ino) != identity: + raise RuntimeError("capture destination changed during publication") temporary.unlink() _fsync_directory(destination.parent) except BaseException: @@ -294,11 +305,13 @@ def _publish_no_clobber(temporary: Path, destination: Path) -> tuple[int, int]: def _unlink_if_identity(path: Path, identity: tuple[int, int]) -> None: """Rollback only a file that is still the inode published by this process.""" + if identity[1] == 0: + return try: - stat = path.stat() + stat = path.stat(follow_symlinks=False) except FileNotFoundError: return - if (stat.st_dev, stat.st_ino) == identity: + if stat.st_ino != 0 and (stat.st_dev, stat.st_ino) == identity: path.unlink() _fsync_directory(path.parent) diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py b/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py index 950a411182f..c8b55c03a3a 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/__init__.py @@ -24,6 +24,7 @@ VerifiedCheckpointManifest, create_checkpoint_manifest, read_stable_file_snapshot, + stable_file_sha256, verify_checkpoint_manifest, ) from .mask_reuse import ( @@ -64,5 +65,6 @@ "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 index 6ce3cadf274..dd79737fb25 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/checkpoint_manifest.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/checkpoint_manifest.py @@ -33,6 +33,7 @@ "VerifiedCheckpointManifest", "create_checkpoint_manifest", "read_stable_file_snapshot", + "stable_file_sha256", "verify_checkpoint_manifest", ] @@ -40,6 +41,7 @@ _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): @@ -93,23 +95,52 @@ def _canonical_json_bytes(value: object) -> bytes: ).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: - descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW) + 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 - opened = os.fstat(descriptor) 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 stat.S_ISLNK(named.st_mode) - or (opened.st_dev, opened.st_ino) != (named.st_dev, named.st_ino) + 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") @@ -124,25 +155,25 @@ def _hash_stable_regular( payload = bytearray() if capture_payload else None observed_size = 0 try: - with os.fdopen(descriptor, "rb") as handle: - for chunk in iter(lambda: handle.read(1024 * 1024), b""): - observed_size += len(chunk) - digest.update(chunk) - if payload is not None: - payload.extend(chunk) - after = os.fstat(handle.fileno()) + 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 - identity_before = (before.st_dev, before.st_ino, before.st_size, before.st_mtime_ns) - identity_after = (after.st_dev, after.st_ino, after.st_size, after.st_mtime_ns) - identity_named = ( - named_after.st_dev, - named_after.st_ino, - named_after.st_size, - named_after.st_mtime_ns, - ) - if identity_before != identity_after or identity_after != identity_named: + 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) @@ -164,20 +195,54 @@ def read_stable_file_snapshot(path: str | Path, *, label: str) -> StableFileSnap 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() - for path in root.rglob("*"): - if path.is_symlink(): + 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"checkpoint contains forbidden symlink {path.relative_to(root).as_posix()!r}" - ) - if path.is_file() and path != manifest_path: - files.add(path.relative_to(root).as_posix()) + 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: - descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) + 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: @@ -222,16 +287,29 @@ def create_checkpoint_manifest(checkpoint: str | Path, *, model: str) -> Verifie 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_dev, published.st_ino) == identity: + if published.st_ino != 0 and (published.st_dev, published.st_ino) == identity: manifest_path.unlink() _fsync_directory(root) finally: diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py index d99f785055b..cfd9e89b586 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py @@ -446,13 +446,23 @@ def _find_skip_softmax_group(raw: Mapping[str, object]) -> Mapping[str, object]: matches = [ group for group in groups.values() - if isinstance(group, Mapping) and group.get("algorithm") == "skip_softmax" + 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" ) - current = matches[0] + 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): diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py b/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py index d1090296e0b..8e90f4feb97 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_checkpoint_manifest.py @@ -15,14 +15,18 @@ """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, ) @@ -70,3 +74,104 @@ def test_verifier_rejects_symlink_manifest(tmp_path): with pytest.raises(CheckpointManifestError, match="without following symlinks"): verify_checkpoint_manifest(root) + + +def test_portable_open_fallback_still_rejects_symlink(tmp_path, monkeypatch): + target = tmp_path / "target.json" + target.write_text("{}\n", encoding="utf-8") + alias = tmp_path / "alias.json" + alias.symlink_to(target) + monkeypatch.delattr(checkpoint_manifest.os, "O_NOFOLLOW", raising=False) + + with pytest.raises(CheckpointManifestError, match="without following symlinks"): + read_stable_file_snapshot(alias, label="fallback input") + + +def test_portable_open_fallback_preserves_exact_binary_bytes(tmp_path, monkeypatch): + path = tmp_path / "binary.dat" + payload = b"line-1\r\nline-2\x00\n" + path.write_bytes(payload) + monkeypatch.delattr(checkpoint_manifest.os, "O_NOFOLLOW", raising=False) + + snapshot = read_stable_file_snapshot(path, label="binary input") + + assert snapshot.payload == payload + assert snapshot.sha256 == sha256(payload).hexdigest() + + +def test_portable_open_uses_binary_flag_when_available(tmp_path, monkeypatch): + path = tmp_path / "binary.dat" + path.write_bytes(b"payload") + binary_flag = 1 << 29 + observed_flags = [] + real_open = checkpoint_manifest.os.open + monkeypatch.setattr(checkpoint_manifest.os, "O_BINARY", binary_flag, raising=False) + + def recording_open(source, flags, *args, **kwargs): + observed_flags.append(flags) + return real_open(source, flags & ~binary_flag, *args, **kwargs) + + monkeypatch.setattr(checkpoint_manifest.os, "open", recording_open) + + snapshot = read_stable_file_snapshot(path, label="binary input") + + assert snapshot.payload == b"payload" + assert observed_flags and observed_flags[0] & binary_flag + + +def test_portable_open_rejects_path_swap_before_read(tmp_path, monkeypatch): + path = tmp_path / "input.dat" + replacement = tmp_path / "replacement.dat" + path.write_bytes(b"expected") + replacement.write_bytes(b"attacker") + real_open = checkpoint_manifest.os.open + swapped = False + monkeypatch.delattr(checkpoint_manifest.os, "O_NOFOLLOW", raising=False) + + def swapping_open(source, flags, *args, **kwargs): + nonlocal swapped + if Path(source) == path and not swapped: + swapped = True + path.unlink() + replacement.rename(path) + return real_open(source, flags, *args, **kwargs) + + monkeypatch.setattr(checkpoint_manifest.os, "open", swapping_open) + + with pytest.raises(CheckpointManifestError, match="stable regular file"): + read_stable_file_snapshot(path, label="swapped input") + + +def test_windows_reparse_attribute_is_link_like(): + observed = SimpleNamespace( + st_mode=0, + st_file_attributes=checkpoint_manifest._FILE_ATTRIBUTE_REPARSE_POINT, + ) + + assert checkpoint_manifest._is_link_like(observed) + + +def test_manifest_creation_without_directory_fsync_support(tmp_path, monkeypatch): + root = _toy_checkpoint(tmp_path / "checkpoint") + monkeypatch.delattr(checkpoint_manifest.os, "O_DIRECTORY", raising=False) + + created = create_checkpoint_manifest(root, model="toy-model") + + assert verify_checkpoint_manifest(root, expected_model="toy-model") == created + + +def test_manifest_publication_preserves_destination_created_by_racer(tmp_path, monkeypatch): + root = _toy_checkpoint(tmp_path / "checkpoint") + manifest = root / CHECKPOINT_MANIFEST_NAME + real_link = checkpoint_manifest.os.link + + def racing_link(source, target, **kwargs): + Path(target).write_text("racer\n", encoding="utf-8") + return real_link(source, target, **kwargs) + + monkeypatch.setattr(checkpoint_manifest.os, "link", racing_link) + + with pytest.raises(CheckpointManifestError, match="appeared during publication"): + create_checkpoint_manifest(root, model="toy-model") + + assert manifest.read_text(encoding="utf-8") == "racer\n" diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py b/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py index 6a005873f5f..44628328086 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py @@ -305,9 +305,9 @@ def test_publish_no_clobber_preserves_destination_created_by_racer(tmp_path, mon temporary.write_text("ours", encoding="utf-8") real_link = os.link - def racing_link(source, target): + def racing_link(source, target, **kwargs): Path(target).write_text("racer", encoding="utf-8") - return real_link(source, target) + return real_link(source, target, **kwargs) monkeypatch.setattr(collect_mask_reuse_cli.os, "link", racing_link) 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 index e43173440c3..0531845ba23 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_mask_reuse_calibration.py @@ -294,3 +294,25 @@ def test_canonicalizes_existing_sparse_attention_config(): "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}, + } From 2cbf00cc08e6df0aed0180187529f77bbcc46886 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 1 Aug 2026 10:07:15 -0700 Subject: [PATCH 15/16] Seal FA4 source for mask reuse capture Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 42 +- examples/vllm_serve/calibrate_mask_reuse.py | 41 +- examples/vllm_serve/collect_mask_reuse.py | 169 +++-- .../vllm_serve/create_fa4_source_witness.py | 57 ++ .../calibration/mask_reuse.py | 2 +- .../calibration/mask_reuse_compact.py | 4 +- .../calibration/source_manifest.py | 598 ++++++++++++++++++ .../test_calibrate_mask_reuse_cli.py | 11 +- .../test_collect_mask_reuse_cli.py | 72 ++- .../test_source_manifest.py | 315 +++++++++ 10 files changed, 1227 insertions(+), 84 deletions(-) create mode 100644 examples/vllm_serve/create_fa4_source_witness.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/calibration/source_manifest.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index bec64b271ef..c5737cd91e2 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -157,6 +157,23 @@ SHA256 and size for every checkpoint file except the manifest itself. It refuses to overwrite an existing manifest, rejects symlinks, and verifies the complete file set. Keep the checkpoint root immutable after creating it. +Create the runtime FA4 source from the exact Git commit outside the serving +container. The generator emits a scoped `git archive` containing only +`flash_attn/` and a canonical witness for every directory and regular file: + +```bash +python create_fa4_source_witness.py /path/to/flash-attention \ + --expected-commit <40-HEX-COMMIT> \ + --archive-output fa4-runtime-source.tar \ + --manifest-output fa4-runtime-source.manifest.json +mkdir fa4-runtime-source +tar -xf fa4-runtime-source.tar -C fa4-runtime-source +``` + +Pin the printed archive and manifest hashes independently. The runtime verifier +does not invoke Git; it rejects missing, extra, changed, executable-mode, link, +reparse-point, and special-file differences from the witnessed archive. + Then collect dense-reference mask-reuse sufficient statistics at a menu of target sparsities under the exact runtime rule `threshold = a * exp(b * target_sparsity) / kv_tokens`: @@ -164,8 +181,12 @@ sparsities under the exact runtime rule ```bash python collect_mask_reuse.py \ --model-id Nemotron-3-Ultra \ + --checkpoint-manifest-sha256 <64-HEX-CHECKPOINT-MANIFEST-SHA256> \ --plan nemotron3_ultra_stride2 \ - --fa4-source /path/to/pinned-calibrated-threshold-flash-attention \ + --fa4-source /path/to/fa4-runtime-source \ + --fa4-source-manifest fa4-runtime-source.manifest.json \ + --fa4-source-manifest-sha256 <64-HEX-SOURCE-MANIFEST-SHA256> \ + --fa4-commit <40-HEX-COMMIT> \ --prompts-jsonl mask_reuse_prompts.jsonl \ --vanilla-config /config.json \ --target-sparsities 0.5 0.6 0.7 \ @@ -185,7 +206,8 @@ vLLM V1 execution through the installed `mask_reuse_fa4` plugin and the explicit pinned FA4 source. It does not load a provisional reuse policy and does not enable quantization. The verified checkpoint identity and group assignment are included in the schema-v2 invocation and its canonical RPC -digest. The checkpoint is reverified after vLLM loads it and before capture. +digest. The checkpoint and FA4 source are reverified after vLLM loads them and +before capture; FA4 is verified once more before evidence publication. Each compact JSONL record includes the target sparsity, sample length, exact binary64 `threshold_log2` and `threshold_lambda`, rectangular query/KV geometry, every @@ -193,13 +215,15 @@ topology anchor/head's self-mask statistics, and every consumer/donor dropped-mass matrix once. It does not expand or repeat the matrix as millions of candidate rows. Fixed-lambda version-4/version-5 captures are not valid inputs to this target-sparsity path. Use the printed compact-capture SHA256 as -the `reuse_bundle_sha256` evidence value. The FA4 source must be a clean Git -checkout; its exact commit is recorded in the capture manifest. The version-3 -capture manifest also retains the canonical vLLM engine arguments, rank-major -TP sentinel evidence, exact prefill/decode attention-call counts, and optional -bitwise dense-shadow coverage. `--validate-dense-output` runs a second pinned -dense FA4 call for every armed final-chunk layer and fails before publication -if its BF16 output differs bitwise from the model-trajectory output. +the `reuse_bundle_sha256` evidence value. The version-4 capture manifest binds +the independently pinned checkpoint and FA4 source witnesses, Git commit/tree, +archive digest, canonical vLLM engine arguments, rank-major TP sentinel +evidence, exact prefill/decode attention-call counts, and optional bitwise +dense-shadow coverage. `--validate-dense-output` runs a second pinned dense FA4 +call for every armed final-chunk layer and fails before publication if its BF16 +output differs bitwise from the model-trajectory output. Schema-version-3 +capture manifests are intentionally rejected because they lack the runtime +source witness; recollect them rather than migrating their provenance labels. Once those observations have been captured, select the per-context target, donor-head map, and exact fallbacks and export the fail-closed candidate: diff --git a/examples/vllm_serve/calibrate_mask_reuse.py b/examples/vllm_serve/calibrate_mask_reuse.py index cb2d95184d9..59c4b5db00b 100644 --- a/examples/vllm_serve/calibrate_mask_reuse.py +++ b/examples/vllm_serve/calibrate_mask_reuse.py @@ -79,6 +79,13 @@ "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", @@ -255,8 +262,8 @@ def _validate_capture_manifest( f"missing={sorted(missing)}, extra={sorted(extra)}" ) expected = { - "capture_manifest_schema_version": 3, - "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v3", + "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, @@ -265,6 +272,32 @@ def _validate_capture_manifest( 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): @@ -298,7 +331,9 @@ def _validate_capture_manifest( def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( - description="Build a fail-closed schema-v3 candidate from compact mask-reuse captures", + 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) diff --git a/examples/vllm_serve/collect_mask_reuse.py b/examples/vllm_serve/collect_mask_reuse.py index 61406fa9718..54d830537c8 100644 --- a/examples/vllm_serve/collect_mask_reuse.py +++ b/examples/vllm_serve/collect_mask_reuse.py @@ -30,8 +30,12 @@ python examples/vllm_serve/collect_mask_reuse.py /path/to/checkpoint \ --model-id Nemotron-3-Ultra \ + --checkpoint-manifest-sha256 012345... \ --plan nemotron3_ultra_stride2 \ - --fa4-source /path/to/flash-attention-at-4c40766b \ + --fa4-source /path/to/extracted-fa4-runtime-source \ + --fa4-source-manifest /path/to/fa4-source-manifest.json \ + --fa4-source-manifest-sha256 abcdef... \ + --fa4-commit 4c40766b... \ --prompts-jsonl prompts.jsonl \ --vanilla-config /path/to/config.json \ --target-sparsities 0.5 0.6 0.7 \ @@ -44,7 +48,7 @@ import importlib.metadata import json import os -import subprocess +import stat import sys import tempfile from hashlib import sha256 @@ -55,6 +59,11 @@ read_stable_file_snapshot, verify_checkpoint_manifest, ) +from modelopt.torch.sparsity.attention_sparsity.calibration.source_manifest import ( + SourceManifestError, + VerifiedSourceManifest, + verify_source_manifest, +) from modelopt.torch.sparsity.attention_sparsity.plugins.mask_reuse_capture import ( MAX_QUERY_CHUNK_TOKENS, CaptureContractError, @@ -105,11 +114,31 @@ def _parser() -> argparse.ArgumentParser: default=None, help="Stable model name stored in observations (default: the model argument)", ) + parser.add_argument( + "--checkpoint-manifest-sha256", + required=True, + help="Separately pinned SHA256 of the checkpoint manifest", + ) parser.add_argument("--plan", required=True, help="Explicit mask-reuse layer-plan preset") parser.add_argument( "--fa4-source", required=True, - help="Pinned FlashAttention source checkout containing flash_attn/cute", + help="FlashAttention source tree extracted from the exact witnessed git archive", + ) + parser.add_argument( + "--fa4-source-manifest", + required=True, + help="Canonical full-tree witness generated with create_fa4_source_witness.py", + ) + parser.add_argument( + "--fa4-source-manifest-sha256", + required=True, + help="Separately pinned SHA256 of --fa4-source-manifest", + ) + parser.add_argument( + "--fa4-commit", + required=True, + help="Separately pinned 40-hex FlashAttention Git commit", ) parser.add_argument("--prompts-jsonl", required=True, help="Strict prompt-plan JSONL") parser.add_argument( @@ -207,43 +236,55 @@ def _canonical_capture_line(capture: dict[str, object]) -> bytes: ).encode() +def _verify_fa4_source( + fa4_source: str, + fa4_source_manifest: str, + fa4_source_manifest_sha256: str, + fa4_commit: str, +) -> VerifiedSourceManifest: + try: + verified = verify_source_manifest( + fa4_source, + fa4_source_manifest, + expected_manifest_sha256=fa4_source_manifest_sha256, + expected_commit=fa4_commit, + expected_source_kind="flash-attention-4", + ) + except SourceManifestError as error: + raise CaptureContractError(f"--fa4-source verification failed: {error}") from error + required = ( + verified.source_root / "flash_attn/cute/interface.py", + verified.source_root / "flash_attn/cute/block_sparsity.py", + ) + try: + required_are_regular = all( + stat.S_ISREG(path.stat(follow_symlinks=False).st_mode) for path in required + ) + except OSError: + required_are_regular = False + if not required_are_regular: + raise CaptureContractError( + "--fa4-source must contain flash_attn/cute/interface.py and block_sparsity.py" + ) + return verified + + def _configure_capture_environment( plan: str, fa4_source: str, + fa4_source_manifest: str, + fa4_source_manifest_sha256: str, + fa4_commit: str, checkpoint_manifest_sha256: str, *, validate_dense_output: bool, -) -> tuple[Path, str]: - source = Path(fa4_source).expanduser().resolve() - required = ( - source / "flash_attn/cute/interface.py", - source / "flash_attn/cute/block_sparsity.py", +) -> VerifiedSourceManifest: + verified = _verify_fa4_source( + fa4_source, + fa4_source_manifest, + fa4_source_manifest_sha256, + fa4_commit, ) - if not source.is_dir() or any(not path.is_file() for path in required): - raise CaptureContractError( - "--fa4-source must contain flash_attn/cute/interface.py and block_sparsity.py" - ) - try: - commit = subprocess.run( - ["git", "-C", str(source), "rev-parse", "HEAD"], - check=True, - capture_output=True, - text=True, - ).stdout.strip() - dirty = subprocess.run( - ["git", "-C", str(source), "status", "--porcelain", "--untracked-files=all"], - check=True, - capture_output=True, - text=True, - ).stdout.strip() - except (OSError, subprocess.CalledProcessError) as error: - raise CaptureContractError("--fa4-source must be a readable Git checkout") from error - if len(commit) != 40 or any(character not in "0123456789abcdef" for character in commit): - raise CaptureContractError("--fa4-source HEAD is not a canonical Git commit") - if dirty: - raise CaptureContractError( - "--fa4-source has tracked or untracked modifications; pin a clean commit" - ) plugins = [ entry for entry in importlib.metadata.entry_points(group="vllm.general_plugins") @@ -259,14 +300,14 @@ def _configure_capture_environment( os.environ[DENSE_SHADOW_ENV] = "1" if validate_dense_output else "0" os.environ["PYTHONDONTWRITEBYTECODE"] = "1" sys.dont_write_bytecode = True - os.environ["MASK_REUSE_FA4_SOURCE"] = str(source) + os.environ["MASK_REUSE_FA4_SOURCE"] = str(verified.source_root) os.environ["MASK_REUSE_FA4_FORCE_DENSE"] = "0" os.environ["VLLM_PLUGINS"] = "mask_reuse_fa4" # Collection is deliberately policy-free. Remove inherited serving # settings so a stale deployment policy cannot become threshold authority. for name in _POLICY_ENVS: os.environ.pop(name, None) - return source, commit + return verified def _fsync_directory(path: Path) -> None: @@ -325,6 +366,10 @@ def run(args: argparse.Namespace) -> tuple[Path, Path]: threshold_scale_factor = parse_vanilla_prefill_fit(vanilla_snapshot.payload) targets = _target_menu(args.target_sparsities) checkpoint = verify_checkpoint_manifest(args.model, expected_model=args.model_id) + if checkpoint.sha256 != args.checkpoint_manifest_sha256: + raise CaptureContractError( + "checkpoint manifest does not match --checkpoint-manifest-sha256" + ) model_id = checkpoint.model output_path = Path(args.output) manifest_path = ( @@ -339,9 +384,12 @@ def run(args: argparse.Namespace) -> tuple[Path, Path]: output_path.parent.mkdir(parents=True, exist_ok=True) manifest_path.parent.mkdir(parents=True, exist_ok=True) - fa4_source, fa4_source_commit = _configure_capture_environment( + fa4_source = _configure_capture_environment( args.plan, args.fa4_source, + args.fa4_source_manifest, + args.fa4_source_manifest_sha256, + args.fa4_commit, checkpoint.sha256, validate_dense_output=args.validate_dense_output, ) @@ -356,6 +404,16 @@ def run(args: argparse.Namespace) -> tuple[Path, Path]: raise CaptureContractError( "checkpoint identity changed while vLLM loaded the model; no capture was started" ) + loaded_fa4_source = _verify_fa4_source( + args.fa4_source, + args.fa4_source_manifest, + args.fa4_source_manifest_sha256, + args.fa4_commit, + ) + if loaded_fa4_source != fa4_source: + raise CaptureContractError( + "FA4 source identity changed while vLLM loaded the model; no capture was started" + ) statuses = llm.collective_rpc("mask_reuse_capture_status") validate_capture_statuses(statuses) tokenizer = llm.get_tokenizer() @@ -428,27 +486,50 @@ def run(args: argparse.Namespace) -> tuple[Path, Path]: temporary_path.unlink(missing_ok=True) raise - prompt_final = read_stable_file_snapshot(args.prompts_jsonl, label="prompt plan") - vanilla_final = read_stable_file_snapshot(args.vanilla_config, label="vanilla config") - if prompt_final.sha256 != prompt_plan_sha256 or vanilla_final.sha256 != vanilla_config_sha256: + try: + prompt_final = read_stable_file_snapshot(args.prompts_jsonl, label="prompt plan") + vanilla_final = read_stable_file_snapshot(args.vanilla_config, label="vanilla config") + final_fa4_source = _verify_fa4_source( + args.fa4_source, + args.fa4_source_manifest, + args.fa4_source_manifest_sha256, + args.fa4_commit, + ) + except BaseException: + assert temporary_path is not None + temporary_path.unlink(missing_ok=True) + raise + if ( + prompt_final.sha256 != prompt_plan_sha256 + or vanilla_final.sha256 != vanilla_config_sha256 + or final_fa4_source != fa4_source + ): assert temporary_path is not None temporary_path.unlink(missing_ok=True) raise CaptureContractError( - "prompt plan or vanilla calibration changed during capture; evidence was discarded" + "prompt plan, vanilla calibration, or FA4 source changed during capture; " + "evidence was discarded" ) capture_sha256 = capture_digest.hexdigest() manifest = { - "capture_manifest_schema_version": 3, - "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v3", + "capture_manifest_schema_version": 4, + "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v4", "model": model_id, "checkpoint_manifest_sha256": checkpoint.sha256, "checkpoint_manifest_path": str(checkpoint.manifest_path), "checkpoint_file_count": checkpoint.file_count, "checkpoint_total_size_bytes": checkpoint.total_size_bytes, "plan": args.plan, - "fa4_source": str(fa4_source), - "fa4_source_commit": fa4_source_commit, + "fa4_source": str(fa4_source.source_root), + "fa4_source_commit": fa4_source.git_commit, + "fa4_source_git_tree": fa4_source.git_tree, + "fa4_source_git_archive_sha256": fa4_source.git_archive_sha256, + "fa4_source_manifest_path": str(fa4_source.manifest_path), + "fa4_source_manifest_sha256": fa4_source.manifest_sha256, + "fa4_source_directory_count": fa4_source.directory_count, + "fa4_source_file_count": fa4_source.file_count, + "fa4_source_total_size_bytes": fa4_source.total_size_bytes, "engine_kwargs": engine_kwargs, "dense_shadow_validation_requested": args.validate_dense_output, "target_sparsity_hex": [target.hex() for target in targets], diff --git a/examples/vllm_serve/create_fa4_source_witness.py b/examples/vllm_serve/create_fa4_source_witness.py new file mode 100644 index 00000000000..e32cc4f3823 --- /dev/null +++ b/examples/vllm_serve/create_fa4_source_witness.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Create an exact FlashAttention source archive and its runtime witness.""" + +from __future__ import annotations + +import argparse + +from modelopt.torch.sparsity.attention_sparsity.calibration.source_manifest import ( + create_source_manifest_from_git_archive, +) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Create a Git-free runtime witness for one exact FlashAttention commit" + ) + parser.add_argument("checkout", help="Clean FlashAttention Git checkout") + parser.add_argument("--expected-commit", required=True, help="Required 40-hex HEAD commit") + parser.add_argument("--archive-output", required=True, help="New exact git-archive tar path") + parser.add_argument( + "--manifest-output", required=True, help="New canonical source-witness JSON path" + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + generated = create_source_manifest_from_git_archive( + args.checkout, + expected_commit=args.expected_commit, + source_kind="flash-attention-4", + archive_output=args.archive_output, + manifest_output=args.manifest_output, + ) + print(f"git_commit={generated.git_commit}") + print(f"git_tree={generated.git_tree}") + print(f"git_archive_sha256={generated.git_archive_sha256}") + print(f"source_manifest_sha256={generated.manifest_sha256}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py index cfd9e89b586..397d61c152e 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse.py @@ -692,7 +692,7 @@ def _validate_thresholds( expected_log2 = ( math.log2(a) + b * target * math.log2(math.e) - math.log2(observation.sample_length) ) - expected_lambda = float(getattr(math, "exp2")(expected_log2)) + 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 " diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py index 6977134bb73..9d549a7496b 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/mask_reuse_compact.py @@ -248,7 +248,7 @@ def from_mapping(cls, raw: Mapping[str, object]) -> CompactMaskReuseCapture: 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 float(getattr(math, "exp2")(threshold_log2)).hex() != threshold_lambda.hex(): + 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") @@ -454,7 +454,7 @@ def _validate_threshold(capture: CompactMaskReuseCapture, fit: Mapping[str, obje + float(params["b"]) * capture.target_sparsity * math.log2(math.e) - math.log2(capture.sample_length) ) - expected_lambda = float(getattr(math, "exp2")(expected_log2)) + 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(): diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/source_manifest.py b/modelopt/torch/sparsity/attention_sparsity/calibration/source_manifest.py new file mode 100644 index 00000000000..9748b086c47 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/source_manifest.py @@ -0,0 +1,598 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Git-free verification of runtime source extracted from an exact Git archive.""" + +from __future__ import annotations + +import json +import os +import shutil +import stat +import subprocess # nosec B404 - used only by the explicit outside-container generator +import tarfile +import tempfile +import unicodedata +from dataclasses import dataclass +from hashlib import sha256 +from io import BytesIO +from pathlib import Path, PurePosixPath +from typing import cast + +from .checkpoint_manifest import ( + CheckpointManifestError, + read_stable_file_snapshot, + stable_file_sha256, +) + +__all__ = [ + "GeneratedSourceManifest", + "SourceManifestError", + "VerifiedSourceManifest", + "create_source_manifest_from_git_archive", + "verify_source_manifest", +] + +_ARCHIVE_SCOPE = "flash_attn" +_MANIFEST_FIELDS = frozenset( + { + "source_manifest_schema_version", + "source_kind", + "git_commit", + "git_tree", + "git_archive_sha256", + "archive_scope", + "directories", + "files", + } +) +_FILE_FIELDS = frozenset({"path", "mode", "size_bytes", "sha256"}) +_FILE_ATTRIBUTE_REPARSE_POINT = getattr(stat, "FILE_ATTRIBUTE_REPARSE_POINT", 0x400) + + +class SourceManifestError(ValueError): + """Raised when source bytes do not match their sealed Git-archive witness.""" + + +def _strict_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise SourceManifestError(f"source manifest repeats JSON key {key!r}") + result[key] = value + return result + + +def _exact_fields(raw: dict[str, object], expected: frozenset[str], label: str) -> None: + missing = expected - raw.keys() + extra = raw.keys() - expected + if missing or extra: + raise SourceManifestError( + f"{label} fields do not match the schema; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + + +def _canonical_json_bytes(value: object) -> bytes: + return ( + json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + + +def _canonical_text(value: object, label: str) -> str: + if ( + not isinstance(value, str) + or not value + or value != value.strip() + or unicodedata.normalize("NFC", value) != value + or any(ord(character) < 32 for character in value) + ): + raise SourceManifestError(f"{label} must be non-empty canonical NFC text") + return value + + +def _hex_digest(value: object, length: int, label: str) -> str: + if ( + not isinstance(value, str) + or len(value) != length + or any(character not in "0123456789abcdef" for character in value) + ): + raise SourceManifestError(f"{label} must be {length} lowercase hexadecimal characters") + return value + + +def _relative_path(value: object, label: str) -> str: + text = _canonical_text(value, label) + path = PurePosixPath(text) + if ( + path.is_absolute() + or text != path.as_posix() + or "\\" in text + or any(part in {"", ".", ".."} for part in path.parts) + or path.parts[0] == ".git" + ): + raise SourceManifestError(f"{label} must be a canonical non-.git relative POSIX path") + return text + + +def _scoped_path(value: object, label: str) -> str: + text = _relative_path(value, label) + if PurePosixPath(text).parts[0] != _ARCHIVE_SCOPE: + raise SourceManifestError(f"{label} must be within {_ARCHIVE_SCOPE!r}") + return text + + +def _is_reparse_point(value: os.stat_result) -> bool: + return bool(getattr(value, "st_file_attributes", 0) & _FILE_ATTRIBUTE_REPARSE_POINT) + + +def _file_mode(value: os.stat_result) -> str: + return "100755" if value.st_mode & 0o111 else "100644" + + +@dataclass(frozen=True, slots=True) +class _TreeSnapshot: + directories: tuple[str, ...] + files: tuple[dict[str, object], ...] + + +def _source_tree_snapshot(root: Path) -> _TreeSnapshot: + directories: set[str] = set() + files: dict[str, dict[str, object]] = {} + pending: list[tuple[Path, str]] = [(root, "")] + while pending: + directory, directory_relative = pending.pop() + try: + observed_directory = directory.stat(follow_symlinks=False) + if not stat.S_ISDIR(observed_directory.st_mode) or _is_reparse_point( + observed_directory + ): + raise SourceManifestError( + f"source directory {directory_relative or '.'!r} is not stable" + ) + with os.scandir(directory) as iterator: + children = sorted(iterator, key=lambda entry: entry.name) + except OSError as error: + raise SourceManifestError( + f"could not traverse source directory {directory_relative or '.'!r}" + ) from error + for child in children: + if not directory_relative and child.name == ".git": + continue + relative = _relative_path( + f"{directory_relative}/{child.name}".lstrip("/"), "source path" + ) + path = Path(child.path) + try: + observed = child.stat(follow_symlinks=False) + except OSError as error: + raise SourceManifestError(f"could not inspect source path {relative!r}") from error + if stat.S_ISLNK(observed.st_mode) or _is_reparse_point(observed): + raise SourceManifestError(f"source path {relative!r} must not be a link") + if stat.S_ISDIR(observed.st_mode): + directories.add(relative) + pending.append((path, relative)) + elif stat.S_ISREG(observed.st_mode): + try: + digest = stable_file_sha256(path, label=f"source file {relative!r}") + except CheckpointManifestError as error: + raise SourceManifestError( + f"could not hash stable source file {relative!r}" + ) from error + files[relative] = { + "path": relative, + "mode": _file_mode(observed), + "size_bytes": observed.st_size, + "sha256": digest, + } + else: + raise SourceManifestError(f"source path {relative!r} is not regular") + return _TreeSnapshot( + directories=tuple(sorted(directories)), + files=tuple(files[path] for path in sorted(files)), + ) + + +def _parse_manifest(payload: bytes) -> dict[str, object]: + try: + raw = json.loads(payload, object_pairs_hook=_strict_object) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise SourceManifestError("source manifest is not strict UTF-8 JSON") from error + if not isinstance(raw, dict): + raise SourceManifestError("source manifest must be a JSON object") + _exact_fields(raw, _MANIFEST_FIELDS, "source manifest") + if ( + type(raw["source_manifest_schema_version"]) is not int + or raw["source_manifest_schema_version"] != 1 + ): + raise SourceManifestError("source_manifest_schema_version must be 1") + if payload != _canonical_json_bytes(raw): + raise SourceManifestError("source manifest bytes are not canonical JSON") + _canonical_text(raw["source_kind"], "source manifest.source_kind") + _hex_digest(raw["git_commit"], 40, "source manifest.git_commit") + _hex_digest(raw["git_tree"], 40, "source manifest.git_tree") + _hex_digest(raw["git_archive_sha256"], 64, "source manifest.git_archive_sha256") + if raw["archive_scope"] != _ARCHIVE_SCOPE: + raise SourceManifestError(f"source manifest.archive_scope must be {_ARCHIVE_SCOPE!r}") + + raw_directories = raw["directories"] + if not isinstance(raw_directories, list): + raise SourceManifestError("source manifest.directories must be a list") + directories = [ + _scoped_path(value, f"source manifest.directories[{index}]") + for index, value in enumerate(raw_directories) + ] + if directories != sorted(set(directories)): + raise SourceManifestError("source manifest directories must be unique and sorted") + + raw_files = raw["files"] + if not isinstance(raw_files, list) or not raw_files: + raise SourceManifestError("source manifest.files must be a non-empty list") + paths: list[str] = [] + for index, item in enumerate(raw_files): + label = f"source manifest.files[{index}]" + if not isinstance(item, dict): + raise SourceManifestError(f"{label} must be an object") + _exact_fields(item, _FILE_FIELDS, label) + paths.append(_scoped_path(item["path"], f"{label}.path")) + if item["mode"] not in {"100644", "100755"}: + raise SourceManifestError(f"{label}.mode must be 100644 or 100755") + size = item["size_bytes"] + if isinstance(size, bool) or not isinstance(size, int) or size < 0: + raise SourceManifestError(f"{label}.size_bytes must be an integer >= 0") + _hex_digest(item["sha256"], 64, f"{label}.sha256") + if paths != sorted(set(paths)): + raise SourceManifestError("source manifest file paths must be unique and sorted") + if set(paths) & set(directories): + raise SourceManifestError("source manifest paths cannot be both files and directories") + return raw + + +@dataclass(frozen=True, slots=True) +class VerifiedSourceManifest: + """Identity of an exact, fully enumerated runtime source tree.""" + + source_root: Path + manifest_path: Path + source_kind: str + git_commit: str + git_tree: str + git_archive_sha256: str + manifest_sha256: str + directory_count: int + file_count: int + total_size_bytes: int + + +def verify_source_manifest( + source: str | Path, + manifest: str | Path, + *, + expected_manifest_sha256: str, + expected_commit: str, + expected_source_kind: str, +) -> VerifiedSourceManifest: + """Verify every non-``.git`` source path without invoking Git.""" + root = Path(source).expanduser().resolve() + if not root.is_dir(): + raise SourceManifestError("source root must be a local directory") + expected_digest = _hex_digest(expected_manifest_sha256, 64, "expected source manifest SHA256") + expected_git_commit = _hex_digest(expected_commit, 40, "expected source Git commit") + expected_kind = _canonical_text(expected_source_kind, "expected source kind") + try: + manifest_snapshot = read_stable_file_snapshot(manifest, label="source manifest") + except CheckpointManifestError as error: + raise SourceManifestError("could not read stable source manifest") from error + if manifest_snapshot.sha256 != expected_digest: + raise SourceManifestError("source manifest does not match its expected SHA256") + raw = _parse_manifest(manifest_snapshot.payload) + if raw["source_kind"] != expected_kind: + raise SourceManifestError("source manifest kind does not match the expected source kind") + if raw["git_commit"] != expected_git_commit: + raise SourceManifestError("source manifest commit does not match the expected Git commit") + + actual = _source_tree_snapshot(root) + if list(actual.directories) != raw["directories"] or list(actual.files) != raw["files"]: + raise SourceManifestError("source tree does not exactly match its sealed manifest") + if _source_tree_snapshot(root) != actual: + raise SourceManifestError("source tree changed during verification") + return VerifiedSourceManifest( + source_root=root, + manifest_path=manifest_snapshot.path.resolve(), + source_kind=expected_kind, + git_commit=expected_git_commit, + git_tree=str(raw["git_tree"]), + git_archive_sha256=str(raw["git_archive_sha256"]), + manifest_sha256=manifest_snapshot.sha256, + directory_count=len(actual.directories), + file_count=len(actual.files), + total_size_bytes=sum(cast("int", item["size_bytes"]) for item in actual.files), + ) + + +@dataclass(frozen=True, slots=True) +class GeneratedSourceManifest: + """Hashes emitted by the outside-container Git-archive generator.""" + + git_commit: str + git_tree: str + git_archive_sha256: str + manifest_sha256: str + + +def _run_git(checkout: Path, arguments: list[str]) -> bytes: + executable = shutil.which("git") + if executable is None: + raise SourceManifestError("Git is required by the outside-container generator") + try: + return subprocess.run( + [executable, "-C", str(checkout), *arguments], # nosec B603 + check=True, + capture_output=True, + ).stdout + except (OSError, subprocess.CalledProcessError) as error: + raise SourceManifestError(f"Git command failed: {' '.join(arguments)}") from error + + +def _fsync_directory(path: Path) -> None: + directory_flag = getattr(os, "O_DIRECTORY", None) + if directory_flag is None: + return + descriptor = os.open(path, os.O_RDONLY | directory_flag) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def _unlink_if_identity(path: Path, identity: tuple[int, int]) -> None: + if identity[1] == 0: + return + try: + observed = path.stat(follow_symlinks=False) + except FileNotFoundError: + return + if ( + observed.st_ino != 0 + and stat.S_ISREG(observed.st_mode) + and not _is_reparse_point(observed) + and (observed.st_dev, observed.st_ino) == identity + ): + path.unlink() + _fsync_directory(path.parent) + + +def _temporary_payload(destination: Path, payload: bytes) -> tuple[Path, tuple[int, int], int]: + destination.parent.mkdir(parents=True, exist_ok=True) + descriptor, name = tempfile.mkstemp( + dir=destination.parent, + prefix=f".{destination.name}.", + suffix=".tmp", + ) + temporary = Path(name) + identity: tuple[int, int] | None = None + try: + with os.fdopen(os.dup(descriptor), "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + opened = os.fstat(descriptor) + observed = temporary.stat(follow_symlinks=False) + identity = (opened.st_dev, opened.st_ino) + if ( + opened.st_ino == 0 + or not stat.S_ISREG(opened.st_mode) + or _is_reparse_point(observed) + or (observed.st_dev, observed.st_ino) != identity + ): + raise SourceManifestError("source artifact temporary file has no stable identity") + return temporary, identity, descriptor + except BaseException: + if identity is None: + try: + opened = os.fstat(descriptor) + if opened.st_ino != 0 and stat.S_ISREG(opened.st_mode): + identity = (opened.st_dev, opened.st_ino) + except OSError: + pass + if identity is not None: + _unlink_if_identity(temporary, identity) + os.close(descriptor) + raise + + +def _publish_no_clobber( + temporary: Path, destination: Path, identity: tuple[int, int], descriptor: int +) -> tuple[int, int]: + observed = temporary.stat(follow_symlinks=False) + opened = os.fstat(descriptor) + if ( + observed.st_ino == 0 + or (observed.st_dev, observed.st_ino) != identity + or (opened.st_dev, opened.st_ino) != identity + ): + raise SourceManifestError("source artifact temporary file changed before publication") + os.link(temporary, destination, follow_symlinks=False) + try: + published = destination.stat(follow_symlinks=False) + opened_after = os.fstat(descriptor) + if ( + published.st_ino == 0 + or not stat.S_ISREG(published.st_mode) + or _is_reparse_point(published) + or (published.st_dev, published.st_ino) != identity + or (opened_after.st_dev, opened_after.st_ino) != identity + ): + raise SourceManifestError("source artifact destination changed during publication") + _unlink_if_identity(temporary, identity) + _fsync_directory(destination.parent) + except BaseException: + _unlink_if_identity(destination, identity) + raise + return identity + + +def _publish_source_artifacts( + archive_path: Path, + archive: bytes, + manifest_path: Path, + manifest: bytes, +) -> None: + archive_temporary_path, archive_temporary_identity, archive_descriptor = _temporary_payload( + archive_path, archive + ) + archive_temporary: Path | None = archive_temporary_path + manifest_temporary: Path | None = None + manifest_temporary_identity: tuple[int, int] | None = None + manifest_descriptor: int | None = None + archive_identity: tuple[int, int] | None = None + manifest_identity: tuple[int, int] | None = None + try: + manifest_temporary, manifest_temporary_identity, manifest_descriptor = _temporary_payload( + manifest_path, manifest + ) + assert archive_temporary is not None + archive_identity = _publish_no_clobber( + archive_temporary, archive_path, archive_temporary_identity, archive_descriptor + ) + archive_temporary = None + manifest_identity = _publish_no_clobber( + manifest_temporary, + manifest_path, + manifest_temporary_identity, + manifest_descriptor, + ) + manifest_temporary = None + if ( + stable_file_sha256(archive_path, label="published source archive") + != sha256(archive).hexdigest() + or stable_file_sha256(manifest_path, label="published source manifest") + != sha256(manifest).hexdigest() + ): + raise SourceManifestError("published source artifacts failed stable rehash") + except BaseException as error: + if manifest_identity is not None: + _unlink_if_identity(manifest_path, manifest_identity) + if archive_identity is not None: + _unlink_if_identity(archive_path, archive_identity) + if manifest_temporary is not None and manifest_temporary_identity is not None: + _unlink_if_identity(manifest_temporary, manifest_temporary_identity) + if archive_temporary is not None: + _unlink_if_identity(archive_temporary, archive_temporary_identity) + if isinstance(error, FileExistsError): + raise SourceManifestError( + "source artifact destination appeared during publication" + ) from error + if isinstance(error, CheckpointManifestError): + raise SourceManifestError("could not rehash published source artifacts") from error + raise + finally: + if manifest_descriptor is not None: + os.close(manifest_descriptor) + os.close(archive_descriptor) + + +def _manifest_from_git_archive( + archive: bytes, *, source_kind: str, git_commit: str, git_tree: str +) -> bytes: + directories: set[str] = set() + files: dict[str, dict[str, object]] = {} + try: + with tarfile.open(fileobj=BytesIO(archive), mode="r:") as handle: + for member in handle.getmembers(): + relative = _relative_path(member.name.rstrip("/"), "Git archive path") + if member.isdir(): + directories.add(relative) + elif member.isfile(): + extracted = handle.extractfile(member) + if extracted is None: + raise SourceManifestError(f"could not read Git archive file {relative!r}") + payload = extracted.read() + if len(payload) != member.size: + raise SourceManifestError(f"Git archive file {relative!r} is truncated") + files[relative] = { + "path": relative, + "mode": "100755" if member.mode & 0o111 else "100644", + "size_bytes": len(payload), + "sha256": sha256(payload).hexdigest(), + } + else: + raise SourceManifestError( + f"Git archive path {relative!r} is not a regular file or directory" + ) + except (tarfile.TarError, OSError) as error: + raise SourceManifestError("could not parse exact Git archive") from error + if not files: + raise SourceManifestError("Git archive contains no source files") + return _canonical_json_bytes( + { + "source_manifest_schema_version": 1, + "source_kind": _canonical_text(source_kind, "source kind"), + "git_commit": _hex_digest(git_commit, 40, "Git commit"), + "git_tree": _hex_digest(git_tree, 40, "Git tree"), + "git_archive_sha256": sha256(archive).hexdigest(), + "archive_scope": _ARCHIVE_SCOPE, + "directories": sorted(directories), + "files": [files[path] for path in sorted(files)], + } + ) + + +def create_source_manifest_from_git_archive( + checkout: str | Path, + *, + expected_commit: str, + source_kind: str, + archive_output: str | Path, + manifest_output: str | Path, +) -> GeneratedSourceManifest: + """Create a sealed runtime-source witness and its exact scoped Git archive.""" + root = Path(checkout).expanduser().resolve() + if not root.is_dir(): + raise SourceManifestError("Git checkout must be a local directory") + commit = _hex_digest(expected_commit, 40, "expected Git commit") + top_level = _run_git(root, ["rev-parse", "--show-toplevel"]).decode().strip() + if Path(top_level).resolve() != root: + raise SourceManifestError("Git checkout must be the repository root") + if _run_git(root, ["rev-parse", "HEAD"]).decode().strip() != commit: + raise SourceManifestError("Git checkout HEAD does not match the expected commit") + status_arguments = ["status", "--porcelain=v1", "--untracked-files=all"] + if _run_git(root, status_arguments).strip(): + raise SourceManifestError("Git checkout has tracked or untracked modifications") + tree = _hex_digest( + _run_git(root, ["rev-parse", "HEAD^{tree}"]).decode().strip(), 40, "Git tree" + ) + archive = _run_git(root, ["archive", "--format=tar", commit, "--", _ARCHIVE_SCOPE]) + manifest = _manifest_from_git_archive( + archive, source_kind=source_kind, git_commit=commit, git_tree=tree + ) + if ( + _run_git(root, ["rev-parse", "HEAD"]).decode().strip() != commit + or _run_git(root, ["rev-parse", "HEAD^{tree}"]).decode().strip() != tree + or _run_git(root, status_arguments).strip() + ): + raise SourceManifestError("Git checkout changed while its archive was generated") + + archive_path = Path(archive_output).expanduser().resolve() + manifest_path = Path(manifest_output).expanduser().resolve() + if archive_path == manifest_path: + raise SourceManifestError("archive and source-manifest outputs must differ") + _publish_source_artifacts(archive_path, archive, manifest_path, manifest) + return GeneratedSourceManifest( + git_commit=commit, + git_tree=tree, + git_archive_sha256=sha256(archive).hexdigest(), + manifest_sha256=sha256(manifest).hexdigest(), + ) diff --git a/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 index feacf8fda53..a1ec161804e 100644 --- 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 @@ -62,8 +62,8 @@ def _base_args(tmp_path: Path): capture_manifest.write_bytes( _canonical( { - "capture_manifest_schema_version": 3, - "capture_protocol": "modelopt_vllm_mask_reuse_target_sparsity_v3", + "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), @@ -72,6 +72,13 @@ def _base_args(tmp_path: Path): "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()], diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py b/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py index 44628328086..c5754b43e76 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_collect_mask_reuse_cli.py @@ -25,6 +25,7 @@ import pytest +from modelopt.torch.sparsity.attention_sparsity.calibration import source_manifest from modelopt.torch.sparsity.attention_sparsity.calibration.checkpoint_manifest import ( create_checkpoint_manifest, ) @@ -38,6 +39,29 @@ collect_mask_reuse_cli = importlib.util.module_from_spec(_SPEC) _SPEC.loader.exec_module(collect_mask_reuse_cli) +_FA4_COMMIT = "a" * 40 +_FA4_TREE = "b" * 40 +_FA4_ARCHIVE_SHA256 = "c" * 64 + + +def _write_source_witness(source: Path, destination: Path) -> str: + snapshot = source_manifest._source_tree_snapshot(source) + raw = { + "source_manifest_schema_version": 1, + "source_kind": "flash-attention-4", + "git_commit": _FA4_COMMIT, + "git_tree": _FA4_TREE, + "git_archive_sha256": _FA4_ARCHIVE_SHA256, + "archive_scope": "flash_attn", + "directories": list(snapshot.directories), + "files": list(snapshot.files), + } + payload = ( + json.dumps(raw, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + destination.write_bytes(payload) + return sha256(payload).hexdigest() + class _Tokenizer: def encode(self, prompt, *, add_special_tokens): @@ -187,6 +211,8 @@ def test_main_bootstraps_policy_free_backend_and_writes_normalized_evidence(tmp_ cute.mkdir(parents=True) (cute / "interface.py").write_text("# pinned\n", encoding="utf-8") (cute / "block_sparsity.py").write_text("# pinned\n", encoding="utf-8") + fa4_source_manifest = tmp_path / "fa4-source-manifest.json" + fa4_source_manifest_sha256 = _write_source_witness(fa4_source, fa4_source_manifest) fake_vllm = SimpleNamespace(LLM=_LLM, SamplingParams=_SamplingParams) monkeypatch.setitem(sys.modules, "vllm", fake_vllm) monkeypatch.setattr( @@ -196,13 +222,6 @@ def test_main_bootstraps_policy_free_backend_and_writes_normalized_evidence(tmp_ SimpleNamespace(name="mask_reuse_fa4", value="mask_reuse_vllm.plugin:register") ], ) - monkeypatch.setattr( - collect_mask_reuse_cli.subprocess, - "run", - lambda command, **kwargs: SimpleNamespace( - stdout="a" * 40 + "\n" if command[-2:] == ["rev-parse", "HEAD"] else "" - ), - ) monkeypatch.setenv("MASK_REUSE_FA4_POLICY", "stale-policy.json") monkeypatch.setenv("MASK_REUSE_FA4_POLICY_SHA256", "stale") parsed_payloads = {} @@ -226,10 +245,18 @@ def parse_vanilla(payload): str(checkpoint), "--model-id", "test-model", + "--checkpoint-manifest-sha256", + checkpoint_manifest.sha256, "--plan", "test_stride2", "--fa4-source", str(fa4_source), + "--fa4-source-manifest", + str(fa4_source_manifest), + "--fa4-source-manifest-sha256", + fa4_source_manifest_sha256, + "--fa4-commit", + _FA4_COMMIT, "--prompts-jsonl", str(prompts), "--vanilla-config", @@ -282,12 +309,16 @@ def parse_vanilla(payload): } assert {capture["invocation"]["target_sparsity_hex"] for capture in captures} == {(0.7).hex()} report = json.loads(manifest.read_text()) - assert report["capture_manifest_schema_version"] == 3 - assert report["capture_protocol"] == "modelopt_vllm_mask_reuse_target_sparsity_v3" + assert report["capture_manifest_schema_version"] == 4 + assert report["capture_protocol"] == "modelopt_vllm_mask_reuse_target_sparsity_v4" assert report["checkpoint_manifest_sha256"] == checkpoint_manifest.sha256 assert report["prompt_plan_file_sha256"] == sha256(parsed_payloads["prompts"]).hexdigest() assert report["vanilla_config_file_sha256"] == sha256(parsed_payloads["vanilla"]).hexdigest() - assert report["fa4_source_commit"] == "a" * 40 + assert report["fa4_source_commit"] == _FA4_COMMIT + assert report["fa4_source_git_tree"] == _FA4_TREE + assert report["fa4_source_git_archive_sha256"] == _FA4_ARCHIVE_SHA256 + assert report["fa4_source_manifest_sha256"] == fa4_source_manifest_sha256 + assert report["fa4_source_file_count"] == 2 assert report["dense_shadow_validation_requested"] is True assert report["engine_kwargs"]["tensor_parallel_size"] == 1 assert report["capture_count"] == len(captures) @@ -336,31 +367,26 @@ def fail_fsync(path): assert not temporary.exists() -def test_capture_environment_rejects_untracked_fa4_source(tmp_path, monkeypatch): +def test_capture_environment_rejects_extra_fa4_source_path(tmp_path, monkeypatch): fa4_source = tmp_path / "flash-attention" cute = fa4_source / "flash_attn/cute" cute.mkdir(parents=True) (cute / "interface.py").write_text("# pinned\n", encoding="utf-8") (cute / "block_sparsity.py").write_text("# pinned\n", encoding="utf-8") - monkeypatch.setattr( - collect_mask_reuse_cli.subprocess, - "run", - lambda command, **kwargs: SimpleNamespace( - stdout=( - "a" * 40 + "\n" - if command[-2:] == ["rev-parse", "HEAD"] - else "?? flash_attn/cute/local_override.py\n" - ) - ), - ) + manifest = tmp_path / "fa4-source-manifest.json" + digest = _write_source_witness(fa4_source, manifest) + (cute / "local_override.py").write_text("# extra\n", encoding="utf-8") with pytest.raises( collect_mask_reuse_cli.CaptureContractError, - match="tracked or untracked modifications", + match="does not exactly match", ): collect_mask_reuse_cli._configure_capture_environment( "test_stride2", str(fa4_source), + str(manifest), + digest, + _FA4_COMMIT, "0" * 64, validate_dense_output=True, ) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py b/tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py new file mode 100644 index 00000000000..993916a1bce --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py @@ -0,0 +1,315 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for exact, Git-free source-tree witnesses.""" + +import json +import shutil +import subprocess +import tarfile +from hashlib import sha256 +from pathlib import Path + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration import source_manifest +from modelopt.torch.sparsity.attention_sparsity.calibration.source_manifest import ( + SourceManifestError, + create_source_manifest_from_git_archive, + verify_source_manifest, +) + +_COMMIT = "a" * 40 +_TREE = "b" * 40 +_ARCHIVE_SHA256 = "c" * 64 + + +def _toy_source(root: Path) -> Path: + (root / "flash_attn/cute").mkdir(parents=True) + (root / "flash_attn/empty-dir").mkdir() + (root / "flash_attn/cute/interface.py").write_text("# interface\n", encoding="utf-8") + (root / "flash_attn/cute/block_sparsity.py").write_text("# sparse\n", encoding="utf-8") + return root + + +def _write_witness(source: Path, manifest: Path) -> str: + snapshot = source_manifest._source_tree_snapshot(source) + raw = { + "source_manifest_schema_version": 1, + "source_kind": "flash-attention-4", + "git_commit": _COMMIT, + "git_tree": _TREE, + "git_archive_sha256": _ARCHIVE_SHA256, + "archive_scope": "flash_attn", + "directories": list(snapshot.directories), + "files": list(snapshot.files), + } + payload = ( + json.dumps(raw, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + manifest.write_bytes(payload) + return sha256(payload).hexdigest() + + +def _verify(source: Path, manifest: Path, digest: str): + return verify_source_manifest( + source, + manifest, + expected_manifest_sha256=digest, + expected_commit=_COMMIT, + expected_source_kind="flash-attention-4", + ) + + +def _symlink_or_skip(path: Path, target: str) -> None: + try: + path.symlink_to(target) + except OSError as error: + pytest.skip(f"symbolic links unavailable: {error}") + + +def test_source_witness_accepts_exact_files_and_directories(tmp_path): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + digest = _write_witness(source, manifest) + + verified = _verify(source, manifest, digest) + + assert verified.git_commit == _COMMIT + assert verified.git_tree == _TREE + assert verified.manifest_sha256 == digest + assert verified.file_count == 2 + assert verified.directory_count == 3 + + +@pytest.mark.parametrize("change", ["mutated", "missing", "extra", "mode"]) +def test_source_witness_rejects_file_tree_changes(tmp_path, change): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + digest = _write_witness(source, manifest) + interface = source / "flash_attn/cute/interface.py" + if change == "mutated": + interface.write_text("# modified\n", encoding="utf-8") + elif change == "missing": + interface.unlink() + elif change == "extra": + (source / "flash_attn/cute/shadow.py").write_text("# extra\n", encoding="utf-8") + else: + interface.chmod(interface.stat().st_mode | 0o111) + + with pytest.raises(SourceManifestError, match="does not exactly match"): + _verify(source, manifest, digest) + + +def test_source_witness_rejects_regular_file_replaced_by_symlink(tmp_path): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + digest = _write_witness(source, manifest) + interface = source / "flash_attn/cute/interface.py" + interface.unlink() + _symlink_or_skip(interface, "block_sparsity.py") + + with pytest.raises(SourceManifestError, match="must not be a link"): + _verify(source, manifest, digest) + + +def test_source_witness_rejects_extra_symlink(tmp_path): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + digest = _write_witness(source, manifest) + _symlink_or_skip(source / "flash_attn/unexpected-link", "cute/interface.py") + with pytest.raises(SourceManifestError, match="must not be a link"): + _verify(source, manifest, digest) + + +def test_source_witness_requires_independent_manifest_and_commit_pins(tmp_path): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + digest = _write_witness(source, manifest) + + with pytest.raises(SourceManifestError, match="expected SHA256"): + _verify(source, manifest, "0" * 64) + with pytest.raises(SourceManifestError, match="expected Git commit"): + verify_source_manifest( + source, + manifest, + expected_manifest_sha256=digest, + expected_commit="0" * 40, + expected_source_kind="flash-attention-4", + ) + + +def test_source_witness_rejects_boolean_schema_version(tmp_path): + source = _toy_source(tmp_path / "source") + manifest = tmp_path / "source-manifest.json" + _write_witness(source, manifest) + raw = json.loads(manifest.read_bytes()) + raw["source_manifest_schema_version"] = True + payload = ( + json.dumps(raw, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + "\n" + ).encode() + manifest.write_bytes(payload) + + with pytest.raises(SourceManifestError, match="schema_version must be 1"): + _verify(source, manifest, sha256(payload).hexdigest()) + + +def test_source_artifact_temporary_is_cleaned_when_write_setup_fails(tmp_path, monkeypatch): + destination = tmp_path / "source.tar" + real_dup = source_manifest.os.dup + + def fail_dup(descriptor): + assert descriptor >= 0 + raise OSError("injected dup failure") + + monkeypatch.setattr(source_manifest.os, "dup", fail_dup) + + with pytest.raises(OSError, match="injected dup failure"): + source_manifest._temporary_payload(destination, b"archive") + + monkeypatch.setattr(source_manifest.os, "dup", real_dup) + assert list(tmp_path.iterdir()) == [] + + +def test_source_artifact_publication_preserves_destination_racer(tmp_path, monkeypatch): + archive = tmp_path / "source.tar" + manifest = tmp_path / "source-manifest.json" + real_link = source_manifest.os.link + link_count = 0 + + def racing_link(source, destination, **kwargs): + nonlocal link_count + link_count += 1 + if link_count == 2: + Path(destination).write_bytes(b"racer") + return real_link(source, destination, **kwargs) + + monkeypatch.setattr(source_manifest.os, "link", racing_link) + + with pytest.raises(SourceManifestError, match="destination appeared"): + source_manifest._publish_source_artifacts(archive, b"archive", manifest, b"manifest") + + assert not archive.exists() + assert manifest.read_bytes() == b"racer" + assert sorted(path.name for path in tmp_path.iterdir()) == [manifest.name] + + +def test_source_artifact_publication_preserves_post_publish_replacement(tmp_path, monkeypatch): + archive = tmp_path / "source.tar" + manifest = tmp_path / "source-manifest.json" + real_stable_hash = source_manifest.stable_file_sha256 + + def replace_archive_before_rehash(path, *, label): + if label == "published source archive": + Path(path).unlink() + Path(path).write_bytes(b"racer") + return real_stable_hash(path, label=label) + + monkeypatch.setattr(source_manifest, "stable_file_sha256", replace_archive_before_rehash) + + with pytest.raises(SourceManifestError, match="failed stable rehash"): + source_manifest._publish_source_artifacts(archive, b"archive", manifest, b"manifest") + + assert archive.read_bytes() == b"racer" + assert not manifest.exists() + assert sorted(path.name for path in tmp_path.iterdir()) == [archive.name] + + +@pytest.mark.skipif(shutil.which("git") is None, reason="Git is required by the generator") +def test_generator_witnesses_the_exact_archive_without_runtime_git(tmp_path): + checkout = tmp_path / "checkout" + checkout.mkdir() + subprocess.run(["git", "init", "-q", str(checkout)], check=True) + subprocess.run(["git", "-C", str(checkout), "config", "user.name", "Test"], check=True) + subprocess.run( + ["git", "-C", str(checkout), "config", "user.email", "test@example.com"], check=True + ) + _toy_source(checkout) + subprocess.run(["git", "-C", str(checkout), "add", "."], check=True) + subprocess.run(["git", "-C", str(checkout), "commit", "-qm", "source"], check=True) + commit = subprocess.run( + ["git", "-C", str(checkout), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + archive = tmp_path / "source.tar" + manifest = tmp_path / "source-manifest.json" + + generated = create_source_manifest_from_git_archive( + checkout, + expected_commit=commit, + source_kind="flash-attention-4", + archive_output=archive, + manifest_output=manifest, + ) + extracted = tmp_path / "extracted" + extracted.mkdir() + with tarfile.open(archive, "r:") as handle: + for member in handle.getmembers(): + path = extracted / member.name + if member.isdir(): + path.mkdir(parents=True, exist_ok=True) + else: + payload = handle.extractfile(member) + assert payload is not None + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(payload.read()) + path.chmod(member.mode) + + verified = verify_source_manifest( + extracted, + manifest, + expected_manifest_sha256=generated.manifest_sha256, + expected_commit=commit, + expected_source_kind="flash-attention-4", + ) + assert verified.git_archive_sha256 == generated.git_archive_sha256 + + +@pytest.mark.skipif(shutil.which("git") is None, reason="Git is required by the generator") +def test_generator_archive_excludes_ignored_checkout_artifacts(tmp_path): + checkout = tmp_path / "checkout" + checkout.mkdir() + subprocess.run(["git", "init", "-q", str(checkout)], check=True) + subprocess.run(["git", "-C", str(checkout), "config", "user.name", "Test"], check=True) + subprocess.run( + ["git", "-C", str(checkout), "config", "user.email", "test@example.com"], check=True + ) + (checkout / ".gitignore").write_text("*.pyc\n", encoding="utf-8") + (checkout / "flash_attn").mkdir() + (checkout / "flash_attn/source.py").write_text("# source\n", encoding="utf-8") + subprocess.run(["git", "-C", str(checkout), "add", "."], check=True) + subprocess.run(["git", "-C", str(checkout), "commit", "-qm", "source"], check=True) + (checkout / "flash_attn/shadow.pyc").write_bytes(b"ignored") + commit = subprocess.run( + ["git", "-C", str(checkout), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + archive = tmp_path / "source.tar" + manifest = tmp_path / "source-manifest.json" + generated = create_source_manifest_from_git_archive( + checkout, + expected_commit=commit, + source_kind="flash-attention-4", + archive_output=archive, + manifest_output=manifest, + ) + assert generated.manifest_sha256 == sha256(manifest.read_bytes()).hexdigest() + with tarfile.open(archive, "r:") as handle: + assert "flash_attn/shadow.pyc" not in handle.getnames() From 86c0306088b59c00c1bd661140cc4c3be604a707 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 1 Aug 2026 11:00:25 -0700 Subject: [PATCH 16/16] Fix source witness publication on Windows Signed-off-by: Kai Xu --- .../calibration/source_manifest.py | 59 ++++++++++++------- .../test_source_manifest.py | 14 ++++- 2 files changed, 50 insertions(+), 23 deletions(-) diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/source_manifest.py b/modelopt/torch/sparsity/attention_sparsity/calibration/source_manifest.py index 9748b086c47..837bbf43ca3 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/source_manifest.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/source_manifest.py @@ -407,15 +407,18 @@ def _temporary_payload(destination: Path, payload: bytes) -> tuple[Path, tuple[i identity = (opened.st_dev, opened.st_ino) except OSError: pass + if os.name == "nt": + os.close(descriptor) if identity is not None: _unlink_if_identity(temporary, identity) - os.close(descriptor) + if os.name != "nt": + os.close(descriptor) raise def _publish_no_clobber( temporary: Path, destination: Path, identity: tuple[int, int], descriptor: int -) -> tuple[int, int]: +) -> None: observed = temporary.stat(follow_symlinks=False) opened = os.fstat(descriptor) if ( @@ -425,23 +428,16 @@ def _publish_no_clobber( ): raise SourceManifestError("source artifact temporary file changed before publication") os.link(temporary, destination, follow_symlinks=False) - try: - published = destination.stat(follow_symlinks=False) - opened_after = os.fstat(descriptor) - if ( - published.st_ino == 0 - or not stat.S_ISREG(published.st_mode) - or _is_reparse_point(published) - or (published.st_dev, published.st_ino) != identity - or (opened_after.st_dev, opened_after.st_ino) != identity - ): - raise SourceManifestError("source artifact destination changed during publication") - _unlink_if_identity(temporary, identity) - _fsync_directory(destination.parent) - except BaseException: - _unlink_if_identity(destination, identity) - raise - return identity + published = destination.stat(follow_symlinks=False) + opened_after = os.fstat(descriptor) + if ( + published.st_ino == 0 + or not stat.S_ISREG(published.st_mode) + or _is_reparse_point(published) + or (published.st_dev, published.st_ino) != identity + or (opened_after.st_dev, opened_after.st_ino) != identity + ): + raise SourceManifestError("source artifact destination changed during publication") def _publish_source_artifacts( @@ -464,16 +460,28 @@ def _publish_source_artifacts( manifest_path, manifest ) assert archive_temporary is not None - archive_identity = _publish_no_clobber( + archive_identity = archive_temporary_identity + _publish_no_clobber( archive_temporary, archive_path, archive_temporary_identity, archive_descriptor ) + if os.name == "nt": + os.close(archive_descriptor) + archive_descriptor = None + _unlink_if_identity(archive_temporary, archive_temporary_identity) + _fsync_directory(archive_path.parent) archive_temporary = None - manifest_identity = _publish_no_clobber( + manifest_identity = manifest_temporary_identity + _publish_no_clobber( manifest_temporary, manifest_path, manifest_temporary_identity, manifest_descriptor, ) + if os.name == "nt": + os.close(manifest_descriptor) + manifest_descriptor = None + _unlink_if_identity(manifest_temporary, manifest_temporary_identity) + _fsync_directory(manifest_path.parent) manifest_temporary = None if ( stable_file_sha256(archive_path, label="published source archive") @@ -483,6 +491,12 @@ def _publish_source_artifacts( ): raise SourceManifestError("published source artifacts failed stable rehash") except BaseException as error: + if os.name == "nt" and manifest_descriptor is not None: + os.close(manifest_descriptor) + manifest_descriptor = None + if os.name == "nt" and archive_descriptor is not None: + os.close(archive_descriptor) + archive_descriptor = None if manifest_identity is not None: _unlink_if_identity(manifest_path, manifest_identity) if archive_identity is not None: @@ -501,7 +515,8 @@ def _publish_source_artifacts( finally: if manifest_descriptor is not None: os.close(manifest_descriptor) - os.close(archive_descriptor) + if archive_descriptor is not None: + os.close(archive_descriptor) def _manifest_from_git_archive( diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py b/tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py index 993916a1bce..c784d32308f 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_source_manifest.py @@ -16,6 +16,7 @@ """Tests for exact, Git-free source-tree witnesses.""" import json +import os import shutil import subprocess import tarfile @@ -94,7 +95,18 @@ def test_source_witness_accepts_exact_files_and_directories(tmp_path): assert verified.directory_count == 3 -@pytest.mark.parametrize("change", ["mutated", "missing", "extra", "mode"]) +@pytest.mark.parametrize( + "change", + [ + "mutated", + "missing", + "extra", + pytest.param( + "mode", + marks=pytest.mark.skipif(os.name == "nt", reason="Windows has no POSIX execute bit"), + ), + ], +) def test_source_witness_rejects_file_tree_changes(tmp_path, change): source = _toy_source(tmp_path / "source") manifest = tmp_path / "source-manifest.json"