### Install Example Dependencies

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/README.md

Install the necessary dependencies for running examples, including the package itself.

```bash
pip install .[examples]
```

--------------------------------

### Run Training Example

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/README.md

Execute the training script for the enwik8 language modeling example. It is recommended to log in to wandb first.

```bash
python train.py
```

--------------------------------

### Install Native Sparse Attention PyTorch

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/README.md

Install the library using pip. This command installs the core package.

```bash
pip install native-sparse-attention-pytorch
```

--------------------------------

### Basic Sparse Attention Layer Setup

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Demonstrates how to initialize and use the core SparseAttention layer with common parameters for attention mechanisms.

```python
import torch
from native_sparse_attention_pytorch import SparseAttention

# Create attention layer
attn = SparseAttention(
    dim=512,
    dim_head=64,
    heads=8,
    sliding_window_size=32,
    compress_block_size=4,
    compress_block_sliding_stride=2,
    selection_block_size=4,
    num_selected_blocks=4,
    causal=True
)

# Forward pass
x = torch.randn(2, 256, 512)
output = attn(x)  # Shape: (2, 256, 512)
```

--------------------------------

### Example: Using CompressTransformer with SparseAttention

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Demonstrates how to instantiate and use the CompressTransformer with SparseAttention for a compression task. Requires torch and native_sparse_attention_pytorch.

```python
from native_sparse_attention_pytorch.compress_networks import CompressTransformer
from native_sparse_attention_pytorch import SparseAttention

# Create a transformer-based compressor
compress_net = CompressTransformer(
    num_layers=2,
    dim=512,  # 8 heads * 64 dim_head
    num_heads=8,
    ff_hidden_dim=1024,
    dropout=0.1
)

# Use with SparseAttention
attn = SparseAttention(
    dim=512,
    dim_head=64,
    heads=8,
    sliding_window_size=32,
    compress_block_size=4,
    compress_block_sliding_stride=2,
    selection_block_size=4,
    num_selected_blocks=4,
    compress_mlp=compress_net
)

x = torch.randn(2, 256, 512)
output = attn(x)
```

--------------------------------

### Initialize and Use SparseAttention Layer

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/SparseAttention.md

Demonstrates how to create a SparseAttention layer and perform a training forward pass. Also shows an example of inference with caching.

```python
import torch
from native_sparse_attention_pytorch import SparseAttention

# Create a sparse attention layer
attn = SparseAttention(
    dim=512,
    dim_head=64,
    heads=8,
    sliding_window_size=32,
    compress_block_size=4,
    compress_block_sliding_stride=2,
    selection_block_size=4,
    num_selected_blocks=4,
    causal=True
)

# Training forward pass
x = torch.randn(2, 256, 512)  # batch=2, seq_len=256, dim=512
output = attn(x)  # shape: (2, 256, 512)

# Inference with cache
cache = None
for i in range(100):
    next_token = torch.randn(2, 1, 512)
    output, cache = attn(
        next_token,
        cache=cache,
        return_cache=True
    )
```

--------------------------------

### Install Triton from Nightly Build

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/TritonKernels.md

Provides the command to install the nightly build of the Triton library, which may be necessary for compatibility with the latest features or bug fixes.

```bash
pip install -U --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/ triton-nightly
```

--------------------------------

### GroupedMLP Example

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Demonstrates the use of GroupedMLP for compressing key-value sequences with independent parameters per head. The expand_factor can be tuned for model capacity.

```python
from native_sparse_attention_pytorch.compress_networks import GroupedMLP

mlp = GroupedMLP(
    dim_head=64,
    compress_window_size=4,
    heads=8,
    expand_factor=2.0
)
kv = torch.randn(2, 8, 10, 4, 64)
compressed = mlp(kv)  # (2, 8, 10, 64)
```

--------------------------------

### ConvLinearCompress Example

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Demonstrates the usage of ConvLinearCompress for compressing key-value sequences. Ensure the input tensor matches the expected shape.

```python
from native_sparse_attention_pytorch.compress_networks import ConvLinearCompress

compress = ConvLinearCompress(heads=8, dim_head=64, compress_window_size=4)
kv = torch.randn(2, 8, 10, 4, 64)  # batch, heads, windows, block_size, dim_head
compressed = compress(kv)  # (2, 8, 10, 64)
```

--------------------------------

### Minimal Example of SparseAttention Usage

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/README.md

Demonstrates how to create and use the SparseAttention layer with sample input tensors. Requires PyTorch and the native-sparse-attention-pytorch library.

```python
import torch
from native_sparse_attention_pytorch import SparseAttention

# Create a sparse attention layer
attn = SparseAttention(
    dim=512,
    dim_head=64,
    heads=8,
    sliding_window_size=32,
    compress_block_size=4,
    compress_block_sliding_stride=2,
    selection_block_size=4,
    num_selected_blocks=4,
    causal=True
)

# Process sequences
x = torch.randn(2, 256, 512)  # (batch, seq_len, dim)
output = attn(x)              # (batch, seq_len, dim)
```

--------------------------------

### Tensor Type Hinting Example

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Example of using Jaxtyping-based tensor type hints in a function signature.

```python
def forward(q: Float['b h n d'], k: Float['b h n d']) -> Float['b h n d']:
    ...
```

--------------------------------

### SingleProjection Example

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Illustrates the use of SingleProjection for compressing key-value sequences via linear projection. Setting heads to a value greater than 1 enables per-head projections.

