这是一份针对 EchoLLMServing 推理引擎架构及核心算子优化的工程笔记梳理。内容已针对技术深度、架构权衡（Trade-off）以及实现细节进行了重构，，并补充了底层逻辑分析。

EchoLLMServing 推理框架笔记

1. 架构设计与分布式执行逻辑

当前项目采用了多进程架构支持张量并行（Tensor Parallelism），控制流与数据流分离：

  • 控制平面：Rank 0 作为主节点（LLMEngine），负责请求接收、调度（Scheduler）和元数据管理（可结合 Gloo 进行跨进程协调）。
  • 数据平面：副卡初始化 ModelRunner 并阻塞于共享内存（shm）或类似 ZeroMQ 的异步队列等待指令。
  • 执行流：Rank 0 在 generate 中调用 run，将计算元信息写入 shm，触发所有 Rank 并行执行前向传播。在 LMHead 阶段，各卡计算出局部的 logits 后，通过 All-Reduce 或 Gather 汇总到 Rank 0，由采样器（Sampler）根据策略输出 Token。

架构审视：这种主从架构简化了调度的复杂度，但在大规模并发下，Rank 0 的调度开销可能成为瓶颈。引入 CUDA Graph 是降低调度开销的有效手段，通过捕获算子执行序列，能够显著减少 Kernel Launch Overhead。

2. 核心算子实现与优化 (@Triton)

FlashAttention

  • 定位：解决 Prefill 阶段的 Compute Bound 问题。
  • 实现逻辑：对输入 Q 的 seq_len 按 BLOCK_M 分块，对外层 Head 分块。在 Triton program 内处理 tl.arange(0, head_dim) 的向量化计算。
  • Trade-off：虽然极大地提升了计算密度，但对 SRAM 容量要求高。BLOCK_M 的大小受限于硬件的 Shared Memory 上限。

PagedAttention (V1 -> V2 演进)

  • V1 的问题：最初尝试将 context 中的全部 KV 取出再计算，导致严重的 Memory Bound。显存带宽被非连续的访存彻底打满，吞吐极低。
  • V2 的改进：将 context_len 按 BLOCK_M 划分，对 KV 矩阵分块并行计算 QK，进行 Online Softmax 增量更新。减少了不必要的显存搬运，将数据尽可能留在 SRAM 中处理。

FlashDecoding (长序列 Decode 优化)

在长上下文 Decode 时，传统的 PagedAttention 仍面临带宽瓶颈。引入 Split-K 机制与双 Kernel 架构优化：

  • OOM 与 Reduce 代价分析：最初将 split size 固定为 64，导致每个 block 都会产生临时结果（m, l, acc）。如果并行度过高且 BLOCK_N 很小，申请的 Global Memory 空间会急剧膨胀导致 OOM，且最终的 Reduce Kernel 需要遍历巨大的临时数组，抵消了并行带来的收益。
  • 解决方案：将 seq_len 按照预设的 split 划分为几个 chunk，每个 program 处理一个 split 内的所有 KV Cache。

