SGLang Notes(主要是 Awesome-ML-SYS-Tutorial Notes)

• notes 主要来自：https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/README-cn.md
• 以及阅读中拓展出来的一些东西

Radix Attention

• 参考:https://github.com/CalvinXKY/InfraTech/blob/main/llm_infer/sglang_radix_attention.ipynb
• 加了 kv cache 的 attentin
• 用 radix tree 来压缩 prefix，有点像 trie
• 引用计数法来做生命周期管理

ML System 基本功

• torch-memory-savor
  • cuda graph,可以手动把若干个 torch 指令编成 cuda graph，这样减少 gpu<->cpu 的指令传输（否则对于复杂拓扑如 fsdp 可能有 cpu bottleneck），整体思路和 bigdata 里 spark/ray 把固定的 operator 编排做成 dag 是类似的原理；和 torch.compile 有点像，但是 torch.compile 更加面向开发者，提供通用的性能加速而非 cuda。
  • torch-memory-savor 是 sglang 的 cuda 内存管理（为了更好的维护显存从而保证 cuda graph 可用），调用 cuda 底层 api 来实现「逻辑内存」和「物理内存」的映射和管理。有 pause/resume/malloc 等方法。这里没太理解为什么要频繁释放物理内存，直接申请一大块出来反复用，做 gc 不是更好？
• nccl(nvidia collective communication library)
  • AllReduce：所有 GPU 的数据先进行规约（如求和），然后广播到所有 GPU,
  • Broadcast：从一个源 GPU 向其他所有 GPU 广播数据
  • Reduce：将所有 GPU 的数据规约到一个目标 GPU
  • AllGather: 收集所有 GPU 的数据并分发给所有 GPU, dim * world_size
  • ReduceScatter: 规约后将结果分散到所有 GPU
  • Send/Recv: 点对点通信
  • AllToAll: 将数据分发到所有 GPU
  • send/recv
  • isend/irecv/batch_isend_irecv: 异步的，要调用 wait
• 通信算法，ring algo vs tree algo
  • ring: allreduce 通信次数要 2*(n-1) (n-1 次求和 + n-1 次广播)
  • tree: allreduce 通信次数 2*log(N)
  • 从 NCCL 2.4 版本开始，对于node数较多的跨节点通信使用 Double Binary Tree 算法。相比于传统的 Tree Algorithm，构造了两棵互补的二叉树用于平衡通信开销。
• rl-memory-management(https://hebiao064.github.io/rl-memory-management)
  • 还没看，等学一下 rl 再看