```python
from native_sparse_attention_pytorch.compress_networks import SingleProjection

proj = SingleProjection(dim_head=64, compress_window_size=4, heads=8)
kv = torch.randn(2, 8, 10, 4, 64)
compressed = proj(kv)  # (2, 8, 10, 64)
```

--------------------------------

### Compression Block Layout Example

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Architecture.md

Illustrates the division of an input sequence into blocks with a specified stride for compression. It shows how blocks can overlap and calculates the number of blocks.

```python
Example: seq_len=256, block_size=4, stride=2

Tokens:     [0  1  2  3  4  5  6  7  8  9  ...]
Padding:    [         pad 0-3 ...               ] ← For first block
Block 0:    [X  X  X  X]  (tokens 0-3)
Block 1:    [   X  X  X  X]  (tokens 2-5)        ← 50% overlap
Block 2:    [      X  X  X  X]  (tokens 4-7)
...

num_blocks = ceil((seq_len - block_size) / stride) + 1
           = ceil((256 - 4) / 2) + 1 = 127
```

--------------------------------

### AttentionPool Example

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Shows how to use AttentionPool to compress key-value sequences using softmax-weighted pooling. The input tensor shape is critical for correct operation.

```python
from native_sparse_attention_pytorch.compress_networks import AttentionPool

pool = AttentionPool(dim_head=64, compress_window_size=4)
kv = torch.randn(2, 8, 10, 4, 64)
compressed = pool(kv)  # (2, 8, 10, 64)
```

--------------------------------

### FeedForward Example

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Transformer.md

Instantiates and uses the FeedForward module. Requires the FeedForward function. Input tensor shape is (batch, seq_len, dim).

```python
from native_sparse_attention_pytorch.transformer import FeedForward

ff = FeedForward(dim=512, expansion_factor=4.0)
x = torch.randn(2, 256, 512)
output = ff(x)
```

--------------------------------

### Attention Example

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Transformer.md

Instantiates and uses the Attention module. Requires PyTorch and the Attention class. Input tensor shape is (batch, seq_len, dim).

```python
import torch
from native_sparse_attention_pytorch.transformer import Attention

attn = Attention(dim=512, dim_head=64, heads=8, causal=True)
x = torch.randn(2, 256, 512)
output = attn(x)
```

--------------------------------

### Custom SparseAttention kwargs Configuration

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/configuration.md

Example of how to provide a custom dictionary of keyword arguments to the SparseAttention constructor for fine-grained control over sparse attention mechanisms.

```python
sparse_kwargs = {
    'sliding_window_size': 64,
    'compress_block_size': 8,
    'compress_block_sliding_stride': 4,
    'selection_block_size': 8,
    'num_selected_blocks': 6,
    'num_compressed_mem_kv': 2,
    'use_diff_topk': True,
}

model = Transformer(
    num_tokens=256,
    dim=512,
    depth=12,
    use_sparse_attn=True,
    sparse_attn_kwargs=sparse_kwargs
)
```

--------------------------------

### Install Triton Kernel >= 3.0.0

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/errors.md

If using the Triton kernel, ensure you have Triton version 3.0.0 or higher installed. Update Triton using pip.

```bash
pip install -U triton
# Or from nightly
pip install -U --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/ triton-nightly
```

--------------------------------

### Grouped-Query Attention Configuration

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/README.md

Configure SparseAttention with fewer key/value heads than query heads to reduce memory usage. This example uses 2 KV heads for 8 query heads.

```python
attn = SparseAttention(
    dim=512, dim_head=64,
    heads=8,      # 8 query heads
    kv_heads=2,   # 2 key/value heads (4x reduction in KV)
    ...
)
```

--------------------------------

### Example Type Annotations in Triton Kernels

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/types.md

Comprehensive type annotations used in Triton kernels for native sparse attention. Demonstrates complex type unions for indices and masks.

```python
def native_sparse_attend(
    fq: Float['b qh n d'],           # Query: (batch, query_heads, seq_len, dim)
    fk: Float['b kh n d'],           # Key: (batch, kv_heads, seq_len, dim)
    fv: Float['b kh n d'],           # Value: (batch, kv_heads, seq_len, dim)
    block_size: int,
    selected_block_indices: Int['b qh n sel'] | Int['b kh n sel'],  # Block indices
    fmask: Bool['b qh n sel'] | Bool['b kh n sel'],                 # Selection mask
    sel_scale: Float['b kh n sel'] | Float['b qh n sel'] | None = None,
    ...
) -> Tensor | tuple[Tensor, Tensor]:
```

--------------------------------

### Basic Usage of SparseAttention Module

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/README.md

Demonstrates how to initialize and use the SparseAttention module with sample tokens. Ensure torch is imported.

```python
import torch
from native_sparse_attention_pytorch import SparseAttention

attn = SparseAttention(
    dim = 512,
    dim_head = 64,
    heads = 8,
    sliding_window_size = 2,
    compress_block_size = 4,
    compress_block_sliding_stride = 2,
    selection_block_size = 4,
    num_selected_blocks = 2
)

tokens = torch.randn(2, 31, 512)

attended = attn(tokens)

assert tokens.shape == attended.shape
```

--------------------------------

### SparseAttention Initialization with Triton Kernel

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/TritonKernels.md

