NATTEN API完全参考：轻松调用多维稀疏注意力的关键接口与参数

【免费下载链接】NATTENFast Multi-dimensional Sparse Attention项目地址: https://gitcode.com/gh_mirrors/na/NATTEN

NATTEN（Fast Multi-dimensional Sparse Attention）是一个高性能的多维稀疏注意力库，专为处理复杂数据结构中的注意力机制设计。本指南将全面解析NATTEN的核心API接口与参数，帮助开发者快速上手并高效调用多维稀疏注意力功能。

核心功能概述

NATTEN提供了一系列优化的注意力操作，支持1D、2D和3D等多维稀疏注意力计算，特别适用于处理序列数据、图像和视频等复杂输入。其核心优势在于：

• 多维支持：原生支持1D（序列）、2D（图像）和3D（视频）数据的稀疏注意力计算
• 高性能核：基于CUDA的优化实现，包括针对Hopper和Blackwell架构的专用 kernels
• 灵活接口：提供函数式和模块化两种调用方式，适配不同使用场景

NATTEN多维稀疏注意力机制可视化，展示了不同维度下的注意力计算模式

主要API接口详解

1. 函数式接口：attention

attention函数是NATTEN的核心接口，用于执行标准点积注意力计算，支持多维稀疏模式。

def attention( query: Tensor, key: Tensor, value: Tensor, is_causal: bool = False, scale: Optional[float] = None, # varlen parameters seqlens_Q: Optional[Tensor] = None, seqlens_KV: Optional[Tensor] = None, cumulative_seqlen_Q: Optional[Tensor] = None, cumulative_seqlen_KV: Optional[Tensor] = None, max_seqlen_Q: Optional[int] = None, max_seqlen_KV: Optional[int] = None, # backend parameters backend: Optional[str] = None, q_tile_size: Optional[int] = None, kv_tile_size: Optional[int] = None, backward_q_tile_size: Optional[int] = None, backward_kv_tile_size: Optional[int] = None, backward_kv_splits: Optional[int] = None, backward_use_pt_reduction: bool = False, run_persistent_kernel: bool = True, kernel_schedule: Optional[Union[str, KernelSchedule]] = None, torch_compile: bool = False, return_lse: bool = False, ) -> Union[Tensor, Tuple[Tensor, Tensor]]:

关键参数说明：

• 输入张量：

  • query: 4D查询张量，形状为[batch, seqlen, heads, head_dim]
  • key: 4D键张量，形状为[batch, seqlen_kv, heads_kv, head_dim]
  • value: 4D值张量，形状为[batch, seqlen_kv, heads_kv, head_dim_v]

• 注意力控制：

  • is_causal: 是否启用因果掩码，默认为False（双向注意力）
  • scale: 注意力缩放因子，默认为head_dim ** -0.5

• 变长序列支持：

  • seqlens_Q/seqlens_KV: 变长序列长度张量（适用于非编译场景）
  • cumulative_seqlen_Q/cumulative_seqlen_KV: 累积序列长度（编译友好）
  • max_seqlen_Q/max_seqlen_KV: 最大序列长度

• 性能优化：

  • backend: 指定后端（如cutlass-fmha、hopper-fmha、blackwell-fmha）
  • q_tile_size/kv_tile_size: 前向计算的分块大小
  • run_persistent_kernel: 是否使用持久化核函数，默认为True

2. 模块化接口：NeighborhoodAttentionGeneric

对于PyTorch模型集成，NATTEN提供了NeighborhoodAttentionGeneric模块，可直接作为神经网络层使用。

class NeighborhoodAttentionGeneric(nn.Module): def __init__( self, dim: int, kernel_size: Union[int, Tuple[int, ...]], num_heads: int, dilation: Union[int, Tuple[int, ...]] = 1, stride: Union[int, Tuple[int, ...]] = 1, padding: Optional[Union[int, Tuple[int, ...]]] = None, qkv_bias: bool = True, proj_bias: bool = True, attn_drop: float = 0.0, proj_drop: float = 0.0, is_causal: bool = False, kernel_size_kv: Optional[Union[int, Tuple[int, ...]]] = None, dilation_kv: Optional[Union[int, Tuple[int, ...]]] = None, stride_kv: Optional[Union[int, Tuple[int, ...]]] = None, padding_kv: Optional[Union[int, Tuple[int, ...]]] = None, additional_heads: int = 0, additional_kv_heads: int = 0, share_additional_kv: bool = True, # backend parameters backend: Optional[str] = None, torch_compile: bool = False, ) -> None:

核心参数：

• 维度与头数：

  • dim: 输入特征维度
  • num_heads: 注意力头数

• 邻域控制：

  • kernel_size: 注意力核大小（决定邻域范围）
  • dilation: 膨胀率，控制感受野大小
  • stride: 步幅，控制下采样
  • padding: 填充大小

• 高级配置：

  • qkv_bias/proj_bias: 是否使用偏置
  • attn_drop/proj_drop: Dropout比率
  • is_causal: 是否启用因果注意力

后端选择与性能优化

NATTEN提供多种后端实现，针对不同硬件架构优化：

可用后端

• cutlass-fmha: 基础CUTLASS实现，兼容大多数NVIDIA GPU
• hopper-fmha: 针对Hopper架构优化（如H100）
• blackwell-fmha: 针对Blackwell架构优化（如B200）
• flex-fmha: 灵活的纯PyTorch实现，便于调试和扩展

性能对比

不同后端在典型任务上的性能表现（吞吐量，越高越好）：

Hopper架构下不同配置的性能对比

Blackwell架构下不同配置的性能对比

后端选择策略

1. 开发与调试：使用flex-fmha，便于调试和原型验证
2. Hopper GPU (H100)：使用hopper-fmha，利用Tensor Cores优化
3. Blackwell GPU (B200)：使用blackwell-fmha，支持最新硬件特性
4. 兼容性优先：使用cutlass-fmha，兼容大多数NVIDIA GPU

多维注意力操作

NATTEN支持多种维度的稀疏注意力操作，适应不同数据类型：

1D注意力（序列数据）

适用于文本、时间序列等1D数据：

import natten from natten import functional as F # 1D序列注意力示例 query = torch.randn(2, 1024, 12, 64).cuda() # [batch, seqlen, heads, head_dim] key = torch.randn(2, 1024, 12, 64).cuda() value = torch.randn(2, 1024, 12, 64).cuda() # 执行1D稀疏注意力 output = F.attention( query, key, value, backend="blackwell-fmha", q_tile_size=128, kv_tile_size=64 )

1D稀疏注意力示意图

2D注意力（图像数据）

适用于图像、视频帧等2D数据：

# 2D图像注意力示例 model = natten.NeighborhoodAttentionGeneric( dim=256, kernel_size=7, num_heads=8, dilation=1, stride=1, padding=3 ).cuda() input = torch.randn(2, 32, 32, 256).cuda() # [batch, height, width, dim] output = model(input)

2D稀疏注意力示意图

3D注意力（视频数据）

适用于视频、3D医学图像等3D数据：

# 3D视频注意力示例 model = natten.NeighborhoodAttentionGeneric( dim=512, kernel_size=(3, 7, 7), # 时间×高度×宽度 num_heads=16, dilation=(1, 2, 2), stride=1, padding=(1, 7, 7) ).cuda() input = torch.randn(2, 16, 32, 32, 512).cuda() # [batch, time, height, width, dim] output = model(input)

3D稀疏注意力示意图

实用工具与最佳实践

变长序列处理

NATTEN提供工具函数处理变长序列，优化内存使用：

from natten.utils.varlen import generate_varlen_parameters # 生成变长序列参数 cumulative_seqlen_Q, cumulative_seqlen_KV, max_seqlen_Q, max_seqlen_KV = generate_varlen_parameters( q, k, v, seqlens_Q, seqlens_KV ) # 使用变长参数调用注意力 output = F.attention( query, key, value, cumulative_seqlen_Q=cumulative_seqlen_Q, cumulative_seqlen_KV=cumulative_seqlen_KV, max_seqlen_Q=max_seqlen_Q, max_seqlen_KV=max_seqlen_KV, backend="blackwell-fmha" )

上下文配置

NATTEN提供上下文配置功能，控制内存使用和确定性：

import natten.context as ctx # 设置内存使用偏好 ctx.set_memory_usage_preference("strict") # 严格控制内存使用 # 启用确定性算法 ctx.use_deterministic_algorithms(True) # 启用KV并行 ctx.use_kv_parallelism_in_fused_na(True)

性能调优建议

1. 分块大小选择：

   • 大序列（>2048）：使用较大分块（q_tile_size=128, kv_tile_size=64）
   • 小序列（<512）：使用较小分块（q_tile_size=64, kv_tile_size=32）

2. 精度选择：

   • 优先使用FP16或BF16，显著提升性能
   • Blackwell GPU可尝试FP8，进一步提高吞吐量

3. 核函数调度：

   • 大模型：使用kernel_schedule="persistent"
   • 小模型：使用kernel_schedule="ephemeral"

安装与快速开始

安装步骤

# 从源码安装 git clone https://gitcode.com/gh_mirrors/na/NATTEN cd NATTEN pip install -e .

验证安装

import natten print(f"NATTEN version: {natten.__version__}") # 输出: NATTEN version: x.x.x # 验证CUDA后端 natten.utils.checks.can_run_cutlass_fna(torch.randn(1, 16, 16, 64).cuda()) # 输出: True

总结

NATTEN提供了强大而灵活的多维稀疏注意力API，通过优化的后端实现和丰富的配置选项，能够高效处理各种复杂数据结构的注意力计算。无论是序列数据、图像还是视频，NATTEN都能提供卓越的性能和易用性，是构建先进深度学习模型的理想选择。

通过本指南，您应该已经掌握了NATTEN核心API的使用方法和最佳实践。如需深入了解更多高级功能，请参考官方文档：docs/index.md。

祝您在项目中充分发挥NATTEN的强大能力！🚀

【免费下载链接】NATTENFast Multi-dimensional Sparse Attention项目地址: https://gitcode.com/gh_mirrors/na/NATTEN