CctoctoFX (Yuxiao Huang)

本文从数学原理出发，深入分析FlashAttention的核心思想、算法设计和各版本演进，通过详实的数学推导、直观的流程图表和具体的数值示例，帮助读者真正掌握这一革命性的Attention优化技术。 1. 问题的本质：传统Attention的根本瓶颈 1.1 传统Attention机制的计算模式 传统的Self-Attention机制遵循如下计算流程： $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$ 让我们用具体数值来理解这个过程的复杂性： 示例场景：考虑一个典型的语言模型场景 序列长度：$n = 2048$（如GPT-2的上下文长度） 特征维度：$d_k = 64$（每个attention head的维度） 输入张量形状：$Q, K, V \in \mathbb{R}^{2048 \times 64}$ 第一步：计算注意力得分矩阵 $$S = \frac{QK^T}{\sqrt{d_k}} \in \mathbb{R}^{2048 \times 2048}$$ 这一步产生了一个 $2048 \times 2048 = 4,194,304$ 个元素的矩阵，以FP16精度存储需要约8MB内存。 第二步：Softmax归一化 $$P = \text{softmax}(S) \in \mathbb{R}^{2048 \times 2048}$$ Softmax计算需要： 计算每行的最大值：$m_i = \max_j S_{i,j}$ 计算指数和：$l_i = \sum_j e^{S_{i,j} - m_i}$ 归一化：$P_{i,j} = \frac{e^{S_{i,j} - m_i}}{l_i}$ 这又需要存储另一个 $2048 \times 2048$ 的矩阵。 ...

Verl DataProto 实现原理与数据流动分析 目录 1. 概述 2. DataProto 核心架构 3. HybridFlow 设计理念 4. 控制流与计算流分离 5. 数据流动机制 6. Dispatch 模式详解 7. 性能优化策略 8. 总结 1. 概述 Verl 是一个基于 HybridFlow 论文的开源强化学习训练框架，专门为大语言模型的后训练优化而设计。其核心创新在于将控制流和计算流分离，通过 DataProto 协议实现高效的数据交换。 2. DataProto 核心架构 2.1 数据结构设计 DataProto 是 verl 框架中用于数据交换的核心协议，所有在 Worker 之间流转的数据，都被统一封装在一个名为 DataProto 的数据结构中。它不仅仅是一个字典，更承载着 RLHF 流程中所有的信息演变, 基于 PyTorch 的 TensorDict 构建： 1 2 3 4 5 @dataclass class DataProto: batch: TensorDict = None # 张量数据容器 non_tensor_batch: dict = field(default_factory=dict) # 非张量数据 meta_info: dict = field(default_factory=dict) # 元信息 核心特性： 统一接口: 提供标准化的数据容器，支持张量和非张量数据 设备管理: 自动处理 GPU/CPU 设备间的数据移动 内存优化: 支持分块处理和内存复用 序列化: 支持高效的序列化和反序列化 2.2 数据一致性检查 1 2 3 4 5 6 7 8 9 10 11 12 13 14 def check_consistency(self): """检查 DataProto 的一致性""" if self.batch is not None: assert len(self.batch.batch_size) == 1, "只支持 num_batch_dims=1" if self.non_tensor_batch is not None: for key, val in self.non_tensor_batch.items(): assert isinstance(val, np.ndarray) # 检查批次大小一致性 if self.batch is not None and self.non_tensor_batch is not None: batch_size = self.batch.batch_size[0] for key, val in self.non_tensor_batch.items(): assert val.shape[0] == batch_size 3. HybridFlow 设计理念 3.1 设计动机 传统 RL 系统面临的问题： ...

最近 RL sys 圈子的吴锡斌老师在 verl 上设计了将 rollout 与 tool 调用解耦的 AgentLoop，实现了自由灵活的 mutli-turn RL。在每个 AgentLoop 内部，rollout engine 只对外提供一个 token-in-token-out 的接口，而 tool 调用则通过 ToolAgentLoop 来实现。我个人比较喜欢这样解耦的设计，同时，AgentLoop 的代码结构也比较清晰。我个人学习了一次整个代码后，觉着 AgentLoop 的设计甚是不错，但是 ActorRolloutRefWorker 的历史包袱还是很重。 本文简单分析了 agent loop 的源码，并给出了一些自己的看法。 如果我们把整个 ActorRolloutRefWorker 当做一个 sgl.Engine 的话，AgentLoop 里面包装的两层 AsyncSGLangServer 和 AsyncLLMServerManager。AsyncSGLangServer 相当于在 sgl.Engine 上包装了 fastapi 成了 server，而 AsyncLLMServerManager 是在 server 上包了一层 router 做 load balance，相当于 sglang 的 router。这两层设计都是合理的，主要麻烦的是 ActorRolloutRefWorker，层层调用，最后一共经过 7 个 class 才调到 sgl.Engine，最近 verl 团队也在致力于对这块 worker class 的重构，敬请期待。最后，AgentLoopManager，AgentLoopWorker 和 AgentLoop 这三层，我觉得 AgentLoopWorker 可能未必有必要，其他两层挺合理的。 ...