Demonstrates how to initialize the SparseAttention module with the Triton kernel optimization enabled. This is the primary way to leverage the performance benefits of the Triton implementation.

```python
from native_sparse_attention_pytorch import SparseAttention

attn = SparseAttention(
    dim=512,
    dim_head=64,
    heads=8,
    sliding_window_size=32,
    compress_block_size=4,
    compress_block_sliding_stride=2,
    selection_block_size=4,
    num_selected_blocks=4,
    causal=True,
    use_triton_kernel=True  # Enable Triton optimization
)

x = torch.randn(2, 256, 512).cuda()
output = attn(x)
```

--------------------------------

### SingleProjection Initialization

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Initializes a SingleProjection module for linear projection-based compression. Can be configured for per-head projection by setting heads > 1.

```python
class SingleProjection(nn.Module):
    def __init__(
        self,
        dim_head: int,
        compress_window_size: int,
        heads: int = 1
    )
```

--------------------------------

### Importing Compression Network Components

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Lists the import statements for various compression network building blocks.

```python
# Compression networks
from native_sparse_attention_pytorch.compress_networks import (
    ConvLinearCompress,
    AttentionPool,
    GroupedMLP,
    SingleProjection,
    CompressTransformer
)
```

--------------------------------

### Get maximum negative value for dtype

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/HelperFunctions.md

Retrieves the largest negative representable number for a given tensor's data type. Used for masking operations.

```python
def max_neg_value(t):
    return -torch.finfo(t.dtype).max
```

--------------------------------

### Integrating SparseAttention with Custom Masks

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/MaskFunctions.md

Demonstrates how to instantiate SparseAttention and use pre-created sliding and fine-grained masks during the forward pass. Ensure masks are created once and reused for efficiency.

```python
from native_sparse_attention_pytorch import SparseAttention
from native_sparse_attention_pytorch.native_sparse_attention import (
    create_sliding_mask,
    create_fine_mask
)

attn = SparseAttention(
    dim=512, dim_head=64, heads=8,
    sliding_window_size=32,
    compress_block_size=4,
    compress_block_sliding_stride=2,
    selection_block_size=4,
    num_selected_blocks=4,
    causal=True
)

seq_len = 256
x = torch.randn(2, seq_len, 512)

# Create masks once
sliding_mask = create_sliding_mask(seq_len, 32, causal=True)
fine_mask_factory = create_fine_mask(seq_len, 4, causal=True)

# Use during forward pass
output = attn(
    x,
    sliding_window_flex_mask=sliding_mask,
    fine_selection_flex_mask=fine_mask_factory
)
```

--------------------------------

### Small Model (Efficient) SparseAttention Configuration

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/configuration.md

Configuration for a small and efficient SparseAttention model, suitable for scenarios prioritizing speed and lower resource usage.

```python
from native_sparse_attention_pytorch import SparseAttention

attn = SparseAttention(
    dim=256,
    dim_head=32,
    heads=8,
    sliding_window_size=16,
    compress_block_size=4,
    compress_block_sliding_stride=2,
    selection_block_size=4,
    num_selected_blocks=2,
    causal=True
)
```

--------------------------------

### Enabling Verbose Triton Compilation Logs

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/errors.md

To get detailed logs during Triton kernel compilation, set the verbose flag to True using `triton.testing.set_verbose(True)`.

```python
import triton
triton.testing.set_verbose(True)
```

--------------------------------

### Medium Model Configuration Preset

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/README.md

A balanced configuration preset for SparseAttention, enabling Triton kernel optimization. Recommended for general use.

```python
SparseAttention(
    dim=512, dim_head=64, heads=8,
    sliding_window_size=32,
    compress_block_size=4, compress_block_sliding_stride=2,
    selection_block_size=4, num_selected_blocks=4,
    causal=True, use_triton_kernel=True
)
```

--------------------------------

### Pack tensor with einops and get inverse function

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/HelperFunctions.md

Packs a tensor using einops and returns a function to unpack it later. Useful for simplifying tensor reshaping operations.

```python
def pack_one_with_inverse(t, pattern):
    packed, ps = pack([t], pattern)
    def inverse(out):
        return unpack(out, ps, pattern)[0]
    return packed, inverse
```

--------------------------------

### File Organization Structure

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/INDEX.md

Illustrates the directory structure of the native-sparse-attention-pytorch project, showing the location of key documentation files and modules.

```markdown
/workspace/home/output/
├── README.md               ← Start here
├── INDEX.md               ← This file
├── API_Overview.md        ← API reference
├── SparseAttention.md     ← Core module
├── Transformer.md         ← Models
├── CompressionNetworks.md ← Compression
├── TritonKernels.md       ← GPU acceleration
├── MaskFunctions.md       ← FlexAttention
├── Architecture.md        ← Internal design
├── configuration.md       ← Parameters
├── HelperFunctions.md     ← Utilities
├── errors.md              ← Debugging
└── types.md               ← Type annotations

```

--------------------------------

### Importing Mask Creation Functions

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Shows how to import functions for creating different types of attention masks.

```python
# Mask creation
from native_sparse_attention_pytorch.native_sparse_attention import (
    create_sliding_mask,
    create_compress_mask,
    create_fine_mask
)
```

--------------------------------

### Validating SparseAttention Configuration with Quick Forward Pass

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/errors.md

Perform a quick forward pass with random data after initializing `SparseAttention` to validate its configuration and ensure it runs without immediate errors.

