# SPDX-FileCopyrightText: Copyright (c) 2024 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. """Conversion and restoration utilities for sparse attention.""" import fnmatch from collections.abc import Callable from typing import Any import torch.nn as nn from modelopt import __version__ as mo_version from modelopt.torch.opt.conversion import ModelLikeModule, ModeloptStateManager from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict from modelopt.torch.utils import atomic_print, get_unwrapped_name from .config import SparseAttentionConfig from .plugins import register_custom_model_plugins_on_the_fly from .sparse_attention import SparseAttentionModule, SparseAttentionRegistry from .utils import get_named_sparse_attention_modules, get_sparse_attention_modules SPARSE_SOFTMAX_METADATA_KEYS = ( "sparsity_n", "sparsity_m", "dense_sink_tokens", "dense_recent_tokens", ) 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. - ``backend="triton"``: registers the Triton kernel with HF and sets ``attn_implementation="modelopt_triton"``. - ``backend="pytorch"`` (default): sets ``attn_implementation="eager"`` so that softmax-patching methods (e.g. skip-softmax) work correctly. FlashAttention and SDPA bypass ``F.softmax``, so eager is required. - ``method="vsa"``: no-op. VSA patches ``F.scaled_dot_product_attention`` directly in ``SparseAttentionModule.forward()``, so no ``attn_implementation`` change is needed. This is called automatically during ``mts.sparsify()`` so users never need to manually set ``attn_implementation``. """ sparse_cfg = config.sparse_cfg if hasattr(config, "sparse_cfg") else {} # Collect methods and backends only from layer configs (identified by having a "method" key). # Other dict entries (e.g. "calibration") are not layer configs. layer_cfgs = [v for v in sparse_cfg.values() if isinstance(v, dict) and "method" in v] methods = {v.get("method") for v in layer_cfgs} backends = {v.get("backend", "pytorch") for v in layer_cfgs} # VSA patches F.scaled_dot_product_attention directly — it does not change # attn_implementation. Skip the rest for VSA-only configs. if methods == {"vsa"}: return # Reject mixed VSA + non-VSA configs (VSA patches SDPA globally per-module, # while softmax-patching methods need attn_implementation="eager"). non_vsa_methods = methods - {"vsa"} if "vsa" in methods and non_vsa_methods: raise ValueError( f"Cannot mix VSA with other sparse attention methods ({non_vsa_methods}). " f"VSA patches F.scaled_dot_product_attention, which is incompatible " f"with softmax-patching or triton methods." ) model_config = getattr(model, "config", None) if "triton" in backends and "pytorch" in backends: raise ValueError( "Mixed backends ('triton' and 'pytorch') in the same model are not " "supported. All sparse attention layers must use the same backend." ) if "triton" in backends: from modelopt.torch.kernels.sparsity.attention import register_triton_attention if register_triton_attention is None: raise ImportError( "Triton backend requires 'triton' package. Install with: pip install triton" ) if not register_triton_attention(): raise RuntimeError( "Failed to register Triton attention with HuggingFace. " "Check that your transformers version supports ALL_ATTENTION_FUNCTIONS." ) # Set attn_implementation so HF dispatches to the Triton kernel. # HF's ALL_ATTENTION_FUNCTIONS is checked at forward time, not construction # time, so this works even after the model is already loaded. if model_config is not None: model_config._attn_implementation = "modelopt_triton" elif model_config is not None: # For pytorch backend, force eager for softmax patching. model_config._attn_implementation = "eager" def is_attn_sparsified(model: nn.Module) -> bool: """Check if a model has sparse attention applied. Similar to quantization's is_quantized for API consistency. Args: model: Model to check Returns: True if model contains any SparseAttentionModule instances """ return any(isinstance(module, SparseAttentionModule) for module in model.modules()) def _register_diffusers_backends_if_needed(model: nn.Module) -> None: """Register diffusers/LTX Triton attention backends if the model needs them. Called before plugin registration so that the backends are available when ``SparseAttentionModule.forward()`` activates the skip-softmax context. """ import contextlib # Register the diffusers Triton backend if the model is a diffusers ModelMixin try: from diffusers.models.modeling_utils import ModelMixin if isinstance(model, ModelMixin): from modelopt.torch.kernels.sparsity.attention import ( register_diffusers_triton_attention, ) if register_diffusers_triton_attention is not None: register_diffusers_triton_attention() except (ImportError, Exception): pass # Patch ltx_core Attention modules if present (independent of diffusers) try: from modelopt.torch.kernels.sparsity.attention import register_ltx_triton_attention except (ImportError, RuntimeError): return if register_ltx_triton_attention is not None: with contextlib.suppress(Exception): register_ltx_triton_attention(model) def convert_to_sparse_attention_model( model: ModelLikeModule, config: SparseAttentionConfig ) -> ConvertReturnType: """Convert model to use sparse attention. Args: model: Model to convert config: Sparse attention configuration Returns: Tuple of (converted_model, metadata) """ # Initialize the true module if necessary model = model.init_modellike() if isinstance(model, ModelLikeModule) else model # Register diffusers backends for diffusion models _register_diffusers_backends_if_needed(model) # Set the correct attn_implementation for the chosen backend _set_attn_implementation(model, config) # Apply custom model plugins register_custom_model_plugins_on_the_fly(model) # Replace attention modules with sparse versions replace_sparse_attention_modules(model, version=ModeloptStateManager(model).state_version) # Apply configuration to sparse attention modules sparse_cfg = config.sparse_cfg if hasattr(config, "sparse_cfg") else {} set_sparse_attention_by_cfg(model, sparse_cfg) # Create metadata metadata = {} update_sparse_attention_metadata(model, config, metadata) return model, metadata def replace_sparse_attention_modules(model: nn.Module, version=None): """Replace regular attention modules with sparse attention modules. Recursively replace all attention modules in the model with their sparse attention counterparts. Args: model: Model to process version: State version for tracking (optional) """ # Recursively replace modules _replace_sparse_attention_modules(model, version=version) # Count and report replaced modules replaced_count = len(get_sparse_attention_modules(model)) if replaced_count > 0: print(f"Inserted {replaced_count} sparse attention modules") def _replace_sparse_attention_modules(model: nn.Module, version=None): """Helper function for replace_sparse_attention_modules.""" for name, child in model.named_children(): if type(child) in SparseAttentionRegistry: # REPLACE on the parent (model), not on child sparse_module = SparseAttentionRegistry.convert(child) setattr(model, name, sparse_module) # Now recurse into whichever module is now at `model.name` _replace_sparse_attention_modules(getattr(model, name), version=version) def set_sparse_attention_by_cfg(model: nn.Module, sparse_cfg: dict): """Apply sparse attention configuration to model. Similar to quantization's set_quantizer_by_cfg. Args: model: Model with sparse attention modules sparse_cfg: Sparse configuration dictionary mapping patterns to attributes """ sparse_cfg = sparse_cfg.copy() # Apply default first if exists if "default" in sparse_cfg: set_sparse_attention_attribute(model, "*", sparse_cfg["default"]) sparse_cfg.pop("default") # Apply pattern-specific configs for pattern, cfg in sparse_cfg.items(): set_sparse_attention_attribute(model, pattern, cfg) def set_sparse_attention_attribute( model: nn.Module, wildcard_or_filter: str | Callable, attribute_cfg: dict[str, Any], ): """Set sparse attention attributes for modules matching pattern. Similar to quantization's set_quantizer_attributes_partial. Args: model: Model to configure wildcard_or_filter: Pattern to match module names attribute_cfg: Attributes to apply (must include 'method') """ # Filter out model-level configs that shouldn't be passed to modules module_cfg = {k: v for k, v in attribute_cfg.items() if k != "calibration"} for name, module in get_named_sparse_attention_modules(model): # Check pattern match matched = False if isinstance(wildcard_or_filter, str): matched = fnmatch.fnmatch(name, wildcard_or_filter) elif callable(wildcard_or_filter): matched = wildcard_or_filter(name) else: raise NotImplementedError(f"Unsupported type {type(wildcard_or_filter)}") if matched: # Apply config using the same method as TensorQuantizer module.set_from_attribute_config(module_cfg) def restore_sparse_attention_model( model: ModelLikeModule, config: SparseAttentionConfig, metadata: MetadataDict ) -> nn.Module: """Restore sparse attention model from saved state. Args: model: Model to restore config: Sparse attention configuration metadata: Saved metadata Returns: Restored model """ # Convert to sparse attention model model, _ = convert_to_sparse_attention_model(model, config) # Restore sparse attention state from metadata if "sparse_attention_state" in metadata: restore_sparse_attention_state(model, metadata["sparse_attention_state"]) return model def restore_sparse_attention_state(model: nn.Module, state_dict: dict[str, Any]): """Restore sparse attention state from state dict. Args: model: Model with sparse attention modules state_dict: Saved state dictionary """ for name, module in get_named_sparse_attention_modules(model): module_name = get_unwrapped_name(name, model) if module_name in state_dict: module.set_from_modelopt_state(state_dict[module_name]) def update_sparse_attention_metadata( model: nn.Module, config: SparseAttentionConfig, metadata: MetadataDict ) -> None: """Update metadata with sparse attention state. Args: model: Model with sparse attention config: Configuration used metadata: Metadata dict to update """ metadata["sparse_attention_state"] = { get_unwrapped_name(name, model): module.get_modelopt_state() for name, module in get_named_sparse_attention_modules(model) } metadata["sparse_attention_config"] = ( config.model_dump() if hasattr(config, "model_dump") else vars(config) ) def get_sparse_attention_mode_state( model: nn.Module, ) -> tuple[SparseAttentionConfig, MetadataDict] | None: """Return the config and metadata of the latest ``sparse_attention`` mode applied to ``model``.""" # Wrappers and containers of converted models have no mode state of their own. if not hasattr(model, ModeloptStateManager._state_key): return None mode_states = [ (config, metadata) for mode, config, metadata in ModeloptStateManager(model).modes_with_states() if str(mode) == "sparse_attention" ] return mode_states[-1] if mode_states else None def _get_sparse_softmax_export_config(module: SparseAttentionModule) -> dict[str, Any] | None: """Return N:M sparse-softmax export metadata for one sparse attention module.""" method_config = getattr(module, "_method_config", {}) method_name = module._sparse_method_instance.name should_export = method_name == "triton_sparse_softmax" or method_config.get( "export_sparse_softmax", False ) if not should_export: return None sparse_softmax_config = { key: method_config[key] for key in SPARSE_SOFTMAX_METADATA_KEYS if key in method_config } if sparse_softmax_config.get("sparsity_n", 0) <= 0: return None return sparse_softmax_config def export_sparse_attention_config( model: nn.Module, *, for_diffusion: bool = False ) -> dict[str, Any] | None: """Extract sparse attention config for export to config.json. Extracts calibrated skip-softmax parameters and N:M sparse-softmax metadata from sparse attention modules. The exported config allows computing threshold at runtime: scale_factor = a * exp(b * target_sparsity) threshold = scale_factor / seqlen Args: model: Model with sparse attention applied for_diffusion: Emit the single-phase schema read by the TRT-LLM VisualGen runtime: one ``threshold_scale_factor.coefficients`` block and a scalar ``target_sparsity``. The default is the LLM schema, with per-phase (``prefill``/``decode``) blocks and a ``{phase: ratio}`` dict. Returns: Dictionary with sparse attention config for HuggingFace config.json export. Returns None if no exportable sparse attention metadata is found. Example output:: { "config_groups": { "group_0": { "algorithm": "skip_softmax", "targets": ["LlamaAttention"], "threshold_scale_factor": { "formula": "a * exp(b * target_sparsity)", "prefill": {"a": 7.93, "b": 8.61}, "decode": {"a": 0.12, "b": 9.85}, }, "target_sparsity": {"prefill": 0.5}, }, "group_1": { "algorithm": "sparse_softmax", "targets": ["LlamaAttention"], "sparsity_n": 2, "sparsity_m": 4, "dense_sink_tokens": 0, "dense_recent_tokens": 64, }, }, "producer": {"name": "modelopt", "version": "0.37.0"}, } """ # Collect sparse attention module info calibration_params = None target_sparse_ratio = None sparse_softmax_config = None target_classes: set[str] = set() disabled_layer_names: list[str] = [] initial_disabled_steps = 0 for name, module in get_named_sparse_attention_modules(model): # Get the original wrapped module's class name if hasattr(module, "get_original_cls_by_level"): original_cls = module.get_original_cls_by_level(level=0) if original_cls is not None: target_classes.add(original_cls.__name__) # Record layers kept dense (e.g. cross-attention, first/last blocks) so the # deployment side sparsifies the same subset that was calibrated, rather than # every instance of the target class. if not module.is_enabled: disabled_layer_names.append(get_unwrapped_name(name, model)) continue # Get calibration params from first enabled module that has them if calibration_params is None: calibration_params = getattr(module._sparse_method_instance, "calibration_params", None) if target_sparse_ratio is None: target_sparse_ratio = getattr( module._sparse_method_instance, "target_sparse_ratio", None ) if sparse_softmax_config is None: sparse_softmax_config = _get_sparse_softmax_export_config(module) # A single run-wide value: take it from the first enabled module that sets it # (same harvesting pattern as calibration_params / target_sparse_ratio above). if not initial_disabled_steps: initial_disabled_steps = module._method_config.get("initial_disabled_steps", 0) if calibration_params is None and sparse_softmax_config is None: return None targets = sorted(target_classes) if target_classes else ["Attention"] config_groups: dict[str, Any] = {} group_idx = 0 # Each algorithm's parameters live inside its own config group. if calibration_params is not None: skip_group: dict[str, Any] = { "algorithm": "skip_softmax", "targets": targets, } if disabled_layer_names: skip_group["ignore"] = disabled_layer_names if initial_disabled_steps: 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. 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 if for_diffusion: # Diffusion calibrates a single (prefill) phase; VisualGen rejects phase-keyed dicts. tsf = skip_group["threshold_scale_factor"] phases = {p: tsf.pop(p) for p in ("prefill", "decode") if p in tsf} if phases: tsf["coefficients"] = phases.get("prefill") or next(iter(phases.values())) ratio = skip_group.get("target_sparsity") if isinstance(ratio, dict): scalar = ( ratio["prefill"] if "prefill" in ratio else next(iter(ratio.values()), None) ) if scalar is None: del skip_group["target_sparsity"] # VisualGen needs a scalar else: skip_group["target_sparsity"] = scalar config_groups[f"group_{group_idx}"] = skip_group group_idx += 1 if sparse_softmax_config is not None: # N:M sparse-softmax params live in this group. sparse_group: dict[str, Any] = { "algorithm": "sparse_softmax", "targets": targets, } sparse_group.update(sparse_softmax_config) config_groups[f"group_{group_idx}"] = sparse_group return { "config_groups": config_groups, "producer": export_config_producer(), } def disable_sparse_attention(model: nn.Module, wildcard_or_filter_func: str | Callable): """Disable sparse attention for matching modules. Similar to mtq.disable_quantizer for API consistency. Args: model: Model with sparse attention applied wildcard_or_filter_func: Wildcard string or filter function to match module names. For example: "*lm_head*", "*layer_0*", etc. Example: >>> import modelopt.torch.sparsity as mts >>> import modelopt.torch.sparsity.attention_sparsity as sparse_attn >>> model = mts.sparsify(model, config) >>> # Disable sparse attention for lm_head >>> sparse_attn.disable_sparse_attention(model, "*lm_head*") """ for name, module in get_named_sparse_attention_modules(model): matched = False if isinstance(wildcard_or_filter_func, str): matched = fnmatch.fnmatch(name, wildcard_or_filter_func) elif callable(wildcard_or_filter_func): matched = wildcard_or_filter_func(name) if matched: module.disable() def enable_sparse_attention(model: nn.Module, wildcard_or_filter_func: str | Callable): """Enable sparse attention for matching modules. Similar to mtq.enable_quantizer for API consistency. Args: model: Model with sparse attention applied wildcard_or_filter_func: Wildcard string or filter function to match module names. For example: "*attention*", "*attn*", etc. Example: >>> import modelopt.torch.sparsity as mts >>> import modelopt.torch.sparsity.attention_sparsity as sparse_attn >>> model = mts.sparsify(model, config) >>> # Re-enable sparse attention for all attention modules >>> sparse_attn.enable_sparse_attention(model, "*attention*") """ for name, module in get_named_sparse_attention_modules(model): matched = False if isinstance(wildcard_or_filter_func, str): matched = fnmatch.fnmatch(name, wildcard_or_filter_func) elif callable(wildcard_or_filter_func): matched = wildcard_or_filter_func(name) if matched: module.enable() def _format_threshold(info: dict) -> str: """Format threshold info for display.""" t = info.get("type") if t == "dynamic_calibrated": # Exponential model: threshold = a * exp(b * sparsity) / seqlen params = info.get("calibration_params", {}) target = info.get("target_sparse_ratio", {}) parts = [] for phase in ["prefill", "decode"]: if phase in params: a, b = params[phase]["a"], params[phase]["b"] s = target.get(phase, 0.5) parts.append(f"{phase}: a={a:.4f}, b={b:.2f}, target={s:.0%}") return f"calibrated({', '.join(parts)})" if t == "static": v = info.get("value") if isinstance(v, dict): return f"threshold={v}" return f"threshold={v:.2e}" if isinstance(v, float) else f"threshold={v}" return "threshold=N/A" @atomic_print def print_sparse_attention_summary(model: nn.Module): """Print summary of sparse attention modules in the model. Args: model: Model with sparse attention applied """ sparse_modules = get_named_sparse_attention_modules(model) if not sparse_modules: print("No sparse attention modules found") return enabled = sum(1 for _, m in sparse_modules if m.is_enabled) print(f"Sparse attention: {enabled}/{len(sparse_modules)} modules enabled") # Group by (method, threshold) groups: dict[tuple[str, str], int] = {} for _, module in sparse_modules: if not module.is_enabled: continue method = getattr(module, "_method", "unknown") threshold = _format_threshold(module.get_threshold_info()) groups[(method, threshold)] = groups.get((method, threshold), 0) + 1 for (method, threshold), count in sorted(groups.items()): print(f" {method}: {count} layers, {threshold}")