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

NATTEN API完全参考轻松调用多维稀疏注意力的关键接口与参数【免费下载链接】NATTENFast Multi-dimensional Sparse Attention项目地址: https://gitcode.com/gh_mirrors/na/NATTENNATTENFast Multi-dimensional Sparse Attention是一个高性能的多维稀疏注意力库专为处理复杂数据结构中的注意力机制设计。本指南将全面解析NATTEN的核心API接口与参数帮助开发者快速上手并高效调用多维稀疏注意力功能。核心功能概述NATTEN提供了一系列优化的注意力操作支持1D、2D和3D等多维稀疏注意力计算特别适用于处理序列数据、图像和视频等复杂输入。其核心优势在于多维支持原生支持1D序列、2D图像和3D视频数据的稀疏注意力计算高性能核基于CUDA的优化实现包括针对Hopper和Blackwell架构的专用 kernels灵活接口提供函数式和模块化两种调用方式适配不同使用场景NATTEN多维稀疏注意力机制可视化展示了不同维度下的注意力计算模式主要API接口详解1. 函数式接口attentionattention函数是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-fmhaq_tile_size/kv_tile_size: 前向计算的分块大小run_persistent_kernel: 是否使用持久化核函数默认为True2. 模块化接口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 GPUhopper-fmha: 针对Hopper架构优化如H100blackwell-fmha: 针对Blackwell架构优化如B200flex-fmha: 灵活的纯PyTorch实现便于调试和扩展性能对比不同后端在典型任务上的性能表现吞吐量越高越好Hopper架构下不同配置的性能对比Blackwell架构下不同配置的性能对比后端选择策略开发与调试使用flex-fmha便于调试和原型验证Hopper GPU (H100)使用hopper-fmha利用Tensor Cores优化Blackwell GPU (B200)使用blackwell-fmha支持最新硬件特性兼容性优先使用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, backendblackwell-fmha, q_tile_size128, kv_tile_size64 )1D稀疏注意力示意图2D注意力图像数据适用于图像、视频帧等2D数据# 2D图像注意力示例 model natten.NeighborhoodAttentionGeneric( dim256, kernel_size7, num_heads8, dilation1, stride1, padding3 ).cuda() input torch.randn(2, 32, 32, 256).cuda() # [batch, height, width, dim] output model(input)2D稀疏注意力示意图3D注意力视频数据适用于视频、3D医学图像等3D数据# 3D视频注意力示例 model natten.NeighborhoodAttentionGeneric( dim512, kernel_size(3, 7, 7), # 时间×高度×宽度 num_heads16, dilation(1, 2, 2), stride1, 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_Qcumulative_seqlen_Q, cumulative_seqlen_KVcumulative_seqlen_KV, max_seqlen_Qmax_seqlen_Q, max_seqlen_KVmax_seqlen_KV, backendblackwell-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)性能调优建议分块大小选择大序列2048使用较大分块q_tile_size128, kv_tile_size64小序列512使用较小分块q_tile_size64, kv_tile_size32精度选择优先使用FP16或BF16显著提升性能Blackwell GPU可尝试FP8进一步提高吞吐量核函数调度大模型使用kernel_schedulepersistent小模型使用kernel_scheduleephemeral安装与快速开始安装步骤# 从源码安装 git clone https://gitcode.com/gh_mirrors/na/NATTEN cd NATTEN pip install -e .验证安装import natten print(fNATTEN 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考