```python
attn = SparseAttention(...)
x_test = torch.randn(2, 256, 512)
output = attn(x_test)  # Quick validation
```

--------------------------------

### Transformer Model Initialization and Training

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Transformer.md

Initializes a Transformer model with sparse attention enabled and demonstrates a forward pass for training, including loss calculation and backpropagation. Requires PyTorch and the native_sparse_attention_pytorch library.

```python
import torch
from native_sparse_attention_pytorch.transformer import Transformer

model = Transformer(
    num_tokens=256,
    dim=512,
    depth=6,
    heads=8,
    use_sparse_attn=True,
    sparse_attn_kwargs={
        'sliding_window_size': 32,
        'compress_block_size': 4,
        'compress_block_sliding_stride': 2,
        'selection_block_size': 4,
        'num_selected_blocks': 4,
    }
)

# Forward pass with loss
token_ids = torch.randint(0, 256, (4, 512))  # batch of 4, length 512
loss = model(token_ids, return_loss=True)
loss.backward()
```

--------------------------------

### ConvLinearCompress Initialization

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Initializes a ConvLinearCompress module for grouped convolution-based compression. Requires specifying the number of heads, dimension per head, and the window size for compression.

```python
class ConvLinearCompress(nn.Module):
    def __init__(
        self,
        heads: int,
        dim_head: int,
        compress_window_size: int
    )
```

--------------------------------

### Import SparseAttention Class

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/README.md

Import the core SparseAttention module from the library. This is typically the first step when using the library in your project.

```python
from native_sparse_attention_pytorch import SparseAttention
```

--------------------------------

### Large Model Configuration Preset

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/README.md

A comprehensive configuration preset for SparseAttention, featuring high dimensions, more heads, and Triton kernel optimization. Suitable for demanding tasks.

```python
SparseAttention(
    dim=1024, dim_head=64, heads=16,
    sliding_window_size=64,
    compress_block_size=8, compress_block_sliding_stride=4,
    selection_block_size=8, num_selected_blocks=8,
    kv_heads=4, causal=True, use_triton_kernel=True
)
```

--------------------------------

### sample

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Transformer.md

Generates tokens autoregressively with optional top-k sampling. This method allows for controlled text generation by specifying parameters like temperature and filtering thresholds.

```APIDOC
## sample

### Description
Generates tokens autoregressively with optional top-k sampling.

### Method
`sample`

### Parameters
- `prompt` (Tensor): Initial prompt tokens of shape `(batch, prompt_len)`
- `seq_len` (int): Total target sequence length (includes prompt)
- `temperature` (float): Softmax temperature (0 for greedy, >0 for stochastic)
- `filter_thres` (float): Top-k threshold (0.9 = keep top 10%)
- `use_cache_kv` (bool): Whether to cache key/values during generation

### Returns
- Generated tokens of shape `(batch, seq_len - prompt_len)` beyond the prompt
```

--------------------------------

### AttentionPool Initialization

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Initializes an AttentionPool module for pooling-based compression. The compress_window_size parameter is kept for interface compatibility but is not used.

```python
class AttentionPool(nn.Module):
    def __init__(
        self,
        dim_head: int,
        compress_window_size: int
    )
```

--------------------------------

### Transformer Constructor

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Initialize a full transformer language model. Configure token count, dimensions, depth, attention heads, and options for sparse or flexible attention mechanisms.

```python
Transformer(
    num_tokens, dim, depth,
    dim_head=64, heads=8, kv_heads=None,
    ff_expansion_factor=4.0,
    use_sparse_attn=False, causal=True,
    use_flex_sliding_window=False,
    use_flex_fine_selection=False,
    use_triton_fine_selection=False,
    sparse_attn_kwargs={...}
)
```

--------------------------------

### SparseAttention Constructor

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Initialize the main sparse attention layer. Configure dimensions, attention heads, windowing, compression, and selection parameters.

```python
SparseAttention(
    dim, dim_head, heads, sliding_window_size,
    compress_block_size, compress_block_sliding_stride,
    selection_block_size, num_selected_blocks,
    kv_heads=None, num_compressed_mem_kv=1, causal=False,
    norm=True, use_diff_topk=False, use_triton_kernel=False,
    query_heads_share_selected_kv=True,
    compress_mlp=None, compress_mlp_expand_factor=1.0,
    strategy_combine_mlp=None
)
```

--------------------------------

### Medium Model (Balanced) SparseAttention Configuration

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/configuration.md

Configuration for a medium-sized SparseAttention model, offering a balance between performance and resource consumption. Includes an option for Triton kernels on compatible hardware.

```python
attn = SparseAttention(
    dim=512,
    dim_head=64,
    heads=8,
    sliding_window_size=32,
    compress_block_size=4,
    compress_block_sliding_stride=2,
    selection_block_size=4,
    num_selected_blocks=4,
    causal=True,
    use_triton_kernel=True  # If on A100+
)
```

--------------------------------

### GroupedMLP Initialization

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Initializes a GroupedMLP module for per-head MLP compression. Allows configuration of hidden layer expansion factor. Ensure heads and dim_head match attention configuration.

```python
class GroupedMLP(nn.Module):
    def __init__(
        self,
        dim_head: int,
        compress_window_size: int,
        heads: int,
        expand_factor: float = 1.0
    )
```

--------------------------------

