sparse_attention#

Extensible sparse attention module.

Classes

SparseAttentionModule

Generic sparse attention module wrapper for applying sparsity to attention layers.
class SparseAttentionModule#

Bases: DynamicModule

Generic sparse attention module wrapper for applying sparsity to attention layers.

This module wraps existing attention implementations to add sparse attention capabilities. The activation mechanism is delegated to the configured method via method.get_sparse_context(module), so each method defines how it integrates with the forward pass (e.g. softmax patching, kernel flags).

Forward Flow:#

  1. Check if sparse attention is enabled (pass-through if disabled)

  2. Obtain method-specific context via _sparse_method_instance.get_sparse_context(self)

  3. Run the original forward inside the context

  4. Collect statistics if stats manager is enabled

Attributes:#

_enabled: bool

Whether sparse attention is enabled

_method: str

The sparse attention method to use (e.g., “flash_skip_softmax”)

_method_config: dict

Configuration dictionary for the sparse method (threshold, br, bc, etc.)

_sparse_method_instance: SparseAttentionMethod

Instance of the configured sparse attention method

disable()#

Disable sparse attention for this module.

enable()#

Enable sparse attention for this module.

forward(*args, **kwargs)#

Forward with selected sparse attention method.

  • VSA: patches F.scaled_dot_product_attention to intercept the SDPA call inside the original forward. Cross-attention is skipped.

  • Softmax-patching methods (e.g. flash_skip_softmax): use the context manager path below.

get_modelopt_state()#

Get the module state to save in the ModelOpt checkpoint metadata.

Similar to TensorQuantizer.get_modelopt_state. Calibration results live on the method instance rather than in the state dict, so they are saved here.

Return type:

dict[str, Any]

get_stats()#

Get sparsity statistics from the stats manager.

Returns:

Dictionary with sparsity statistics including ‘average_sparsity’ if available. Returns empty dict if stats manager is not enabled.

Return type:

dict

get_threshold_info()#

Get threshold information from the sparse method instance.

Returns:

Dictionary with threshold information from the sparse method.

Return type:

dict[str, Any]

property is_enabled: bool#

Check if sparse attention is enabled.

set_from_attribute_config(attribute_cfg=None)#

Set sparse attention attributes from configuration.

Similar to TensorQuantizer.set_from_attribute_config.

Parameters:

attribute_cfg (SparseAttentionAttributeConfig | dict | None) – Sparse attention attribute configuration.

set_from_modelopt_state(modelopt_state)#

Restore the state saved by get_modelopt_state().

Similar to TensorQuantizer.set_from_modelopt_state. The method and its config are rebuilt from the mode config before this is called; this restores the runtime state on top of them.

Parameters:

modelopt_state (dict[str, Any])