sgl-project/sglang

[Kernel] cuDNN attention backend

Open

#2,272 创建于 2024年11月30日

在 GitHub 查看
 (5 评论) (1 反应) (2 负责人)Python (6,216 fork)auto 404
enhancementgood first issuehelp wantedhigh priorityinactive

仓库指标

Star
 (28,442 star)
PR 合并指标
 (平均合并 2天 1小时) (30 天内合并 1,000 个 PR)

描述

cuDNN provides very fast attention implementation and it is well maintained by NVIDIA. We would like to add a new attention backend based on cudnn.

Steps

  1. Learn this cudnn paged attention python api. https://github.com/NVIDIA/cudnn-frontend/blob/v1.8.0/samples/python/52_scaled_dot_product_attention_with_paged_caches.ipynb
  2. Add a new attention backend "cudnn" here https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/layers/attention
  3. We should be able to use it with python3 -m sglang.launch_server --model meta-llama/Llama-3.1-8B-Instruct --attention-backend cudnn

贡献者指南

研究方向
研究 sglang/srt/layers/attention 中现有的注意力后端，理解其接口，并探索步骤 1 中链接的 cuDNN 分页注意力 Python API。然后按照现有后端的模式实现一个新的后端。
技术栈
python
领域
backendmachine learning
议题类型
功能
难度面向新贡献者的预计实现难度，1 表示很小改动，5 表示专家级工作。
3
预计时间有经验贡献者完成调查、实现、测试并准备 pull request 的粗略时间范围。
1-2 天
活动状态议题当前的可参与程度：新鲜、活跃、陈旧、阻塞或等待维护者输入。
活跃
清晰度议题是否清楚说明期望改动、验收标准和下一步。
清晰
前置要求
PythonCUDAcuDNN
新手友好度1-100 的估计分数，表示该议题对首次贡献者的友好程度。
60