From 4bedb5de67d7276541551f676ccd1ba79e4f6a45 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 20:49:46 -0700 Subject: [PATCH 01/14] 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 a787c2363277dca7fd1814a13c13b66619032534 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 20:56:22 -0700 Subject: [PATCH 02/14] Add paged vLLM skip-softmax calibration - attention_calibrate: read K/V through a paged cache (vLLM NHD layout) via block_table, reusing the shared paged tile loaders. Contiguous-KV behavior is unchanged; paged and contiguous launches produce identical counters and output. - vLLM adapters: calibration mode (enable_calibration / disable_calibration / iter_sparse_impls). When active, forward routes to the paged calibration kernel -- full dense attention plus multi-threshold tile-skip counting -- so generation is numerically unchanged while stats accumulate. Each scheduled request is measured independently (decode vs prefill phase per request), and raw per-threshold tile counts are recorded so tensor-parallel ranks can be aggregated by summing counts before the fit. - FlashInfer adapter writes the current K/V to the cache before the calibrate kernel reads it (releases that update the cache inside forward); fp16/bf16-only and NHD-only, both validated explicitly. Signed-off-by: Kai Xu --- .../kernels/sparsity/attention/calibrate.py | 132 +++++++++++-- .../attention_sparsity/plugins/vllm.py | 183 ++++++++++++++++++ 2 files changed, 301 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/kernels/sparsity/attention/calibrate.py b/modelopt/torch/kernels/sparsity/attention/calibrate.py index d26e781d48c..68748e7c50c 100644 --- a/modelopt/torch/kernels/sparsity/attention/calibrate.py +++ b/modelopt/torch/kernels/sparsity/attention/calibrate.py @@ -28,7 +28,12 @@ import triton import triton.language as tl -from modelopt.torch.kernels.common.attention.triton_fa import LOG2E, _apply_mask +from modelopt.torch.kernels.common.attention.triton_fa import ( + LOG2E, + _apply_mask, + _load_paged_k_tile, + _load_paged_v_tile, +) # --------------------------------------------------------------------------- @@ -64,6 +69,18 @@ def _attn_fwd_calibrate( HEAD_DIM: tl.constexpr, NUM_THRESHOLDS: tl.constexpr, PADDED_THRESHOLDS: tl.constexpr, # next_power_of_2(NUM_THRESHOLDS) for tl.arange + IS_PAGED: tl.constexpr = False, # Whether K/V are read from a paged KV cache + K_cache=None, # [num_blocks, page_size, num_kv_heads, head_dim] paged K + V_cache=None, # [num_blocks, page_size, num_kv_heads, head_dim] paged V + Block_table=None, # [batch, max_blocks_per_seq] page table + stride_kc_block=0, + stride_kc_pos=0, + stride_kc_head=0, + stride_vc_block=0, + stride_vc_pos=0, + stride_vc_head=0, + PAGE_SIZE: tl.constexpr = 16, + max_blocks_per_seq=0, ): """Forward kernel with multi-threshold sparsity measurement. @@ -126,12 +143,33 @@ def _attn_fwd_calibrate( for kv_start in range(0, kv_bound, BLOCK_N): kv_start = tl.multiple_of(kv_start, BLOCK_N) - k_offs = (kv_offset + kv_start + kv_pos[None, :]) * stride_kbs + dim_pos[:, None] - k = tl.load( - k_base + k_offs, - mask=((kv_start + kv_pos[None, :]) < seq_len_kv) & d_mask[:, None], - other=0.0, - ) + # Load K^T [BLOCK_D, BLOCK_N] from paged cache or contiguous K. + if IS_PAGED: + k = _load_paged_k_tile( + K_cache, + Block_table, + batch_idx, + kv_head_idx, + kv_start, + kv_pos, + dim_pos, + seq_len_kv, + stride_kc_block, + stride_kc_pos, + stride_kc_head, + PAGE_SIZE, + BLOCK_N, + BLOCK_D, + HEAD_DIM, + max_blocks_per_seq, + ) + else: + k_offs = (kv_offset + kv_start + kv_pos[None, :]) * stride_kbs + dim_pos[:, None] + k = tl.load( + k_base + k_offs, + mask=((kv_start + kv_pos[None, :]) < seq_len_kv) & d_mask[:, None], + other=0.0, + ) scores = tl.dot(q, k) * qk_scale scores = _apply_mask(scores, q_pos, kv_pos, seq_len_q, seq_len_kv, kv_start, IS_CAUSAL) @@ -164,12 +202,32 @@ def _attn_fwd_calibrate( row_sum = row_sum * correction + l_new acc = acc * correction[:, None] - v_offs = (kv_offset + kv_start + kv_pos[:, None]) * stride_vbs + dim_pos[None, :] - v = tl.load( - v_base + v_offs, - mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], - other=0.0, - ) + if IS_PAGED: + v = _load_paged_v_tile( + V_cache, + Block_table, + batch_idx, + kv_head_idx, + kv_start, + kv_pos, + dim_pos, + seq_len_kv, + stride_vc_block, + stride_vc_pos, + stride_vc_head, + PAGE_SIZE, + BLOCK_N, + BLOCK_D, + HEAD_DIM, + max_blocks_per_seq, + ) + else: + v_offs = (kv_offset + kv_start + kv_pos[:, None]) * stride_vbs + dim_pos[None, :] + v = tl.load( + v_base + v_offs, + mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], + other=0.0, + ) acc = tl.dot(p.to(v.dtype), v, acc) row_max = m_new @@ -212,6 +270,10 @@ def attention_calibrate( max_input_len_k: int | None = None, *, threshold_trials: list[float] | None = None, + k_cache: torch.Tensor | None = None, + v_cache: torch.Tensor | None = None, + block_table: torch.Tensor | None = None, + page_size: int = 16, ) -> tuple[torch.Tensor, torch.Tensor]: """Flash attention with multi-threshold skip-softmax sparsity measurement. @@ -219,12 +281,21 @@ def attention_calibrate( measuring how many KV tiles would be skipped at each threshold in ``threshold_trials``. No autograd — forward only. - All arguments except ``threshold_trials`` match + All positional arguments match :func:`modelopt.torch.kernels.common.attention.attention`. Args: threshold_trials: List of threshold values to measure sparsity for. Each value is converted to log2-scaled space for the kernel. + k_cache: Paged K cache ``[num_blocks, page_size, num_kv_heads, head_dim]``. + When provided, K/V are read from the paged cache via ``block_table`` + (vLLM NHD layout) instead of from the contiguous ``k``/``v`` tensors. + ``k``/``v`` are then dummies whose only meaningful dimension is + ``shape[1] == num_kv_heads`` (used to compute the GQA ratio). + v_cache: Paged V cache ``[num_blocks, page_size, num_kv_heads, head_dim]``. + block_table: Page table ``[batch, max_blocks_per_seq]`` mapping each + sequence's block indices to global page IDs. + page_size: Number of tokens per page in the KV cache. Returns: Tuple of ``(output, sparsity_counters)``: @@ -237,6 +308,10 @@ def attention_calibrate( if threshold_trials is None or len(threshold_trials) == 0: raise ValueError("threshold_trials must be a non-empty list") + is_paged = k_cache is not None + if is_paged and block_table is None: + raise ValueError("block_table is required when k_cache/v_cache are provided.") + # Calibration has only been validated with uniform-length batches (current # diffusion + RULER paths). Varlen inputs would exercise code paths in the # kernel that have not been tested — fail loudly rather than silently @@ -281,6 +356,11 @@ def attention_calibrate( b_seq_len_k = b_seq_len b_start_loc_k = b_start_loc + # Paged mode: KV positions come from block_table, so the contiguous KV + # offsets are unused. Provide a dummy so Triton can compile the tl.load. + if b_start_loc_k is None: + b_start_loc_k = torch.zeros_like(b_start_loc) + num_thresholds = len(threshold_trials) # Scores already include sm_scale and LOG2E; convert lambda to log2 space only. @@ -304,6 +384,18 @@ def attention_calibrate( num_programs * num_thresholds, dtype=torch.int32, device=q.device ) + # Paged KV cache strides (zeros when not paged; computed here so the type + # narrowing of k_cache/v_cache/block_table is explicit for the kernel call). + if is_paged: + assert k_cache is not None and v_cache is not None and block_table is not None + kc_strides = (k_cache.stride(0), k_cache.stride(1), k_cache.stride(2)) + vc_strides = (v_cache.stride(0), v_cache.stride(1), v_cache.stride(2)) + max_blocks_per_seq = block_table.shape[1] + else: + kc_strides = (0, 0, 0) + vc_strides = (0, 0, 0) + max_blocks_per_seq = 0 + # Triton launches on torch.cuda.current_device(), which is not necessarily # the device the tensors live on (e.g. under accelerate device_map="auto" # sharding). Activate the tensor's device so the kernel dereferences the @@ -338,6 +430,18 @@ def attention_calibrate( HEAD_DIM=HEAD_DIM, NUM_THRESHOLDS=num_thresholds, PADDED_THRESHOLDS=triton.next_power_of_2(num_thresholds), + IS_PAGED=is_paged, + K_cache=k_cache, + V_cache=v_cache, + Block_table=block_table, + stride_kc_block=kc_strides[0], + stride_kc_pos=kc_strides[1], + stride_kc_head=kc_strides[2], + stride_vc_block=vc_strides[0], + stride_vc_pos=vc_strides[1], + stride_vc_head=vc_strides[2], + PAGE_SIZE=page_size, + max_blocks_per_seq=max_blocks_per_seq, num_warps=4, num_stages=1, ) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 243db16a2bc..a3a5ebe1271 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -45,6 +45,7 @@ ) from modelopt.torch.kernels.common.attention.triton_fa import attention as triton_attention from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite +from modelopt.torch.kernels.sparsity.attention.calibrate import attention_calibrate @functools.cache @@ -252,6 +253,104 @@ def _resolve_forward( ) +def _calibration_active(impl) -> bool: + """Return whether skip-softmax calibration mode is enabled on an impl.""" + return bool(getattr(impl, "_calibrate", False)) and bool( + getattr(impl, "_calib_threshold_trials", None) + ) + + +def _forward_calibrate( + impl, + *, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + num_actual_tokens: int, + output: torch.Tensor, +) -> torch.Tensor: + """Measure per-request tile-skip stats via the paged Triton calibration kernel. + + Each scheduled request is calibrated independently (batch=1) so its KV + length is the per-sample length the exponential fit needs, and so the + kernel keeps the uniform-length contract it was validated against. The + kernel computes full attention, so ``output`` is written densely and the + forward pass is numerically unchanged. + + Phase and causality are decided per request: ``q_len == 1`` is a decode + step (full-cache, non-causal); ``q_len > 1`` is (chunked) prefill (causal + — the kernel offsets the query into the KV span). A mixed prefill/decode + batch therefore contributes correctly to both phase fits. + + Records raw per-threshold tile counts (not ratios) on + ``impl._calib_records`` so tensor-parallel workers can be aggregated by + summing counts before the fit. + """ + if key_cache.dtype not in (torch.float16, torch.bfloat16): + raise NotImplementedError( + f"skip-softmax calibration requires an fp16/bf16 KV cache, got {key_cache.dtype}" + ) + if key_cache.shape[2] != impl.num_kv_heads: + # NHD is the only supported paged layout: [blocks, page, kv_heads, dim]. + # An HND FlashInfer cache would put kv_heads on axis 1. + raise NotImplementedError( + f"KV cache layout is not NHD (shape {tuple(key_cache.shape)}, " + f"expected axis 2 == num_kv_heads == {impl.num_kv_heads}); " + "HND caches are unsupported for calibration" + ) + page_size = key_cache.shape[1] + trials = impl._calib_threshold_trials + batch = seq_lens.shape[0] + b_start_loc = cu_seqlens_q[:batch] + b_seq_len = cu_seqlens_q[1 : batch + 1] - cu_seqlens_q[:batch] + + q = query[:num_actual_tokens].contiguous() + # Dummy K/V: in paged mode KV is read from the cache via block_table. + # Only shape[1] (num_kv_heads) is consulted, to compute the GQA ratio. + k_dummy = torch.empty(0, impl.num_kv_heads, impl.head_size, device=q.device, dtype=q.dtype) + + for i in range(batch): + q_len = int(b_seq_len[i].item()) + if q_len <= 0: + continue + q_start = int(b_start_loc[i].item()) + seq_k = int(seq_lens[i].item()) + phase = "decode" if q_len <= 1 else "prefill" + + oi, counters = attention_calibrate( + q[q_start : q_start + q_len], + k_dummy, + k_dummy, + b_start_loc=torch.zeros(1, device=q.device, dtype=torch.int32), + b_seq_len=b_seq_len[i : i + 1].to(torch.int32), + max_input_len=q_len, + is_causal=q_len > 1, + softmax_scale=impl.scale, + b_seq_len_k=seq_lens[i : i + 1].to(torch.int32), + max_input_len_k=seq_k, + threshold_trials=trials, + k_cache=key_cache, + v_cache=value_cache, + block_table=block_table[i : i + 1], + page_size=page_size, + ) + output[q_start : q_start + q_len] = oi + + impl._calib_records.append( + { + "phase": phase, + "sample_length": seq_k, + "total_tiles": counters[:, 0].tolist(), + "skipped_tiles": counters[:, 1].tolist(), + } + ) + + return output + + # Resolution guards raw configured transforms; dispatch rechecks effective # sparse work after calibration and decode-only pruning. def _forward_modelopt( @@ -517,6 +616,26 @@ def native_forward(): ) return native_result + if _calibration_active(self): + if getattr(attn_metadata, "use_cascade", False): + # Cascade splits shared prefixes across requests, so per-request + # KV lengths are unavailable; skip measurement for this launch. + return native_forward() + # vLLM >= 0.15 writes the current K/V to the paged cache before + # impl.forward, so the calibrate kernel reads a complete cache. + key_cache, value_cache = kv_cache.unbind(0) + return _forward_calibrate( + self, + query=query, + key_cache=key_cache, + value_cache=value_cache, + block_table=attn_metadata.block_table, + seq_lens=attn_metadata.seq_lens, + cu_seqlens_q=attn_metadata.query_start_loc, + num_actual_tokens=attn_metadata.num_actual_tokens, + output=output, + ) + resolved = _resolve_forward( self, layer, @@ -705,6 +824,36 @@ def prepare_modelopt(): _maybe_update_flashinfer_cache(layer, key, value, kv_cache, attn_metadata, impl) cache_prepared = True + if _calibration_active(impl): + if getattr(attn_metadata, "use_cascade", False): + # Cascade splits shared prefixes across requests, so per-request + # KV lengths are unavailable; skip measurement for this launch. + return dense_fallback() + missing = [name for name in _FLASHINFER_METADATA_FIELDS if not hasattr(attn_metadata, name)] + if missing: + raise NotImplementedError( + "FlashInfer metadata is missing the ModelOpt calibration " + f"fields: {', '.join(missing)}" + ) + if kv_cache.ndim != 5 or kv_cache.shape[1] != 2: + raise ValueError( + "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" + ) + # Order matters: releases that update the KV cache inside forward must + # write the current K/V before the calibrate kernel reads the cache. + prepare_modelopt() + return _forward_calibrate( + impl, + query=query, + key_cache=kv_cache[:, 0], + value_cache=kv_cache[:, 1], + block_table=attn_metadata._modelopt_block_table, + seq_lens=attn_metadata._modelopt_seq_lens, + cu_seqlens_q=attn_metadata._modelopt_query_start_loc, + num_actual_tokens=attn_metadata._modelopt_num_actual_tokens, + output=output, + ) + resolved = _resolve_forward( impl, layer, @@ -833,3 +982,37 @@ def _clone_sparse_impl(old_impl, new_cls=None): new_impl = object.__new__(new_cls) new_impl.__dict__.update(old_state) return new_impl + + +def iter_sparse_impls(model): + """Yield every ModelOpt sparse attention impl reachable from a vLLM model. + + Walks ``model.named_modules()`` and returns the swapped ``impl`` of each + attention layer (FlashAttention or FlashInfer adapter). Used by the + calibration installer and RPC methods to toggle calibration mode and + harvest stats without knowing vLLM's module layout. + """ + for _, module in model.named_modules(): + impl = getattr(module, "impl", None) + if impl is None: + continue + if isinstance(impl, ModelOptSparseAttentionImpl) or ( + _FLASHINFER_IMPL_CLS is not None and isinstance(impl, _FLASHINFER_IMPL_CLS) + ): + yield impl + + +def enable_calibration(impls, threshold_trials: list[float]) -> None: + """Put a set of sparse impls into calibration mode and clear prior records.""" + if not threshold_trials: + raise ValueError("threshold_trials must be a non-empty list for calibration.") + for impl in impls: + impl._calibrate = True + impl._calib_threshold_trials = list(threshold_trials) + impl._calib_records = [] + + +def disable_calibration(impls) -> None: + """Turn off calibration mode (collected records are left intact).""" + for impl in impls: + impl._calibrate = False From 32245bf9aa9055b04364f100089a8d352660e367 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 21:05:21 -0700 Subject: [PATCH 03/14] Add skip-softmax calibration installer, TP-count aggregation, and example - install_vllm_skip_softmax_calibration: validation-before-mutation install of calibration adapters on every attention layer -- requires eager execution, fp16/bf16 model and KV dtypes, no active attention Q/K/P/V fakequant, and a FlashAttention/FlashInfer backend; disables cascade. Measurement starts separately (enable_calibration RPC), so warmup launches are never recorded. - DynamicThresholdCalibrator.calibrate_from_stats: backend-agnostic fit stage extracted from calibrate(), preserving result fields (fit_logspace, log_a) and reporting per-sample sparsity; the HF path delegates to it unchanged. - plugins/sparse_attn_calibration (vLLM-free): raw tile-count merging across layers and tensor-parallel ranks (counts are additive; ratios are formed only after the global merge, one fit per phase -- never per rank, never by averaging fitted coefficients), plus a canonical sparse_attention_config builder matching export_sparse_attention_config so the serving loader round-trips it; existing N:M groups are preserved. - Thin example: SkipSoftmaxCalibWorker (load hook + enable/status/counts RPCs) and calibrate_sparse_attn.py driver (CLI, prompts, rank-count aggregation, fit, checkpoint-config writing). Signed-off-by: Kai Xu --- examples/vllm_serve/calibrate_sparse_attn.py | 246 ++++++++++++++++++ examples/vllm_serve/sparse_attn_worker.py | 63 ++++- .../calibration/calibrator.py | 52 +++- .../plugins/sparse_attn_calibration.py | 236 +++++++++++++++++ .../attention_sparsity/plugins/vllm.py | 22 ++ .../plugins/vllm_runtime.py | 72 +++++ .../test_sparse_attn_worker.py | 6 +- 7 files changed, 693 insertions(+), 4 deletions(-) create mode 100644 examples/vllm_serve/calibrate_sparse_attn.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py new file mode 100644 index 00000000000..6c36f8f65ec --- /dev/null +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Calibrate skip-softmax thresholds *through vLLM* and write the serving config. + +Runs calibration prompts through a vLLM ``LLM`` whose attention layers carry +the ModelOpt calibration adapters (installed by +``sparse_attn_worker.SkipSoftmaxCalibWorker`` via +``install_vllm_skip_softmax_calibration``). The paged Triton calibration +kernel measures, per candidate threshold, how many KV tiles would be skipped — +over the paged KV cache, for both prefill and decode — then this driver +aggregates the raw counts from every tensor-parallel rank and fits the +exponential model ``scale_factor = a * exp(b * sparsity)`` once per phase. + +The fitted ``(a, b)`` are written as a canonical ``sparse_attention_config`` +block (the same schema ModelOpt's HF export produces), so the serving path +(``vllm_serve_sparse_attn.py`` / ``install_vllm_sparse_attention_from_checkpoint``) +loads it without changes. Any exported N:M sparse-softmax groups already in +the checkpoint config are preserved. + +Usage: + python calibrate_sparse_attn.py \ + --prompts_file prompts.txt \ + --target_sparse_ratio 0.5 \ + --decode_tokens 32 \ + --update_checkpoint_config + +``--prompts_file`` is one prompt per line; longer, varied-length prompts give a +better fit. With no file, a tiny built-in demo set is used (fine for a smoke +test, not for a real fit). +""" + +import argparse +import json +import os +import sys +from pathlib import Path + +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_calibration import ( + DEFAULT_THRESHOLD_TRIALS, + build_sparse_attention_config, + fit_from_counts, + merge_phase_counts, +) + +_DEMO_PROMPTS = [ + "Summarize the history of computing in a few paragraphs. " * 40, + "Explain how attention works in transformer models. " * 60, + "Write a detailed essay about renewable energy sources. " * 80, +] + + +def _load_prompts(prompts_file: str | None) -> list[str]: + if prompts_file is None: + print( + "[ModelOpt] No --prompts_file given; using a tiny built-in demo set. " + "Pass real, varied-length prompts for a usable fit." + ) + return _DEMO_PROMPTS + lines = [ln.strip() for ln in Path(prompts_file).read_text().splitlines() if ln.strip()] + if not lines: + raise ValueError(f"No prompts found in {prompts_file}") + print(f"[ModelOpt] Loaded {len(lines)} calibration prompts from {prompts_file}") + return lines + + +def _existing_sparse_config(ckpt: str) -> dict | None: + """Read the checkpoint's sparse_attention_config so non-skip groups survive.""" + config_json = Path(ckpt) / "config.json" + if not config_json.is_file(): + return None + existing = json.loads(config_json.read_text()).get("sparse_attention_config") + return existing if isinstance(existing, dict) else None + + +def _write_config(ckpt: str, sparse_config: dict, update_checkpoint: bool) -> None: + """Dump the sparse_attention_config and optionally merge into config.json.""" + out_path = Path("sparse_attention_config.json") + out_path.write_text(json.dumps(sparse_config, indent=2)) + print(f"[ModelOpt] Wrote calibrated config to {out_path.resolve()}") + + if not update_checkpoint: + print( + "[ModelOpt] Re-run with --update_checkpoint_config to merge this into " + f"{ckpt}/config.json (required for vllm_serve_sparse_attn.py to pick it up)." + ) + return + + config_json = Path(ckpt) / "config.json" + config = json.loads(config_json.read_text()) + config["sparse_attention_config"] = sparse_config + config_json.write_text(json.dumps(config, indent=2)) + print(f"[ModelOpt] Merged sparse_attention_config into {config_json}") + + +def main(): + parser = argparse.ArgumentParser(description="Calibrate skip-softmax thresholds via vLLM") + parser.add_argument("model", type=str, help="Path to the HF checkpoint to calibrate") + parser.add_argument("--prompts_file", type=str, default=None, help="One prompt per line") + parser.add_argument( + "--target_sparse_ratio", + type=float, + default=0.5, + help="Target sparsity baked into the exported config (applied to both phases)", + ) + parser.add_argument( + "--decode_tokens", + type=int, + default=32, + help="Decode tokens to generate per prompt (drives decode-phase calibration)", + ) + parser.add_argument( + "--max_model_len", type=int, default=None, help="vLLM max_model_len override" + ) + parser.add_argument( + "--tensor_parallel_size", type=int, default=1, help="vLLM tensor-parallel size" + ) + parser.add_argument( + "--gpu_memory_utilization", + type=float, + default=None, + help="vLLM GPU memory utilization fraction", + ) + parser.add_argument( + "--trust_remote_code", + action="store_true", + help="Trust remote code for custom model classes (e.g. NemotronH)", + ) + parser.add_argument("--dtype", type=str, default=None, help="Model dtype, e.g. bfloat16") + parser.add_argument( + "--attention_backend", + type=str, + default=None, + help="Force the vLLM attention backend, e.g. FLASH_ATTN or FLASHINFER. " + "Default: let vLLM choose (the installer supports whichever of FlashAttention " + "/ FlashInfer is selected).", + ) + parser.add_argument( + "--engine_kwargs", + type=str, + default=None, + help="JSON dict of extra vLLM engine kwargs, e.g. " + '\'{"enable_expert_parallel": true, "mamba_cache_mode": "align"}\' ' + "for hybrid MoE/Mamba models", + ) + parser.add_argument( + "--fit_logspace", + action="store_true", + help="Fit the exponential model in log space (wide scale_factor ranges)", + ) + parser.add_argument( + "--update_checkpoint_config", + action="store_true", + help="Merge the calibrated config into /config.json in place", + ) + args = parser.parse_args() + + # Workers run in separate processes and must import the calibration worker. + repo_root = str(Path(__file__).resolve().parent) + if repo_root not in sys.path: + sys.path.insert(0, repo_root) + current = os.environ.get("PYTHONPATH") + os.environ["PYTHONPATH"] = os.pathsep.join([current, repo_root]) if current else repo_root + + from vllm import LLM, SamplingParams + + prompts = _load_prompts(args.prompts_file) + + llm_kwargs = { + "model": args.model, + "worker_cls": "sparse_attn_worker.SkipSoftmaxCalibWorker", + # The calibration installer requires eager execution: the per-request + # calibration loop cannot be CUDA-graph captured. + "enforce_eager": True, + # Shared-prefix reuse would make prefill measurements cover only the + # non-cached suffix of each prompt; the installer rejects it. + "enable_prefix_caching": False, + } + if args.max_model_len is not None: + llm_kwargs["max_model_len"] = args.max_model_len + if args.tensor_parallel_size and args.tensor_parallel_size > 1: + llm_kwargs["tensor_parallel_size"] = args.tensor_parallel_size + if args.gpu_memory_utilization is not None: + llm_kwargs["gpu_memory_utilization"] = args.gpu_memory_utilization + if args.trust_remote_code: + llm_kwargs["trust_remote_code"] = True + if args.dtype is not None: + llm_kwargs["dtype"] = args.dtype + if args.attention_backend is not None: + llm_kwargs["attention_backend"] = args.attention_backend + if args.engine_kwargs: + extra = json.loads(args.engine_kwargs) + if not isinstance(extra, dict): + raise ValueError("--engine_kwargs must be a JSON object") + llm_kwargs.update(extra) + llm = LLM(**llm_kwargs) + + trials = list(DEFAULT_THRESHOLD_TRIALS) + n_layers = llm.collective_rpc("sparse_calib_enable", args=(trials,))[0] + status = llm.collective_rpc("sparse_calib_status")[0] + print(f"[ModelOpt] Calibration enabled on {n_layers} attention layers") + print(f"[ModelOpt] Active sparse impls: {status['impl_types']}") + + # generate() drives prefill (prefill-phase stats) then decode_tokens decode + # steps (decode-phase stats). The calibration kernel computes full attention, + # so the generated text is unaffected — only tile-skip counts are recorded. + sampling = SamplingParams(temperature=0.0, max_tokens=args.decode_tokens) + llm.generate(prompts, sampling) + + # Aggregate RAW counts from every TP rank (each rank only measures its + # attention-head shard), then fit once per phase on the global counts. + rank_counts = llm.collective_rpc("sparse_calib_counts") + merged = merge_phase_counts(rank_counts) + calibration_params = fit_from_counts(merged, trials, fit_logspace=args.fit_logspace) + + if not calibration_params: + print( + "[ModelOpt] Calibration produced no valid fit — try more/longer prompts " + "so observed sparsity spans the (10%, 90%) fitting window." + ) + return + + sparse_config = build_sparse_attention_config( + calibration_params, + {"prefill": args.target_sparse_ratio, "decode": args.target_sparse_ratio}, + existing_config=_existing_sparse_config(args.model), + ) + print("[ModelOpt] Calibrated threshold_scale_factor:") + print(json.dumps(sparse_config["config_groups"]["group_0"]["threshold_scale_factor"], indent=2)) + _write_config(args.model, sparse_config, args.update_checkpoint_config) + + +if __name__ == "__main__": + main() diff --git a/examples/vllm_serve/sparse_attn_worker.py b/examples/vllm_serve/sparse_attn_worker.py index 7b5fb28bf68..27e1eebdf97 100644 --- a/examples/vllm_serve/sparse_attn_worker.py +++ b/examples/vllm_serve/sparse_attn_worker.py @@ -17,12 +17,22 @@ from vllm.v1.worker.gpu_worker import Worker as BaseWorker +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_calibration import ( + DEFAULT_THRESHOLD_TRIALS, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( + collect_calibration_counts, + disable_calibration, + enable_calibration, + iter_sparse_impls, +) from modelopt.torch.sparsity.attention_sparsity.plugins.vllm_runtime import ( install_vllm_nvfp4_attention, + install_vllm_skip_softmax_calibration, install_vllm_sparse_attention_from_checkpoint, ) -__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker"] # noqa: RUF022 +__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker", "SkipSoftmaxCalibWorker"] # noqa: RUF022 _QUANT_FORMAT_KEYS = ("q_format", "k_format", "p_format", "v_format") @@ -69,6 +79,57 @@ def load_model(self, *args, **kwargs) -> None: _print_install_report("Sparse attention", report) +class SkipSoftmaxCalibWorker(BaseWorker): + """Calibrate skip-softmax thresholds through the engine. + + Unlike :class:`SparseAttnWorker` (which serves an already-calibrated + ``sparse_attention_config``), this worker *produces* that config. The + library installer swaps calibration-capable adapters onto every attention + layer at load; measurement starts only when the driver calls + ``sparse_calib_enable`` (so warmup launches are never recorded) and raw + per-threshold tile counts are harvested with ``sparse_calib_counts`` for + the driver to aggregate across TP ranks and fit. + """ + + def load_model(self, *args, **kwargs) -> None: + """Load the model, then install calibration adapters on every layer.""" + super().load_model(*args, **kwargs) + report = install_vllm_skip_softmax_calibration(self.model_runner) + print( + f"[ModelOpt] Skip-softmax calibration installed on {report.installed_count} " + f"attention layers: {dict(report.backend_counts)}" + ) + + # -- RPC methods (invoked via LLM.collective_rpc) ---------------------- + + def sparse_calib_enable(self, threshold_trials: list[float] | None = None) -> int: + """Enter calibration mode on all installed impls; returns layer count.""" + impls = list(iter_sparse_impls(_unwrapped_model(self))) + enable_calibration(impls, list(threshold_trials or DEFAULT_THRESHOLD_TRIALS)) + return len(impls) + + def sparse_calib_status(self) -> dict: + """Report active impls and record counts, so the backend is verifiable.""" + impls = list(iter_sparse_impls(_unwrapped_model(self))) + impl_types: dict[str, int] = {} + total_records = 0 + for impl in impls: + impl_types[type(impl).__name__] = impl_types.get(type(impl).__name__, 0) + 1 + total_records += len(getattr(impl, "_calib_records", [])) + return { + "num_sparse_layers": len(impls), + "impl_types": impl_types, + "calibrating": any(getattr(impl, "_calibrate", False) for impl in impls), + "total_records": total_records, + } + + def sparse_calib_counts(self) -> dict[str, list[dict]]: + """Stop measuring and return this rank's layer-merged raw tile counts.""" + model = _unwrapped_model(self) + disable_calibration(list(iter_sparse_impls(model))) + return collect_calibration_counts(model) + + class QuantSparseAttnWorker(BaseWorker): """Install quantized attention plus optional checkpoint sparsity. diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py index aded26fefdc..a37b447a686 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py @@ -130,8 +130,6 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic # with one entry per threshold, eliminating the need for repeated forward passes. print(f"\nStage 1: Collecting {phase} sparsity data for all thresholds in one pass...") - all_data_points = [] # List of {"threshold", "length", "scale_factor", "sparsity"} - self._set_thresholds(attention_modules, self.threshold_trials) self._enable_calibration_mode(attention_modules) with torch.no_grad(): @@ -139,6 +137,29 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic per_sample_stats = self._extract_calibration_stats(attention_modules, phase=phase) self._disable_calibration_mode(attention_modules) + return self.calibrate_from_stats(per_sample_stats, phase) + + def calibrate_from_stats(self, per_sample_stats: list[dict], phase: str) -> dict[str, Any]: + """Fit the exponential model from already-collected per-sample stats. + + This is the backend-agnostic Stage 2/3 of :meth:`calibrate`. The HF and + diffusion paths reach it through :meth:`calibrate` (which runs a + ``forward_loop`` to collect the stats first); the vLLM path collects the + stats itself — one record per scheduled request — and calls this directly + so both paths share the same exponential fit. + + Args: + per_sample_stats: List of ``{"sparsity": [s_0, ..., s_n], "sample_length": L}`` + records, one per calibration sample. ``sparsity`` holds the + skipped-tile fraction at each threshold in ``threshold_trials`` + (same order, same length). + phase: Phase being calibrated ('prefill' or 'decode'). + + Returns: + Dict with calibration results including a, b, r_squared, and num_data_points. + """ + all_data_points = [] # List of {"threshold", "length", "scale_factor", "sparsity"} + for sample_stat in per_sample_stats: length = sample_stat["sample_length"] sparsity_list = sample_stat["sparsity"] @@ -153,6 +174,12 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic } ) + # Per-sample measured sparsity (one row per calibration sample: its + # skipped-tile fraction at every threshold). Printed before the fit- + # validity guard so the raw per-sample data is visible even when the fit + # bails (e.g. degenerate near-zero sparsity). + self._print_per_sample_sparsity(per_sample_stats, phase) + if len(all_data_points) < 10: warnings.warn( f"Not enough data points for {phase} calibration. " @@ -286,11 +313,32 @@ def exponential(sparsity, a, b): "fit_logspace": self.fit_logspace, "min_observed_sparsity": min_observed_sparsity, "max_observed_sparsity": max_observed_sparsity, + # Raw per-sample measured sparsity, so callers can audit the spread + # across samples (not just the fitted average). + "per_sample_sparsity": [ + { + "sample_length": s.get("sample_length", 0), + "sparsity": list(s.get("sparsity", [])), + } + for s in per_sample_stats + ], } if self.fit_logspace: result["log_a"] = float(log_a) return result + def _print_per_sample_sparsity(self, per_sample_stats: list[dict], phase: str) -> None: + """Print each sample's measured skipped-tile fraction at every threshold.""" + if not per_sample_stats: + return + print(f"\nPer-sample {phase} sparsity (skipped-tile fraction per threshold):") + header = " ".join(f"{t:>7.0e}" for t in self.threshold_trials) + print(f" {'sample':>6} {'length':>8} {header}") + for idx, stat in enumerate(per_sample_stats): + sparsity = stat.get("sparsity", []) + row = " ".join(f"{s:>7.2%}" for s in sparsity) + print(f" {idx:>6} {stat.get('sample_length', 0):>8} {row}") + def _enable_calibration_mode(self, modules: list[nn.Module]): """Enable calibration mode on sparse attention modules.""" for idx, module in enumerate(modules): diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py new file mode 100644 index 00000000000..f60bef1b867 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""vLLM-free helpers for skip-softmax calibration through a serving engine. + +The serving adapters (``plugins/vllm.py``) record **raw per-threshold tile +counts** per scheduled request. These helpers merge those counts — across the +layers of one rank and across tensor-parallel ranks — then fit the exponential +threshold model and build the canonical ``sparse_attention_config`` block. + +Counts are additive, so aggregation is a plain sum: every layer of a rank and +every TP rank observes the same launches in the same order (TP ranks each see +their head shard), which makes records align by index within a phase. Fitting +happens once, per phase, on the globally merged counts — never per rank +(head-sharded counts are incomplete) and never by averaging independently +fitted coefficients (the fit is nonlinear). + +Everything here operates on plain Python data and is unit-testable without +vLLM installed. +""" + +from typing import Any + +import modelopt + +__all__ = [ + "DEFAULT_THRESHOLD_TRIALS", + "build_sparse_attention_config", + "fit_from_counts", + "merge_count_records", + "merge_phase_counts", + "split_records_by_phase", + "stats_from_counts", +] + +# Default threshold sweep — should span sparsities from ~10% to ~95%. +DEFAULT_THRESHOLD_TRIALS = [ + 1e-4, + 1e-3, + 5e-3, + 1e-2, + 3e-2, + 5e-2, + 1e-1, + 2e-1, + 3e-1, + 5e-1, + 7e-1, + 9e-1, +] + +_PHASES = ("prefill", "decode") + + +def split_records_by_phase(records: list[dict]) -> dict[str, list[dict]]: + """Group one impl's ordered calibration records by phase, preserving order.""" + per_phase: dict[str, list[dict]] = {phase: [] for phase in _PHASES} + for record in records: + per_phase.setdefault(record["phase"], []).append(record) + return per_phase + + +def merge_count_records(sources: list[list[dict]]) -> list[dict]: + """Element-wise sum aligned raw-count records from multiple sources. + + ``sources`` is a list over sources — the layers of one rank, or the + already layer-merged records of each TP rank — where each source is an + ordered list of ``{"sample_length", "total_tiles", "skipped_tiles"}`` + records for one phase. All sources observe the same launches in the same + order, so records align by index; tile counts are additive across both + layers and head-sharded TP ranks. + """ + sources = [source for source in sources if source] + if not sources: + return [] + num_samples = min(len(source) for source in sources) + merged = [] + for i in range(num_samples): + base = sources[0][i] + total = [0] * len(base["total_tiles"]) + skipped = [0] * len(base["skipped_tiles"]) + for source in sources: + record = source[i] + if record["sample_length"] != base["sample_length"]: + raise ValueError( + "Misaligned calibration records: sample lengths differ across " + f"sources at index {i} ({record['sample_length']} vs " + f"{base['sample_length']})" + ) + total = [a + b for a, b in zip(total, record["total_tiles"])] + skipped = [a + b for a, b in zip(skipped, record["skipped_tiles"])] + merged.append( + { + "sample_length": base["sample_length"], + "total_tiles": total, + "skipped_tiles": skipped, + } + ) + return merged + + +def merge_phase_counts(rank_counts: list[dict[str, list[dict]]]) -> dict[str, list[dict]]: + """Merge per-phase raw-count records collected from every TP rank. + + ``rank_counts`` is the list of per-rank results (one + ``{"prefill": [...], "decode": [...]}`` dict per rank, as returned by + ``collect_calibration_counts``). Use ALL ranks: with tensor parallelism + each rank only measures its attention-head shard, so any single rank's + counts are incomplete. + """ + phases = {phase for rank in rank_counts for phase in rank} + return { + phase: merge_count_records([rank.get(phase, []) for rank in rank_counts]) + for phase in phases + } + + +def stats_from_counts(count_records: list[dict]) -> list[dict]: + """Convert merged raw-count records into per-sample sparsity-ratio stats. + + Returns ``{"sample_length", "sparsity"}`` records in the shape + :meth:`DynamicThresholdCalibrator.calibrate_from_stats` consumes. + """ + stats = [] + for record in count_records: + sparsity = [ + (skipped / total if total else 0.0) + for skipped, total in zip(record["skipped_tiles"], record["total_tiles"]) + ] + stats.append({"sample_length": record["sample_length"], "sparsity": sparsity}) + return stats + + +def fit_from_counts( + per_phase_counts: dict[str, list[dict]], + threshold_trials: list[float], + *, + fit_logspace: bool = False, +) -> dict[str, dict[str, float]]: + """Fit the exponential skip-softmax model from globally merged counts. + + Reuses :class:`DynamicThresholdCalibrator` so vLLM-calibrated ``(a, b)`` + are identical in form to the HF path and export unchanged via + ``threshold_scale_factor``. One fit per phase, on counts already merged + across all TP ranks and layers. + + Returns: + ``{phase: {"a", "b", "min_observed_sparsity", "max_observed_sparsity"}}`` + for each phase that produced a valid fit. + """ + from ..calibration.calibrator import DynamicThresholdCalibrator + + calibration_params: dict[str, dict[str, float]] = {} + for phase, records in per_phase_counts.items(): + if not records: + continue + calibrator = DynamicThresholdCalibrator( + threshold_trials=list(threshold_trials), fit_logspace=fit_logspace + ) + result = calibrator.calibrate_from_stats(stats_from_counts(records), phase=phase) + if "a" in result and "b" in result: + params = {"a": result["a"], "b": result["b"]} + for key in ("min_observed_sparsity", "max_observed_sparsity"): + if key in result: + params[key] = result[key] + calibration_params[phase] = params + return calibration_params + + +def _normalize_target_sparsity(target_sparsity: dict[str, float] | float) -> dict[str, float]: + if isinstance(target_sparsity, (int, float)): + return {phase: float(target_sparsity) for phase in _PHASES} + return {phase: float(target_sparsity.get(phase, 0.5)) for phase in _PHASES} + + +def build_sparse_attention_config( + calibration_params: dict[str, dict[str, float]], + target_sparsity: dict[str, float] | float = 0.5, + *, + existing_config: dict | None = None, +) -> dict[str, Any]: + """Build the canonical ``sparse_attention_config`` block for a checkpoint. + + Emits the same schema as + ``modelopt.torch.sparsity.attention_sparsity.conversion.export_sparse_attention_config`` + — a ``config_groups`` entry with ``algorithm: skip_softmax`` holding the + group-local ``threshold_scale_factor`` and ``target_sparsity`` — so + ``load_from_checkpoint_metadata`` (the serving loader) round-trips it + without changes. + + Non-skip groups from ``existing_config`` (e.g. exported N:M + ``sparse_softmax`` metadata) are preserved after the skip group; an + existing ``skip_softmax`` group is replaced by the new calibration. + """ + threshold_scale_factor: dict[str, Any] = {"formula": "a * exp(b * target_sparsity)"} + for phase in _PHASES: + if phase in calibration_params: + threshold_scale_factor[phase] = { + "a": float(calibration_params[phase]["a"]), + "b": float(calibration_params[phase]["b"]), + } + + skip_group: dict[str, Any] = { + "algorithm": "skip_softmax", + "targets": ["Attention"], + "threshold_scale_factor": threshold_scale_factor, + "target_sparsity": _normalize_target_sparsity(target_sparsity), + } + + config_groups: dict[str, Any] = {"group_0": skip_group} + existing_groups = (existing_config or {}).get("config_groups") + if isinstance(existing_groups, dict): + preserved = [ + group + for group in existing_groups.values() + if isinstance(group, dict) and group.get("algorithm") != "skip_softmax" + ] + for idx, group in enumerate(preserved, start=1): + config_groups[f"group_{idx}"] = group + + return { + "config_groups": config_groups, + "producer": {"name": "modelopt", "version": modelopt.__version__}, + } diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index a3a5ebe1271..80e30e5fc18 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -1016,3 +1016,25 @@ def disable_calibration(impls) -> None: """Turn off calibration mode (collected records are left intact).""" for impl in impls: impl._calibrate = False + + +def collect_calibration_counts(model) -> dict[str, list[dict]]: + """Harvest one rank's raw per-phase tile counts from every calibrating impl. + + Sums counts across the rank's layers per aligned sample (every layer sees + the same launches in the same order), keeping raw + ``{"sample_length", "total_tiles", "skipped_tiles"}`` records per phase. + The driver merges these across TP ranks with + :func:`~.sparse_attn_calibration.merge_phase_counts` and fits once per + phase with :func:`~.sparse_attn_calibration.fit_from_counts` — sparsity + ratios are only formed after the global merge. + """ + from .sparse_attn_calibration import merge_count_records, split_records_by_phase + + per_phase_layers: dict[str, list[list[dict]]] = {} + for impl in iter_sparse_impls(model): + split = split_records_by_phase(getattr(impl, "_calib_records", [])) + for phase, records in split.items(): + if records: + per_phase_layers.setdefault(phase, []).append(records) + return {phase: merge_count_records(layers) for phase, layers in per_phase_layers.items()} diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index b5a578585d8..04d703f46a1 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -32,6 +32,7 @@ __all__ = [ "VllmAttentionInstallReport", "install_vllm_nvfp4_attention", + "install_vllm_skip_softmax_calibration", "install_vllm_sparse_attention_from_checkpoint", ] @@ -495,6 +496,77 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor return _build_report(plan) +def _attention_quant_error(module) -> str | None: + """Reject calibration on layers with any active attention Q/K/P/V fakequant.""" + for attr in ("q_bmm_quantizer", "k_bmm_quantizer", "p_bmm_quantizer", "v_bmm_quantizer"): + if getattr(getattr(module, attr, None), "is_enabled", False): + return f"{attr} is enabled; skip-softmax calibration requires unquantized attention" + if getattr(module, "_query_quant_in_kernel", False) or getattr( + module, "_value_quant_in_kernel", False + ): + return ( + "in-kernel attention quantization flags are set; skip-softmax " + "calibration requires unquantized attention" + ) + return None + + +def install_vllm_skip_softmax_calibration(model_runner) -> VllmAttentionInstallReport: + """Install skip-softmax calibration adapters into a loaded vLLM model. + + Swaps the backend-matched ModelOpt adapter onto every attention layer and + disables cascade attention, following validation-before-mutation: every + known compatibility error — across all layers — is collected and raised + before any module is changed. Calibration itself starts separately via + :func:`~.vllm.enable_calibration` (typically over a worker RPC), so engine + warmup/profiling launches after install are served natively and never + pollute the measurement; until then the adapters delegate every forward to + the backend's native implementation. + + Requirements validated here: eager execution (``enforce_eager=True`` — + the per-request calibration loop cannot be CUDA-graph captured), fp16/bf16 + model and KV-cache dtypes, no active attention Q/K/P/V fakequant, and a + FlashAttention or FlashInfer backend per layer. + """ + from vllm.config.compilation import CUDAGraphMode + + model = _unwrapped_model(model_runner) + candidates = [ + (name, module) + for name, module in model.named_modules() + if isinstance(module, _VLLM_ATTENTION) + ] + + _require_supported_vllm() + errors = _global_errors(model_runner) + if _cudagraph_mode(model_runner) != CUDAGraphMode.NONE: + errors.append( + "skip-softmax calibration requires eager execution (enforce_eager=True); " + "the per-request calibration loop cannot be CUDA-graph captured" + ) + if not candidates: + errors.append("no attention layers were found") + + plans = [] + for name, module in candidates: + reasons = _layer_errors(module) + if quant_error := _attention_quant_error(module): + reasons.append(quant_error) + new_impl, requires_flashinfer_patch, backend_error = _select_new_impl(module) + if backend_error: + reasons.append(backend_error) + if reasons: + errors.extend(f"{name or ''}: {reason}" for reason in reasons) + continue + plans.append( + _AttentionPlan(name, module, new_impl, {}, None, None, requires_flashinfer_patch) + ) + _raise_unsupported(errors, "skip-softmax calibration") + + plan = _InstallPlan(model_runner, tuple(plans), False, "SKIP_SOFTMAX_CALIBRATION") + return _apply_vllm_attention_plans(plan) + + def install_vllm_sparse_attention_from_checkpoint( model_runner, ) -> VllmAttentionInstallReport: diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 578922db077..f87a8cbb26b 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -80,7 +80,11 @@ def guarded_import(name, *args, **kwargs): monkeypatch.setattr(builtins, "__import__", guarded_import) worker_module = _load_worker_module("sparse_attn_worker_import_test") - assert worker_module.__all__ == ["SparseAttnWorker", "QuantSparseAttnWorker"] + assert worker_module.__all__ == [ + "SparseAttnWorker", + "QuantSparseAttnWorker", + "SkipSoftmaxCalibWorker", + ] @pytest.mark.parametrize( From 966a214a5598088ea398541ecd180f9537520e77 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 17 Jul 2026 21:17:29 -0700 Subject: [PATCH 04/14] Add vLLM skip-softmax calibration tests and docs Tests: - Paged calibrate kernel: paged == contiguous (counters bit-equal, output match) across aligned and non-128-aligned lengths with shuffled block tables; decode-shaped measurement covers the full cache; high-block-ID pointer regression (int64 paged offsets, memory-gated). - Calibration/serving tile contract: an active skip-softmax launch skips exactly the tiles the calibration kernel counted at the same threshold (same 128x128 geometry, same prefix-max criterion); skip composes with P/V QDQ (sm89+). - Adapter forward: mixed prefill/decode batches record per-request phases and raw counts while output stays dense vs an SDPA reference; NHD and fp16/bf16 cache validation; layer count summing; FlashInfer cache write ordered before the calibrate-kernel read. - Installer: install-without-measure lifecycle, eager/quantizer/fp8-KV/ no-layer rejections with validation-before-mutation, and the sparse-only CUDA-graph guard for calibrated decode configs. - vLLM-free helpers: count merging/alignment, TP-rank aggregation, synthetic exponential fit recovery, calibrate_from_stats field contract, canonical config round trip through the serving loader with N:M group preservation. Docs: vLLM calibration workflow section in examples/vllm_serve/README.md. Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 14 + .../attention/test_paged_calibrate.py | 236 ++++++++++ .../test_vllm_calibration.py | 411 ++++++++++++++++++ .../test_sparse_attn_calibration.py | 203 +++++++++ 4 files changed, 864 insertions(+) create mode 100644 tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py create mode 100644 tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index fc4e8a0ebcc..424163d5371 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -178,6 +178,20 @@ Workflow: If the checkpoint has no `sparse_attention_config`, the sparse-only installer passes through and vLLM runs unchanged. Whole-model fakequant flows remain handled by `vllm_serve_fakequant.py`; the compact attention-only path is below. +### Calibrate skip-softmax thresholds through vLLM + +Instead of the HF path in step 1, thresholds can be calibrated directly through vLLM — over the paged KV cache, for both prefill and decode, with tensor parallelism: + +```bash +python calibrate_sparse_attn.py \ + --prompts_file prompts.txt --target_sparse_ratio 0.5 \ + --decode_tokens 32 --tensor_parallel_size 8 --update_checkpoint_config +``` + +`install_vllm_skip_softmax_calibration` (called by `sparse_attn_worker.SkipSoftmaxCalibWorker` at model load) swaps calibration adapters onto every attention layer after validating all of them — eager execution is required, model and KV-cache dtypes must be fp16/bf16, and no attention Q/K/P/V fakequant may be active. During `llm.generate`, the paged Triton calibration kernel computes full dense attention (generation is numerically unchanged) while counting, per candidate threshold, how many KV tiles the skip criterion would drop. The driver then collects **raw tile counts from every TP rank** (each rank only measures its head shard), merges them, fits `scale_factor = a * exp(b * sparsity)` once per phase, and writes the same canonical `sparse_attention_config` block the HF export produces — preserving any exported N:M sparse-softmax groups — so the serving workflow above picks it up unchanged. + +Calibration and serving measure skipping at the same fixed 128x128 tile geometry (active skip-softmax launches bypass the autotuner), so the sparsity realized at serve time matches the calibrated `(a, b)` model. + The reusable serving policies live in `modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py`. `install_vllm_sparse_attention_from_checkpoint` installs checkpoint-driven sparse-only attention, while `install_vllm_nvfp4_attention` installs fixed NVFP4 Q/K/P/V with optional checkpoint sparsity. Both validate every selected layer before publishing any replacement implementation and return a `VllmAttentionInstallReport` with the installed layer names and backend counts. `sparse_attn_worker.py` only invokes these APIs after vLLM loads the model. It retains `SparseAttnWorker` as the launcher's default and provides `QuantSparseAttnWorker` for the compact NVFP4 policy. Other vLLM integrations can invoke the same library APIs directly: diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py new file mode 100644 index 00000000000..9e55879dbb6 --- /dev/null +++ b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Paged-cache calibration kernel tests and the calibration/serving tile contract.""" + +import pytest +import torch +from conftest import make_qkv, make_varlen_meta + +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE + +if TRITON_KERNEL_AVAILABLE: + from modelopt.torch.kernels.common.attention import attention + from modelopt.torch.kernels.sparsity.attention.calibrate import attention_calibrate + +pytestmark = pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") + +TRIALS = [1e-3, 1e-2, 1e-1, 3e-1] + + +def _pack_paged(k, v, page_size, *, shuffle=True, num_spare_blocks=7): + """Pack one sequence's contiguous K/V into a (optionally shuffled) paged cache.""" + seq = k.shape[0] + num_blocks = (seq + page_size - 1) // page_size + order = torch.randperm(num_blocks) if shuffle else torch.arange(num_blocks) + k_cache = k.new_zeros(num_blocks + num_spare_blocks, page_size, k.shape[1], k.shape[2]) + v_cache = torch.zeros_like(k_cache) + block_table = torch.zeros(1, num_blocks, device=k.device, dtype=torch.int32) + for i in range(num_blocks): + page = int(order[i]) + num_spare_blocks # keep low pages unused + ts, te = i * page_size, min((i + 1) * page_size, seq) + k_cache[page, : te - ts] = k[ts:te] + v_cache[page, : te - ts] = v[ts:te] + block_table[0, i] = page + return k_cache, v_cache, block_table + + +class TestPagedCalibrate: + @pytest.mark.parametrize("seq_len", [256, 300, 512]) # 300: non-128-aligned padding + def test_paged_matches_contiguous_prefill(self, seq_len): + """Paged and contiguous calibration agree exactly on counters and output.""" + torch.manual_seed(0) + num_heads, num_kv_heads, head_dim, page_size = 8, 2, 64, 16 + q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + locs, lens = make_varlen_meta([seq_len]) + + out_ref, counters_ref = attention_calibrate( + q, k, v, locs, lens, seq_len, is_causal=True, threshold_trials=TRIALS + ) + + k_cache, v_cache, block_table = _pack_paged(k, v, page_size) + k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) + out_paged, counters_paged = attention_calibrate( + q, + k_dummy, + k_dummy, + locs, + lens, + seq_len, + is_causal=True, + threshold_trials=TRIALS, + b_seq_len_k=lens, + max_input_len_k=seq_len, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + page_size=page_size, + ) + + assert torch.equal(counters_ref, counters_paged) + torch.testing.assert_close(out_paged, out_ref, rtol=1e-3, atol=1e-3) + + def test_paged_decode_measures_full_cache(self): + """A one-row decode query measures every KV tile of the paged cache.""" + torch.manual_seed(1) + num_heads, num_kv_heads, head_dim, page_size = 8, 2, 64, 16 + ctx = 384 + q, k, v = make_qkv(ctx, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + k_cache, v_cache, block_table = _pack_paged(k, v, page_size) + k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) + locs = torch.zeros(1, device="cuda", dtype=torch.int32) + + _, counters = attention_calibrate( + q[:1], + k_dummy, + k_dummy, + locs, + torch.ones(1, device="cuda", dtype=torch.int32), + 1, + is_causal=False, + threshold_trials=TRIALS, + b_seq_len_k=torch.tensor([ctx], device="cuda", dtype=torch.int32), + max_input_len_k=ctx, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + page_size=page_size, + ) + + num_kv_tiles = -(-ctx // 128) + assert counters[:, 0].tolist() == [num_heads * num_kv_tiles] * len(TRIALS) + + def test_high_block_id_pointer_arithmetic(self): + """Block IDs whose int32 byte offsets would wrap still read correctly.""" + num_kv_heads, head_dim, page_size = 2, 64, 16 + block_elems = page_size * num_kv_heads * head_dim + # Smallest block ID whose element offset exceeds int32. + high_block = (2**31) // block_elems + 1 + bytes_needed = 2 * (high_block + 1) * block_elems * 2 # K and V caches, bf16 + free, _ = torch.cuda.mem_get_info() + if free < bytes_needed + (2 << 30): + pytest.skip(f"needs ~{bytes_needed / 2**30:.1f} GiB free GPU memory") + + torch.manual_seed(2) + num_heads = 4 + q, k, v = make_qkv(page_size, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + k_cache = torch.zeros( + high_block + 1, page_size, num_kv_heads, head_dim, device="cuda", dtype=torch.bfloat16 + ) + v_cache = torch.zeros_like(k_cache) + k_cache[high_block] = k + v_cache[high_block] = v + block_table = torch.tensor([[high_block]], device="cuda", dtype=torch.int32) + locs, lens = make_varlen_meta([page_size]) + + out_ref, counters_ref = attention_calibrate( + q, k, v, locs, lens, page_size, is_causal=True, threshold_trials=TRIALS + ) + k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) + out_paged, counters_paged = attention_calibrate( + q, + k_dummy, + k_dummy, + locs, + lens, + page_size, + is_causal=True, + threshold_trials=TRIALS, + b_seq_len_k=lens, + max_input_len_k=page_size, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + page_size=page_size, + ) + del k_cache, v_cache + + assert torch.equal(counters_ref, counters_paged) + torch.testing.assert_close(out_paged, out_ref, rtol=1e-3, atol=1e-3) + + +class TestCalibrationServingTileContract: + """Active skip launches and calibration must count identically (same tiles).""" + + def _contrasty_qkv(self, seq_len, num_heads, num_kv_heads, head_dim): + """K with a dominant head-of-sequence so later tiles are skippable.""" + torch.manual_seed(3) + q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + k = k * 0.05 + k[:32] = k[:32] * 600.0 # first tile dominates the running max by >> log2(threshold) + return q, k, v + + @pytest.mark.parametrize("threshold", [1e-3, 1e-2]) + def test_serve_skip_counts_equal_calibrate_counts(self, threshold): + seq_len, num_heads, num_kv_heads, head_dim = 512, 8, 2, 64 + q, k, v = self._contrasty_qkv(seq_len, num_heads, num_kv_heads, head_dim) + locs, lens = make_varlen_meta([seq_len]) + scale = 1.0 / (head_dim**0.5) + + _, counters = attention_calibrate( + q, + k, + v, + locs, + lens, + seq_len, + is_causal=True, + softmax_scale=scale, + threshold_trials=[threshold], + ) + + out = attention( + q, + k, + v, + locs, + lens, + seq_len, + is_causal=True, + softmax_scale=scale, + skip_softmax_threshold=threshold, + measure_sparsity=True, + ) + + calib_total, calib_skipped = int(counters[0, 0]), int(counters[0, 1]) + assert calib_skipped > 0, "test data must produce skippable tiles" + # Same 128x128 tile geometry and same prefix-max criterion => the serve + # kernel must skip exactly the tiles calibration predicted. + assert out._sparsity_total == calib_total + assert out._sparsity_skipped == calib_skipped + + def test_skip_composes_with_pv_qdq(self): + """Active skip at the fixed tile compiles and runs with P/V QDQ enabled.""" + if torch.cuda.get_device_capability() < (8, 9): + pytest.skip("NVFP4/FP8 QDQ needs compute capability >= 8.9 (E4M3 casts)") + seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 + q, k, v = self._contrasty_qkv(seq_len, num_heads, num_kv_heads, head_dim) + locs, lens = make_varlen_meta([seq_len]) + + out = attention( + q, + k, + v, + locs, + lens, + seq_len, + is_causal=True, + skip_softmax_threshold=1e-2, + p_qdq="nvfp4", + p_qdq_amax=1.0, + v_qdq="nvfp4", + v_qdq_amax=float(v.abs().amax()), + ) + assert torch.isfinite(out).all() diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py new file mode 100644 index 00000000000..e5192e7e7f3 --- /dev/null +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -0,0 +1,411 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for skip-softmax calibration through the vLLM adapters and installer.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl + +from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as attention_plugin +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm_runtime +from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( + ModelOptSparseAttentionImpl, + collect_calibration_counts, + disable_calibration, + enable_calibration, + iter_sparse_impls, +) + +TRIALS = [1e-3, 1e-1, 5e-1] + + +# --------------------------------------------------------------------------- +# Installer fakes (mirroring test_vllm_runtime.py) +# --------------------------------------------------------------------------- +def _bare_attention(impl_cls=FlashAttentionImpl): + module = object.__new__(vllm_runtime._VLLM_ATTENTION) + nn.Module.__init__(module) + module.attn_type = "decoder" + module.head_size = 64 + module.device = torch.device("cpu") + module.dtype = torch.float16 + module.impl = object.__new__(impl_cls) + module.impl.sinks = None + return module + + +def _model_runner(model, *, sparse_metadata=None, cudagraph_mode=CUDAGraphMode.NONE): + hf_config = SimpleNamespace(sparse_attention_config=sparse_metadata) + model_config = SimpleNamespace(hf_config=hf_config, dtype=torch.float16) + return SimpleNamespace( + model=model, + model_config=model_config, + cascade_attn_enabled=True, + vllm_config=SimpleNamespace( + model_config=model_config, + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + enable_dbo=False, + use_ubatching=False, + ), + cache_config=SimpleNamespace(enable_prefix_caching=False, cache_dtype="auto"), + compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), + kv_transfer_config=None, + speculative_config=None, + ), + ) + + +class TestCalibrationInstaller: + def test_installs_adapters_without_enabling_measurement(self): + first = _bare_attention() + second = _bare_attention() + runner = _model_runner(nn.ModuleDict({"a_attn": first, "b_attn": second})) + + report = vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + assert report.installed_count == 2 + assert report.sparse_algorithm == "SKIP_SOFTMAX_CALIBRATION" + assert report.cascade_disabled and runner.cascade_attn_enabled is False + for module in (first, second): + assert isinstance(module.impl, ModelOptSparseAttentionImpl) + # Measurement starts only via enable_calibration, so warmup + # launches after install are never recorded. + assert not attention_plugin._calibration_active(module.impl) + + def test_rejects_non_eager_execution(self): + runner = _model_runner( + nn.ModuleDict({"attn": _bare_attention()}), + cudagraph_mode=CUDAGraphMode.PIECEWISE, + ) + with pytest.raises(NotImplementedError, match="enforce_eager"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + def test_rejects_active_attention_quantizers_atomically(self): + quantized = _bare_attention() + quantized.q_bmm_quantizer = SimpleNamespace(is_enabled=True) + clean = _bare_attention() + clean_impl = clean.impl + runner = _model_runner(nn.ModuleDict({"q_attn": quantized, "c_attn": clean})) + + with pytest.raises(NotImplementedError, match="requires unquantized attention"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + # Validation-before-mutation: the clean layer must not be touched either. + assert clean.impl is clean_impl + assert runner.cascade_attn_enabled is True + + def test_rejects_fp8_kv_cache(self): + attention = _bare_attention() + attention.kv_cache_dtype = "fp8" + runner = _model_runner(nn.ModuleDict({"attn": attention})) + with pytest.raises(NotImplementedError, match="FP8 KV cache"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + def test_rejects_model_without_attention_layers(self): + runner = _model_runner(nn.ModuleDict({})) + with pytest.raises(NotImplementedError, match="no attention layers"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + +class TestSparseOnlyGraphGuard: + """Commit-contract: the calibrated-decode graph guard is not quantize-gated.""" + + _CALIBRATED_META = { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "threshold_scale_factor": { + "prefill": {"a": 7.9, "b": 8.6}, + "decode": {"a": 0.12, "b": 9.8}, + }, + } + } + } + + def test_sparse_only_install_rejects_full_decode_graph(self): + attention = _bare_attention() + runner = _model_runner( + nn.ModuleDict({"attn": attention}), + sparse_metadata=self._CALIBRATED_META, + cudagraph_mode=CUDAGraphMode.FULL, + ) + with pytest.raises(NotImplementedError, match="non-FULL CUDA graph"): + vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert not isinstance(attention.impl, ModelOptSparseAttentionImpl) + + def test_sparse_only_install_allows_eager(self): + attention = _bare_attention() + runner = _model_runner( + nn.ModuleDict({"attn": attention}), + sparse_metadata=self._CALIBRATED_META, + cudagraph_mode=CUDAGraphMode.NONE, + ) + report = vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert report.installed_count == 1 + + +# --------------------------------------------------------------------------- +# Calibration forward through the FlashAttention adapter (GPU) +# --------------------------------------------------------------------------- +def _make_impl(num_heads, head_dim, num_kv_heads): + return ModelOptSparseAttentionImpl( + num_heads=num_heads, + head_size=head_dim, + scale=1.0 / (head_dim**0.5), + num_kv_heads=num_kv_heads, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + ) + + +def _paged_cache_for(seqs_kv, num_kv_heads, head_dim, page_size, device, dtype): + """Scatter per-request contiguous K/V lists into a stacked paged cache.""" + blocks_per_seq = [(kv.shape[0] + page_size - 1) // page_size for kv, _ in seqs_kv] + num_blocks = sum(blocks_per_seq) + max_blocks = max(blocks_per_seq) + k_cache = torch.zeros(num_blocks, page_size, num_kv_heads, head_dim, device=device, dtype=dtype) + v_cache = torch.zeros_like(k_cache) + block_table = torch.zeros(len(seqs_kv), max_blocks, device=device, dtype=torch.int32) + g = 0 + for b, (k, v) in enumerate(seqs_kv): + for blk in range(blocks_per_seq[b]): + block_table[b, blk] = g + ts, te = blk * page_size, min((blk + 1) * page_size, k.shape[0]) + k_cache[g, : te - ts] = k[ts:te] + v_cache[g, : te - ts] = v[ts:te] + g += 1 + return torch.stack([k_cache, v_cache], dim=0), block_table + + +def _sdpa_reference(q, k, v, is_causal): + # [tokens, heads, dim] -> [1, heads, tokens, dim] + qh, kh, vh = (t.transpose(0, 1).unsqueeze(0).float() for t in (q, k, v)) + kh = kh.repeat_interleave(q.shape[1] // k.shape[1], dim=1) + vh = vh.repeat_interleave(q.shape[1] // v.shape[1], dim=1) + if is_causal and q.shape[0] < k.shape[0]: + # Suffix-causal mask for decode/chunked prefill shapes. + mask = torch.ones(q.shape[0], k.shape[0], dtype=torch.bool, device=q.device).tril( + diagonal=k.shape[0] - q.shape[0] + ) + out = torch.nn.functional.scaled_dot_product_attention(qh, kh, vh, attn_mask=mask) + else: + out = torch.nn.functional.scaled_dot_product_attention(qh, kh, vh, is_causal=is_causal) + return out.squeeze(0).transpose(0, 1).to(q.dtype) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestCalibrationForward: + def test_mixed_batch_records_phases_and_stays_dense(self): + torch.manual_seed(0) + device, dtype = "cuda", torch.bfloat16 + num_heads, num_kv_heads, head_dim, page_size = 4, 2, 64, 16 + prefill_len, decode_ctx = 64, 48 + + k0 = torch.randn(prefill_len, num_kv_heads, head_dim, device=device, dtype=dtype) + v0 = torch.randn_like(k0) + k1 = torch.randn(decode_ctx, num_kv_heads, head_dim, device=device, dtype=dtype) + v1 = torch.randn_like(k1) + q = torch.randn(prefill_len + 1, num_heads, head_dim, device=device, dtype=dtype) + + kv_cache, block_table = _paged_cache_for( + [(k0, v0), (k1, v1)], num_kv_heads, head_dim, page_size, device, dtype + ) + attn_metadata = SimpleNamespace( + num_actual_tokens=prefill_len + 1, + max_query_len=prefill_len, + max_seq_len=max(prefill_len, decode_ctx), + query_start_loc=torch.tensor( + [0, prefill_len, prefill_len + 1], device=device, dtype=torch.int32 + ), + seq_lens=torch.tensor([prefill_len, decode_ctx], device=device, dtype=torch.int32), + block_table=block_table, + ) + + impl = _make_impl(num_heads, head_dim, num_kv_heads) + impl.sparse_kw = {} + enable_calibration([impl], TRIALS) + output = torch.empty_like(q) + out = impl.forward( + layer=None, + query=q, + key=q[:, :num_kv_heads], + value=q[:, :num_kv_heads], + kv_cache=kv_cache, + attn_metadata=attn_metadata, + output=output, + ) + + # Two records: one per request, phases decided per request. + records = impl._calib_records + assert [r["phase"] for r in records] == ["prefill", "decode"] + assert [r["sample_length"] for r in records] == [prefill_len, decode_ctx] + for record in records: + assert len(record["total_tiles"]) == len(TRIALS) + assert all(t > 0 for t in record["total_tiles"]) + assert all(0 <= s <= t for s, t in zip(record["skipped_tiles"], record["total_tiles"])) + + # Output is full dense attention (calibration never skips). + ref_prefill = _sdpa_reference(q[:prefill_len], k0, v0, is_causal=True) + ref_decode = _sdpa_reference(q[prefill_len:], k1, v1, is_causal=False) + torch.testing.assert_close(out[:prefill_len], ref_prefill, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(out[prefill_len:], ref_decode, rtol=2e-2, atol=2e-2) + + def test_collect_calibration_counts_sums_layers(self): + class FakeModel(nn.Module): + def __init__(self, impls): + super().__init__() + self._impls = impls + self.layers = nn.ModuleList([nn.Identity() for _ in impls]) + for identity, impl in zip(self.layers, impls): + identity.impl = impl + + impls = [object.__new__(ModelOptSparseAttentionImpl) for _ in range(2)] + enable_calibration(impls, TRIALS) + for idx, impl in enumerate(impls): + impl._calib_records = [ + { + "phase": "prefill", + "sample_length": 128, + "total_tiles": [4, 4, 4], + "skipped_tiles": [idx, idx + 1, idx + 2], + } + ] + model = FakeModel(impls) + assert len(list(iter_sparse_impls(model))) == 2 + disable_calibration(impls) + + counts = collect_calibration_counts(model) + assert counts["prefill"] == [ + {"sample_length": 128, "total_tiles": [8, 8, 8], "skipped_tiles": [1, 3, 5]} + ] + + def test_rejects_non_nhd_cache_layout(self): + num_heads, num_kv_heads, head_dim = 4, 2, 64 + impl = _make_impl(num_heads, head_dim, num_kv_heads) + enable_calibration([impl], TRIALS) + # HND-shaped cache: [blocks, kv_heads, page, dim] -> axis 2 != num_kv_heads. + kv_cache = torch.zeros(2, 1, num_kv_heads, 16, head_dim, dtype=torch.bfloat16) + attn_metadata = SimpleNamespace( + num_actual_tokens=1, + max_query_len=1, + max_seq_len=8, + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=torch.tensor([8], dtype=torch.int32), + block_table=torch.zeros(1, 1, dtype=torch.int32), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + with pytest.raises(NotImplementedError, match="not NHD"): + impl.forward( + layer=None, + query=q, + key=q[:, :num_kv_heads], + value=q[:, :num_kv_heads], + kv_cache=kv_cache, + attn_metadata=attn_metadata, + output=torch.empty_like(q), + ) + + def test_rejects_non_16bit_cache(self): + num_heads, num_kv_heads, head_dim = 4, 2, 64 + impl = _make_impl(num_heads, head_dim, num_kv_heads) + enable_calibration([impl], TRIALS) + kv_cache = torch.zeros(2, 1, 16, num_kv_heads, head_dim, dtype=torch.uint8) + attn_metadata = SimpleNamespace( + num_actual_tokens=1, + max_query_len=1, + max_seq_len=8, + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=torch.tensor([8], dtype=torch.int32), + block_table=torch.zeros(1, 1, dtype=torch.int32), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + with pytest.raises(NotImplementedError, match="fp16/bf16 KV cache"): + impl.forward( + layer=None, + query=q, + key=q[:, :num_kv_heads], + value=q[:, :num_kv_heads], + kv_cache=kv_cache, + attn_metadata=attn_metadata, + output=torch.empty_like(q), + ) + + +# --------------------------------------------------------------------------- +# FlashInfer adapter: cache write must precede the calibrate-kernel read +# --------------------------------------------------------------------------- +class TestFlashInferCalibrationOrdering: + def test_cache_write_happens_before_calibrate_read(self, monkeypatch): + calls = [] + monkeypatch.setattr( + attention_plugin, + "_maybe_update_flashinfer_cache", + lambda *args, **kwargs: calls.append("cache_write"), + ) + + def fake_calibrate(q, *args, **kwargs): + calls.append("calibrate") + counters = torch.zeros(len(TRIALS), 2, dtype=torch.int64) + return torch.zeros_like(q), counters + + monkeypatch.setattr(attention_plugin, "attention_calibrate", fake_calibrate) + + num_heads, num_kv_heads, head_dim, page = 4, 2, 64, 16 + impl = SimpleNamespace( + num_kv_heads=num_kv_heads, + head_size=head_dim, + scale=1.0 / (head_dim**0.5), + _calibrate=True, + _calib_threshold_trials=list(TRIALS), + _calib_records=[], + ) + kv_cache = torch.zeros(3, 2, page, num_kv_heads, head_dim, dtype=torch.bfloat16) + attn_metadata = SimpleNamespace( + _modelopt_block_table=torch.zeros(1, 1, dtype=torch.int32), + _modelopt_seq_lens=torch.tensor([8], dtype=torch.int32), + _modelopt_query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + _modelopt_num_actual_tokens=1, + _modelopt_max_query_len=1, + _modelopt_max_seq_len=8, + _modelopt_causal=False, + slot_mapping=torch.zeros(1, dtype=torch.int64), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + + out = attention_plugin._flashinfer_forward( + impl, + None, # native_forward is unused on the calibration path + None, # layer + q, + q[:, :num_kv_heads], + q[:, :num_kv_heads], + kv_cache, + attn_metadata, + output=torch.empty_like(q), + ) + + assert calls == ["cache_write", "calibrate"] + assert torch.isfinite(out).all() + assert len(impl._calib_records) == 1 + assert impl._calib_records[0]["phase"] == "decode" diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py new file mode 100644 index 00000000000..77f836c9205 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for vLLM-free skip-softmax calibration helpers (no vLLM needed).""" + +import math +from types import SimpleNamespace + +import pytest + +from modelopt.torch.sparsity.attention_sparsity.calibration.calibrator import ( + DynamicThresholdCalibrator, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_calibration import ( + DEFAULT_THRESHOLD_TRIALS, + build_sparse_attention_config, + fit_from_counts, + merge_count_records, + merge_phase_counts, + split_records_by_phase, + stats_from_counts, +) +from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_config import ( + load_from_checkpoint_metadata, +) + + +def _record(phase, length, totals, skipped): + return { + "phase": phase, + "sample_length": length, + "total_tiles": list(totals), + "skipped_tiles": list(skipped), + } + + +class TestCountMerging: + def test_split_records_by_phase_preserves_order(self): + records = [ + _record("prefill", 100, [4], [1]), + _record("decode", 101, [2], [0]), + _record("prefill", 200, [8], [3]), + ] + split = split_records_by_phase(records) + assert [r["sample_length"] for r in split["prefill"]] == [100, 200] + assert [r["sample_length"] for r in split["decode"]] == [101] + + def test_merge_sums_counts_elementwise(self): + layer_a = [_record("prefill", 128, [10, 10], [2, 4])] + layer_b = [_record("prefill", 128, [10, 10], [1, 3])] + merged = merge_count_records([layer_a, layer_b]) + assert merged == [{"sample_length": 128, "total_tiles": [20, 20], "skipped_tiles": [3, 7]}] + + def test_merge_truncates_to_shortest_source(self): + long = [_record("prefill", 1, [1], [0]), _record("prefill", 2, [1], [1])] + short = [_record("prefill", 1, [1], [1])] + merged = merge_count_records([long, short]) + assert len(merged) == 1 + assert merged[0]["skipped_tiles"] == [1] + + def test_merge_rejects_misaligned_sample_lengths(self): + with pytest.raises(ValueError, match="Misaligned calibration records"): + merge_count_records( + [[_record("prefill", 100, [1], [0])], [_record("prefill", 200, [1], [0])]] + ) + + def test_merge_phase_counts_across_ranks(self): + rank0 = {"prefill": [_record("prefill", 64, [5], [1])], "decode": []} + rank1 = {"prefill": [_record("prefill", 64, [5], [2])]} + merged = merge_phase_counts([rank0, rank1]) + assert merged["prefill"][0]["total_tiles"] == [10] + assert merged["prefill"][0]["skipped_tiles"] == [3] + assert merged["decode"] == [] + + def test_stats_from_counts_forms_ratios_after_merge(self): + stats = stats_from_counts( + [{"sample_length": 64, "total_tiles": [8, 0], "skipped_tiles": [2, 0]}] + ) + assert stats == [{"sample_length": 64, "sparsity": [0.25, 0.0]}] + + +class TestFitFromCounts: + def test_fit_recovers_synthetic_exponential(self): + a_true, b_true = 5.0, 8.0 + trials = DEFAULT_THRESHOLD_TRIALS + + def counts(length, total): + sparsity = [ + min(0.95, max(0.0, math.log(max(t * length, 1e-9) / a_true) / b_true)) + for t in trials + ] + return { + "sample_length": length, + "total_tiles": [total] * len(trials), + "skipped_tiles": [int(s * total) for s in sparsity], + } + + per_phase = {"prefill": [counts(length, 4000) for length in (2048, 4096, 8192, 16384)]} + params = fit_from_counts(per_phase, trials) + assert abs(params["prefill"]["a"] - a_true) / a_true < 0.3 + assert abs(params["prefill"]["b"] - b_true) / b_true < 0.15 + assert 0.0 <= params["prefill"]["min_observed_sparsity"] <= 1.0 + + def test_empty_phase_produces_no_fit(self): + assert fit_from_counts({"decode": []}, DEFAULT_THRESHOLD_TRIALS) == {} + + +class TestCalibrateFromStats: + def _stats(self, trials): + a_true, b_true = 3.0, 9.0 + stats = [] + for length in (1024, 2048, 4096, 8192): + sparsity = [ + min(0.95, max(0.0, math.log(max(t * length, 1e-9) / a_true) / b_true)) + for t in trials + ] + stats.append({"sample_length": length, "sparsity": sparsity}) + return stats + + def test_linear_fit_reports_fit_logspace_false(self): + calibrator = DynamicThresholdCalibrator(threshold_trials=list(DEFAULT_THRESHOLD_TRIALS)) + result = calibrator.calibrate_from_stats(self._stats(DEFAULT_THRESHOLD_TRIALS), "prefill") + assert result["fit_logspace"] is False + assert "log_a" not in result + assert len(result["per_sample_sparsity"]) == 4 + + def test_logspace_fit_preserves_log_a(self): + calibrator = DynamicThresholdCalibrator( + threshold_trials=list(DEFAULT_THRESHOLD_TRIALS), fit_logspace=True + ) + result = calibrator.calibrate_from_stats(self._stats(DEFAULT_THRESHOLD_TRIALS), "prefill") + assert result["fit_logspace"] is True + assert math.isclose(math.exp(result["log_a"]), result["a"], rel_tol=1e-9) + + +class TestBuildSparseAttentionConfig: + _PARAMS = {"prefill": {"a": 7.9, "b": 8.6}, "decode": {"a": 0.12, "b": 9.8}} + + def test_canonical_schema(self): + config = build_sparse_attention_config(self._PARAMS, 0.4) + group = config["config_groups"]["group_0"] + assert group["algorithm"] == "skip_softmax" + assert group["threshold_scale_factor"]["prefill"] == {"a": 7.9, "b": 8.6} + assert group["threshold_scale_factor"]["formula"] == "a * exp(b * target_sparsity)" + assert group["target_sparsity"] == {"prefill": 0.4, "decode": 0.4} + assert config["producer"]["name"] == "modelopt" + + def test_preserves_nm_groups_and_replaces_old_skip_group(self): + existing = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4}, + "group_1": { + "algorithm": "skip_softmax", + "threshold_scale_factor": {"prefill": {"a": 1.0, "b": 1.0}}, + }, + } + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + groups = config["config_groups"] + assert len(groups) == 2 + assert groups["group_0"]["algorithm"] == "skip_softmax" + assert groups["group_0"]["threshold_scale_factor"]["prefill"]["a"] == 7.9 + assert groups["group_1"]["algorithm"] == "sparse_softmax" + assert groups["group_1"]["sparsity_n"] == 2 + + def test_round_trips_through_serving_loader(self): + config = build_sparse_attention_config(self._PARAMS, {"prefill": 0.5, "decode": 0.3}) + hf_config = SimpleNamespace(sparse_attention_config=config) + loaded = load_from_checkpoint_metadata(hf_config) + assert loaded is not None + sparse_cfg, preset = loaded + assert preset == "CHECKPOINT_CALIBRATED_SOFTMAX_SKIP" + layer_cfg = sparse_cfg["sparse_cfg"]["*attn*"] + assert layer_cfg["method"] == "triton_skip_softmax" + assert layer_cfg["threshold_scale_factor"]["decode"] == {"a": 0.12, "b": 9.8} + assert layer_cfg["target_sparse_ratio"] == {"prefill": 0.5, "decode": 0.3} + + def test_round_trip_with_preserved_nm_group_activates_both(self): + existing = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4} + } + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + loaded = load_from_checkpoint_metadata(SimpleNamespace(sparse_attention_config=config)) + assert loaded is not None + sparse_cfg, preset = loaded + assert preset == "CHECKPOINT_CALIBRATED_SOFTMAX_SKIP_SPARSE_SOFTMAX" + layer_cfg = sparse_cfg["sparse_cfg"]["*attn*"] + assert layer_cfg["sparsity_n"] == 2 + assert "threshold_scale_factor" in layer_cfg From c5def509294fc440cb4a42d2e1af84117fcdfdad Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 18 Jul 2026 21:27:13 -0700 Subject: [PATCH 05/14] Add changelog entry for vLLM skip-softmax calibration Signed-off-by: Kai Xu --- CHANGELOG.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 95e3f12ef58..71af7779b53 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,10 @@ Changelog **New Features** +*Sparsity* + +- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. See ``examples/vllm_serve/calibrate_sparse_attn.py``. + *Quantization* - Add a Muse Glimmer AutoQuantize recipe that searches language-model MLP projections, self-attention projections, and ``lm_head`` over W4A16 NVFP4 Four-Over-Six, FP8, and BF16 fallback at 5.5 effective bits while leaving the vision tower unquantized. From ef1ea216d0eb2d6cd2e7e62b1e49cc0083e6a64b Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 18 Jul 2026 21:27:17 -0700 Subject: [PATCH 06/14] Default vLLM calibration prompts to the RULER dataset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Build calibration prompts with RulerDatasetBuilder — the same dataset the PyTorch (HF) calibration path uses — so vLLM- and HF-calibrated thresholds are fit on identical data (defaults mirror the HF path: 24 samples, max_seqlen 32768; NIAH essay haystack via download_ruler_data.sh and --calib_data_dir). The built-in demo prompt fallback is removed; --prompts_file remains as an explicit override for custom calibration data. Signed-off-by: Kai Xu --- examples/vllm_serve/README.md | 8 +- examples/vllm_serve/calibrate_sparse_attn.py | 91 ++++++++++++++------ 2 files changed, 74 insertions(+), 25 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 424163d5371..43628fbb7a3 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -183,11 +183,17 @@ If the checkpoint has no `sparse_attention_config`, the sparse-only installer pa Instead of the HF path in step 1, thresholds can be calibrated directly through vLLM — over the paged KV cache, for both prefill and decode, with tensor parallelism: ```bash +# One-time: fetch the RULER essay haystack +bash ../llm_sparsity/attention_sparsity/download_ruler_data.sh + python calibrate_sparse_attn.py \ - --prompts_file prompts.txt --target_sparse_ratio 0.5 \ + --calib_data_dir ../llm_sparsity/attention_sparsity/data \ + --target_sparse_ratio 0.5 \ --decode_tokens 32 --tensor_parallel_size 8 --update_checkpoint_config ``` +Calibration prompts default to the **RULER dataset** via the same `RulerDatasetBuilder` the HF calibration path uses (`--calib_samples` / `--calib_max_seqlen` mirror the HF defaults of 24 / 32768), so vLLM- and PyTorch-calibrated thresholds are fit on identical data. `--prompts_file` (one prompt per line) substitutes custom calibration data. + `install_vllm_skip_softmax_calibration` (called by `sparse_attn_worker.SkipSoftmaxCalibWorker` at model load) swaps calibration adapters onto every attention layer after validating all of them — eager execution is required, model and KV-cache dtypes must be fp16/bf16, and no attention Q/K/P/V fakequant may be active. During `llm.generate`, the paged Triton calibration kernel computes full dense attention (generation is numerically unchanged) while counting, per candidate threshold, how many KV tiles the skip criterion would drop. The driver then collects **raw tile counts from every TP rank** (each rank only measures its head shard), merges them, fits `scale_factor = a * exp(b * sparsity)` once per phase, and writes the same canonical `sparse_attention_config` block the HF export produces — preserving any exported N:M sparse-softmax groups — so the serving workflow above picks it up unchanged. Calibration and serving measure skipping at the same fixed 128x128 tile geometry (active skip-softmax launches bypass the autotuner), so the sparsity realized at serve time matches the calibrated `(a, b)` model. diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py index 6c36f8f65ec..65091485981 100644 --- a/examples/vllm_serve/calibrate_sparse_attn.py +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -32,14 +32,16 @@ Usage: python calibrate_sparse_attn.py \ - --prompts_file prompts.txt \ --target_sparse_ratio 0.5 \ --decode_tokens 32 \ --update_checkpoint_config -``--prompts_file`` is one prompt per line; longer, varied-length prompts give a -better fit. With no file, a tiny built-in demo set is used (fine for a smoke -test, not for a real fit). +Calibration prompts default to the RULER dataset — the same +``RulerDatasetBuilder`` the PyTorch (HF) calibration path uses — so both paths +calibrate on identical data. NIAH tasks need the essay haystack downloaded by +``examples/llm_sparsity/attention_sparsity/download_ruler_data.sh`` (point +``--calib_data_dir`` at its ``data`` directory). ``--prompts_file`` (one prompt +per line) overrides the RULER set with custom calibration data. """ import argparse @@ -55,25 +57,39 @@ merge_phase_counts, ) -_DEMO_PROMPTS = [ - "Summarize the history of computing in a few paragraphs. " * 40, - "Explain how attention works in transformer models. " * 60, - "Write a detailed essay about renewable energy sources. " * 80, -] +def _load_prompts(llm, args) -> list[str]: + """Load override prompts from a file, or build the default RULER set.""" + if args.prompts_file is not None: + lines = [ + ln.strip() for ln in Path(args.prompts_file).read_text().splitlines() if ln.strip() + ] + if not lines: + raise ValueError(f"No prompts found in {args.prompts_file}") + print(f"[ModelOpt] Loaded {len(lines)} calibration prompts from {args.prompts_file}") + return lines -def _load_prompts(prompts_file: str | None) -> list[str]: - if prompts_file is None: - print( - "[ModelOpt] No --prompts_file given; using a tiny built-in demo set. " - "Pass real, varied-length prompts for a usable fit." - ) - return _DEMO_PROMPTS - lines = [ln.strip() for ln in Path(prompts_file).read_text().splitlines() if ln.strip()] - if not lines: - raise ValueError(f"No prompts found in {prompts_file}") - print(f"[ModelOpt] Loaded {len(lines)} calibration prompts from {prompts_file}") - return lines + # Same dataset as the HF calibration path (calibration/calibrate.py), so the + # vLLM- and PyTorch-calibrated thresholds are fit on identical data. + from modelopt.torch.sparsity.attention_sparsity.calibration.ruler_dataset import ( + RulerDatasetBuilder, + ) + + builder = RulerDatasetBuilder( + samples=args.calib_samples, + max_seqlen=args.calib_max_seqlen, + tokenizer_name_or_path=llm.get_tokenizer(), + max_length_filter=int(args.calib_max_seqlen * 1.5), + data_dir=args.calib_data_dir, + ) + samples = builder.build_calibration_dataset() + prompts = [sample["input"] for sample in samples] + lengths = sorted(sample["length"] for sample in samples) + print( + f"[ModelOpt] Built {len(prompts)} RULER calibration prompts " + f"(token lengths {lengths[0]}..{lengths[-1]})" + ) + return prompts def _existing_sparse_config(ckpt: str) -> dict | None: @@ -108,7 +124,33 @@ def _write_config(ckpt: str, sparse_config: dict, update_checkpoint: bool) -> No def main(): parser = argparse.ArgumentParser(description="Calibrate skip-softmax thresholds via vLLM") parser.add_argument("model", type=str, help="Path to the HF checkpoint to calibrate") - parser.add_argument("--prompts_file", type=str, default=None, help="One prompt per line") + parser.add_argument( + "--prompts_file", + type=str, + default=None, + help="Optional custom calibration prompts (one per line), overriding the " + "default RULER dataset", + ) + parser.add_argument( + "--calib_samples", + type=int, + default=24, + help="Total RULER samples, distributed across length bins (HF-path default: 24)", + ) + parser.add_argument( + "--calib_max_seqlen", + type=int, + default=32768, + help="Maximum RULER sequence length; length bins descend in powers of 2. " + "Must fit within --max_model_len together with --decode_tokens.", + ) + parser.add_argument( + "--calib_data_dir", + type=str, + default=None, + help="RULER data directory containing the 'essays' haystack (populated by " + "examples/llm_sparsity/attention_sparsity/download_ruler_data.sh)", + ) parser.add_argument( "--target_sparse_ratio", type=float, @@ -176,8 +218,6 @@ def main(): from vllm import LLM, SamplingParams - prompts = _load_prompts(args.prompts_file) - llm_kwargs = { "model": args.model, "worker_cls": "sparse_attn_worker.SkipSoftmaxCalibWorker", @@ -207,6 +247,9 @@ def main(): llm_kwargs.update(extra) llm = LLM(**llm_kwargs) + # Built after engine init so the RULER builder reuses the engine's tokenizer. + prompts = _load_prompts(llm, args) + trials = list(DEFAULT_THRESHOLD_TRIALS) n_layers = llm.collective_rpc("sparse_calib_enable", args=(trials,))[0] status = llm.collective_rpc("sparse_calib_status")[0] From ecb8965b07a298560360ccccf5e592da82c969a2 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sat, 18 Jul 2026 23:11:04 -0700 Subject: [PATCH 07/14] 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 71af7779b53..5df67102c48 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog *Sparsity* -- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. See ``examples/vllm_serve/calibrate_sparse_attn.py``. +- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. Configurations that cannot compile the 128x128 tile are rejected rather than re-tiled, and skip-softmax cannot be combined with attention quantization. See ``examples/vllm_serve/calibrate_sparse_attn.py``. *Quantization* diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 43628fbb7a3..34a802c85e2 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -194,7 +194,7 @@ python calibrate_sparse_attn.py \ Calibration prompts default to the **RULER dataset** via the same `RulerDatasetBuilder` the HF calibration path uses (`--calib_samples` / `--calib_max_seqlen` mirror the HF defaults of 24 / 32768), so vLLM- and PyTorch-calibrated thresholds are fit on identical data. `--prompts_file` (one prompt per line) substitutes custom calibration data. -`install_vllm_skip_softmax_calibration` (called by `sparse_attn_worker.SkipSoftmaxCalibWorker` at model load) swaps calibration adapters onto every attention layer after validating all of them — eager execution is required, model and KV-cache dtypes must be fp16/bf16, and no attention Q/K/P/V fakequant may be active. During `llm.generate`, the paged Triton calibration kernel computes full dense attention (generation is numerically unchanged) while counting, per candidate threshold, how many KV tiles the skip criterion would drop. The driver then collects **raw tile counts from every TP rank** (each rank only measures its head shard), merges them, fits `scale_factor = a * exp(b * sparsity)` once per phase, and writes the same canonical `sparse_attention_config` block the HF export produces — preserving any exported N:M sparse-softmax groups — so the serving workflow above picks it up unchanged. +`install_vllm_skip_softmax_calibration` (called by `sparse_attn_worker.SkipSoftmaxCalibWorker` at model load) swaps calibration adapters onto every attention layer after validating all of them — eager execution is required, model and KV-cache dtypes must be fp16/bf16, and no attention Q/K/P/V fakequant may be active. During `llm.generate`, the paged Triton calibration kernel computes full dense attention — no sparsification is applied to generation, though the dense kernel's numerics differ slightly from the native backend's — while counting, per candidate threshold, how many KV tiles the skip criterion would drop. The driver then collects **raw tile counts from every TP rank** (each rank only measures its head shard), merges them, fits `scale_factor = a * exp(b * sparsity)` once per phase, and writes the same canonical `sparse_attention_config` block the HF export produces — preserving any exported N:M sparse-softmax groups — so the serving workflow above picks it up unchanged. Calibration and serving measure skipping at the same fixed 128x128 tile geometry (active skip-softmax launches bypass the autotuner), so the sparsity realized at serve time matches the calibrated `(a, b)` model. @@ -212,7 +212,7 @@ report = install_vllm_nvfp4_attention(model_runner, sparse_cfg="checkpoint") Limitations: -- vLLM V1 chunked prefill and prefix-cache suffix attention are supported by offsetting query positions into the longer KV span. +- vLLM V1 chunked prefill and prefix-cache suffix attention are supported by offsetting query positions into the longer KV span. This applies to sparse-only serving; quantized attention installs and skip-softmax calibration reject `enable_prefix_caching` (quantize-on-write and per-request measurement both require uncached prefills). - `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`. ### Compact NVFP4 attention worker @@ -229,7 +229,7 @@ python vllm_serve_sparse_attn.py -tp 8 \ The installer supports both FlashInfer and FlashAttention, and the worker prints the installed adapter counts. Pass `--attention-backend FLASHINFER` or `--attention-backend FLASH_ATTN` only when an explicit override is needed. -This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant format to Q/K/P/V. Q is dynamic; missing K/V scales default to global scale 1.0, and P defaults to amax 1.0. Existing scalar attention amax values are preserved, but this path does not calibrate or restore them itself. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored. +This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant format to Q/K/P/V. Q is dynamic; missing K/V scales default to global scale 1.0, and P defaults to amax 1.0. Existing scalar attention amax values are preserved, but this path does not calibrate or restore them itself. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored for N:M sparse softmax; calibrated skip-softmax groups are rejected in combination with attention quantization, because quantized Q/K/P change the score distribution the skip thresholds were calibrated on. Decode uses a fixed 32-split, 128-key-tile schedule. P QDQ consumes split-local, unnormalized online-softmax probabilities, so changing that schedule can change diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py index 65091485981..bc5ed8281c9 100644 --- a/examples/vllm_serve/calibrate_sparse_attn.py +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -257,9 +257,11 @@ def main(): print(f"[ModelOpt] Active sparse impls: {status['impl_types']}") # generate() drives prefill (prefill-phase stats) then decode_tokens decode - # steps (decode-phase stats). The calibration kernel computes full attention, - # so the generated text is unaffected — only tile-skip counts are recorded. - sampling = SamplingParams(temperature=0.0, max_tokens=args.decode_tokens) + # steps (decode-phase stats). No sparsification is applied during + # calibration — the kernel computes full dense attention while recording + # tile-skip counts. ignore_eos forces the full decode length so early EOS + # cannot thin the decode-phase statistics. + sampling = SamplingParams(temperature=0.0, max_tokens=args.decode_tokens, ignore_eos=True) llm.generate(prompts, sampling) # Aggregate RAW counts from every TP rank (each rank only measures its @@ -268,12 +270,16 @@ def main(): merged = merge_phase_counts(rank_counts) calibration_params = fit_from_counts(merged, trials, fit_logspace=args.fit_logspace) - if not calibration_params: + requested_phases = ["prefill"] + (["decode"] if args.decode_tokens > 0 else []) + missing = [phase for phase in requested_phases if phase not in calibration_params] + if missing: print( - "[ModelOpt] Calibration produced no valid fit — try more/longer prompts " + f"[ModelOpt] Calibration FAILED: no valid fit for phase(s) {', '.join(missing)}. " + "No config was written — a partially calibrated export would silently serve " + "the missing phase dense. Try more/longer prompts (and more decode tokens) " "so observed sparsity spans the (10%, 90%) fitting window." ) - return + sys.exit(1) sparse_config = build_sparse_attention_config( calibration_params, diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 4d69a13a4cb..145de3a51e2 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -23,7 +23,6 @@ """ import math -import warnings from typing import Any import torch @@ -1042,36 +1041,27 @@ def grid(META): # (e.g. BLOCK_N=32) would realize a different sparsity than # calibrated, so skip launches always use the calibration tile. # - # If the 128-row tile exceeds the device's shared memory (fp32 - # inputs on ~100KB-smem GPUs), halve BLOCK_M and retry. BLOCK_N - # — the calibrated KV skip granularity — is never reduced. - # Smaller BLOCK_M splits the per-tile skip vote into narrower - # row groups: per-row numerics stay protected by the same - # threshold, but realized sparsity can exceed the calibrated - # target, so warn when stepping down. - block_m = _P_QDQ_MEASURE_BLOCK_M if p_qdq_mode else _MEASURE_BLOCK_M - while True: - try: - _attn_fwd.fn[grid]( - *fwd_args, - **fwd_kwargs, - BLOCK_M=block_m, - BLOCK_N=_MEASURE_BLOCK_N, - num_warps=_MEASURE_NUM_WARPS, - num_stages=_MEASURE_NUM_STAGES, - ) - break - except triton.runtime.errors.OutOfResources: - if block_m <= 16: - raise - block_m //= 2 - if apply_skip: - warnings.warn( - f"skip-softmax tile BLOCK_M reduced to {block_m} " - "(shared memory limit); realized sparsity may " - "exceed the calibrated target on this device/dtype", - stacklevel=2, - ) + # The tile is a contract, not a preference: configurations that + # cannot compile it (e.g. fp32 inputs on ~100KB-shared-memory + # GPUs) are rejected rather than re-tiled, because a different + # tile realizes a different sparsity than was calibrated. + try: + _attn_fwd.fn[grid]( + *fwd_args, + **fwd_kwargs, + BLOCK_M=_P_QDQ_MEASURE_BLOCK_M if p_qdq_mode else _MEASURE_BLOCK_M, + BLOCK_N=_MEASURE_BLOCK_N, + num_warps=_MEASURE_NUM_WARPS, + num_stages=_MEASURE_NUM_STAGES, + ) + except triton.runtime.errors.OutOfResources as err: + raise RuntimeError( + "skip-softmax requires the fixed 128x128 calibration tile, " + f"which exceeds this GPU's shared memory for {q.dtype} " + f"inputs ({err}). Use fp16/bf16 inputs or a device with " + "more shared memory; re-tiling would change the " + "calibrated sparsity contract." + ) from err else: _attn_fwd[grid]( *fwd_args, diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py index a37b447a686..5e604c039ca 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py @@ -28,6 +28,33 @@ from ..stats_manager import SparseAttentionStatsManager from ..utils import get_sparse_attention_modules +# Canonical skip-softmax threshold sweep — should span sparsities from ~10% to +# ~95%. Shared by the HF calibration path (this class's default) and the vLLM +# calibration path (``plugins/sparse_attn_calibration.py``), so both fit on the +# same trial grid. +DEFAULT_THRESHOLD_TRIALS = [ + 1e-6, + 5e-6, + 1e-5, + 5e-5, + 1e-4, + 5e-4, + 1e-3, + 5e-3, + 1e-2, + 2e-2, + 5e-2, + 1e-1, + 2e-1, + 3e-1, + 5e-1, + 7e-1, + 8e-1, + 9e-1, + 9.5e-1, + 9.9e-1, +] + class DynamicThresholdCalibrator: """Dynamic threshold calibrator using Exponential model. @@ -67,28 +94,7 @@ def __init__( where scale_factors span many orders of magnitude. """ # Default threshold trials if not provided - self.threshold_trials = threshold_trials or [ - 1e-6, - 5e-6, - 1e-5, - 5e-5, - 1e-4, - 5e-4, - 1e-3, - 5e-3, - 1e-2, - 2e-2, - 5e-2, - 1e-1, - 2e-1, - 3e-1, - 5e-1, - 7e-1, - 8e-1, - 9e-1, - 9.5e-1, - 9.9e-1, - ] + self.threshold_trials = threshold_trials or list(DEFAULT_THRESHOLD_TRIALS) self.fit_logspace = fit_logspace def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dict[str, Any]: diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py index f60bef1b867..b66a5c27436 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -35,6 +35,10 @@ import modelopt +# One canonical sweep for both calibration paths: re-exported from the HF-path +# calibrator so the vLLM path fits on the identical trial grid. +from ..calibration.calibrator import DEFAULT_THRESHOLD_TRIALS + __all__ = [ "DEFAULT_THRESHOLD_TRIALS", "build_sparse_attention_config", @@ -45,22 +49,6 @@ "stats_from_counts", ] -# Default threshold sweep — should span sparsities from ~10% to ~95%. -DEFAULT_THRESHOLD_TRIALS = [ - 1e-4, - 1e-3, - 5e-3, - 1e-2, - 3e-2, - 5e-2, - 1e-1, - 2e-1, - 3e-1, - 5e-1, - 7e-1, - 9e-1, -] - _PHASES = ("prefill", "decode") @@ -81,24 +69,40 @@ def merge_count_records(sources: list[list[dict]]) -> list[dict]: records for one phase. All sources observe the same launches in the same order, so records align by index; tile counts are additive across both layers and head-sharded TP ranks. + + Alignment is a contract: every source must report the same number of + samples, the same per-sample lengths, and the same threshold-vector width. + Any mismatch indicates a collection bug and raises rather than silently + dropping records. """ - sources = [source for source in sources if source] if not sources: return [] - num_samples = min(len(source) for source in sources) + sample_counts = {len(source) for source in sources} + if len(sample_counts) != 1: + raise ValueError( + "Misaligned calibration records: sources disagree on sample count " + f"({sorted(sample_counts)})" + ) + num_samples = sample_counts.pop() merged = [] for i in range(num_samples): base = sources[0][i] - total = [0] * len(base["total_tiles"]) - skipped = [0] * len(base["skipped_tiles"]) - for source in sources: - record = source[i] + width = len(base["total_tiles"]) + total = [0] * width + skipped = [0] * width + for record in (source[i] for source in sources): if record["sample_length"] != base["sample_length"]: raise ValueError( "Misaligned calibration records: sample lengths differ across " f"sources at index {i} ({record['sample_length']} vs " f"{base['sample_length']})" ) + if len(record["total_tiles"]) != width or len(record["skipped_tiles"]) != width: + raise ValueError( + "Misaligned calibration records: threshold-vector widths differ " + f"across sources at index {i} " + f"({len(record['total_tiles'])}/{len(record['skipped_tiles'])} vs {width})" + ) total = [a + b for a, b in zip(total, record["total_tiles"])] skipped = [a + b for a, b in zip(skipped, record["skipped_tiles"])] merged.append( @@ -118,13 +122,21 @@ def merge_phase_counts(rank_counts: list[dict[str, list[dict]]]) -> dict[str, li ``{"prefill": [...], "decode": [...]}`` dict per rank, as returned by ``collect_calibration_counts``). Use ALL ranks: with tensor parallelism each rank only measures its attention-head shard, so any single rank's - counts are incomplete. + counts are incomplete. A phase recorded by some ranks but not others + indicates a collection bug and raises. """ phases = {phase for rank in rank_counts for phase in rank} - return { - phase: merge_count_records([rank.get(phase, []) for rank in rank_counts]) - for phase in phases - } + merged: dict[str, list[dict]] = {} + for phase in phases: + sources = [rank.get(phase, []) for rank in rank_counts] + empty = sum(1 for source in sources if not source) + if empty and empty != len(sources): + raise ValueError( + f"Misaligned calibration records: {empty}/{len(sources)} rank(s) " + f"recorded no {phase!r} samples while others did" + ) + merged[phase] = merge_count_records([] if empty else sources) + return merged def stats_from_counts(count_records: list[dict]) -> list[dict]: @@ -180,7 +192,7 @@ def fit_from_counts( def _normalize_target_sparsity(target_sparsity: dict[str, float] | float) -> dict[str, float]: - if isinstance(target_sparsity, (int, float)): + if isinstance(target_sparsity, int | float): return {phase: float(target_sparsity) for phase in _PHASES} return {phase: float(target_sparsity.get(phase, 0.5)) for phase in _PHASES} @@ -230,7 +242,14 @@ def build_sparse_attention_config( for idx, group in enumerate(preserved, start=1): config_groups[f"group_{idx}"] = group - return { + result: dict[str, Any] = { "config_groups": config_groups, "producer": {"name": "modelopt", "version": modelopt.__version__}, } + # Legacy checkpoints carry N:M parameters as a top-level ``sparse_softmax`` + # dict (read by the serving loader ahead of group params) — preserve it so + # recalibration does not silently reset N:M settings to defaults. + legacy_sparse_softmax = (existing_config or {}).get("sparse_softmax") + if isinstance(legacy_sparse_softmax, dict): + result["sparse_softmax"] = legacy_sparse_softmax + return result diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 80e30e5fc18..2ebf2aa39ac 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -28,6 +28,7 @@ """ import functools +import importlib import inspect import math import warnings @@ -260,6 +261,32 @@ def _calibration_active(impl) -> bool: ) +def _flashinfer_kv_cache_layout() -> str | None: + """Best-effort query of vLLM's configured FlashInfer KV-cache layout. + + Returns ``"NHD"`` / ``"HND"`` when the running vLLM exposes the layout + (newer releases select HND for some Blackwell FlashInfer paths), or + ``None`` when it cannot be determined — callers then fall back to the + shape-based check in :func:`_forward_calibrate`. + """ + for module_name in ( + "vllm.v1.attention.backends.utils", + "vllm.attention.backends.utils", + ): + try: + module = importlib.import_module(module_name) + except ImportError: + continue + getter = getattr(module, "get_kv_cache_layout", None) + if getter is None: + continue + try: + return str(getter()) + except Exception: + return None + return None + + def _forward_calibrate( impl, *, @@ -277,8 +304,9 @@ def _forward_calibrate( Each scheduled request is calibrated independently (batch=1) so its KV length is the per-sample length the exponential fit needs, and so the kernel keeps the uniform-length contract it was validated against. The - kernel computes full attention, so ``output`` is written densely and the - forward pass is numerically unchanged. + kernel computes full attention, so ``output`` is written densely — no + sparsification is applied to generation (the dense Triton kernel's + numerics may differ slightly from the native backend's). Phase and causality are decided per request: ``q_len == 1`` is a decode step (full-cache, non-causal); ``q_len > 1`` is (chunked) prefill (causal @@ -839,6 +867,14 @@ def prepare_modelopt(): raise ValueError( "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" ) + layout = _flashinfer_kv_cache_layout() + if layout is not None and layout.upper() != "NHD": + # Authoritative layout metadata beats the shape heuristic (which is + # ambiguous when page_size equals the per-rank KV-head count). + raise NotImplementedError( + f"FlashInfer KV-cache layout {layout!r} is unsupported for " + "skip-softmax calibration; only NHD is supported" + ) # Order matters: releases that update the KV cache inside forward must # write the current K/V before the calibrate kernel reads the cache. prepare_modelopt() @@ -1031,10 +1067,21 @@ def collect_calibration_counts(model) -> dict[str, list[dict]]: """ from .sparse_attn_calibration import merge_count_records, split_records_by_phase - per_phase_layers: dict[str, list[list[dict]]] = {} - for impl in iter_sparse_impls(model): - split = split_records_by_phase(getattr(impl, "_calib_records", [])) - for phase, records in split.items(): - if records: - per_phase_layers.setdefault(phase, []).append(records) - return {phase: merge_count_records(layers) for phase, layers in per_phase_layers.items()} + splits = [ + split_records_by_phase(getattr(impl, "_calib_records", [])) + for impl in iter_sparse_impls(model) + ] + phases = {phase for split in splits for phase, records in split.items() if records} + counts: dict[str, list[dict]] = {} + for phase in phases: + sources = [split.get(phase, []) for split in splits] + empty = sum(1 for source in sources if not source) + if empty: + # Every layer sees every launch, so a layer with no records for a + # phase others measured indicates a collection bug. + raise ValueError( + f"Misaligned calibration records: {empty}/{len(sources)} attention " + f"layer(s) recorded no {phase!r} samples while others did" + ) + counts[phase] = merge_count_records(sources) + return counts diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index 04d703f46a1..a00a7ca012e 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -249,6 +249,13 @@ def _device_capability_error(device) -> str | None: return None +def _skip_softmax_active(sparse_kw: dict[str, Any]) -> bool: + """Return whether a layer's sparse config contains skip-softmax work.""" + return bool(sparse_kw) and ( + "skip_softmax_threshold" in sparse_kw or "threshold_scale_factor" in sparse_kw + ) + + def _sparse_graph_error(sparse_kw: dict[str, Any], mode) -> str | None: from vllm.config.compilation import CUDAGraphMode @@ -354,6 +361,16 @@ def _plan_vllm_attention( plans = [] for name, module, sparse_kw in candidates: reasons = _layer_errors(module) + if quantize and _skip_softmax_active(sparse_kw): + # Quantized Q/K/P change the attention-score distribution the skip + # thresholds were calibrated on, so the calibrated sparsity contract + # no longer holds. N:M sparse softmax has no calibrated threshold + # and composes with quantization. + reasons.append( + "skip-softmax cannot be combined with attention quantization; " + "serve skip-softmax unquantized or drop the skip_softmax group " + "(N:M sparse softmax composes with quantization)" + ) device = dtype = None if quantize: device, dtype = quant_plugin._get_device_dtype(module) @@ -561,6 +578,12 @@ def install_vllm_skip_softmax_calibration(model_runner) -> VllmAttentionInstallR plans.append( _AttentionPlan(name, module, new_impl, {}, None, None, requires_flashinfer_patch) ) + if any(plan.requires_flashinfer_patch for plan in plans): + layout = attention_plugin._flashinfer_kv_cache_layout() + if layout is not None and layout.upper() != "NHD": + errors.append( + f"FlashInfer KV-cache layout {layout!r} is unsupported for calibration (NHD only)" + ) _raise_unsupported(errors, "skip-softmax calibration") plan = _InstallPlan(model_runner, tuple(plans), False, "SKIP_SOFTMAX_CALIBRATION") diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py index 9e55879dbb6..7020f90d7c5 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py @@ -116,27 +116,27 @@ def test_high_block_id_pointer_arithmetic(self): """Block IDs whose int32 byte offsets would wrap still read correctly.""" num_kv_heads, head_dim, page_size = 2, 64, 16 block_elems = page_size * num_kv_heads * head_dim - # Smallest block ID whose element offset exceeds int32. + # Smallest block ID whose element offset exceeds int32. V aliases the K + # cache storage (same values on both operands), halving the allocation. high_block = (2**31) // block_elems + 1 - bytes_needed = 2 * (high_block + 1) * block_elems * 2 # K and V caches, bf16 + bytes_needed = (high_block + 1) * block_elems * 2 # one shared K/V cache, bf16 free, _ = torch.cuda.mem_get_info() if free < bytes_needed + (2 << 30): pytest.skip(f"needs ~{bytes_needed / 2**30:.1f} GiB free GPU memory") torch.manual_seed(2) num_heads = 4 - q, k, v = make_qkv(page_size, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) + q, k, _ = make_qkv(page_size, num_heads, num_kv_heads, head_dim, dtype=torch.bfloat16) k_cache = torch.zeros( high_block + 1, page_size, num_kv_heads, head_dim, device="cuda", dtype=torch.bfloat16 ) - v_cache = torch.zeros_like(k_cache) + v_cache = k_cache # alias: V reads the same storage (and the same values) k_cache[high_block] = k - v_cache[high_block] = v block_table = torch.tensor([[high_block]], device="cuda", dtype=torch.int32) locs, lens = make_varlen_meta([page_size]) out_ref, counters_ref = attention_calibrate( - q, k, v, locs, lens, page_size, is_causal=True, threshold_trials=TRIALS + q, k, k, locs, lens, page_size, is_causal=True, threshold_trials=TRIALS ) k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) out_paged, counters_paged = attention_calibrate( diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py index 2806e39b4a5..598f97a3fda 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py @@ -356,11 +356,16 @@ class TestBackwardWithSparsity: """Backward pass with skip-softmax (covers _attn_bwd_dq / _attn_bwd_dkdv).""" def test_backward_with_skip_softmax(self): - """Backward pass runs without error when skip-softmax is active.""" + """Backward pass runs without error when skip-softmax is active. + + fp16 rather than fp32: active skip launches require the fixed 128x128 + calibration tile, which fp32 inputs cannot compile on ~100KB-shared- + memory GPUs (such configurations are rejected by design). + """ seq_len, num_heads, head_dim = 128, 4, 64 scale = 1.0 / (head_dim**0.5) torch.manual_seed(7) - q, k, v = make_qkv(seq_len, num_heads, num_heads, head_dim, dtype=torch.float32) + q, k, v = make_qkv(seq_len, num_heads, num_heads, head_dim, dtype=torch.float16) q.requires_grad_(True) k.requires_grad_(True) v.requires_grad_(True) diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py index caefb467555..033952d34eb 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py @@ -217,17 +217,26 @@ def test_triton_matches_pytorch_reference(self): locs = torch.arange(batch, device="cuda", dtype=torch.int32) * seq_len lens = torch.full((batch,), seq_len, device="cuda", dtype=torch.int32) - triton_out = attention( - q_flat, - k_flat, - v_flat, - locs, - lens, - seq_len, - is_causal=True, - softmax_scale=scale, - skip_softmax_threshold=threshold, - ) + try: + triton_out = attention( + q_flat, + k_flat, + v_flat, + locs, + lens, + seq_len, + is_causal=True, + softmax_scale=scale, + skip_softmax_threshold=threshold, + ) + except RuntimeError as err: + if "shared memory" in str(err): + # Active skip requires the fixed 128x128 calibration tile; fp32 + # inputs cannot compile it on ~100KB-shared-memory GPUs and the + # configuration is rejected by design. fp32 is kept here for a + # tight reference comparison on GPUs that support it. + pytest.skip("fp32 skip tile exceeds this GPU's shared memory") + raise triton_out_4d = triton_out.view(batch, seq_len, num_heads, head_dim).permute(0, 2, 1, 3) # Both outputs should be close — same algorithm, different implementations. diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py index e5192e7e7f3..62a75b92d8f 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -22,6 +22,7 @@ from torch import nn from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.attention.backends.flashinfer import FlashInferImpl from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as attention_plugin @@ -126,6 +127,106 @@ def test_rejects_model_without_attention_layers(self): vllm_runtime.install_vllm_skip_softmax_calibration(runner) +class TestQuantSkipRejection: + """Skip-softmax cannot be combined with attention quantization.""" + + _CALIBRATED_META = { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "threshold_scale_factor": {"prefill": {"a": 7.9, "b": 8.6}}, + } + } + } + _NM_META = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4} + } + } + + def test_quantized_install_rejects_calibrated_skip(self): + attention = _bare_attention() + runner = _model_runner( + nn.ModuleDict({"attn": attention}), sparse_metadata=self._CALIBRATED_META + ) + with pytest.raises( + NotImplementedError, match="cannot be combined with attention quantization" + ): + vllm_runtime.install_vllm_nvfp4_attention(runner) + assert not isinstance(attention.impl, ModelOptSparseAttentionImpl) + + def test_quantized_plan_allows_nm_sparsity(self, monkeypatch): + from modelopt.torch.quantization.plugins import vllm as quant_plugin + + monkeypatch.setattr( + quant_plugin, + "_get_device_dtype", + lambda module: (torch.device("cpu"), torch.float16), + ) + runner = _model_runner( + nn.ModuleDict({"attn": _bare_attention()}), sparse_metadata=self._NM_META + ) + plan = vllm_runtime._plan_vllm_attention(runner, quantize=True, sparse_cfg="checkpoint") + assert len(plan.layers) == 1 + assert plan.layers[0].sparse_kw.get("sparsity_n") == 2 + + +class TestFlashInferLayoutGuard: + """HND FlashInfer caches are rejected via layout metadata, pre-measurement.""" + + def test_installer_rejects_hnd_layout(self, monkeypatch): + monkeypatch.setattr(attention_plugin, "_flashinfer_kv_cache_layout", lambda: "HND") + attention = _bare_attention(FlashInferImpl) + original_impl = attention.impl + runner = _model_runner(nn.ModuleDict({"attn": attention})) + with pytest.raises(NotImplementedError, match="HND"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + assert attention.impl is original_impl + + def test_forward_rejects_hnd_layout_before_cache_write(self, monkeypatch): + monkeypatch.setattr(attention_plugin, "_flashinfer_kv_cache_layout", lambda: "HND") + writes = [] + monkeypatch.setattr( + attention_plugin, + "_maybe_update_flashinfer_cache", + lambda *args, **kwargs: writes.append(1), + ) + num_heads, num_kv_heads, head_dim, page = 4, 2, 64, 16 + impl = SimpleNamespace( + num_kv_heads=num_kv_heads, + head_size=head_dim, + scale=1.0 / (head_dim**0.5), + _calibrate=True, + _calib_threshold_trials=list(TRIALS), + _calib_records=[], + ) + kv_cache = torch.zeros(3, 2, page, num_kv_heads, head_dim, dtype=torch.bfloat16) + attn_metadata = SimpleNamespace( + _modelopt_block_table=torch.zeros(1, 1, dtype=torch.int32), + _modelopt_seq_lens=torch.tensor([8], dtype=torch.int32), + _modelopt_query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + _modelopt_num_actual_tokens=1, + _modelopt_max_query_len=1, + _modelopt_max_seq_len=8, + _modelopt_causal=False, + slot_mapping=torch.zeros(1, dtype=torch.int64), + ) + q = torch.zeros(1, num_heads, head_dim, dtype=torch.bfloat16) + with pytest.raises(NotImplementedError, match="HND"): + attention_plugin._flashinfer_forward( + impl, + None, + None, + q, + q[:, :num_kv_heads], + q[:, :num_kv_heads], + kv_cache, + attn_metadata, + output=torch.empty_like(q), + ) + assert not writes, "layout must be validated before the cache write" + + class TestSparseOnlyGraphGuard: """Commit-contract: the calibrated-decode graph guard is not quantize-gated.""" diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py index 77f836c9205..85e459dd650 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py @@ -63,12 +63,11 @@ def test_merge_sums_counts_elementwise(self): merged = merge_count_records([layer_a, layer_b]) assert merged == [{"sample_length": 128, "total_tiles": [20, 20], "skipped_tiles": [3, 7]}] - def test_merge_truncates_to_shortest_source(self): + def test_merge_rejects_ragged_sources(self): long = [_record("prefill", 1, [1], [0]), _record("prefill", 2, [1], [1])] short = [_record("prefill", 1, [1], [1])] - merged = merge_count_records([long, short]) - assert len(merged) == 1 - assert merged[0]["skipped_tiles"] == [1] + with pytest.raises(ValueError, match="disagree on sample count"): + merge_count_records([long, short]) def test_merge_rejects_misaligned_sample_lengths(self): with pytest.raises(ValueError, match="Misaligned calibration records"): @@ -76,6 +75,18 @@ def test_merge_rejects_misaligned_sample_lengths(self): [[_record("prefill", 100, [1], [0])], [_record("prefill", 200, [1], [0])]] ) + def test_merge_rejects_threshold_width_mismatch(self): + with pytest.raises(ValueError, match="threshold-vector widths"): + merge_count_records( + [[_record("prefill", 100, [1, 2], [0, 1])], [_record("prefill", 100, [1], [0])]] + ) + + def test_merge_phase_counts_rejects_rank_phase_mismatch(self): + rank0 = {"prefill": [_record("prefill", 64, [5], [1])]} + rank1 = {"prefill": []} + with pytest.raises(ValueError, match="recorded no 'prefill' samples"): + merge_phase_counts([rank0, rank1]) + def test_merge_phase_counts_across_ranks(self): rank0 = {"prefill": [_record("prefill", 64, [5], [1])], "decode": []} rank1 = {"prefill": [_record("prefill", 64, [5], [2])]} @@ -187,6 +198,22 @@ def test_round_trips_through_serving_loader(self): assert layer_cfg["threshold_scale_factor"]["decode"] == {"a": 0.12, "b": 9.8} assert layer_cfg["target_sparse_ratio"] == {"prefill": 0.5, "decode": 0.3} + def test_preserves_legacy_toplevel_sparse_softmax(self): + existing = { + "config_groups": { + "group_0": {"algorithm": "sparse_softmax", "sparsity_n": 2, "sparsity_m": 4} + }, + "sparse_softmax": {"sparsity_n": 1, "sparsity_m": 4, "dense_recent_tokens": 128}, + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + # The serving loader reads the legacy top-level key ahead of group params. + assert config["sparse_softmax"] == existing["sparse_softmax"] + loaded = load_from_checkpoint_metadata(SimpleNamespace(sparse_attention_config=config)) + assert loaded is not None + layer_cfg = loaded[0]["sparse_cfg"]["*attn*"] + assert layer_cfg["sparsity_n"] == 1 + assert layer_cfg["dense_recent_tokens"] == 128 + def test_round_trip_with_preserved_nm_group_activates_both(self): existing = { "config_groups": { From e5c080a8824c1fbda516497f63ca812e9b99446d Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Sun, 19 Jul 2026 15:42:38 -0700 Subject: [PATCH 08/14] Address review round 2: close remaining skip+quant routes and edge cases - Skip-softmax + attention quantization is now unreachable from every direction: sparse-only installs reject calibrated skip onto layers with active attention quantizers (not just quantized installs adding skip), and the raw kernel API itself rejects P/V QDQ with an active skip threshold -- which makes the P-QDQ 16-row measurement tile dead code, so the fixed skip tile is unconditionally 128x128. - FlashInfer layout helper preserves a genuine None from get_kv_cache_layout (str(None) became the truthy string None and both guards hard-rejected valid configurations instead of using the shape fallback). - Counter vectors are validated against len(threshold_trials) in both calibrate_from_stats and fit_from_counts: consistently short vectors previously zipped silently and could misattribute sparsities. - Driver decode semantics match vLLM: the first output token comes from the prefill forward, so --decode_tokens now means decode-attention steps and generation runs decode_tokens + 1 output tokens. - Calibrate kernel mirrors the serving kernel IEEE fp32 QK dot so raw fp32 calibration and serving round near-threshold scores identically. - target_sparsity validated to [0, 1] (same range as the HF config). Signed-off-by: Kai Xu --- examples/vllm_serve/calibrate_sparse_attn.py | 18 +++++---- .../kernels/common/attention/triton_fa.py | 15 ++++++- .../kernels/sparsity/attention/calibrate.py | 10 ++++- .../calibration/calibrator.py | 6 +++ .../plugins/sparse_attn_calibration.py | 18 ++++++++- .../attention_sparsity/plugins/vllm.py | 5 ++- .../plugins/vllm_runtime.py | 40 ++++++++++++------- .../attention/test_paged_calibrate.py | 34 +++++++--------- .../test_vllm_calibration.py | 32 +++++++++++++++ .../test_sparse_attn_calibration.py | 15 +++++++ 10 files changed, 147 insertions(+), 46 deletions(-) diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py index bc5ed8281c9..28f6a7eed9e 100644 --- a/examples/vllm_serve/calibrate_sparse_attn.py +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -161,7 +161,9 @@ def main(): "--decode_tokens", type=int, default=32, - help="Decode tokens to generate per prompt (drives decode-phase calibration)", + help="Decode attention steps per prompt (drives decode-phase calibration). " + "Generation runs decode_tokens + 1 output tokens: the first output token " + "comes from the prefill forward and performs no decode attention.", ) parser.add_argument( "--max_model_len", type=int, default=None, help="vLLM max_model_len override" @@ -256,12 +258,14 @@ def main(): print(f"[ModelOpt] Calibration enabled on {n_layers} attention layers") print(f"[ModelOpt] Active sparse impls: {status['impl_types']}") - # generate() drives prefill (prefill-phase stats) then decode_tokens decode - # steps (decode-phase stats). No sparsification is applied during - # calibration — the kernel computes full dense attention while recording - # tile-skip counts. ignore_eos forces the full decode length so early EOS - # cannot thin the decode-phase statistics. - sampling = SamplingParams(temperature=0.0, max_tokens=args.decode_tokens, ignore_eos=True) + # generate() drives prefill (prefill-phase stats) then decode steps + # (decode-phase stats). No sparsification is applied during calibration — + # the kernel computes full dense attention while recording tile-skip + # counts. ignore_eos forces the full decode length so early EOS cannot + # thin the decode-phase statistics. max_tokens is decode_tokens + 1: the + # first output token comes from the prefill forward, so decode_tokens + # decode-attention steps need one extra output token. + sampling = SamplingParams(temperature=0.0, max_tokens=args.decode_tokens + 1, ignore_eos=True) llm.generate(prompts, sampling) # Aggregate RAW counts from every TP rank (each rank only measures its diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 145de3a51e2..4a223d4c2ab 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -87,7 +87,6 @@ def _load_qdq_helpers() -> None: ] _MEASURE_BLOCK_M = 128 -_P_QDQ_MEASURE_BLOCK_M = 16 # 128 so the kernel sparsity-measurement block matches the PyTorch # calibration/reference granularity. This is deliberately independent of the # autotuned compute tile. @@ -944,6 +943,16 @@ def forward( else: apply_skip = False skip_threshold_log2 = 0.0 + if apply_skip and (p_qdq_mode or v_qdq_mode): + # Quantized operands change what the calibrated skip thresholds mean, + # and P-QDQ additionally uses a different measurement tile geometry. + # The vLLM installers reject this composition at plan time; the raw + # kernel API rejects it here so no path can serve it. + raise ValueError( + "skip-softmax cannot be combined with attention quantization " + "(P/V QDQ): the calibrated tile-skip contract does not hold " + "under quantized operands" + ) o = torch.empty_like(q) lse = torch.empty(q.shape[0], num_q_heads, device=q.device, dtype=torch.float32) @@ -1046,10 +1055,12 @@ def grid(META): # GPUs) are rejected rather than re-tiled, because a different # tile realizes a different sparsity than was calibrated. try: + # P/V QDQ is rejected above when skip is active, so the tile + # here is unconditionally the 128x128 calibration geometry. _attn_fwd.fn[grid]( *fwd_args, **fwd_kwargs, - BLOCK_M=_P_QDQ_MEASURE_BLOCK_M if p_qdq_mode else _MEASURE_BLOCK_M, + BLOCK_M=_MEASURE_BLOCK_M, BLOCK_N=_MEASURE_BLOCK_N, num_warps=_MEASURE_NUM_WARPS, num_stages=_MEASURE_NUM_STAGES, diff --git a/modelopt/torch/kernels/sparsity/attention/calibrate.py b/modelopt/torch/kernels/sparsity/attention/calibrate.py index 68748e7c50c..b8e9d6475c7 100644 --- a/modelopt/torch/kernels/sparsity/attention/calibrate.py +++ b/modelopt/torch/kernels/sparsity/attention/calibrate.py @@ -69,6 +69,7 @@ def _attn_fwd_calibrate( HEAD_DIM: tl.constexpr, NUM_THRESHOLDS: tl.constexpr, PADDED_THRESHOLDS: tl.constexpr, # next_power_of_2(NUM_THRESHOLDS) for tl.arange + Q_IS_FP32: tl.constexpr = False, # match the serving kernel's IEEE fp32 QK dot IS_PAGED: tl.constexpr = False, # Whether K/V are read from a paged KV cache K_cache=None, # [num_blocks, page_size, num_kv_heads, head_dim] paged K V_cache=None, # [num_blocks, page_size, num_kv_heads, head_dim] paged V @@ -171,7 +172,13 @@ def _attn_fwd_calibrate( other=0.0, ) - scores = tl.dot(q, k) * qk_scale + # Match the serving kernel's QK precision: fp32 Q uses the IEEE dot + # (default tl.dot is TF32 for fp32 inputs), so near-threshold scores + # round to the same skip decisions in calibration and serving. + if Q_IS_FP32: + scores = tl.dot(q, k.to(tl.float32), input_precision="ieee") * qk_scale + else: + scores = tl.dot(q, k) * qk_scale scores = _apply_mask(scores, q_pos, kv_pos, seq_len_q, seq_len_kv, kv_start, IS_CAUSAL) tile_row_max = tl.max(scores, 1) @@ -430,6 +437,7 @@ def attention_calibrate( HEAD_DIM=HEAD_DIM, NUM_THRESHOLDS=num_thresholds, PADDED_THRESHOLDS=triton.next_power_of_2(num_thresholds), + Q_IS_FP32=q.dtype == torch.float32, IS_PAGED=is_paged, K_cache=k_cache, V_cache=v_cache, diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py index 5e604c039ca..ef46758ead0 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py @@ -169,6 +169,12 @@ def calibrate_from_stats(self, per_sample_stats: list[dict], phase: str) -> dict for sample_stat in per_sample_stats: length = sample_stat["sample_length"] sparsity_list = sample_stat["sparsity"] + if len(sparsity_list) != len(self.threshold_trials): + # A silent zip would misattribute sparsities to thresholds. + raise ValueError( + f"per-sample sparsity has {len(sparsity_list)} entries but " + f"{len(self.threshold_trials)} threshold trials are configured" + ) for threshold, sparsity in zip(self.threshold_trials, sparsity_list): scale_factor = threshold * length all_data_points.append( diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py index b66a5c27436..32fb5d67160 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -178,6 +178,12 @@ def fit_from_counts( for phase, records in per_phase_counts.items(): if not records: continue + for record in records: + if len(record["total_tiles"]) != len(threshold_trials): + raise ValueError( + f"{phase} record has {len(record['total_tiles'])} counters but " + f"{len(threshold_trials)} threshold trials are configured" + ) calibrator = DynamicThresholdCalibrator( threshold_trials=list(threshold_trials), fit_logspace=fit_logspace ) @@ -193,8 +199,16 @@ def fit_from_counts( def _normalize_target_sparsity(target_sparsity: dict[str, float] | float) -> dict[str, float]: if isinstance(target_sparsity, int | float): - return {phase: float(target_sparsity) for phase in _PHASES} - return {phase: float(target_sparsity.get(phase, 0.5)) for phase in _PHASES} + values = {phase: float(target_sparsity) for phase in _PHASES} + else: + values = {phase: float(target_sparsity.get(phase, 0.5)) for phase in _PHASES} + for phase, value in values.items(): + # Same range the HF calibration config enforces. + if not 0.0 <= value <= 1.0: + raise ValueError( + f"target_sparsity for phase {phase!r} must be between 0.0 and 1.0, got {value}" + ) + return values def build_sparse_attention_config( diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 2ebf2aa39ac..8bede9ed665 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -281,9 +281,12 @@ def _flashinfer_kv_cache_layout() -> str | None: if getter is None: continue try: - return str(getter()) + value = getter() except Exception: return None + # Preserve a genuine None (layout unset) so the shape fallback runs; + # str(None) would become the truthy string "None" and hard-reject. + return None if value is None else str(value) return None diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index a00a7ca012e..724feb69f90 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -361,16 +361,24 @@ def _plan_vllm_attention( plans = [] for name, module, sparse_kw in candidates: reasons = _layer_errors(module) - if quantize and _skip_softmax_active(sparse_kw): + if _skip_softmax_active(sparse_kw): # Quantized Q/K/P change the attention-score distribution the skip # thresholds were calibrated on, so the calibrated sparsity contract - # no longer holds. N:M sparse softmax has no calibrated threshold - # and composes with quantization. - reasons.append( - "skip-softmax cannot be combined with attention quantization; " - "serve skip-softmax unquantized or drop the skip_softmax group " - "(N:M sparse softmax composes with quantization)" + # no longer holds. This guards both installation directions: quantized + # installs adding skip, and sparse-only installs onto layers that + # already carry active attention quantizers. N:M sparse softmax has + # no calibrated threshold and composes with quantization. + active = ( + "attention quantization is being installed" + if quantize + else _active_attention_quantization(module) ) + if active: + reasons.append( + f"skip-softmax cannot be combined with attention quantization ({active}); " + "serve skip-softmax unquantized or drop the skip_softmax group " + "(N:M sparse softmax composes with quantization)" + ) device = dtype = None if quantize: device, dtype = quant_plugin._get_device_dtype(module) @@ -513,18 +521,22 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor return _build_report(plan) -def _attention_quant_error(module) -> str | None: - """Reject calibration on layers with any active attention Q/K/P/V fakequant.""" +def _active_attention_quantization(module) -> str | None: + """Describe any active attention Q/K/P/V quantization on a layer, or None.""" for attr in ("q_bmm_quantizer", "k_bmm_quantizer", "p_bmm_quantizer", "v_bmm_quantizer"): if getattr(getattr(module, attr, None), "is_enabled", False): - return f"{attr} is enabled; skip-softmax calibration requires unquantized attention" + return f"{attr} is enabled" if getattr(module, "_query_quant_in_kernel", False) or getattr( module, "_value_quant_in_kernel", False ): - return ( - "in-kernel attention quantization flags are set; skip-softmax " - "calibration requires unquantized attention" - ) + return "in-kernel attention quantization flags are set" + return None + + +def _attention_quant_error(module) -> str | None: + """Reject calibration on layers with any active attention Q/K/P/V fakequant.""" + if active := _active_attention_quantization(module): + return f"{active}; skip-softmax calibration requires unquantized attention" return None diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py index 7020f90d7c5..c7390382098 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py @@ -211,26 +211,22 @@ def test_serve_skip_counts_equal_calibrate_counts(self, threshold): assert out._sparsity_total == calib_total assert out._sparsity_skipped == calib_skipped - def test_skip_composes_with_pv_qdq(self): - """Active skip at the fixed tile compiles and runs with P/V QDQ enabled.""" - if torch.cuda.get_device_capability() < (8, 9): - pytest.skip("NVFP4/FP8 QDQ needs compute capability >= 8.9 (E4M3 casts)") + @pytest.mark.parametrize("qdq_kw", [{"p_qdq": "nvfp4"}, {"v_qdq": "nvfp4", "v_qdq_amax": 1.0}]) + def test_skip_rejects_pv_qdq(self, qdq_kw): + """Active skip rejects P/V QDQ: quantized operands break the calibrated contract.""" seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 q, k, v = self._contrasty_qkv(seq_len, num_heads, num_kv_heads, head_dim) locs, lens = make_varlen_meta([seq_len]) - out = attention( - q, - k, - v, - locs, - lens, - seq_len, - is_causal=True, - skip_softmax_threshold=1e-2, - p_qdq="nvfp4", - p_qdq_amax=1.0, - v_qdq="nvfp4", - v_qdq_amax=float(v.abs().amax()), - ) - assert torch.isfinite(out).all() + with pytest.raises(ValueError, match="cannot be combined with attention quantization"): + attention( + q, + k, + v, + locs, + lens, + seq_len, + is_causal=True, + skip_softmax_threshold=1e-2, + **qdq_kw, + ) diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py index 62a75b92d8f..543fbd0a2da 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -155,6 +155,27 @@ def test_quantized_install_rejects_calibrated_skip(self): vllm_runtime.install_vllm_nvfp4_attention(runner) assert not isinstance(attention.impl, ModelOptSparseAttentionImpl) + def test_sparse_only_install_rejects_skip_onto_quantized_layer(self): + """Sparse-only installs must also refuse skip onto live quantizers.""" + attention = _bare_attention() + attention.q_bmm_quantizer = SimpleNamespace(is_enabled=True) + original_impl = attention.impl + runner = _model_runner( + nn.ModuleDict({"attn": attention}), sparse_metadata=self._CALIBRATED_META + ) + with pytest.raises( + NotImplementedError, match="cannot be combined with attention quantization" + ): + vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert attention.impl is original_impl + + def test_sparse_only_install_allows_skip_on_unquantized_layer(self): + runner = _model_runner( + nn.ModuleDict({"attn": _bare_attention()}), sparse_metadata=self._CALIBRATED_META + ) + report = vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) + assert report.installed_count == 1 + def test_quantized_plan_allows_nm_sparsity(self, monkeypatch): from modelopt.torch.quantization.plugins import vllm as quant_plugin @@ -174,6 +195,17 @@ def test_quantized_plan_allows_nm_sparsity(self, monkeypatch): class TestFlashInferLayoutGuard: """HND FlashInfer caches are rejected via layout metadata, pre-measurement.""" + def test_layout_helper_preserves_none(self, monkeypatch): + """A getter returning None must not become the truthy string 'None'.""" + import sys + + fake = SimpleNamespace(get_kv_cache_layout=lambda: None) + monkeypatch.setitem(sys.modules, "vllm.v1.attention.backends.utils", fake) + assert attention_plugin._flashinfer_kv_cache_layout() is None + + fake.get_kv_cache_layout = lambda: "HND" + assert attention_plugin._flashinfer_kv_cache_layout() == "HND" + def test_installer_rejects_hnd_layout(self, monkeypatch): monkeypatch.setattr(attention_plugin, "_flashinfer_kv_cache_layout", lambda: "HND") attention = _bare_attention(FlashInferImpl) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py index 85e459dd650..e85ad3ab77e 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py @@ -127,6 +127,15 @@ def counts(length, total): def test_empty_phase_produces_no_fit(self): assert fit_from_counts({"decode": []}, DEFAULT_THRESHOLD_TRIALS) == {} + def test_fit_rejects_counter_width_vs_trials_mismatch(self): + """Consistent-but-wrong widths must not silently zip against the trials.""" + short = len(DEFAULT_THRESHOLD_TRIALS) - 1 + records = [ + {"sample_length": 4096, "total_tiles": [100] * short, "skipped_tiles": [50] * short} + ] + with pytest.raises(ValueError, match="threshold trials are configured"): + fit_from_counts({"prefill": records}, DEFAULT_THRESHOLD_TRIALS) + class TestCalibrateFromStats: def _stats(self, trials): @@ -198,6 +207,12 @@ def test_round_trips_through_serving_loader(self): assert layer_cfg["threshold_scale_factor"]["decode"] == {"a": 0.12, "b": 9.8} assert layer_cfg["target_sparse_ratio"] == {"prefill": 0.5, "decode": 0.3} + def test_rejects_out_of_range_target_sparsity(self): + with pytest.raises(ValueError, match=r"between 0\.0 and 1\.0"): + build_sparse_attention_config(self._PARAMS, 1.5) + with pytest.raises(ValueError, match=r"between 0\.0 and 1\.0"): + build_sparse_attention_config(self._PARAMS, {"prefill": 0.5, "decode": -0.1}) + def test_preserves_legacy_toplevel_sparse_softmax(self): existing = { "config_groups": { From 9e029d6bbd8f2f0fa97ebfde26343ab56fadfa9b Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Mon, 24 Aug 2026 21:12:38 -0700 Subject: [PATCH 09/14] Address skip-softmax calibration review findings Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 +- examples/vllm_serve/README.md | 3 +- examples/vllm_serve/calibrate_sparse_attn.py | 33 ++++- .../kernels/sparsity/attention/calibrate.py | 38 +++-- modelopt/torch/quantization/tensor_quant.py | 5 +- .../sparsity/attention_sparsity/conversion.py | 39 +++-- .../plugins/sparse_attn_calibration.py | 67 +++++---- .../attention_sparsity/plugins/vllm.py | 135 +++++++++++++----- .../plugins/vllm_runtime.py | 33 ++++- .../common/attention/test_triton_fa_p_qdq.py | 36 ++--- .../quantization/test_tensor_quant_cuda.py | 33 +++++ .../test_vllm_calibration.py | 6 +- .../attention_sparsity/test_vllm_runtime.py | 3 - .../common/attention/test_triton_fa.py | 52 ++++--- .../test_sparse_attn_calibration.py | 25 ++++ 15 files changed, 372 insertions(+), 138 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5df67102c48..8897061c086 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog *Sparsity* -- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. Configurations that cannot compile the 128x128 tile are rejected rather than re-tiled, and skip-softmax cannot be combined with attention quantization. See ``examples/vllm_serve/calibrate_sparse_attn.py``. +- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. Configurations that cannot compile the 128x128 tile are rejected rather than re-tiled, and skip-softmax cannot be combined with attention quantization. **Behavior change for sparse-only serving installs:** ``install_vllm_sparse_attention_from_checkpoint`` now runs the same engine-level compatibility validation as quantized installs (decode context parallelism, DBO, speculative decoding, and FULL mixed-batch CUDA graphs are rejected; prefix caching remains supported), and checkpoints with a calibrated ``decode`` ``threshold_scale_factor`` are rejected under a FULL decode CUDA graph mode — including vLLM's default ``FULL_AND_PIECEWISE`` — because the captured graph would replay one request's stale threshold. Previously such deployments started and silently served miscalibrated sparsity; pass ``--enforce-eager`` (or select a non-FULL decode graph mode). See ``examples/vllm_serve/calibrate_sparse_attn.py``. *Quantization* diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 34a802c85e2..c5596003cd3 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -213,7 +213,8 @@ report = install_vllm_nvfp4_attention(model_runner, sparse_cfg="checkpoint") Limitations: - vLLM V1 chunked prefill and prefix-cache suffix attention are supported by offsetting query positions into the longer KV span. This applies to sparse-only serving; quantized attention installs and skip-softmax calibration reject `enable_prefix_caching` (quantize-on-write and per-request measurement both require uncached prefills). -- `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`. +- `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`. Checkpoints with a calibrated `decode` `threshold_scale_factor` are rejected at install under a FULL decode CUDA graph mode (including vLLM's default `FULL_AND_PIECEWISE`): the captured graph would replay one request's stale threshold. +- Sparse-only installs validate engine-level compatibility like quantized installs do: decode context parallelism, DBO, speculative decoding, and FULL mixed-batch CUDA graphs are rejected (prefix caching remains supported, per the bullet above). ### Compact NVFP4 attention worker diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py index 28f6a7eed9e..42c4d36a972 100644 --- a/examples/vllm_serve/calibrate_sparse_attn.py +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -50,6 +50,7 @@ import sys from pathlib import Path +from modelopt.torch.sparsity.attention_sparsity.calibration.ruler_dataset import RulerDatasetBuilder from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_calibration import ( DEFAULT_THRESHOLD_TRIALS, build_sparse_attention_config, @@ -71,10 +72,6 @@ def _load_prompts(llm, args) -> list[str]: # Same dataset as the HF calibration path (calibration/calibrate.py), so the # vLLM- and PyTorch-calibrated thresholds are fit on identical data. - from modelopt.torch.sparsity.attention_sparsity.calibration.ruler_dataset import ( - RulerDatasetBuilder, - ) - builder = RulerDatasetBuilder( samples=args.calib_samples, max_seqlen=args.calib_max_seqlen, @@ -83,6 +80,12 @@ def _load_prompts(llm, args) -> list[str]: data_dir=args.calib_data_dir, ) samples = builder.build_calibration_dataset() + if not samples: + raise ValueError( + "RULER produced no calibration samples (all candidates exceeded " + f"max_length_filter={int(args.calib_max_seqlen * 1.5)} tokens). " + "Adjust --calib_max_seqlen / --calib_samples, or pass --prompts_file." + ) prompts = [sample["input"] for sample in samples] lengths = sorted(sample["length"] for sample in samples) print( @@ -117,7 +120,11 @@ def _write_config(ckpt: str, sparse_config: dict, update_checkpoint: bool) -> No config_json = Path(ckpt) / "config.json" config = json.loads(config_json.read_text()) config["sparse_attention_config"] = sparse_config - config_json.write_text(json.dumps(config, indent=2)) + # Atomic replace: a crash mid-write must not truncate the checkpoint's + # config.json (write_text would rewrite it in place). + tmp_path = config_json.with_name(config_json.name + ".tmp") + tmp_path.write_text(json.dumps(config, indent=2)) + os.replace(tmp_path, config_json) print(f"[ModelOpt] Merged sparse_attention_config into {config_json}") @@ -211,6 +218,14 @@ def main(): ) args = parser.parse_args() + if args.update_checkpoint_config and not (Path(args.model) / "config.json").is_file(): + # Fail before the (expensive, multi-GPU) calibration run, not after: + # merging requires a local checkpoint directory, not a HF hub ID. + parser.error( + f"--update_checkpoint_config requires a local checkpoint directory " + f"containing config.json; {args.model!r} has none" + ) + # Workers run in separate processes and must import the calibration worker. repo_root = str(Path(__file__).resolve().parent) if repo_root not in sys.path: @@ -218,6 +233,8 @@ def main(): current = os.environ.get("PYTHONPATH") os.environ["PYTHONPATH"] = os.pathsep.join([current, repo_root]) if current else repo_root + # Deferred heavy import: keep argparse/--help (and arg errors) fast, and + # only import vLLM after the PYTHONPATH setup above. from vllm import LLM, SamplingParams llm_kwargs = { @@ -284,6 +301,12 @@ def main(): "so observed sparsity spans the (10%, 90%) fitting window." ) sys.exit(1) + # Export only requested phases: a stray record (e.g. a scheduling corner + # case classified into an unrequested phase) must not bake an + # uncalibrated-by-intent phase into the config. + calibration_params = { + phase: params for phase, params in calibration_params.items() if phase in requested_phases + } sparse_config = build_sparse_attention_config( calibration_params, diff --git a/modelopt/torch/kernels/sparsity/attention/calibrate.py b/modelopt/torch/kernels/sparsity/attention/calibrate.py index b8e9d6475c7..3655c4cbfcd 100644 --- a/modelopt/torch/kernels/sparsity/attention/calibrate.py +++ b/modelopt/torch/kernels/sparsity/attention/calibrate.py @@ -22,6 +22,7 @@ ``modelopt.torch.sparsity.attention_sparsity`` to fit a skip threshold. """ +import functools import math import torch @@ -263,6 +264,22 @@ def _attn_fwd_calibrate( tl.store(Out + o_ptrs, acc, mask=(q_pos[:, None] < seq_len_q) & d_mask[None, :]) +@functools.lru_cache(maxsize=64) +def _log2_threshold_tensor( + threshold_trials: tuple[float, ...], device: torch.device +) -> torch.Tensor: + """Build the log2-space threshold tensor, cached per (trials, device). + + Scores already include sm_scale and LOG2E; convert lambda to log2 space + only. Trials are constant for a whole calibration run, and the vLLM path + calls :func:`attention_calibrate` once per request per layer per step, so + rebuilding (and re-uploading) the tensor per call would be pure waste. + """ + return torch.tensor( + [math.log2(t) for t in threshold_trials], dtype=torch.float32, device=device + ) + + def attention_calibrate( q: torch.Tensor, k: torch.Tensor, @@ -363,19 +380,22 @@ def attention_calibrate( b_seq_len_k = b_seq_len b_start_loc_k = b_start_loc - # Paged mode: KV positions come from block_table, so the contiguous KV - # offsets are unused. Provide a dummy so Triton can compile the tl.load. if b_start_loc_k is None: - b_start_loc_k = torch.zeros_like(b_start_loc) + if not is_paged: + # A zeros dummy here would silently read every sequence's K/V from + # offset 0 — fail loudly instead (contiguous K/V needs real offsets). + raise ValueError( + "b_start_loc_k is required when b_seq_len_k is provided for " + "contiguous (non-paged) K/V" + ) + # Paged mode: KV positions come from block_table, so the contiguous KV + # offsets are unused. Alias b_start_loc (same shape/dtype/device) so + # Triton can compile the tl.load without allocating a dummy per call. + b_start_loc_k = b_start_loc num_thresholds = len(threshold_trials) - # Scores already include sm_scale and LOG2E; convert lambda to log2 space only. - threshold_tensor = torch.tensor( - [math.log2(t) for t in threshold_trials], - dtype=torch.float32, - device=q.device, - ) + threshold_tensor = _log2_threshold_tensor(tuple(threshold_trials), q.device) o = torch.empty_like(q) diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..9e5b6c186ec 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -180,7 +180,10 @@ def _dynamic_block_quantize_impl( and not DISABLE_TRITON_KERNEL and amax is not None ): - return triton_kernel.fp4_fake_quant_block(inputs, amax) + # Forward the configured block size: the kernel defaults to 16, and + # silently ignoring a non-16 NVFP4 block config here would diverge + # from the cuda_ext fallback below (which honors block_size). + return triton_kernel.fp4_fake_quant_block(inputs, amax, block_size=block_size) cuda_ext_mx = get_cuda_ext_mx(raise_if_failed=True) return cuda_ext_mx.fused_amax_convert( inputs, diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index 8c41b895a8a..b2acd137930 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py @@ -39,6 +39,29 @@ ) +def export_threshold_scale_factor(calibration_params: dict[str, Any]) -> dict[str, Any]: + """Build the canonical per-phase ``threshold_scale_factor`` export block. + + Single source of the exported skip-softmax schema fragment, shared by the + HF exporter (:func:`export_sparse_attention_config`) and the vLLM + calibration path (``plugins.sparse_attn_calibration``), so the serving + loader sees one format regardless of which path produced the checkpoint. + """ + block: dict[str, Any] = {"formula": "a * exp(b * target_sparsity)"} + for phase in ("prefill", "decode"): + if phase in calibration_params: + block[phase] = { + "a": float(calibration_params[phase]["a"]), + "b": float(calibration_params[phase]["b"]), + } + return block + + +def export_config_producer() -> dict[str, str]: + """Build the canonical ``producer`` block of ``sparse_attention_config``.""" + return {"name": "modelopt", "version": mo_version} + + def _set_attn_implementation(model: nn.Module, config: SparseAttentionConfig) -> None: """Set the correct attn_implementation based on the sparse attention method/backend. @@ -469,16 +492,7 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: skip_group["initial_disabled_steps"] = initial_disabled_steps # threshold_scale_factor (a * exp(b * target_sparsity)) and target_sparsity are # skip-softmax-specific, so they live in this group. - threshold_scale_factor: dict[str, Any] = { - "formula": "a * exp(b * target_sparsity)", - } - for phase in ["prefill", "decode"]: - if phase in calibration_params: - threshold_scale_factor[phase] = { - "a": calibration_params[phase]["a"], - "b": calibration_params[phase]["b"], - } - skip_group["threshold_scale_factor"] = threshold_scale_factor + skip_group["threshold_scale_factor"] = export_threshold_scale_factor(calibration_params) if target_sparse_ratio is not None: skip_group["target_sparsity"] = target_sparse_ratio config_groups[f"group_{group_idx}"] = skip_group @@ -495,10 +509,7 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: return { "config_groups": config_groups, - "producer": { - "name": "modelopt", - "version": mo_version, - }, + "producer": export_config_producer(), } diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py index 32fb5d67160..e0bd2250d87 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -33,11 +33,13 @@ from typing import Any -import modelopt - # One canonical sweep for both calibration paths: re-exported from the HF-path # calibrator so the vLLM path fits on the identical trial grid. -from ..calibration.calibrator import DEFAULT_THRESHOLD_TRIALS +from ..calibration.calibrator import DEFAULT_THRESHOLD_TRIALS, DynamicThresholdCalibrator + +# One canonical schema for both calibration paths: the exported skip-softmax +# blocks come from the HF exporter's helpers so the formats cannot drift. +from ..conversion import export_config_producer, export_threshold_scale_factor __all__ = [ "DEFAULT_THRESHOLD_TRIALS", @@ -115,7 +117,9 @@ def merge_count_records(sources: list[list[dict]]) -> list[dict]: return merged -def merge_phase_counts(rank_counts: list[dict[str, list[dict]]]) -> dict[str, list[dict]]: +def merge_phase_counts( + rank_counts: list[dict[str, list[dict]]], *, source_desc: str = "rank" +) -> dict[str, list[dict]]: """Merge per-phase raw-count records collected from every TP rank. ``rank_counts`` is the list of per-rank results (one @@ -123,7 +127,9 @@ def merge_phase_counts(rank_counts: list[dict[str, list[dict]]]) -> dict[str, li ``collect_calibration_counts``). Use ALL ranks: with tensor parallelism each rank only measures its attention-head shard, so any single rank's counts are incomplete. A phase recorded by some ranks but not others - indicates a collection bug and raises. + indicates a collection bug and raises. The same merge also sums one + rank's per-layer splits (``collect_calibration_counts`` delegates here + with ``source_desc="attention layer"``). """ phases = {phase for rank in rank_counts for phase in rank} merged: dict[str, list[dict]] = {} @@ -132,10 +138,11 @@ def merge_phase_counts(rank_counts: list[dict[str, list[dict]]]) -> dict[str, li empty = sum(1 for source in sources if not source) if empty and empty != len(sources): raise ValueError( - f"Misaligned calibration records: {empty}/{len(sources)} rank(s) " + f"Misaligned calibration records: {empty}/{len(sources)} {source_desc}(s) " f"recorded no {phase!r} samples while others did" ) - merged[phase] = merge_count_records([] if empty else sources) + # All-empty sources merge to [] (merge_count_records sums zero samples). + merged[phase] = merge_count_records(sources) return merged @@ -172,8 +179,6 @@ def fit_from_counts( ``{phase: {"a", "b", "min_observed_sparsity", "max_observed_sparsity"}}`` for each phase that produced a valid fit. """ - from ..calibration.calibrator import DynamicThresholdCalibrator - calibration_params: dict[str, dict[str, float]] = {} for phase, records in per_phase_counts.items(): if not records: @@ -228,37 +233,49 @@ def build_sparse_attention_config( Non-skip groups from ``existing_config`` (e.g. exported N:M ``sparse_softmax`` metadata) are preserved after the skip group; an - existing ``skip_softmax`` group is replaced by the new calibration. + existing ``skip_softmax`` group is replaced by the new calibration, + carrying over its layer policy (``ignore`` — layers deliberately kept + dense — and ``initial_disabled_steps``): recalibration replaces the + fitted thresholds, not which layers the export sparsifies. """ - threshold_scale_factor: dict[str, Any] = {"formula": "a * exp(b * target_sparsity)"} - for phase in _PHASES: - if phase in calibration_params: - threshold_scale_factor[phase] = { - "a": float(calibration_params[phase]["a"]), - "b": float(calibration_params[phase]["b"]), - } + # target_sparsity covers only fitted phases: claiming a target for a phase + # without calibrated (a, b) would advertise sparsity the serving path + # silently serves dense (it needs the per-phase scale factors). + target_sparsity_by_phase = { + phase: value + for phase, value in _normalize_target_sparsity(target_sparsity).items() + if phase in calibration_params + } skip_group: dict[str, Any] = { "algorithm": "skip_softmax", "targets": ["Attention"], - "threshold_scale_factor": threshold_scale_factor, - "target_sparsity": _normalize_target_sparsity(target_sparsity), + "threshold_scale_factor": export_threshold_scale_factor(calibration_params), + "target_sparsity": target_sparsity_by_phase, } config_groups: dict[str, Any] = {"group_0": skip_group} existing_groups = (existing_config or {}).get("config_groups") if isinstance(existing_groups, dict): - preserved = [ - group - for group in existing_groups.values() - if isinstance(group, dict) and group.get("algorithm") != "skip_softmax" - ] + preserved = [] + for group in existing_groups.values(): + if not isinstance(group, dict): + continue + if group.get("algorithm") == "skip_softmax": + # Keep the replaced group's layer policy: dropping ``ignore`` + # would sparsify layers the original export deliberately kept + # dense (e.g. first/last blocks). + for key in ("ignore", "initial_disabled_steps"): + if key in group and key not in skip_group: + skip_group[key] = group[key] + else: + preserved.append(group) for idx, group in enumerate(preserved, start=1): config_groups[f"group_{idx}"] = group result: dict[str, Any] = { "config_groups": config_groups, - "producer": {"name": "modelopt", "version": modelopt.__version__}, + "producer": export_config_producer(), } # Legacy checkpoints carry N:M parameters as a top-level ``sparse_softmax`` # dict (read by the serving loader ahead of group params) — preserve it so diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 8bede9ed665..a6d3abffa4e 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -48,6 +48,17 @@ from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite from modelopt.torch.kernels.sparsity.attention.calibrate import attention_calibrate +from .sparse_attn_calibration import merge_phase_counts, split_records_by_phase + +__all__ = [ + "ModelOptSparseAttentionBackend", + "ModelOptSparseAttentionImpl", + "collect_calibration_counts", + "disable_calibration", + "enable_calibration", + "iter_sparse_impls", +] + @functools.cache def _flash_attention_kv_cache_layout() -> str: @@ -311,10 +322,17 @@ def _forward_calibrate( sparsification is applied to generation (the dense Triton kernel's numerics may differ slightly from the native backend's). - Phase and causality are decided per request: ``q_len == 1`` is a decode - step (full-cache, non-causal); ``q_len > 1`` is (chunked) prefill (causal - — the kernel offsets the query into the KV span). A mixed prefill/decode - batch therefore contributes correctly to both phase fits. + Phase and causality are decided per request: ``q_len > 1`` is (chunked) + prefill (causal — the kernel offsets the query into the KV span). A + ``q_len == 1`` row is a decode step (full-cache, non-causal) only when + its KV span exceeds the request's prompt (at least one generated token); + a 1-token row still inside the prompt is the final chunk of a chunked + prefill and is recorded as prefill. Prompt lengths come from the runner's + input batch (the installer attaches the runner as ``_calib_model_runner``; + the input batch is resolved per forward because vLLM can rebuild it after + install, same request order as the metadata rows); without it, + ``q_len == 1`` falls back to decode. A mixed prefill/decode batch + therefore contributes correctly to both phase fits. Records raw per-threshold tile counts (not ratios) on ``impl._calib_records`` so tensor-parallel workers can be aggregated by @@ -335,8 +353,20 @@ def _forward_calibrate( page_size = key_cache.shape[1] trials = impl._calib_threshold_trials batch = seq_lens.shape[0] - b_start_loc = cu_seqlens_q[:batch] - b_seq_len = cu_seqlens_q[1 : batch + 1] - cu_seqlens_q[:batch] + # Hoist per-request tensors out of the loop: kernel args are sliced views + # of these, so the loop performs no allocations or casts. + b_seq_len_i32 = (cu_seqlens_q[1 : batch + 1] - cu_seqlens_q[:batch]).to(torch.int32) + seq_lens_i32 = seq_lens[:batch].to(torch.int32) + b_start_loc_zero = torch.zeros(1, device=query.device, dtype=torch.int32) + # vLLM v1 metadata tensors are GPU-resident: take one host copy per launch + # instead of stalling the stream with per-request ``.item()`` syncs. + cu_seqlens_q_cpu = cu_seqlens_q[: batch + 1].cpu() + seq_lens_cpu = seq_lens[:batch].cpu() + # Per-request prompt lengths (same request order as the metadata rows) + # distinguish decode steps from 1-token final chunks of a chunked + # prefill. Resolved from the runner per forward: vLLM can replace + # input_batch after install (KV-cache init for hybrid models). + input_batch = getattr(getattr(impl, "_calib_model_runner", None), "input_batch", None) q = query[:num_actual_tokens].contiguous() # Dummy K/V: in paged mode KV is read from the cache via block_table. @@ -344,23 +374,29 @@ def _forward_calibrate( k_dummy = torch.empty(0, impl.num_kv_heads, impl.head_size, device=q.device, dtype=q.dtype) for i in range(batch): - q_len = int(b_seq_len[i].item()) + q_start = int(cu_seqlens_q_cpu[i]) + q_len = int(cu_seqlens_q_cpu[i + 1]) - q_start if q_len <= 0: continue - q_start = int(b_start_loc[i].item()) - seq_k = int(seq_lens[i].item()) - phase = "decode" if q_len <= 1 else "prefill" + seq_k = int(seq_lens_cpu[i]) + if q_len > 1: + phase = "prefill" + elif input_batch is not None and seq_k <= int(input_batch.num_prompt_tokens[i]): + # 1-token final chunk of a chunked prefill: still inside the prompt. + phase = "prefill" + else: + phase = "decode" oi, counters = attention_calibrate( q[q_start : q_start + q_len], k_dummy, k_dummy, - b_start_loc=torch.zeros(1, device=q.device, dtype=torch.int32), - b_seq_len=b_seq_len[i : i + 1].to(torch.int32), + b_start_loc=b_start_loc_zero, + b_seq_len=b_seq_len_i32[i : i + 1], max_input_len=q_len, is_causal=q_len > 1, softmax_scale=impl.scale, - b_seq_len_k=seq_lens[i : i + 1].to(torch.int32), + b_seq_len_k=seq_lens_i32[i : i + 1], max_input_len_k=seq_k, threshold_trials=trials, k_cache=key_cache, @@ -370,12 +406,14 @@ def _forward_calibrate( ) output[q_start : q_start + q_len] = oi + # One host transfer for both counter columns (counters is GPU-resident). + counters_cpu = counters.cpu() impl._calib_records.append( { "phase": phase, "sample_length": seq_k, - "total_tiles": counters[:, 0].tolist(), - "skipped_tiles": counters[:, 1].tolist(), + "total_tiles": counters_cpu[:, 0].tolist(), + "skipped_tiles": counters_cpu[:, 1].tolist(), } ) @@ -518,6 +556,8 @@ def _dispatch_modelopt( num_prefills: int, num_decode_tokens: int, num_prefill_tokens: int, + max_seq_len_decode: int | None = None, + max_seq_len_prefill: int | None = None, **common_kw, ) -> torch.Tensor: """Run the ModelOpt path, splitting mixed decode+prefill batches by phase. @@ -527,8 +567,14 @@ def _dispatch_modelopt( ``q_len==1`` decode rows with ``q_len>1`` (chunked-)prefill rows, ``max_query_len > 1`` and the whole batch would otherwise take the prefill skip-softmax path. Split so each phase runs its own schedule -- decode rows - always take the fixed decode path. Both the FlashAttention and FlashInfer - adapters share this dispatch. + always take the fixed decode path. + + Both adapters route through this dispatch, but the split is live only on + FlashInfer: its metadata carries the ``num_decodes``/``num_prefills`` + counts (and vLLM reorders those batches decode-first). vLLM's + FlashAttention metadata has no phase counts, so FA mixed batches fall + through to the whole-batch path and are classified by ``max_query_len`` + alone (decode rows then follow the prefill contract for that launch). """ if not (num_decodes and num_prefills): return _forward_modelopt( @@ -554,6 +600,18 @@ def _dispatch_modelopt( if not common_kw.get("quant_active", False): common_kw["dense_fallback"]() + # Each phase derives its skip threshold from its own KV maximum: reusing + # the batch-global max_seq_len (e.g. a co-scheduled 32k prefill next to 2k + # decodes) would shrink the decode threshold far below — much denser than + # — the calibrated target. Fall back to the batch-global value only when + # the builder did not provide per-phase maxima. + decode_kw = dict(common_kw) + if max_seq_len_decode is not None: + decode_kw["max_seq_len"] = max_seq_len_decode + prefill_kw = dict(common_kw) + if max_seq_len_prefill is not None: + prefill_kw["max_seq_len"] = max_seq_len_prefill + _forward_modelopt( impl, query=query[:num_decode_tokens], @@ -563,7 +621,7 @@ def _dispatch_modelopt( num_actual_tokens=num_decode_tokens, max_query_len=num_decode_tokens // num_decodes, output=output[:num_decode_tokens], - **common_kw, + **decode_kw, ) prefill_start = num_decode_tokens prefill_cu_seqlens_q = cu_seqlens_q[num_decodes:] - cu_seqlens_q[num_decodes] @@ -576,7 +634,7 @@ def _dispatch_modelopt( num_actual_tokens=num_prefill_tokens, max_query_len=max_query_len, output=output[prefill_start : prefill_start + num_prefill_tokens], - **common_kw, + **prefill_kw, ) return output @@ -778,6 +836,23 @@ def build(*args, **kwargs): common = build_sig.bind(*args, **kwargs).arguments["common_attn_metadata"] for target, source in _FLASHINFER_METADATA_FIELDS.items(): setattr(metadata, target, getattr(common, source)) + # Per-phase KV maxima for the mixed-batch split (batch is reordered + # decode-first): computed once per build — not per layer forward — so + # the split's threshold derivation neither reuses the batch-global max + # nor syncs the stream inside every layer. + num_decodes = getattr(metadata, "num_decodes", 0) + num_prefills = getattr(metadata, "num_prefills", 0) + max_seq_len_decode = max_seq_len_prefill = None + if num_decodes and num_prefills: + # Prefer the host-resident copy the runner may already carry; + # fall back to one device->host copy per mixed-batch build. + seq_lens_cpu = getattr(common, "_seq_lens_cpu", None) + if seq_lens_cpu is None: + seq_lens_cpu = common.seq_lens.cpu() + max_seq_len_decode = int(seq_lens_cpu[:num_decodes].max()) + max_seq_len_prefill = int(seq_lens_cpu[num_decodes : num_decodes + num_prefills].max()) + metadata._modelopt_max_seq_len_decode = max_seq_len_decode + metadata._modelopt_max_seq_len_prefill = max_seq_len_prefill return metadata setattr(build, "_modelopt_sparse_metadata_patch", True) @@ -940,6 +1015,8 @@ def prepare_modelopt(): num_prefills=getattr(attn_metadata, "num_prefills", 0), num_decode_tokens=getattr(attn_metadata, "num_decode_tokens", 0), num_prefill_tokens=getattr(attn_metadata, "num_prefill_tokens", 0), + max_seq_len_decode=getattr(attn_metadata, "_modelopt_max_seq_len_decode", None), + max_seq_len_prefill=getattr(attn_metadata, "_modelopt_max_seq_len_prefill", None), **common_kw, ) @@ -1068,23 +1145,11 @@ def collect_calibration_counts(model) -> dict[str, list[dict]]: phase with :func:`~.sparse_attn_calibration.fit_from_counts` — sparsity ratios are only formed after the global merge. """ - from .sparse_attn_calibration import merge_count_records, split_records_by_phase - splits = [ split_records_by_phase(getattr(impl, "_calib_records", [])) for impl in iter_sparse_impls(model) ] - phases = {phase for split in splits for phase, records in split.items() if records} - counts: dict[str, list[dict]] = {} - for phase in phases: - sources = [split.get(phase, []) for split in splits] - empty = sum(1 for source in sources if not source) - if empty: - # Every layer sees every launch, so a layer with no records for a - # phase others measured indicates a collection bug. - raise ValueError( - f"Misaligned calibration records: {empty}/{len(sources)} attention " - f"layer(s) recorded no {phase!r} samples while others did" - ) - counts[phase] = merge_count_records(sources) - return counts + # Same merge as the cross-rank aggregation: every layer sees every launch, + # so a layer with no records for a phase others measured indicates a + # collection bug (merge_phase_counts raises). + return merge_phase_counts(splits, source_desc="attention layer") diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index 724feb69f90..6e2ad911288 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -156,7 +156,7 @@ def _cudagraph_mode(model_runner): return mode if mode is not None else CUDAGraphMode.NONE -def _global_errors(model_runner) -> list[str]: +def _global_errors(model_runner, *, sparse_only: bool = False) -> list[str]: config = getattr(model_runner, "vllm_config", None) if config is None: return ["model_runner.vllm_config is required"] @@ -174,10 +174,16 @@ def _global_errors(model_runner) -> list[str]: errors.append("decode_context_parallel_size must be 1") if getattr(parallel, "enable_dbo", False) or getattr(parallel, "use_ubatching", False): errors.append("DBO/ubatching is unsupported") - if getattr(cache_config, "enable_prefix_caching", False): - errors.append("prefix caching is unsupported") - if getattr(config, "kv_transfer_config", None) is not None: - errors.append("KV transfer is unsupported") + if not sparse_only: + # Prefix caching and KV transfer break only flows that quantize the + # cache on write or measure per-request prefills (quantized installs, + # skip-softmax calibration). Sparse-only serving reads the cache + # unmodified and supports prefix-cache suffix attention by offsetting + # query positions (see the vllm_serve README limitations). + if getattr(cache_config, "enable_prefix_caching", False): + errors.append("prefix caching is unsupported") + if getattr(config, "kv_transfer_config", None) is not None: + errors.append("KV transfer is unsupported") if getattr(config, "speculative_config", None) is not None: errors.append("speculative decoding is unsupported") if _cudagraph_mode(model_runner).mixed_mode() == CUDAGraphMode.FULL: @@ -355,7 +361,12 @@ def _plan_vllm_attention( ) _require_supported_vllm() - errors = _global_errors(model_runner) if quantize else [] + # Engine-level checks apply to sparse-only installs too: the ModelOpt + # kernel path silently ignores decode context parallelism, DBO, and + # speculative decoding, and FULL mixed-batch graphs would capture stale + # per-launch thresholds (same rationale as the decode graph guard below). + # Sparse-only installs skip only the cache-mutation checks. + errors = _global_errors(model_runner, sparse_only=not quantize) mode = _cudagraph_mode(model_runner) quant_plugin: Any = _load_quant_plugin() if quantize else None plans = [] @@ -599,6 +610,16 @@ def install_vllm_skip_softmax_calibration(model_runner) -> VllmAttentionInstallR _raise_unsupported(errors, "skip-softmax calibration") plan = _InstallPlan(model_runner, tuple(plans), False, "SKIP_SOFTMAX_CALIBRATION") + # Per-request prompt lengths (same request order as the attention-metadata + # rows, kept aligned by vLLM's in-place batch reorder) let the adapter + # classify q_len == 1 rows: a decode step's KV span exceeds the prompt, + # while a 1-token final chunk of a chunked prefill is still inside it. + # The runner — not its input_batch — is attached: vLLM can rebuild + # input_batch after load_model (may_reinitialize_input_batch during KV + # cache init, e.g. hybrid Mamba/attention models), so the adapter must + # resolve the live object per forward. + for attention_plan in plans: + attention_plan.new_impl._calib_model_runner = model_runner return _apply_vllm_attention_plans(plan) diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py index 60523d44457..99e8a3a83a4 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py @@ -20,7 +20,7 @@ import pytest import torch -from conftest import make_qkv, make_varlen_meta, sdpa_reference +from conftest import make_qkv, make_varlen_meta from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor, e2m1_values @@ -454,9 +454,12 @@ def test_invalid_amax_raises(self): with pytest.raises(ValueError, match="p_qdq_amax"): attention(q, k, v, locs, lens, 8, p_qdq="fp8", p_qdq_amax=0.0) - @requires_native_e4m3 - def test_composes_with_skip_softmax(self): - """p_qdq composes with the skip-softmax feature in one launch.""" + def test_rejects_skip_softmax(self): + """p_qdq cannot combine with active skip-softmax in one launch. + + Quantized P changes the score distribution the skip thresholds were + calibrated on, so the kernel rejects the composition (pre-launch). + """ seq_len, num_heads, num_kv_heads, head_dim = 256, 4, 2, 64 scale = 1.0 / (head_dim**0.5) @@ -464,19 +467,18 @@ def test_composes_with_skip_softmax(self): q, k, v = make_qkv(seq_len, num_heads, num_kv_heads, head_dim, dtype=torch.float16) locs, lens = make_varlen_meta([seq_len]) - o = attention( - q, - k, - v, - locs, - lens, - seq_len, - softmax_scale=scale, - p_qdq="fp8", - skip_softmax_threshold=1e-3, - ) - ref = sdpa_reference(q, k, v, locs, lens) - torch.testing.assert_close(o, ref, rtol=5e-2, atol=5e-2) + with pytest.raises(ValueError, match="cannot be combined with attention quantization"): + attention( + q, + k, + v, + locs, + lens, + seq_len, + softmax_scale=scale, + p_qdq="fp8", + skip_softmax_threshold=1e-3, + ) def test_invalid_mode_raises(self): q, k, v = make_qkv(8, 2, 2, 32, dtype=torch.float16) diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index cf802abaf40..2e2a0105846 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -354,3 +354,36 @@ def test_static_vs_dynamic_fp4_kernels(self, set_torch_dtype, block_size, num_bl f"Mean abs diff: {(output_static - output_dynamic).abs().mean()}\n" f"Max relative diff: {((output_static - output_dynamic).abs() / (output_dynamic.abs() + 1e-8)).max()}" ) + + +def test_dynamic_block_quantize_forwards_block_size(monkeypatch): + """NVFP4 dynamic dispatch must pass the configured block size to the Triton kernel. + + The Triton kernel defaults to block_size=16; dropping the argument would + silently quantize non-16 NVFP4 configs with 16-element blocks while the + cuda_ext fallback honors the configured size (device-dependent numerics). + """ + recorded = {} + + def fake_fp4_fake_quant_block(inputs, amax, block_size=16): + recorded["block_size"] = block_size + return inputs.clone() + + monkeypatch.setattr(tensor_quant, "DISABLE_TRITON_KERNEL", False) + monkeypatch.setattr(triton_kernel, "IS_AVAILABLE", True, raising=False) + monkeypatch.setattr( + triton_kernel, "fp4_fake_quant_block", fake_fp4_fake_quant_block, raising=False + ) + + inputs = torch.randn(4, 64, device="cuda") + amax = inputs.abs().amax() + tensor_quant._dynamic_block_quantize_impl( + inputs, + 32, # non-default NVFP4 block size + amax, + 4, # num_bits total (E2M1) + 2, # exponent_bits -> (2, 1) + 8, # scale_num_bits -> (4, 3) + 4, # scale_exponent_bits + ) + assert recorded["block_size"] == 32 diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py index 543fbd0a2da..d3563d226cf 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -15,6 +15,7 @@ """Tests for skip-softmax calibration through the vLLM adapters and installer.""" +import sys from types import SimpleNamespace import pytest @@ -25,6 +26,7 @@ from vllm.v1.attention.backends.flashinfer import FlashInferImpl from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.quantization.plugins import vllm as quant_plugin from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as attention_plugin from modelopt.torch.sparsity.attention_sparsity.plugins import vllm_runtime from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ( @@ -177,8 +179,6 @@ def test_sparse_only_install_allows_skip_on_unquantized_layer(self): assert report.installed_count == 1 def test_quantized_plan_allows_nm_sparsity(self, monkeypatch): - from modelopt.torch.quantization.plugins import vllm as quant_plugin - monkeypatch.setattr( quant_plugin, "_get_device_dtype", @@ -197,8 +197,6 @@ class TestFlashInferLayoutGuard: def test_layout_helper_preserves_none(self, monkeypatch): """A getter returning None must not become the truthy string 'None'.""" - import sys - fake = SimpleNamespace(get_kv_cache_layout=lambda: None) monkeypatch.setitem(sys.modules, "vllm.v1.attention.backends.utils", fake) assert attention_plugin._flashinfer_kv_cache_layout() is None diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py index b3ceef30eab..9c448c1327c 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_runtime.py @@ -90,7 +90,6 @@ def test_sparse_install_from_checkpoint_is_validation_atomic(): nn.ModuleDict({"valid_attn": valid, "invalid_attn": invalid}), sparse_metadata=_sparse_metadata(), ) - del runner.vllm_config with pytest.raises(NotImplementedError, match="sliding_window"): vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) @@ -107,7 +106,6 @@ def test_installer_rejects_cross_attention_layout_even_if_marked_decoder(): nn.ModuleDict({"cross_attn": attention}), sparse_metadata=_sparse_metadata(), ) - del runner.vllm_config with pytest.raises(NotImplementedError, match="layout CrossAttention"): vllm_runtime.install_vllm_sparse_attention_from_checkpoint(runner) @@ -122,7 +120,6 @@ def test_sparse_install_uses_checkpoint_metadata(monkeypatch, impl_cls): nn.ModuleDict({"attn": attention}), sparse_metadata=_sparse_metadata(), ) - del runner.vllm_config monkeypatch.setattr( vllm_runtime.attention_plugin, "patch_flashinfer_metadata_builder", lambda: True ) diff --git a/tests/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index 62395ff5a7a..13ce5da55ef 100644 --- a/tests/unit/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -139,21 +139,7 @@ def test_forward_routes_every_mode_to_single_autotuner( assert kernel.kwargs["V_QDQ"] == expected_v_qdq -@pytest.mark.parametrize( - ("attention_kwargs", "expected_block_m"), - [ - ({"skip_softmax_threshold": 0.1, "measure_sparsity": True}, 128), - ( - { - "p_qdq": "nvfp4", - "skip_softmax_threshold": 0.1, - "measure_sparsity": True, - }, - 16, - ), - ], -) -def test_forward_measurement_uses_one_fixed_launch(monkeypatch, attention_kwargs, expected_block_m): +def test_forward_measurement_uses_one_fixed_launch(monkeypatch): """Counter measurement bypasses autotuning to avoid repeated atomic updates.""" pytest.importorskip("triton") @@ -173,10 +159,42 @@ def test_forward_measurement_uses_one_fixed_launch(monkeypatch, attention_kwargs starts = torch.tensor([0], dtype=torch.int32) lengths = torch.tensor([seq_len], dtype=torch.int32) - triton_fa.attention(q, k, v, starts, lengths, seq_len, **attention_kwargs) + triton_fa.attention( + q, k, v, starts, lengths, seq_len, skip_softmax_threshold=0.1, measure_sparsity=True + ) assert kernel.fn.launch_count == 1 - assert kernel.fn.kwargs["BLOCK_M"] == expected_block_m + assert kernel.fn.kwargs["BLOCK_M"] == 128 assert kernel.fn.kwargs["BLOCK_N"] == 128 assert kernel.fn.kwargs["num_stages"] == 1 assert kernel.fn.kwargs["num_warps"] == 4 + + +@pytest.mark.parametrize( + "qdq_kwargs", + [{"p_qdq": "nvfp4"}, {"p_qdq": "fp8"}, {"v_qdq": "nvfp4", "v_qdq_amax": 1.0}], +) +def test_forward_rejects_skip_softmax_with_qdq(monkeypatch, qdq_kwargs): + """Active skip-softmax rejects P/V QDQ before any kernel launch.""" + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + kernel = _ForbiddenKernel() + kernel.fn = _ForbiddenKernel() + monkeypatch.setattr(triton_fa, "_attn_fwd", kernel) + monkeypatch.setattr(triton_fa.torch.cuda, "device", lambda _device: nullcontext()) + monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) + monkeypatch.setattr(triton_fa, "_load_qdq_helpers", lambda: None) + + seq_len = 129 + q = torch.empty(seq_len, 2, 16) + k = torch.empty(seq_len, 1, 16) + v = torch.empty_like(k) + starts = torch.tensor([0], dtype=torch.int32) + lengths = torch.tensor([seq_len], dtype=torch.int32) + + with pytest.raises(ValueError, match="cannot be combined with attention quantization"): + triton_fa.attention( + q, k, v, starts, lengths, seq_len, skip_softmax_threshold=0.1, **qdq_kwargs + ) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py index e85ad3ab77e..7ff174952b3 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py @@ -177,6 +177,31 @@ def test_canonical_schema(self): assert group["target_sparsity"] == {"prefill": 0.4, "decode": 0.4} assert config["producer"]["name"] == "modelopt" + def test_target_sparsity_covers_only_fitted_phases(self): + """A phase without calibrated (a, b) must not claim a sparsity target.""" + config = build_sparse_attention_config({"prefill": {"a": 7.9, "b": 8.6}}, 0.4) + group = config["config_groups"]["group_0"] + assert group["target_sparsity"] == {"prefill": 0.4} + assert "decode" not in group["threshold_scale_factor"] + + def test_replaced_skip_group_keeps_layer_policy(self): + """Recalibration replaces thresholds but keeps ignore/initial_disabled_steps.""" + existing = { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "ignore": ["model.layers.0.self_attn"], + "initial_disabled_steps": 4, + "threshold_scale_factor": {"prefill": {"a": 1.0, "b": 1.0}}, + } + } + } + config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) + group = config["config_groups"]["group_0"] + assert group["ignore"] == ["model.layers.0.self_attn"] + assert group["initial_disabled_steps"] == 4 + assert group["threshold_scale_factor"]["prefill"] == {"a": 7.9, "b": 8.6} + def test_preserves_nm_groups_and_replaces_old_skip_group(self): existing = { "config_groups": { From 55a1d4a44a5bfa53727157077edb4179a84c5879 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Tue, 25 Aug 2026 13:33:23 -0700 Subject: [PATCH 10/14] Support FlashInfer HND calibration Signed-off-by: Kai Xu --- .../kernels/sparsity/attention/calibrate.py | 12 +-- .../attention_sparsity/plugins/vllm.py | 48 +---------- .../plugins/vllm_runtime.py | 6 -- .../attention/test_paged_calibrate.py | 12 ++- .../test_vllm_calibration.py | 85 ++++--------------- 5 files changed, 36 insertions(+), 127 deletions(-) diff --git a/modelopt/torch/kernels/sparsity/attention/calibrate.py b/modelopt/torch/kernels/sparsity/attention/calibrate.py index 3655c4cbfcd..30c51df4247 100644 --- a/modelopt/torch/kernels/sparsity/attention/calibrate.py +++ b/modelopt/torch/kernels/sparsity/attention/calibrate.py @@ -311,11 +311,13 @@ def attention_calibrate( 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). + k_cache: Logical paged K-cache view + ``[num_blocks, page_size, num_kv_heads, head_dim]``. Arbitrary + strides support both NHD and HND physical layouts. When provided, + K/V are read via ``block_table`` 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. diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index a6d3abffa4e..267e2390fb5 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -28,7 +28,6 @@ """ import functools -import importlib import inspect import math import warnings @@ -272,35 +271,6 @@ 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: - 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 - - def _forward_calibrate( impl, *, @@ -342,13 +312,11 @@ def _forward_calibrate( 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. + if key_cache.ndim != 4 or key_cache.shape[2] != impl.num_kv_heads: 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" + "skip-softmax calibration requires a logical KV-cache view shaped " + f"[blocks, page, heads, dim], got {tuple(key_cache.shape)} with " + f"num_kv_heads={impl.num_kv_heads}" ) page_size = key_cache.shape[1] trials = impl._calib_threshold_trials @@ -945,14 +913,6 @@ 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() diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index 6e2ad911288..136d7faa34b 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -601,12 +601,6 @@ 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 c7390382098..2c478ed3603 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_paged_calibrate.py @@ -30,12 +30,15 @@ TRIALS = [1e-3, 1e-2, 1e-1, 3e-1] -def _pack_paged(k, v, page_size, *, shuffle=True, num_spare_blocks=7): +def _pack_paged(k, v, page_size, *, layout="NHD", 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]) + shape = (num_blocks + num_spare_blocks, page_size, k.shape[1], k.shape[2]) + k_cache = k.new_zeros(shape) + if layout == "HND": + k_cache = k.new_zeros(shape[0], shape[2], shape[1], shape[3]).permute(0, 2, 1, 3) 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): @@ -48,8 +51,9 @@ def _pack_paged(k, v, page_size, *, shuffle=True, num_spare_blocks=7): class TestPagedCalibrate: + @pytest.mark.parametrize("layout", ["NHD", "HND"]) @pytest.mark.parametrize("seq_len", [256, 300, 512]) # 300: non-128-aligned padding - def test_paged_matches_contiguous_prefill(self, seq_len): + def test_paged_matches_contiguous_prefill(self, seq_len, layout): """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 @@ -60,7 +64,7 @@ def test_paged_matches_contiguous_prefill(self, seq_len): 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_cache, v_cache, block_table = _pack_paged(k, v, page_size, layout=layout) k_dummy = torch.empty(0, num_kv_heads, head_dim, device=q.device, dtype=q.dtype) out_paged, counters_paged = attention_calibrate( q, 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 d3563d226cf..088a91c2ccb 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -15,7 +15,6 @@ """Tests for skip-softmax calibration through the vLLM adapters and installer.""" -import sys from types import SimpleNamespace import pytest @@ -192,69 +191,12 @@ def test_quantized_plan_allows_nm_sparsity(self, monkeypatch): assert plan.layers[0].sparse_kw.get("sparsity_n") == 2 -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'.""" - 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") +class TestFlashInferLayout: + def test_installer_accepts_flashinfer(self): 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" + report = vllm_runtime.install_vllm_skip_softmax_calibration(runner) + assert report.installed_count == 1 class TestSparseOnlyGraphGuard: @@ -431,11 +373,11 @@ def __init__(self, impls): {"sample_length": 128, "total_tiles": [8, 8, 8], "skipped_tiles": [1, 3, 5]} ] - def test_rejects_non_nhd_cache_layout(self): + def test_rejects_non_logical_cache_shape(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. + # A physical HND tensor must be exposed as a logical [blocks, page, heads, dim] view. kv_cache = torch.zeros(2, 1, num_kv_heads, 16, head_dim, dtype=torch.bfloat16) attn_metadata = SimpleNamespace( num_actual_tokens=1, @@ -446,7 +388,7 @@ def test_rejects_non_nhd_cache_layout(self): 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"): + with pytest.raises(NotImplementedError, match="logical KV-cache view"): impl.forward( layer=None, query=q, @@ -487,7 +429,8 @@ def test_rejects_non_16bit_cache(self): # FlashInfer adapter: cache write must precede the calibrate-kernel read # --------------------------------------------------------------------------- class TestFlashInferCalibrationOrdering: - def test_cache_write_happens_before_calibrate_read(self, monkeypatch): + @pytest.mark.parametrize("layout", ["NHD", "HND"]) + def test_cache_write_happens_before_calibrate_read(self, monkeypatch, layout): calls = [] monkeypatch.setattr( attention_plugin, @@ -497,6 +440,7 @@ def test_cache_write_happens_before_calibrate_read(self, monkeypatch): def fake_calibrate(q, *args, **kwargs): calls.append("calibrate") + calls.append(kwargs["k_cache"].stride()) counters = torch.zeros(len(TRIALS), 2, dtype=torch.int64) return torch.zeros_like(q), counters @@ -511,7 +455,12 @@ def fake_calibrate(q, *args, **kwargs): _calib_threshold_trials=list(TRIALS), _calib_records=[], ) - kv_cache = torch.zeros(3, 2, page, num_kv_heads, head_dim, dtype=torch.bfloat16) + shape = (3, 2, page, num_kv_heads, head_dim) + kv_cache = torch.zeros(shape, dtype=torch.bfloat16) + if layout == "HND": + kv_cache = torch.zeros( + shape[0], shape[1], shape[3], shape[2], shape[4], dtype=torch.bfloat16 + ).permute(0, 1, 3, 2, 4) attn_metadata = SimpleNamespace( _modelopt_block_table=torch.zeros(1, 1, dtype=torch.int32), _modelopt_seq_lens=torch.tensor([8], dtype=torch.int32), @@ -536,7 +485,7 @@ def fake_calibrate(q, *args, **kwargs): output=torch.empty_like(q), ) - assert calls == ["cache_write", "calibrate"] + assert calls == ["cache_write", "calibrate", kv_cache[:, 0].stride()] assert torch.isfinite(out).all() assert len(impl._calib_records) == 1 assert impl._calib_records[0]["phase"] == "decode" From b4e72920d187224747d17830d1562334c95e1070 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Thu, 27 Aug 2026 15:37:54 -0700 Subject: [PATCH 11/14] Tighten vLLM skip-softmax calibration PR Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 +- examples/vllm_serve/README.md | 6 ++++ modelopt/torch/quantization/tensor_quant.py | 5 +-- .../quantization/test_tensor_quant_cuda.py | 33 ------------------- 4 files changed, 8 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8897061c086..b38b24cfbc7 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog *Sparsity* -- Add **skip-softmax threshold calibration through vLLM**: ``install_vllm_skip_softmax_calibration`` installs calibration adapters onto every attention layer of a loaded vLLM model (FlashAttention and FlashInfer backends, validation-before-mutation), measures per-request KV-tile skip counts over the paged KV cache with the Triton calibration kernel — full dense attention, so generation is numerically unchanged — for both prefill and decode, aggregates raw counts across tensor-parallel ranks, fits the exponential threshold model once per phase, and writes the same canonical ``sparse_attention_config`` block the HF export produces (existing N:M sparse-softmax groups are preserved). Active skip-softmax serving launches now run on the fixed 128x128 calibration tile instead of autotuned tiles, so the sparsity realized at serve time matches the calibrated ``(a, b)`` model. Configurations that cannot compile the 128x128 tile are rejected rather than re-tiled, and skip-softmax cannot be combined with attention quantization. **Behavior change for sparse-only serving installs:** ``install_vllm_sparse_attention_from_checkpoint`` now runs the same engine-level compatibility validation as quantized installs (decode context parallelism, DBO, speculative decoding, and FULL mixed-batch CUDA graphs are rejected; prefix caching remains supported), and checkpoints with a calibrated ``decode`` ``threshold_scale_factor`` are rejected under a FULL decode CUDA graph mode — including vLLM's default ``FULL_AND_PIECEWISE`` — because the captured graph would replay one request's stale threshold. Previously such deployments started and silently served miscalibrated sparsity; pass ``--enforce-eager`` (or select a non-FULL decode graph mode). See ``examples/vllm_serve/calibrate_sparse_attn.py``. +- Add skip-softmax threshold calibration through vLLM for FlashAttention and FlashInfer, exporting prefill and decode fits as ``sparse_attention_config``. See ``examples/vllm_serve/calibrate_sparse_attn.py`` for usage and compatibility constraints. *Quantization* diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index c5596003cd3..741f515c383 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -192,6 +192,12 @@ python calibrate_sparse_attn.py \ --decode_tokens 32 --tensor_parallel_size 8 --update_checkpoint_config ``` +Calibration always writes `sparse_attention_config.json` in the current directory. +`--update_checkpoint_config` also merges that configuration into `/config.json` in +place, which lets `vllm_serve_sparse_attn.py` load it automatically. This option requires +`` to be a local checkpoint directory; without it, merge the generated configuration +into the checkpoint manually before serving. + 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 — 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. diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 9e5b6c186ec..20e083491aa 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -180,10 +180,7 @@ def _dynamic_block_quantize_impl( and not DISABLE_TRITON_KERNEL and amax is not None ): - # Forward the configured block size: the kernel defaults to 16, and - # silently ignoring a non-16 NVFP4 block config here would diverge - # from the cuda_ext fallback below (which honors block_size). - return triton_kernel.fp4_fake_quant_block(inputs, amax, block_size=block_size) + return triton_kernel.fp4_fake_quant_block(inputs, amax) cuda_ext_mx = get_cuda_ext_mx(raise_if_failed=True) return cuda_ext_mx.fused_amax_convert( inputs, diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index 2e2a0105846..cf802abaf40 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -354,36 +354,3 @@ def test_static_vs_dynamic_fp4_kernels(self, set_torch_dtype, block_size, num_bl f"Mean abs diff: {(output_static - output_dynamic).abs().mean()}\n" f"Max relative diff: {((output_static - output_dynamic).abs() / (output_dynamic.abs() + 1e-8)).max()}" ) - - -def test_dynamic_block_quantize_forwards_block_size(monkeypatch): - """NVFP4 dynamic dispatch must pass the configured block size to the Triton kernel. - - The Triton kernel defaults to block_size=16; dropping the argument would - silently quantize non-16 NVFP4 configs with 16-element blocks while the - cuda_ext fallback honors the configured size (device-dependent numerics). - """ - recorded = {} - - def fake_fp4_fake_quant_block(inputs, amax, block_size=16): - recorded["block_size"] = block_size - return inputs.clone() - - monkeypatch.setattr(tensor_quant, "DISABLE_TRITON_KERNEL", False) - monkeypatch.setattr(triton_kernel, "IS_AVAILABLE", True, raising=False) - monkeypatch.setattr( - triton_kernel, "fp4_fake_quant_block", fake_fp4_fake_quant_block, raising=False - ) - - inputs = torch.randn(4, 64, device="cuda") - amax = inputs.abs().amax() - tensor_quant._dynamic_block_quantize_impl( - inputs, - 32, # non-default NVFP4 block size - amax, - 4, # num_bits total (E2M1) - 2, # exponent_bits -> (2, 1) - 8, # scale_num_bits -> (4, 3) - 4, # scale_exponent_bits - ) - assert recorded["block_size"] == 32 From 15219f3f84f32ee94d6e1eead9fac0eda328ffca Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 28 Aug 2026 15:30:13 -0700 Subject: [PATCH 12/14] Address vLLM calibration review feedback Signed-off-by: Kai Xu --- CHANGELOG.rst | 3 +- examples/vllm_serve/README.md | 3 +- examples/vllm_serve/calibrate_sparse_attn.py | 56 ++++++++-- .../kernels/common/attention/triton_fa.py | 2 +- .../kernels/sparsity/attention/calibrate.py | 20 +++- .../methods/flash_skip_softmax.py | 10 -- .../attention_sparsity/plugins/vllm.py | 32 +++--- .../plugins/vllm_runtime.py | 7 +- .../vllm_serve/test_calibrate_sparse_attn.py | 100 ++++++++++++++++++ .../test_sparse_attn_worker.py | 90 ++++++++++------ .../test_vllm_calibration.py | 14 +++ 11 files changed, 263 insertions(+), 74 deletions(-) create mode 100644 tests/examples/vllm_serve/test_calibrate_sparse_attn.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b38b24cfbc7..1e8860e5870 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,8 @@ Changelog *Sparsity* -- Add skip-softmax threshold calibration through vLLM for FlashAttention and FlashInfer, exporting prefill and decode fits as ``sparse_attention_config``. See ``examples/vllm_serve/calibrate_sparse_attn.py`` for usage and compatibility constraints. +- Add skip-softmax threshold calibration through vLLM for FlashAttention and FlashInfer, exporting prefill and decode fits as ``sparse_attention_config``. Skip-softmax serving now uses fixed 128x128 tiles instead of autotuning so runtime sparsity matches calibration. +- Sparse-only vLLM installs now reject unsupported DCP, DBO/ubatching, speculative decoding, and FULL mixed-batch CUDA graphs; calibrated decode also rejects FULL decode graphs. *Quantization* diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 741f515c383..df6ebdb6461 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -180,7 +180,7 @@ If the checkpoint has no `sparse_attention_config`, the sparse-only installer pa ### 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: +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. Pipeline parallelism is not supported by calibration. ```bash # One-time: fetch the RULER essay haystack @@ -219,6 +219,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. 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). +- Skip-softmax calibration requires pipeline-parallel size 1 because raw count records are merged across tensor-parallel head shards. - `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`. Checkpoints with a calibrated `decode` `threshold_scale_factor` are rejected at install under a FULL decode CUDA graph mode (including vLLM's default `FULL_AND_PIECEWISE`): the captured graph would replay one request's stale threshold. - Sparse-only installs validate engine-level compatibility like quantized installs do: decode context parallelism, DBO, speculative decoding, and FULL mixed-batch CUDA graphs are rejected (prefix caching remains supported, per the bullet above). diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py index 42c4d36a972..316e8a0d9c4 100644 --- a/examples/vllm_serve/calibrate_sparse_attn.py +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -46,6 +46,7 @@ import argparse import json +import math import os import sys from pathlib import Path @@ -58,6 +59,38 @@ merge_phase_counts, ) +_LOCKED_ENGINE_KWARGS = frozenset({"model", "worker_cls", "enforce_eager", "enable_prefix_caching"}) + + +def _sparse_ratio(value: str) -> float: + ratio = float(value) + if not math.isfinite(ratio) or not 0.0 <= ratio <= 1.0: + raise argparse.ArgumentTypeError("must be a finite value between 0.0 and 1.0") + return ratio + + +def _nonnegative_int(value: str) -> int: + result = int(value) + if result < 0: + raise argparse.ArgumentTypeError("must be non-negative") + return result + + +def _engine_kwargs(value: str) -> dict: + try: + kwargs = json.loads(value) + except json.JSONDecodeError as err: + raise argparse.ArgumentTypeError(f"must be a JSON object: {err.msg}") from err + if not isinstance(kwargs, dict): + raise argparse.ArgumentTypeError("must be a JSON object") + if locked := sorted(_LOCKED_ENGINE_KWARGS & kwargs.keys()): + raise argparse.ArgumentTypeError( + "cannot override calibration-controlled option(s): " + ", ".join(locked) + ) + if kwargs.get("pipeline_parallel_size", 1) != 1: + raise argparse.ArgumentTypeError("pipeline_parallel_size must be 1 for calibration") + return kwargs + def _load_prompts(llm, args) -> list[str]: """Load override prompts from a file, or build the default RULER set.""" @@ -112,8 +145,9 @@ def _write_config(ckpt: str, sparse_config: dict, update_checkpoint: bool) -> No 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)." + "[ModelOpt] Checkpoint not modified. Merge the generated configuration as " + f"'sparse_attention_config' in {ckpt}/config.json before serving. On future " + "calibration runs, pass --update_checkpoint_config to do this automatically." ) return @@ -128,7 +162,7 @@ def _write_config(ckpt: str, sparse_config: dict, update_checkpoint: bool) -> No print(f"[ModelOpt] Merged sparse_attention_config into {config_json}") -def main(): +def _build_parser() -> argparse.ArgumentParser: 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( @@ -160,13 +194,13 @@ def main(): ) parser.add_argument( "--target_sparse_ratio", - type=float, + type=_sparse_ratio, default=0.5, help="Target sparsity baked into the exported config (applied to both phases)", ) parser.add_argument( "--decode_tokens", - type=int, + type=_nonnegative_int, default=32, help="Decode attention steps per prompt (drives decode-phase calibration). " "Generation runs decode_tokens + 1 output tokens: the first output token " @@ -200,7 +234,7 @@ def main(): ) parser.add_argument( "--engine_kwargs", - type=str, + type=_engine_kwargs, default=None, help="JSON dict of extra vLLM engine kwargs, e.g. " '\'{"enable_expert_parallel": true, "mamba_cache_mode": "align"}\' ' @@ -216,6 +250,11 @@ def main(): action="store_true", help="Merge the calibrated config into /config.json in place", ) + return parser + + +def main(): + parser = _build_parser() args = parser.parse_args() if args.update_checkpoint_config and not (Path(args.model) / "config.json").is_file(): @@ -260,10 +299,7 @@ def main(): 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_kwargs.update(args.engine_kwargs) llm = LLM(**llm_kwargs) # Built after engine init so the RULER builder reuses the engine's tokenizer. diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 4a223d4c2ab..6f7b83cdf08 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -1039,7 +1039,7 @@ def grid(META): # kernel dereferences the right pointers instead of triggering an # illegal memory access. with torch.cuda.device(q.device): - if do_measure or apply_skip: + if apply_skip: # Fixed-tile launches, bypassing autotune: # - Measurement: runtime counters mutate global tensors, so they # must not run through autotune candidate trials. diff --git a/modelopt/torch/kernels/sparsity/attention/calibrate.py b/modelopt/torch/kernels/sparsity/attention/calibrate.py index 30c51df4247..b508e9d4ab4 100644 --- a/modelopt/torch/kernels/sparsity/attention/calibrate.py +++ b/modelopt/torch/kernels/sparsity/attention/calibrate.py @@ -280,6 +280,23 @@ def _log2_threshold_tensor( ) +def _validate_threshold_trials(threshold_trials) -> list[float]: + """Return finite skip thresholds in the kernel's open interval ``(0, 1)``.""" + if not threshold_trials: + raise ValueError("threshold_trials must be a non-empty list") + try: + trials = [float(value) for value in threshold_trials] + except (TypeError, ValueError) as err: + raise ValueError("threshold_trials must contain only real numbers") from err + invalid = [value for value in trials if not math.isfinite(value) or not 0.0 < value < 1.0] + if invalid: + raise ValueError( + "threshold_trials must contain only finite values strictly between 0 and 1; " + f"got {invalid}" + ) + return trials + + def attention_calibrate( q: torch.Tensor, k: torch.Tensor, @@ -331,8 +348,7 @@ def attention_calibrate( ``[:, 0]`` = total tile evaluations, ``[:, 1]`` = skipped tiles. Sparsity per threshold = ``counters[:, 1] / counters[:, 0]``. """ - if threshold_trials is None or len(threshold_trials) == 0: - raise ValueError("threshold_trials must be a non-empty list") + threshold_trials = _validate_threshold_trials(threshold_trials) is_paged = k_cache is not None if is_paged and block_table is None: 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 baefbd7058d..c1d6465ba66 100644 --- a/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py +++ b/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py @@ -193,16 +193,6 @@ 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.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 267e2390fb5..7b4de426b9d 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -45,7 +45,10 @@ ) 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 +from modelopt.torch.kernels.sparsity.attention.calibrate import ( + _validate_threshold_trials, + attention_calibrate, +) from .sparse_attn_calibration import merge_phase_counts, split_records_by_phase @@ -72,6 +75,16 @@ def _flash_attention_kv_cache_layout() -> str: raise RuntimeError(f"Unsupported vLLM FlashAttention KV cache shape {cache_shape}") +def _flash_attention_kv_cache_views(kv_cache: torch.Tensor, head_size: int): + """Return logical K/V cache views for the installed FlashAttention layout.""" + cache_layout = _flash_attention_kv_cache_layout() + if cache_layout == "kv-first": + return kv_cache.unbind(0) + if cache_layout == "blocks-first": + return kv_cache.unbind(1) + return kv_cache.transpose(1, 2).split(head_size, dim=-1) + + def _target_sparse_ratio_for_phase(target_sparse_ratio, phase: str) -> float: """Return target sparsity for a phase, defaulting old checkpoint metadata.""" if isinstance(target_sparse_ratio, float | int): @@ -326,8 +339,8 @@ def _forward_calibrate( b_seq_len_i32 = (cu_seqlens_q[1 : batch + 1] - cu_seqlens_q[:batch]).to(torch.int32) seq_lens_i32 = seq_lens[:batch].to(torch.int32) b_start_loc_zero = torch.zeros(1, device=query.device, dtype=torch.int32) - # vLLM v1 metadata tensors are GPU-resident: take one host copy per launch - # instead of stalling the stream with per-request ``.item()`` syncs. + # Copy scheduling metadata once per launch. The calibration wrapper and + # counter collection still synchronize once per measured request. cu_seqlens_q_cpu = cu_seqlens_q[: batch + 1].cpu() seq_lens_cpu = seq_lens[:batch].cpu() # Per-request prompt lengths (same request order as the metadata rows) @@ -680,7 +693,7 @@ def native_forward(): 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) + key_cache, value_cache = _flash_attention_kv_cache_views(kv_cache, self.head_size) return _forward_calibrate( self, query=query, @@ -703,13 +716,7 @@ def native_forward(): if resolved is None: return native_forward() - cache_layout = _flash_attention_kv_cache_layout() - if cache_layout == "kv-first": - key_cache, value_cache = kv_cache.unbind(0) - elif cache_layout == "blocks-first": - key_cache, value_cache = kv_cache.unbind(1) - else: - key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) + key_cache, value_cache = _flash_attention_kv_cache_views(kv_cache, self.head_size) is_decode_only = attn_metadata.max_query_len <= 1 common_kw = { "layer": layer, @@ -1080,8 +1087,7 @@ def iter_sparse_impls(model): 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.") + threshold_trials = _validate_threshold_trials(threshold_trials) for impl in impls: impl._calibrate = True impl._calib_threshold_trials = list(threshold_trials) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index 136d7faa34b..51ca1759c65 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -565,8 +565,8 @@ def install_vllm_skip_softmax_calibration(model_runner) -> VllmAttentionInstallR 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. + model and KV-cache dtypes, pipeline-parallel size 1, no active attention + Q/K/P/V fakequant, and a FlashAttention or FlashInfer backend per layer. """ from vllm.config.compilation import CUDAGraphMode @@ -579,6 +579,9 @@ def install_vllm_skip_softmax_calibration(model_runner) -> VllmAttentionInstallR _require_supported_vllm() errors = _global_errors(model_runner) + parallel = getattr(getattr(model_runner, "vllm_config", None), "parallel_config", None) + if getattr(parallel, "pipeline_parallel_size", 1) != 1: + errors.append("pipeline_parallel_size must be 1 for skip-softmax calibration") if _cudagraph_mode(model_runner) != CUDAGraphMode.NONE: errors.append( "skip-softmax calibration requires eager execution (enforce_eager=True); " diff --git a/tests/examples/vllm_serve/test_calibrate_sparse_attn.py b/tests/examples/vllm_serve/test_calibrate_sparse_attn.py new file mode 100644 index 00000000000..ab8bec47a9f --- /dev/null +++ b/tests/examples/vllm_serve/test_calibrate_sparse_attn.py @@ -0,0 +1,100 @@ +# 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 vLLM skip-softmax calibration driver.""" + +import importlib +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "vllm_serve" + + +@pytest.fixture +def calibration_driver(monkeypatch): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + return importlib.import_module("calibrate_sparse_attn") + + +@pytest.mark.parametrize( + "flags", + [ + ["--target_sparse_ratio", "-0.1"], + ["--target_sparse_ratio", "1.1"], + ["--target_sparse_ratio", "nan"], + ["--decode_tokens", "-1"], + ["--engine_kwargs", "[]"], + ["--engine_kwargs", "not-json"], + ["--engine_kwargs", '{"model": "/other"}'], + ["--engine_kwargs", '{"worker_cls": "other.Worker"}'], + ["--engine_kwargs", '{"enforce_eager": false}'], + ["--engine_kwargs", '{"enable_prefix_caching": true}'], + ["--engine_kwargs", '{"pipeline_parallel_size": 2}'], + ], +) +def test_parser_rejects_invalid_inputs_before_engine_start(calibration_driver, flags): + with pytest.raises(SystemExit): + calibration_driver._build_parser().parse_args(["/checkpoint", *flags]) + + +def test_parser_accepts_safe_engine_kwargs(calibration_driver): + args = calibration_driver._build_parser().parse_args( + ["/checkpoint", "--engine_kwargs", '{"enable_expert_parallel": true}'] + ) + assert args.engine_kwargs == {"enable_expert_parallel": True} + + +def test_load_prompts_reads_nonempty_lines(calibration_driver, tmp_path): + prompts_file = tmp_path / "prompts.txt" + prompts_file.write_text(" first prompt\n\nsecond prompt \n") + args = SimpleNamespace(prompts_file=str(prompts_file)) + + assert calibration_driver._load_prompts(None, args) == ["first prompt", "second prompt"] + + +def test_existing_sparse_config_reads_only_dict(calibration_driver, tmp_path): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + config_path = checkpoint / "config.json" + config_path.write_text(json.dumps({"sparse_attention_config": {"config_groups": {}}})) + assert calibration_driver._existing_sparse_config(str(checkpoint)) == {"config_groups": {}} + + config_path.write_text(json.dumps({"sparse_attention_config": ["invalid"]})) + assert calibration_driver._existing_sparse_config(str(checkpoint)) is None + + +@pytest.mark.parametrize("update_checkpoint", [False, True]) +def test_write_config_emits_artifact_and_optionally_updates_checkpoint( + calibration_driver, tmp_path, monkeypatch, update_checkpoint +): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + config_path = checkpoint / "config.json" + config_path.write_text(json.dumps({"model_type": "test"})) + sparse_config = {"config_groups": {"group_0": {"algorithm": "skip_softmax"}}} + monkeypatch.chdir(tmp_path) + + calibration_driver._write_config(str(checkpoint), sparse_config, update_checkpoint) + + assert json.loads((tmp_path / "sparse_attention_config.json").read_text()) == sparse_config + checkpoint_config = json.loads(config_path.read_text()) + assert checkpoint_config["model_type"] == "test" + if update_checkpoint: + assert checkpoint_config["sparse_attention_config"] == sparse_config + else: + assert "sparse_attention_config" not in checkpoint_config 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 f87a8cbb26b..6a909011359 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 @@ -615,6 +615,62 @@ def fake_attention(query, **kwargs): assert captured["page_size"] == 16 +@pytest.mark.parametrize( + ("layout", "backend_shape"), + [ + ("kv-first", (2, 3, 16, 1, 16)), + ("blocks-first", (3, 2, 16, 1, 16)), + ("packed", (3, 1, 16, 32)), + ], +) +def test_flash_attention_calibration_follows_backend_kv_cache_layout( + monkeypatch, layout, backend_shape +): + impl = _make_flash_attention_impl() + impl._calibrate = True + impl._calib_threshold_trials = [1e-3] + if layout == "packed": + shape = [3, impl.num_kv_heads, 16, 2 * impl.head_size] + else: + shape = [3, 16, impl.num_kv_heads, impl.head_size] + shape.insert(0 if layout == "kv-first" else 1, 2) + monkeypatch.setattr( + FlashAttentionBackend, "get_kv_cache_shape", staticmethod(lambda *_args: backend_shape) + ) + vllm_plugin._flash_attention_kv_cache_layout.cache_clear() + kv_cache = torch.zeros(shape, dtype=torch.float16) + query = torch.zeros(4, impl.num_heads, impl.head_size, dtype=torch.float16) + metadata = _flash_attention_metadata(query.shape[0], 16) + captured = {} + + def fake_calibrate(_impl, **kwargs): + captured.update(kwargs) + return kwargs["output"] + + monkeypatch.setattr(vllm_plugin, "_forward_calibrate", fake_calibrate) + + try: + output = torch.empty_like(query) + assert impl.forward(None, query, query, query, kv_cache, metadata, output=output) is output + finally: + vllm_plugin._flash_attention_kv_cache_layout.cache_clear() + + if layout == "packed": + expected_key_cache, expected_value_cache = kv_cache.transpose(1, 2).split( + impl.head_size, dim=-1 + ) + else: + expected_key_cache, expected_value_cache = kv_cache.unbind(0 if layout == "kv-first" else 1) + for name, expected in ( + ("key_cache", expected_key_cache), + ("value_cache", expected_value_cache), + ): + actual = captured[name] + assert actual.shape == expected.shape + assert actual.stride() == expected.stride() + assert actual.data_ptr() == expected.data_ptr() + + def _flash_attention_mixed_metadata(decode_len=1, prefill_len=17): query_lens = (decode_len, prefill_len) seq_lens = (16, 34) @@ -1116,40 +1172,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 = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size) - 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/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py index 088a91c2ccb..4774fdf891c 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 @@ -101,6 +101,12 @@ def test_rejects_non_eager_execution(self): with pytest.raises(NotImplementedError, match="enforce_eager"): vllm_runtime.install_vllm_skip_softmax_calibration(runner) + def test_rejects_pipeline_parallelism(self): + runner = _model_runner(nn.ModuleDict({"attn": _bare_attention()})) + runner.vllm_config.parallel_config.pipeline_parallel_size = 2 + with pytest.raises(NotImplementedError, match="pipeline_parallel_size must be 1"): + 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) @@ -287,6 +293,14 @@ def _sdpa_reference(q, k, v, is_causal): return out.squeeze(0).transpose(0, 1).to(q.dtype) +@pytest.mark.parametrize("trials", [[], [0.0], [-1e-3], [1.0], [float("inf")], [float("nan")]]) +def test_enable_rejects_invalid_threshold_trials(trials): + impl = object.__new__(ModelOptSparseAttentionImpl) + with pytest.raises(ValueError, match="threshold_trials"): + enable_calibration([impl], trials) + assert not hasattr(impl, "_calibrate") + + @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") class TestCalibrationForward: def test_mixed_batch_records_phases_and_stays_dense(self): From 0bcd9222eae2880072c525999b8984428d4d75ac Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 28 Aug 2026 17:06:25 -0700 Subject: [PATCH 13/14] fix: address vLLM calibration review feedback Signed-off-by: Kai Xu --- CHANGELOG.rst | 2 +- examples/vllm_serve/README.md | 8 +- examples/vllm_serve/calibrate_sparse_attn.py | 33 ++++++- .../kernels/common/attention/triton_fa.py | 96 +++++++++++++------ .../plugins/sparse_attn_calibration.py | 17 ++-- .../attention_sparsity/plugins/vllm.py | 2 +- .../plugins/vllm_runtime.py | 42 +++++++- .../vllm_serve/test_calibrate_sparse_attn.py | 35 +++++++ .../test_sparse_attn_worker.py | 31 ++++++ .../test_vllm_calibration.py | 28 ++++++ .../common/attention/test_triton_fa.py | 50 ++++++++++ .../test_sparse_attn_calibration.py | 4 +- 12 files changed, 300 insertions(+), 48 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1e8860e5870..5a305f76c62 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog *Sparsity* -- Add skip-softmax threshold calibration through vLLM for FlashAttention and FlashInfer, exporting prefill and decode fits as ``sparse_attention_config``. Skip-softmax serving now uses fixed 128x128 tiles instead of autotuning so runtime sparsity matches calibration. +- Add skip-softmax threshold calibration through vLLM for FlashAttention and FlashInfer, exporting prefill and decode fits as ``sparse_attention_config``. Skip-softmax serving keeps the calibrated 128-token KV-tile granularity (and 128-row prefill Q tiles), autotunes only its execution schedule, and uses a smaller Q tile for one-token decode. - Sparse-only vLLM installs now reject unsupported DCP, DBO/ubatching, speculative decoding, and FULL mixed-batch CUDA graphs; calibrated decode also rejects FULL decode graphs. *Quantization* diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index df6ebdb6461..6d5abdec3b3 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -180,7 +180,7 @@ If the checkpoint has no `sparse_attention_config`, the sparse-only installer pa ### 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. Pipeline parallelism is not supported by calibration. +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. Pipeline and data parallelism are not supported by calibration. ```bash # One-time: fetch the RULER essay haystack @@ -200,9 +200,9 @@ into the checkpoint manually before serving. 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 — 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. +`install_vllm_skip_softmax_calibration` (called by `sparse_attn_worker.SkipSoftmaxCalibWorker` at model load) swaps calibration adapters onto each attention layer not listed in the checkpoint's existing skip-softmax `ignore` policy after validating all selected layers — 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 the existing skip-softmax layer policy and 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. +Calibration and serving use the same 128-token KV-tile skip granularity and the same 128-row Q tile for prefill, so serving realizes the calibrated skip decision. One-token decode uses a 16-row Q compute tile because its padding rows cannot affect the decision. Serving autotunes only the execution schedule (`num_warps` / `num_stages`); measurement remains a single fixed launch because its counters have side effects. 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. @@ -219,7 +219,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. 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). -- Skip-softmax calibration requires pipeline-parallel size 1 because raw count records are merged across tensor-parallel head shards. +- Skip-softmax calibration requires pipeline- and data-parallel size 1 because raw count records align only across tensor-parallel head shards; data-parallel replicas serve different requests. - `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`. Checkpoints with a calibrated `decode` `threshold_scale_factor` are rejected at install under a FULL decode CUDA graph mode (including vLLM's default `FULL_AND_PIECEWISE`): the captured graph would replay one request's stale threshold. - Sparse-only installs validate engine-level compatibility like quantized installs do: decode context parallelism, DBO, speculative decoding, and FULL mixed-batch CUDA graphs are rejected (prefix caching remains supported, per the bullet above). diff --git a/examples/vllm_serve/calibrate_sparse_attn.py b/examples/vllm_serve/calibrate_sparse_attn.py index 316e8a0d9c4..b132683981f 100644 --- a/examples/vllm_serve/calibrate_sparse_attn.py +++ b/examples/vllm_serve/calibrate_sparse_attn.py @@ -32,6 +32,7 @@ Usage: python calibrate_sparse_attn.py \ + --calib_data_dir \ --target_sparse_ratio 0.5 \ --decode_tokens 32 \ --update_checkpoint_config @@ -89,6 +90,8 @@ def _engine_kwargs(value: str) -> dict: ) if kwargs.get("pipeline_parallel_size", 1) != 1: raise argparse.ArgumentTypeError("pipeline_parallel_size must be 1 for calibration") + if kwargs.get("data_parallel_size", 1) != 1: + raise argparse.ArgumentTypeError("data_parallel_size must be 1 for calibration") return kwargs @@ -128,6 +131,30 @@ def _load_prompts(llm, args) -> list[str]: return prompts +def _preflight_prompt_inputs(args, parser: argparse.ArgumentParser) -> list[str] | None: + """Validate prompt sources before the vLLM engine is initialized.""" + if args.prompts_file is not None: + try: + return _load_prompts(None, args) + except (OSError, ValueError) as err: + parser.error(str(err)) + if args.calib_data_dir is None: + parser.error( + "the default RULER tasks require --calib_data_dir; pass --prompts_file " + "to supply custom prompts instead" + ) + data_dir = Path(args.calib_data_dir) + if not data_dir.is_dir(): + parser.error(f"--calib_data_dir {args.calib_data_dir!r} is not a directory") + essays_dir = data_dir / "essays" + if not essays_dir.is_dir() or next(essays_dir.glob("*.txt"), None) is None: + parser.error( + f"--calib_data_dir {args.calib_data_dir!r} must contain essays/*.txt; " + "run examples/llm_sparsity/attention_sparsity/download_ruler_data.sh first" + ) + return None + + 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" @@ -265,6 +292,9 @@ def main(): f"containing config.json; {args.model!r} has none" ) + # Custom prompts do not need a tokenizer, so read them eagerly as well. + prompts = _preflight_prompt_inputs(args, parser) + # 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: @@ -303,7 +333,8 @@ def main(): llm = LLM(**llm_kwargs) # Built after engine init so the RULER builder reuses the engine's tokenizer. - prompts = _load_prompts(llm, args) + if prompts is None: + prompts = _load_prompts(llm, args) trials = list(DEFAULT_THRESHOLD_TRIALS) n_layers = llm.collective_rpc("sparse_calib_enable", args=(trials,))[0] diff --git a/modelopt/torch/kernels/common/attention/triton_fa.py b/modelopt/torch/kernels/common/attention/triton_fa.py index 6f7b83cdf08..1984ec5e336 100644 --- a/modelopt/torch/kernels/common/attention/triton_fa.py +++ b/modelopt/torch/kernels/common/attention/triton_fa.py @@ -94,6 +94,27 @@ def _load_qdq_helpers() -> None: _MEASURE_NUM_STAGES = 1 _MEASURE_NUM_WARPS = 4 +# Serving keeps the calibrated KV-tile granularity but does not inherit the +# deliberately conservative measurement schedule. A one-token decode has only +# one valid Q row, so padding-row masking makes its skip decision independent +# of BLOCK_M; use the smallest dense-autotune Q tile there. +_SKIP_SERVE_DECODE_BLOCK_M = 16 + +_SKIP_SERVE_CONFIGS = [ + triton.Config({}, num_stages=num_stages, num_warps=num_warps) + for num_stages in (1, 2, 3) + for num_warps in (4, 8) +] + + +def _skip_tile_resource_error(q_dtype, err) -> RuntimeError: + return RuntimeError( + "skip-softmax requires the fixed 128-wide KV calibration tile, " + f"which exceeds this GPU's shared memory for {q_dtype} inputs ({err}). " + "Use fp16/bf16 inputs or a device with more shared memory; re-tiling KV " + "would change the calibrated sparsity contract." + ) + # --------------------------------------------------------------------------- # Paged KV cache helpers @@ -510,6 +531,20 @@ def _attn_fwd( tl.store(Out + o_ptrs, acc, mask=(q_pos[:, None] < seq_len_q) & d_mask[None, :]) +# Serving autotunes only the execution schedule. BLOCK_M/BLOCK_N remain launch +# arguments, so tuning cannot change which attention tiles the calibrated +# threshold skips. Measurement uses _attn_fwd.fn directly because its counters +# must not be incremented by autotune trials. +_attn_fwd_skip_serve = triton.autotune( + configs=( + _SKIP_SERVE_CONFIGS[:1] + if "PYTEST_VERSION" in __import__("os").environ + else _SKIP_SERVE_CONFIGS + ), + key=["N_CTX", "HEAD_DIM", "Q_IS_FP32", "IS_PAGED"], +)(_attn_fwd.fn) + + # --------------------------------------------------------------------------- # Backward kernels # --------------------------------------------------------------------------- @@ -1040,39 +1075,44 @@ def grid(META): # illegal memory access. with torch.cuda.device(q.device): if apply_skip: - # Fixed-tile launches, bypassing autotune: + # Fixed skip-decision geometry: # - 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. + # bypass autotuning to avoid repeated candidate trials. + # - Active prefill skip-softmax: the tile-skip decision depends + # on the (BLOCK_M, BLOCK_N) geometry, and thresholds are + # calibrated at 128x128 granularity (attention_calibrate and + # flash_skip_softmax both use 128x128 blocks). Decode has one + # valid Q row, so its decision is BLOCK_M-invariant and can use + # a 16x128 compute tile. BLOCK_N remains fixed for both phases. # - # 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. + # Serving autotunes only num_warps/num_stages while holding the + # decision geometry fixed. + block_m = ( + _MEASURE_BLOCK_M + if do_measure or max_input_len > 1 + else _SKIP_SERVE_DECODE_BLOCK_M + ) 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=_MEASURE_BLOCK_M, - BLOCK_N=_MEASURE_BLOCK_N, - num_warps=_MEASURE_NUM_WARPS, - num_stages=_MEASURE_NUM_STAGES, - ) + # here always uses the calibrated 128-wide KV granularity. + if do_measure: + _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, + ) + else: + _attn_fwd_skip_serve[grid]( + *fwd_args, + **fwd_kwargs, + BLOCK_M=block_m, + BLOCK_N=_MEASURE_BLOCK_N, + ) 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 + raise _skip_tile_resource_error(q.dtype, err) from err else: _attn_fwd[grid]( *fwd_args, 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 e0bd2250d87..a8d6d34b59a 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_calibration.py @@ -234,13 +234,13 @@ def build_sparse_attention_config( Non-skip groups from ``existing_config`` (e.g. exported N:M ``sparse_softmax`` metadata) are preserved after the skip group; an existing ``skip_softmax`` group is replaced by the new calibration, - carrying over its layer policy (``ignore`` — layers deliberately kept - dense — and ``initial_disabled_steps``): recalibration replaces the - fitted thresholds, not which layers the export sparsifies. + carrying over its layer policy (``targets``, ``ignore`` — layers + deliberately kept dense — and ``initial_disabled_steps``): recalibration + replaces the fitted thresholds, not which layers the export sparsifies. """ - # target_sparsity covers only fitted phases: claiming a target for a phase - # without calibrated (a, b) would advertise sparsity the serving path - # silently serves dense (it needs the per-phase scale factors). + # Keep the emitted metadata scoped to the phases actually fitted. The + # serving loader may default an absent target, but a missing per-phase + # threshold_scale_factor still keeps that phase dense. target_sparsity_by_phase = { phase: value for phase, value in _normalize_target_sparsity(target_sparsity).items() @@ -249,7 +249,6 @@ def build_sparse_attention_config( skip_group: dict[str, Any] = { "algorithm": "skip_softmax", - "targets": ["Attention"], "threshold_scale_factor": export_threshold_scale_factor(calibration_params), "target_sparsity": target_sparsity_by_phase, } @@ -265,7 +264,7 @@ def build_sparse_attention_config( # Keep the replaced group's layer policy: dropping ``ignore`` # would sparsify layers the original export deliberately kept # dense (e.g. first/last blocks). - for key in ("ignore", "initial_disabled_steps"): + for key in ("targets", "ignore", "initial_disabled_steps"): if key in group and key not in skip_group: skip_group[key] = group[key] else: @@ -273,6 +272,8 @@ def build_sparse_attention_config( for idx, group in enumerate(preserved, start=1): config_groups[f"group_{idx}"] = group + skip_group.setdefault("targets", ["Attention"]) + result: dict[str, Any] = { "config_groups": config_groups, "producer": export_config_producer(), diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 7b4de426b9d..b7a2b7fd370 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -821,7 +821,7 @@ def build(*args, **kwargs): if num_decodes and num_prefills: # Prefer the host-resident copy the runner may already carry; # fall back to one device->host copy per mixed-batch build. - seq_lens_cpu = getattr(common, "_seq_lens_cpu", None) + seq_lens_cpu = getattr(common, "seq_lens_cpu", None) if seq_lens_cpu is None: seq_lens_cpu = common.seq_lens.cpu() max_seq_len_decode = int(seq_lens_cpu[:num_decodes].max()) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py index 51ca1759c65..99ffd4a868f 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py @@ -15,6 +15,7 @@ """Install ModelOpt attention transforms into a loaded vLLM model.""" +import fnmatch import importlib from collections import Counter from collections.abc import Mapping @@ -113,6 +114,31 @@ def _model_config(model_runner): return getattr(getattr(model_runner, "vllm_config", None), "model_config", None) +def _calibration_ignore_patterns(model_runner) -> tuple[str, ...]: + """Return the existing skip-softmax layer exclusions, if any.""" + hf_config = getattr(_model_config(model_runner), "hf_config", None) + sparse_meta = getattr(hf_config, "sparse_attention_config", None) + if not isinstance(sparse_meta, dict): + return () + groups = sparse_meta.get("config_groups") + if not isinstance(groups, dict): + return () + + patterns = [] + for group in groups.values(): + if not isinstance(group, dict) or group.get("algorithm") != "skip_softmax": + continue + ignore = group.get("ignore", ()) + if isinstance(ignore, list | tuple): + patterns.extend(name for name in ignore if isinstance(name, str)) + return tuple(patterns) + + +def _is_calibration_ignored(name: str, patterns: tuple[str, ...]) -> bool: + """Match exclusions exactly as the checkpoint serving loader does.""" + return any(fnmatch.fnmatch(name, f"*{pattern}*") for pattern in patterns) + + def _resolve_sparse_config(model_runner, sparse_cfg) -> tuple[dict | None, str | None]: if sparse_cfg is None: return None, None @@ -554,8 +580,9 @@ def _attention_quant_error(module) -> str | 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 + Swaps the backend-matched ModelOpt adapter onto each attention layer that + is not excluded by the checkpoint's existing skip-softmax ``ignore`` + policy 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 @@ -565,16 +592,18 @@ def install_vllm_skip_softmax_calibration(model_runner) -> VllmAttentionInstallR 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, pipeline-parallel size 1, no active attention + model and KV-cache dtypes, pipeline- and data-parallel size 1, 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) + ignore_patterns = _calibration_ignore_patterns(model_runner) candidates = [ (name, module) for name, module in model.named_modules() if isinstance(module, _VLLM_ATTENTION) + and not _is_calibration_ignored(name, ignore_patterns) ] _require_supported_vllm() @@ -582,13 +611,18 @@ def install_vllm_skip_softmax_calibration(model_runner) -> VllmAttentionInstallR parallel = getattr(getattr(model_runner, "vllm_config", None), "parallel_config", None) if getattr(parallel, "pipeline_parallel_size", 1) != 1: errors.append("pipeline_parallel_size must be 1 for skip-softmax calibration") + if getattr(parallel, "data_parallel_size", 1) != 1: + errors.append( + "data_parallel_size must be 1 for skip-softmax calibration: data-parallel " + "replicas serve disjoint requests, so per-rank count records do not align" + ) 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") + errors.append("no attention layers were found after applying the checkpoint ignore policy") plans = [] for name, module in candidates: diff --git a/tests/examples/vllm_serve/test_calibrate_sparse_attn.py b/tests/examples/vllm_serve/test_calibrate_sparse_attn.py index ab8bec47a9f..6b346915c64 100644 --- a/tests/examples/vllm_serve/test_calibrate_sparse_attn.py +++ b/tests/examples/vllm_serve/test_calibrate_sparse_attn.py @@ -45,6 +45,7 @@ def calibration_driver(monkeypatch): ["--engine_kwargs", '{"enforce_eager": false}'], ["--engine_kwargs", '{"enable_prefix_caching": true}'], ["--engine_kwargs", '{"pipeline_parallel_size": 2}'], + ["--engine_kwargs", '{"data_parallel_size": 2}'], ], ) def test_parser_rejects_invalid_inputs_before_engine_start(calibration_driver, flags): @@ -67,6 +68,40 @@ def test_load_prompts_reads_nonempty_lines(calibration_driver, tmp_path): assert calibration_driver._load_prompts(None, args) == ["first prompt", "second prompt"] +@pytest.mark.parametrize("prompts_contents", [None, "\n\n"]) +def test_preflight_rejects_invalid_prompts_before_engine_start( + calibration_driver, tmp_path, prompts_contents +): + prompts_file = tmp_path / "prompts.txt" + if prompts_contents is not None: + prompts_file.write_text(prompts_contents) + parser = calibration_driver._build_parser() + args = parser.parse_args(["/checkpoint", "--prompts_file", str(prompts_file)]) + + with pytest.raises(SystemExit): + calibration_driver._preflight_prompt_inputs(args, parser) + + +def test_preflight_requires_ruler_data_before_engine_start(calibration_driver): + parser = calibration_driver._build_parser() + args = parser.parse_args(["/checkpoint"]) + + with pytest.raises(SystemExit): + calibration_driver._preflight_prompt_inputs(args, parser) + + +def test_preflight_validates_ruler_essay_files_before_engine_start(calibration_driver, tmp_path): + parser = calibration_driver._build_parser() + args = parser.parse_args(["/checkpoint", "--calib_data_dir", str(tmp_path)]) + with pytest.raises(SystemExit): + calibration_driver._preflight_prompt_inputs(args, parser) + + essays = tmp_path / "essays" + essays.mkdir() + (essays / "sample.txt").write_text("essay") + assert calibration_driver._preflight_prompt_inputs(args, parser) is None + + def test_existing_sparse_config_reads_only_dict(calibration_driver, tmp_path): checkpoint = tmp_path / "checkpoint" checkpoint.mkdir() 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 6a909011359..49237f69f4d 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 @@ -263,6 +263,37 @@ def test_flashinfer_metadata_builder_patch_stashes_common_metadata( assert actual == expected +def test_flashinfer_metadata_builder_uses_public_host_seq_lens( + monkeypatch, isolated_flashinfer_builder_patch +): + """Mixed metadata must not copy device sequence lengths back to the host.""" + + class DeviceSeqLens: + def cpu(self): + raise AssertionError("unexpected device-to-host sequence-length copy") + + def build(_self, common_attn_metadata): + return SimpleNamespace(num_decodes=1, num_prefills=1) + + monkeypatch.setattr(FlashInferMetadataBuilder, "build", build) + assert patch_flashinfer_metadata_builder() is True + common = SimpleNamespace( + block_table_tensor=torch.zeros(2, 1, dtype=torch.int32), + seq_lens=DeviceSeqLens(), + seq_lens_cpu=torch.tensor([16, 64], dtype=torch.int32), + query_start_loc=torch.tensor([0, 1, 2], dtype=torch.int32), + num_actual_tokens=2, + max_query_len=1, + max_seq_len=64, + causal=True, + ) + + metadata = FlashInferMetadataBuilder.build(object(), common) + + assert metadata._modelopt_max_seq_len_decode == 16 + assert metadata._modelopt_max_seq_len_prefill == 64 + + def test_select_and_clone_flashinfer_impl_preserves_runtime_state( isolated_flashinfer_builder_patch, ): 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 4774fdf891c..ac4b163185b 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 @@ -107,6 +107,34 @@ def test_rejects_pipeline_parallelism(self): with pytest.raises(NotImplementedError, match="pipeline_parallel_size must be 1"): vllm_runtime.install_vllm_skip_softmax_calibration(runner) + def test_rejects_data_parallelism(self): + runner = _model_runner(nn.ModuleDict({"attn": _bare_attention()})) + runner.vllm_config.parallel_config.data_parallel_size = 2 + with pytest.raises(NotImplementedError, match="data_parallel_size must be 1"): + vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + def test_excludes_checkpoint_ignored_layers_from_measurement(self): + ignored = _bare_attention() + included = _bare_attention() + ignored_impl = ignored.impl + metadata = { + "config_groups": { + "group_0": { + "algorithm": "skip_softmax", + "ignore": ["a_attn"], + } + } + } + runner = _model_runner( + nn.ModuleDict({"a_attn": ignored, "b_attn": included}), sparse_metadata=metadata + ) + + report = vllm_runtime.install_vllm_skip_softmax_calibration(runner) + + assert report.installed_layers == ("b_attn",) + assert ignored.impl is ignored_impl + assert isinstance(included.impl, ModelOptSparseAttentionImpl) + def test_rejects_active_attention_quantizers_atomically(self): quantized = _bare_attention() quantized.q_bmm_quantizer = SimpleNamespace(is_enabled=True) diff --git a/tests/unit/torch/kernels/common/attention/test_triton_fa.py b/tests/unit/torch/kernels/common/attention/test_triton_fa.py index 13ce5da55ef..5985fce4a2b 100644 --- a/tests/unit/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/unit/torch/kernels/common/attention/test_triton_fa.py @@ -99,6 +99,10 @@ def test_forward_uses_minimal_shared_autotune_configs(): ] assert triton_fa._attn_fwd.keys == ["N_CTX", "HEAD_DIM", "Q_IS_FP32", "P_QDQ", "V_QDQ"] + assert {(config.num_stages, config.num_warps) for config in triton_fa._SKIP_SERVE_CONFIGS} == { + (stages, warps) for stages in (1, 2, 3) for warps in (4, 8) + } + assert triton_fa._attn_fwd_skip_serve.keys == ["N_CTX", "HEAD_DIM", "Q_IS_FP32", "IS_PAGED"] @pytest.mark.parametrize( @@ -148,6 +152,7 @@ def test_forward_measurement_uses_one_fixed_launch(monkeypatch): kernel = _ForbiddenKernel() kernel.fn = _CapturingKernel() monkeypatch.setattr(triton_fa, "_attn_fwd", kernel) + monkeypatch.setattr(triton_fa, "_attn_fwd_skip_serve", _ForbiddenKernel()) monkeypatch.setattr(triton_fa.torch.cuda, "device", lambda _device: nullcontext()) monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) monkeypatch.setattr(triton_fa, "_load_qdq_helpers", lambda: None) @@ -170,6 +175,51 @@ def test_forward_measurement_uses_one_fixed_launch(monkeypatch): assert kernel.fn.kwargs["num_warps"] == 4 +@pytest.mark.parametrize( + ("q_len", "kv_len", "expected_block_m"), + [(129, 129, 128), (1, 256, 16)], +) +def test_forward_skip_serving_keeps_kv_tile_and_uses_phase_q_tile( + monkeypatch, q_len, kv_len, expected_block_m +): + """Serving keeps BLOCK_N=128 while decode avoids padded 128-row Q work.""" + pytest.importorskip("triton") + + from modelopt.torch.kernels.common.attention import triton_fa + + kernel = _ForbiddenKernel() + kernel.fn = _ForbiddenKernel() + serving_kernel = _CapturingKernel() + monkeypatch.setattr(triton_fa, "_attn_fwd", kernel) + monkeypatch.setattr(triton_fa, "_attn_fwd_skip_serve", serving_kernel) + monkeypatch.setattr(triton_fa.torch.cuda, "device", lambda _device: nullcontext()) + monkeypatch.setattr(triton_fa, "_load_sparsity_helpers", lambda: None) + monkeypatch.setattr(triton_fa, "_load_qdq_helpers", lambda: None) + + q = torch.empty(q_len, 2, 16) + k = torch.empty(kv_len, 1, 16) + v = torch.empty_like(k) + starts = torch.tensor([0], dtype=torch.int32) + lengths = torch.tensor([q_len], dtype=torch.int32) + kv_lengths = torch.tensor([kv_len], dtype=torch.int32) + + triton_fa.attention( + q, + k, + v, + starts, + lengths, + q_len, + b_start_loc_k=starts, + b_seq_len_k=kv_lengths, + max_input_len_k=kv_len, + skip_softmax_threshold=0.1, + ) + + assert serving_kernel.kwargs["BLOCK_M"] == expected_block_m + assert serving_kernel.kwargs["BLOCK_N"] == 128 + + @pytest.mark.parametrize( "qdq_kwargs", [{"p_qdq": "nvfp4"}, {"p_qdq": "fp8"}, {"v_qdq": "nvfp4", "v_qdq_amax": 1.0}], diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_calibration.py index 7ff174952b3..f7f202108c3 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 @@ -185,11 +185,12 @@ def test_target_sparsity_covers_only_fitted_phases(self): assert "decode" not in group["threshold_scale_factor"] def test_replaced_skip_group_keeps_layer_policy(self): - """Recalibration replaces thresholds but keeps ignore/initial_disabled_steps.""" + """Recalibration replaces thresholds but keeps the existing layer policy.""" existing = { "config_groups": { "group_0": { "algorithm": "skip_softmax", + "targets": ["LlamaAttention"], "ignore": ["model.layers.0.self_attn"], "initial_disabled_steps": 4, "threshold_scale_factor": {"prefill": {"a": 1.0, "b": 1.0}}, @@ -198,6 +199,7 @@ def test_replaced_skip_group_keeps_layer_policy(self): } config = build_sparse_attention_config(self._PARAMS, 0.5, existing_config=existing) group = config["config_groups"]["group_0"] + assert group["targets"] == ["LlamaAttention"] assert group["ignore"] == ["model.layers.0.self_attn"] assert group["initial_disabled_steps"] == 4 assert group["threshold_scale_factor"]["prefill"] == {"a": 7.9, "b": 8.6} From b84cc4c70a5b0c17f374cd03990679a2adbe77a6 Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Fri, 28 Aug 2026 23:37:53 -0700 Subject: [PATCH 14/14] test: make vLLM calibration cache fixtures portable Signed-off-by: Kai Xu --- .../test_vllm_calibration.py | 29 ++++++++++++++----- 1 file changed, 21 insertions(+), 8 deletions(-) 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 ac4b163185b..41d2141657b 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 @@ -21,7 +21,7 @@ import torch from torch import nn from vllm.config.compilation import CUDAGraphMode -from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl from vllm.v1.attention.backends.flashinfer import FlashInferImpl from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE @@ -287,12 +287,15 @@ def _make_impl(num_heads, head_dim, num_kv_heads): 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.""" + """Scatter per-request K/V lists into the installed backend's paged layout.""" 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) + cache_shape = FlashAttentionBackend.get_kv_cache_shape( + num_blocks, page_size, num_kv_heads, head_dim + ) + kv_cache = torch.zeros(cache_shape, device=device, dtype=dtype) + k_cache, v_cache = attention_plugin._flash_attention_kv_cache_views(kv_cache, head_dim) block_table = torch.zeros(len(seqs_kv), max_blocks, device=device, dtype=torch.int32) g = 0 for b, (k, v) in enumerate(seqs_kv): @@ -302,7 +305,7 @@ def _paged_cache_for(seqs_kv, num_kv_heads, head_dim, page_size, device, dtype): 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 + return kv_cache, block_table def _sdpa_reference(q, k, v, is_causal): @@ -415,12 +418,17 @@ def __init__(self, impls): {"sample_length": 128, "total_tiles": [8, 8, 8], "skipped_tiles": [1, 3, 5]} ] - def test_rejects_non_logical_cache_shape(self): + def test_rejects_non_logical_cache_shape(self, monkeypatch): num_heads, num_kv_heads, head_dim = 4, 2, 64 impl = _make_impl(num_heads, head_dim, num_kv_heads) enable_calibration([impl], TRIALS) - # A physical HND tensor must be exposed as a logical [blocks, page, heads, dim] view. + # Inject a malformed logical view to exercise the adapter's shape guard. kv_cache = torch.zeros(2, 1, num_kv_heads, 16, head_dim, dtype=torch.bfloat16) + monkeypatch.setattr( + attention_plugin, + "_flash_attention_kv_cache_views", + lambda cache, _head_size: cache.unbind(0), + ) attn_metadata = SimpleNamespace( num_actual_tokens=1, max_query_len=1, @@ -441,11 +449,16 @@ def test_rejects_non_logical_cache_shape(self): output=torch.empty_like(q), ) - def test_rejects_non_16bit_cache(self): + def test_rejects_non_16bit_cache(self, monkeypatch): 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) + monkeypatch.setattr( + attention_plugin, + "_flash_attention_kv_cache_views", + lambda cache, _head_size: cache.unbind(0), + ) attn_metadata = SimpleNamespace( num_actual_tokens=1, max_query_len=1,