### Customizing SparseAttention with CompressTransformer

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Shows how to provide a custom CompressTransformer instance to SparseAttention, overriding the default MLP. Ensure the CompressTransformer is initialized with appropriate parameters.

```python
from native_sparse_attention_pytorch import SparseAttention
from native_sparse_attention_pytorch.compress_networks import CompressTransformer

compress_mlp = CompressTransformer(
    num_layers=3,
    dim=512,
    num_heads=8,
    ff_hidden_dim=2048
)

attn = SparseAttention(
    dim=512,
    dim_head=64,
    heads=8,
    sliding_window_size=32,
    compress_block_size=4,
    compress_block_sliding_stride=2,
    selection_block_size=4,
    num_selected_blocks=4,
    compress_mlp=compress_mlp,
    compress_mlp_expand_factor=1.0  # ignored when compress_mlp is provided
)
```

--------------------------------

### Configure Grouped Query Attention (GQA)

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/SparseAttention.md

Illustrates how to configure the SparseAttention layer for regular multi-head attention versus grouped-query attention using the `kv_heads` parameter.

```python
# Regular multi-head attention (8 heads for both Q and KV)
attn = SparseAttention(
    dim=512, dim_head=64, heads=8,
    ...
)

# Grouped query attention (8 Q heads, 2 KV heads)
attn = SparseAttention(
    dim=512, dim_head=64, heads=8, kv_heads=2,
    ...
)
```

--------------------------------

### Long Context (8k+ tokens) SparseAttention Configuration

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/configuration.md

Configuration optimized for handling very long contexts (8k+ tokens) using SparseAttention. Employs larger block sizes and Triton kernels for efficiency.

```python
attn = SparseAttention(
    dim=512,
    dim_head=64,
    heads=8,
    sliding_window_size=128,
    compress_block_size=16,
    compress_block_sliding_stride=8,
    selection_block_size=16,
    num_selected_blocks=8,
    num_compressed_mem_kv=2,
    causal=True,
    use_triton_kernel=True
)
```

--------------------------------

### Importing Triton Kernel Function

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Demonstrates the import statement for the custom Triton kernel used for sparse attention computation.

```python
# Triton kernels
from native_sparse_attention_pytorch.triton_native_sparse_attention import (
    native_sparse_attend
)
```

--------------------------------

### Query Head Sharing Options

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Architecture.md

Illustrates the difference between sharing and not sharing query heads for block selection in attention mechanisms. Sharing reduces memory and compute at the potential cost of quality, while independent selection may improve quality but increase cost.

```text
With sharing (query_heads_share_selected_kv=True):
  Average importance scores across H/KV_H query heads
  All queries in group attend to same selected blocks
  → Lower memory and compute

Without sharing (query_heads_share_selected_kv=False):
  Each query head selects independently
  → Potentially better quality, higher cost
```

--------------------------------

### Type Annotations

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Jaxtyping-based tensor type hints for improved code clarity and static analysis.

```APIDOC
## Type Annotations

**File:** `tensor_typing.py`

Jaxtyping-based tensor type hints:

```python
from native_sparse_attention_pytorch.tensor_typing import Float, Int, Bool
```

Usage:
```python
def forward(q: Float['b h n d'], k: Float['b h n d']) -> Float['b h n d']:
    ...
```
```

--------------------------------

### Custom Compression Network for Sparse Attention

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Illustrates how to integrate a custom compression network, such as CompressTransformer, into the SparseAttention layer.

```python
from native_sparse_attention_pytorch import SparseAttention
from native_sparse_attention_pytorch.compress_networks import CompressTransformer

compress = CompressTransformer(
    num_layers=2,
    dim=512,
    num_heads=8
)

attn = SparseAttention(
    dim=512, dim_head=64, heads=8,
    sliding_window_size=32,
    compress_block_size=4,
    compress_block_sliding_stride=2,
    selection_block_size=4,
    num_selected_blocks=4,
    compress_mlp=compress
)

output = attn(x)
```

--------------------------------

### Import Tensor Types

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Import type hints for tensors from the tensor_typing module.

```python
from native_sparse_attention_pytorch.tensor_typing import Float, Int, Bool
```

--------------------------------

### Importing Transformer Components

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Details the import statements for Transformer, Attention, and FeedForward modules.

```python
# Transformer models
from native_sparse_attention_pytorch.transformer import (
    Transformer, Attention, FeedForward
)
```

--------------------------------

### Importing Core Sparse Attention Class

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Shows the import statement for the main SparseAttention class.

```python
# Main class
from native_sparse_attention_pytorch import SparseAttention
```

--------------------------------

### Transformer Model with Sparse Attention

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Shows how to instantiate and train a Transformer model that leverages sparse attention for efficiency.

```python
from native_sparse_attention_pytorch.transformer import Transformer

model = Transformer(
    num_tokens=256,
    dim=512,
    depth=6,
    heads=8,
    use_sparse_attn=True
)

# Training
token_ids = torch.randint(0, 256, (4, 256))
loss = model(token_ids, return_loss=True)

# Generation
prompt = torch.randint(0, 256, (1, 10))
generated = model.sample(prompt, seq_len=100, temperature=0.8)
```

--------------------------------

### Sliding Window Path: Attention Calculation

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Architecture.md

Illustrates the attention mechanism for the Sliding Window Path, where queries attend to a fixed-size window of keys and values.

