大模型深度学习算子库后端高性能计算【免费下载链接】flashinferFlashInfer: Kernel Library for LLM Serving项目地址https://gitcode.com/gh_mirrors/fl/flashinfer点击查看免费下载flashinfer.cudnn是 FlashInfer 中直接调用 NVIDIA cuDNN 运行时的一组 Python 封装它覆盖两条完全不同的技术线——一条是 cuDNN 的 SDPAScaled Dot-Product Attention图引擎为批量 Prefill 与 Decode 注意力提供backendcudnn后端另一条是 cuDNN 面向 SM100 家族SM100–SM103、SM107的融合线性注意力引擎为 Gated Delta RuleGDN、GDN-2、Delta ProductGDP与 Kimi Delta AttentionKDA提供图节点。读完本文你将掌握该后端的版本门槛、六个公开 API 的参数语义、多 token Decode 与滑动窗口的掩码约定、状态张量跨边界的布局规则以及它明确不支持的功能边界从而能直接在BatchPrefillWithPagedKVCacheWrapper/BatchDecodeWithPagedKVCacheWrapper或底层函数上安全地使用 cuDNN 后端。cuDNN 后端在 FlashInfer 中的定位cuDNN 后端是一组薄封装这些封装本身只负责把 PyTorch 张量声明成 cuDNN 图graph的输入输出端口、绑定 UID、执行图并做契约转换如 LSE 的底数换算真正的计算由 cuDNN 运行时完成。它通过两条路径暴露给用户作为 wrapper 的 backend在BatchPrefillWithPagedKVCacheWrapper/BatchDecodeWithPagedKVCacheWrapper上传backendcudnncuDNN 在宿主机 GPU 上可用时或直接调用底层函数。作为线性注意力的 engineflashinfer.chunk_gated_delta_rule、flashinfer.chunk_gated_delta_rule2、flashinfer.chunk_gated_delta_product、flashinfer.recurrent_kda都可以直接传backendcudnn。整个模块由五个文件构成flashinfer/cudnn/decode.py、prefill.py、linear_attention.py分别实现三条功能线utils.py提供按设备缓存的 cuDNN handle 与 cubin 回退加载器__init__.py导出全部六个公开函数函数所在文件对应功能cudnn_batch_prefill_with_kv_cacheflashinfer/cudnn/prefill.py批量 Prefill SDPA支持 Paged / 非 Paged KVcudnn_batch_decode_with_kv_cacheflashinfer/cudnn/decode.py批量 Decode SDPAPaged KVcudnn_chunk_gated_delta_ruleflashinfer/cudnn/linear_attention.pyGDN 分块 PrefillSM100 融合引擎cudnn_chunk_gated_delta_rule2flashinfer/cudnn/linear_attention.pyGDN-2 分块 PrefillSM100 融合引擎cudnn_chunk_gated_delta_productflashinfer/cudnn/linear_attention.pyGDP 分块 PrefillSM100 融合引擎cudnn_recurrent_kdaflashinfer/cudnn/linear_attention.pyKDA 分块 PrefillSM100 融合引擎从源码结构看flashinfer/cudnn/init.py这六个函数是模块对外唯一的公共面。何时被自动选中cudnn从不作为 GDN 或 KDA 的隐式默认后端——这两者 FlashInfer 都有自己的内核而 GDN-2 与 GDP 在 FlashInfer 中没有对应内核因此它们的auto会解析到 cuDNN。这一点在 Decode 侧同样成立backendcudnn只覆盖普通多 token Prefill子集会在到达 Decode 之前被解析并处理见 flashinfer/kda_decode.py 与 flashinfer/kda_decode.py 的报错信息。环境与版本门槛cuDNN 后端依赖 NVIDIA 的 cudnn Python frontendcudnn包。模块在导入时用try: import cudnn探测可用性并设置CUDNN_AVAILABLE标志不可用时SDPA 路径会回退到仓库自带的 cubin 内核功能受限线性注意力路径则直接抛出RuntimeError。两条功能线的版本要求不同门槛差异如下功能线最低版本要求说明线性注意力全部四个入口cudnn-frontend 1.29 且安装cutedslextra无 cuDNN backend 版本下限背后的 FROST 引擎是 CuTeDSL 内核由 frontend 自行编译Decode SDPA图路径cuDNN graph 后端可用cudnn包可导入无额外版本检查回退 cubin 路径仅支持 bf16Prefill 直接 seqlensfp16/bf16cuDNN backend 9.24 且 frontend 1.25见_cudnn_supports_direct_seqlensPrefill 直接 seqlensfp8e4m3/e5m2cuDNN backend 9.25 且 frontend 1.27同时是sdpa_fp8暴露cu_seq_len_q/kv的首个版本Prefill 混合形式token 单位 Q 每请求 KV即 Paged 路径cuDNN backend 9.25 且 frontend 1.27frontend PR #430见_cudnn_supports_direct_seqlens(mixedTrue)Prefill 形状覆盖shape overridecuDNN backend 9.22 且 frontend 1.29额外受环境变量FLASHINFER_CUDNN_PREFILL_SHAPE_OVERRIDE默认1控制线性注意力的安装命令在 flashinfer/cudnn/linear_attention.py 中给出pip install -U nvidia-cudnn-frontend[cutedsl]。frontend 版本不足时会快速失败并提示升级。此外代码中明确检查了 FP8 Prefill 需要 cuDNN backend 版本 9.17.1见 flashinfer/cudnn/prefill.py。cuDNN SDPA批量 Prefill 与 Decode 注意力SDPA 路径通过cudnn.graph(handle)构建图Q/K/V、block tables、序列长度、sink 等张量被赋予固定的 UID见decode.py与prefill.py中的UIDs枚举执行时通过var_map把当前调用的张量指针绑定到对应 UID 上再graph.execute。图的构建受cudnn.graph_cache(key_fn...)缓存控制——缓存键把 shape、stride、dtype、attn_scale、mask 参数与各 presence 标志全部烘焙进图描述符因此相同形状但不同 dtype 或不同 scale 的调用绝不会复用同一张图否则重放会静默地把缓冲区按首个调用者的 dtype 重新解释或重放过期 scale 的图。Decode 能力矩阵与参数cudnn_batch_decode_with_kv_cacheflashinfer/cudnn/decode.py覆盖的能力集合如下dtypefp16 / bf16输出沿用q.dtypefp16 需要 cuDNN graph 后端回退 cubin 路径仅 bf16。GQA / MQAnum_heads_qo num_heads_kvQ/KV 头数可以不同。return_lse返回 SDPA 的Stats输出需要 cuDNN graph 后端cubin 回退路径不支持。CUDA graphsis_cuda_graph_compatibleTrue时以可被 CUDA graph 捕获的方式规划执行。多 token Decodeq_len_per_req 1推测解码 / MTP 验证场景掩码为右下因果对角线bottom-right causal diagonal。左滑动窗口window_left控制行只关注对角位置之前的window_left个 key 加上该位置本身。Attention sinkssinks参数形状(num_heads_qo,)、fp32、与 q 同设备语义与 fa2 / trtllm-gen 的sinks一致Streaming-LLM / gpt-oss 风格每头在 softmax 分母中追加一个零值行的额外 logit。核心参数及约定如下表参数形状 / 取值说明q(batch_size * q_len_per_req, num_heads_qo, head_dim)支持任意 batch/head stride例如从打包 QKV 投影切片而来只要 head_dim 是最内层且稠密k_cache/v_cache(total_num_pages, num_heads_kv, page_size, head_dim)Paged KVK/V 完整布局shapestridedtype都进入缓存键scalefloatsoftmax 缩放通常1/sqrt(head_dim)烘焙进图同形状不同 scale 不共享图workspace_buffer张量随 batch 规模增长典型 Decode 负载 128 MB 足够max_sequence_kvint每请求 KV 最大 token 数s_kv_maxactual_seq_lens_kv(batch_size,)每请求 KV 长度cuDNN 路径必须在与 q 相同的 CUDA 设备上block_tables(batch_size, num_pages_per_seq)Paged 页表页宽不同的同 batch 表不共享图重放会以过期行步长寻址q_len_per_reqint默认 1请求内连续行数行 i 关注 key0 .. kv_len - q_len_per_req i要求kv_len q_len_per_reqwindow_leftint默认 -1滑动窗口左界-1 关闭窗口sinks(num_heads_qo,)fp32每头一个额外 softmax logit关于窗口与对角线的对齐源码中有明确注释flashinfer/cudnn/decode.pyFlashInfer 的window_left计数的是严格位于对角线之前的 key而 cuDNN 的左界把对角线本身也计入因此跨边界时做了diagonal_band_left_bound window_left 1的换算当s_qo 1或window_left 0时启用use_causal_mask_bottom_right。sinks在q_len_per_req 1时是否真正可用取决于 cuDNN 栈的 SDPA 引擎cudnn-frontend 1.30 且启用 FROST 引擎时可以服务否则后端引擎会在首次运行时抛出 not-supported 错误这与线性注意力引擎拒绝它无法服务的图是同一套机制。LSE 契约cuDNN 原生输出的是自然对数natural-logsoftmax statsFlashInfer 的契约是以 2 为底lse[b,h] log2(sum_j(exp(scale * q[b,h]·k[...])))对所有有效 KV 位置求和因此执行后统一执行lse.mul_(log2e)换算flashinfer/cudnn/decode.py这样 LSE 可以直接喂给 cascade-merge 内核等下游。Decode 侧还有一个值得注意的校验cuDNN 图路径只接受与张量 batch stride 匹配的稠密 offsets_check_dense_decode_offsets非稠密的batch_offsets_q/o会触发异步设备断言——包括在 CUDA graph 重放前被改动的情况。Prefill 参数与两条长度表达路径cudnn_batch_prefill_with_kv_cacheflashinfer/cudnn/prefill.py接受打包的q(total_qo_tokens, num_heads_qo, head_dim_qk)KV 既可以是 Paged 四维(total_num_pages, num_heads_kv, page_size, head_dim)也可以是非 Paged 三维。注意约束head_dim_qk必须是 128 或 192head_dim_vo必须是 128。与 Decode 不同Prefill 支持fp8e4m3 / e5m2此时走g.sdpa_fp8(...)分支需要q_scale/k_scale/v_scale三个标量反量化张量形状(1,1,1,1)的 fp32 标量张量并输出amax统计。此外还有o_data_type输出 dtype默认等于q.dtype与lse_baselog2默认ln时跳过换算直接返回 cuDNN 原生自然对数 stats。Prefill 的 batch 寻址有两条长度表达路径通过batch_offsets_units参数选择elements默认历史行为batch_offsets_*是预缩放的张量元素偏移例如查询的cumsum(seq_lens) * num_heads_qo * head_dim_qkactual_seq_lens_q/kv通常显式给出。tokensFlashInfer 惯例offsets 是纯 token 单位的 prefix sumqo_indptr/kv_indptr风格。当 cuDNN backend 9.24 / frontend 1.25fp16/bf16或 9.25 / 1.27fp8时token 单位的 indptr 会被直接消费direct path无需转换预扫描否则 FlashInfer 内部把它们缩放成 element 单位。在 direct path 下cu_seq_len_q占据 Q 的 seq-len UID 槽位KV 侧形式独立选择非 Paged 时两侧都用累积长度cu_seq_len_q cu_seq_len_kvPaged 时是混合形式token 单位cu_seq_len_q 每请求seq_len_kvKV 通过 block tables 寻址混合形式需要 cuDNN 9.25。batch_size 1时可以省略 offsetsbatch_size 1时batch_offsets_q与batch_offsets_o必填否则图虽然结构良好但会静默错读每个 batch 的 token源码以显式ValueError拒绝见 flashinfer/cudnn/prefill.py。wrapper 侧传入的qo_indptr本来就是 token 单位因此直接透传为cu_seq_len_q与 Q/O ragged offsets并由引擎按每张量(num_heads * head_dim)乘数自行换算见 flashinfer/prefill.py。Prefill 的 shape override避免逐形状重建计划cuDNN 对非 Paged 路径会按声明的(batch, max_seq_len)固化执行计划——SM100 上构建一次约 55–70 ms若触发内核变体编译则约 1 秒而服务端每次请求的 batch 与长度都在变。为此 Prefill 实现了一套shape override机制flashinfer/cudnn/prefill.py图只在缓存形状下构建一次batch 缓存为 4096超出时按 2 的幂增长序列长度按类别划分——s_q 1cuDNN 的 Decode 内核override 不允许跨越该边界、 128短行引擎_OVERRIDE_SHORT_SEQ、 128归入 65536 或更大的 2 的幂_OVERRIDE_CACHE_SEQ_LONG。短/长引擎的翻转点实测在 128 与 256 之间SM100 与 SM107 上与 batch、LSE、head dim 无关。每次 execute 通过override_uids / override_shapes / override_strides传入真实(b, s_q, s_kv)没有逐形状计划也没有 padding 行。override 图按声明 batch 预留 TMA 描述符batch 4096 时约 1 MiB workspace因此调用方 workspace 不足时prepare_cudnn_batch_prefill会退回精确形状图而非越界执行。可通过环境变量FLASHINFER_CUDNN_PREFILL_SHAPE_OVERRIDE0关闭。cubin 回退路径当cudnn包不可导入或backendcubin时SDPA 走仓库自带的 cubin 内核gen_cudnn_fmha_module见 flashinfer/cudnn/utils.pyDecode仅 bf16、仅单 tokenq_len_per_req 1、window_left、sinks均抛NotImplementedError、不支持return_lseactual_seq_lens_kv可以传 CPU 张量内部拷贝到 GPU。Prefill强制return_lseTrue仅支持d_qk 192 且非 Paged或d_qk 128 且 Paged两种组合必须提供actual_seq_lens_q/kv仅返回 base-2 LSE。SM100 线性注意力引擎这是flashinfer.cudnn的第二个技术面cuDNN 的graph.gdn/graph.gdn2/graph.gdp/graph.kda图节点。它们只被一个条件门控cudnn-frontend 1.29 且带cutedslextra。除此之外没有任何 cuDNN backend 版本下限——FROST 引擎是 CuTeDSL 内核由 frontend 自行编译。其余所有约束SM100 家族、head dim、输入 dtype、头数关系都归属引擎自身引擎会拒绝它无法服务的图具体拒绝原因记录在 frontend 日志中而抛出的cudnnGraphNotSupportedError本身是通用的。四个入口的参数语义与其 FlashInfer 原生版本一一对应chunk_gated_delta_rule、chunk_gated_delta_rule2、chunk_gated_delta_product、recurrent_kda共性包括cu_seqlens必填[num_seqs 1]的 int32/int64 打包序列边界。state 布局[num_seqs, num_sab_heads, head_size, head_size]V-majorfp32 或 bf16output_state不得与initial_state别名引擎会把一个序列拆分到多个 CTAchunk-0 CTA 读入态会与末 chunk CTA 写出的态竞争。batch_invariant禁用 split-K 分区使归约顺序从而结果不依赖序列如何分批代价是长序列上失去 split-K 的并行度换来短序列上省去固定调度开销。use_qk_l2norm_in_kernel把 q/k 的 L2 归一化融合进内核。scale默认1/sqrt(head_dim)与原生路径一致。output_final_state仅控制返回值而 KDA 遵循 Cake / CuTe DSL Prefill 后端的惯例——只要给了initial_state且未单独提供output_state就无条件把 state 推进到 final与output_final_state无关后者只决定返回与否。gate domain线性域与对数域遗忘门的跨边界表示是双方约定的一部分GDN 与 GDP的 forget gate 在此边界上是线性空间 alphaalpha exp(log_g)以gate_domainlinear跨入 cuDNNg为None时全 1。GDN-2 与 KDA的 gate 在两侧都是对数空间GDN-2 的g是逐通道自然对数衰减gate_domainlogNone时全 0KDA 的g默认也是 log 空间只有use_gate_in_kernelTrue时才作为原始 pre-activation 交给 cuDNN 用A_log/dt_bias/lower_bound做 safe-gate 变换此时A_log与dt_bias必须同时提供。g张量按自身 dtype 直接透传fp32 / bf16 / fp16 均可cuDNN 三种都读只在 KDA 的 fp16 场景下存在一个数值注意点内核的 chunk 累积衰减逆运算会约束衰减强度大约每个 token 每通道alpha 0.9否则溢出bf16 携带 fp32 量级的指数没有此约束。beta默认全 1 且按语义区分——GDN/GDP 是后 sigmoid 的更新门fp32 或q.dtypeKDA 还支持beta_is_logitTrue时在核内做 sigmoid。状态张量跨边界零拷贝循环状态在跨边界时不做转置FlashInfer 以 V-major 持有[N, H, V, K]cuDNN 也是同样的布局因此initial_state/output_state缓冲直接透传。cuDNN 的算子以 fp32 或 bf16 接收状态并按输入精度返回final_state——因此bf16 状态池跨边界完全零拷贝。被路由层拒绝的参数FlashInfer 有而 cuDNN 入口没有的参数——状态检查点state checkpointing、索引状态池indexed state pools、上下文并行 delta rule、推测解码speculative decode——会在调用前被路由层直接拒绝不会到达 cuDNN。KDA 还有一条头数约束cuDNN 以max(H, HV)个头承载 KDA 状态而 FlashInfer 以HV个因此当H HV时抛NotImplementedError见 flashinfer/cudnn/linear_attention.py。GDP 的 Householder 展开GDP 是 GDN 递推在展开子 token 时间线上的推广每个 token 应用n num_householder次 beta 门控的 Householder 更新token 的衰减作用在更新之前读出跟随最后一次更新。因此k/v的形状是[total_seq_len * num_householder, num_heads, head_size]token t 的 n 次更新占据行t*n .. t*nn-1g与cu_seqlens仍按真实 token 行给出beta按展开时间线给出。num_householder 1时 GDP 恰好等价于 GDN。限制与注意事项汇总综合文档与源码cuDNN 后端的功能边界如下Decode 不支持RoPE、soft-caplogits 软上限、fp8/NVFP4 KVsinks 仅在q_len_per_req 1且 cuDNN 栈frontend 1.30 / FROST 引擎支持时可用。cuDNN Paged Prefill 不支持sinks、NVFP4 KV、skip-softmax见 flashinfer/prefill.py。cuDNN 单 token GQA Prefill在return_lseTrue且q_len_per_req 1时不会写出每个 LSE 头NVBug 6783545wrapper 会显式抛NotImplementedError提示换后端或关闭 LSEflashinfer/prefill.py。cuDNN Prefill 要求out/lse连续KV 布局若为NHD会先转置为 cuDNN 期望的形式flashinfer/prefill.py。所有绑定到图的张量含 workspace必须与 q 同设备且每次重绑定时都会重新校验设备图元数据匹配并不能保证运行时设备一致。线性注意力全部入口都要求 SM100 家族SM100–SM103、SM107以及引擎可接受的 head dim / dtype / 头数关系否则引擎拒绝服务原因在 frontend 日志。源码指引cuDNN 模块入口与导出flashinfer/cudnn/init.pyDecode 图构建、缓存键与 LSE 换算flashinfer/cudnn/decode.pyPrefill 图构建、direct seqlens 版本探测、shape override 与 cubin 回退flashinfer/cudnn/prefill.py线性注意力四个入口、gate domain 与状态透传flashinfer/cudnn/linear_attention.py每设备 handle 缓存与 cubin 加载flashinfer/cudnn/utils.pywrapper 侧的 cuDNN 分支backend 选择、参数校验、计划复用flashinfer/prefill.py、flashinfer/decode.pycuDNN 后端的 trace 模板与参考实现flashinfer/trace/templates/attention.py若需在自有环境中验证版本兼容性可参考 flashinfer/collect_env.py 对libcudnn的环境探测方式要深入 FROST 引擎侧的 CuTeDSL 内核形态可阅读 flashinfer/experimental/cudnn_frost_selected_kernels_moe_grouped_gemm/README.md 了解其生成内核的组织方式。赞分享大模型深度学习算子库后端高性能计算【免费下载链接】flashinferFlashInfer: Kernel Library for LLM Serving项目地址https://gitcode.com/gh_mirrors/fl/flashinfer点击查看免费下载相关推荐深入解析MMCV视觉注意力机制从空间注意力到通道注意力的完整指南深入解析MMCV视觉注意力机制从空间注意力到通道注意力的完整指南 MMCV作为OpenMMLab计算机视觉基础库提供了强大的视觉注意力机制实现。本文将为您全人工智能计算机视觉深度学习PyTorch torch.nn.attention 模块全指南SDPA 后端控制、FlashAttention 注册与 FlexAttention 可编程注意力PyTorch torch.nn.attention 模块全指南SDPA 后端控制、FlashAttention 注册与 FlexAttention 可编程注人工智能机器学习深度学习分布式训练模型编译Mini-SGLang注意力后端如何快速配置FlashAttention与FlashInferMini SGLang注意力后端如何快速配置FlashAttention与FlashInfer Mini SGLang是一个轻量级高性能的大语言模型推理框架大模型人工智能推理引擎模型推理服务本地部署上一篇如何快速将B站视频转为文字bili2text工具完整指南下一篇如何用3个简单步骤将B站视频变成可搜索的文字稿创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考