methods

Sparse attention methods package.

Classes

SparseAttentionMethod

Base class for sparse attention methods.

Functions

get_sparse_method

Get sparse attention method by name and optional version.

register_sparse_method

Decorator to register sparse attention methods with version support.
class SparseAttentionMethod

Bases: ABC

Base class for sparse attention methods.

abstract apply_sparsity(query=None, key=None, value=None, attention_scores=None)

Apply sparsity to attention computation.

Parameters:

  • query (Tensor | None) – Query tensor

  • key (Tensor | None) – Key tensor

  • value (Tensor | None) – Value tensor

  • attention_scores (Tensor | None) – Pre-computed attention scores

Returns:

Tuple of (query, key, value, attention_scores) with sparsity applied

Return type:

tuple[Tensor | None, Tensor | None, Tensor | None, Tensor | None]

abstract property name: str

Method name identifier.

get_sparse_method(name, version=None)

Get sparse attention method by name and optional version.

Parameters:

  • name (str) – Method name to retrieve

  • version (str | None) – Optional version string. If None, uses latest version.

Returns:

Method class

Raises:

ValueError – If method name or version is not registered

Return type:

type[SparseAttentionMethod]

Example

>>> get_sparse_method("flash_skip_softmax")  # Latest version
>>> get_sparse_method("flash_skip_softmax", "v1")  # Specific version
register_sparse_method(name, version='v1')

Decorator to register sparse attention methods with version support.

Parameters:

  • name (str) – Method name to register

  • version (str) – Version string (default: “v1”)

Example:

@register_sparse_method("my_method", version="v3")
class MyMethodV3(SparseAttentionMethod): ...