```text
Q: (B, H, N, D_head)
K: (B, KV_H, N, D_head)
    ↓ Extract last window_size tokens
K_window: (B, KV_H, window_size, D_head)
V_window: (B, KV_H, window_size, D_head)
    ↓ Compute local attention
Sliding_attn_out: (B, H, N, D_head)
    [All query positions attend to same window]
```

--------------------------------

### Language Modeling with Transformer

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/README.md

Instantiate and use the Transformer model for language modeling tasks, including training and generation. Requires PyTorch.

```python
from native_sparse_attention_pytorch.transformer import Transformer

model = Transformer(
    num_tokens=10000,
    dim=512,
    depth=12,
    heads=8,
    use_sparse_attn=True,
    causal=True
)

# Training
token_ids = torch.randint(0, 10000, (4, 512))
loss = model(token_ids, return_loss=True)
loss.backward()

# Generation
prompt = torch.randint(0, 10000, (1, 50))
generated = model.sample(prompt, seq_len=200, temperature=0.8)
```

--------------------------------

### Selection Block Layout

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Architecture.md

Describes the layout of fine selection blocks after compression. It shows how many fine blocks are created and how selection is performed per query position.

```python
Compressed blocks: 127 (from above)
Fine block size: 8
Fine blocks: ceil(127 / 8) = 16

Selection: Per query position, pick top K of 16 fine blocks
```

--------------------------------

### Create Compression Mask

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Function to create a compression attention mask for PyTorch FlexAttention backend.

```python
create_compress_mask(seq_len, kv_seq_len, compress_block_sliding_stride, mem_kv_len=0, causal=True)
```

--------------------------------

### Provide a default value

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/HelperFunctions.md

Returns the given value if it is not None, otherwise returns a specified default value. Useful for optional parameters.

```python
def default(v, d):
    return v if exists(v) else d
```

--------------------------------

### Custom Compression Network

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/README.md

Define and use a custom compression network with SparseAttention. This allows for more efficient attention mechanisms by compressing key-value blocks.

```python
from native_sparse_attention_pytorch import SparseAttention
from native_sparse_attention_pytorch.compress_networks import CompressTransformer

compress = CompressTransformer(
    num_layers=2,
    dim=512,
    num_heads=8
)

attn = SparseAttention(
    dim=512, dim_head=64, heads=8,
    sliding_window_size=32,
    compress_block_size=4,
    compress_block_sliding_stride=2,
    selection_block_size=4,
    num_selected_blocks=4,
    compress_mlp=compress
)
```

--------------------------------

### Importing Type Annotations

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Shows the import statements for type annotations used within the library.

```python
# Type annotations
from native_sparse_attention_pytorch.tensor_typing import Float, Int, Bool
```

--------------------------------

### Large Model (Comprehensive) SparseAttention Configuration

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/configuration.md

Configuration for a large and comprehensive SparseAttention model, designed for maximum capacity and performance. Supports grouped-query attention and Triton kernels.

```python
attn = SparseAttention(
    dim=1024,
    dim_head=64,
    heads=16,
    sliding_window_size=64,
    compress_block_size=8,
    compress_block_sliding_stride=4,
    selection_block_size=8,
    num_selected_blocks=8,
    kv_heads=4,  # Grouped-query attention
    causal=True,
    use_triton_kernel=True
)
```

--------------------------------

### Attention Constructor

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Initialize a standard multi-head self-attention layer with rotary embeddings.

```python
Attention(dim, dim_head=64, heads=8, causal=True, kv_heads=None)
```

--------------------------------

### Ensuring Same Device for Components

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/errors.md

When encountering device mismatches, ensure that both the attention module and the input tensor reside on the same device (e.g., CUDA). Explicitly move components to the correct device if necessary.

```python
# Ensure all components on same device
attn = attn.cuda()
x = x.cuda()
output = attn(x)

# Problem: Mixed devices
attn = attn.cuda()
x = x  # Left on CPU
output = attn(x)  # RuntimeError
```

--------------------------------

### Low-Level Forward Kernel for Sparse Attention

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/TritonKernels.md

Internal kernel launcher for the forward pass of sparse attention. Used by `native_sparse_attend`. Requires block indices and mask for key-value blocks.

```python
def native_sparse_attn_forward(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    kv_block_indices: Tensor,
    kv_block_mask: Tensor,
    block_size: int = 128,
    include_block_causal: bool = True,
    return_sliding_window_out: bool = False
) -> tuple[Tensor, Tensor, Tensor, Tensor]
```

--------------------------------

### Native Sparse Attention Pathways

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Architecture.md

Illustrates the three parallel attention pathways: Compressed, Fine Selection, and Sliding Window, which are combined to produce the final output.

```text
Input (batch, seq_len, dim)
    │
    ├─→ [Compressed Attention Path]  ─┐
    ├─→ [Fine Selection Path]         ├─→ [Weighted Combine] → Output
    ├─→ [Sliding Window Path]         ┘
```

--------------------------------

### SparseAttention Class

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/SparseAttention.md

Initializes the SparseAttention module with various parameters controlling the attention mechanisms, dimensions, and strategies.