针对笔记中遗留问题的解答：

  • 问题：Split_num 是否可以根据 context_len 动态调整？
    • 解答：必须动态调整。固定的 split_num 会导致短序列的 kernel 切分过细（增加 kernel 启动和 reduce 开销），而长序列切分不足（无法充分利用 SM）。通常的启发式策略是：split_num = max(1, context_len // optimal_chunk_size)，其中 optimal_chunk_size 根据 GPU 的 SM 数量和 Cache 大小在 Profiling 中动态标定（例如 256 或 512）。
  • 问题：调度时是否优先调度上下文长度均匀的 Seq 以避免 Kernel 空转？
    • 解答：是的。在未实现完全 Unified Attention 的情况下，同 Batch 内序列长度差异过大会导致严重的 Padding 计算浪费或线程束（Warp）分化。通过持续批处理（Continuous Batching）结合细粒度的 Block 管理，可以最大程度缓解这个问题。

Fused Add-RMSNorm 与 Triton 拓展

  • 问题分析：使用 PyTorch 原生的 x = x + residual 然后做 RMSNorm 时，发现总吞吐下降，通过 Nsight Systems (nsys) 分析证实，内存拷贝与读写开销（Memory Access）是主因。
  • Triton 实现权衡：
    • 写融合算子，将读、Add、写 Res、Norm 全部在一个 Kernel 内完成。
    • 关于分块与 Kernel Launch Overhead：对输入向量的最后一个维度（hidden_dim）进行处理。如果 hidden_dim 过大（如 Llama 的 4096），可以将其拆分为几个 Block。并行度高导致的 Launch Overhead 可以通过 CUDA Graph 解决。关键在于利用 tl.load 时的向量化读写（保证首地址对齐）。

3. 调度、显存管理与特性支持

异步推理与 Continuous Batching

AsyncLLMEngine 通过 asyncio.queue 实现了请求的非阻塞生成。Continuous Batching 的核心价值在于打破了静态 Batch 的生命周期绑定，每次 Step 后立即回收完成状态的 Seq 所占用的 Block 资源，并打入新请求，极大提升了吞吐。

Chunked Prefill 调度规则

  • 核心动机：解决长短 Prompt 混杂时的 Time to First Token (TTFT) 延迟毛刺，以及分离阶段导致 Decode 计算资源利用率低的问题。
  • Block Manager 状态不一致问题：Chunked Prefill 会导致 num_cached_tokens 与实际输入的 input_ids 长度在单次 forward 中不匹配，必须在 post-process 中进行增量状态维护。

针对笔记中遗留问题的解答：

  • 问题：混合 Batch 时，分开做两种 Attn 和合成一种（Unified Attention）有什么区别？
    • 解答：
      1. 分开做（分离 Kernel）：工程实现简单。Prefill 的请求走 FlashAttention，Decode 的请求走 FlashDecoding。缺点是 Batch 较小时，两个 Kernel 依次 Launch，可能无法跑满 GPU 的 SM，存在硬件资源的波谷。
      2. 合成一种（Unified Attention）：将 Decode 视为长度为 1 的特殊 Prefill 阶段，使用一套底层机制处理。这要求极高的指针与索引管理技巧（类似 FlashInfer 的设计）。优点是只需 Launch 一个 Kernel，SM 占用率极高；缺点是分支预测和 Warp 分化处理非常复杂。前期建议分开做，后期追求极限性能时再考虑 Unified。

KV Cache INT8 量化

  • 工程经验：最初全量反量化（Dequantization）导致严重的显存读写墙，耗时极高。
  • 正确路径：将 Per-token 反量化逻辑直接 Fuse 进 FlashAttention/FlashDecoding Kernel 内。在读取 INT8 Cache 到 SRAM 后，立即使用伴随的 Scale 因子还原为 FP16 进行计算。这一操作以极小的计算开销（Compute）换取了 50% 的显存带宽（Bandwidth）与容量收益。

4. 底层架构思考 (Thinking & Trade-offs)

Decode 阶段 复杂度缓解策略

Decode 本质上是 Memory Bound，复杂度随上下文线性增长。常见缓解方案：

  • 算法层：Sliding Window Attention (SWA)、KV 压缩/池化（如 H2O）、稀疏注意力。
  • 系统层：前缀缓存（RadixTree LRU Prefix Cache），将公共的前缀（如 System Prompt）复用，避免重复计算并节省 KV Cache 空间。

为什么必须用 Triton/CUDA 代替 PyTorch 张量广播？

PyTorch 的张量广播发生在逻辑层和 Global Memory 层面。它会产生大量隐式的中间张量，导致 HBM 被反复读写（Memory Wall）。Triton 和 CUDA 的核心优势在于显存层级控制：将高复杂度的 计算约束在 L1/L2 Cache 或 SRAM 中，将对外部 HBM 的访问降低到 级别。在算力远大于带宽的 Ampere/Hopper 架构上，这是唯一解。

还有什么值得用 Triton 优化的算子？

  • Fused RoPE + QKV 投影：将旋转位置编码直接融合进 QKV 的 Linear 计算之后，避免显存的来回搬运。
  • Fused SwiGLU (或 GeLU)：激活函数通常是单纯的 element-wise 操作，典型的 Memory Bound，非常适合融合。
  • MoE Top-K Routing：如果未来支持 Qwen3-MoE，专家路由的 Top-K 计算、Softmax 及 token 重排，使用 Triton 可以大幅降低开销。