SGLang 算子源码分析

飞书用户2030

2024年2月20日修改

💡

SGLang 源码分析按顺序：

1.
​SGLang 后端源码分析 ​

2.
​SGLang 算子源码分析（建设中） ​

3.
​SGLang 前端源码分析（建设中） ​

基础

•
必学 Triton，暂不考虑 SGLang 中 Flashinfer 相关的实现​
◦
参考我的 Triton 学习笔记：​Triton ​

•

必学 KVCache

◦
需要理解推理过程有 Prefill （代码里为 Extend）和 Decode 两个阶段，并且理解 MHA、MQA 和 GQA​

◦

推荐：NLP（二十）：漫谈 KV Cache 优化方法，深度理解 StreamingLLM，里面第一部分的动图很好，可以比较一下 no-cache 和 cached 的两种情况，说明了为什么 Prefill 阶段需要 attention_mask。

•

可学 FlashAttention

◦
参考学习笔记：​2024.1.25 FlashAttention ​

◦
主要使用 Online Safe Softmax​

common.docs_name - LarkCCM_Docs_Menu_Image

Llama2

模型架构

SGLang 模型架构基本和 Vllm 一样，且基本使用 Vllm 的算子，其自己主要实现的是 RadixAttention，我们主要分析一下 RadixAttention。其他的模型部分请参考 ​VLLM 算子源码分析（建设中）​

RadixAttention

根据 ForwardMode 有 EXTEND 和 DECODE 两种情况（PREFILL 暂时没有被调用过）。​

Extend

Extend Attention 和普通 Prefill 阶段的 Attention 算子的区别：

1.
在 Extend 阶段（相当于正常情况的 Prefill 阶段），一条 Request 的 Prefix 部分已经有 KV Cache，保存在 K_Buffer/V_Buffer 中（之后 Extend 的部分计算好 KV 之后也会保存到里面）​

2.
 Extend 阶段输入的 input_ids 只有 extend 部分的 Token，所以在计算 Attention 时，输入值是 extend 部分的 token 的 QKV（Q_Extend/K_Extend/V_Extend），prefix 部分的 token 会复用 K_Buffer/V_Buffer。同时输入的 input_ids 是 batch 内所有 request flatten 得到的，所以 request 的长度不同没有影响。​

3.
Extend 阶段的输出也是只有 extend 部分的 Token 对应的 hidden_states​

推导：

SGLang 算子源码分析​

SGLang 算子源码分析