```APIDOC
## Class SparseAttention

### Description
Initializes the SparseAttention module with various parameters controlling the attention mechanisms, dimensions, and strategies.

### Parameters
#### Constructor Parameters
- **dim** (int) - Required - Model embedding dimension
- **dim_head** (int) - Required - Dimension per attention head (e.g., 64)
- **heads** (int) - Required - Total number of query attention heads
- **sliding_window_size** (int) - Required - Size of the local sliding window for local attention strategy
- **compress_block_size** (int) - Required - Size of blocks in the key/value sequence for compression
- **compress_block_sliding_stride** (int) - Required - Stride when creating compress blocks; must be ≤ compress_block_size
- **selection_block_size** (int) - Required - Size of blocks for fine-grained attention selection
- **num_selected_blocks** (int) - Required - Number of top-scoring compress blocks to attend to in fine attention (≥ 0)
- **kv_heads** (int) - Optional - Number of key/value heads for grouped-query attention. Defaults to `heads` if not specified. Must divide evenly into `heads`
- **num_compressed_mem_kv** (int) - Optional - Number of learnable memory tokens for compressed attention. Default: 1
- **causal** (bool) - Optional - Whether to apply causal masking (autoregressive mode). Default: False
- **norm** (bool) - Optional - Whether to apply RMSNorm before computing Q, K, V. Default: True
- **use_diff_topk** (bool) - Optional - Whether to use differentiable top-k for block selection with straight-through estimators. Default: False
- **use_triton_kernel** (bool) - Optional - Whether to use Triton-optimized kernels for fine attention (requires GPU). Default: False
- **query_heads_share_selected_kv** (bool) - Optional - If True, importance scores are averaged across query heads to select top-k blocks per kv head. If False, each query head within a group can select different blocks (higher memory/compute). Default: True
- **compress_mlp** (nn.Module) - Optional - Custom compression MLP for compress blocks. If None, a default 2-layer MLP is created. Default: None
- **compress_mlp_expand_factor** (float) - Optional - Expansion factor for the hidden dimension of the default compression MLP. Default: 1.0
- **strategy_combine_mlp** (nn.Module) - Optional - Custom MLP for combining the three attention strategies. If None, a learned weighted combination is created. Default: None
```

--------------------------------

### Configuring num_selected_blocks for Fine Attention

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/errors.md

When fine attention is experimental and `num_selected_blocks` is set to 0, it is disabled. Set `num_selected_blocks` to a positive integer to enable it.

```python
# Better configuration
attn = SparseAttention(
    ...,
    num_selected_blocks=4  # Select top-4 blocks per query
)
```

--------------------------------

### SimpleMultiheadSelfAttention Module

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Defines a standard multi-head self-attention layer. Use this for basic self-attention mechanisms.

```python
class SimpleMultiheadSelfAttention(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        dropout: float = 0.0
    )
```

--------------------------------

### SimpleTransformerFeedForward Module

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Implements a two-layer feed-forward network with GELU activation. Suitable for transformer architectures.

```python
class SimpleTransformerFeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        dropout: float = 0.0
    )
```

--------------------------------

### Strategy Combination: Weighted Sum

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Architecture.md

Shows how the outputs from the three attention pathways are combined using learned weights.

```text
Compressed_attn_out: (B, H, N, D_head)
Fine_attn_out:       (B, H, N, D_head)
Sliding_attn_out:    (B, H, N, D_head)
    ↓ Stack
Stack: (3, B, H, N, D_head)

Weights: (B, N, H, 3)  [from strategy_combine_mlp]
    ↓ Sigmoid activation
Weights: (B, N, H, 3)  [in range (0, 1)]
    ↓ Weighted sum
Output: (B, H, N, D_head) = ∑(Weights[:,:,:,i] * Path[i])
```

--------------------------------

### Inference with Caching in Transformer

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Demonstrates an efficient inference loop for a Transformer model using a cache to store past states, speeding up token generation.

```python
model = Transformer(
    num_tokens=256,
    dim=512,
    depth=6,
    causal=True
)

cache = None
prompt = torch.randint(0, 256, (1, 10))

for _ in range(50):
    logits, cache = model(prompt[:, -1:], cache=cache, return_cache=True)
    next_token = logits[:, -1].argmax(-1, keepdim=True)
    prompt = torch.cat([prompt, next_token], dim=1)
```

--------------------------------

### default(v, d)

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/HelperFunctions.md

Returns the value `v` if it is not None; otherwise, returns a default value `d`. Simplifies setting default parameters.

```APIDOC
## default(v, d)

### Description
Return value `v` if not None, otherwise return default `d`.

### Parameters
* **v**: The value to check.
* **d**: The default value to return if `v` is None.

### Usage Example
```python
kv_heads = default(kv_heads, heads)  # Use heads if kv_heads not provided
```
```

--------------------------------

### Generic Scaled Dot-Product Attention

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/HelperFunctions.md

Implements scaled dot-product attention with support for grouped-query attention. Use this for standard attention calculations.

```python
def attend(
    q,
    k,
    v,
    mask=None,
    return_sim=False,
    scale=None
)
```

```python
from native_sparse_attention_pytorch.native_sparse_attention import attend

q = torch.randn(2, 8, 256, 64)
k = torch.randn(2, 8, 256, 64)
v = torch.randn(2, 8, 256, 64)

output = attend(q, k, v, scale=64**-0.5)
```

--------------------------------

### Fine Selection Path: Selection Phase

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Architecture.md

Details the block selection process in the Fine Selection Path, including aggregation and Top-K selection of importance scores.