VeRL框架的参数众多，基于当前（2025.8.5）主线分支整理，附带了相关的理解，一些描述不一定完全正确，供学习参考。 Batch Size 参数名称 详细解释 data.train_batch_size 作用：定义了单次训练发送给 Rollout Engine 的样本数量，也即这是在每个 PPO 迭代开始时，从训练数据集中采样的提示 （Prompt）数量。 详细解释：这个值是 RL 训练中的基本样本数量。例如，设置为 1024 意味着在一次迭代中会： 1. 从数据集中随机抽取 1024 个 prompt。 2. 将这 1024 个 prompt 发送给当前的 Rollout Engine 中，从而得到 1024 组完整的 trajectories（prompt, response）。 3. 接下来，这 1024 个 trajectories 进行经验计算（make experience），后续用于 Actor 和 Critic 模型的更新。 影响与权衡：影响总共训练的样本量。 data.val_batch_size （Deprecated) 作用：在 Validation 阶段使用的批次大小。 详细解释：这与 train_batch_size 类似，但仅用于评估模型性能，不参与训练。如果设置为 null，会使用验证集的大小作为默认值。Note: 已经deprecated，推荐设置为 null。此时，整个 validation dataset 一次性发给 SGLang engines，自行进行内存管理。 actor_rollout_ref.actor.ppo_mini_batch_size critic.ppo_mini_batch_size 作用：定义了 PPO 训练更新中的 mini-batch 大小。 详细解释：data.train_batch_size 收集到的全部经验数据将被分割成多个 mini-batch，每块的大小就是 ppo_mini_batch_size。模型每处理完一个 mini-batch，才会进行一次参数更新。 例如，如果 train_batch_size = 1024，ppo_mini_batch_size = 256，那么在一个 PPO Epoch 中，模型会进行 1024 / 256 = 4 次参数更新。 影响与权衡：增大 mini-batch，单次更新的梯度更稳定，但更新频率更低，更新次数减少。 actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu critic.ppo_micro_batch_size_per_gpu 作用：定义了在单个 GPU 上进行一次 forward/backward 的数据大小。 详细解释：这是实现梯度累积的核心参数。mini-batch 会被再次切分为若干个 micro-batch。例如，在单卡上，ppo_mini_batch_size = 256，ppo_micro_batch_size_per_gpu = 32，那么梯度累积的步数就是 256 / 32 = 8。这意味着模型会运行 8 次 forward 得到 loss，然后 backward 的到 gradient。每次处理 32 个样本，直到累积完整个 mini-batch 计算出的梯度。此时，使用累积的总梯度，对模型参数进行一次更新（optimizer.step()）。这个值必须根据显存大小来严格调整，是防止 OOM 的关键。 影响与权衡：增大此值，减少了梯度累积的次数，可以提高训练的吞吐量，增大显存消耗。 actor_rollout_ref.actor.ppo_micro_batch_size critic.ppo_micro_batch_size（Deprecated) 作用：已弃用，被 per_gpu 版本取代，因为它能更好地适应分布式训练环境。 Dynamic Batch Size 当样本长度差异很大时，按样本数量划分批次可能导致不同批次的计算量极不均衡，而基于 token 总数来控制 batch size 是一种平衡每个 batch 训练时间的方案。 ...

本文档为开发者提供 SGLang 后端代码的代码梳理，按照一个请求从输入到最后输出的顺序进行讲解。下图简要介绍了这一流程： 具体而言，请求的处理过程如下： 用户启动 Server ，初始化 FastAPI App、TokenizerManager、DetokenizerManager 和 Scheduler，每个组件运行各自的无限事件循环（infinite event loop）。 用户向 FastAPI Server 发送 /v1/chat/completions 请求，Server 通过 v1_chat_completions endpoint 将请求转发到 TokenizerManager。 v1_chat_completions 函数将请求转换为 ChatCompletionRequest，再转换为 GenerateReqInput，并调用 TokenizerManager 的 generate_request 方法。 TokenizerManager 对请求进行 tokenization，并以 Python 对象（pyobj）形式将其转发给 Scheduler，同时调用 TokenizerManager 的 _wait_one_response 方法。 Scheduler 在事件循环 event_loop_normal 中处理请求： Scheduler 通过 recv_requests 接收请求，调用 process_input_requests 处理输入，通过 handle_generate_request 管理生成请求的逻辑，并将其加入 waiting_queue。 从 waiting_queue 中，Scheduler 使用 get_next_batch_to_run 为即将处理的请求创建 ScheduleBatch。 Scheduler 执行 run_batch 函数，将 ScheduleBatch 转换为 ModelWorkerBatch。 Scheduler 调用 TpModelWorker 的 forward_batch_generation，等待 logits_output 和 next_token_ids。 TpModelWorker 初始化 ForwardBatch，将其转发至 ModelRunner，并等待 logits_output。 ModelRunner 处理 ForwardBatch，调用 forward_extend 执行模型的前向计算（forward pass）。 模型通过 AttentionBackend 加速生成 logits，返回给 ModelRunner，进而返回给 TpModelWorker。 TpModelWorker 从 ModelRunner 接收 logits_output，调用 ModelRunner 的 sample 方法生成 next_token_ids，并将其发送回 Scheduler。 Scheduler 通过 process_batch_result 处理批次结果，使用 tree_cache.cache_finished_req(req) 缓存请求，并通过 check_finished 验证完成状态。对于未完成的请求，Scheduler 继续其事件循环，直到这个请求满足结束条件；对于已完成的请求，则转发到 Scheduler 的 stream_output。 在 stream_output 函数中，Scheduler 处理输出，将其包装成 BatchTokenIDOut，并发送给 DetokenizerManager。 DetokenizerManager 在其事件循环中接收 BatchTokenIDOut，处理后生成 BatchStrOut 并返回给 TokenizerManager。 ...