```text
Importance_scores: (B, H, N, num_blocks)
    ↓ [Optional] Average across grouped queries
Importance_scores: (B, KV_H, N, num_blocks)  [if query_heads_share_selected_kv=True]
    ↓ Aggregate compressed/fine block scores (if stride ≠ block_size)
Importance_scores: (B, KV_H, N, num_fine_blocks)
    ↓ Softmax + Top-K selection
Selected_indices: (B, KV_H, N, num_selected_blocks)
Selected_scores: (B, KV_H, N, num_selected_blocks)
```

--------------------------------

### Resolving Flex and Triton Fine Selection Conflict

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/errors.md

Avoid using both `use_flex_fine_selection` and `use_triton_fine_selection` simultaneously. Choose only one backend for fine selection to prevent assertion errors.

```python
# Wrong
model = Transformer(
    ...,
    use_flex_fine_selection=True,
    use_triton_fine_selection=True
)

# Correct: Choose one
model = Transformer(..., use_flex_fine_selection=True)
# OR
model = Transformer(..., use_triton_fine_selection=True)
```

--------------------------------

### SingleProjection

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Implements a simple linear projection for sequence compression. It can optionally use per-head projections if the number of heads is greater than 1.

```APIDOC
## SingleProjection

### Description
Simple linear projection-based compression. Optionally grouped per head.

### Constructor
```python
SingleProjection(dim_head: int, compress_window_size: int, heads: int = 1)
```

### Constructor Parameters
#### Parameters
- **dim_head** (int) - Dimension per head
- **compress_window_size** (int) - Size of block to compress
- **heads** (int) - Number of heads. If > 1, uses per-head projection (default: 1)

### Forward Method
```python
forward(self, kv: Tensor) -> Tensor
```

#### Parameters
- **kv** (Tensor) - Input of shape `(batch, heads, windows, block_size, dim_head)`

#### Returns
- Tensor of shape `(batch, heads, windows, dim_head)`

### Example
```python
from native_sparse_attention_pytorch.compress_networks import SingleProjection

proj = SingleProjection(dim_head=64, compress_window_size=4, heads=8)
kv = torch.randn(2, 8, 10, 4, 64)
compressed = proj(kv)  # (2, 8, 10, 64)
```
```

--------------------------------

### Transformer Model Generation

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/Transformer.md

Generates new tokens autoregressively using the Transformer model's sample method. Supports temperature-based stochasticity and top-k filtering. Caching of key/value pairs can be enabled for efficiency.

```python
# Generate 100 tokens from a 10-token prompt
prompt = torch.randint(0, 256, (1, 10))
generated = model.sample(
    prompt=prompt,
    seq_len=110,
    temperature=0.8,
    filter_thres=0.9,
    use_cache_kv=True
)
# generated shape: (1, 100)
```

--------------------------------

### Debugging Shape Mismatches in SparseAttention

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/errors.md

Verify that the total dimension calculated from `dim_head` and `heads` matches the specified `dim` in `SparseAttention`. Incorrect dimensions will lead to failures.

```python
# Typical issue: Dimension mismatch
attn = SparseAttention(dim=512, dim_head=64, heads=8)
# Total: 64 * 8 = 512 ✓

# Wrong: Dimension mismatch
attn = SparseAttention(dim=512, dim_head=64, heads=7)
# Total: 64 * 7 = 448 ✗ (will fail)
```

--------------------------------

### Helper Functions

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/API_Overview.md

Utility functions for tensor operations, numeric computation, and sampling.

```APIDOC
## Helper Functions

**File:** Various modules

Utility functions for tensor operations, numeric computation, and sampling:

- **Existence**: `exists(v)`, `default(v, d)`
- **Numeric**: `round_down_mult()`, `round_up_mult()`, `divisible_by()`
- **Tensor Ops**: `pad_at_dim()`, `pack_one_with_inverse()`, `straight_through()`
- **Inspection**: `is_empty()`, `max_neg_value()`, `is_contiguous()`
- **Sampling**: `log()`, `gumbel_noise()`, `gumbel_sample()`, `top_k()`
```

--------------------------------

### AttentionPool

Source: https://github.com/lucidrains/native-sparse-attention-pytorch/blob/main/_autodocs/CompressionNetworks.md

Implements pooling-based compression inspired by Enformer. It uses softmax-weighted pooling over the block dimension to compress sequences. The compress_window_size parameter is kept for interface compatibility but is not used.

```APIDOC
## AttentionPool

### Description
Pooling-based compression inspired by Enformer (DeepMind). Uses softmax-weighted pooling over the block dimension.

### Constructor
```python
AttentionPool(dim_head: int, compress_window_size: int)
```

### Constructor Parameters
#### Parameters
- **dim_head** (int) - Dimension per head
- **compress_window_size** (int) - Not used; kept for interface compatibility

### Forward Method
```python
forward(self, kv: Tensor) -> Tensor
```

#### Parameters
- **kv** (Tensor) - Input of shape `(batch, heads, windows, block_size, dim_head)`

#### Returns
- Tensor of shape `(batch, heads, windows, dim_head)` — attention-pooled values

### Example
```python
from native_sparse_attention_pytorch.compress_networks import AttentionPool

pool = AttentionPool(dim_head=64, compress_window_size=4)
kv = torch.randn(2, 8, 10, 4, 64)
compressed = pool(kv)  # (2, 8, 10, 64)
```
```